Upload folder using huggingface_hub

README.md

@@ -0,0 +1,151 @@
+---
+datasets:
+- SEVIR
+frameworks:
+- PyTorch
+language:
+- en
+- zh
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Weather Forecast
+- Spatiotemporal Forecast
+- SEVIR
+- Earthformer
+tasks: []
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">Earthformer</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+Earthformer was proposed by researchers from Amazon Web Services (AWS) in collaboration with the Hong Kong University of Science and Technology. It aims to address the prohibitive computational cost of conventional Transformers when processing high-dimensional geophysical data.
+
+Earthformer: Exploring Space-Time Transformers for Earth System Forecasting
+
+https://arxiv.org/abs/2207.05833
+
+# Model Description
+
+Earthformer is a space-time Transformer model designed for Earth system forecasting, such as weather and climate. Its core component is a novel attention mechanism called Cuboid Attention.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Weather Forecast Training | Train the Earthformer weather forecasting model using SEVIR structured data. |
+| Local Quick Validation | Use synthetic data to verify data loading, model training, inference, and inference result visualization. |
+| ModelScope / OneCode Execution | Download as a standalone model package, install dependencies, and run scripts directly. |
+| Multi-GPU Training | Launch multi-process training via `torchrun`. |
+
+# Usage Guide
+
+## 1. OneCode Usage
+
+Experience intelligent one-click AI4S programming through the OneCode online environment:
+
+[Click to Experience Intelligent One-Click AI4S Programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+The commands below assume they are run from the Earthformer project root directory.
+
+**Hardware Requirements**
+
+- Training and inference require a GPU or DCU recognized by PyTorch; CPU can be used to generate synthetic data and verify configuration, but cannot run the current training and inference scripts.
+- Multi-GPU training uses the NCCL backend. Please make sure the device driver, communication libraries, and PyTorch version are compatible.
+- DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching your cluster, is recommended.
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/Earthformer --local-dir ./Earthformer
+cd Earthformer
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Please activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+# Please activate CONDA first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+### Generate Synthetic Data
+
+Synthetic data is only used to verify the data protocol and program flow; it does not represent real SEVIR data or the model's forecast quality:
+
+```bash
+python script/fake_data.py
+```
+
+By default this generates `data/synthetic_sevir/{train,val,test}.npz` and `metadata.json`. To generate synthetic data matching the official SEVIR spatial size:
+
+```bash
+python script/fake_data.py --output-dir data/synthetic_sevir_384 --height 384 --width 384
+```
+
+### Training
+
+Single GPU:
+
+```bash
+python script/train.py
+```
+
+Multi-GPU:
+
+```bash
+torchrun --nproc_per_node=8 script/train.py
+```
+
+Training starts from random initialization and saves to `data/checkpoint/earthformer.pt` by default.
+
+### Training Weights
+
+This repository provides weights trained on SEVIR data in the `weights/` folder. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+Inference reads the training checkpoint by default and writes the output to `output/predictions.npz`.
+
+```bash
+python script/inference.py
+```
+
+### Evaluation and Visualization
+
+```bash
+python script/result.py
+```
+
+# OneScience Official Information
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation & License
+
+- This repository is the OneScience reproduction of the original Earthformer paper. The official Earthformer implementation is released under the Apache License 2.0; the use of this repository's code and the SEVIR data remains subject to the licenses and terms of use of the respective projects.

conf/config.yaml

@@ -0,0 +1,54 @@
+data:
+  protocol: synthetic_sevir
+  # Reference dataset: official SEVIR VIL (earthformer_sevir_v1.yaml) uses
+  # 384x384 frames at 5-minute intervals, seq_len = 25 (13 input + 12 output),
+  # "sequent" sampling with stride 12, single channel. The small spatial shape
+  # below is intentionally the default for CPU smoke verification.
+  input_length: 13
+  output_length: 12
+  height: 32
+  width: 32
+  channels: 1
+  frame_interval_minutes: 5
+  normalization: unit
+  data_dir: data/synthetic_sevir
+  train_samples: 8
+  val_samples: 2
+  test_samples: 2
+  train_npz: data/synthetic_sevir/train.npz
+  val_npz: data/synthetic_sevir/val.npz
+  test_npz: data/synthetic_sevir/test.npz
+  fallback_if_missing: true
+
+model:
+  dims: [4, 8]
+  depths: [1, 1]
+  heads: 1
+  pattern: [[2, 4, 4]]
+  num_global_vectors: 1
+  ff_ratio: 2.0
+  dropout: 0.0
+
+train:
+  seed: 42
+  # Training always starts from a randomly initialized model (from scratch).
+  # The --resume flag only restores a locally trained checkpoint for continued
+  # training; no official pretrained Earthformer weights are ever downloaded.
+  batch_size: 1
+  epochs: 5
+  validation_steps: 1
+  # Per-epoch validation computes only cheap MSE/MAE (the training loss is
+  # already MSE). Enable this to also compute SSIM/CSI every epoch; the
+  # authoritative full evaluation lives in script/result.py.
+  compute_full_metrics: false
+  learning_rate: 0.001
+  weight_decay: 0.00001
+  # Trained weights for inference are saved here as earthformer.pt.
+  output_dir: data/checkpoint
+
+dataloader:
+  num_workers: 0
+  pin_memory: false
+
+distributed:
+  backend: auto

config.json

@@ -0,0 +1,77 @@
+{
+  "model_name": "Earthformer",
+  "model_type": "earthformer",
+  "architectures": [
+    "Earthformer"
+  ],
+  "framework": "PyTorch",
+  "domain": "climate-and-atmosphere",
+  "task": "spatiotemporal-forecasting",
+  "implementation": {
+    "entry_point": "model/earthformer.py",
+    "scope": "two-level hierarchical Cuboid Attention space-time Transformer with BTHWC input and output"
+  },
+  "architecture": {
+    "family": "hierarchical space-time Transformer with Cuboid Attention",
+    "attention_mechanism": "cuboid-window attention with optional global vectors",
+    "input_format": "BTHWC",
+    "encoder": {
+      "levels": 2,
+      "stem": "Conv2d 3x3 stride 2",
+      "downsample": "Conv2d 3x3 stride 2",
+      "cuboid_pattern": [
+        [
+          2,
+          4,
+          4
+        ]
+      ]
+    },
+    "decoder": {
+      "cuboid_pattern": "axial",
+      "cross_attention": "CuboidCross(T,1,1): future queries attend to history at each spatial site",
+      "prediction": "learned future query with per-frame up-projection and encoder skip connection"
+    },
+    "activation": "GELU",
+    "normalization": "LayerNorm",
+    "repository_default_config": {
+      "purpose": "CPU smoke verification with synthetic SEVIR data",
+      "dims": [
+        4,
+        8
+      ],
+      "depths": [
+        1,
+        1
+      ],
+      "heads": 1,
+      "num_global_vectors": 1,
+      "ff_ratio": 2.0,
+      "dropout": 0.0
+    }
+  },
+  "data": {
+    "dataset": "SEVIR",
+    "variable": "vertically_integrated_liquid",
+    "frame_interval_minutes": 5,
+    "input_length": 13,
+    "output_length": 12,
+    "channels": 1,
+    "official_spatial_size": [
+      384,
+      384
+    ],
+    "official_sampling": "sequent with stride 12",
+    "default_smoke_spatial_size": [
+      32,
+      32
+    ],
+    "normalization": "unit"
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/earthformer.py",
+    "script/data_loader.py",
+    "configuration.json"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "spatiotemporal_forecasting",
+  "model": "Earthformer",
+  "input_format": "BTHWC",
+  "protocol": "synthetic_sevir",
+  "default_config": "conf/config.yaml",
+  "train": "script/train.py",
+  "inference": "script/inference.py",
+  "evaluation": "script/result.py",
+  "visualization": "script/result.py"
+}

model/earthformer.py

@@ -0,0 +1,213 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+@dataclass
+class CuboidMeta:
+    batch: int
+    shape: tuple[int, int, int]
+    padded: tuple[int, int, int]
+    cuboid: tuple[int, int, int]
+
+
+def cuboid_partition(x: torch.Tensor, cuboid: tuple[int, int, int]) -> tuple[torch.Tensor, CuboidMeta]:
+    b, t, h, w, c = x.shape
+    bt, bh, bw = (min(size, dim) for size, dim in zip(cuboid, (t, h, w)))
+    pt, ph, pw = (-t) % bt, (-h) % bh, (-w) % bw
+    padded = F.pad(x.permute(0, 4, 1, 2, 3), (0, pw, 0, ph, 0, pt)).permute(0, 2, 3, 4, 1)
+    tp, hp, wp = padded.shape[1:4]
+    windows = padded.reshape(b, tp // bt, bt, hp // bh, bh, wp // bw, bw, c)
+    windows = windows.permute(0, 1, 3, 5, 2, 4, 6, 7).reshape(-1, bt * bh * bw, c)
+    return windows, CuboidMeta(b, (t, h, w), (tp, hp, wp), (bt, bh, bw))
+
+
+def cuboid_merge(windows: torch.Tensor, meta: CuboidMeta) -> torch.Tensor:
+    b, (t, h, w), (tp, hp, wp), (bt, bh, bw) = meta.batch, meta.shape, meta.padded, meta.cuboid
+    c = windows.shape[-1]
+    x = windows.reshape(b, tp // bt, hp // bh, wp // bw, bt, bh, bw, c)
+    x = x.permute(0, 1, 4, 2, 5, 3, 6, 7).reshape(b, tp, hp, wp, c)
+    return x[:, :t, :h, :w]
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim: int, ratio: float, dropout: float):
+        super().__init__()
+        hidden = int(dim * ratio)
+        self.net = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class CuboidAttentionLayer(nn.Module):
+    def __init__(self, dim: int, heads: int, cuboid: tuple[int, int, int], ff_ratio: float, dropout: float, use_global: bool):
+        super().__init__()
+        self.cuboid = cuboid
+        self.use_global = use_global
+        self.local_norm = nn.LayerNorm(dim)
+        self.global_norm = nn.LayerNorm(dim) if use_global else None
+        self.local_attention = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
+        self.global_attention = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True) if use_global else None
+        self.local_ff_norm = nn.LayerNorm(dim)
+        self.local_ff = FeedForward(dim, ff_ratio, dropout)
+        self.global_ff_norm = nn.LayerNorm(dim) if use_global else None
+        self.global_ff = FeedForward(dim, ff_ratio, dropout) if use_global else None
+
+    def forward(self, x: torch.Tensor, global_vectors: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor | None]:
+        normalized = self.local_norm(x)
+        windows, meta = cuboid_partition(normalized, self.cuboid)
+        if self.use_global:
+            if global_vectors is None:
+                raise ValueError("global vectors are required when use_global=True")
+            windows_per_batch = windows.shape[0] // x.shape[0]
+            repeated_global = self.global_norm(global_vectors).repeat_interleave(windows_per_batch, dim=0)
+            key_value = torch.cat((windows, repeated_global), dim=1)
+        else:
+            key_value = windows
+        attended = self.local_attention(windows, key_value, key_value, need_weights=False)[0]
+        x = x + cuboid_merge(attended, meta)
+        x = x + self.local_ff(self.local_ff_norm(x))
+        if self.use_global:
+            global_query = self.global_norm(global_vectors)
+            all_tokens = self.local_norm(x).reshape(x.shape[0], -1, x.shape[-1])
+            global_kv = torch.cat((global_query, all_tokens), dim=1)
+            global_vectors = global_vectors + self.global_attention(global_query, global_kv, global_kv, need_weights=False)[0]
+            global_vectors = global_vectors + self.global_ff(self.global_ff_norm(global_vectors))
+        return x, global_vectors
+
+
+def resolve_pattern(pattern: str | list[list[int]], shape: tuple[int, int, int]) -> list[tuple[int, int, int]]:
+    if pattern == "axial":
+        t, h, w = shape
+        return [(t, 1, 1), (1, h, 1), (1, 1, w)]
+    return [tuple(int(value) for value in item) for item in pattern]
+
+
+class CuboidBlock(nn.Module):
+    def __init__(self, dim: int, heads: int, pattern: str | list[list[int]], shape: tuple[int, int, int], ff_ratio: float, dropout: float, use_global: bool):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            CuboidAttentionLayer(dim, heads, cuboid, ff_ratio, dropout, use_global)
+            for cuboid in resolve_pattern(pattern, shape)
+        )
+
+    def forward(self, x: torch.Tensor, global_vectors: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor | None]:
+        for layer in self.layers:
+            x, global_vectors = layer(x, global_vectors)
+        return x, global_vectors
+
+
+class CuboidCrossAttention(nn.Module):
+    """CuboidCross(T,1,1): future queries attend to history at each spatial site."""
+
+    def __init__(self, dim: int, heads: int, ff_ratio: float, dropout: float):
+        super().__init__()
+        self.query_norm = nn.LayerNorm(dim)
+        self.memory_norm = nn.LayerNorm(dim)
+        self.attention = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
+        self.ff_norm = nn.LayerNorm(dim)
+        self.ff = FeedForward(dim, ff_ratio, dropout)
+
+    def forward(self, query: torch.Tensor, memory: torch.Tensor) -> torch.Tensor:
+        b, k, h, w, c = query.shape
+        if memory.shape[0] != b or memory.shape[2:4] != (h, w):
+            raise ValueError("cross-attention memory must match batch and spatial dimensions")
+        q = self.query_norm(query).permute(0, 2, 3, 1, 4).reshape(b * h * w, k, c)
+        m = self.memory_norm(memory).permute(0, 2, 3, 1, 4).reshape(b * h * w, memory.shape[1], c)
+        attended = self.attention(q, m, m, need_weights=False)[0]
+        attended = attended.reshape(b, h, w, k, c).permute(0, 3, 1, 2, 4)
+        query = query + attended
+        return query + self.ff(self.ff_norm(query))
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, dim: int, heads: int, pattern: str | list[list[int]], shape: tuple[int, int, int], ff_ratio: float, dropout: float, use_global: bool):
+        super().__init__()
+        self.self_block = CuboidBlock(dim, heads, pattern, shape, ff_ratio, dropout, use_global)
+        self.cross = CuboidCrossAttention(dim, heads, ff_ratio, dropout)
+
+    def forward(self, x: torch.Tensor, memory: torch.Tensor, global_vectors: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor | None]:
+        x, global_vectors = self.self_block(x, global_vectors)
+        return self.cross(x, memory), global_vectors
+
+
+class Earthformer(nn.Module):
+    """Two-level Cuboid Attention Earthformer with BTHWC input and output."""
+
+    def __init__(self, config: dict):
+        super().__init__()
+        data, model = config["data"], config["model"]
+        self.input_length = int(data["input_length"])
+        self.output_length = int(data["output_length"])
+        self.height, self.width = int(data["height"]), int(data["width"])
+        channels = int(data["channels"])
+        d0, d1 = (int(value) for value in model["dims"])
+        depths = model["depths"]
+        heads = int(model["heads"])
+        pattern = model.get("pattern", "axial")
+        ff_ratio, dropout = float(model.get("ff_ratio", 2.0)), float(model.get("dropout", 0.0))
+        self.num_global = int(model.get("num_global_vectors", 0))
+        use_global = self.num_global > 0
+        h0, w0, h1, w1 = self.height // 2, self.width // 2, self.height // 4, self.width // 4
+        self.stem = nn.Conv2d(channels, d0, 3, stride=2, padding=1)
+        self.downsample = nn.Conv2d(d0, d1, 3, stride=2, padding=1)
+        self.encoder_pos0 = nn.Parameter(torch.zeros(1, self.input_length, h0, w0, d0))
+        self.encoder_pos1 = nn.Parameter(torch.zeros(1, self.input_length, h1, w1, d1))
+        self.future_query = nn.Parameter(torch.empty(1, self.output_length, h1, w1, d1))
+        nn.init.trunc_normal_(self.future_query, std=0.02)
+        self.encoder0 = nn.ModuleList(CuboidBlock(d0, heads, pattern, (self.input_length, h0, w0), ff_ratio, dropout, use_global) for _ in range(depths[0]))
+        self.encoder1 = nn.ModuleList(CuboidBlock(d1, heads, pattern, (self.input_length, h1, w1), ff_ratio, dropout, use_global) for _ in range(depths[1]))
+        self.decoder1 = nn.ModuleList(DecoderBlock(d1, heads, "axial", (self.output_length, h1, w1), ff_ratio, dropout, use_global) for _ in range(depths[1]))
+        self.decoder0 = nn.ModuleList(DecoderBlock(d0, heads, "axial", (self.output_length, h0, w0), ff_ratio, dropout, use_global) for _ in range(depths[0]))
+        self.up_project = nn.Conv2d(d1, d0, 3, padding=1)
+        self.skip_project = nn.Linear(d0, d0)
+        self.head = nn.Conv2d(d0, channels, 3, padding=1)
+        if use_global:
+            self.encoder_global0 = nn.Parameter(torch.zeros(1, self.num_global, d0))
+            self.encoder_global1 = nn.Parameter(torch.zeros(1, self.num_global, d1))
+            self.decoder_global1 = nn.Parameter(torch.zeros(1, self.num_global, d1))
+            self.decoder_global0 = nn.Parameter(torch.zeros(1, self.num_global, d0))
+
+    @staticmethod
+    def _frames(module: nn.Module, x: torch.Tensor) -> torch.Tensor:
+        b, t, h, w, c = x.shape
+        result = module(x.permute(0, 1, 4, 2, 3).reshape(b * t, c, h, w))
+        return result.reshape(b, t, result.shape[1], result.shape[2], result.shape[3]).permute(0, 1, 3, 4, 2)
+
+    def _global(self, name: str, batch: int) -> torch.Tensor | None:
+        value = getattr(self, name, None)
+        return value.expand(batch, -1, -1) if value is not None else None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        expected = (self.input_length, self.height, self.width)
+        if x.ndim != 5 or x.shape[1:4] != expected:
+            raise ValueError(f"expected input [B,{expected[0]},{expected[1]},{expected[2]},C], got {tuple(x.shape)}")
+        batch = x.shape[0]
+        e0 = self._frames(self.stem, x) + self.encoder_pos0
+        g0 = self._global("encoder_global0", batch)
+        for block in self.encoder0:
+            e0, g0 = block(e0, g0)
+        e1 = self._frames(self.downsample, e0) + self.encoder_pos1
+        g1 = self._global("encoder_global1", batch)
+        for block in self.encoder1:
+            e1, g1 = block(e1, g1)
+        d1 = self.future_query.expand(batch, -1, -1, -1, -1)
+        gd1 = self._global("decoder_global1", batch)
+        for block in self.decoder1:
+            d1, gd1 = block(d1, e1, gd1)
+        b, k, h, w, c = d1.shape
+        up = F.interpolate(d1.permute(0, 1, 4, 2, 3).reshape(b * k, c, h, w), scale_factor=2, mode="nearest")
+        d0 = self.up_project(up).reshape(b, k, -1, h * 2, w * 2).permute(0, 1, 3, 4, 2)
+        d0 = d0 + self.skip_project(e0.mean(dim=1, keepdim=True)).expand(-1, k, -1, -1, -1)
+        gd0 = self._global("decoder_global0", batch)
+        for block in self.decoder0:
+            d0, gd0 = block(d0, e0, gd0)
+        b, k, h, w, c = d0.shape
+        full = F.interpolate(d0.permute(0, 1, 4, 2, 3).reshape(b * k, c, h, w), scale_factor=2, mode="nearest")
+        output = self.head(full)
+        return output.reshape(b, k, -1, self.height, self.width).permute(0, 1, 3, 4, 2)

