In this tutorial, we build an ensemble weather forecasting workflow with NVIDIA Earth2Studio. We install the required Earth2Studio components while preserving Colab’s existing CUDA-enabled PyTorch environment, load the FCN prognostic model, and retrieve atmospheric initial conditions from GFS. We then implement a custom wind-power diagnostic that converts 10-meter wind components into turbine capacity factors, along with a variable-scaled perturbation system that applies physically appropriate noise amplitudes to different atmospheric variables while retaining an unperturbed control member. Using Earth2Studio’s low-level iterator, coordinate-mapping, batching, and Zarr APIs, we construct our own ensemble execution pipeline, write forecast and diagnostic fields to a coordinate-aware data store, and 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. Copy CodeCopiedUse a different Browser 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}”) We install Earth2Studio while preserving Colab’s existing CUDA-enabled PyTorch and NumPy environment through package constraints. We configure the model cache, import the forecasting, data, statistics, plotting, and coordinate-management utilities, and detect the available compute device. We also define the ensemble size, batch size, forecast duration, saved variables, verification variables, initialization time, and New Delhi point of interest. Copy CodeCopiedUse a different Browser 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 We create a custom diagnostic model that converts 10-meter wind components into hub-height wind speed and turbine capacity factor. We validate coordinate compatibility through Earth2Studio’s handshake utilities and support batched inputs with the provided decorators. We also implement variable-specific spatial perturbations that retain member zero as an unperturbed control forecast. Copy CodeCopiedUse a different Browser def write_vars(io, x, coords, names): “””Write selected channels of a (…, variable, lat, lon) tensor to the IO backend.””” vax = list(coords).index(“variable”) sub = OrderedDict((k, v) for k, v in coords.items() if k != “variable”) for name in names: hit = np.where(np.asarray(coords[“variable”]) == name)[0] if hit.size: io.write(x.select(vax, int(hit[0])).cpu(), sub, name) def run_ensemble(time, nsteps, nensemble, batch_size, prognostic, diagnostic, perturbation, data, io, save_vars, device): time = to_time_array(time) ic = prognostic.input_coords() x0, c0 = fetch_data(source=data, time=time, lead_time=ic[“lead_time”], variable=ic[“variable”], device=device) print(f”Initial condition tensor: {tuple(x0.shape)} dims={list(c0)}”) oc = prognostic.output_coords(ic) dt = oc[“lead_time”] prog_vars = [v for v in save_vars if v in set(map(str, oc[“variable”]))] total = OrderedDict({ “ensemble”: np.arange(nensemble), “time”: time, “lead_time”: np.asarray([dt * i for i in range(nsteps + 1)]).flatten(), “lat”: oc[“lat”], “lon”: oc[“lon”], }) io.add_array(total, prog_vars + [“wind_cf”]) dx_target = OrderedDict((k, v) for k, v in diagnostic.input_coords().items() if k != “batch”) nbatch = int(np.ceil(nensemble / batch_size)) with torch.inference_mode(): for b in tqdm(range(nbatch), desc=”ensemble batches”): lo = b * batch_size n = min(batch_size, nensemble – lo) x = x0.unsqueeze(0).repeat(n, *([1] * x0.ndim)) coords = OrderedDict({“ensemble”: np.arange(lo, lo + n), **c0}) x, coords = perturbation(x, coords) x, coords = map_coords(x, coords, ic) for step, (xs, cs) in enumerate(prognostic.create_iterator(x, coords)): write_vars(io, xs, cs, prog_vars) xw, cw = map_coords(xs, cs, dx_target) xw, cw = diagnostic(xw, cw) write_vars(io, xw, cw, [“wind_cf”]) if step >= nsteps: break torch.cuda.empty_cache() if device.type == “cuda” else None return io model = FCN.load_model(FCN.load_default_package()).to(DEVICE) grid = model.output_coords(model.input_coords()) LAT, LON = grid[“lat”], grid[“lon”] diagnostic = WindPowerCF(LAT, LON).to(DEVICE) pert = VariableScaledNoise( amplitudes={“t2m”: 0.20, “t850”: 0.20, “z500”: 40.0, “z850”: 25.0, “u10m”: 0.25, “v10m”: 0.25,