Install

Base Install Limitations

The base pip install does not guarantee all functionality and/or examples are operational due to optional dependencies. We encourage users that face package issues to familiarize themselves with the optional model installs and suggested environment set up for the most complete experience.

Install using Pip

To get the latest release of Earth2Studio, install from the Python index:

pip install earth2studio

Install using uv (recommended)

This package is developed using uv and it’s recommended that users use an uv project for the best install experience:

mkdir earth2studio-project && cd earth2studio-project
uv init --python=3.12
uv add "earth2studio @ git+https://github.com/NVIDIA/earth2studio.git@0.6.0"

uv install

The use of the latest git release tag for the package install with uv is intentional. This will allow uv to handle any complicated dependency conditions and automatically handle url based dependencies. This is not achievable using the pypi registry but makes installing optional packages much easier down the line.

Install Main Branch

To install the latest main branch version of Earth2Studio main branch:

pip install "earth2studio @ git+https://github.com/NVIDIA/earth2studio"

or if you are using uv:

uv add "earth2studio @ git+https://github.com/NVIDIA/earth2studio.git"

Verify Installation

python
# or when using uv
uv run python

>>> import earth2studio
>>> earth2studio.__version__

Optional Dependencies

uv Package Manager

For developers uv package manager should be used. uv is not required for just using Earth2Studio thus both pip and uv commands are included. uv commands assume Earth2Studio has already been added to the project using git source used in the above sections.

uv Package Manager

uv commands assume Earth2Studio has already been added to the project with the git link used in the above sections.

Model Dependencies

Models typically require additional dependencies which are not installed by default. Use the optional install commands to add these dependencies.

Prognostics

Aurora

Note: The shipped Aurora package has a restricted dependency which is incompatible with other Earth2Studio dependencies, thus it is suggested to use the forked variant.

pip

pip install "microsoft-aurora @ git+https://github.com/NickGeneva/aurora.git@ab41cf1de67d5dcc723b96fc9a6219e4b548d181"
pip install earth2studio[aurora]

uv

# Patched fork
uv add earth2studio --extra aurora-fork
# Original package from msc
uv add earth2studio --extra aurora

FourCastNet

pip

pip install earth2studio[fcn]

uv

uv add earth2studio --extra fcn

FengWu

Notes: Requires ONNX GPU Runtime. May need manual install depending on CUDA and Python version.

pip

pip install earth2studio[fengwu]

uv

uv add earth2studio --extra fengwu

FuXi

Notes: Requires ONNX GPU Runtime. May need manual install depending on CUDA version.

pip

pip install earth2studio[fuxi]

uv

uv add earth2studio --extra fuxi

Pangu

Notes: Requires ONNX GPU Runtime. May need manual install depending on CUDA version.

pip

pip install earth2studio[pangu]

uv

uv add earth2studio --extra pangu

SFNO

Notes: Requires Modulus-Makani to be installed manually.

pip

pip install "makani[all] @ git+https://github.com/NickGeneva/modulus-makani.git@3da09f9e52a6393839d73d44262779ac7279bc2f"
pip install earth2studio[sfno]

uv

uv add earth2studio --extra sfno

StormCast

pip

pip install earth2studio[stormcast]

uv

uv add earth2studio --extra stormcast

InterpModAFNO

Notes: Requires a base prognostic model to be installed.

pip

pip install earth2studio[interp-modafno]

uv

uv add earth2studio --extra interp-modafno

Diagnostics

ClimateNet

Notes: No additional dependencies are needed for ClimateNet but included for completeness.

pip

pip install earth2studio[climatenet]

uv

uv add earth2studio --extra climatenet

CorrDiff

Notes: Additional dependencies for all CorrDiff models.

pip

pip install earth2studio[corrdiff]

uv

uv add earth2studio --extra corrdiff

Cyclone Trackers

Notes: Additional dependencies for all cyclone tracking models. Only Python 3.12 and below support.

uv

uv pip install earth2studio --extra cyclone

pip

pip install earth2studio[cyclone]

Derived

Notes: Additional dependencies for all derived diagnostic models. No additional dependencies are needed for the derived models at the moment but included for completeness.

