
.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/01_deterministic_workflow.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_01_deterministic_workflow.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_01_deterministic_workflow.py:


Running Deterministic Inference
===============================

Basic deterministic inference workflow.

This example will demonstrate how to run a simple inference workflow to generate a
basic determinstic forecast using one of the built in models of Earth-2 Inference
Studio.

In this example you will learn:

- How to instantiate a built in prognostic model
- Creating a data source and IO object
- Running a simple built in workflow
- Post-processing results

.. GENERATED FROM PYTHON SOURCE LINES 35-42

.. code-block:: Python

    # /// script
    # dependencies = [
    #   "earth2studio[dlwp] @ git+https://github.com/NVIDIA/earth2studio.git",
    #   "cartopy",
    # ]
    # ///


.. GENERATED FROM PYTHON SOURCE LINES 43-48

Set Up
------
All workflows inside Earth2Studio require constructed components to be
handed to them. In this example, let's take a look at the most basic:
:py:meth:`earth2studio.run.deterministic`.

.. GENERATED FROM PYTHON SOURCE LINES 50-54

.. literalinclude:: ../../earth2studio/run.py
   :language: python
   :start-after: # sphinx - deterministic start
   :end-before: # sphinx - deterministic end

.. GENERATED FROM PYTHON SOURCE LINES 56-61

Thus, we need the following:

- Prognostic Model: Use the built in FourCastNet Model :py:class:`earth2studio.models.px.FCN`.
- Datasource: Pull data from the GFS data api :py:class:`earth2studio.data.GFS`.
- IO Backend: Let's save the outputs into a Zarr store :py:class:`earth2studio.io.ZarrBackend`.

.. GENERATED FROM PYTHON SOURCE LINES 63-84

.. code-block:: Python

    import os

    os.makedirs("outputs", exist_ok=True)
    from dotenv import load_dotenv

    load_dotenv()  # TODO: make common example prep function

    from earth2studio.data import GFS
    from earth2studio.io import ZarrBackend
    from earth2studio.models.px import DLWP

    # Load the default model package which downloads the check point from NGC
    package = DLWP.load_default_package()
    model = DLWP.load_model(package)

    # Create the data source
    data = GFS()

    # Create the IO handler, store in memory
    io = ZarrBackend()


.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Downloading dlwp_cubesphere.zip: 0%|          | 0.00/67.2M [00:00<?, ?B/s]    Downloading dlwp_cubesphere.zip: 57%|█████▋    | 38.3M/67.2M [00:00<00:00, 401MB/s]    Downloading dlwp_cubesphere.zip: 100%|██████████| 67.2M/67.2M [00:00<00:00, 482MB/s]


.. GENERATED FROM PYTHON SOURCE LINES 85-94

Execute the Workflow
--------------------
With all components initialized, running the workflow is a single line of Python code.
Workflow will return the provided IO object back to the user, which can be used to
then post process. Some have additional APIs that can be handy for post-processing or
saving to file. Check the API docs for more information.

For the forecast we will predict for two days (these will get executed as a batch) for
20 forecast steps which is 5 days.

.. GENERATED FROM PYTHON SOURCE LINES 96-103

.. code-block:: Python

    import earth2studio.run as run

    nsteps = 20
    io = run.deterministic(["2024-01-01"], nsteps, model, data, io)

    print(io.root.tree())