script/data_loader.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from script.fake_data import SPLIT_OFFSETS, generate_sequence
+
+
+class SyntheticSEVIRDataset(Dataset):
+    def __init__(self, size: int, config: dict, seed: int):
+        data = config["data"]
+        self.size = size
+        self.input_length = int(data["input_length"])
+        self.output_length = int(data["output_length"])
+        self.height = int(data["height"])
+        self.width = int(data["width"])
+        self.seed = seed
+
+    def __len__(self) -> int:
+        return self.size
+
+    def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]:
+        sequence = generate_sequence(
+            self.height, self.width, self.input_length, self.output_length, self.seed + index
+        )
+        tensor = torch.from_numpy(sequence)
+        return tensor[: self.input_length], tensor[self.input_length :]
+
+
+class NPZSequenceDataset(Dataset):
+    def __init__(self, path: str | Path, config: dict):
+        data = config["data"]
+        with np.load(path) as payload:
+            if "inputs" not in payload or "targets" not in payload:
+                raise ValueError("NPZ must contain 'inputs' and 'targets'")
+            inputs, targets = payload["inputs"], payload["targets"]
+        expected_input = (int(data["input_length"]), int(data["height"]), int(data["width"]), int(data["channels"]))
+        expected_target = (int(data["output_length"]), int(data["height"]), int(data["width"]), int(data["channels"]))
+        if inputs.ndim != 5 or tuple(inputs.shape[1:]) != expected_input:
+            raise ValueError(f"inputs must have shape [N,{','.join(map(str, expected_input))}], got {inputs.shape}")
+        if targets.ndim != 5 or tuple(targets.shape[1:]) != expected_target:
+            raise ValueError(f"targets must have shape [N,{','.join(map(str, expected_target))}], got {targets.shape}")
+        if len(inputs) != len(targets) or len(inputs) == 0:
+            raise ValueError("inputs and targets must have the same non-zero sample count")
+        normalization = data.get("normalization", "unit")
+        if normalization == "uint8_255":
+            if inputs.dtype != np.uint8 or targets.dtype != np.uint8:
+                raise ValueError("uint8_255 normalization requires uint8 NPZ arrays")
+            inputs, targets = inputs.astype(np.float32) / 255.0, targets.astype(np.float32) / 255.0
+        else:
+            if not np.issubdtype(inputs.dtype, np.floating) or not np.issubdtype(targets.dtype, np.floating):
+                raise ValueError("unit normalization requires floating-point NPZ arrays")
+            inputs, targets = inputs.astype(np.float32), targets.astype(np.float32)
+            if not np.isfinite(inputs).all() or not np.isfinite(targets).all():
+                raise ValueError("NPZ arrays contain non-finite values")
+            if inputs.min() < 0 or inputs.max() > 1 or targets.min() < 0 or targets.max() > 1:
+                raise ValueError("unit-normalized NPZ arrays must be within [0,1]; float values are never implicitly divided by 255")
+        self.inputs = inputs
+        self.targets = targets
+
+    def __len__(self) -> int:
+        return len(self.inputs)
+
+    def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.from_numpy(self.inputs[index]), torch.from_numpy(self.targets[index])
+
+
+def _seed_worker(worker_id: int) -> None:
+    del worker_id
+    worker_seed = torch.initial_seed() % 2**32
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)
+
+
+def make_loader(
+    config: dict,
+    split: str,
+    distributed: bool = False,
+    rank: int = 0,
+    world_size: int = 1,
+    shuffle: bool | None = None,
+) -> tuple[DataLoader, DistributedSampler | None]:
+    if split not in SPLIT_OFFSETS:
+        raise ValueError(f"unknown split: {split}")
+    data, train = config["data"], config["train"]
+    path_value = data.get(f"{split}_npz")
+    path = Path(path_value) if path_value else None
+    if path is not None and path.is_file():
+        dataset: Dataset = NPZSequenceDataset(path, config)
+    elif bool(data.get("fallback_if_missing", True)):
+        dataset = SyntheticSEVIRDataset(
+            int(data[f"{split}_samples"]), config, int(train["seed"]) + SPLIT_OFFSETS[split]
+        )
+    else:
+        raise FileNotFoundError(f"configured {split} NPZ does not exist: {path}")
+    should_shuffle = split == "train" if shuffle is None else shuffle
+    sampler = None
+    if distributed:
+        sampler = DistributedSampler(
+            dataset, num_replicas=world_size, rank=rank, shuffle=should_shuffle, seed=int(train["seed"]), drop_last=False
+        )
+    generator = torch.Generator().manual_seed(int(train["seed"]) + SPLIT_OFFSETS[split] + rank)
+    options = config["dataloader"]
+    loader = DataLoader(
+        dataset,
+        batch_size=int(train["batch_size"]),
+        shuffle=should_shuffle and sampler is None,
+        sampler=sampler,
+        num_workers=int(options.get("num_workers", 0)),
+        pin_memory=bool(options.get("pin_memory", False)),
+        worker_init_fn=_seed_worker,
+        generator=generator,
+    )
+    return loader, sampler

