
.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/02_diagnostic_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_02_diagnostic_workflow.py>`
        to download the full example code.

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

.. _sphx_glr_examples_02_diagnostic_workflow.py:


Running Diagnostic Inference
============================

Basic prognostic + diagnostic inference workflow.

This example will demonstrate how to run a deterministic inference workflow that couples
a prognostic model with a diagnostic model. This diagnostic model will predict a new
atmospheric quantity from the predicted fields of the prognostic.

In this example you will learn:

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

.. GENERATED FROM PYTHON SOURCE LINES 36-43

.. code-block:: Python

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


.. GENERATED FROM PYTHON SOURCE LINES 44-48

Set Up
------
For this example, the built in diagnostic workflow :py:meth:`earth2studio.run.diagnostic`
will be used.

.. GENERATED FROM PYTHON SOURCE LINES 50-54

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

.. GENERATED FROM PYTHON SOURCE LINES 57-63

Thus, we need the following:

- Prognostic Model: Use the built in FourCastNet Model :py:class:`earth2studio.models.px.FCN`.
- Diagnostic Model: Use the built in precipitation AFNO model :py:class:`earth2studio.models.dx.PrecipitationAFNO`.
- Datasource: Pull data from the GFS data api :py:class:`earth2studio.data.GFS`.
- IO Backend: Save the outputs into a Zarr store :py:class:`earth2studio.io.ZarrBackend`.

.. GENERATED FROM PYTHON SOURCE LINES 65-90

.. 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.dx import PrecipitationAFNO
    from earth2studio.models.px import FCN

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

    package = PrecipitationAFNO.load_default_package()
    diagnostic_model = PrecipitationAFNO.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 fcn.zip: 0%|          | 0.00/267M [00:00<?, ?B/s]    Downloading fcn.zip: 14%|█▍        | 37.2M/267M [00:00<00:00, 389MB/s]    Downloading fcn.zip: 34%|███▍      | 90.9M/267M [00:00<00:00, 491MB/s]    Downloading fcn.zip: 54%|█████▍    | 144M/267M [00:00<00:00, 521MB/s]     Downloading fcn.zip: 73%|███████▎  | 194M/267M [00:00<00:00, 467MB/s]    Downloading fcn.zip: 95%|█████████▍| 253M/267M [00:00<00:00, 517MB/s]    Downloading fcn.zip: 100%|██████████| 267M/267M [00:00<00:00, 504MB/s]
    Downloading precipitation_afno.zip: 0%|          | 0.00/261M [00:00<?, ?B/s]    Downloading precipitation_afno.zip: 17%|█▋        | 44.6M/261M [00:00<00:00, 467MB/s]    Downloading precipitation_afno.zip: 34%|███▍      | 89.3M/261M [00:00<00:00, 434MB/s]    Downloading precipitation_afno.zip: 51%|█████     | 133M/261M [00:00<00:00, 443MB/s]     Downloading precipitation_afno.zip: 72%|███████▏  | 189M/261M [00:00<00:00, 499MB/s]    Downloading precipitation_afno.zip: 94%|█████████▎| 244M/261M [00:00<00:00, 527MB/s]    Downloading precipitation_afno.zip: 100%|██████████| 261M/261M [00:00<00:00, 504MB/s]


.. GENERATED FROM PYTHON SOURCE LINES 91-97

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.

.. GENERATED FROM PYTHON SOURCE LINES 99-108

.. code-block:: Python

    import earth2studio.run as run

    nsteps = 8
    io = run.diagnostic(
        ["2021-06-01"], nsteps, prognostic_model, diagnostic_model, data, io
    )

    print(io.root.tree())


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

 .. code-block:: none

    2025-11-12 18:28:40.635 | INFO     | earth2studio.run:diagnostic:190 - Running diagnostic workflow!
    2025-11-12 18:28:40.635 | INFO     | earth2studio.run:diagnostic:197 - Inference device: cuda
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.072 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 343243269-852573
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.074 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 433893232-1229275
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.076 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 265295671-828915
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.078 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 209303792-746874
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.080 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 422428446-507899
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.081 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 210050666-761020
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.083 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 349147153-959449
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.085 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 263723951-825006
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.087 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 269727465-946420
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.088 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 155987833-763786
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.090 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 342333604-909665
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.092 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 472734498-976207
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.093 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 350106602-970343
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.095 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 427023039-950918
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.096 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 215895048-622148
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.098 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 407717702-959428
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.100 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 414253228-840220
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.101 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 471746333-988165
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.103 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 412154325-999310
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.104 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 426047081-975958
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.106 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 406739323-978379
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.107 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 270673885-954744
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.109 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 0-1002486
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.111 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 344095842-882412
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.112 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 264548957-746714
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]                                                             2025-11-12 18:28:41.114 | DEBUG    | earth2studio.data.gfs:fetch_array:380 - Fetching GFS grib file: noaa-gfs-bdp-pds/gfs.20210601/00/atmos/gfs.t00z.pgrb2.0p25.f000 215284398-610650
    Fetching GFS data:   0%|          | 0/26 [00:00<?, ?it/s]    Fetching GFS data:   4%|▍         | 1/26 [00:00<00:13,  1.87it/s]    Fetching GFS data:   8%|▊         | 2/26 [00:00<00:07,  3.09it/s]    Fetching GFS data:  12%|█▏        | 3/26 [00:00<00:05,  4.08it/s]    Fetching GFS data:  42%|████▏     | 11/26 [00:00<00:00, 20.00it/s]    Fetching GFS data:  58%|█████▊    | 15/26 [00:01<00:00, 24.00it/s]    Fetching GFS data:  81%|████████  | 21/26 [00:01<00:00, 27.34it/s]    Fetching GFS data: 100%|██████████| 26/26 [00:01<00:00, 20.35it/s]
    2025-11-12 18:28:42.474 | SUCCESS  | earth2studio.run:diagnostic:220 - Fetched data from GFS
    2025-11-12 18:28:42.489 | INFO     | earth2studio.run:diagnostic:252 - Inference starting!

    Running inference:   0%|          | 0/9 [00:00<?, ?it/s]
    Running inference:  11%|█         | 1/9 [00:00<00:01,  4.22it/s]
    Running inference:  22%|██▏       | 2/9 [00:00<00:01,  4.39it/s]
    Running inference:  33%|███▎      | 3/9 [00:00<00:01,  4.51it/s]
    Running inference:  44%|████▍     | 4/9 [00:00<00:01,  4.54it/s]
    Running inference:  56%|█████▌    | 5/9 [00:01<00:00,  4.53it/s]
    Running inference:  67%|██████▋   | 6/9 [00:01<00:00,  4.48it/s]
    Running inference:  78%|███████▊  | 7/9 [00:01<00:00,  4.45it/s]
    Running inference:  89%|████████▉ | 8/9 [00:01<00:00,  4.46it/s]
    Running inference: 100%|██████████| 9/9 [00:02<00:00,  4.45it/s]    Running inference: 100%|██████████| 9/9 [00:02<00:00,  4.46it/s]
    2025-11-12 18:28:44.507 | SUCCESS  | earth2studio.run:diagnostic:266 - Inference complete
    /
    ├── lat (720,) float64
    ├── lead_time (9,) timedelta64
    ├── lon (1440,) float64
    ├── time (1,) datetime64
    └── tp (1, 9, 720, 1440) float32


.. GENERATED FROM PYTHON SOURCE LINES 109-120

Post Processing
---------------
The last step is to plot the resulting predicted total precipitation. The power of
diagnostic models is that they allow the prediction of any variable from a pre-trained
prognostic model.

.. note::
  The built in workflow will only save the direct outputs of the diagnostic. In this
  example only total precipitation is accessible for plotting. If you wish to save
  outputs of both the prognostic and diagnostic, we recommend writing a custom
  workflow.

.. GENERATED FROM PYTHON SOURCE LINES 122-164

.. code-block:: Python

    from datetime import datetime

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

    forecast = datetime(2021, 6, 1)
    variable = "tp"
    step = 8  # lead time = 48 hrs

    plt.close("all")
    # Create a Orthographic projection of USA
    projection = ccrs.Orthographic(-100, 40)

    # 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
    levels = np.arange(0.0, 0.01, 0.001)
    im = ax.contourf(
        io["lon"][:],
        io["lat"][:],
        io[variable][0, step],
        levels,
        transform=ccrs.PlateCarree(),
        vmax=0.01,
        vmin=0.00,
        cmap="terrain",
    )

    # Set title
    ax.set_title(f"{forecast.strftime('%Y-%m-%d')} - Lead time: {6*step}hrs")

    # Add coastlines and gridlines6
    ax.set_extent([220, 340, 20, 70])  # [lat min, lat max, lon min, lon max]
    ax.coastlines()
    ax.gridlines()
    plt.colorbar(
        im, ax=ax, ticks=levels, shrink=0.75, pad=0.04, label="Total precipitation (m)"
    )

    plt.savefig("outputs/02_tp_prediction.jpg")


.. image-sg:: /examples/images/sphx_glr_02_diagnostic_workflow_001.png
   :alt: 2021-06-01 - Lead time: 48hrs
   :srcset: /examples/images/sphx_glr_02_diagnostic_workflow_001.png
   :class: sphx-glr-single-img


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

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


.. _sphx_glr_download_examples_02_diagnostic_workflow.py:

.. only:: html

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

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

      :download:`Download Jupyter notebook: 02_diagnostic_workflow.ipynb <02_diagnostic_workflow.ipynb>`

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

      :download:`Download Python source code: 02_diagnostic_workflow.py <02_diagnostic_workflow.py>`

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

      :download:`Download zipped: 02_diagnostic_workflow.zip <02_diagnostic_workflow.zip>`


.. only:: html

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

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