.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    2025-11-12 18:27:49.993 | INFO     | earth2studio.run:deterministic:75 - Running simple workflow!
    2025-11-12 18:27:49.993 | INFO     | earth2studio.run:deterministic:82 - Inference device: cuda
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]                                                            2025-11-12 18:27:50.406 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20231231/18/atmos/gfs.t18z.pgrb2.0p25.f000 408062467-879185
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]                                                            2025-11-12 18:27:50.409 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20231231/18/atmos/gfs.t18z.pgrb2.0p25.f000 397402829-996456
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]                                                            2025-11-12 18:27:50.411 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20231231/18/atmos/gfs.t18z.pgrb2.0p25.f000 294691465-856457
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]                                                            2025-11-12 18:27:50.413 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20231231/18/atmos/gfs.t18z.pgrb2.0p25.f000 208052937-721817
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]                                                            2025-11-12 18:27:50.415 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20231231/18/atmos/gfs.t18z.pgrb2.0p25.f000 420029701-1181204
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]                                                            2025-11-12 18:27:50.416 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20231231/18/atmos/gfs.t18z.pgrb2.0p25.f000 251230645-803982
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]                                                            2025-11-12 18:27:50.418 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20231231/18/atmos/gfs.t18z.pgrb2.0p25.f000 329116923-847018
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]    Fetching GFS data:  14%|█▍        | 1/7 [00:01<00:07,  1.22s/it]    Fetching GFS data:  29%|██▊       | 2/7 [00:01<00:02,  1.67it/s]    Fetching GFS data: 100%|██████████| 7/7 [00:01<00:00,  5.05it/s]
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]                                                            2025-11-12 18:27:52.219 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20240101/00/atmos/gfs.t00z.pgrb2.0p25.f000 246334297-805355
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]                                                            2025-11-12 18:27:52.222 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20240101/00/atmos/gfs.t00z.pgrb2.0p25.f000 391722290-987401
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]                                                            2025-11-12 18:27:52.224 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20240101/00/atmos/gfs.t00z.pgrb2.0p25.f000 204118947-720169
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]                                                            2025-11-12 18:27:52.226 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20240101/00/atmos/gfs.t00z.pgrb2.0p25.f000 414179964-1179422
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]                                                            2025-11-12 18:27:52.228 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20240101/00/atmos/gfs.t00z.pgrb2.0p25.f000 289307267-851916
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]                                                            2025-11-12 18:27:52.230 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20240101/00/atmos/gfs.t00z.pgrb2.0p25.f000 402321768-876246
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]                                                            2025-11-12 18:27:52.232 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20240101/00/atmos/gfs.t00z.pgrb2.0p25.f000 323956279-837771
    Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]    Fetching GFS data:  14%|█▍        | 1/7 [00:00<00:03,  1.70it/s]    Fetching GFS data:  29%|██▊       | 2/7 [00:00<00:01,  3.19it/s]    Fetching GFS data:  57%|█████▋    | 4/7 [00:00<00:00,  6.43it/s]    Fetching GFS data: 100%|██████████| 7/7 [00:00<00:00,  8.18it/s]
    2025-11-12 18:27:53.182 | SUCCESS  | earth2studio.run:deterministic:106 - Fetched data from GFS
    2025-11-12 18:27:53.200 | INFO     | earth2studio.run:deterministic:136 - Inference starting!

    Running inference:   0%|          | 0/21 [00:00<?, ?it/s]
    Running inference:   5%|▍         | 1/21 [00:00<00:15,  1.32it/s]
    Running inference:  10%|▉         | 2/21 [00:01<00:18,  1.05it/s]
    Running inference:  14%|█▍        | 3/21 [00:02<00:15,  1.14it/s]
    Running inference:  19%|█▉        | 4/21 [00:03<00:14,  1.16it/s]
    Running inference:  24%|██▍       | 5/21 [00:04<00:13,  1.19it/s]
    Running inference:  29%|██▊       | 6/21 [00:05<00:12,  1.21it/s]
    Running inference:  33%|███▎      | 7/21 [00:05<00:11,  1.21it/s]
    Running inference:  38%|███▊      | 8/21 [00:06<00:10,  1.22it/s]
    Running inference:  43%|████▎     | 9/21 [00:07<00:09,  1.23it/s]
    Running inference:  48%|████▊     | 10/21 [00:08<00:08,  1.25it/s]
    Running inference:  52%|█████▏    | 11/21 [00:09<00:07,  1.26it/s]
    Running inference:  57%|█████▋    | 12/21 [00:09<00:07,  1.24it/s]
    Running inference:  62%|██████▏   | 13/21 [00:10<00:06,  1.24it/s]
    Running inference:  67%|██████▋   | 14/21 [00:11<00:05,  1.24it/s]
    Running inference:  71%|███████▏  | 15/21 [00:12<00:04,  1.25it/s]
    Running inference:  76%|███████▌  | 16/21 [00:13<00:04,  1.24it/s]
    Running inference:  81%|████████  | 17/21 [00:13<00:03,  1.25it/s]
    Running inference:  86%|████████▌ | 18/21 [00:14<00:02,  1.25it/s]
    Running inference:  90%|█████████ | 19/21 [00:15<00:01,  1.25it/s]
    Running inference:  95%|█████████▌| 20/21 [00:16<00:00,  1.25it/s]
    Running inference: 100%|██████████| 21/21 [00:17<00:00,  1.25it/s]    Running inference: 100%|██████████| 21/21 [00:17<00:00,  1.23it/s]
    2025-11-12 18:28:10.289 | SUCCESS  | earth2studio.run:deterministic:146 - Inference complete
    /
    ├── lat (721,) float64
    ├── lead_time (21,) timedelta64
    ├── lon (1440,) float64
    ├── t2m (1, 21, 721, 1440) float32
    ├── t850 (1, 21, 721, 1440) float32
    ├── tcwv (1, 21, 721, 1440) float32
    ├── time (1,) datetime64
    ├── z1000 (1, 21, 721, 1440) float32
    ├── z300 (1, 21, 721, 1440) float32
    ├── z500 (1, 21, 721, 1440) float32
    └── z700 (1, 21, 721, 1440) float32


.. GENERATED FROM PYTHON SOURCE LINES 104-111

Post Processing
---------------
The last step is to post process our results. Cartopy is a great library for plotting
fields on projections of a sphere. Here we will just plot the temperature at 2 meters
(t2m) 1 day into the forecast.

Notice that the Zarr IO function has additional APIs to interact with the stored data.

.. GENERATED FROM PYTHON SOURCE LINES 113-143

.. code-block:: Python

    import cartopy.crs as ccrs
    import matplotlib.pyplot as plt

    forecast = "2024-01-01"
    variable = "t2m"
    step = 4  # lead time = 24 hrs

    plt.close("all")
    # Create a Robinson projection
    projection = ccrs.Robinson()

    # Create a figure and axes with the specified projection
    fig, ax = plt.subplots(subplot_kw={"projection": projection}, figsize=(10, 6))

    # Plot the field using pcolormesh
    im = ax.pcolormesh(
        io["lon"][:],
        io["lat"][:],
        io[variable][0, step],
        transform=ccrs.PlateCarree(),
        cmap="Spectral_r",
    )

    # Set title
    ax.set_title(f"{forecast} - Lead time: {6*step}hrs")

    # Add coastlines and gridlines
    ax.coastlines()
    ax.gridlines()
    plt.savefig("outputs/01_t2m_prediction.jpg")


.. image-sg:: /examples/images/sphx_glr_01_deterministic_workflow_001.png
   :alt: 2024-01-01 - Lead time: 24hrs
   :srcset: /examples/images/sphx_glr_01_deterministic_workflow_001.png
   :class: sphx-glr-single-img


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 45.186 seconds)


.. _sphx_glr_download_examples_01_deterministic_workflow.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: 01_deterministic_workflow.ipynb <01_deterministic_workflow.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: 01_deterministic_workflow.py <01_deterministic_workflow.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: 01_deterministic_workflow.zip <01_deterministic_workflow.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_