script/fake_data.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+if __package__ in (None, ""):
+    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from script.utils import DEFAULT_CONFIG, load_config
+
+
+SPLIT_OFFSETS = {"train": 0, "val": 100_000, "test": 200_000}
+
+# Matches the official Earthformer SEVIR config
+# (scripts/cuboid_transformer/sevir/earthformer_sevir_v1.yaml):
+#   dataset.img_height/img_width = 384, in_len = 13, out_len = 12,
+#   seq_len = 25, interval_real_time = 5, sample_mode = "sequent",
+#   stride = 12, metrics_list = ['csi', 'pod', 'sucr', 'bias'],
+#   threshold_list = [16, 74, 133, 160, 181, 219].
+SEVIR_VIL_REFERENCE = {
+    "dataset": "SEVIR VIL",
+    "spatial_shape": [384, 384, 1],
+    "sequence": {"input_frames": 13, "output_frames": 12},
+    "seq_len": 25,
+    "sample_mode": "sequent",
+    "stride": 12,
+    "frame_interval_minutes": 5,
+    "thresholds": [16, 74, 133, 160, 181, 219],
+    "layout": "train/val/test NPZ splits; each holds inputs [N,13,H,W,1] and targets [N,12,H,W,1]",
+}
+
+
+def generate_sequence(height: int, width: int, input_length: int, output_length: int, seed: int) -> np.ndarray:
+    """Generate a continuous, temporally coherent 25-frame SEVIR-like VIL window.
+
+    The window is sampled exactly like official SEVIR "sequent" mode: a single
+    continuous sequence of seq_len = input_length + output_length frames is
+    produced and later split into the first 13 input and the last 12 target
+    frames, matching the official 13 -> 12 (65 -> 60 minute) task at 5-minute
+    intervals.
+    """
+    rng = np.random.default_rng(seed)
+    total = input_length + output_length
+    yy, xx = np.mgrid[:height, :width]
+    background = rng.normal(0.0, 0.008, (height, width)).astype(np.float32)
+    background = (background + np.roll(background, 1, 0) + np.roll(background, 1, 1)) / 3.0
+    cell_count = int(rng.integers(2, 5))
+    cells = []
+    for _ in range(cell_count):
+        cells.append(
+            (
+                rng.uniform(0.15 * width, 0.85 * width),
+                rng.uniform(0.15 * height, 0.85 * height),
+                rng.uniform(-0.45, 0.45),
+                rng.uniform(-0.45, 0.45),
+                rng.uniform(max(1.2, width / 18), max(2.0, width / 8)),
+                rng.uniform(max(1.2, height / 18), max(2.0, height / 8)),
+                rng.uniform(0.45, 0.95),
+                rng.uniform(-0.035, 0.035),
+                rng.uniform(0, np.pi),
+            )
+        )
+    frames = np.empty((total, height, width, 1), dtype=np.float32)
+    for time in range(total):
+        frame = np.maximum(background * (0.8 + 0.2 * np.sin(time / 5)), 0.0)
+        for cx, cy, vx, vy, sx, sy, amplitude, growth, angle in cells:
+            dx, dy = xx - (cx + vx * time), yy - (cy + vy * time)
+            ca, sa = np.cos(angle), np.sin(angle)
+            xr, yr = ca * dx + sa * dy, -sa * dx + ca * dy
+            scale = np.clip(1.0 + growth * time, 0.55, 1.6)
+            intensity = amplitude * np.exp(-0.5 * ((xr / (sx * scale)) ** 2 + (yr / (sy * scale)) ** 2))
+            lifecycle = np.clip(1.0 + growth * time, 0.35, 1.25)
+            frame += intensity.astype(np.float32) * lifecycle
+        noise = rng.normal(0.0, 0.004, (height, width)).astype(np.float32)
+        frames[time, ..., 0] = np.clip(frame + noise, 0.0, 1.0)
+    return frames
+
+
+def generate_split(samples: int, height: int, width: int, input_length: int, output_length: int, seed: int) -> tuple[np.ndarray, np.ndarray]:
+    sequences = np.stack(
+        [generate_sequence(height, width, input_length, output_length, seed + index) for index in range(samples)]
+    )
+    return sequences[:, :input_length], sequences[:, input_length:]
+
+
+def parse_args() -> argparse.Namespace:
+    config_parser = argparse.ArgumentParser(add_help=False)
+    config_parser.add_argument("--config", default=str(DEFAULT_CONFIG))
+    config_args, _ = config_parser.parse_known_args()
+    defaults = load_config(config_args.config)
+    data = defaults["data"]
+    parser = argparse.ArgumentParser(description="Generate deterministic synthetic SEVIR-like VIL sequences")
+    parser.add_argument("--config", default=config_args.config)
+    parser.add_argument("--output-dir", default=data["data_dir"])
+    parser.add_argument("--height", type=int, default=int(data["height"]))
+    parser.add_argument("--width", type=int, default=int(data["width"]))
+    parser.add_argument("--train-samples", type=int, default=int(data["train_samples"]))
+    parser.add_argument("--val-samples", type=int, default=int(data["val_samples"]))
+    parser.add_argument("--test-samples", type=int, default=int(data["test_samples"]))
+    parser.add_argument("--seed", type=int, default=int(defaults["train"]["seed"]))
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    config = load_config(args.config)
+    data = config["data"]
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    split_sizes = {"train": args.train_samples, "val": args.val_samples, "test": args.test_samples}
+    for split, samples in split_sizes.items():
+        if samples <= 0:
+            raise ValueError(f"{split} samples must be positive")
+        inputs, targets = generate_split(
+            samples,
+            args.height,
+            args.width,
+            int(data["input_length"]),
+            int(data["output_length"]),
+            args.seed + SPLIT_OFFSETS[split],
+        )
+        np.savez_compressed(output_dir / f"{split}.npz", inputs=inputs, targets=targets)
+    metadata = {
+        "synthetic": True,
+        "official_sevir": False,
+        "description": "Deterministic synthetic SEVIR-like VIL; not official SEVIR data",
+        "protocol": "synthetic_sevir",
+        "reference": SEVIR_VIL_REFERENCE,
+        "seq_len": int(data["input_length"]) + int(data["output_length"]),
+        "sample_mode": "sequent",
+        "stride": SEVIR_VIL_REFERENCE["stride"],
+        "frame_interval_minutes": int(data["frame_interval_minutes"]),
+        "input_frames": int(data["input_length"]),
+        "output_frames": int(data["output_length"]),
+        "shape": [args.height, args.width, 1],
+        "normalization": "unit [0,1] float32",
+        "seed": args.seed,
+        "splits": split_sizes,
+    }
+    (output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")
+    print(json.dumps({"output_dir": str(output_dir), "metadata": metadata}, indent=2))
+
+
+if __name__ == "__main__":
+    main()

script/inference.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+if __package__ in (None, ""):
+    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from model.earthformer import Earthformer
+from script.data_loader import make_loader
+from script.utils import clean_state_dict, load_checkpoint_payload, load_config, resolve_cli_path, resolve_device
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run Earthformer inference on the first batch")
+    parser.add_argument("--config", help="Optional data/config override; checkpoint config is used by default")
+    parser.add_argument("--checkpoint", default="data/checkpoint/earthformer.pt")
+    parser.add_argument("--split", choices=("train", "val", "test"), default="test")
+    parser.add_argument("--output", default="output/predictions.npz")
+    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    device = resolve_device(args.device)
+    payload = load_checkpoint_payload(resolve_cli_path(args.checkpoint), device)
+    config = load_config(args.config) if args.config else payload["config"]
+    model = Earthformer(config).to(device)
+    model.load_state_dict(clean_state_dict(payload["model"]))
+    model.eval()
+    loader, _ = make_loader(config, args.split, shuffle=False)
+    inputs, targets = next(iter(loader))
+    with torch.no_grad():
+        predictions = model(inputs.to(device)).clamp(0.0, 1.0).cpu().numpy()
+    output = Path(resolve_cli_path(args.output))
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, inputs=inputs.numpy(), targets=targets.numpy(), predictions=predictions)
+    print(json.dumps({"output": str(output), "shape": list(predictions.shape)}, indent=2))
+
+
+if __name__ == "__main__":
+    main()

