Building Custom Batched Ensemble Weather Forecasting with NVIDIA Earth2Studio

Summary

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. 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…