File size: 4,675 Bytes
57e06a2 09af8e7 57e06a2 09af8e7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | # GenCast — native PyTorch denoiser (official architecture, official Mini weights load) + EDM/DPM-Solver++ sampler
[← Model zoo](../README.md) · [All model pages](README.md) · **Paper:** *Nature* 2025, [doi:10.1038/s41586-024-08252-9](https://doi.org/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.Module`s 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` |
```python
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](https://github.com/GISWLH/WeatherAI/blob/main/docs/model_status.md).
|