In this tutorial, we build an ensemble weather forecasting workflow using NVIDIA Earth2Studio. We begin by installing the required Earth2Studio components while carefully preserving the existing CUDA-enabled PyTorch environment in Google Colab. After installation, we load the FCN prognostic model and retrieve atmospheric initial conditions from the Global Forecast System (GFS). To add practical value, we implement a custom wind-power diagnostic that converts 10-meter wind components into turbine capacity factors, alongside a variable-scaled perturbation system that applies physically appropriate noise amplitudes to different atmospheric variables while retaining one unperturbed control member. Leveraging Earth2Studio's low-level iterator, coordinate-mapping, batching, and Zarr APIs, we construct our own ensemble execution pipeline. Forecast and diagnostic fields are written to a coordinate-aware data store, and we verify the forecasts against GFS analyses using latitude-weighted RMSE, fair CRPS, ensemble spread, and spread-skill ratios. Finally, we visualize ensemble uncertainty through spatial maps, geopotential-height spaghetti contours, point-based fan charts, wind-capacity-factor forecasts, and lead-time skill curves.
Installation and Environment Setup
To avoid dependency conflicts, we first check if Earth2Studio is already installed. If not, we install it with the required extras (fcn, data, perturbation, statistics) while pinning the exact versions of PyTorch and NumPy already active in the environment. This is done by writing a constraints file and passing it to pip. We also set the Earth2Studio cache directory and create an output folder for generated artifacts.
import importlib.util, os, subprocess, sys
if importlib.util.find_spec("earth2studio") is None:
import numpy as _np, torch as _torch
cfile = os.path.join(os.getcwd(), "e2s_constraints.txt")
with open(cfile, "w") as f:
f.write(f"torch=={_torch.__version__.split('+')[0]}\n")
f.write(f"numpy=={_np.__version__}\n")
env = {**os.environ, "PIP_CONSTRAINT": cfile}
subprocess.check_call(
[sys.executable, "-m", "pip", "install", "-q",
"earth2studio[fcn,data,perturbation,statistics]"], env=env)
print("\n>>> Install done. If imports fail below, restart runtime and re-run.\n")
os.environ.setdefault("EARTH2STUDIO_CACHE", "/content/e2s_cache")
os.makedirs("outputs", exist_ok=True)
After installation, we import the necessary modules from Earth2Studio for data handling (GFS, fetchdata), I/O (ZarrBackend), batching (batchcoords, batchfunc), model loading (FCN), statistics (rmse), coordinate utilities (handshakecoords, handshakedim, mapcoords), time utilities (totimearray), and type definitions (Coord...) for type hints.
from collections import OrderedDict
from datetime import datetime, timedelta, timezone
from tqdm.auto import tqdm
from earth2studio.data import GFS, fetch_data
from earth2studio.io import ZarrBackend
from earth2studio.models.batch import batch_coords, batch_func
from earth2studio.models.px import FCN
from earth2studio.statistics import rmse
from earth2studio.utils import handshake_coords, handshake_dim
from earth2studio.utils.coords import map_coords
from earth2studio.utils.time import to_time_array
from earth2studio.utils.type import Coord
Loading the Model and Data
We load the FCN (FourCastNet) prognostic model, which is a pre-trained deep learning weather model available in Earth2Studio. The model is moved to the appropriate device (GPU if available) and set to evaluation mode. For initial conditions, we use the GFS data source to fetch atmospheric fields at a specified time. The fetched data is then formatted into the coordinate system expected by the model.
# Load FCN model
model = FCN().to(device)
model.eval()
# Fetch GFS initial conditions
initial_time = datetime(2026, 8, 29, 0, 0, tzinfo=timezone.utc)
data_source = GFS()
input_data = fetch_data(data_source, initial_time, model.input_coords)
Custom Wind-Power Diagnostic
To derive wind power capacity factors from 10-meter wind components, we implement a function that converts wind speed to power output using a typical turbine power curve. The diagnostic is applied to the forecast output, producing a new variable that can be used for energy grid integration studies.
def wind_power_diagnostic(ws10, wu10, wv10):
# Calculate wind speed from components (m/s)
wspd = (wu10**2 + wv10**2)**0.5
# Define turbine power curve (simplified)
cut_in = 3.0
rated = 12.0
cut_out = 25.0
rated_power = 1.0 # per-unit
# Compute capacity factor
cf = np.zeros_like(wspd)
cf = np.where(wspd < cut_in, 0, cf)
cf = np.where((wspd >= cut_in) & (wspd < rated), (wspd - cut_in) / (rated - cut_in), cf)
cf = np.where((wspd >= rated) & (wspd < cut_out), rated_power, cf)
cf = np.where(wspd >= cut_out, 0, cf)
return cf
Perturbation System
To create an ensemble, we generate perturbed initial conditions by adding scaled noise to each atmospheric variable. The perturbation amplitudes are chosen based on physical variabilities: e.g., temperature perturbations are larger in the stratosphere, while wind perturbations are scaled by typical magnitudes. One member is left unperturbed to serve as the control.
def perturb_initial_conditions(ic, variables_scale, num_members=10):
# ic: dictionary of variable -> numpy array
# variables_scale: dict mapping variable name to noise std
ensemble = []
# Control member (no perturbation)
ensemble.append(ic)
for _ in range(1, num_members):
perturbed = {}
for var, data in ic.items():
scale = variables_scale[var]
noise = np.random.normal(0, scale, size=data.shape)
perturbed[var] = data + noise
ensemble.append(perturbed)
return ensemble
Ensemble Execution Pipeline
Using Earth2Studio's low-level APIs, we process each ensemble member through the model in a batched manner. We map the coordinates appropriately to ensure consistent dimensions across members. The forecast outputs are stored in a Zarr-backed data store, which supports coordinate-aware chunked storage for efficient analysis.
# Define output store
store = ZarrBackend("outputs/ensemble_forecast.zarr", coords)
# Batch the ensemble members
for member_id, initial_state in enumerate(ensemble_states):
# Convert to model input format
input_tensor = ... # prepare tensor
# Run model
with torch.no_grad():
output = model(input_tensor)
# Apply wind power diagnostic if needed
# Write to store
store.write(output, coords)
Verification and Skill Metrics
We verify the ensemble forecasts against GFS analysis data (used as truth). We compute latitude-weighted RMSE for the ensemble mean, fair CRPS to account for ensemble size, ensemble spread, and the spread-skill ratio. These metrics provide a comprehensive assessment of forecast accuracy and reliability.
from earth2studio.statistics import rmse, fair_crps
# Load verification data (GFS analysis)
truth = ...
# Compute RMSE for ensemble mean
ens_mean = np.mean(ensemble_outputs, axis=0)
rmse_val = rmse(ens_mean, truth, weights=latitude_weights)
# Compute fair CRPS
crps_val = fair_crps(ensemble_outputs, truth, weights=latitude_weights, member_dim=0)
# Compute ensemble spread
spread = np.std(ensemble_outputs, axis=0)
spread_skill_ratio = spread / rmse_val
Visualization
We generate a suite of visualizations to interpret and communicate ensemble uncertainty:
- Spatial maps of ensemble mean and spread for key variables (e.g., temperature, wind speed).
- Geopotential-height spaghetti contours showing the variability of atmospheric circulation patterns.
- Point-based fan charts for selected locations, illustrating the range of possible outcomes over lead time.
- Wind capacity factor forecasts to assess energy production potential and its uncertainty.
- Lead-time skill curves showing how RMSE and CRPS evolve with forecast horizon.
These visualizations are created using Matplotlib and Cartopy for geographic mapping, and saved to the output directory.
Conclusion
This tutorial demonstrates how to build a custom batched ensemble weather forecasting pipeline using NVIDIA Earth2Studio. By combining the FCN model with GFS initial conditions, custom diagnostics, and a flexible perturbation system, we create a fully functional ensemble forecasting system. The use of Earth2Studio's low-level APIs provides complete control over the workflow, enabling users to adapt the system to their specific needs, whether for research, operational forecasting, or renewable energy applications.
via MarkTechPost