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.
# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
# SPDX-FileCopyrightText: All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from collections import OrderedDict
from datetime import datetime
from math import ceil
import numpy as np
import torch
from loguru import logger
from tqdm import tqdm
from earth2studio.data import DataSource, ForecastSource, fetch_data
from earth2studio.io import IOBackend
from earth2studio.models.dx import DiagnosticModel
from earth2studio.models.px import PrognosticModel
from earth2studio.perturbation import Perturbation
from earth2studio.utils.checkpoint import (
Checkpoint,
CheckpointSession,
NullCheckpoint,
)
from earth2studio.utils.coords import CoordSystem, map_coords, split_coords
from earth2studio.utils.time import to_time_array
logger.remove()
logger.add(lambda msg: tqdm.write(msg, end=""), colorize=True)
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,
verbose: bool = True,
checkpoint: Checkpoint | CheckpointSession | NullCheckpoint = NullCheckpoint(),
) -> 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
verbose : bool, optional
Print inference progress, by default True
checkpoint : Checkpoint, optional
Checkpoint manager or checkpoint session used to record and resume workflow
progress, by default no checkpoint
Returns
-------
IOBackend
Output IO object
"""
logger.info("Running simple workflow!")
# Load model onto the device
device = (
device
if device is not None
else torch.device("cuda" if torch.cuda.is_available() else "cpu")
)
logger.info(f"Inference device: {device}")
prognostic = prognostic.to(device)
prognostic_ic = prognostic.input_coords()
time = to_time_array(time)
# Set up IO backend
total_coords = prognostic.output_coords(prognostic.input_coords()).copy()
for key, value in prognostic.output_coords(
prognostic.input_coords()
).items(): # Scrub batch dims
if value.shape == (0,):
del total_coords[key]
total_coords["time"] = time
total_coords["lead_time"] = np.asarray(
[
prognostic.output_coords(prognostic.input_coords())["lead_time"] * i
for i in range(nsteps + 1)
]
).flatten()
total_coords.move_to_end("lead_time", last=False)
total_coords.move_to_end("time", last=False)
for key, value in total_coords.items():
total_coords[key] = output_coords.get(key, value)
var_names = total_coords.pop("variable")
io.add_array(total_coords, var_names)
with checkpoint as ckpt:
restart_step = None
if ckpt.exists and ckpt.write_count > 0:
if ckpt.catalog.level < 2:
logger.warning(
"deterministic received checkpoint level "
f"{ckpt.catalog.level}; component state may not be "
"complete enough to resume a rollout. Re-running from "
"lead time zero."
)
else:
restart_step = ckpt.write_count - 1
if restart_step >= nsteps:
logger.success("\nInference complete")
return io
# Fetch data from data source and load onto device
if hasattr(prognostic, "interp_method"):
interp_to = prognostic_ic
interp_method = prognostic.interp_method
else:
interp_to = None
interp_method = "nearest"
x, coords = fetch_data(
source=data,
time=time,
variable=prognostic_ic["variable"],
lead_time=prognostic_ic["lead_time"],
device=device,
interp_to=interp_to,
interp_method=interp_method,
)
logger.success(f"Fetched data from {data.__class__.__name__}")
# Map lat and lon if needed
x, coords = map_coords(x, coords, prognostic.input_coords())
# Create prognostic iterator
model = prognostic.create_iterator(x, coords)
logger.info("Inference starting!")
initial_progress = 0 if restart_step is None else restart_step + 1
with tqdm(
total=nsteps + 1,
initial=initial_progress,
desc="Running inference",
position=1,
disable=(not verbose),
) as pbar:
for local_step, (x, coords) in enumerate(model):
step = (
local_step
if restart_step is None
else restart_step + local_step + 1
)
current_lead_time = coords["lead_time"][-1]
# Subselect domain/variables as indicated in output_coords
x, coords = map_coords(x, coords, output_coords)
io.write(*split_coords(x, coords))
ckpt.write(lead_time=current_lead_time)
pbar.update(1)
if step == nsteps:
break
ckpt.flush()
logger.success("\nInference complete")
return io
def diagnostic(
time: list[str] | list[datetime] | list[np.datetime64],
nsteps: int,
prognostic: PrognosticModel,
diagnostic: DiagnosticModel,
data: DataSource | ForecastSource,
io: IOBackend,
output_coords: CoordSystem = OrderedDict({}),
device: torch.device | None = None,
verbose: bool = True,
checkpoint: Checkpoint | CheckpointSession | NullCheckpoint = NullCheckpoint(),
) -> 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 | ForecastSource
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
verbose : bool, optional
Print inference progress, by default True
checkpoint : Checkpoint, optional
Checkpoint manager or checkpoint session used to record and resume workflow
progress, by default no checkpoint
Returns
-------
IOBackend
Output IO object
"""
logger.info("Running diagnostic workflow!")
device = (
device
if device is not None
else torch.device("cuda" if torch.cuda.is_available() else "cpu")
)
logger.info(f"Inference device: {device}")
prognostic = prognostic.to(device)
diagnostic = diagnostic.to(device)
prognostic_ic = prognostic.input_coords()
diagnostic_ic = diagnostic.input_coords()
time = to_time_array(time)
total_coords = prognostic.output_coords(prognostic.input_coords())
for key, value in prognostic.output_coords(
prognostic.input_coords()
).items(): # Scrub batch dims
if key in diagnostic.output_coords(diagnostic_ic):
total_coords[key] = diagnostic.output_coords(diagnostic_ic)[key]
if value.shape == (0,):
del total_coords[key]
total_coords["time"] = time
total_coords["lead_time"] = np.asarray(
[
prognostic.output_coords(prognostic.input_coords())["lead_time"] * i
for i in range(nsteps + 1)
]
).flatten()
total_coords.move_to_end("lead_time", last=False)
total_coords.move_to_end("time", last=False)
for key, value in total_coords.items():
total_coords[key] = output_coords.get(key, value)
var_names = total_coords.pop("variable")
io.add_array(total_coords, var_names)
with checkpoint as ckpt:
restart_step = None
if ckpt.exists and ckpt.write_count > 0:
if ckpt.catalog.level < 2:
logger.warning(
"diagnostic received checkpoint level "
f"{ckpt.catalog.level}; component state may not be "
"complete enough to resume a rollout. Re-running from "
"lead time zero."
)
else:
restart_step = ckpt.write_count - 1
if restart_step >= nsteps:
logger.success("\nInference complete")
return io
if hasattr(prognostic, "interp_method"):
interp_to = prognostic_ic
interp_method = prognostic.interp_method
else:
interp_to = None
interp_method = "nearest"
x, coords = fetch_data(
source=data,
time=time,
variable=prognostic_ic["variable"],
lead_time=prognostic_ic["lead_time"],
device=device,
interp_to=interp_to,
interp_method=interp_method,
)
logger.success(f"Fetched data from {data.__class__.__name__}")
x, coords = map_coords(x, coords, prognostic_ic)
model = prognostic.create_iterator(x, coords)
logger.info("Inference starting!")
initial_progress = 0 if restart_step is None else restart_step + 1
with tqdm(
total=nsteps + 1,
initial=initial_progress,
desc="Running inference",
position=1,
disable=(not verbose),
) as pbar:
for local_step, (x, coords) in enumerate(model):
step = (
local_step
if restart_step is None
else restart_step + local_step + 1
)
current_lead_time = coords["lead_time"][-1]
x, coords = map_coords(x, coords, diagnostic_ic)
x, coords = diagnostic(x, coords)
x, coords = map_coords(x, coords, output_coords)
io.write(*split_coords(x, coords))
ckpt.write(lead_time=current_lead_time)
pbar.update(1)
if step == nsteps:
break
ckpt.flush()
logger.success("\nInference complete")
return io
def ensemble(
time: list[str] | list[datetime] | list[np.datetime64],
nsteps: int,
nensemble: int,
prognostic: PrognosticModel,
data: DataSource,
io: IOBackend,
perturbation: Perturbation,
batch_size: int | None = None,
output_coords: CoordSystem = OrderedDict({}),
device: torch.device | None = None,
verbose: bool = True,
checkpoint: Checkpoint | CheckpointSession | NullCheckpoint = NullCheckpoint(),
) -> IOBackend:
"""Built in ensemble workflow.
Parameters
----------
time : list[str] | list[datetime] | list[np.datetime64]
List of string, datetimes or np.datetime64
nsteps : int
Number of forecast steps
nensemble : int
Number of ensemble members to run inference for.
prognostic : PrognosticModel
Prognostic models
data : DataSource
Data source
io : IOBackend
IO object
perturbation : Perturbation
Method to perturb the initial condition to create an ensemble.
batch_size: int, optional
Number of ensemble members to run in a single batch,
by default None.
output_coords: CoordSystem, optional
IO output coordinate system override, by default OrderedDict({})
device : torch.device, optional
Device to run inference on, by default None
verbose : bool, optional
Print inference progress, by default True
checkpoint : Checkpoint, optional
Checkpoint manager or checkpoint session used to record and resume workflow
progress, by default no checkpoint
Returns
-------
IOBackend
Output IO object
"""
logger.info("Running ensemble inference!")
device = (
device
if device is not None
else torch.device("cuda" if torch.cuda.is_available() else "cpu")
)
logger.info(f"Inference device: {device}")
prognostic = prognostic.to(device)
prognostic_ic = prognostic.input_coords()
time = to_time_array(time)
if hasattr(prognostic, "interp_method"):
interp_to = prognostic_ic
interp_method = prognostic.interp_method
else:
interp_to = None
interp_method = "nearest"
x0, coords0 = fetch_data(
source=data,
time=time,
variable=prognostic_ic["variable"],
lead_time=prognostic_ic["lead_time"],
device=device,
interp_to=interp_to,
interp_method=interp_method,
)
logger.success(f"Fetched data from {data.__class__.__name__}")
total_coords = prognostic.output_coords(prognostic.input_coords()).copy()
if "batch" in total_coords:
del total_coords["batch"]
total_coords["time"] = time
total_coords["lead_time"] = np.asarray(
[
prognostic.output_coords(prognostic.input_coords())["lead_time"] * i
for i in range(nsteps + 1)
]
).flatten()
total_coords.move_to_end("lead_time", last=False)
total_coords.move_to_end("time", last=False)
total_coords = {"ensemble": np.arange(nensemble)} | total_coords
for key, value in total_coords.items():
total_coords[key] = output_coords.get(key, value)
variables_to_save = total_coords.pop("variable")
io.add_array(total_coords, variables_to_save)
if batch_size is None:
batch_size = nensemble
batch_size = min(nensemble, batch_size)
with checkpoint as ckpt:
completed_ensembles = []
if ckpt.exists and not isinstance(ckpt, NullCheckpoint):
completed_ensembles = [
int(value) for value in ckpt.metadata.get("completed_ensembles", [])
]
completed = set(completed_ensembles)
start_batch_id = next(
(index for index in range(nensemble) if index not in completed),
nensemble,
)
number_of_batches = ceil((nensemble - start_batch_id) / batch_size)
restart_first_batch = (
ckpt.exists
and ckpt.write_count > 0
and start_batch_id < nensemble
and ckpt.lead_time != total_coords["lead_time"][-1]
)
logger.info(f"Starting {nensemble} Member Ensemble Inference with \
{number_of_batches} number of batches.")
for batch_index, batch_id in enumerate(
tqdm(
range(start_batch_id, nensemble, batch_size),
total=number_of_batches,
desc="Total Ensemble Batches",
position=2,
disable=(not verbose),
)
):
mini_batch_size = min(batch_size, nensemble - batch_id)
ensemble_coords = np.arange(batch_id, batch_id + mini_batch_size)
ensemble_members = [int(value) for value in ensemble_coords]
restart_step = None
if batch_index == 0 and restart_first_batch:
if ckpt.catalog.level < 2:
logger.warning(
"ensemble received checkpoint level "
f"{ckpt.catalog.level}; component state may not be "
"complete enough to resume a rollout. Re-running from "
"lead time zero."
)
ckpt.write_count = 0
else:
restart_step = ckpt.write_count - 1
if restart_step >= nsteps:
continue
elif not isinstance(ckpt, NullCheckpoint):
ckpt.write_count = 0
x = x0.to(device)
coords = OrderedDict({"ensemble": ensemble_coords}) | coords0.copy()
x = x.unsqueeze(0).repeat(mini_batch_size, *([1] * x.ndim))
x, coords = map_coords(x, coords, prognostic_ic)
x, coords = perturbation(x, coords)
model = prognostic.create_iterator(x, coords)
initial_progress = 0 if restart_step is None else restart_step + 1
with tqdm(
total=nsteps + 1,
initial=initial_progress,
desc=f"Running batch {batch_id} inference",
position=1,
leave=False,
disable=(not verbose),
) as pbar:
for local_step, (x, coords) in enumerate(model):
step = (
local_step
if restart_step is None
else restart_step + local_step + 1
)
current_lead_time = coords["lead_time"][-1]
x, coords = map_coords(x, coords, output_coords)
io.write(*split_coords(x, coords))
if step == nsteps:
completed.update(ensemble_members)
completed_ensembles = sorted(completed)
ckpt.write(
lead_time=current_lead_time,
completed_ensembles=completed_ensembles,
)
pbar.update(1)
if step == nsteps:
break
ckpt.flush()
logger.success("\nInference complete")
return io
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
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 output14 lines
/__w/earth2studio/earth2studio/.venv/lib/python3.13/site-packages/torch/cuda/__init__.py:64: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml # type: ignore[import]
WARNING[XFORMERS]: xFormers can't load C++/CUDA extensions. xFormers was built for:
PyTorch 2.10.0+cu128 with CUDA 1208 (you have 2.13.0+cu130)
Python 3.10.19 (you have 3.13.13)
Please reinstall xformers (see https://github.com/facebookresearch/xformers#installing-xformers)
Memory-efficient attention, SwiGLU, sparse and more won't be available.
Set XFORMERS_MORE_DETAILS=1 for more details
CuPy distance computation test failed with error: cuVS >= 24.12 or pylibraft < 24.12 should be installed to use this feature
Downloading dlwp_cubesphere.zip: 0%| | 0.00/67.2M [00:00<?, ?B/s]
Downloading dlwp_cubesphere.zip: 41%|████▏ | 27.8M/67.2M [00:00<00:00, 291MB/s]
Downloading dlwp_cubesphere.zip: 83%|████████▎ | 55.6M/67.2M [00:00<00:00, 270MB/s]
Downloading dlwp_cubesphere.zip: 100%|██████████| 67.2M/67.2M [00:00<00:00, 272MB/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, 2.05it/s]
Fetching GFS data: 29%|██▊ | 4/14 [00:00<00:01, 8.20it/s]
Fetching GFS data: 57%|█████▋ | 8/14 [00:00<00:00, 11.49it/s]
Fetching GFS data: 100%|██████████| 14/14 [00:00<00:00, 15.96it/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 output72 lines
2026-08-15 04:34:53.211 | INFO | earth2studio.run:deterministic:85 - Running simple workflow!
2026-08-15 04:34:53.211 | 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, 5.31it/s]
Fetching GFS data: 100%|██████████| 7/7 [00:00<00:00, 37.12it/s]
Fetching GFS data: 0%| | 0/7 [00:00<?, ?it/s]
Fetching GFS data: 14%|█▍ | 1/7 [00:00<00:01, 5.58it/s]
Fetching GFS data: 100%|██████████| 7/7 [00:00<00:00, 37.25it/s]
2026-08-15 04:34:53.686 | SUCCESS | earth2studio.run:deterministic:154 - Fetched data from GFS
2026-08-15 04:34:53.687 | INFO | earth2studio.run:deterministic:162 - Inference starting!
Running inference: 0%| | 0/21 [00:00<?, ?it/s]
Running inference: 5%|▍ | 1/21 [00:00<00:04, 4.61it/s]
Running inference: 10%|▉ | 2/21 [00:00<00:09, 2.09it/s]
Running inference: 14%|█▍ | 3/21 [00:01<00:06, 2.61it/s]
Running inference: 19%|█▉ | 4/21 [00:01<00:05, 2.87it/s]
Running inference: 24%|██▍ | 5/21 [00:01<00:05, 3.12it/s]
Running inference: 29%|██▊ | 6/21 [00:02<00:04, 3.23it/s]
Running inference: 33%|███▎ | 7/21 [00:02<00:04, 3.38it/s]
Running inference: 38%|███▊ | 8/21 [00:02<00:03, 3.42it/s]
Running inference: 43%|████▎ | 9/21 [00:02<00:03, 3.38it/s]
Running inference: 48%|████▊ | 10/21 [00:03<00:03, 3.33it/s]
Running inference: 52%|█████▏ | 11/21 [00:03<00:03, 3.15it/s]
Running inference: 57%|█████▋ | 12/21 [00:03<00:02, 3.18it/s]
Running inference: 62%|██████▏ | 13/21 [00:04<00:02, 3.22it/s]
Running inference: 67%|██████▋ | 14/21 [00:04<00:02, 3.20it/s]
Running inference: 71%|███████▏ | 15/21 [00:04<00:01, 3.34it/s]
Running inference: 76%|███████▌ | 16/21 [00:05<00:01, 3.37it/s]
Running inference: 81%|████████ | 17/21 [00:05<00:01, 3.44it/s]
Running inference: 86%|████████▌ | 18/21 [00:05<00:00, 3.25it/s]
Running inference: 90%|█████████ | 19/21 [00:05<00:00, 3.18it/s]
Running inference: 95%|█████████▌| 20/21 [00:06<00:00, 3.23it/s]
Running inference: 100%|██████████| 21/21 [00:06<00:00, 3.23it/s]
Running inference: 100%|██████████| 21/21 [00:06<00:00, 3.19it/s]
2026-08-15 04:35:00.261 | 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 runtime1m 18s
Setup
Duration45.8 sTagged cells
CPU load9%Average · Peak 12%
Process memory2.1 GiBPeak resident set
Network received81.8 MiBHost-wide estimate
Network sent241.0 KiBHost-wide estimate
GPU utilization0%Peak 2%
GPU memory0.5 GiBof 79.6 GiB
GPU power81 WPeak draw
Inference
Duration7.4 sTagged cells
CPU load11%Average · Peak 14%
Process memory3.2 GiBPeak resident set
Network received10.3 MiBHost-wide estimate
Network sent75.9 KiBHost-wide estimate
GPU utilization2%Peak 6%
GPU memory0.9 GiBof 79.6 GiB
GPU power82 WPeak draw
Plotting
Duration19.8 sTagged cells
CPU load11%Average · Peak 14%
Process memory3.4 GiBPeak resident set
Network received98.6 KiBHost-wide estimate
Network sent3.9 KiBHost-wide estimate
GPU utilization0%Peak 0%
GPU memory0.9 GiBof 79.6 GiB
GPU power81 WPeak draw
Execution environment
CPUAMD EPYC 7313P 16-Core Processor
GPUNVIDIA H100 PCIe · 79.6 GiB
System RAM58.5 GiB
PlatformLinux 6.8.0-136-generic
Python3.13.13
GPU driver / CUDADriver 595.84 · CUDA support 13.2