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

Set Up#

For this example, the built in diagnostic workflow earth2studio.run.diagnostic() will be used.

def diagnostic(
    time: list[str] | list[datetime] | list[np.datetime64],
    nsteps: int,
    prognostic: PrognosticModel,
    diagnostic: DiagnosticModel,
    data: DataSource,
    io: IOBackend,
    output_coords: CoordSystem = OrderedDict({}),
    device: torch.device | None = None,
) -> IOBackend:
    """Built in diagnostic workflow.
    This workflow creates a determinstic inference pipeline that couples a prognostic
    model with a diagnostic model.

    Parameters
    ----------
    time : list[str] | list[datetime] | list[np.datetime64]
        List of string, datetimes or np.datetime64
    nsteps : int
        Number of forecast steps
    prognostic : PrognosticModel
        Prognostic model
    diagnostic: DiagnosticModel
        Diagnostic model, must be on same coordinate axis as prognostic
    data : DataSource
        Data source
    io : IOBackend
        IO object
    output_coords: CoordSystem, optional
        IO output coordinate system override, by default OrderedDict({})
    device : torch.device, optional
        Device to run inference on, by default None

    Returns
    -------
    IOBackend
        Output IO object
    """

Thus, we need the following:

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

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()

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.

import earth2studio.run as run

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

print(io.root.tree())

2024-06-25 13:56:04.790 | INFO     | earth2studio.run:diagnostic:179 - Running diagnostic workflow!
2024-06-25 13:56:04.790 | INFO     | earth2studio.run:diagnostic:186 - Inference device: cuda
2024-06-25 13:56:04.836 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:149 - Fetching GFS index file: 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:   0%|          | 0/26 [00:00<?, ?it/s]

2024-06-25 13:56:05.160 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u10m at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:   0%|          | 0/26 [00:00<?, ?it/s]
Fetching GFS for 2021-06-01 00:00:00:   4%|▍         | 1/26 [00:06<02:38,  6.33s/it]

2024-06-25 13:56:11.486 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v10m at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:   4%|▍         | 1/26 [00:06<02:38,  6.33s/it]
Fetching GFS for 2021-06-01 00:00:00:   8%|▊         | 2/26 [00:09<01:53,  4.73s/it]

2024-06-25 13:56:15.104 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t2m at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:   8%|▊         | 2/26 [00:09<01:53,  4.73s/it]
Fetching GFS for 2021-06-01 00:00:00:  12%|█▏        | 3/26 [00:12<01:25,  3.71s/it]

2024-06-25 13:56:17.586 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: sp at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  12%|█▏        | 3/26 [00:12<01:25,  3.71s/it]
Fetching GFS for 2021-06-01 00:00:00:  15%|█▌        | 4/26 [00:14<01:04,  2.91s/it]

2024-06-25 13:56:19.276 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: msl at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  15%|█▌        | 4/26 [00:14<01:04,  2.91s/it]
Fetching GFS for 2021-06-01 00:00:00:  19%|█▉        | 5/26 [00:15<00:48,  2.31s/it]

2024-06-25 13:56:20.512 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t850 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  19%|█▉        | 5/26 [00:15<00:48,  2.31s/it]
Fetching GFS for 2021-06-01 00:00:00:  23%|██▎       | 6/26 [00:16<00:37,  1.86s/it]

2024-06-25 13:56:21.493 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u1000 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  23%|██▎       | 6/26 [00:16<00:37,  1.86s/it]
Fetching GFS for 2021-06-01 00:00:00:  27%|██▋       | 7/26 [00:17<00:29,  1.54s/it]

2024-06-25 13:56:22.373 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v1000 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  27%|██▋       | 7/26 [00:17<00:29,  1.54s/it]
Fetching GFS for 2021-06-01 00:00:00:  31%|███       | 8/26 [00:17<00:22,  1.27s/it]

2024-06-25 13:56:23.071 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z1000 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  31%|███       | 8/26 [00:17<00:22,  1.27s/it]
Fetching GFS for 2021-06-01 00:00:00:  35%|███▍      | 9/26 [00:18<00:18,  1.07s/it]

2024-06-25 13:56:23.696 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u850 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  35%|███▍      | 9/26 [00:18<00:18,  1.07s/it]
Fetching GFS for 2021-06-01 00:00:00:  38%|███▊      | 10/26 [00:19<00:14,  1.07it/s]

2024-06-25 13:56:24.324 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v850 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  38%|███▊      | 10/26 [00:19<00:14,  1.07it/s]
Fetching GFS for 2021-06-01 00:00:00:  42%|████▏     | 11/26 [00:19<00:12,  1.22it/s]

2024-06-25 13:56:24.884 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z850 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  42%|████▏     | 11/26 [00:19<00:12,  1.22it/s]
Fetching GFS for 2021-06-01 00:00:00:  46%|████▌     | 12/26 [00:20<00:10,  1.33it/s]

2024-06-25 13:56:25.478 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u500 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  46%|████▌     | 12/26 [00:20<00:10,  1.33it/s]
Fetching GFS for 2021-06-01 00:00:00:  50%|█████     | 13/26 [00:20<00:09,  1.39it/s]

2024-06-25 13:56:26.131 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v500 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  50%|█████     | 13/26 [00:20<00:09,  1.39it/s]
Fetching GFS for 2021-06-01 00:00:00:  54%|█████▍    | 14/26 [00:21<00:08,  1.43it/s]

