Earth2Studio: a Python inference toolkit for swapping AI weather models

The problem Earth2Studio solves for AI weather work

An AI weather model is rarely usable on its own. Each one expects its own input variables on its own grid, ships its own checkpoint format, and produces output you then have to write somewhere and compare against something. Swap FourCastNet3 for GraphCast Operational and the surrounding code changes too. Earth2Studio's stated goal is to remove that surrounding code: it describes itself as an "AI inference pipeline toolkit" that rides on top of different AI frameworks, model architectures and data sources while exposing a unified API. The intended user is a researcher or engineer who wants to run several models over the same initial conditions, or build a pipeline that chains data sources and models, without maintaining a separate adapter per model. It is not a training framework. The README directs anyone wanting training recipes for the Earth-2 open models to the PhysicsNeMo repository.

How the model, data and IO objects fit together

The architecture is visible in the quick-start snippets: three interchangeable objects are handed to a run function. A prognostic model comes from earth2studio.models.px, a data source from earth2studio.data, and an output store from earth2studio.io. The run module supplies the loop. In the FourCastNet3 example the model is built with FCN3.load_model(FCN3.load_default_package()), the data source is GFS(), the store is ZarrBackend("outputs/fcn3_forecast.zarr"), and deterministic is called with a list of initialisation times, a step count, and those three objects. The same shape holds for the ECMWF AIFS example, which pairs AIFS with the IFS data source, and for GraphCast Operational, which pairs with GFS and runs four steps. That is the whole composability claim: the run signature does not change when the components do, so a pipeline can mix data sources, chain models, or write to a different backend by substituting one argument. The repository also ships a serve/ directory and a recipes/ directory alongside the earth2studio package, and the examples tree is organised by workflow rather than by model: 01_getting_started, 02_medium_range, 03_downscaling, 04_nowcasting, 05_data_assimilation, 06_seasonal, 07_misc and 08_extend.

Installing Earth2Studio and running a first forecast

The repository's Makefile defines the install target as two uv sync calls, the second pulling the all and aifs extras. The package requires Python 3.11 or newer and below 3.15, and torch is a base dependency, so a GPU environment is the expected default; the release notes state that as of version 0.14.0 the TOML default installs target CUDA 13.

make install

If you prefer to drive uv yourself, the same extras can be requested directly. The data extra is where the heavier source-specific clients live; pyproject.toml shows arraylake gated on python_version>='3.12' inside that group, so the extra you can install depends on your interpreter.

uv sync --extra all --extra aifs

The README also documents an agent-assisted path, which installs a set of Earth2Studio skills and then lets a coding agent recommend a model or scaffold a forecast. The commands are published as npx invocations against NVIDIA/skills.

npx skills add NVIDIA/skills --skill earth2studio-install
npx skills add NVIDIA/skills --skill earth2studio-deterministic-forecast

Once the environment is up, the first real use is the FourCastNet3 snippet from the README. It loads the default package, fetches initial conditions from GFS for one timestamp, and writes ten steps into a Zarr store on disk. Expect a model download on first run, since load_default_package pulls from a hosted package, and expect the output directory to be created relative to your working directory.

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)

Third-party licences are the real constraint

Earth2Studio itself is Apache-2.0, and the pyproject.toml declares that licence and the matching classifier. That covers the framework code only. The README carries an explicit warning that Earth2Studio is an interface to third-party models, checkpoints and datasets, that licences for those assets are owned by their providers, and that you must ensure you have the rights to download, use and redistribute each model and dataset. Links to the original licence and source are often provided in the API docs for each model or data source. So the licence question is per component, not per framework, and it is not something a package installer resolves for you. If you plan to redistribute derived outputs or ship a hosted service, the component list is the thing to audit. The .licenses/ directory at the repository root suggests the project tracks this internally, but it does not transfer any rights to you.

Where Earth2Studio is the wrong tool

Training is out of scope. The README points to PhysicsNeMo for training recipes for the Earth-2 open models, and nothing in the pyproject.toml or the repository layout suggests a training loop. If your goal is to fine-tune a checkpoint on your own reanalysis, this toolkit is the inference half of the problem. Second, the abstraction has a cost: every model you want that is not already wrapped means writing an adapter against the same interface, and the value of the unified API drops sharply if you only ever run one model with one data source. Third, the dependency surface is wide and partly optional. The base install pulls in torch, xarray, zarr, netCDF4, h5py, pygrib, s3fs and gcsfs, and pyproject.toml pins netCDF4 below 1.7.3 with a link to an upstream issue, which is the kind of pin that can collide with another package in a shared environment. The data extra adds source-specific clients, and at least one of them is gated by Python version. A minimal container that only needs one model will still carry all of this.

Earth2Studio compared with a single-model repository

The obvious alternative is to use a model vendor's own inference repository directly. GraphCast and AIFS both have upstream code, and going straight to it means you get the reference implementation with its own assumptions about input preparation and output format, plus its own release cadence. Earth2Studio's difference is the interface boundary: it treats the model as one object among three, so the same run call drives FCN3, AIFS or GraphCast Operational, and the data source and output store are equally swappable. That matters when the comparison itself is the work, for instance when you want the same initial conditions from GFS feeding two different models into two Zarr stores. It matters much less when you have already chosen a model and only need it to run. The second axis of difference is data access: the project lists cloud-optimised sources including the Dynamical.org suite reading from anonymous Icechunk repositories and EarthMover ERA5 and IFS sources hosted by BrightBand, which is a different proposition from downloading GRIB files yourself.

Maintenance, releases and upgrade cost

The last push to the default branch was on 2026-07-30, the same day as the 0.17.0 release, and the two releases before it landed on 2026-06-29 (0.16.0) and 2026-05-26 (0.15.0). That is a roughly monthly minor-release cadence across the visible window. The repository is not archived. For an adopter, the practical consequence of that cadence is that the API surface is still moving: a minor version bump is where new models and data sources appear, and the changelog is the file to read before upgrading rather than assuming compatibility. The project keeps a CHANGELOG.md at the root and a CITATION.cff, and the version is dynamic in pyproject.toml, so the installed package version is the reliable thing to record. Upgrades also mean re-checking extras, because the optional dependency groups are where model-specific clients live and where Python-version gating appears.

Editorial conclusion

Adopt Earth2Studio if you need to compare AI weather models on the same data and the same output store without rewriting glue code for each one, and if you have a CUDA-capable GPU plus the rights to the checkpoints you download. Skip it if you want a trained model of your own, since the repository points to PhysicsNeMo for training recipes, or if a single fixed model already satisfies your pipeline. Before committing, verify the licence and source link for each model and data source you plan to use, and check whether the data extra you need (for example aifs) is available for your Python version, because the repository pins some extras by Python version.

Frequently asked questions

Official sources

1. Official documentation
2. Official README
3. Project repository
4. Release notes