Building Custom Batched Ensemble Weather Forecasting with NVIDIA Earth2Studio

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.

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.

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.

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.

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.