2024-06-25 13:56:26.787 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z500 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  54%|█████▍    | 14/26 [00:21<00:08,  1.43it/s]
Fetching GFS for 2021-06-01 00:00:00:  58%|█████▊    | 15/26 [00:22<00:07,  1.47it/s]

2024-06-25 13:56:27.419 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t500 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  58%|█████▊    | 15/26 [00:22<00:07,  1.47it/s]
Fetching GFS for 2021-06-01 00:00:00:  62%|██████▏   | 16/26 [00:22<00:06,  1.55it/s]

2024-06-25 13:56:27.985 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z50 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  62%|██████▏   | 16/26 [00:22<00:06,  1.55it/s]
Fetching GFS for 2021-06-01 00:00:00:  65%|██████▌   | 17/26 [00:23<00:06,  1.48it/s]

2024-06-25 13:56:28.730 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: r500 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  65%|██████▌   | 17/26 [00:23<00:06,  1.48it/s]
Fetching GFS for 2021-06-01 00:00:00:  69%|██████▉   | 18/26 [00:24<00:05,  1.47it/s]

2024-06-25 13:56:29.419 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: r850 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  69%|██████▉   | 18/26 [00:24<00:05,  1.47it/s]
Fetching GFS for 2021-06-01 00:00:00:  73%|███████▎  | 19/26 [00:24<00:04,  1.46it/s]

2024-06-25 13:56:30.110 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: tcwv at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  73%|███████▎  | 19/26 [00:24<00:04,  1.46it/s]
Fetching GFS for 2021-06-01 00:00:00:  77%|███████▋  | 20/26 [00:25<00:04,  1.36it/s]

2024-06-25 13:56:30.965 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u100m at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  77%|███████▋  | 20/26 [00:25<00:04,  1.36it/s]
Fetching GFS for 2021-06-01 00:00:00:  81%|████████  | 21/26 [00:26<00:03,  1.36it/s]

2024-06-25 13:56:31.701 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v100m at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  81%|████████  | 21/26 [00:26<00:03,  1.36it/s]
Fetching GFS for 2021-06-01 00:00:00:  85%|████████▍ | 22/26 [00:27<00:02,  1.38it/s]

2024-06-25 13:56:32.405 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: u250 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  85%|████████▍ | 22/26 [00:27<00:02,  1.38it/s]
Fetching GFS for 2021-06-01 00:00:00:  88%|████████▊ | 23/26 [00:28<00:02,  1.35it/s]

2024-06-25 13:56:33.174 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: v250 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  88%|████████▊ | 23/26 [00:28<00:02,  1.35it/s]
Fetching GFS for 2021-06-01 00:00:00:  92%|█████████▏| 24/26 [00:28<00:01,  1.40it/s]

2024-06-25 13:56:33.833 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: z250 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  92%|█████████▏| 24/26 [00:28<00:01,  1.40it/s]
Fetching GFS for 2021-06-01 00:00:00:  96%|█████████▌| 25/26 [00:29<00:00,  1.41it/s]

2024-06-25 13:56:34.525 | DEBUG    | earth2studio.data.gfs:fetch_gfs_dataarray:196 - Fetching GFS grib file for variable: t250 at 2021-06-01 00:00:00

Fetching GFS for 2021-06-01 00:00:00:  96%|█████████▌| 25/26 [00:29<00:00,  1.41it/s]
Fetching GFS for 2021-06-01 00:00:00: 100%|██████████| 26/26 [00:30<00:00,  1.46it/s]
Fetching GFS for 2021-06-01 00:00:00: 100%|██████████| 26/26 [00:30<00:00,  1.15s/it]
2024-06-25 13:56:35.234 | SUCCESS  | earth2studio.run:diagnostic:200 - Fetched data from GFS
2024-06-25 13:56:35.260 | INFO     | earth2studio.run:diagnostic:231 - Inference starting!

Running inference:   0%|          | 0/9 [00:00<?, ?it/s]
Running inference:  11%|█         | 1/9 [00:00<00:00,  8.85it/s]
Running inference:  22%|██▏       | 2/9 [00:00<00:01,  6.94it/s]
Running inference:  33%|███▎      | 3/9 [00:00<00:00,  6.60it/s]
Running inference:  44%|████▍     | 4/9 [00:00<00:00,  6.41it/s]
Running inference:  56%|█████▌    | 5/9 [00:00<00:00,  6.30it/s]
Running inference:  67%|██████▋   | 6/9 [00:00<00:00,  6.14it/s]
Running inference:  78%|███████▊  | 7/9 [00:01<00:00,  6.04it/s]
Running inference:  89%|████████▉ | 8/9 [00:01<00:00,  5.94it/s]
Running inference: 100%|██████████| 9/9 [00:01<00:00,  5.91it/s]
Running inference: 100%|██████████| 9/9 [00:01<00:00,  6.19it/s]
2024-06-25 13:56:36.713 | SUCCESS  | earth2studio.run:diagnostic:245 - Inference complete
/
 ├── lat (720,) float64
 ├── lead_time (9,) timedelta64[h]
 ├── lon (1440,) float64
 ├── time (1,) datetime64[ns]
 └── tp (1, 9, 720, 1440) float32

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.

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")

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

Gallery generated by Sphinx-Gallery