Blame for models/aurora.md

# Aurora — native PyTorch implementation

[← Model zoo](../README.md) · [All model pages](README.md) · **Paper:** *Nature* 2025, [doi:10.1038/s41586-025-09005-y](https://doi.org/10.1038/s41586-025-09005-y)

`weatherai.models.aurora` re-implements Aurora (*Nature* 2025; MIT,
[microsoft/aurora](https://github.com/microsoft/aurora)) as explicit `nn.Module`s with readable `forward()`s,
in the style of the Pangu/FuXi code (no dependency on the `aurora` pip package):

| file | contents |
|---|---|
| `patch_embed.py` | `LevelPatchEmbed` – per-variable 3-D patch conv over (history × P × P) |
| `perceiver.py` | `PerceiverAttention`, `PerceiverResampler` (post-norm cross-attention + MLP) |
| `encoder.py` | `Perceiver3DEncoder` – patch-embed surface(+static) and every level → level aggregation → + position / patch-scale / lead-time / absolute-time Fourier embeddings |
| `swin3d.py` | `Swin3DBackbone`: `SwinStage`s of `Swin3DBlock` (shifted 3-D windows with periodic-longitude mask, AdaLN on lead time), `PatchMerging3D`, `PatchSplitting3D` (linear1 → reshape/permute → crop padding → norm → linear2) |
| `decoder.py` | `Perceiver3DDecoder` – level de-aggregation + per-variable linear unpatchify |
| `aurora.py` | `Aurora` (plain-tensor I/O, official normalisation statistics, `Aurora_lite`, `Aurora_small`) |
| `official.py` | the old thin wrapper over `microsoft-aurora` – **reference oracle only** (tests / parity) |

```python

import torch

from weatherai.models import Aurora, Aurora_lite, Aurora_small



m = Aurora_lite().eval()                      # ~3 M params, random init, CPU friendly

surf   = torch.randn(1, 2, 4, 16, 32) + torch.tensor([278., 0, 0, 1e5]).view(1, 1, 4, 1, 1)  # (B,T=2,[2t,10u,10v,msl],H,W)

static = torch.randn(3, 16, 32)               # [lsm, z, slt]

atmos  = torch.randn(1, 2, 5, 4, 16, 32)      # (B,T,[z,u,v,t,q],L=4,H,W)

with torch.no_grad():

    surf_next, atmos_next = m(surf, static, atmos)   # physical units, +6 h



# fully configurable / trainable: any depths, heads, window, patch size, embed dim

my = Aurora(embed_dim=128, encoder_depths=(2, 4, 2), decoder_depths=(2, 4, 2),

            encoder_num_heads=(4, 8, 16), decoder_num_heads=(16, 8, 4), num_heads=4, window_size=(2, 6, 12))

loss = (my(surf, static, atmos)[0] - surf[:, 1]).pow(2).mean(); loss.backward()   # see tests for a full training step



# Official small pretrained checkpoint (~450 MB from HF microsoft/aurora), strict=True load into the native model

small = Aurora_small(pretrained=True).eval()  # or Aurora_small(checkpoint_path="aurora-0.25-small-pretrained.ckpt")

```

**Verified (numerically, vs. the official `microsoft-aurora` package at b628d7c, same weights, same inputs; fp32,

`atol=1e-5, rtol=1e-4` – observed max abs difference was exactly 0.0 on CPU):** random-weight state-dict transfer
`strict=True` for 3 configs (window (2,2,2)/(2,3,5)/(2,6,12), uneven depths, v3 lead-time embedding), grids
16×32 … 32×64 incl. the odd-latitude crop (H%4==1) and patch grids that need merge-padding / split-cropping
(25×44); the official `aurora-0.25-small-pretrained.ckpt` strict-loads into the native `Aurora_small` and
reproduces the official output on 33×64 inputs. HF ZeroGPU results: see [docs/model_status.md](https://github.com/GISWLH/WeatherAI/blob/main/docs/model_status.md).
Unit tests: `tests/models/aurora` (shapes, odd grids, merge/split + window helpers, param-name layout,
25-step training demo, grads, parity).
**Not verified / not included:** skill on real ERA5/HRES data (random inputs only), multi-step rollout;
upstream features that are *not* ported: LoRA, stochastic/ensemble mode, level-conditioned embeddings,
dynamic/atmos-static variables, separate/modulation heads (so the fine-tuned 1.3 B, HRES-0.1°, air-pollution,
wave and v1.5 checkpoints are **not** loadable; `aurora-0.25-pretrained` and `-small-pretrained` are); the 1.3 B
checkpoint was not run; built-in normalisation stats only for the 13 standard ERA5 levels.