models/gencast.md

GenCast — native PyTorch denoiser (official architecture, official Mini weights load) + EDM/DPM-Solver++ sampler

← Model zoo · All model pages · Paper: Nature 2025, doi:10.1038/s41586-024-08252-9

GenCast (Nature 2025) is a conditional EDM diffusion model whose denoiser is a GraphCast-style grid→mesh→grid network with a 16-layer k-hop mesh transformer; official code is JAX (google-deepmind/weathernext). weatherai.models.gencast contains explicit nn.Modules of the official denoiser (parameter names mirror the Haiku names) plus the EDM wrapper / DPM-Solver++ 2S sampler. No wrapper around JAX is involved at runtime.

file	contents
denoiser.py	FourierFeaturesMLP (noise-level encoder), LinearNormConditioning, ConditionedMLP (Linear→swish→Linear→LN→cond-affine), Grid2MeshGNN, MeshTransformer (+ MeshAttention: dense masked or banded tri-block-diagonal k-hop attention), Mesh2GridGNN, GenCastDenoiser, convert_official_params, GenCastDenoiser_from_official, channel stacking helpers
graphs.py	grid2mesh radius graph, finest-mesh banded (reverse Cuthill-McKee) permutation, mesh2grid containing-triangle graph, k-hop mask (re-uses graphcast.mesh)
gencast.py	EDM preconditioning/loss, noise schedule, DPM-Solver++ 2S + churn sampler, DenoiserNet (config-driven adapter), GenCast, GenCast_lite

import numpy as np, torch
from weatherai.models.gencast import (GenCast_lite, GenCastDenoiser_from_official, denoiser_inputs, unstack_variables)

# (a) small config-driven, trainable instance of the SAME native architecture (random init)
m = GenCast_lite(img_size=(16, 32), out_channels=4, cond_channels=8, latent=32, transformer_layers=2, mesh_level=1)
cond, target = torch.randn(2, 8, 16, 32), torch.randn(2, 4, 16, 32)
loss = m.loss(target, cond); loss.backward()                     # EDM-weighted denoising loss (training step)
members = m.eval().sample(cond, num_members=4)                   # (4, 2, 4, 16, 32) DPM-Solver++ 2S + churn

# (b) official GenCast-1p0deg-Mini weights (public GCS bucket, 230 MB) in the native denoiser
#     gs://dm_graphcast/gencast/params/GenCast 1p0deg Mini <2019.npz
lat, lon = np.linspace(-90, 90, 181), np.arange(360)             # official 1° grid
net = GenCastDenoiser_from_official("mini.npz", lat, lon)        # strict load; mesh_size=4, 16 layers, 57.5 M params
x = denoiser_inputs(inputs, forcings, noisy_targets, (181, 360)) # dict-of-arrays -> (B, 181*360, 264), official channel order
raw = net(x, noise_levels)                                       # (B, Ng, 84) raw network output

Status — verified (CPU, fp32):

• Denoiser vs official JAX with the official Mini weights (reference from the official weathernext1_gen.denoiser.Denoiser, scripts/gencast_denoiser_reference.py, random inputs, official plain-mha attention because the TPU splash kernel is unavailable):
  • trained configuration (mesh_size=4 / 2562 mesh nodes, attention_k_hop=16, 181×360 grid, 102 k grid→mesh edges, batch 1): max |Δ| = 4.8e-5 on all 84 outputs (output magnitude up to 22; max rel. to output max 3.2e-6);
  • small configuration (mesh_size=2, k_hop=3, 13×24, batch 2, two noise levels): max |Δ| = 7.4e-5; banded tri-block-diagonal attention gives the same (6.9e-5).
  • graph construction (permuted mesh, grid2mesh/mesh2grid indices and edge features) equals the official graphs (indices exactly, features ≤3e-8).
  • all all official parameter tensors (57,492,612 weights) strict-load; unit tests use atol 2e-4 / rtol 1e-4 (small) and atol 5e-4 (full).
• EDM coefficients, noise/churn schedules and the DPM-Solver++ 2S sampler (± churn) match the official JAX sampler (rtol 2e-4, closed-form toy denoiser, fixed noise).
• Training step example/test: loss decreases on a fixed batch, all parameters receive gradients (tests/models/gencast).

Not verified / not ported: the network only was compared against the official network — not the full model (no end-to-end official ensemble forecast comparison; the official InputsAndResiduals normalisation, NaN-cleaning of SST and the autoregressive rollout are not ported); the sampling noise is iid Gaussian on the grid, official is spherical-harmonic white noise (needs dinosaur); the 0.25° and Operational checkpoints were not run (same code path, larger graphs); the lite config is not a trained model. Details: docs/model_status.md.