A tutorial demonstrates how to build a custom batched ensemble weather forecasting workflow using NVIDIA Earth2Studio. The guide installs the required components while keeping Colab’s existing CUDA-enabled PyTorch environment intact. It loads the FCN prognostic model and pulls atmospheric initial conditions from GFS. A custom wind-power diagnostic converts 10-meter wind components into turbine capacity factors. A variable-scaled perturbation system applies physically appropriate noise amplitudes to different atmospheric variables while keeping an unperturbed control member. The process uses Earth2Studio’s low-level iterator, coordinate-mapping, batching, and Zarr APIs to construct the ensemble execution pipeline. Forecasts and diagnostic fields write to a coordinate-aware data store. Verification against GFS analyses uses latitude-weighted RMSE, fair CRPS, ensemble spread, and spread-skill ratios. Visualisations include spatial maps, geopotential-height spaghetti contours, point-based fan charts, wind-capacity-factor forecasts, and lead-time skill curves.
Setup and configuration
The installation preserves Colab’s CUDA-enabled PyTorch and NumPy environment through package constraints. Configuration sets the model cache and imports forecasting, data, statistics, plotting, and coordinate-management utilities. The code detects the available compute device. Variables define the ensemble size, batch size, forecast duration, saved variables, verification variables, initialization time, and a point of interest in New Delhi.
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 the imports below fail: Runtime > Restart session, re-run.\n")
os.environ.setdefault("EARTH2STUDIO_CACHE", "/content/e2s_cache")
os.makedirs("outputs", exist_ok=True)
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 CoordSystem
if DEVICE.type == "cpu":
print("!! No GPU detected — this will be very slow. Runtime > Change runtime type > T4 GPU")
NENSEMBLE = 8
BATCH_SIZE = 2
NSTEPS = 8
SAVE_VARS = ["t2m", "z500", "u10m", "v10m", "tcwv"]
VERIFY_VARS = ["t2m", "z500", "u10m"]
INIT = (datetime.now(timezone.utc) - timedelta(days=7)).replace()
INIT_STR = INIT.strftime("%Y-%m-%dT%H:%M:%S")
POI = ("New Delhi", 28.61, 77.21)
print(f"Initialization: {INIT_STR} | device: {DEVICE}")
Wind power diagnostics and noise
A custom diagnostic model converts 10-meter wind components into hub-height wind speed and turbine capacity factor. Coordinate compatibility validation uses Earth2Studio’s handshake utilities. The setup supports batched inputs via provided decorators. Variable-specific spatial perturbations retain member zero as an unperturbed control forecast.
class WindPowerCF(torch.nn.Module):
"""Turbine capacity factor [0,1] from 10 m winds via power-law shear + power curve."""
def __init__(self, lat, lon, hub=100.0, alpha=0.143,
cut_in=3.0, rated=12.0, cut_out=25.0):
super().__init__()
self.lat, self.lon = lat, lon
self.hub, self.alpha = hub, alpha
self.cut_in, self.rated, self.cut_out = cut_in, rated, cut_out
def input_coords(self) -> CoordSystem:
return OrderedDict({
"batch": np.empty(0),
"variable": np.array(["u10m", "v10m"]),
"lat": self.lat,
"lon": self.lon,
})
@batch_coords()
def output_coords(self, input_coords: CoordSystem) -> CoordSystem:
target = self.input_coords()
for i, (key, _) in enumerate(target.items()):
if key != "batch":
handshake_dim(input_coords, key, i)
handshake_coords(input_coords, target, key)
oc = OrderedDict({
"batch": np.empty(0),
"variable": np.array(["wind_cf"]),
"lat": self.lat,
"lon": self.lon,
})
oc["batch"] = input_coords["batch"]
return oc
@batch_func()
def __call__(self, x: torch.Tensor, coords: CoordSystem):
oc = self.output_coords(coords)
u, v = x[..., 0:1, :, :], x[..., 1:2, :, :]
ws10 = torch.sqrt(u * u + v * v)
ws = ws10 * (self.hub / 10.0) ** self.alpha
ramp = (ws ** 3 - self.cut_in ** 3) / (self.rated ** 3 - self.cut_in ** 3)
cf = torch.zeros_like(ws)
cf = torch.where((ws >= self.cut_in) & (ws < self.rated), ramp.clamp(0, 1), cf)
cf = torch.where((ws >= self.rated) & (ws <= self.cut_out), torch.ones_like(cf), cf)
return cf, oc
class VariableScaledNoise:
"""Spatially correlated noise with per-variable amplitudes + control member."""
def __init__(self, amplitudes: dict, default: float = 0.0, control_member: bool = True):
self.amplitudes, self.default, self.control = amplitudes, default, control_member
try:
from earth2studio.perturbation import SphericalGaussian
self.sampler, self.kind = SphericalGaussian(noise_amplitude=1.0), "SphericalGaussian"
except Exception:
from earth2studio.perturbation import Brown
self.sampler, self.kind = Brown(noise_amplitude=1.0), "Brown"
def __call__(self, x: torch.Tensor, coords: CoordSystem):
noise, _ = self.sampler(torch.zeros_like(x), coords)
vax = list(coords).index("variable")
amps = torch.tensor([self.amplitudes.get(str(v), self.default)
for v in coords["variable"]], device=x.device, dtype=x.dtype)
shape = [1] * x.ndim; shape[vax] = amps.numel()
pert = noise * amps.reshape(shape)
if self.control and "ensemble" in coords:
eax = list(coords).index("ensemble")
mask = torch.tensor((np.asarray(coords["ensemble"]) != 0).astype(np.float32),
device=x.device, dtype=x.dtype)
mshape = [1] * x.ndim; mshape[eax] = mask.numel()
pert = pert * mask.reshape(mshape)
return x + pert, coords
The code defines a class for turbine capacity factor calculation using a power-law shear and power curve. It takes latitude, longitude, hub height, and wind speed thresholds. The output coordinates map wind components to a single capacity factor variable. A second class handles spatially correlated noise with per-variable amplitudes. It attempts to use SphericalGaussian sampling but falls back to Brown noise if unavailable. The function adds noise to the input tensor while preserving the control member where the ensemble index is zero.
Running the ensemble
The workflow creates a function to write selected channels of a tensor to the IO backend. It loops through variables and writes data to the Zarr store. The main execution function takes time, steps, ensemble size, batch size, prognostic, diagnostic, perturbation, data, IO backend, saved variables, and device as arguments. It converts the time input to an array. Initial conditions fetch from the data source using the specified time and lead time. The code prints the initial condition tensor shape and dimensions. It calculates the output coordinates and determines the time step. Prognostic variables filter the saved variables list. A total coordinate dictionary builds the ensemble, time, lead time, latitude, and longitude arrays. The IO backend adds these arrays plus the prognostic and wind capacity factor variables. The batch
Related articles
