How to Integrate WeatherNext into Your Own Applications: A Developer's Guide

WeatherNext integrates into applications through a unified Predictor API that abstracts JAX model complexity behind an xarray-based interface, allowing you to swap GraphCast and GenCast models without changing your data pipeline.

The google-deepmind/weathernext repository provides production-ready weather forecasting models with a deliberately model-agnostic design. Whether you're building climate dashboards, research pipelines, or operational forecasting systems, the library's clean abstraction layer lets you focus on your application logic rather than model internals.

Understanding the WeatherNext Architecture

The Predictor Base Class

All WeatherNext models implement the abstract Predictor class defined in [weathernext/utils/predictor_base.py](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/predictor_base.py#L27-L45). This base class specifies three key methods:

• __call__ – Required. Performs inference and returns forecasts as an xarray.Dataset.
• loss – Optional. Computes training losses.
• loss_and_predictions – Optional. Returns both predictions and losses in one call.

Concrete predictor subclasses handle the JAX-specific implementation while exposing a framework-agnostic interface.

Available WeatherNext Model Families

The repository ships two production-grade model implementations:

Model Family	Source File	Architecture
WeatherNext 1 Graph (GraphCast)	[weathernext/weathernext1_graph/graphcast.py](https://github.com/google-deepmind/weathernext/blob/main/weathernext/weathernext1_graph/graphcast.py)	Graph neural network on an icosahedral mesh
WeatherNext 1 Gen (GenCast)	[weathernext/weathernext1_gen/gencast.py](https://github.com/google-deepmind/weathernext/blob/main/weathernext/weathernext1_gen/gencast.py)	Autoregressive generative model for probabilistic forecasts

Both models share common utilities for data handling, checkpointing, and diagnostics—enabling seamless substitution in your integration code.

Step-by-Step Integration Workflow

1. Install the WeatherNext Package

The repository includes a standard setup.py for pip installation:

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

2. Download a Pretrained Checkpoint

WeatherNext requires pretrained weights to produce meaningful forecasts. The repository provides download scripts in the docs/ folder. The checkpoint loading utilities live in [utils/checkpoint.py](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/checkpoint.py).

3. Prepare Input Data as xarray

WeatherNext expects xarray.Dataset inputs matching the model's required variables, pressure levels, and lead-times. Helper functions in [utils/data_utils.py](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/data_utils.py) handle NetCDF conversion and JAX array preparation.

4. Instantiate Your Predictor

Choose between GraphCastPredictor or GenCastPredictor based on your forecasting needs—probabilistic vs. deterministic, generative vs. graph-based.

5. Run Inference

Call the predictor with inputs, a targets template (specifying output structure), and forcings. The method returns an xarray.Dataset ready for downstream processing.

Complete Integration Example

This end-to-end example demonstrates integrating WeatherNext GraphCast into a custom application. The same pattern applies to GenCast—only the import and class name change.

# Step 1: Import WeatherNext components

from weathernext.weathernext1_graph.graphcast import GraphCastPredictor

# from weathernext.weathernext1_gen.gencast import GenCastPredictor  # Alternative

from weathernext.utils import checkpoint, data_utils


# Step 2: Load pretrained parameters

ckpt_path = "/path/to/your/checkpoint.npz"
params = checkpoint.load_checkpoint(ckpt_path)


# Step 3: Initialize the predictor

predictor = GraphCastPredictor(params=params)


# Step 4: Load and prepare input data

ds_inputs = data_utils.load_dataset("sample_input.nc")
ds_forcings = data_utils.load_dataset("sample_forcings.nc")

# Targets template defines output structure—shape matters, values don't

targets_template = ds_inputs.isel(time=0).drop_vars("time")


# Step 5: Generate forecast

forecast = predictor(
    inputs=ds_inputs,
    targets_template=targets_template,
    forcings=ds_forcings,
)


# Step 6: Integrate into your application

forecast.to_netcdf("operational_forecast.nc")
processed_output = your_custom_processing(forecast)

Key Integration Files Reference

File	Purpose
[weathernext/utils/predictor_base.py](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/predictor_base.py)	Abstract Predictor class defining the integration contract
[weathernext/weathernext1_graph/graphcast.py](https://github.com/google-deepmind/weathernext/blob/main/weathernext/weathernext1_graph/graphcast.py)	Concrete GraphCastPredictor implementation
[weathernext/weathernext1_gen/gencast.py](https://github.com/google-deepmind/weathernext/blob/main/weathernext/weathernext1_gen/gencast.py)	Concrete GenCastPredictor implementation
[weathernext/utils/checkpoint.py](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/checkpoint.py)	load_checkpoint() and checkpoint management
[weathernext/utils/data_utils.py](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/data_utils.py)	NetCDF/xarray loading and JAX conversion helpers
[weathernext/utils/losses.py](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/losses.py)	Standard loss functions for training scenarios
[setup.py](https://github.com/google-deepmind/weathernext/blob/main/setup.py)	Package installation entry point

Design Patterns for WeatherNext Integration

Model-Agnostic Pipelines

The Predictor abstraction enables elegant model swapping. Your data preparation, inference orchestration, and post-processing remain identical regardless of which WeatherNext model you deploy:

def generate_forecast(predictor_class, params, inputs, forcings):
    """Model-agnostic forecast generation."""
    predictor = predictor_class(params=params)
    targets_template = inputs.isel(time=0).drop_vars("time")
    return predictor(inputs, targets_template, forcings)

# Switch models without code changes

if use_probabilistic:
    forecast = generate_forecast(GenCastPredictor, params, inputs, forcings)
else:
    forecast = generate_forecast(GraphCastPredictor, params, inputs, forcings)

xarray-Native Workflows

WeatherNext's choice of xarray as the primary data structure integrates naturally with the scientific Python ecosystem. Forecast outputs work directly with:

• Cartopy and matplotlib for geospatial visualization
• Dask for distributed computing on large forecast ensembles
• Zarr and NetCDF for storage and archival
• Pangeo tools for cloud-native climate analytics

Summary

• WeatherNext integration centers on the Predictor base class in utils/predictor_base.py, which abstracts all JAX complexity behind an xarray interface.

• Two production models are available: GraphCast (deterministic graph neural network) and GenCast (probabilistic generative model), with identical integration patterns.

• The five-step workflow covers installation, checkpoint download, data preparation, predictor instantiation, and inference—returning standard xarray datasets for downstream use.

• Model-agnostic design lets you swap GraphCast for GenCast by changing a single class import, keeping your data pipeline untouched.

• Core utilities in utils/checkpoint.py, utils/data_utils.py, and utils/losses.py provide reusable components for checkpoint management, I/O, and training.

Frequently Asked Questions

What data format does WeatherNext require?

WeatherNext consumes and produces xarray.Dataset objects. According to the source code in utils/data_utils.py, the library includes helpers to load NetCDF files and convert between xarray structures and JAX arrays. Your input datasets must contain the specific variables, pressure levels, and temporal dimensions expected by your chosen model checkpoint.

Can I switch between GraphCast and GenCast without rewriting my code?

Yes. Both GraphCastPredictor and GenCastPredictor inherit from the same Predictor abstract base class in utils/predictor_base.py. They share identical __call__ signatures, accepting inputs, targets_template, and forcings parameters. Only the class instantiation changes—the surrounding data preparation and result handling remain identical.

How do I load pretrained WeatherNext weights?

Use checkpoint.load_checkpoint() from utils/checkpoint.py. The function accepts a path to a .npz file (or compatible format) and returns model parameters suitable for passing to a predictor's params= argument. The repository provides checkpoint download scripts in the docs/ folder to obtain official pretrained weights.

Is WeatherNext suitable for real-time operational forecasting?

The Predictor API design in weathernext supports operational use cases through its clean separation of concerns. Inference calls return standard xarray datasets that integrate with scheduling systems, monitoring tools, and downstream pipelines. Performance characteristics depend on your hardware (JAX accelerates on GPU/TPU) and the autoregressive depth for GenCast probabilistic ensembles.

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