Running Ensemble Inference¶
Simple ensemble inference workflow.
This example will demonstrate how to run a simple inference workflow to generate a ensemble 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
- Select a perturbation method
- Running a simple built in workflow for ensembling
- 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 ensemble workflow:
earth2studio.run.ensemble.
Thus, we need the following:
- Prognostic Model: Use the built in FourCastNet model
earth2studio.models.px.FCN. - Perturbation Method: Use the Spherical Gaussian Method
earth2studio.perturbation.SphericalGaussian. - 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
import numpy as np
from earth2studio.data import GFS
from earth2studio.io import ZarrBackend
from earth2studio.models.px import FCN
from earth2studio.perturbation import SphericalGaussian
from earth2studio.run import ensemble
from earth2studio.utils.time import to_time_array
# Load the default model package which downloads the check point from NGC
package = FCN.load_default_package()
model = FCN.load_model(package)
# Instantiate the pertubation method
sg = SphericalGaussian(noise_amplitude=0.15)
# Create the data source
data = GFS()
# Create the IO handler, store in memory
chunks = {"ensemble": 1, "time": 1, "lead_time": 1}
io = ZarrBackend(
file_name="outputs/03_ensemble_sg.zarr",
chunks=chunks,
backend_kwargs={"overwrite": True},
)
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(["2024-01-01"]),
model.input_coords()["variable"],
)
print(f"Cached GFS input shape: {sample.shape}")
Console output4 lines
Fetching GFS data: 0%| | 0/26 [00:00<?, ?it/s]
Fetching GFS data: 4%|▍ | 1/26 [00:00<00:18, 1.33it/s]
Fetching GFS data: 100%|██████████| 26/26 [00:00<00:00, 34.10it/s]
Cached GFS input shape: (1, 26, 721, 1440)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 10 steps (for FCN, this is 60 hours) with 8 ensemble members which will be ran in 2 batches with batch size 4.
nsteps = 10
nensemble = 8
batch_size = 2
io = ensemble(
["2024-01-01"],
nsteps,
nensemble,
model,
data,
io,
sg,
batch_size=batch_size,
output_coords={"variable": np.array(["t2m", "tcwv"])},
)
Console output119 lines
2026-10-24 19:45:36.794 | INFO | earth2studio.run:ensemble:394 - Running ensemble inference!
2026-10-24 19:45:36.794 | INFO | earth2studio.run:ensemble:401 - Inference device: cuda
Fetching GFS data: 0%| | 0/26 [00:00<?, ?it/s]
Fetching GFS data: 4%|▍ | 1/26 [00:00<00:18, 1.32it/s]
Fetching GFS data: 100%|██████████| 26/26 [00:00<00:00, 33.53it/s]
2026-10-24 19:45:37.882 | SUCCESS | earth2studio.run:ensemble:422 - Fetched data from GFS
2026-10-24 19:45:37.950 | INFO | earth2studio.run:ensemble:466 - Starting 8 Member Ensemble Inference with 4 number of batches.
Total Ensemble Batches: 0%| | 0/4 [00:00<?, ?it/s]
Running batch 0 inference: 0%| | 0/11 [00:00<?, ?it/s]
Running batch 0 inference: 9%|▉ | 1/11 [00:00<00:01, 9.24it/s]
Running batch 0 inference: 18%|█▊ | 2/11 [00:00<00:01, 4.58it/s]
Running batch 0 inference: 27%|██▋ | 3/11 [00:00<00:01, 5.89it/s]
Running batch 0 inference: 36%|███▋ | 4/11 [00:00<00:01, 6.84it/s]
Running batch 0 inference: 45%|████▌ | 5/11 [00:00<00:00, 7.48it/s]
Running batch 0 inference: 55%|█████▍ | 6/11 [00:00<00:00, 7.91it/s]
Running batch 0 inference: 64%|██████▎ | 7/11 [00:00<00:00, 8.40it/s]
Running batch 0 inference: 73%|███████▎ | 8/11 [00:01<00:00, 8.60it/s]
Running batch 0 inference: 82%|████████▏ | 9/11 [00:01<00:00, 8.72it/s]
Running batch 0 inference: 91%|█████████ | 10/11 [00:01<00:00, 8.81it/s]
Running batch 0 inference: 100%|██████████| 11/11 [00:01<00:00, 8.87it/s]
Total Ensemble Batches: 25%|██▌ | 1/4 [00:02<00:08, 2.81s/it]
Running batch 2 inference: 0%| | 0/11 [00:00<?, ?it/s]
Running batch 2 inference: 18%|█▊ | 2/11 [00:00<00:00, 12.98it/s]
Running batch 2 inference: 36%|███▋ | 4/11 [00:00<00:00, 9.05it/s]
Running batch 2 inference: 45%|████▌ | 5/11 [00:00<00:00, 8.83it/s]
Running batch 2 inference: 55%|█████▍ | 6/11 [00:00<00:00, 8.36it/s]
Running batch 2 inference: 64%|██████▎ | 7/11 [00:00<00:00, 7.91it/s]
Running batch 2 inference: 73%|███████▎ | 8/11 [00:00<00:00, 7.76it/s]
Running batch 2 inference: 82%|████████▏ | 9/11 [00:01<00:00, 8.27it/s]
Running batch 2 inference: 91%|█████████ | 10/11 [00:01<00:00, 8.07it/s]
Running batch 2 inference: 100%|██████████| 11/11 [00:01<00:00, 7.92it/s]
Total Ensemble Batches: 50%|█████ | 2/4 [00:05<00:05, 2.61s/it]
Running batch 4 inference: 0%| | 0/11 [00:00<?, ?it/s]
Running batch 4 inference: 18%|█▊ | 2/11 [00:00<00:00, 13.07it/s]
Running batch 4 inference: 36%|███▋ | 4/11 [00:00<00:00, 9.33it/s]
Running batch 4 inference: 45%|████▌ | 5/11 [00:00<00:00, 8.75it/s]
Running batch 4 inference: 55%|█████▍ | 6/11 [00:00<00:00, 8.49it/s]
Running batch 4 inference: 64%|██████▎ | 7/11 [00:00<00:00, 8.15it/s]
Running batch 4 inference: 73%|███████▎ | 8/11 [00:00<00:00, 8.27it/s]
Running batch 4 inference: 82%|████████▏ | 9/11 [00:01<00:00, 8.03it/s]
Running batch 4 inference: 91%|█████████ | 10/11 [00:01<00:00, 7.95it/s]
Running batch 4 inference: 100%|██████████| 11/11 [00:01<00:00, 8.30it/s]
Total Ensemble Batches: 75%|███████▌ | 3/4 [00:07<00:02, 2.53s/it]
Running batch 6 inference: 0%| | 0/11 [00:00<?, ?it/s]
Running batch 6 inference: 18%|█▊ | 2/11 [00:00<00:00, 11.10it/s]
Running batch 6 inference: 36%|███▋ | 4/11 [00:00<00:00, 9.06it/s]
Running batch 6 inference: 45%|████▌ | 5/11 [00:00<00:00, 8.60it/s]
Running batch 6 inference: 55%|█████▍ | 6/11 [00:00<00:00, 8.29it/s]
Running batch 6 inference: 64%|██████▎ | 7/11 [00:00<00:00, 7.98it/s]
Running batch 6 inference: 73%|███████▎ | 8/11 [00:00<00:00, 7.84it/s]
Running batch 6 inference: 82%|████████▏ | 9/11 [00:01<00:00, 7.76it/s]
Running batch 6 inference: 91%|█████████ | 10/11 [00:01<00:00, 8.18it/s]
Running batch 6 inference: 100%|██████████| 11/11 [00:01<00:00, 8.12it/s]
Total Ensemble Batches: 100%|██████████| 4/4 [00:10<00:00, 2.50s/it]
Total Ensemble Batches: 100%|██████████| 4/4 [00:10<00:00, 2.54s/it]
2026-10-24 19:45:48.119 | SUCCESS | earth2studio.run:ensemble:536 -
Inference completePost Processing¶
The last step is to post process our results. Cartopy is a great library for plotting fields on projections of a sphere.
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"
def plot_(axi, data, title, cmap):
"""Convenience function for plotting pcolormesh."""
# Plot the field using pcolormesh
im = axi.pcolormesh(
io["lon"][:],
io["lat"][:],
data,
transform=ccrs.PlateCarree(),
cmap=cmap,
)
plt.colorbar(im, ax=axi, shrink=0.6, pad=0.04)
# Set title
axi.set_title(title)
# Add coastlines and gridlines
axi.coastlines()
axi.gridlines()
for variable, cmap in zip(["tcwv"], ["Blues"]):
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, (ax1, ax2, ax3) = plt.subplots(
nrows=1, ncols=3, subplot_kw={"projection": projection}, figsize=(16, 3)
)
plot_(
ax1,
io[variable][0, 0, step],
f"{forecast} - Lead time: {6*step}hrs - Member: {0}",
cmap,
)
plot_(
ax2,
io[variable][1, 0, step],
f"{forecast} - Lead time: {6*step}hrs - Member: {1}",
cmap,
)
plot_(
ax3,
np.std(io[variable][:, 0, step], axis=0),
f"{forecast} - Lead time: {6*step}hrs - Std",
cmap,
)
plt.savefig(f"outputs/03_{forecast}_{variable}_{step}_ensemble.jpg")
Execution profile
Runtime telemetry
Profiled phases
Total runtime1m 9s
Setup
Duration18.4 sTagged cells
CPU load11%Average · Peak 55%
Process memory1.8 GiBPeak resident set
Network received0 BHost-wide estimate
Network sent0 BHost-wide estimate
GPU utilization0%Peak 0%
GPU memory0.0 GiBof 79.6 GiB
GPU power48 WPeak draw
Inference
Duration11.3 sTagged cells
CPU load9%Average · Peak 11%
Process memory4.8 GiBPeak resident set
Network received0 BHost-wide estimate
Network sent0 BHost-wide estimate
GPU utilization26%Peak 73%
GPU memory3.8 GiBof 79.6 GiB
GPU power258 WPeak draw
Plotting
Duration36.1 sTagged cells
CPU load8%Average · Peak 10%
Process memory5.3 GiBPeak resident set
Network received0 BHost-wide estimate
Network sent0 BHost-wide estimate
GPU utilization0%Peak 0%
GPU memory3.8 GiBof 79.6 GiB
GPU power94 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