NVIDIA Earth2Studio: Custom Batched Ensemble Forecasting Pipeline
In this article
NVIDIA's Earth2Studio library lets practitioners bypass its high-level convenience runners and compose fully custom ensemble pipelines from individual building blocks — perturbation samplers, coordinate-aware iterators, diagnostic chains, and Zarr I/O — all within a single Colab session. The construction pattern, where pipeline architecture rather than model substitution drives capability, is increasingly the dominant leverage point in applied scientific AI.
Environment Setup and Configuration
Installation targets the fcn, data, perturbation, and statistics extras of earth2studio, writing a constraint file that pins the pre-existing torch and numpy versions before invoking pip, which prevents Colab's CUDA-enabled PyTorch from being overwritten. The model cache is pointed to /content/e2s_cache. Core configuration sets the ensemble size at 8 members, batch size at 2, autoregressive steps at 8, and the point of interest at New Delhi (28.61°N, 77.21°E). Initialization time is computed as seven days prior to the current UTC clock, and a runtime warning fires if no GPU is detected.
Custom Diagnostic and Perturbation Components
The WindPowerCF module is a torch.nn.Module that accepts u10m and v10m fields and outputs a turbine capacity factor between 0 and 1. It extrapolates 10-meter wind speed to hub height using a power-law shear exponent of 0.143 with a configurable hub height defaulting to 100 m, then applies a cubic ramp between a cut-in speed of 3.0 m/s and a rated speed of 12.0 m/s, a unity plateau up to cut-out at 25.0 m/s, and zero outside those bounds. Coordinate compatibility is enforced through handshake_dim and handshake_coords before any tensor operation, and both the forward pass and coordinate mapping are decorated with @batch_func() and @batch_coords() to stay inside Earth2Studio's batching contract.
The VariableScaledNoise class wraps either SphericalGaussian or Brown from earth2studio.perturbation (falling back to the latter if the former is unavailable) with per-variable amplitude scaling. The amplitude dictionary assigns physically motivated magnitudes: 0.20 K for t2m and t850, 40.0 m²/s² for z500, 25.0 m²/s² for z850, 0.25 m/s for u10m and v10m, 0.40 m/s for u500 and v500, and 0.30 kg/m² for tcwv. A mask zeroes perturbations on ensemble member zero, preserving it as an unperturbed control.
Ensemble Execution Loop and I/O
| Pipeline Stage | Earth2Studio API Used | Key Parameter / Output |
|---|---|---|
| Initial condition fetch | fetch_data + GFS() |
Variables defined by FCN input_coords() |
| Perturbation | VariableScaledNoise (custom) |
8 members, member 0 unperturbed control |
| Coordinate alignment | map_coords |
Aligns batch coords before each model call |
| Autoregressive stepping | FCN.create_iterator |
8 steps; iterator yields (tensor, CoordSystem) |
| Wind diagnostic | WindPowerCF (custom) |
Outputs wind_cf field per step |
| Persistence | ZarrBackend |
Chunks: ensemble=1, time=1, lead_time=1 |
| Verification | Lat-weighted RMSE, fair CRPS, spread-skill ratio | Variables: t2m, z500, u10m |
The run_ensemble function batches the 8-member ensemble into 4 iterations of 2 members each. For each batch it repeats the initial condition tensor, applies perturbations, then calls prognostic.create_iterator and iterates up to NSTEPS. At every step, selected channels from SAVE_VARS (t2m, z500, u10m, v10m, tcwv) are written to the Zarr store, and the wind diagnostic is chained immediately after via a second map_coords call. CUDA cache is flushed between batches. The Zarr backend is configured with overwrite=True and chunked at unit size along the ensemble, time, and lead-time axes.
Verification and Visualization
GFS analyses for every forecast-valid time are fetched at zero lead time and used as ground truth. Latitude weighting is cosine of latitude. RMSE is computed on the ensemble mean using earth2studio.statistics.rmse with a fallback to a manual implementation when the built-in signature is incompatible. Fair (unbiased) CRPS subtracts ensemble spread from the mean absolute error, scaled by 1 / (2 * M * (M − 1)) where M is the ensemble size. Spread-skill ratios are printed per lead hour for t2m, z500, and u10m.
Visualization covers four panels: ensemble mean 2 m temperature, ensemble spread, GFS analysis, and mean error — all at the final lead time. A spaghetti contour plot draws the 5520 m geopotential-height contour for each member over the 25°–75°N, 280°E–40°E belt, with the control in black, perturbed members in blue, and the GFS analysis in crimson. A fan chart for New Delhi shows ensemble min/max, ensemble mean, the control member, and GFS analysis for 2 m temperature alongside a separate wind capacity factor panel. The completed dataset is opened with xr.open_zarr for downstream inspection.
Earth2Studio's component API is now expressive enough that the primary skill required is interface composition rather than model training — a pattern consistent with systems-engineering gains that now rival scaling. For renewable-energy and climate-risk teams, attaching physics-based diagnostics like WindPowerCF directly inside the ensemble loop — rather than in a post-processing step — reduces latency between forecast generation and decision-relevant output. As ensemble sizes grow, the batching and chunking conventions established here become the load-bearing infrastructure.