NVIDIA Earth2Studio

The HENS recipe implements a multi-checkpoint inference pipeline designed for large-scale ensemble weather forecasting. It is specifically built to recover the HENS method, which uses Spherical Fourier Neural Operators (SFNO) and customized bred vectors for initial condition perturbations to provide calibrated ensemble forecasting.

Key capabilities include:

• Multi-GPU inference: Parallel processing of model checkpoints.
• Uncertainty Quantification: Improved representation of extreme weather events and sampling of forecast distribution tails.
• Reproducibility: Ability to reproduce individual inference batches.
• Integrated Diagnostics: Support for one or more Diagnostic models within the pipeline.
• Specialized Tracking: Includes tropical cyclone tracking.
• Efficient Storage: Supports regional output to minimize data footprint.

The Earth2Studio Python Client SDK provides two primary interfaces for interacting with the Earth2Studio REST API server:

1. RemoteEarth2Workflow: A high-level interface designed for seamless integration with Earth2Studio and Xarray. This is the recommended interface for most users.
2. Earth2StudioClient: A low-level interface providing direct API access and result file management. Use this for custom workflows, manual file downloads, or asynchronous request handling.

Earth2Studio provides examples for two primary downscaling workflows to produce high-resolution outputs from global model forecasts:

1. Generative Downscaling: Using generative models to transform low-resolution global forecasts into high-resolution local representations.
2. Ensemble Downscaling: Utilizing ensemble methods to account for uncertainty when producing high-resolution outputs.

These examples are located in the examples/03_downscaling/ directory of the repository.

The earth2studio.statistics module provides various statistical and metric calculation functions for analyzing inference data. These functions are used to evaluate the performance and accuracy of weather and climate models.

Available metrics include:

• Error Metrics: mae (Mean Absolute Error), rmse (Root Mean Square Error), lsd (Logarithmic Score Distance).
• Probabilistic Scores: brier (Brier Score), crps (Continuous Ranked Probability Score), energy_score (Energy Score), fss (Fractions Skill Score).
• Distributional Metrics: acc (Accuracy), rank_histogram, spread_skill_ratio.
• Descriptive Statistics: mean, std (Standard Deviation), variance.

import earth2studio.statistics as stats

# Example usage (conceptual based on available symbols):
# error = stats.rmse(prediction, ground_truth)
# score = stats.crps(prediction, ground_truth)

The earth2studio.data module provides several data sources that deliver tabular data in the form of DataFrames. These sources cover various domains including analysis, reanalysis, observations, and simulations, and support products such as wind, precipitation, temperature, and more.

Available DataFrame data sources include:

• Global Climate/Weather Data: data.GHCNDaily, data.GHCNHourly, data.IEM_ASOS, data.ISD
• Satellite/Remote Sensing: data.GOESGLM, data.JPSS_ATMS, data.JPSS_CRIS, data.MetOpAMSUA, data.MetOpAVHRR, data.MetOpIASI, data.MetOpMHS
• Tropical Cyclone Data: data.IBTrACS
• Observation Converters: data.NNJAObsConv, data.NNJAObsSat, data.NomadsGDASObsConv, data.UFSObsConv, data.UFSObsSat
• Testing/Mock Data: data.RandomDataFrame

IO backends in earth2studio.io are used to write model outputs (forecast results, ensemble data, etc.) to disk or memory. Unlike input data handling (managed by earth2studio.data), output handling requires providing an array_name to distinguish between different arrays within the backend. This name is required for compatibility with xarray, zarr, and netcdf formats.

Key methods in the IO backend interface include:

• add_array: Adds an array and its associated coordinates to the store.
• write: Explicitly stores the provided data in the backend. This may trigger synchronization if the input tensor is on the GPU.

Earth2Studio provides various data sources used for downloading, caching, and reading weather and climate data. These sources act as APIs into Xarray DataArray objects. They are primarily used for:

1. Fetching initial conditions for model inference.
2. Providing validation data for scoring model performance.

Note on Licensing: Each data source may have its own respective license. Users should familiarize themselves with the specific licenses and limitations of the data source they are using.

Earth2Studio provides six core functional areas for weather and climate AI research:

1. Built-in Workflows: Pre-defined inference workflows to accelerate development.
2. Prognostic Models: A coherent interface for accessing the latest AI weather forecast models.
3. Diagnostic Models: Models used for mapping forecast outputs to other quantities of interest.
4. Datasources: Connectors for integrating both on-premise and remote data stores into your workflows.
5. IO: Utilities for exporting data for post-processing.
6. Statistical Operators: Methods for performing complex uncertainty analysis directly within an inference workflow.

Earth2Studio provides a framework for weather and climate modeling workflows. You can navigate the project through several key documentation areas:

• User Guide: Instructions for installation, setup, and core concepts.
• API Reference: Detailed documentation of the exposed earth2studio modules and functions.
• Examples: Practical demonstrations of common usage patterns.
• Recipes: Complex, end-to-end workflows (e.g., Model Evaluation, Huge Ensembles, Tropical Cyclone Tracking) located in the recipes/ directory of the repository.

To minimize incompatibilities, the following environment is recommended:

Software

• OS: Ubuntu 24.04 LTS
• Python: 3.13
• CUDA: 13.0

Hardware Most models require a GPU with sufficient memory and compute. The recommended baseline for most supported models is:

• GPU: NVIDIA GPU with compute capability $\ge$ 8.9 (e.g., L40S, RTX A6000, H100, B200)
• GPU Memory: $\ge$ 40 GB
• Precision: FP32
• Disk Space: 128 GB