Earth2studio
Open-source deep-learning framework for exploring, building and deploying AI weather/climate workflows.
Install / Use
/learn @NVIDIA/Earth2studioREADME
NVIDIA Earth2Studio
[![python version][e2studio_python_img]][e2studio_python_url] [![license][e2studio_license_img]][e2studio_license_url] [![coverage][e2studio_cov_img]][e2studio_cov_url] [![mypy][e2studio_mypy_img]][e2studio_mypy_url] [![format][e2studio_format_img]][e2studio_format_url] [![ruff][e2studio_ruff_img]][e2studio_ruff_url] [![uv][e2studio_uv_img]][e2studio_uv_url]
Earth2Studio is a Python-based package designed to get users up and running with AI Earth system models fast. Our mission is to enable everyone to build, research and explore AI driven weather and climate science.
<!-- markdownlint-disable MD036 -->- Earth2Studio Documentation -
<!-- markdownlint-enable MD036 -->[Install][e2studio_install_url] | [User-Guide][e2studio_userguide_url] | [Examples][e2studio_examples_url] | [API][e2studio_api_url]
Quick start
Running AI weather prediction can be done with just a few lines of code.
- For detailed installation steps, including model-specific installations, see the [install guide][e2studio_install_url].
- See the [examples][e2studio_examples_url] gallery providing different inference workflow samples.
- Swap out [data sources][e2studio_data_api] or [models][e2studio_px_api] depending on your use case!
NVIDIA FourCastNet3
from earth2studio.models.px import FCN3
from earth2studio.data import GFS
from earth2studio.io import ZarrBackend
from earth2studio.run import deterministic as run
model = FCN3.load_model(FCN3.load_default_package())
data = GFS()
io = ZarrBackend("outputs/fcn3_forecast.zarr")
run(["2025-01-01T00:00:00"], 10, model, data, io)
ECMWF AIFS
from earth2studio.models.px import AIFS
from earth2studio.data import IFS
from earth2studio.io import ZarrBackend
from earth2studio.run import deterministic as run
model = AIFS.load_model(AIFS.load_default_package())
data = IFS()
io = ZarrBackend("outputs/aifs_forecast.zarr")
run(["2025-01-01T00:00:00"], 10, model, data, io)
Google Graphcast
from earth2studio.models.px import GraphCastOperational
from earth2studio.data import GFS
from earth2studio.io import ZarrBackend
from earth2studio.run import deterministic as run
package = GraphCastOperational.load_default_package()
model = GraphCastOperational.load_model(package)
data = GFS()
io = ZarrBackend("outputs/graphcast_operational_forecast.zarr")
run(["2025-01-01T00:00:00"], 4, model, data, io)
[!IMPORTANT] Earth2Studio is an interface to third‑party models, checkpoints, and datasets. Licenses for these assets are owned by their providers. Ensure you have the rights to download, use, and (if applicable) redistribute each model and dataset. Links to the original license and source are often provided in the API docs for each model/data source.
Latest News
- StormScope, NVIDIA's latest regional nowcasting model, is now available in Earth2Studio with a dedicated inference example to demonstrate its use for predicting CONUS satellite and radar data.
- Atlas, NVIDIA's latest medium range model, is now available in Earth2Studio providing state-of-the-art accuracy for medium range global forecasts.
- CorrDiff for CMIP6 to ERA5 is a novel generative downscaling model to generate ERA5 fields from CMIP data enabling users to run ERA5 based prognostic and diagnostic models on future climate simulations.
- ECMWF AIFSENS model wrapper added, an ensemble-based probablistic data driven forecast model developed by the European Centre for Medium-Range Weather Forecasts (ECMWF).
For a complete list of latest features and improvements see the changelog.
Overview
Earth2Studio is an AI inference pipeline toolkit focused on weather and climate applications that is designed to ride on top of different AI frameworks, model architectures, data sources and SciML tooling while providing a unified API.
<div align="center">
The composability of the different core components in Earth2Studio easily allows the development and deployment of increasingly complex pipelines that may chain multiple data sources, AI models and other modules together.
<div align="center">
The unified ecosystem of Earth2Studio provides users the opportunity to rapidly swap out components for alternatives. In addition to the largest model zoo of weather/climate AI models, Earth2Studio is packed with useful functionality such as optimized data access to cloud data stores, statistical operations and more to accelerate your pipelines.
<div align="center">
Earth-2 Open Models
Access state of the art Nvidia open models for climate and weather: Earth-2 Open Models. For training recipes for these models, see the [PhysicsNeMo repository][physicsnemo_repo_url].
Features
Earth2Studio package focuses on supplying you the tools to build your own workflows, pipelines, APIs, or packages using modular components including:
<details> <summary>Prognostic Models</summary>[Prognostic models][e2studio_px_url] in Earth2Studio perform time integration, taking atmospheric fields at a specific time and auto-regressively predicting the same fields into the future (typically 6 hours per step), enabling both single time-step predictions and extended time-series forecasting.
Earth2Studio maintains the largest collection of pre-trained state-of-the-art AI weather/climate models ranging from global forecast models to regional specialized models, covering various resolutions, architectures, and forecasting capabilities to suit different computational and accuracy requirements.
Available models include but are not limited to:
| Model | Resolution | Architecture | Time Step | Coverage | |-------|------------|--------------|-----------|----------| | GraphCast Small | 1.0° | Graph Neural Network | 6h | Global | | GraphCast Operational | 0.25° | Graph Neural Network | 6h | Global | | Pangu 3hr | 0.25° | Transformer | 3h | Global | | Pangu 6hr | 0.25° | Transformer | 6h | Global | | Pangu 24hr | 0.25° | Transformer | 24h | Global | | Aurora | 0.25° | Transformer | 6h | Global | | FuXi | 0.25° | Transformer | 6h | Global | | AIFS | 0.25° | Transformer | 6h | Global | | AIFS Ensemble | 0.25° | Transformer Ensemble | 6h | Global | | StormCast | 3km | Diffusion + Regression | 1h | Regional (US) | | SFNO | 0.25° | Neural Operator | 6h | Global | | DLESyM | 0.25° | Convolutional | 6h | Global |
For a complete list, see the [prognostic model API docs][e2studio_px_api].
</details> <details> <summary>Diagnostic Models</summary>[Diagnostic models][e2studio_dx_url] in Earth2Studio perform time-independent transformations, typically taking geospatial fields at a specific time and predicting new derived quantities without performing time integration enabling users to build pipelines to predict specific quantities of interest that may not be provided by forecasting models.
Earth2Studio contains a growing collection of specialized diagnostic models for various phenomena including precipitation prediction, tropical cyclone tracking, solar radiation estimation, wind gust forecasting, and more.
Available diagnostics include but are not limited to:
| Model | Resolution | Architecture | Coverage | Output | |-------|------------|--------------|----------|--------| | PrecipitationAFNO | 0.25° | Neural Operator | Global | Total precipitation | | SolarRadiationAFNO1H | 0.25° | Neural Operator | Global | Surface solar radiation | | WindgustAFNO | 0.25° | AFNO | Global | Maximum wind gust | | TCTrackerVitart | 0.25° | Algorithmic | Global | TC tracks & properties | | CBottleInfill | 100km | Diffusion | Global | Global climate sample | | CBottleSR | 5km | Diffusion | Regional / Global | High-res climate | | CorrDiff | Variable | Diffusion | Regional | Fine-scale weather | | CorrDiffTaiwan | 2km | Diffusion | Regional (Taiwan) | Taiwan fine-scale weather |
For a complete list, see the [diagnostic model API docs][e2studio_dx_api].
</details> <details> <summary>Datasources</summary>[Data sources][e2studio_data_url] in Earth2Studio provide a standardized API for accessing weather and climate datasets from various providers (numerical models, data assimilation results, and AI-generated data), enabling seamless integration of initial conditions for model inference and validation data for scoring across different data formats and storage systems.
Earth2Studio includes data sources ranging from operational weather models (GFS, HRRR, IFS) and reanalysis datasets (ERA5 via ARCO, CDS) to AI-generated climate data (cBottle) and local file systems. Fetching data is
