{"slug": "building-custom-batched-ensemble-weather-forecasting-with-nvidia-earth2studio", "title": "Building Custom Batched Ensemble Weather Forecasting with NVIDIA Earth2Studio", "summary": "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.", "body_md": "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.\n\n``` python\nimport importlib.util, os, subprocess, sys\nif importlib.util.find_spec(\"earth2studio\") is None:\n   import numpy as _np, torch as _torch\n   cfile = os.path.join(os.getcwd(), \"e2s_constraints.txt\")\n   with open(cfile, \"w\") as f:\n       f.write(f\"torch=={_torch.__version__.split('+')[0]}\\n\")\n       f.write(f\"numpy=={_np.__version__}\\n\")\n   env = {**os.environ, \"PIP_CONSTRAINT\": cfile}\n   subprocess.check_call(\n       [sys.executable, \"-m\", \"pip\", \"install\", \"-q\",\n        \"earth2studio[fcn,data,perturbation,statistics]\"], env=env)\n   print(\"\\n>>> Install done. If the imports below fail: Runtime > Restart session, re-run.\\n\")\nos.environ.setdefault(\"EARTH2STUDIO_CACHE\", \"/content/e2s_cache\")\nos.makedirs(\"outputs\", exist_ok=True)\nfrom collections import OrderedDict\nfrom datetime import datetime, timedelta, timezone\nfrom tqdm.auto import tqdm\nfrom earth2studio.data import GFS, fetch_data\nfrom earth2studio.io import ZarrBackend\nfrom earth2studio.models.batch import batch_coords, batch_func\nfrom earth2studio.models.px import FCN\nfrom earth2studio.statistics import rmse\nfrom earth2studio.utils import handshake_coords, handshake_dim\nfrom earth2studio.utils.coords import map_coords\nfrom earth2studio.utils.time import to_time_array\nfrom earth2studio.utils.type import CoordSystem\nif DEVICE.type == \"cpu\":\n   print(\"!! No GPU detected — this will be very slow. Runtime > Change runtime type > T4 GPU\")\nNENSEMBLE  = 8\nBATCH_SIZE = 2\nNSTEPS     = 8\nSAVE_VARS  = [\"t2m\", \"z500\", \"u10m\", \"v10m\", \"tcwv\"]\nVERIFY_VARS = [\"t2m\", \"z500\", \"u10m\"]\nINIT = (datetime.now(timezone.utc) - timedelta(days=7)).replace()\nINIT_STR = INIT.strftime(\"%Y-%m-%dT%H:%M:%S\")\nPOI = (\"New Delhi\", 28.61, 77.21)\nprint(f\"Initialization: {INIT_STR}  |  device: {DEVICE}\")\n```\n\nWe 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.\n\n```\nclass WindPowerCF(torch.nn.Module):\n   \"\"\"Turbine capacity factor [0,1] from 10 m winds via power-law shear + power curve.\"\"\"\n   def __init__(self, lat, lon, hub=100.0, alpha=0.143,\n                cut_in=3.0, rated=12.0, cut_out=25.0):\n       super().__init__()\n       self.lat, self.lon = lat, lon\n       self.hub, self.alpha = hub, alpha\n       self.cut_in, self.rated, self.cut_out = cut_in, rated, cut_out\n   def input_coords(self) -> CoordSystem:\n       return OrderedDict({\n           \"batch\": np.empty(0),\n           \"variable\": np.array([\"u10m\", \"v10m\"]),\n           \"lat\": self.lat,\n           \"lon\": self.lon,\n       })\n   @batch_coords()\n   def output_coords(self, input_coords: CoordSystem) -> CoordSystem:\n       target = self.input_coords()\n       for i, (key, _) in enumerate(target.items()):\n           if key != \"batch\":\n               handshake_dim(input_coords, key, i)\n               handshake_coords(input_coords, target, key)\n       oc = OrderedDict({\n           \"batch\": np.empty(0),\n           \"variable\": np.array([\"wind_cf\"]),\n           \"lat\": self.lat,\n           \"lon\": self.lon,\n       })\n       oc[\"batch\"] = input_coords[\"batch\"]\n       return oc\n   @batch_func()\n   def __call__(self, x: torch.Tensor, coords: CoordSystem):\n       oc = self.output_coords(coords)\n       u, v = x[..., 0:1, :, :], x[..., 1:2, :, :]\n       ws10 = torch.sqrt(u * u + v * v)\n       ws = ws10 * (self.hub / 10.0) ** self.alpha\n       ramp = (ws ** 3 - self.cut_in ** 3) / (self.rated ** 3 - self.cut_in ** 3)\n       cf = torch.zeros_like(ws)\n       cf = torch.where((ws >= self.cut_in) & (ws < self.rated), ramp.clamp(0, 1), cf)\n       cf = torch.where((ws >= self.rated) & (ws <= self.cut_out), torch.ones_like(cf), cf)\n       return cf, oc\nclass VariableScaledNoise:\n   \"\"\"Spatially correlated noise with per-variable amplitudes + control member.\"\"\"\n   def __init__(self, amplitudes: dict, default: float = 0.0, control_member: bool = True):\n       self.amplitudes, self.default, self.control = amplitudes, default, control_member\n       try:\n           from earth2studio.perturbation import SphericalGaussian\n           self.sampler, self.kind = SphericalGaussian(noise_amplitude=1.0), \"SphericalGaussian\"\n       except Exception:\n           from earth2studio.perturbation import Brown\n           self.sampler, self.kind = Brown(noise_amplitude=1.0), \"Brown\"\n   def __call__(self, x: torch.Tensor, coords: CoordSystem):\n       noise, _ = self.sampler(torch.zeros_like(x), coords)\n       vax = list(coords).index(\"variable\")\n       amps = torch.tensor([self.amplitudes.get(str(v), self.default)\n                            for v in coords[\"variable\"]], device=x.device, dtype=x.dtype)\n       shape = [1] * x.ndim; shape[vax] = amps.numel()\n       pert = noise * amps.reshape(shape)\n       if self.control and \"ensemble\" in coords:\n           eax = list(coords).index(\"ensemble\")\n           mask = torch.tensor((np.asarray(coords[\"ensemble\"]) != 0).astype(np.float32),\n                               device=x.device, dtype=x.dtype)\n           mshape = [1] * x.ndim; mshape[eax] = mask.numel()\n           pert = pert * mask.reshape(mshape)\n       return x + pert, coords\n```\n\nWe 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.\n\n``` python\ndef write_vars(io, x, coords, names):\n   \"\"\"Write selected channels of a (…, variable, lat, lon) tensor to the IO backend.\"\"\"\n   vax = list(coords).index(\"variable\")\n   sub = OrderedDict((k, v) for k, v in coords.items() if k != \"variable\")\n   for name in names:\n       hit = np.where(np.asarray(coords[\"variable\"]) == name)[0]\n       if hit.size:\n           io.write(x.select(vax, int(hit[0])).cpu(), sub, name)\ndef run_ensemble(time, nsteps, nensemble, batch_size, prognostic, diagnostic,\n                perturbation, data, io, save_vars, device):\n   time = to_time_array(time)\n   ic = prognostic.input_coords()\n   x0, c0 = fetch_data(source=data, time=time, lead_time=ic[\"lead_time\"],\n                       variable=ic[\"variable\"], device=device)\n   print(f\"Initial condition tensor: {tuple(x0.shape)}  dims={list(c0)}\")\n   oc = prognostic.output_coords(ic)\n   dt = oc[\"lead_time\"]\n   prog_vars = [v for v in save_vars if v in set(map(str, oc[\"variable\"]))]\n   total = OrderedDict({\n       \"ensemble\": np.arange(nensemble),\n       \"time\": time,\n       \"lead_time\": np.asarray([dt * i for i in range(nsteps + 1)]).flatten(),\n       \"lat\": oc[\"lat\"],\n       \"lon\": oc[\"lon\"],\n   })\n   io.add_array(total, prog_vars + [\"wind_cf\"])\n   dx_target = OrderedDict((k, v) for k, v in diagnostic.input_coords().items() if k != \"batch\")\n   nbatch = int(np.ceil(nensemble / batch_size))\n   with torch.inference_mode():\n       for b in tqdm(range(nbatch), desc=\"ensemble batches\"):\n           lo = b * batch_size\n           n = min(batch_size, nensemble - lo)\n           x = x0.unsqueeze(0).repeat(n, *([1] * x0.ndim))\n           coords = OrderedDict({\"ensemble\": np.arange(lo, lo + n), **c0})\n           x, coords = perturbation(x, coords)\n           x, coords = map_coords(x, coords, ic)\n           for step, (xs, cs) in enumerate(prognostic.create_iterator(x, coords)):\n               write_vars(io, xs, cs, prog_vars)\n               xw, cw = map_coords(xs, cs, dx_target)\n               xw, cw = diagnostic(xw, cw)\n               write_vars(io, xw, cw, [\"wind_cf\"])\n               if step >= nsteps:\n                   break\n           torch.cuda.empty_cache() if device.type == \"cuda\" else None\n   return io\nmodel = FCN.load_model(FCN.load_default_package()).to(DEVICE)\ngrid = model.output_coords(model.input_coords())\nLAT, LON = grid[\"lat\"], grid[\"lon\"]\ndiagnostic = WindPowerCF(LAT, LON).to(DEVICE)\npert = VariableScaledNoise(\n   amplitudes={\"t2m\": 0.20, \"t850\": 0.20, \"z500\": 40.0, \"z850\": 25.0,\n               \"u10m\": 0.25, \"v10m\": 0.25, \"u500\": 0.40, \"v500\": 0.40, \"tcwv\": 0.30},\n   default=0.0, control_member=True)\nprint(f\"Perturbation sampler: {pert.kind}\")\nio = ZarrBackend(file_name=\"outputs/e2s_ensemble.zarr\",\n                chunks={\"ensemble\": 1, \"time\": 1, \"lead_time\": 1},\n                backend_kwargs={\"overwrite\": True})\nio = run_ensemble([INIT_STR], NSTEPS, NENSEMBLE, BATCH_SIZE,\n                 model, diagnostic, pert, GFS(), io, SAVE_VARS, DEVICE)\nprint(io.root.tree())\n```\n\nWe 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.\n\n```\nleads = np.asarray(io[\"lead_time\"][:]).astype(\"timedelta64[ns]\")\nlead_h = leads.astype(\"timedelta64[h]\").astype(int)\nvalid = to_time_array([INIT_STR])[0] + leads\ntruth, tc = fetch_data(source=GFS(), time=valid,\n                      lead_time=np.array([np.timedelta64(0, \"h\")]),\n                      variable=np.array(VERIFY_VARS), device=\"cpu\")\ntruth = truth[:, 0]\nw = torch.cos(torch.deg2rad(torch.as_tensor(np.asarray(LAT), dtype=torch.float32)))\nw2d = w[:, None].expand(len(LAT), len(LON)).contiguous()\nmcoords = OrderedDict({\"lead_time\": leads, \"lat\": np.asarray(LAT), \"lon\": np.asarray(LON)})\ndef fair_crps(ens, obs, weights):\n   \"\"\"Fair (unbiased) CRPS, lat-weighted. ens: (M, lat, lon), obs: (lat, lon).\"\"\"\n   M = ens.shape[0]\n   wn = weights / weights.sum()\n   skill = ((ens - obs).abs() * wn).sum(dim=(-2, -1)).mean()\n   spread = torch.zeros((), dtype=ens.dtype)\n   for i in range(M):\n       spread = spread + ((ens[i] - ens).abs() * wn).sum(dim=(-2, -1)).sum()\n   return (skill - spread / (2 * M * (M - 1))).item()\nscores = {}\nfor k, var in enumerate(VERIFY_VARS):\n   fc = torch.as_tensor(np.asarray(io[var][:]))[:, 0].float()\n   ob = truth[:, k].float()\n   mean = fc.mean(0)\n   try:\n       metric = rmse(reduction_dimensions=[\"lat\", \"lon\"], weights=w2d)\n       r, _ = metric(mean, mcoords, ob, mcoords)\n       r = r.numpy()\n   except Exception as e:\n       print(f\"(built-in rmse unavailable: {e})\")\n       wn = (w2d / w2d.sum())\n       r = torch.sqrt((((mean - ob) ** 2) * wn).sum(dim=(-2, -1))).numpy()\n   wn = w2d / w2d.sum()\n   spread = torch.sqrt((fc.var(0, unbiased=True) * wn).sum(dim=(-2, -1))).numpy()\n   crps = np.array([fair_crps(fc[:, t], ob[t], w2d) for t in range(fc.shape[1])])\n   scores[var] = dict(rmse=r, spread=spread, crps=crps, fc=fc, obs=ob, mean=mean)\n   print(f\"\\n=== {var} ===\")\n   print(f\"{'lead[h]':>8}{'RMSE':>12}{'spread':>12}{'ratio':>9}{'CRPS':>12}\")\n   for t in range(len(lead_h)):\n       ratio = spread[t] / r[t] if r[t] > 0 else np.nan\n       print(f\"{lead_h[t]:>8}{r[t]:>12.3f}{spread[t]:>12.3f}{ratio:>9.2f}{crps[t]:>12.3f}\")\n```\n\nWe 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.\n\n```\nlat_np, lon_np = np.asarray(LAT), np.asarray(LON)\nilat = int(np.argmin(np.abs(lat_np - POI[1])))\nilon = int(np.argmin(np.abs(lon_np - (POI[2] % 360))))\nlast = -1\nd = scores[\"t2m\"]\nfields = [(d[\"mean\"][last].numpy() - 273.15, \"ensemble mean t2m [C]\", \"RdBu_r\", None),\n         (d[\"fc\"][:, last].std(0).numpy(), \"ensemble spread [K]\", \"magma\", None),\n         (d[\"obs\"][last].numpy() - 273.15, \"GFS analysis [C]\", \"RdBu_r\", None),\n         ((d[\"mean\"][last] - d[\"obs\"][last]).numpy(), \"mean error [K]\", \"coolwarm\", 5)]\nfig, axs = plt.subplots(2, 2, figsize=(15, 7), constrained_layout=True)\nfor ax, (f, title, cmap, lim) in zip(axs.ravel(), fields):\n   kw = dict(vmin=-lim, vmax=lim) if lim else {}\n   im = ax.pcolormesh(lon_np, lat_np, f, cmap=cmap, shading=\"auto\", **kw)\n   ax.set_title(f\"{title} — +{lead_h[last]} h\"); plt.colorbar(im, ax=ax, shrink=0.85)\nplt.show()\nz = scores[\"z500\"][\"fc\"][:, last].numpy() / 9.81\nla = (lat_np > 25) & (lat_np < 75)\nlo = (lon_np > 280) | (lon_np < 40)\nlon_shift = np.where(lon_np > 180, lon_np - 360, lon_np)\norder = np.argsort(lon_shift[lo])\nplt.figure(figsize=(11, 5))\nfor m in range(z.shape[0]):\n   sub = z[m][np.ix_(la, lo)][:, order]\n   plt.contour(lon_shift[lo][order], lat_np[la], sub, levels=[5520],\n               colors=[\"k\" if m == 0 else \"C0\"], linewidths=[2.0 if m == 0 else 0.8])\nzo = scores[\"z500\"][\"obs\"][last].numpy() / 9.81\nplt.contour(lon_shift[lo][order], lat_np[la], zo[np.ix_(la, lo)][:, order],\n           levels=[5520], colors=\"crimson\", linewidths=2.5)\nplt.title(f\"z500 5520 m spaghetti at +{lead_h[last]} h \"\n         f\"(black=control, blue=members, red=GFS analysis)\")\nplt.xlabel(\"lon\"); plt.ylabel(\"lat\"); plt.show()\nt2m_pt = scores[\"t2m\"][\"fc\"][:, :, ilat, ilon].numpy() - 273.15\nobs_pt = scores[\"t2m\"][\"obs\"][:, ilat, ilon].numpy() - 273.15\ncf_pt = np.asarray(io[\"wind_cf\"][:])[:, 0, :, ilat, ilon]\nfig, (a1, a2) = plt.subplots(1, 2, figsize=(14, 4))\na1.fill_between(lead_h, t2m_pt.min(0), t2m_pt.max(0), alpha=0.25, label=\"member range\")\na1.plot(lead_h, t2m_pt.mean(0), \"o-\", label=\"ensemble mean\")\na1.plot(lead_h, t2m_pt[0], \"k--\", label=\"control\")\na1.plot(lead_h, obs_pt, \"r^-\", label=\"GFS analysis\")\na1.set_title(f\"2 m temperature — {POI[0]}\"); a1.set_xlabel(\"lead [h]\"); a1.set_ylabel(\"C\")\na1.legend(); a1.grid(alpha=.3)\na2.fill_between(lead_h, cf_pt.min(0), cf_pt.max(0), alpha=0.25, color=\"seagreen\")\na2.plot(lead_h, cf_pt.mean(0), \"o-\", color=\"seagreen\")\na2.set_title(f\"wind capacity factor (custom diagnostic) — {POI[0]}\")\na2.set_xlabel(\"lead [h]\"); a2.set_ylim(0, 1); a2.grid(alpha=.3)\nplt.tight_layout(); plt.show()\nfig, axs = plt.subplots(1, len(VERIFY_VARS), figsize=(5 * len(VERIFY_VARS), 3.6))\nfor ax, var in zip(np.atleast_1d(axs), VERIFY_VARS):\n   s = scores[var]\n   ax.plot(lead_h, s[\"rmse\"], \"o-\", label=\"RMSE (ens. mean)\")\n   ax.plot(lead_h, s[\"spread\"], \"s--\", label=\"spread\")\n   ax.plot(lead_h, s[\"crps\"], \"^:\", label=\"fair CRPS\")\n   ax.set_title(var); ax.set_xlabel(\"lead [h]\"); ax.grid(alpha=.3); ax.legend(fontsize=8)\nplt.tight_layout(); plt.show()\nimport xarray as xr\nds = xr.open_zarr(\"outputs/e2s_ensemble.zarr\")\nprint(ds)\n```\n\nWe 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.\n\nIn 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.\n\nCheck out the ** FULL CODES here.** Also, feel free to follow us on\n\n**and don’t forget to join our**[Twitter](https://x.com/intent/follow?screen_name=marktechpost)\n\n**and Subscribe to**\n\n[150k+ML SubReddit](https://www.reddit.com/r/machinelearningnews/)**. Wait! are you on telegram?**\n\n[our Newsletter](https://magic.beehiiv.com/v1/f5e63dd4-5653-4f09-83e2-321a8b1ba526?email={{email}})\n\n[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)\n\nSana 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.", "url": "https://wpnews.pro/news/building-custom-batched-ensemble-weather-forecasting-with-nvidia-earth2studio", "canonical_source": "https://www.marktechpost.com/2026/08/29/building-custom-batched-ensemble-weather-forecasting-with-nvidia-earth2studio/", "published_at": "2026-08-29 18:57:26+00:00", "updated_at": "2026-08-29 19:18:20.605550+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "ai-tools", "ai-research"], "entities": ["NVIDIA", "Earth2Studio", "Marktechpost", "FCN", "GFS"], "alternates": {"html": "https://wpnews.pro/news/building-custom-batched-ensemble-weather-forecasting-with-nvidia-earth2studio", "markdown": "https://wpnews.pro/news/building-custom-batched-ensemble-weather-forecasting-with-nvidia-earth2studio.md", "text": "https://wpnews.pro/news/building-custom-batched-ensemble-weather-forecasting-with-nvidia-earth2studio.txt", "jsonld": "https://wpnews.pro/news/building-custom-batched-ensemble-weather-forecasting-with-nvidia-earth2studio.jsonld"}}