pip

pip install earth2studio[derived]

uv

uv add earth2studio --extra derived

Precipitation AFNO

pip

pip install earth2studio[precip-afno]

uv

uv add earth2studio --extra precip-afno

Precipitation AFNO V2

Notes: Improved version of the Precipitation AFNO model with enhanced accuracy.

pip

pip install earth2studio[precip-afno-v2]

uv

uv add earth2studio --extra precip-afno-v2

Windgust AFNO

pip

pip install earth2studio[windgust-afno]

uv

uv add earth2studio --extra windgust-afno

Submodule Dependencies

A few features in various submodules require some specific dependencies that have been deemed too specific to warrant an addition to the core dependencies. These can be installed with a submodule wide install group:

Data

pip

pip install earth2studio[data]

uv

uv add earth2studio --extra data

Perturbation

pip

pip install earth2studio[perturbation]

uv

uv add earth2studio --extra perturbation

Statistics

pip

pip install earth2studio[statistics]

uv

uv add earth2studio --extra statistics

Install All Optional Dependencies

In Earth2Studio, it’s recommended that users pick and choose the optional dependencies that are needed for their use case. Installing everything at once and for all models is only expected to work in a few specific environments and may not include support for every model depending on conflicts. This is only supported using uv and when using github as the source, not pypi registry. To install a best effort all optional dependencies group, use the following:

uv

uv add earth2studio --extra all

Environments

For the best experience, we recommend creating a fresh environment whether that be using uv, a Docker container or even a Conda environment. Below are some recipes for creating a handful of environments for setting up Earth2Studio in an isolated environment. For developer environments, please refer to the Overview.

uv Project

Using uv is the recommend way to set up a Python environment for Earth2Studio. Assuming uv is installed, use the following commands:

mkdir earth2studio-project && cd earth2studio-project
uv init --python=3.12
uv add "earth2studio @ git+https://github.com/NVIDIA/earth2studio.git@0.6.0"

or if you are already inside an existing uv project:

uv venv --python=3.12
uv add "earth2studio @ git+https://github.com/NVIDIA/earth2studio.git@0.6.0"

PyTorch Docker Container

For a docker environment the Nvidia PyTorch container provides a good base.

docker run -i -t nvcr.io/nvidia/pytorch:25.03-py3

>>> apt-get update && apt-get install -y git make curl && rm -rf /var/lib/apt/lists/*
>>> unset PIP_CONSTRAINT
>>> curl -LsSf https://astral.sh/uv/install.sh | sh
>>> uv venv --python=3.12
>>> uv pip install "earth2studio[all] @ git+https://github.com/NVIDIA/earth2studio.git@0.6.0"

Custom Container

For a dedicated docker container the following can be used to get started. There is some complexity to undo the pip constraints from the PyTorch container, but otherwise the install process is the same.

FROM nvcr.io/nvidia/pytorch:25.03-py3
COPY --from=ghcr.io/astral-sh/uv:0.6.13 /uv /uvx /bin/

# Install system dependencies
RUN apt-get update && apt-get install -y \
    make \
    curl \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app

# Disable contraint files in the container
ENV PIP_CONSTRAINT=

RUN uv venv --python=3.12
RUN uv pip install "earth2studio[all] @ git+https://github.com/NVIDIA/earth2studio.git@0.6.0"

Conda Environment

It is no longer recommend to use any conda environment manager for Earth2Studio in favor of uv. This is because the virtual environments set up by uv makes the system-wide conda environments not needed. However this demonstrates that in principle Earth2Studio can be installed using standard package tooling.

conda create -n earth2studio python=3.12
conda activate earth2studio

pip install earth2studio
# Manually follow up with optional dependencies needed [all] will not work

Configuration

Earth2Studio uses a few environment variables to configure various parts of the package. The important ones are:

• EARTH2STUDIO_CACHE: The location of the cache used for Earth2Studio. This is a file path where things like models and cached data from data sources will be stored.

• EARTH2STUDIO_PACKAGE_TIMEOUT: The max number of seconds for a download operation of a model package file from a remote store such as NGC, Huggingface or S3.