Running Deterministic Inference¶
Basic deterministic inference workflow.
This example will demonstrate how to run a simple inference workflow to generate a basic deterministic 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.
Thus, we need the following:
- Prognostic Model: Use the built in Deep Learning Weather Prediction (DLWP) model
earth2studio.models.px.DLWP. - 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
from earth2studio.utils.time import to_time_array
# 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()
Console output8 lines
Downloading dlwp_cubesphere.zip: 0%| | 0.00/67.2M [00:00<?, ?B/s]
Downloading dlwp_cubesphere.zip: 2%|▏ | 1.36M/67.2M [00:00<00:04, 14.2MB/s]
Downloading dlwp_cubesphere.zip: 12%|█▏ | 8.21M/67.2M [00:00<00:01, 47.9MB/s]
Downloading dlwp_cubesphere.zip: 37%|███▋ | 24.8M/67.2M [00:00<00:00, 86.3MB/s]
Downloading dlwp_cubesphere.zip: 62%|██████▏ | 41.5M/67.2M [00:00<00:00, 117MB/s]
Downloading dlwp_cubesphere.zip: 79%|███████▊ | 52.9M/67.2M [00:00<00:00, 101MB/s]
Downloading dlwp_cubesphere.zip: 98%|█████████▊| 65.5M/67.2M [00:00<00:00, 102MB/s]
Downloading dlwp_cubesphere.zip: 100%|██████████| 67.2M/67.2M [00:00<00:00, 95.4MB/s]Fetch Data¶
You can easily fetch raw Xarray data from an initial condition data source with a simple call. By default, this caches the data locally on your machine, so you won't have to re-download it if you access it again or use it in an inference pipeline.
sample = data(
to_time_array(["2023-12-31T18:00:00", "2024-01-01"]),
model.input_coords()["variable"],
)
print(sample)
Console output53 lines
Fetching GFS data: 0%| | 0/14 [00:00<?, ?it/s]
Fetching GFS data: 7%|▋ | 1/14 [00:00<00:06, 1.90it/s]
Fetching GFS data: 21%|██▏ | 3/14 [00:00<00:02, 3.84it/s]
Fetching GFS data: 29%|██▊ | 4/14 [00:00<00:02, 4.67it/s]
Fetching GFS data: 100%|██████████| 14/14 [00:00<00:00, 14.01it/s]
<xarray.DataArray (time: 2, variable: 7, lat: 721, lon: 1440)> Size: 116MB
array([[[[ 2.53431367e+02, 2.53431367e+02, 2.53431367e+02, ...,
2.53431367e+02, 2.53431367e+02, 2.53431367e+02],
[ 2.54021367e+02, 2.54021367e+02, 2.54021367e+02, ...,
2.54021367e+02, 2.54021367e+02, 2.54021367e+02],
[ 2.54761367e+02, 2.54761367e+02, 2.54761367e+02, ...,
2.54751367e+02, 2.54761367e+02, 2.54761367e+02],
...,
[ 2.63911367e+02, 2.63911367e+02, 2.63901367e+02, ...,
2.63911367e+02, 2.63911367e+02, 2.63911367e+02],
[ 2.64881367e+02, 2.64881367e+02, 2.64881367e+02, ...,
2.64881367e+02, 2.64881367e+02, 2.64881367e+02],
[ 2.65661367e+02, 2.65661367e+02, 2.65661367e+02, ...,
2.65661367e+02, 2.65661367e+02, 2.65661367e+02]],
[[ 1.44381158e+03, 1.44381158e+03, 1.44381158e+03, ...,
1.44381158e+03, 1.44381158e+03, 1.44381158e+03],
[ 1.40975126e+03, 1.40975126e+03, 1.40975126e+03, ...,
1.40959430e+03, 1.40975126e+03, 1.40975126e+03],
[ 1.37035430e+03, 1.37035430e+03, 1.37035430e+03, ...,
1.37035430e+03, 1.37035430e+03, 1.37035430e+03],
...
[ 1.94594834e+00, 1.94594834e+00, 1.94434834e+00, ...,
1.94914834e+00, 1.94914834e+00, 1.94754834e+00],
[ 1.99074834e+00, 1.99074834e+00, 1.99074834e+00, ...,
1.99234834e+00, 1.99234834e+00, 1.99234834e+00],
[ 2.04194834e+00, 2.04194834e+00, 2.04194834e+00, ...,
2.04194834e+00, 2.04194834e+00, 2.04194834e+00]],
[[ 2.44794863e+02, 2.44794863e+02, 2.44794863e+02, ...,
2.44794863e+02, 2.44794863e+02, 2.44794863e+02],
[ 2.44884863e+02, 2.44884863e+02, 2.44874863e+02, ...,
2.44884863e+02, 2.44884863e+02, 2.44884863e+02],
[ 2.46354863e+02, 2.46344863e+02, 2.46344863e+02, ...,
2.46354863e+02, 2.46354863e+02, 2.46354863e+02],
...,
[ 2.54594863e+02, 2.54574863e+02, 2.54564863e+02, ...,
2.54634863e+02, 2.54624863e+02, 2.54604863e+02],
[ 2.55794863e+02, 2.55794863e+02, 2.55794863e+02, ...,
2.55804863e+02, 2.55804863e+02, 2.55794863e+02],
[ 2.56544863e+02, 2.56544863e+02, 2.56544863e+02, ...,
2.56544863e+02, 2.56544863e+02, 2.56544863e+02]]]])
Coordinates:
* time (time) datetime64[ns] 16B 2023-12-31T18:00:00 2024-01-01
* variable (variable) <U5 140B 't850' 'z1000' 'z700' ... 'z300' 'tcwv' 't2m'
* lat (lat) float64 6kB 90.0 89.75 89.5 89.25 ... -89.5 -89.75 -90.0
* lon (lon) float64 12kB 0.0 0.25 0.5 0.75 ... 359.0 359.2 359.5 359.8
lead_time timedelta64[ns] 8B 00:00:00Execute 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())
Console output46 lines
2026-10-24 19:34:34.765 | INFO | earth2studio.run:deterministic:85 - Running simple workflow!
2026-10-24 19:34:34.765 | INFO | earth2studio.run:deterministic:92 - Inference device: cuda
Fetching GFS data: 0%| | 0/7 [00:00<?, ?it/s]
Fetching GFS data: 14%|█▍ | 1/7 [00:00<00:01, 4.28it/s]
Fetching GFS data: 100%|██████████| 7/7 [00:00<00:00, 29.37it/s]
Fetching GFS data: 0%| | 0/7 [00:00<?, ?it/s]
Fetching GFS data: 14%|█▍ | 1/7 [00:00<00:01, 5.14it/s]
Fetching GFS data: 100%|██████████| 7/7 [00:00<00:00, 33.85it/s]
2026-10-24 19:34:35.551 | SUCCESS | earth2studio.run:deterministic:154 - Fetched data from GFS
2026-10-24 19:34:35.552 | INFO | earth2studio.run:deterministic:162 - Inference starting!
Running inference: 0%| | 0/21 [00:00<?, ?it/s]
Running inference: 10%|▉ | 2/21 [00:00<00:04, 4.55it/s]
Running inference: 19%|█▉ | 4/21 [00:00<00:02, 7.36it/s]
Running inference: 29%|██▊ | 6/21 [00:00<00:01, 10.12it/s]
Running inference: 43%|████▎ | 9/21 [00:00<00:00, 13.64it/s]
Running inference: 52%|█████▏ | 11/21 [00:00<00:00, 15.09it/s]
Running inference: 67%|██████▋ | 14/21 [00:01<00:00, 16.92it/s]
Running inference: 81%|████████ | 17/21 [00:01<00:00, 18.55it/s]
Running inference: 95%|█████████▌| 20/21 [00:01<00:00, 19.24it/s]
Running inference: 100%|██████████| 21/21 [00:01<00:00, 14.84it/s]
2026-10-24 19:34:36.968 | SUCCESS | earth2studio.run:deterministic:189 -
Inference complete
/
├── lat (721,) float64
├── lead_time (21,) timedelta64[h]
├── lon (1440,) float64
├── t2m (1, 21, 721, 1440) float32
├── t850 (1, 21, 721, 1440) float32
├── tcwv (1, 21, 721, 1440) float32
├── time (1,) datetime64[ns]
├── z1000 (1, 21, 721, 1440) float32
├── z300 (1, 21, 721, 1440) float32
├── z500 (1, 21, 721, 1440) float32
└── z700 (1, 21, 721, 1440) float32Post 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")
Execution profile
Runtime telemetry
Profiled phases
Total runtime55.5 s
Setup
Duration30.2 sTagged cells
CPU load8%Average · Peak 11%
Process memory1.6 GiBPeak resident set
Network received80.0 MiBHost-wide estimate
Network sent187.7 KiBHost-wide estimate
GPU utilization0%Peak 0%
GPU memory0.0 GiBof 79.6 GiB
GPU power47 WPeak draw
Inference
Duration2.6 sTagged cells
CPU load11%Average · Peak 12%
Process memory2.2 GiBPeak resident set
Network received42 BHost-wide estimate
Network sent42 BHost-wide estimate
GPU utilization4%Peak 7%
GPU memory0.9 GiBof 79.6 GiB
GPU power81 WPeak draw
Plotting
Duration18.9 sTagged cells
CPU load9%Average · Peak 11%
Process memory2.8 GiBPeak resident set
Network received97.1 KiBHost-wide estimate
Network sent6.3 KiBHost-wide estimate
GPU utilization0%Peak 0%
GPU memory0.9 GiBof 79.6 GiB
GPU power82 WPeak draw
Execution environment
CPUAMD EPYC 7313P 16-Core Processor
GPUNVIDIA H100 PCIe · 79.6 GiB
System RAM58.5 GiB
PlatformLinux 6.8.0-142-generic
Python3.13.13
GPU driver / CUDADriver 595.91.07 · CUDA support 13.2