Building Custom Batched Ensemble Weather Forecasting with NVIDIA Earth2Studio

NVIDIA's Earth2Studio tutorial demonstrates a custom batched ensemble weather forecasting workflow, using the FCN prognostic model and GFS initial conditions to generate 8-member ensembles with wind-power diagnostics and verification metrics. The tutorial, published by Marktechpost, provides code for installing Earth2Studio in Colab, implementing perturbation systems, and visualizing forecast uncertainty.

In this tutorial https://github.com/MARKTECHPOST-AI-MEDIA-INC/AI-Agents-Projects-Tutorials/blob/main/Deep%20Learning/NVIDIA Earth2Studio Custom Ensemble Forecasting Marktechpost.ipynb , we build an ensemble weather forecasting workflow with NVIDIA Earth2Studio https://github.com/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. python 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. 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. python 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, "u500": 0.40, "v500": 0.40, "tcwv": 0.30}, default=0.0, control member=True print f"Perturbation sampler: {pert.kind}" io = ZarrBackend file name="outputs/e2s ensemble.zarr", chunks={"ensemble": 1, "time": 1, "lead time": 1}, backend kwargs={"overwrite": True} io = run ensemble INIT STR , NSTEPS, NENSEMBLE, BATCH SIZE, model, diagnostic, pert, GFS , io, SAVE VARS, DEVICE print io.root.tree We define helper functions that select atmospheric channels and write them into a coordinate-aware Zarr backend. We build a custom batched ensemble loop that fetches GFS initial conditions, perturbs ensemble members, aligns coordinates, iterates the FCN model, and chains the wind-power diagnostic. We then load the model, initialize the diagnostic and perturbation components, execute the forecast, and inspect the resulting Zarr structure. leads = np.asarray io "lead time" : .astype "timedelta64 ns " lead h = leads.astype "timedelta64 h " .astype int valid = to time array INIT STR 0 + leads truth, tc = fetch data source=GFS , time=valid, lead time=np.array np.timedelta64 0, "h" , variable=np.array VERIFY VARS , device="cpu" truth = truth :, 0 w = torch.cos torch.deg2rad torch.as tensor np.asarray LAT , dtype=torch.float32 w2d = w :, None .expand len LAT , len LON .contiguous mcoords = OrderedDict {"lead time": leads, "lat": np.asarray LAT , "lon": np.asarray LON } def fair crps ens, obs, weights : """Fair unbiased CRPS, lat-weighted. ens: M, lat, lon , obs: lat, lon .""" M = ens.shape 0 wn = weights / weights.sum skill = ens - obs .abs wn .sum dim= -2, -1 .mean spread = torch.zeros , dtype=ens.dtype for i in range M : spread = spread + ens i - ens .abs wn .sum dim= -2, -1 .sum return skill - spread / 2 M M - 1 .item scores = {} for k, var in enumerate VERIFY VARS : fc = torch.as tensor np.asarray io var : :, 0 .float ob = truth :, k .float mean = fc.mean 0 try: metric = rmse reduction dimensions= "lat", "lon" , weights=w2d r, = metric mean, mcoords, ob, mcoords r = r.numpy except Exception as e: print f" built-in rmse unavailable: {e} " wn = w2d / w2d.sum r = torch.sqrt mean - ob 2 wn .sum dim= -2, -1 .numpy wn = w2d / w2d.sum spread = torch.sqrt fc.var 0, unbiased=True wn .sum dim= -2, -1 .numpy crps = np.array fair crps fc :, t , ob t , w2d for t in range fc.shape 1 scores var = dict rmse=r, spread=spread, crps=crps, fc=fc, obs=ob, mean=mean print f"\n=== {var} ===" print f"{'lead h ': 8}{'RMSE': 12}{'spread': 12}{'ratio': 9}{'CRPS': 12}" for t in range len lead h : ratio = spread t / r t if r t 0 else np.nan print f"{lead h t : 8}{r t : 12.3f}{spread t : 12.3f}{ratio: 9.2f}{crps t : 12.3f}" We retrieve GFS analyses for every forecast-valid time and use them as the reference data for verification. We calculate latitude-weighted RMSE, ensemble spread, fair CRPS, and spread-to-error ratios for temperature, geopotential height, and wind variables. We store the forecast fields and evaluation metrics in a structured dictionary and print lead-time skill summaries for each variable. lat np, lon np = np.asarray LAT , np.asarray LON ilat = int np.argmin np.abs lat np - POI 1 ilon = int np.argmin np.abs lon np - POI 2 % 360 last = -1 d = scores "t2m" fields = d "mean" last .numpy - 273.15, "ensemble mean t2m C ", "RdBu r", None , d "fc" :, last .std 0 .numpy , "ensemble spread K ", "magma", None , d "obs" last .numpy - 273.15, "GFS analysis C ", "RdBu r", None , d "mean" last - d "obs" last .numpy , "mean error K ", "coolwarm", 5 fig, axs = plt.subplots 2, 2, figsize= 15, 7 , constrained layout=True for ax, f, title, cmap, lim in zip axs.ravel , fields : kw = dict vmin=-lim, vmax=lim if lim else {} im = ax.pcolormesh lon np, lat np, f, cmap=cmap, shading="auto", kw ax.set title f"{title} — +{lead h last } h" ; plt.colorbar im, ax=ax, shrink=0.85 plt.show z = scores "z500" "fc" :, last .numpy / 9.81 la = lat np 25 & lat np < 75 lo = lon np 280 | lon np < 40 lon shift = np.where lon np 180, lon np - 360, lon np order = np.argsort lon shift lo plt.figure figsize= 11, 5 for m in range z.shape 0 : sub = z m np.ix la, lo :, order plt.contour lon shift lo order , lat np la , sub, levels= 5520 , colors= "k" if m == 0 else "C0" , linewidths= 2.0 if m == 0 else 0.8 zo = scores "z500" "obs" last .numpy / 9.81 plt.contour lon shift lo order , lat np la , zo np.ix la, lo :, order , levels= 5520 , colors="crimson", linewidths=2.5 plt.title f"z500 5520 m spaghetti at +{lead h last } h " f" black=control, blue=members, red=GFS analysis " plt.xlabel "lon" ; plt.ylabel "lat" ; plt.show t2m pt = scores "t2m" "fc" :, :, ilat, ilon .numpy - 273.15 obs pt = scores "t2m" "obs" :, ilat, ilon .numpy - 273.15 cf pt = np.asarray io "wind cf" : :, 0, :, ilat, ilon fig, a1, a2 = plt.subplots 1, 2, figsize= 14, 4 a1.fill between lead h, t2m pt.min 0 , t2m pt.max 0 , alpha=0.25, label="member range" a1.plot lead h, t2m pt.mean 0 , "o-", label="ensemble mean" a1.plot lead h, t2m pt 0 , "k--", label="control" a1.plot lead h, obs pt, "r^-", label="GFS analysis" a1.set title f"2 m temperature — {POI 0 }" ; a1.set xlabel "lead h " ; a1.set ylabel "C" a1.legend ; a1.grid alpha=.3 a2.fill between lead h, cf pt.min 0 , cf pt.max 0 , alpha=0.25, color="seagreen" a2.plot lead h, cf pt.mean 0 , "o-", color="seagreen" a2.set title f"wind capacity factor custom diagnostic — {POI 0 }" a2.set xlabel "lead h " ; a2.set ylim 0, 1 ; a2.grid alpha=.3 plt.tight layout ; plt.show fig, axs = plt.subplots 1, len VERIFY VARS , figsize= 5 len VERIFY VARS , 3.6 for ax, var in zip np.atleast 1d axs , VERIFY VARS : s = scores var ax.plot lead h, s "rmse" , "o-", label="RMSE ens. mean " ax.plot lead h, s "spread" , "s--", label="spread" ax.plot lead h, s "crps" , "^:", label="fair CRPS" ax.set title var ; ax.set xlabel "lead h " ; ax.grid alpha=.3 ; ax.legend fontsize=8 plt.tight layout ; plt.show import xarray as xr ds = xr.open zarr "outputs/e2s ensemble.zarr" print ds We visualize ensemble behavior through temperature mean, spread, analysis, and error maps at the final forecast lead time. We generate geopotential-height spaghetti contours, a New Delhi temperature fan chart, a wind-capacity-factor forecast, and lead-time skill curves. We finally open the Zarr output with Xarray so that we can inspect, analyze, or export the complete ensemble dataset. In conclusion, we established a flexible and extensible Earth2Studio workflow that goes beyond running a predefined ensemble function. We directly controlled initial-condition perturbation, member batching, model iteration, diagnostic chaining, coordinate alignment, data persistence, verification, and visualization within a single Colab environment. We also demonstrated how physically scaled perturbations and an unperturbed control member help us interpret ensemble spread. At the same time, RMSE, fair CRPS, and spread-skill diagnostics allow us to evaluate forecast accuracy and calibration across lead times. The resulting Zarr dataset preserves the complete ensemble structure and remains accessible through Xarray for further analysis or conversion. Because the workflow follows Earth2Studio’s component interfaces, we can extend it by replacing the prognostic model, changing the atmospheric data source, adding new diagnostics, increasing the ensemble size, or adopting asynchronous storage without redesigning the full forecasting pipeline. Check out the FULL CODES here. Also, feel free to follow us on and don’t forget to join our Twitter https://x.com/intent/follow?screen name=marktechpost and Subscribe to 150k+ML SubReddit https://www.reddit.com/r/machinelearningnews/ . Wait are you on telegram? our Newsletter https://magic.beehiiv.com/v1/f5e63dd4-5653-4f09-83e2-321a8b1ba526?email={{email}} now you can join us on telegram as well. https://t.me/machinelearningresearchnews Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us https://forms.gle/wbash1wF6efRj8G58 Sana Hassan, a consulting intern at Marktechpost and dual-degree student at IIT Madras, is passionate about applying technology and AI to address real-world challenges. With a keen interest in solving practical problems, he brings a fresh perspective to the intersection of AI and real-life solutions.