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

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: earth2studio.run.deterministic().

def deterministic(
    time: list[str] | list[datetime] | list[np.datetime64],
    nsteps: int,
    prognostic: PrognosticModel,
    data: DataSource,
    io: IOBackend,
    output_coords: CoordSystem = OrderedDict({}),
    device: torch.device | None = None,
) -> IOBackend:
    """Built in deterministic workflow.
    This workflow creates a determinstic inference pipeline to produce a forecast
    prediction using a prognostic 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
    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.
Datasource: Pull data from the GFS data api earth2studio.data.GFS.
IO Backend: Let’s 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.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()

2025-05-15 03:00:52.136 | WARNING  | earth2studio.models.auto.ngc:__init__:110 - API key found but no org found, using the org xznhlxbpr3qw

Downloading dlwp_cubesphere.zip: 0%|          | 0.00/67.2M [00:00<?, ?B/s]
Downloading dlwp_cubesphere.zip: 2%|▏         | 1.23M/67.2M [00:00<00:05, 12.9MB/s]
Downloading dlwp_cubesphere.zip: 26%|██▌       | 17.5M/67.2M [00:00<00:00, 106MB/s]
Downloading dlwp_cubesphere.zip: 54%|█████▍    | 36.4M/67.2M [00:00<00:00, 148MB/s]
Downloading dlwp_cubesphere.zip: 82%|████████▏ | 55.3M/67.2M [00:00<00:00, 168MB/s]
Downloading dlwp_cubesphere.zip: 100%|██████████| 67.2M/67.2M [00:00<00:00, 154MB/s]

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.

import earth2studio.run as run

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

print(io.root.tree())

2025-05-15 03:00:54.904 | INFO     | earth2studio.run:deterministic:75 - Running simple workflow!
2025-05-15 03:00:54.905 | INFO     | earth2studio.run:deterministic:82 - Inference device: cuda

Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]

2025-05-15 03:00:55.353 | DEBUG    | earth2studio.data.gfs:fetch_array:352 - 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]

2025-05-15 03:00:55.354 | DEBUG    | earth2studio.data.gfs:fetch_array:352 - 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-05-15 03:00:55.355 | DEBUG    | earth2studio.data.gfs:fetch_array:352 - 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-05-15 03:00:55.355 | DEBUG    | earth2studio.data.gfs:fetch_array:352 - 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-05-15 03:00:55.356 | DEBUG    | earth2studio.data.gfs:fetch_array:352 - 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-05-15 03:00:55.357 | DEBUG    | earth2studio.data.gfs:fetch_array:352 - 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-05-15 03:00:55.357 | DEBUG    | earth2studio.data.gfs:fetch_array:352 - 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]
Fetching GFS data:  14%|█▍        | 1/7 [00:00<00:05,  1.19it/s]
Fetching GFS data:  57%|█████▋    | 4/7 [00:01<00:00,  4.79it/s]
Fetching GFS data:  86%|████████▌ | 6/7 [00:03<00:00,  1.51it/s]
Fetching GFS data: 100%|██████████| 7/7 [00:03<00:00,  1.82it/s]
Fetching GFS data: 100%|██████████| 7/7 [00:03<00:00,  1.88it/s]

Fetching GFS data:   0%|          | 0/7 [00:00<?, ?it/s]

2025-05-15 03:00:59.474 | DEBUG    | earth2studio.data.gfs:fetch_array:352 - 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-05-15 03:00:59.477 | DEBUG    | earth2studio.data.gfs:fetch_array:352 - 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-05-15 03:00:59.478 | DEBUG    | earth2studio.data.gfs:fetch_array:352 - 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-05-15 03:00:59.480 | DEBUG    | earth2studio.data.gfs:fetch_array:352 - 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-05-15 03:00:59.481 | DEBUG    | earth2studio.data.gfs:fetch_array:352 - 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-05-15 03:00:59.483 | DEBUG    | earth2studio.data.gfs:fetch_array:352 - 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]

2025-05-15 03:00:59.484 | DEBUG    | earth2studio.data.gfs:fetch_array:352 - 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]
Fetching GFS data:  14%|█▍        | 1/7 [00:00<00:03,  1.79it/s]
Fetching GFS data:  29%|██▊       | 2/7 [00:00<00:01,  3.20it/s]
Fetching GFS data:  86%|████████▌ | 6/7 [00:00<00:00, 10.38it/s]
Fetching GFS data: 100%|██████████| 7/7 [00:00<00:00,  8.47it/s]
2025-05-15 03:01:00.340 | SUCCESS  | earth2studio.run:deterministic:106 - Fetched data from GFS
2025-05-15 03:01:00.344 | WARNING  | earth2studio.io.zarr:add_array:192 - Datetime64 not supported in zarr 3.0, converting to int64 nanoseconds since epoch
2025-05-15 03:01:00.348 | WARNING  | earth2studio.io.zarr:add_array:198 - Timedelta64 not supported in zarr 3.0, converting to int64 nanoseconds since epoch
2025-05-15 03:01:00.369 | INFO     | earth2studio.run:deterministic:136 - Inference starting!

Running inference:   0%|          | 0/21 [00:00<?, ?it/s]
Running inference:   5%|▍         | 1/21 [00:01<00:22,  1.12s/it]
Running inference:  10%|▉         | 2/21 [00:03<00:34,  1.82s/it]
Running inference:  14%|█▍        | 3/21 [00:05<00:35,  1.96s/it]
Running inference:  19%|█▉        | 4/21 [00:07<00:35,  2.07s/it]
Running inference:  24%|██▍       | 5/21 [00:10<00:34,  2.18s/it]
Running inference:  29%|██▊       | 6/21 [00:12<00:34,  2.27s/it]
Running inference:  33%|███▎      | 7/21 [00:15<00:33,  2.36s/it]
Running inference:  38%|███▊      | 8/21 [00:17<00:32,  2.47s/it]
Running inference:  43%|████▎     | 9/21 [00:20<00:30,  2.56s/it]
Running inference:  48%|████▊     | 10/21 [00:23<00:29,  2.65s/it]
Running inference:  52%|█████▏    | 11/21 [00:26<00:27,  2.74s/it]
Running inference:  57%|█████▋    | 12/21 [00:29<00:25,  2.84s/it]
Running inference:  62%|██████▏   | 13/21 [00:32<00:23,  2.92s/it]
Running inference:  67%|██████▋   | 14/21 [00:35<00:21,  3.02s/it]
Running inference:  71%|███████▏  | 15/21 [00:39<00:18,  3.11s/it]
Running inference:  76%|███████▌  | 16/21 [00:42<00:16,  3.21s/it]
Running inference:  81%|████████  | 17/21 [00:46<00:13,  3.30s/it]
Running inference:  86%|████████▌ | 18/21 [00:49<00:10,  3.40s/it]
Running inference:  90%|█████████ | 19/21 [00:53<00:06,  3.49s/it]
Running inference:  95%|█████████▌| 20/21 [00:57<00:03,  3.58s/it]
Running inference: 100%|██████████| 21/21 [01:01<00:00,  3.65s/it]
Running inference: 100%|██████████| 21/21 [01:01<00:00,  2.91s/it]
2025-05-15 03:02:01.425 | SUCCESS  | earth2studio.run:deterministic:146 - Inference complete
/
├── lat (721,) float64
├── lead_time (21,) int64
├── lon (1440,) float64
├── t2m (1, 21, 721, 1440) float32
├── t850 (1, 21, 721, 1440) float32
├── tcwv (1, 21, 721, 1440) float32
├── time (1,) int64
├── z1000 (1, 21, 721, 1440) float32
├── z300 (1, 21, 721, 1440) float32
├── z500 (1, 21, 721, 1440) float32
└── z700 (1, 21, 721, 1440) float32

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.

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

Total running time of the script: (2 minutes 16.216 seconds)

Gallery generated by Sphinx-Gallery