script/metrics.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import torch
+
+
+CSI_THRESHOLDS = torch.tensor([16, 74, 133, 160, 181, 219], dtype=torch.float32) / 255.0
+
+
+def mse(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    return torch.mean((prediction - target) ** 2)
+
+
+def _ssim_per_sample(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    prediction, target = prediction.clamp(0.0, 1.0), target.clamp(0.0, 1.0)
+    mu_x, mu_y = prediction.mean((2, 3)), target.mean((2, 3))
+    centered_x = prediction - mu_x[:, :, None, None, :]
+    centered_y = target - mu_y[:, :, None, None, :]
+    var_x = centered_x.square().mean((2, 3))
+    var_y = centered_y.square().mean((2, 3))
+    covariance = (centered_x * centered_y).mean((2, 3))
+    c1, c2 = 0.01**2, 0.03**2
+    score = ((2 * mu_x * mu_y + c1) * (2 * covariance + c2)) / (
+        (mu_x.square() + mu_y.square() + c1) * (var_x + var_y + c2)
+    )
+    return score.mean((1, 2))
+
+
+def metric_sums(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    """Return additive sample-level MSE/MAE/SSIM sums and CSI event counts."""
+    prediction, target = prediction.float(), target.float()
+    reduce_dims = tuple(range(1, prediction.ndim))
+    values = [
+        (prediction - target).square().mean(reduce_dims).sum(),
+        (prediction - target).abs().mean(reduce_dims).sum(),
+        _ssim_per_sample(prediction, target).sum(),
+        prediction.new_tensor(prediction.shape[0]),
+    ]
+    for threshold in CSI_THRESHOLDS.to(prediction.device):
+        predicted, observed = prediction >= threshold, target >= threshold
+        values.extend([(predicted & observed).sum(), (predicted & ~observed).sum(), (~predicted & observed).sum()])
+    return torch.stack(values).to(torch.float64)
+
+
+def metrics_from_sums(sums: torch.Tensor) -> dict[str, float]:
+    count = sums[3].clamp_min(1)
+    csi = []
+    for index in range(len(CSI_THRESHOLDS)):
+        hits, false_alarms, misses = sums[4 + index * 3 : 7 + index * 3]
+        csi.append((hits / (hits + false_alarms + misses).clamp_min(1)).item())
+    return {
+        "mse": (sums[0] / count).item(),
+        "mae": (sums[1] / count).item(),
+        "ssim": (sums[2] / count).item(),
+        "mean_csi": sum(csi) / len(csi),
+    }
+
+
+def metric_sums_light(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    """Additive MSE/MAE sums for cheap per-epoch validation: [mse_sum, mae_sum, count]."""
+    prediction, target = prediction.float(), target.float()
+    reduce_dims = tuple(range(1, prediction.ndim))
+    return torch.stack(
+        [
+            (prediction - target).square().mean(reduce_dims).sum(),
+            (prediction - target).abs().mean(reduce_dims).sum(),
+            prediction.new_tensor(prediction.shape[0]),
+        ]
+    ).to(torch.float64)
+
+
+def metrics_from_light_sums(sums: torch.Tensor) -> dict[str, float]:
+    count = sums[2].clamp_min(1)
+    return {"mse": (sums[0] / count).item(), "mae": (sums[1] / count).item()}
+
+
+def compute_metrics(prediction: torch.Tensor, target: torch.Tensor) -> dict[str, float]:
+    return metrics_from_sums(metric_sums(prediction.detach(), target.detach()))

script/result.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+if __package__ in (None, ""):
+    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from model.earthformer import Earthformer
+from script.data_loader import make_loader
+from script.metrics import metric_sums, metrics_from_sums
+from script.utils import clean_state_dict, load_checkpoint_payload, resolve_cli_path, resolve_device
+
+
+def _squeeze(frame: np.ndarray) -> np.ndarray:
+    return np.squeeze(frame)
+
+
+def plot_lead_time(truth: np.ndarray, prediction: np.ndarray, output_path: Path, title: str) -> Path:
+    """Three-panel Truth / Prediction / Difference image, fuxi/fengwu style."""
+    truth, prediction = _squeeze(truth), _squeeze(prediction)
+    difference = truth - prediction
+    rmse = float(np.sqrt(np.mean(difference**2)))
+    vmin, vmax = min(truth.min(), prediction.min()), max(truth.max(), prediction.max())
+    diff_abs_max = max(float(np.abs(difference).max()), 1e-8)
+
+    panels = [
+        {"data": truth, "title": "Truth", "cmap": "viridis", "vmin": vmin, "vmax": vmax},
+        {"data": prediction, "title": "Prediction", "cmap": "viridis", "vmin": vmin, "vmax": vmax},
+        {"data": difference, "title": f"Difference (RMSE={rmse:.2f})", "cmap": "RdBu_r", "vmin": -diff_abs_max, "vmax": diff_abs_max},
+    ]
+    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
+    for ax, panel in zip(axes, panels):
+        image = ax.imshow(panel["data"], cmap=panel["cmap"], vmin=panel["vmin"], vmax=panel["vmax"])
+        ax.set_title(panel["title"], fontsize=12, pad=4)
+        ax.set_xlabel("Pixel")
+        ax.set_ylabel("Pixel")
+        plt.colorbar(image, ax=ax, orientation="horizontal")
+    fig.suptitle(title, fontsize=14, fontweight="bold", y=0.98)
+    fig.savefig(output_path, dpi=300, bbox_inches="tight")
+    plt.close(fig)
+    return output_path
+
+
+def visualize_predictions(
+    truth: np.ndarray,
+    prediction: np.ndarray,
+    output_dir: str | Path,
+    max_samples: int = 2,
+    stride: int = 2,
+) -> list[str]:
+    """Write per-sample/lead-time PNG images and a per-sample mosaic into output_dir."""
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    if truth.shape != prediction.shape:
+        raise ValueError(f"truth and prediction must have the same shape, got {truth.shape} vs {prediction.shape}")
+    samples = min(max_samples, int(truth.shape[0]))
+    lead_times = list(range(0, int(truth.shape[1]), stride))
+    generated: list[str] = []
+    for sample in range(samples):
+        for lead_time in lead_times:
+            path = output_dir / f"earthformer_sample_{sample:03d}_t{lead_time:02d}.png"
+            plot_lead_time(
+                truth[sample, lead_time],
+                prediction[sample, lead_time],
+                path,
+                f"Earthformer VIL sample {sample} lead time {lead_time} (+{lead_time * 5} min)",
+            )
+            generated.append(str(path))
+        mosaic = output_dir / f"earthformer_overview_sample_{sample:03d}.png"
+        _plot_mosaic(truth[sample], prediction[sample], lead_times, mosaic)
+        generated.append(str(mosaic))
+    return generated
+
+
+def _plot_mosaic(truth: np.ndarray, prediction: np.ndarray, lead_times: list[int], output_path: Path) -> Path:
+    rows, columns = len(lead_times) * 2, len(lead_times)
+    fig, axes = plt.subplots(rows, columns, figsize=(columns * 2.6, rows * 2.2))
+    for column, lead_time in enumerate(lead_times):
+        truth_frame = _squeeze(truth[lead_time])
+        pred_frame = _squeeze(prediction[lead_time])
+        vmin, vmax = min(truth_frame.min(), pred_frame.min()), max(truth_frame.max(), pred_frame.max())
+        axes[0, column].imshow(truth_frame, cmap="viridis", vmin=vmin, vmax=vmax)
+        axes[0, column].set_title(f"+{lead_time * 5} min", fontsize=9)
+        axes[len(lead_times), column].imshow(pred_frame, cmap="viridis", vmin=vmin, vmax=vmax)
+        for row in range(rows):
+            axes[row, column].set_xticks([])
+            axes[row, column].set_yticks([])
+    axes[0, 0].set_ylabel("Truth", fontsize=10)
+    axes[len(lead_times), 0].set_ylabel("Prediction", fontsize=10)
+    fig.suptitle("Earthformer VIL sample overview", fontsize=13, fontweight="bold", y=0.99)
+    fig.tight_layout(rect=(0, 0, 1, 0.97))
+    fig.savefig(output_path, dpi=200, bbox_inches="tight")
+    plt.close(fig)
+    return output_path
+
+
+def evaluate(model: torch.nn.Module, loader, device: torch.device) -> dict[str, float]:
+    sums = torch.zeros(22, dtype=torch.float64, device=device)
+    with torch.no_grad():
+        for inputs, targets in loader:
+            prediction = model(inputs.to(device)).clamp(0.0, 1.0)
+            sums += metric_sums(prediction, targets.to(device))
+    result = metrics_from_sums(sums.cpu())
+    result["note"] = "Lightweight metrics on configured data; mean CSI is not the official complete SEVIR evaluation"
+    return result
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Evaluate Earthformer and visualize predictions (fuxi/fengwu style PNG images)"
+    )
+    parser.add_argument("--checkpoint", default="data/checkpoint/earthformer.pt")
+    parser.add_argument("--split", choices=("train", "val", "test"), default="test")
+    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
+    parser.add_argument("--predictions", default="output/predictions.npz", help="Inference arrays for visualization")
+    parser.add_argument("--output-dir", default="output/visualization", help="Directory for visualization PNG images")
+    parser.add_argument("--max-samples", type=int, default=2)
+    parser.add_argument("--stride", type=int, default=2, help="Lead-time stride, aligned with official plot_stride")
+    parser.add_argument("--skip-visualization", action="store_true", help="Only evaluate, do not render PNG images")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    device = resolve_device(args.device)
+    payload = load_checkpoint_payload(resolve_cli_path(args.checkpoint), device)
+    config = payload["config"]
+    model = Earthformer(config).to(device)
+    model.load_state_dict(clean_state_dict(payload["model"]))
+    model.eval()
+    loader, _ = make_loader(config, args.split, shuffle=False)
+    result = evaluate(model, loader, device)
+    report = {"metrics": result}
+    if args.skip_visualization:
+        print(json.dumps(report, indent=2))
+        return
+    predictions_path = Path(resolve_cli_path(args.predictions))
+    with np.load(predictions_path) as payload:
+        truth = payload["targets"]
+        prediction = payload["predictions"]
+    output_dir = resolve_cli_path(args.output_dir)
+    images = visualize_predictions(truth, prediction, output_dir, args.max_samples, args.stride)
+    report["output_dir"] = output_dir
+    report["images"] = images
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()

script/train.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+
+if __package__ in (None, ""):
+    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from model.earthformer import Earthformer
+from script.data_loader import make_loader
+from script.metrics import metric_sums, metric_sums_light, metrics_from_light_sums, metrics_from_sums, mse
+from script.utils import DEFAULT_CONFIG, atomic_torch_save, clean_state_dict, load_checkpoint_payload, load_config, resolve_device
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def validate(
+    model: torch.nn.Module,
+    loader,
+    device: torch.device,
+    max_batches: int | None = None,
+    full_metrics: bool = False,
+) -> dict[str, float]:
+    """Run validation with only cheap MSE/MAE by default.
+
+    The training loss is already the per-sample MSE, so recomputing MSE during
+    training is redundant monitoring. SSIM and per-threshold CSI are expensive
+    on full-resolution 384x384 data, so they are computed only when
+    `train.compute_full_metrics` is enabled; the authoritative full evaluation
+    lives in script/result.py.
+    """
+    model.eval()
+    sums = torch.zeros(22 if full_metrics else 3, dtype=torch.float64, device=device)
+    with torch.no_grad():
+        for batch_index, (inputs, targets) in enumerate(loader):
+            predictions = model(inputs.to(device, non_blocking=True)).clamp(0.0, 1.0)
+            targets_device = targets.to(device, non_blocking=True)
+            if full_metrics:
+                sums += metric_sums(predictions, targets_device)
+            else:
+                sums += metric_sums_light(predictions, targets_device)
+            if max_batches is not None and batch_index + 1 >= max_batches:
+                break
+    if dist.is_initialized():
+        dist.all_reduce(sums, op=dist.ReduceOp.SUM)
+    count_index = 3 if full_metrics else 2
+    if sums[count_index].item() == 0:
+        raise ValueError("validation loader is empty")
+    return (metrics_from_sums if full_metrics else metrics_from_light_sums)(sums.cpu())
+
+
+def run_training(config: dict, requested_device: str = "auto", resume: str | None = None) -> tuple[Path, dict[str, float]]:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    rank = int(os.environ.get("RANK", "0"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    distributed = world_size > 1
+    device = resolve_device(requested_device, local_rank)
+    backend = config["distributed"].get("backend", "auto")
+    if backend == "auto":
+        backend = "nccl" if device.type == "cuda" else "gloo"
+    if distributed:
+        dist.init_process_group(backend=backend, rank=rank, world_size=world_size)
+    try:
+        seed_everything(int(config["train"]["seed"]))
+        model = Earthformer(config).to(device)
+        optimizer = torch.optim.AdamW(
+            model.parameters(),
+            lr=float(config["train"]["learning_rate"]),
+            weight_decay=float(config["train"]["weight_decay"]),
+        )
+        start_epoch, step = 0, 0
+        if resume:
+            payload = load_checkpoint_payload(resume, device)
+            model.load_state_dict(clean_state_dict(payload["model"]))
+            if "optimizer" in payload:
+                optimizer.load_state_dict(payload["optimizer"])
+            start_epoch = int(payload.get("epoch", -1)) + 1
+            step = int(payload.get("step", 0))
+        if distributed:
+            model = DistributedDataParallel(
+                model,
+                device_ids=[local_rank] if device.type == "cuda" else None,
+                find_unused_parameters=True,
+            )
+        train_loader, train_sampler = make_loader(config, "train", distributed, rank, world_size)
+        val_loader, _ = make_loader(config, "val", distributed, rank, world_size, shuffle=False)
+        full_metrics = bool(config["train"].get("compute_full_metrics", False))
+        last_epoch = max(start_epoch - 1, 0)
+        metrics: dict[str, float] = {}
+        for epoch in range(start_epoch, int(config["train"]["epochs"])):
+            last_epoch = epoch
+            if train_sampler is not None:
+                train_sampler.set_epoch(epoch)
+            model.train()
+            epoch_loss_sum, epoch_steps = 0.0, 0
+            for inputs, targets in train_loader:
+                optimizer.zero_grad(set_to_none=True)
+                loss = mse(model(inputs.to(device, non_blocking=True)), targets.to(device, non_blocking=True))
+                if not torch.isfinite(loss):
+                    raise FloatingPointError("training loss is not finite")
+                loss.backward()
+                optimizer.step()
+                epoch_loss_sum += float(loss.detach().item())
+                epoch_steps += 1
+                step += 1
+            if epoch_steps == 0:
+                break
+            metrics = validate(model, val_loader, device, int(config["train"].get("validation_steps", 1)), full_metrics)
+            if rank == 0:
+                epoch_line = {"epoch": epoch, "train_loss": epoch_loss_sum / epoch_steps, "validation": metrics}
+                print(json.dumps(epoch_line))
+        checkpoint = Path(config["train"]["output_dir"]) / "earthformer.pt"
+        if rank == 0:
+            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+            atomic_torch_save(
+                {
+                    "model": raw_model.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                    "config": config,
+                    "metrics": metrics,
+                    "epoch": last_epoch,
+                    "step": step,
+                    "world_size": world_size,
+                },
+                checkpoint,
+            )
+            print(json.dumps({"checkpoint": str(checkpoint), "step": step, "world_size": world_size, "metrics": metrics}, indent=2))
+        if distributed:
+            dist.barrier()
+        return checkpoint, metrics
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Train Earthformer with one process or DDP")
+    parser.add_argument("--config", default=str(DEFAULT_CONFIG))
+    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
+    parser.add_argument("--resume")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_training(load_config(args.config), args.device, args.resume)

script/utils.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any
+
+import torch
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_CONFIG = ROOT / "conf/config.yaml"
+
+
+def load_config(path: str | Path = DEFAULT_CONFIG) -> dict[str, Any]:
+    config_path = Path(path).expanduser().resolve()
+    with config_path.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    if not isinstance(config, dict):
+        raise ValueError("configuration must be a YAML mapping")
+    for section in ("data", "model", "train", "dataloader", "distributed"):
+        if section not in config:
+            raise ValueError(f"missing configuration section: {section}")
+    data, model = config["data"], config["model"]
+    for key in ("input_length", "output_length", "height", "width", "channels"):
+        if int(data.get(key, 0)) <= 0:
+            raise ValueError(f"data.{key} must be positive")
+    if int(data["height"]) % 4 or int(data["width"]) % 4:
+        raise ValueError("data height and width must be divisible by four")
+    dims, depths = model.get("dims"), model.get("depths")
+    if not isinstance(dims, list) or len(dims) != 2 or int(dims[1]) != 2 * int(dims[0]):
+        raise ValueError("model.dims must be [D, 2*D]")
+    if not isinstance(depths, list) or len(depths) != 2 or min(int(x) for x in depths) < 1:
+        raise ValueError("model.depths must contain two positive integers")
+    heads = int(model.get("heads", 0))
+    if heads < 1 or any(int(dim) % heads for dim in dims):
+        raise ValueError("model.heads must divide both hidden dimensions")
+    normalization = data.get("normalization", "unit")
+    if normalization not in ("unit", "uint8_255"):
+        raise ValueError("data.normalization must be 'unit' or 'uint8_255'")
+    for key in ("data_dir", "train_npz", "val_npz", "test_npz"):
+        if data.get(key):
+            value = Path(data[key]).expanduser()
+            data[key] = str(value if value.is_absolute() else ROOT / value)
+    output_dir = Path(config["train"]["output_dir"]).expanduser()
+    config["train"]["output_dir"] = str(output_dir if output_dir.is_absolute() else ROOT / output_dir)
+    return config
+
+
+def resolve_device(requested: str, local_rank: int = 0) -> torch.device:
+    if requested not in ("auto", "cpu", "cuda"):
+        raise ValueError("device must be auto, cpu, or cuda")
+    use_accelerator = requested == "cuda" or (requested == "auto" and torch.cuda.is_available())
+    if use_accelerator:
+        if not torch.cuda.is_available():
+            raise RuntimeError("CUDA/ROCm device requested but torch.cuda.is_available() is false")
+        torch.cuda.set_device(local_rank)
+        return torch.device("cuda", local_rank)
+    return torch.device("cpu")
+
+
+def resolve_cli_path(value: str | None) -> str | None:
+    if not value:
+        return value
+    path = Path(value).expanduser()
+    return str(path if path.is_absolute() else ROOT / path)
+
+
+def load_checkpoint_payload(path: str | Path, device: torch.device) -> dict[str, Any]:
+    payload = torch.load(path, map_location=device, weights_only=False)
+    if not isinstance(payload, dict) or "model" not in payload or "config" not in payload:
+        raise ValueError("checkpoint must contain model and config")
+    return payload
+
+
+def clean_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    return {key.removeprefix("module."): value for key, value in state_dict.items()}
+
+
+def atomic_torch_save(payload: dict[str, Any], path: str | Path) -> Path:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp")
+    torch.save(payload, temporary)
+    os.replace(temporary, path)
+    return path