How to Install WeatherNext: Complete Setup Guide for DeepMind's Weather Prediction Model

WeatherNext can be installed via pip directly from GitHub or cloned locally, with pretrained weights downloaded separately from Google Cloud Storage.

This guide covers both installation methods for the google-deepmind/weathernext repository, explains the package structure, and shows how to verify your setup with a quick inference test.

Two Ways to Install WeatherNext

WeatherNext is distributed as a pure-Python package with core functionality in the weathernext directory and the flagship WeatherNext 2 model implemented through the FGN (Functional Generative Network) in weathernext/weathernext2/fgn.py [1†L15-L30].

Method 1: Pip Install from GitHub (Recommended)

The fastest way to install WeatherNext is using pip with a version tag. The repository README specifies this one-line command [3†L36-L38]:

pip install git+https://github.com/google-deepmind/weathernext.git@v0.3.0

This approach automatically:

• Pulls the weathernext package and all submodules
• Installs dependencies defined in setup.py [4†L32-L53]
• Includes JSON configuration files as package data [5†L30-L31]

Method 2: Clone and Install Locally

For development, modifications, or bleeding-edge features, clone the repository and install with setuptools [4†L14-L22]:

git clone https://github.com/google-deepmind/weathernext.git
cd weathernext
pip install .

Why choose this method?

• Edit source files directly (e.g., weathernext/utils/losses.py, weathernext/utils/predictor_base.py)
• Run tests against the latest main branch
• Build custom model configurations in weathernext/weathernext2/configs/

Dependencies Installed Automatically

Both installation methods resolve dependencies via setup.py [4†L32-L53]:

Package	Purpose
jax	Numerical computing and automatic differentiation
dm-haiku	Neural network library for model construction
xarray	Labeled multi-dimensional arrays for weather data
dinosaur-dycore	Dynamical core for atmospheric simulations

Additional utilities in weathernext/utils/ (e.g., xarray_jax helpers) provide data-handling primitives [2†L21-L30].

Download Pretrained Model Weights

The Python package installation does not include weights. Download these separately from the public Google Cloud bucket referenced in the README [3†L40-L45]:

# Example: WeatherNext Cyclones Mini model

gsutil cp gs://dm_graphcast/WeatherNextCyclones_Mini_<2024>.npz /local/path/

Replace the bucket path with the specific model variant you need. The README contains direct links for WeatherNext 2, Cyclones, and other specialized checkpoints.

Verify Installation with a Test Inference

Confirm your WeatherNext install works by loading a configuration and running the FGN predictor:

import weathernext.weathernext2 as wn2
from weathernext.utils import predictor_base

# Load installed configuration JSON

config_path = wn2.package_data_path('configs/WeatherNext2.json')
config = wn2.load_config(config_path)

# Build the FGN predictor from `weathernext/weathernext2/fgn.py`

predictor = wn2.construct_predictor(config)

# Create dummy inputs (replace with ERA5/HRES data in practice)

inputs = predictor_base.fake_inputs()
targets = predictor_base.fake_targets()

# Run inference

predictions = predictor(inputs, targets_template=targets, is_training=False)
print(predictions)

This exercises the full stack: configuration loading, model construction via weathernext/weathernext2/architecture.py, and the base predictor API from weathernext/utils/predictor_base.py.

Key Source Files to Know

Path	Role
setup.py	Package metadata, dependency specifications, and data file inclusion [4†L14-L53]
weathernext/weathernext2/fgn.py	Core FGN predictor implementation [1†L15-L30]
weathernext/weathernext2/architecture.py	Model assembly utilities
weathernext/utils/predictor_base.py	Base class for all predictors with standardized API [2†L21-L30]
weathernext/weathernext2/configs/WeatherNext2.json	Full-resolution model configuration [5†L30-L31]

Summary

• Install WeatherNext via pip for stable releases: pip install git+https://github.com/google-deepmind/weathernext.git@v0.3.0
• Clone and install locally for development work and source modifications
• Dependencies (JAX, Haiku, xarray) resolve automatically through setup.py
• Download weights separately from Google Cloud Storage—installation excludes these large files
• Verify with a quick inference test using wn2.construct_predictor() and dummy inputs

Frequently Asked Questions

What Python version does WeatherNext require?

WeatherNext requires Python 3.9 or later, as specified in setup.py classifiers. The JAX dependency further constrains this to Python versions compatible with your local CUDA installation if using GPU acceleration.

Can I install WeatherNext without GPU support?

Yes. The base installation includes CPU-only JAX. For GPU support, install the appropriate jax[cuda] variant separately after installing WeatherNext, following the official JAX CUDA installation guide.

Where are the configuration files stored after installation?

JSON configs live in weathernext/weathernext2/configs/ and install as package data. Access them programmatically via wn2.package_data_path('configs/WeatherNext2.json') rather than hardcoding paths [5†L30-L31].

How do I update WeatherNext to a newer version?

Run the pip install command with the desired tag: pip install --upgrade git+https://github.com/google-deepmind/weathernext.git@v0.X.Y. For local clones, pull latest changes and rerun pip install . to refresh the installation.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}