# How to Run WeatherNext Locally: A Complete Setup and Inference Guide

> Run WeatherNext locally with this complete setup and inference guide. Discover the JAX/Flax framework for powerful weather forecasting via a simple command.

- Repository: [Google DeepMind/weathernext](https://github.com/google-deepmind/weathernext)
- Tags: how-to-guide
- Published: 2026-08-16

---

**WeatherNext is a JAX/Flax-based weather forecasting framework that runs locally with a mini checkpoint and a single command-line call.**

Running **WeatherNext locally** requires installing JAX and Flax dependencies, downloading a pre-trained mini checkpoint, and invoking the [`gencast.py`](https://github.com/google-deepmind/weathernext/blob/main/gencast.py) entry point. The repository provides CPU-compatible configurations for laptop use and optional GPU/TPU acceleration for larger workloads.

## Install the WeatherNext Package

Begin by setting up a Python environment with JAX, Flax, and the WeatherNext library.

### Step 1: Clone and Navigate

```bash
git clone https://github.com/google-deepmind/weathernext.git
cd weathernext

```

### Step 2: Install JAX and Flax

Choose the appropriate JAX build for your hardware:

```bash

# For CPU-only machines:

pip install "jax[cpu]" flax

# For CUDA-enabled GPUs (CUDA 11+):

pip install "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html flax

```

### Step 3: Install WeatherNext

The repository README specifies versioned installation via pip:

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

```

This command installs the library and registers the `weathernext` namespace for imports.

## Download a Pre-Trained Checkpoint

WeatherNext includes small test checkpoints designed for local development. The **WeatherNextCyclones_Mini.json** configuration fits in standard laptop memory.

Download the mini checkpoint configuration:

```bash
wget https://github.com/google-deepmind/weathernext/raw/main/weathernext/weathernext2/configs/WeatherNextCyclones_Mini.json -O mini_config.json

```

The config file resides at [`weathernext/weathernext2/configs/WeatherNextCyclones_Mini.json`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/weathernext2/configs/WeatherNextCyclones_Mini.json) in the repository tree and points to model weights compatible with limited compute.

## Run Inference with GenCast

The primary entry point for **running WeatherNext locally** is [`weathernext/weathernext1_gen/gencast.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/weathernext1_gen/gencast.py). This module provides both a command-line interface and a programmatic API.

### Command-Line Execution

Generate a forecast with a single command:

```bash
python -m weathernext.weathernext1_gen.gencast \
    --config_path=mini_config.json \
    --output_dir=./forecast_output

```

The script loads the model, initializes the state, and writes the forecast as a **NetCDF file** to `./forecast_output/`. Progress logs print to stdout during the autoregressive generation loop.

### Programmatic Usage

For custom pipelines, import the modules directly:

```python
from weathernext.weathernext1_gen import gencast
from weathernext.utils import checkpoint

# Load configuration and model weights

config = checkpoint.load_config('mini_config.json')
model = checkpoint.load_model(config)

# Initialize with random state (or replace with real ERA5 data)

initial_state = model.sample_initial_state(batch_size=1)

# Generate 24-hour forecast (6-hour steps × 4)

forecast = gencast.generate(model, initial_state, steps=4)

# Persist to NetCDF format

forecast.save('my_forecast.nc')

```

The `checkpoint.load_model()` function handles weight loading from paths specified in the JSON config, while `gencast.generate()` executes the forward pass through the graph-based transformer architecture defined in [`weathernext/weathernext2/architecture.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/weathernext2/architecture.py).

## Explore Interactive Notebooks

The `docs/` directory contains runnable Jupyter notebooks that demonstrate end-to-end workflows:

- **`docs/weathernext2/wn2_demo.ipynb`** — WeatherNext 2 architecture walkthrough
- **`docs/weathernext1_gen/gencast_mini_demo.ipynb`** — Interactive mini checkpoint demo

Launch a notebook with:

```bash
jupyter notebook docs/weathernext1_gen/gencast_mini_demo.ipynb

```

These notebooks automatically resolve checkpoint paths and include visualization utilities for forecast analysis.

## Optional: Enable GPU or TPU Acceleration

WeatherNext performance scales with available accelerators. Set JAX platform flags before running inference.

### GPU Configuration

```bash
export JAX_PLATFORM_NAME=gpu
python -m weathernext.weathernext1_gen.gencast --config_path=mini_config.json ...

```

### TPU Configuration

```bash
export XLA_FLAGS=--xla_gpu_cuda_data_dir=/usr/lib/cuda
export JAX_PLATFORM_NAME=tpu
python -m weathernext.weathernext1_gen.gencast --config_path=mini_config.json ...

```

JAX automatically dispatches linear algebra kernels to the selected backend. The mini checkpoint runs efficiently on CPU; full-resolution models require accelerator memory.

## Summary

- **Installation**: Use `pip install git+https://github.com/google-deepmind/weathernext.git@v0.3.0` with JAX[cpu] or JAX[cuda]
- **Checkpoint**: Download [`WeatherNextCyclones_Mini.json`](https://github.com/google-deepmind/weathernext/blob/main/WeatherNextCyclones_Mini.json) for local testing
- **Inference**: Execute `python -m weathernext.weathernext1_gen.gencast` with `--config_path` and `--output_dir`
- **API**: Import from `weathernext.weathernext1_gen.gencast` and `weathernext.utils.checkpoint` for custom scripts
- **Notebooks**: Run `docs/weathernext1_gen/gencast_mini_demo.ipynb` for interactive exploration

## Frequently Asked Questions

### What hardware do I need to run WeatherNext locally?

A standard laptop with 8GB+ RAM runs the **mini checkpoint** comfortably. The CPU-only JAX build handles inference in minutes for short-range forecasts. For operational-scale predictions or training, a CUDA GPU or Cloud TPU reduces runtime significantly.

### Where are the model weights stored?

Checkpoint weights download automatically based on URLs in the JSON config files. The `weathernext.utils.checkpoint` module manages caching and loading. The mini checkpoint (~100MB) suits development; full checkpoints require 10GB+ storage.

### Can I modify the forecast lead time?

Yes. The `gencast.generate()` function accepts a `steps` parameter controlling autoregressive rollout. Each step typically represents 6 hours, so `steps=4` produces a 24-hour forecast. Modify this in Python or extend the CLI script in [`weathernext/weathernext1_gen/gencast.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/weathernext1_gen/gencast.py).

### How do I load real ERA5 data instead of random initialization?

Replace `model.sample_initial_state()` with data loading from `weathernext.data_loaders.era5`. The repository includes ERA5 preprocessing utilities in `weathernext/data_loaders/` that convert GRIB or NetCDF inputs to the internal JAX array format expected by the model.