Upload folder using huggingface_hub

README.md

@@ -0,0 +1,158 @@
+---
+datasets:
+- OneScience/ERA5
+frameworks:
+- PyTorch
+language:
+- en
+- zh
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Weather Forecast
+- Regional Forecast
+- Diffusion Model
+- ERA5
+- HRRR
+tasks: []
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">StormCast</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+StormCast is a generative regional weather forecasting model proposed by NVIDIA, targeting high-resolution nowcasting of mesoscale convective weather.
+
+Paper: StormCast: A Machine Learning Method for Meso-β-Scale Convection-resolving Weather Forecasting
+
+https://arxiv.org/abs/2408.10958
+
+# Model Description
+
+StormCast constrains the evolution of regional states with large-scale weather backgrounds, and uses a generative diffusion approach to supplement the fine-scale structures that deterministic forecasts struggle to represent.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Two-Stage Weather Forecast Training | Train a deterministic regression model and a conditional residual diffusion model in sequence. |
+| 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
+
+**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/StormCast --local-dir ./StormCast
+cd StormCast
+```
+
+### 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
+```
+
+### Training Data Introduction
+
+The OneScience community provides ERA5 data for training (due to file size limits, the current repository contains a slice of the full dataset). Users can download it with the command below and confirm that the data path in `conf/config.yaml` is set correctly:
+
+```bash
+hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data
+```
+
+### Generate Synthetic Data for Pipeline Validation
+
+```bash
+python scripts/fake_data.py
+```
+
+Synthetic data is only used to verify the data protocol and program flow; it does not represent the model's scientific forecasting capability.
+
+### Training
+
+Single GPU:
+
+```bash
+# Train the deterministic regression model; weights are saved to data/checkpoint/regression/model_bak.pt by default
+python scripts/train.py --stage regression
+# Train the residual diffusion model; weights are saved to data/checkpoint/diffusion/model_bak.pt by default
+python scripts/train.py --stage diffusion
+```
+
+### Multi-GPU
+
+```bash
+# Train the deterministic regression model
+torchrun --nproc_per_node=2 scripts/train.py --stage regression
+# Train the residual diffusion model
+torchrun --nproc_per_node=2 scripts/train.py --stage diffusion
+```
+
+### Training Weights
+
+This repository provides weights trained on ERA5 reanalysis data in the `weights/` folder. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+Run autoregressive prediction with the default configuration and the two sets of weights saved during training:
+
+```bash
+python scripts/inference.py
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Plots are saved to `outputs/inference/plots/` by default.
+
+# 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 a reproduction of the original StormCast paper.

conf/config.yaml

@@ -0,0 +1,191 @@
+project:
+  name: stormcast
+  seed: 42
+  output_dir: ./outputs
+
+data:
+  root_dir: ./data
+  train_years: [2019]
+  val_years: [2021]
+  test_years: [2022]
+  time_step_hours: 1
+  input_steps: 1
+  output_steps: 1
+  era5_image_size: [721, 1440]
+  image_size: [512, 640]
+  num_timesteps: 4
+  normalize: true
+  era5_variables:
+    - u10m
+    - v10m
+    - t2m
+    - tcwv
+    - sp
+    - msl
+    - u1000
+    - u850
+    - u500
+    - u250
+    - v1000
+    - v850
+    - v500
+    - v250
+    - z1000
+    - z850
+    - z500
+    - z250
+    - t1000
+    - t850
+    - t500
+    - t250
+    - q1000
+    - q850
+    - q500
+    - q250
+  state_variables:
+    - u10m
+    - v10m
+    - t2m
+    - msl
+    - u1hl
+    - u2hl
+    - u3hl
+    - u4hl
+    - u5hl
+    - u6hl
+    - u7hl
+    - u8hl
+    - u9hl
+    - u10hl
+    - u11hl
+    - u13hl
+    - u15hl
+    - u20hl
+    - u25hl
+    - u30hl
+    - v1hl
+    - v2hl
+    - v3hl
+    - v4hl
+    - v5hl
+    - v6hl
+    - v7hl
+    - v8hl
+    - v9hl
+    - v10hl
+    - v11hl
+    - v13hl
+    - v15hl
+    - v20hl
+    - v25hl
+    - v30hl
+    - t1hl
+    - t2hl
+    - t3hl
+    - t4hl
+    - t5hl
+    - t6hl
+    - t7hl
+    - t8hl
+    - t9hl
+    - t10hl
+    - t11hl
+    - t13hl
+    - t15hl
+    - t20hl
+    - t25hl
+    - t30hl
+    - q1hl
+    - q2hl
+    - q3hl
+    - q4hl
+    - q5hl
+    - q6hl
+    - q7hl
+    - q8hl
+    - q9hl
+    - q10hl
+    - q11hl
+    - q13hl
+    - q15hl
+    - q20hl
+    - q25hl
+    - q30hl
+    - Z1hl
+    - Z2hl
+    - Z3hl
+    - Z4hl
+    - Z5hl
+    - Z6hl
+    - Z7hl
+    - Z8hl
+    - Z9hl
+    - Z10hl
+    - Z11hl
+    - Z13hl
+    - Z15hl
+    - Z20hl
+    - Z25hl
+    - Z30hl
+    - p1hl
+    - p2hl
+    - p3hl
+    - p4hl
+    - p5hl
+    - p6hl
+    - p7hl
+    - p8hl
+    - p9hl
+    - p10hl
+    - p11hl
+    - p13hl
+    - p15hl
+    - p20hl
+    - refc
+  invariant_variables: [lsm, orography]
+
+dataloader:
+  batch_size: 1
+  num_workers: 0
+  pin_memory: true
+
+model:
+  regression_conditions: [state, background, invariant]
+  diffusion_conditions: [state, regression, invariant]
+  image_size: [512, 640]
+  model_channels: 128
+  channel_mult: [1, 2, 2, 2, 2]
+  num_blocks: 4
+  attention_resolutions: []
+  regression_weights: ./data/checkpoint/regression/model_bak.pt
+  diffusion_weights: ./data/checkpoint/diffusion/model_bak.pt
+
+training:
+  stage: regression
+  from_scratch: true
+  initial_weights: null
+  resume_checkpoint: null
+  regression_weights: ./data/checkpoint/regression/model_bak.pt
+  epochs: 10
+  max_steps: 10
+  learning_rate: 0.0002
+  betas: [0.9, 0.999]
+  weight_decay: 0.0
+  gradient_clip_norm: null
+  checkpoint_interval: 1
+  log_interval: 1
+  sigma_data: 0.5
+  P_mean: -1.2
+  P_std: 1.2
+  checkpoint_dir: ./data/checkpoint
+
+inference:
+  split: test
+  num_steps: 1
+  diffusion_steps: 18
+  sigma_min: 0.002
+  sigma_max: 800.0
+  rho: 7.0
+  output_dir: ./outputs/inference
+  plot_state_variable: refc
+  plot_background_variable: t2m

config.json

@@ -0,0 +1,119 @@
+{
+  "model_name": "StormCast",
+  "model_type": "stormcast",
+  "architectures": [
+    "StormCast"
+  ],
+  "framework": "PyTorch",
+  "domain": "atmosphere",
+  "task": "regional-convection-resolving-weather-forecasting",
+  "implementation": {
+    "entry_point": "model/stormer.py",
+    "scope": "two-stage generative forecasting: a deterministic regression UNet followed by an EDM-preconditioned conditional residual diffusion UNet, both built on the SongUNet backbone"
+  },
+  "architecture": {
+    "family": "two-stage UNet with conditional residual diffusion",
+    "stages": [
+      "regression",
+      "diffusion"
+    ],
+    "backbone": "SongUNet",
+    "image_size": [
+      512,
+      640
+    ],
+    "model_channels": 128,
+    "channel_mult": [
+      1,
+      2,
+      2,
+      2,
+      2
+    ],
+    "num_blocks": 4,
+    "attention_resolutions": [],
+    "state_channels": 99,
+    "background_channels": 26,
+    "invariant_channels": 2,
+    "regression": {
+      "condition": [
+        "state",
+        "background",
+        "invariant"
+      ],
+      "input_channels": 127,
+      "output_channels": 99,
+      "additive_pos_embed": false
+    },
+    "diffusion": {
+      "condition": [
+        "state",
+        "regression",
+        "invariant"
+      ],
+      "input_channels": 299,
+      "output_channels": 99,
+      "additive_pos_embed": true,
+      "preconditioning": "EDM",
+      "sigma_data": 0.5,
+      "P_mean": -1.2,
+      "P_std": 1.2
+    }
+  },
+  "sampling": {
+    "diffusion_steps": 18,
+    "sigma_min": 0.002,
+    "sigma_max": 800.0,
+    "rho": 7.0
+  },
+  "data": {
+    "dataset": "ERA5 backgrounds with local high-resolution state",
+    "background_source": "ERA5",
+    "state_source": "convection-resolving regional reanalysis (HRRR-style)",
+    "regional_grid": {
+      "projection": "Lambert conformal conic",
+      "image_size": [
+        512,
+        640
+      ],
+      "horizontal_spacing_km": 3.0,
+      "standard_latitude_degrees": 38.5,
+      "origin_latitude_degrees": 38.5,
+      "central_longitude_degrees": -97.5
+    },
+    "background_grid": [
+      721,
+      1440
+    ],
+    "background_spatial_resolution_degrees": 0.25,
+    "time_step_hours": 1,
+    "input_steps": 1,
+    "output_steps": 1,
+    "state_surface_variables": [
+      "10m_u_component_of_wind",
+      "10m_v_component_of_wind",
+      "2m_temperature",
+      "mean_sea_level_pressure",
+      "composite_reflectivity"
+    ],
+    "state_model_level_variables": [
+      "u_component_of_wind",
+      "v_component_of_wind",
+      "temperature",
+      "specific_humidity",
+      "geopotential",
+      "pressure"
+    ],
+    "invariant_variables": [
+      "land_sea_mask",
+      "orography"
+    ]
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/stormer.py",
+    "scripts/train.py",
+    "scripts/data_loader.py",
+    "scripts/grid.py"
+  ]
+}

configuration.json

@@ -0,0 +1 @@
+{"framework":"PyTorch","task":"weather-forecasting"}

model/stormer.py

@@ -0,0 +1,406 @@
+from __future__ import annotations
+
+import contextlib
+import sys
+import types
+from typing import Any
+
+import torch
+from torch import nn
+
+
+def _install_nvtx_fallback() -> None:
+    try:
+        __import__("nvtx")
+        return
+    except ModuleNotFoundError:
+        pass
+
+    class _Annotate(contextlib.ContextDecorator):
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            pass
+
+        def __enter__(self) -> _Annotate:
+            return self
+
+        def __exit__(self, *args: Any) -> bool:
+            return False
+
+    module = types.ModuleType("nvtx")
+    module.annotate = _Annotate
+    sys.modules["nvtx"] = module
+
+
+_install_nvtx_fallback()
+
+try:
+    __import__("onescience.models.module")
+except ModuleNotFoundError:
+    from onescience.modules import module as _onescience_module
+
+    sys.modules["onescience.models.module"] = _onescience_module
+
+from onescience.models.diffusion.song_unet import SongUNet  # noqa: E402
+
+
+class _StormCastSongUNet(SongUNet):
+    """OneScience SongUNet with StormCast's optional learned spatial embedding."""
+
+    def __init__(
+        self,
+        *args: Any,
+        additive_pos_embed: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.additive_pos_embed = additive_pos_embed
+        if additive_pos_embed:
+            model_channels = self.enc[next(iter(self.enc))].out_channels
+            self.spatial_emb = nn.Parameter(
+                torch.empty(1, model_channels, self.img_shape_y, self.img_shape_x)
+            )
+            nn.init.trunc_normal_(self.spatial_emb, std=0.02)
+            first_layer = self.enc[next(iter(self.enc))]
+            first_layer.register_forward_hook(self._add_spatial_embedding)
+
+    def _add_spatial_embedding(
+        self,
+        module: nn.Module,
+        inputs: tuple[torch.Tensor, ...],
+        output: torch.Tensor,
+    ) -> torch.Tensor:
+        if output.shape[-2:] != self.spatial_emb.shape[-2:]:
+            raise ValueError(
+                "Input grid must match img_resolution when additive_pos_embed is enabled"
+            )
+        return output + self.spatial_emb.to(dtype=output.dtype)
+
+
+class StormCastRegressionUNet(nn.Module):
+    """Deterministic regression wrapper."""
+
+    def __init__(
+        self,
+        img_resolution: int | list[int] | tuple[int, int],
+        img_in_channels: int,
+        img_out_channels: int,
+        use_fp16: bool = False,
+        sigma_min: float = 0.0,
+        sigma_max: float = float("inf"),
+        sigma_data: float = 0.5,
+        model_type: str = "SongUNet",
+        **model_kwargs: Any,
+    ) -> None:
+        super().__init__()
+        if model_type != "SongUNet":
+            raise ValueError("Regression requires model_type='SongUNet'")
+        self.register_buffer("device_buffer", torch.empty(0))
+        self.img_resolution = img_resolution
+        self.img_in_channels = img_in_channels
+        self.img_out_channels = img_out_channels
+        self.use_fp16 = use_fp16
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.sigma_data = sigma_data
+        self.model = _StormCastSongUNet(
+            img_resolution=img_resolution,
+            in_channels=img_in_channels,
+            out_channels=img_out_channels,
+            **model_kwargs,
+        )
+
+    def forward(self, x: torch.Tensor, force_fp32: bool = False) -> torch.Tensor:
+        _validate_image(x, "x", self.img_in_channels)
+        dtype = _model_dtype(x, self.use_fp16, force_fp32)
+        output = self.model(
+            x.to(dtype),
+            torch.zeros(x.shape[0], dtype=x.dtype, device=x.device),
+            class_labels=None,
+        )
+        _validate_output_dtype(output, dtype)
+        return output.to(torch.float32)
+
+
+class StormCastEDMPrecond(nn.Module):
+    """EDM preconditioner for conditional residual diffusion."""
+
+    def __init__(
+        self,
+        img_resolution: int | list[int] | tuple[int, int],
+        img_channels: int,
+        label_dim: int = 0,
+        use_fp16: bool = False,
+        sigma_min: float = 0.0,
+        sigma_max: float = float("inf"),
+        sigma_data: float = 0.5,
+        model_type: str = "SongUNet",
+        img_in_channels: int | None = None,
+        img_out_channels: int | None = None,
+        **model_kwargs: Any,
+    ) -> None:
+        super().__init__()
+        if model_type != "SongUNet":
+            raise ValueError("Diffusion requires model_type='SongUNet'")
+        if label_dim != 0:
+            raise ValueError("Diffusion does not use class labels")
+        self.register_buffer("device_buffer", torch.empty(0))
+        self.img_resolution = img_resolution
+        self.img_channels = img_channels
+        self.img_in_channels = img_channels if img_in_channels is None else img_in_channels
+        self.img_out_channels = img_channels if img_out_channels is None else img_out_channels
+        self.label_dim = label_dim
+        self.use_fp16 = use_fp16
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.sigma_data = sigma_data
+        self.model = _StormCastSongUNet(
+            img_resolution=img_resolution,
+            in_channels=self.img_in_channels,
+            out_channels=self.img_out_channels,
+            label_dim=label_dim,
+            **model_kwargs,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        sigma: torch.Tensor,
+        condition: torch.Tensor | None = None,
+        force_fp32: bool = False,
+    ) -> torch.Tensor:
+        _validate_image(x, "x", self.img_out_channels)
+        sigma = torch.as_tensor(sigma, device=x.device, dtype=torch.float32).reshape(
+            -1, 1, 1, 1
+        )
+        if sigma.shape[0] not in (1, x.shape[0]):
+            raise ValueError("sigma must contain one value or one value per batch item")
+
+        c_skip = self.sigma_data**2 / (sigma.square() + self.sigma_data**2)
+        c_out = sigma * self.sigma_data / (sigma.square() + self.sigma_data**2).sqrt()
+        c_in = 1 / (self.sigma_data**2 + sigma.square()).sqrt()
+        c_noise = sigma.log() / 4
+        model_input = c_in * x.to(torch.float32)
+        if condition is not None:
+            _validate_condition(condition, x)
+            model_input = torch.cat((model_input, condition.to(torch.float32)), dim=1)
+        if model_input.shape[1] != self.img_in_channels:
+            raise ValueError(
+                f"Diffusion model expects {self.img_in_channels} total channels, "
+                f"got {model_input.shape[1]}"
+            )
+
+        dtype = _model_dtype(x, self.use_fp16, force_fp32)
+        output = self.model(
+            model_input.to(dtype), c_noise.flatten(), class_labels=None
+        )
+        _validate_output_dtype(output, dtype)
+        return c_skip * x.to(torch.float32) + c_out * output.to(torch.float32)
+
+    @staticmethod
+    def round_sigma(sigma: float | list[float] | torch.Tensor) -> torch.Tensor:
+        return torch.as_tensor(sigma)
+
+
+class StormCast(nn.Module):
+    """Compose regression and conditional residual diffusion stages."""
+
+    def __init__(
+        self,
+        regression: StormCastRegressionUNet,
+        diffusion: StormCastEDMPrecond,
+    ) -> None:
+        super().__init__()
+        self.regression = regression
+        self.diffusion = diffusion
+
+    def regression_condition(
+        self,
+        state: torch.Tensor,
+        background: torch.Tensor,
+        invariant: torch.Tensor,
+    ) -> torch.Tensor:
+        invariant = _expand_invariant(invariant, state.shape[0])
+        return torch.cat((state, background, invariant), dim=1)
+
+    def diffusion_condition(
+        self,
+        state: torch.Tensor,
+        regression: torch.Tensor,
+        invariant: torch.Tensor,
+    ) -> torch.Tensor:
+        invariant = _expand_invariant(invariant, state.shape[0])
+        return torch.cat((state, regression, invariant), dim=1)
+
+    def predict_regression(
+        self,
+        state: torch.Tensor,
+        background: torch.Tensor,
+        invariant: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.regression(self.regression_condition(state, background, invariant))
+
+    def denoise_residual(
+        self,
+        noisy_residual: torch.Tensor,
+        sigma: torch.Tensor,
+        state: torch.Tensor,
+        regression: torch.Tensor,
+        invariant: torch.Tensor,
+    ) -> torch.Tensor:
+        condition = self.diffusion_condition(state, regression, invariant)
+        return self.diffusion(noisy_residual, sigma, condition=condition)
+
+
+@torch.no_grad()
+def edm_heun_sample(
+    model: StormCastEDMPrecond,
+    condition: torch.Tensor,
+    output_channels: int,
+    num_steps: int = 18,
+    sigma_min: float = 0.002,
+    sigma_max: float = 800.0,
+    rho: float = 7.0,
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
+    """Sample a residual with the deterministic EDM Heun path."""
+    if num_steps < 1:
+        raise ValueError("num_steps must be at least 1")
+    if not 0 < sigma_min < sigma_max:
+        raise ValueError("Expected 0 < sigma_min < sigma_max")
+    if rho <= 0:
+        raise ValueError("rho must be positive")
+    if condition.ndim != 4:
+        raise ValueError("condition must have shape (B, C, H, W)")
+
+    step_indices = torch.arange(
+        num_steps, device=condition.device, dtype=torch.float64
+    )
+    denominator = max(num_steps - 1, 1)
+    sigma_steps = (
+        sigma_max ** (1 / rho)
+        + step_indices
+        / denominator
+        * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))
+    ).pow(rho)
+    sigma_steps = torch.cat((sigma_steps, sigma_steps.new_zeros(1)))
+    shape = (condition.shape[0], output_channels, *condition.shape[-2:])
+    latent = torch.randn(
+        shape,
+        device=condition.device,
+        dtype=torch.float32,
+        generator=generator,
+    )
+    x_next = latent.to(torch.float64) * sigma_steps[0]
+
+    for index, (sigma_cur, sigma_next) in enumerate(
+        zip(sigma_steps[:-1], sigma_steps[1:])
+    ):
+        x_cur = x_next
+        denoised = model(
+            x_cur.to(torch.float32),
+            sigma_cur.to(torch.float32),
+            condition=condition,
+        ).to(torch.float64)
+        derivative = (x_cur - denoised) / sigma_cur
+        x_next = x_cur + (sigma_next - sigma_cur) * derivative
+
+        if index < num_steps - 1:
+            denoised_next = model(
+                x_next.to(torch.float32),
+                sigma_next.to(torch.float32),
+                condition=condition,
+            ).to(torch.float64)
+            derivative_next = (x_next - denoised_next) / sigma_next
+            x_next = x_cur + (sigma_next - sigma_cur) * (
+                0.5 * derivative + 0.5 * derivative_next
+            )
+    return x_next.to(torch.float32)
+
+
+def build_regression_model(
+    image_size: list[int] | tuple[int, int] = (512, 640),
+    state_channels: int = 99,
+    background_channels: int = 26,
+    invariant_channels: int = 2,
+    model_channels: int = 128,
+    channel_mult: list[int] | tuple[int, ...] = (1, 2, 2, 2, 2),
+    attn_resolutions: list[int] | tuple[int, ...] = (),
+    **kwargs: Any,
+) -> StormCastRegressionUNet:
+    return StormCastRegressionUNet(
+        img_resolution=list(image_size),
+        img_in_channels=state_channels + background_channels + invariant_channels,
+        img_out_channels=state_channels,
+        model_type="SongUNet",
+        model_channels=model_channels,
+        channel_mult=list(channel_mult),
+        attn_resolutions=list(attn_resolutions),
+        embedding_type="zero",
+        additive_pos_embed=False,
+        **kwargs,
+    )
+
+
+def build_diffusion_model(
+    image_size: list[int] | tuple[int, int] = (512, 640),
+    state_channels: int = 99,
+    invariant_channels: int = 2,
+    model_channels: int = 128,
+    channel_mult: list[int] | tuple[int, ...] = (1, 2, 2, 2, 2),
+    attn_resolutions: list[int] | tuple[int, ...] = (),
+    **kwargs: Any,
+) -> StormCastEDMPrecond:
+    condition_channels = state_channels + state_channels + invariant_channels
+    return StormCastEDMPrecond(
+        img_resolution=list(image_size),
+        img_channels=state_channels + condition_channels,
+        img_in_channels=state_channels + condition_channels,
+        img_out_channels=state_channels,
+        model_type="SongUNet",
+        model_channels=model_channels,
+        channel_mult=list(channel_mult),
+        attn_resolutions=list(attn_resolutions),
+        additive_pos_embed=True,
+        **kwargs,
+    )
+
+
+def _model_dtype(
+    x: torch.Tensor, use_fp16: bool, force_fp32: bool
+) -> torch.dtype:
+    return (
+        torch.float16
+        if use_fp16 and not force_fp32 and x.device.type == "cuda"
+        else torch.float32
+    )
+
+
+def _validate_image(x: torch.Tensor, name: str, channels: int) -> None:
+    if x.ndim != 4:
+        raise ValueError(f"{name} must have shape (B, C, H, W), got {tuple(x.shape)}")
+    if x.shape[1] != channels:
+        raise ValueError(f"{name} must have {channels} channels, got {x.shape[1]}")
+
+
+def _validate_condition(condition: torch.Tensor, x: torch.Tensor) -> None:
+    if condition.ndim != 4:
+        raise ValueError("condition must have shape (B, C, H, W)")
+    if condition.shape[0] != x.shape[0] or condition.shape[-2:] != x.shape[-2:]:
+        raise ValueError("condition batch and spatial dimensions must match x")
+
+
+def _validate_output_dtype(output: torch.Tensor, dtype: torch.dtype) -> None:
+    if output.dtype != dtype and not torch.is_autocast_enabled():
+        raise ValueError(f"Expected model output dtype {dtype}, got {output.dtype}")
+
+
+def _expand_invariant(invariant: torch.Tensor, batch_size: int) -> torch.Tensor:
+    if invariant.ndim == 3:
+        invariant = invariant.unsqueeze(0)
+    if invariant.ndim != 4:
+        raise ValueError("invariant must have shape (C, H, W) or (B, C, H, W)")
+    if invariant.shape[0] == 1 and batch_size != 1:
+        invariant = invariant.expand(batch_size, -1, -1, -1)
+    if invariant.shape[0] != batch_size:
+        raise ValueError("invariant batch dimension must be 1 or match state")
+    return invariant

scripts/data_loader.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import h5py
+import torch
+from onescience.datapipes.climate.era5 import ERA5Dataset
+from torch.utils.data import Dataset
+
+from grid import lambert_grid
+
+
+class StormCastDataset(Dataset):
+    """Pair OneScience ERA5 backgrounds with synchronized local state targets."""
+
+    def __init__(
+        self,
+        data_root: str | Path,
+        years: list[int],
+        era5_variables: list[str],
+        state_variables: list[str],
+        invariant_variables: list[str],
+        image_size: list[int] | tuple[int, int],
+        input_steps: int = 1,
+        output_steps: int = 1,
+        normalize: bool = True,
+    ) -> None:
+        if input_steps != 1 or output_steps != 1:
+            raise ValueError("StormCast pairing currently requires one input and one target step")
+
+        self.data_root = Path(data_root)
+        self.years = years
+        self.era5_variables = era5_variables
+        self.state_variables = state_variables
+        self.invariant_variables = invariant_variables
+        self.image_size = tuple(image_size)
+        self.normalize = normalize
+        self.era5 = ERA5Dataset(
+            dataset_dir=str(self.data_root / "era5"),
+            used_years=years,
+            used_variables=era5_variables,
+            input_steps=input_steps,
+            output_steps=output_steps,
+            normalize=normalize,
+        )
+        self.samples_per_year = self.era5.samples_per_year
+        self._validate_era5_grid()
+        self._validate_local_files()
+        self.invariants = self._load_invariants()
+        self._initialize_background_regrid()
+
+    def _validate_era5_grid(self) -> None:
+        if self.era5.H < 2 or self.era5.W < 2:
+            raise ValueError("ERA5 grid must have at least two points per dimension")
+        expected = (721, 1440)
+        if (self.era5.H, self.era5.W) != expected:
+            raise ValueError(
+                f"StormCast expects ERA5 on the global {expected} grid, "
+                f"got {(self.era5.H, self.era5.W)}"
+            )
+
+    def _validate_local_files(self) -> None:
+        for year in self.years:
+            path = self.data_root / "hrrr" / "data" / f"{year}.h5"
+            if not path.is_file():
+                raise FileNotFoundError(f"Missing local state file: {path}")
+            with h5py.File(path, "r") as handle:
+                fields = handle["fields"]
+                variables = [
+                    value.decode() if isinstance(value, bytes) else str(value)
+                    for value in fields.attrs["variables"]
+                ]
+                if variables != self.state_variables:
+                    raise ValueError(
+                        "Local state channel order differs from data.state_variables"
+                    )
+                expected_steps = self.samples_per_year + 1
+                if fields.shape[0] != expected_steps:
+                    raise ValueError(
+                        f"{path} has {fields.shape[0]} steps, expected {expected_steps}"
+                    )
+                if tuple(fields.shape[-2:]) != self.image_size:
+                    raise ValueError(
+                        f"Local state grid is {tuple(fields.shape[-2:])}, "
+                        f"expected regional grid {self.image_size}"
+                    )
+
+    def _load_invariants(self) -> torch.Tensor:
+        path = self.data_root / "hrrr" / "invariants.h5"
+        with h5py.File(path, "r") as handle:
+            fields = handle["fields"]
+            variables = [
+                value.decode() if isinstance(value, bytes) else str(value)
+                for value in fields.attrs["variables"]
+            ]
+            if variables != self.invariant_variables:
+                raise ValueError(
+                    "Invariant channel order differs from data.invariant_variables"
+                )
+            invariants = torch.as_tensor(fields[:], dtype=torch.float32)
+            if tuple(invariants.shape[-2:]) != self.image_size:
+                raise ValueError(
+                    f"Invariant grid is {tuple(invariants.shape[-2:])}, "
+                    f"expected {self.image_size}"
+                )
+            return invariants
+
+    def _initialize_background_regrid(self) -> None:
+        with h5py.File(self.data_root / "hrrr" / "invariants.h5", "r") as handle:
+            if "lat" in handle and "lon" in handle:
+                target_lat = torch.as_tensor(handle["lat"][:], dtype=torch.float32)
+                target_lon = torch.as_tensor(handle["lon"][:], dtype=torch.float32)
+            else:
+                target_lat_np, target_lon_np = lambert_grid(self.image_size)
+                target_lat = torch.from_numpy(target_lat_np)
+                target_lon = torch.from_numpy(target_lon_np)
+        if target_lat.shape != self.image_size or target_lon.shape != self.image_size:
+            raise ValueError("StormCast target latitude/longitude grid has wrong shape")
+
+        lat_position = (90.0 - target_lat) / (180.0 / (self.era5.H - 1))
+        lon_position = torch.remainder(target_lon, 360.0) / (360.0 / self.era5.W)
+        self.lat0 = lat_position.floor().long().clamp(0, self.era5.H - 2)
+        self.lat1 = self.lat0 + 1
+        self.lon0 = lon_position.floor().long().remainder(self.era5.W)
+        self.lon1 = (self.lon0 + 1).remainder(self.era5.W)
+        self.lat_weight = lat_position - self.lat0
+        self.lon_weight = lon_position - lon_position.floor()
+
+    def _regrid_background(self, background: torch.Tensor) -> torch.Tensor:
+        f00 = background[..., self.lat0, self.lon0]
+        f01 = background[..., self.lat0, self.lon1]
+        f10 = background[..., self.lat1, self.lon0]
+        f11 = background[..., self.lat1, self.lon1]
+        lon_weight = self.lon_weight.to(background.dtype)
+        lat_weight = self.lat_weight.to(background.dtype)
+        top = torch.lerp(f00, f01, lon_weight)
+        bottom = torch.lerp(f10, f11, lon_weight)
+        return torch.lerp(top, bottom, lat_weight)
+
+    def __len__(self) -> int:
+        return len(self.era5)
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        background, _, _, step_index, time_index = self.era5[index]
+        background = self._regrid_background(background)
+        year_index = index // self.samples_per_year
+        year = self.years[year_index]
+        path = self.data_root / "hrrr" / "data" / f"{year}.h5"
+
+        with h5py.File(path, "r") as handle:
+            state = torch.as_tensor(
+                handle["fields"][step_index : step_index + 2], dtype=torch.float32
+            )
+            if self.normalize:
+                means = torch.as_tensor(handle["global_means"][:], dtype=torch.float32)
+                stds = torch.as_tensor(handle["global_stds"][:], dtype=torch.float32)
+                state = (state - means) / stds
+
+        return {
+            "background": background,
+            "state": (state[0], state[1]),
+            "invariant": self.invariants,
+            "step_index": step_index,
+            "time_index": time_index,
+        }

scripts/fake_data.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import h5py
+import numpy as np
+import yaml
+
+from grid import lambert_grid
+
+
+def load_config(path: str | Path) -> dict:
+    with Path(path).open("r", encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def write_temporal_fields(
+    path: Path,
+    variables: list[str],
+    num_timesteps: int,
+    image_size: tuple[int, int],
+    time_step_hours: int,
+) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    channels = len(variables)
+    height, width = image_size
+    means = np.zeros((1, channels, 1, 1), dtype=np.float32)
+    stds = np.ones((1, channels, 1, 1), dtype=np.float32)
+    with h5py.File(path, "w") as handle:
+        fields = handle.create_dataset(
+            "fields",
+            shape=(num_timesteps, channels, height, width),
+            dtype=np.float32,
+            chunks=(1, channels, height, width),
+            fillvalue=0.0,
+        )
+        fields.attrs["variables"] = variables
+        fields.attrs["time_step"] = time_step_hours
+        handle.create_dataset("global_means", data=means)
+        handle.create_dataset("global_stds", data=stds)
+
+
+def write_invariants(
+    path: Path,
+    variables: list[str],
+    image_size: tuple[int, int],
+) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    height, width = image_size
+    target_lat, target_lon = lambert_grid(image_size)
+    invariants = np.zeros((len(variables), height, width), dtype=np.float32)
+    with h5py.File(path, "w") as handle:
+        fields = handle.create_dataset(
+            "fields",
+            shape=(len(variables), height, width),
+            dtype=np.float32,
+            chunks=(1, height, width),
+            data=invariants,
+        )
+        fields.attrs["variables"] = variables
+        handle.create_dataset("lat", data=target_lat)
+        handle.create_dataset("lon", data=target_lon)
+
+
+def generate(config: dict) -> None:
+    data = config["data"]
+    root = Path(data["root_dir"])
+    years = sorted(
+        set(data["train_years"] + data["val_years"] + data["test_years"])
+    )
+    era5_image_size = tuple(data["era5_image_size"])
+    image_size = tuple(data["image_size"])
+    if era5_image_size != (721, 1440):
+        raise ValueError("ERA5 grid must be 721 x 1440")
+    if image_size != (512, 640):
+        raise ValueError("Regional grid must be 512 x 640")
+
+    if len(data["era5_variables"]) != 26:
+        raise ValueError("The configured ERA5 input must contain 26 channels")
+    if len(data["state_variables"]) != 99:
+        raise ValueError("The configured local state must contain 99 channels")
+    if data["invariant_variables"] != ["lsm", "orography"]:
+        raise ValueError("Invariant order must be [lsm, orography]")
+
+    for year in years:
+        write_temporal_fields(
+            root / "era5" / "data" / f"{year}.h5",
+            data["era5_variables"],
+            data["num_timesteps"],
+            era5_image_size,
+            data["time_step_hours"],
+        )
+        write_temporal_fields(
+            root / "hrrr" / "data" / f"{year}.h5",
+            data["state_variables"],
+            data["num_timesteps"],
+            image_size,
+            data["time_step_hours"],
+        )
+
+    write_invariants(
+        root / "hrrr" / "invariants.h5",
+        data["invariant_variables"],
+        image_size,
+    )
+    print(f"Generated project validation data under {root}")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate lightweight StormCast data")
+    parser.add_argument("--config", default="conf/config.yaml")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    config_path = Path(args.config).resolve()
+    config = load_config(config_path)
+    project_root = config_path.parent.parent
+    for key in ("root_dir",):
+        path = Path(config["data"][key])
+        if not path.is_absolute():
+            config["data"][key] = str((project_root / path).resolve())
+    generate(config)

scripts/grid.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import numpy as np
+
+
+def lambert_grid(image_size: tuple[int, int]) -> tuple[np.ndarray, np.ndarray]:
+    """Create the project's regional latitude/longitude grid in pure NumPy."""
+    if image_size != (512, 640):
+        raise ValueError("The regional grid must have shape (512, 640)")
+
+    radius = 6371229.0
+    standard_latitude = np.deg2rad(38.5)
+    origin_latitude = np.deg2rad(38.5)
+    central_longitude = np.deg2rad(-97.5)
+    x = -2697520.1425219304 + 3000.0 * np.arange(1799, dtype=np.float64)
+    y = -1587306.1525566636 + 3000.0 * np.arange(1059, dtype=np.float64)
+    x = x[579:1219]
+    y = y[273:785]
+    xx, yy = np.meshgrid(x, y)
+
+    n = np.sin(standard_latitude)
+    f = np.cos(standard_latitude) * np.tan(np.pi / 4 + standard_latitude / 2) ** n / n
+    rho0 = radius * f / np.tan(np.pi / 4 + origin_latitude / 2) ** n
+    rho = np.hypot(xx, rho0 - yy)
+    theta = np.arctan2(xx, rho0 - yy)
+    latitude = 2 * np.arctan((radius * f / rho) ** (1 / n)) - np.pi / 2
+    longitude = central_longitude + theta / n
+    return np.rad2deg(latitude).astype(np.float32), np.mod(np.rad2deg(longitude), 360).astype(np.float32)

scripts/inference.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+from typing import Any
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+MODEL_ROOT = PROJECT_ROOT / "model"
+SCRIPT_ROOT = PROJECT_ROOT / "scripts"
+for path in (PROJECT_ROOT, MODEL_ROOT, SCRIPT_ROOT):
+    if str(path) not in sys.path:
+        sys.path.insert(0, str(path))
+
+import h5py
+import numpy as np
+import torch
+import yaml
+
+from stormer import StormCast, edm_heun_sample
+from data_loader import StormCastDataset
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run StormCast inference")
+    parser.add_argument("--config", type=Path, default=Path("conf/config.yaml"))
+    parser.add_argument("--regression-weights", type=Path)
+    parser.add_argument("--diffusion-weights", type=Path)
+    parser.add_argument("--output", type=Path)
+    parser.add_argument("--num-steps", type=int)
+    parser.add_argument("--diffusion-steps", type=int)
+    parser.add_argument("--seed", type=int)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    config_path = args.config.resolve()
+    with config_path.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    _resolve_paths(config, config_path.parent.parent)
+    inference_config = config["inference"]
+    regression_weights = args.regression_weights or Path(
+        config["model"]["regression_weights"]
+    )
+    diffusion_weights = args.diffusion_weights or Path(
+        config["model"]["diffusion_weights"]
+    )
+    output = args.output or Path(inference_config["output_dir"]) / "forecast.h5"
+    run_inference(
+        config=config,
+        regression_weights=regression_weights,
+        diffusion_weights=diffusion_weights,
+        output=output,
+        num_steps=args.num_steps or inference_config["num_steps"],
+        diffusion_steps=args.diffusion_steps or inference_config["diffusion_steps"],
+        seed=config["project"]["seed"] if args.seed is None else args.seed,
+    )
+
+
+@torch.no_grad()
+def run_inference(
+    config: dict[str, Any],
+    regression_weights: Path,
+    diffusion_weights: Path,
+    output: Path,
+    num_steps: int,
+    diffusion_steps: int,
+    seed: int,
+) -> Path:
+    if not torch.cuda.is_available():
+        raise RuntimeError("StormCast inference requires a CUDA/HIP device")
+    if num_steps < 1:
+        raise ValueError("num_steps must be at least 1")
+    device = torch.device("cuda", 0)
+    torch.cuda.set_device(device)
+    data_config = config["data"]
+    inference_config = config["inference"]
+    split = inference_config["split"]
+    years_key = f"{split}_years"
+    if years_key not in data_config:
+        raise ValueError(f"Unknown inference split: {split}")
+    dataset = StormCastDataset(
+        data_root=data_config["root_dir"],
+        years=data_config[years_key],
+        era5_variables=data_config["era5_variables"],
+        state_variables=data_config["state_variables"],
+        invariant_variables=data_config["invariant_variables"],
+        image_size=data_config["image_size"],
+        input_steps=data_config["input_steps"],
+        output_steps=data_config["output_steps"],
+        normalize=data_config["normalize"],
+    )
+    if num_steps > len(dataset):
+        raise ValueError(f"Requested {num_steps} steps but split contains {len(dataset)}")
+
+    regression = _load_weights(regression_weights, "regression").to(device).eval()
+    diffusion = _load_weights(diffusion_weights, "diffusion").to(device).eval()
+    model = StormCast(regression, diffusion)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    output.parent.mkdir(parents=True, exist_ok=True)
+
+    first = dataset[0]
+    state = first["state"][0].unsqueeze(0).to(device, dtype=torch.float32)
+    invariant = first["invariant"].to(device, dtype=torch.float32)
+    state_channels = len(data_config["state_variables"])
+    background_channels = len(data_config["era5_variables"])
+    height, width = data_config["image_size"]
+    if [height, width] != list(config["model"]["image_size"]):
+        raise ValueError("Data and model image sizes must match")
+    if list(data_config["era5_image_size"]) != [721, 1440]:
+        raise ValueError("ERA5 grid must be 721 x 1440")
+
+    with h5py.File(output, "w") as handle:
+        handle.attrs["normalized"] = bool(data_config["normalize"])
+        handle.attrs["seed"] = seed
+        handle.attrs["diffusion_steps"] = diffusion_steps
+        handle.attrs["sigma_min"] = inference_config["sigma_min"]
+        handle.attrs["sigma_max"] = inference_config["sigma_max"]
+        handle.attrs["rho"] = inference_config["rho"]
+        handle.attrs["regression_weights"] = str(regression_weights.resolve())
+        handle.attrs["diffusion_weights"] = str(diffusion_weights.resolve())
+        handle.attrs["state_variables"] = np.asarray(
+            data_config["state_variables"], dtype="S"
+        )
+        handle.attrs["background_variables"] = np.asarray(
+            data_config["era5_variables"], dtype="S"
+        )
+        prediction_store = handle.create_dataset(
+            "prediction", (num_steps, state_channels, height, width), dtype="f4"
+        )
+        regression_store = handle.create_dataset(
+            "regression", (num_steps, state_channels, height, width), dtype="f4"
+        )
+        target_store = handle.create_dataset(
+            "target", (num_steps, state_channels, height, width), dtype="f4"
+        )
+        background_store = handle.create_dataset(
+            "background", (num_steps, background_channels, height, width), dtype="f4"
+        )
+        time_store = handle.create_dataset("time_index", (num_steps,), dtype="i8")
+
+        for index in range(num_steps):
+            sample = dataset[index]
+            background = sample["background"].unsqueeze(0).to(
+                device, dtype=torch.float32
+            )
+            regression_prediction = model.predict_regression(
+                state, background, invariant
+            )
+            condition = model.diffusion_condition(
+                state, regression_prediction, invariant
+            )
+            residual = edm_heun_sample(
+                diffusion,
+                condition,
+                output_channels=state_channels,
+                num_steps=diffusion_steps,
+                sigma_min=inference_config["sigma_min"],
+                sigma_max=inference_config["sigma_max"],
+                rho=inference_config["rho"],
+                generator=generator,
+            )
+            prediction = regression_prediction + residual
+
+            prediction_store[index] = prediction[0].cpu().numpy()
+            regression_store[index] = regression_prediction[0].cpu().numpy()
+            target_store[index] = sample["state"][1].numpy()
+            background_store[index] = sample["background"].numpy()
+            time_store[index] = int(np.asarray(sample["time_index"]).reshape(-1)[-1])
+            state = prediction
+            print(f"forecast_step={index + 1}/{num_steps}")
+    print(f"output={output}")
+    return output
+
+
+def _load_weights(path: Path, kind: str) -> torch.nn.Module:
+    if not path.is_file():
+        raise FileNotFoundError(
+            f"Missing {kind} weights: {path}. "
+            "Train the corresponding stage first or pass an explicit weight path."
+        )
+    checkpoint = torch.load(path, map_location="cpu", weights_only=False)
+    saved_config = checkpoint.get("config")
+    if saved_config is None:
+        raise ValueError("Project .pt weights must include their training config")
+    model = _build_stage_models_from_config(saved_config, kind)
+    model.load_state_dict(checkpoint["model_state_dict"], strict=True)
+    return model
+
+
+def _build_stage_models_from_config(config: dict[str, Any], kind: str) -> torch.nn.Module:
+    from stormer import build_diffusion_model, build_regression_model
+
+    data_config = config["data"]
+    model_config = config["model"]
+    common = {
+        "image_size": model_config["image_size"],
+        "state_channels": len(data_config["state_variables"]),
+        "invariant_channels": len(data_config["invariant_variables"]),
+        "model_channels": model_config["model_channels"],
+        "channel_mult": model_config["channel_mult"],
+        "num_blocks": model_config["num_blocks"],
+        "attn_resolutions": model_config["attention_resolutions"],
+    }
+    if kind == "regression":
+        return build_regression_model(
+            **common, background_channels=len(data_config["era5_variables"])
+        )
+    return build_diffusion_model(**common)
+
+
+def _resolve_paths(config: dict[str, Any], project_root: Path) -> None:
+    for section, key in (("data", "root_dir"), ("inference", "output_dir")):
+        path = Path(config[section][key])
+        if not path.is_absolute():
+            config[section][key] = str((project_root / path).resolve())
+    for key in ("regression_weights", "diffusion_weights"):
+        path = Path(config["model"][key])
+        if not path.is_absolute():
+            config["model"][key] = str((project_root / path).resolve())
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import argparse
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import h5py
+import matplotlib
+import numpy as np
+import yaml
+
+matplotlib.use("Agg")
+from matplotlib import pyplot as plt
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Visualize StormCast HDF5 forecasts")
+    parser.add_argument("--config", type=Path, default=Path("conf/config.yaml"))
+    parser.add_argument("--input", type=Path, default="./outputs/inference/forecast.h5")
+    parser.add_argument("--output-dir", type=Path)
+    parser.add_argument("--state-variable")
+    parser.add_argument("--background-variable")
+    parser.add_argument("--step", type=int, action="append")
+    parser.add_argument(
+        "--normalized",
+        action="store_true",
+        help="Plot model-space values instead of applying dataset statistics",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    config_path = args.config.resolve()
+    with config_path.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    project_root = config_path.parent.parent
+    data_root = Path(config["data"]["root_dir"])
+    if not data_root.is_absolute():
+        data_root = (project_root / data_root).resolve()
+    output_dir = args.output_dir or Path(config["inference"]["output_dir"]) / "plots"
+    if not output_dir.is_absolute():
+        output_dir = (project_root / output_dir).resolve()
+    visualize(
+        input_path=args.input,
+        output_dir=output_dir,
+        data_root=data_root,
+        state_variable=args.state_variable
+        or config["inference"]["plot_state_variable"],
+        background_variable=args.background_variable
+        or config["inference"]["plot_background_variable"],
+        steps=args.step,
+        denormalize=not args.normalized,
+    )
+
+
+def visualize(
+    input_path: Path,
+    output_dir: Path,
+    data_root: Path,
+    state_variable: str,
+    background_variable: str,
+    steps: list[int] | None = None,
+    denormalize: bool = True,
+) -> list[Path]:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    outputs: list[Path] = []
+    with h5py.File(input_path, "r") as handle:
+        state_variables = _decode_strings(handle.attrs["state_variables"])
+        background_variables = _decode_strings(handle.attrs["background_variables"])
+        state_index = _variable_index(state_variables, state_variable, "state")
+        background_index = _variable_index(
+            background_variables, background_variable, "background"
+        )
+        selected_steps = steps or list(range(handle["prediction"].shape[0]))
+        for step in selected_steps:
+            if not 0 <= step < handle["prediction"].shape[0]:
+                raise IndexError(f"Step {step} is outside the forecast range")
+
+        source_normalized = bool(handle.attrs.get("normalized", False))
+        stats_by_year: dict[
+            int,
+            tuple[
+                tuple[np.ndarray, np.ndarray],
+                tuple[np.ndarray, np.ndarray],
+            ],
+        ] = {}
+
+        for step in selected_steps:
+            prediction = handle["prediction"][step, state_index]
+            target = handle["target"][step, state_index]
+            background = handle["background"][step, background_index]
+            time_index = int(handle["time_index"][step])
+            if denormalize and source_normalized:
+                year = int(str(time_index)[:4])
+                if year not in stats_by_year:
+                    stats_by_year[year] = (
+                        _read_stats(
+                            data_root / "hrrr" / "data" / f"{year}.h5",
+                            state_variables,
+                        ),
+                        _read_stats(
+                            data_root / "era5" / "data" / f"{year}.h5",
+                            background_variables,
+                        ),
+                    )
+                state_stats, background_stats = stats_by_year[year]
+                prediction = _denormalize(prediction, state_stats, state_index)
+                target = _denormalize(target, state_stats, state_index)
+                background = _denormalize(
+                    background, background_stats, background_index
+                )
+            output = output_dir / f"forecast_{step:03d}_{state_variable}.png"
+            _save_four_panel(
+                prediction,
+                target,
+                background,
+                state_variable,
+                background_variable,
+                time_index,
+                output,
+                normalized=source_normalized and not denormalize,
+            )
+            outputs.append(output)
+            print(f"plot={output}")
+    return outputs
+
+
+def _save_four_panel(
+    prediction: np.ndarray,
+    target: np.ndarray,
+    background: np.ndarray,
+    state_variable: str,
+    background_variable: str,
+    time_index: int,
+    output: Path,
+    normalized: bool,
+) -> None:
+    error = prediction - target
+    state_min = float(min(np.nanmin(prediction), np.nanmin(target)))
+    state_max = float(max(np.nanmax(prediction), np.nanmax(target)))
+    error_limit = max(float(np.nanmax(np.abs(error))), np.finfo(np.float32).eps)
+    time_label = datetime.strptime(str(time_index), "%Y%m%d%H").strftime(
+        "%Y-%m-%d %H:00"
+    )
+    units = " (normalized)" if normalized else ""
+    figure, axes = plt.subplots(1, 4, figsize=(19, 4.8), constrained_layout=True)
+    panels = (
+        (
+            prediction,
+            f"StormCast {state_variable}{units}",
+            "viridis",
+            state_min,
+            state_max,
+        ),
+        (target, f"Target {state_variable}{units}", "viridis", state_min, state_max),
+        (background, f"ERA5 {background_variable}{units}", "magma", None, None),
+        (
+            error,
+            f"Error {state_variable}{units}",
+            "RdBu_r",
+            -error_limit,
+            error_limit,
+        ),
+    )
+    for axis, (data, title, cmap, vmin, vmax) in zip(axes, panels):
+        image = axis.imshow(
+            data, origin="lower", cmap=cmap, vmin=vmin, vmax=vmax, aspect="auto"
+        )
+        axis.set_title(title, fontsize=10)
+        axis.set_xticks([])
+        axis.set_yticks([])
+        figure.colorbar(image, ax=axis, fraction=0.046, pad=0.03)
+    figure.suptitle(f"StormCast valid time: {time_label}", fontsize=13)
+    figure.savefig(output, dpi=160)
+    plt.close(figure)
+
+
+def _read_stats(
+    path: Path, expected_variables: list[str]
+) -> tuple[np.ndarray, np.ndarray]:
+    if not path.is_file():
+        raise FileNotFoundError(f"Missing statistics file: {path}")
+    with h5py.File(path, "r") as handle:
+        variables = _decode_strings(handle["fields"].attrs["variables"])
+        if variables != expected_variables:
+            raise ValueError(f"Variable order in {path} differs from inference output")
+        means = np.asarray(handle["global_means"][:], dtype=np.float32).reshape(-1)
+        stds = np.asarray(handle["global_stds"][:], dtype=np.float32).reshape(-1)
+    return means, stds
+
+
+def _denormalize(
+    data: np.ndarray,
+    stats: tuple[np.ndarray, np.ndarray],
+    index: int,
+) -> np.ndarray:
+    means, stds = stats
+    return data * stds[index] + means[index]
+
+
+def _decode_strings(values: Any) -> list[str]:
+    return [
+        value.decode() if isinstance(value, bytes) else str(value) for value in values
+    ]
+
+
+def _variable_index(variables: list[str], name: str, kind: str) -> int:
+    try:
+        return variables.index(name)
+    except ValueError as error:
+        raise ValueError(f"Unknown {kind} variable {name!r}") from error
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,474 @@
+from __future__ import annotations
+
+import argparse
+import os
+import random
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+MODEL_ROOT = PROJECT_ROOT / "model"
+SCRIPT_ROOT = PROJECT_ROOT / "scripts"
+for path in (PROJECT_ROOT, MODEL_ROOT, SCRIPT_ROOT):
+    if str(path) not in sys.path:
+        sys.path.insert(0, str(path))
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+import yaml
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+
+from stormer import build_diffusion_model, build_regression_model
+from data_loader import StormCastDataset
+
+
+@dataclass
+class DistributedContext:
+    device: torch.device
+    rank: int
+    local_rank: int
+    world_size: int
+
+    @property
+    def distributed(self) -> bool:
+        return self.world_size > 1
+
+    @property
+    def is_main(self) -> bool:
+        return self.rank == 0
+
+
+def regression_loss(
+    model: nn.Module,
+    condition: torch.Tensor,
+    target: torch.Tensor,
+) -> torch.Tensor:
+    return F.mse_loss(model(condition), target)
+
+
+def edm_residual_loss(
+    model: nn.Module,
+    residual: torch.Tensor,
+    condition: torch.Tensor,
+    sigma_data: float = 0.5,
+    p_mean: float = -1.2,
+    p_std: float = 1.2,
+) -> torch.Tensor:
+    sigma = torch.exp(
+        torch.randn(residual.shape[0], device=residual.device) * p_std + p_mean
+    )
+    noise = torch.randn_like(residual) * sigma[:, None, None, None]
+    denoised = model(residual + noise, sigma, condition=condition)
+    weight = (sigma.square() + sigma_data**2) / (sigma * sigma_data).square()
+    return (weight[:, None, None, None] * (denoised - residual).square()).mean()
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Train the two-stage StormCast model")
+    parser.add_argument("--config", type=Path, default=Path("conf/config.yaml"))
+    parser.add_argument("--stage", choices=("regression", "diffusion"))
+    parser.add_argument("--resume", type=Path)
+    parser.add_argument("--initial-weights", type=Path)
+    parser.add_argument("--regression-weights", type=Path)
+    parser.add_argument("--max-steps", type=int)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    config_path = args.config.resolve()
+    with config_path.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    _resolve_config_paths(config, config_path.parent.parent)
+
+    training = config["training"]
+    stage = args.stage or training["stage"]
+    resume = args.resume or _optional_path(training.get("resume_checkpoint"))
+    initial_weights = args.initial_weights or _optional_path(training.get("initial_weights"))
+    regression_weights = args.regression_weights or _optional_path(
+        training.get("regression_weights")
+    )
+    max_steps = args.max_steps if args.max_steps is not None else training["max_steps"]
+    if resume is None and initial_weights is None and not training["from_scratch"]:
+        checkpoint_key = f"{stage}_checkpoint"
+        initial_weights = _optional_path(config["model"].get(checkpoint_key))
+        if initial_weights is None:
+            raise ValueError(
+                f"training.from_scratch is false but model.{checkpoint_key} is not set"
+            )
+
+    context = initialize_distributed()
+    _seed_everything(config["project"]["seed"], context.rank)
+    try:
+        train(
+            config=config,
+            stage=stage,
+            context=context,
+            resume=resume,
+            initial_weights=initial_weights,
+            regression_weights=regression_weights,
+            max_steps=max_steps,
+        )
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
+def train(
+    config: dict[str, Any],
+    stage: str,
+    context: DistributedContext,
+    resume: Path | None,
+    initial_weights: Path | None,
+    regression_weights: Path | None,
+    max_steps: int | None,
+) -> None:
+    if stage not in ("regression", "diffusion"):
+        raise ValueError("training.stage must be 'regression' or 'diffusion'")
+    data_config = config["data"]
+    loader_config = config["dataloader"]
+    training_config = config["training"]
+    if list(data_config["image_size"]) != list(config["model"]["image_size"]):
+        raise ValueError("Data and model image sizes must match")
+    if list(data_config["era5_image_size"]) != [721, 1440]:
+        raise ValueError("ERA5 grid must be 721 x 1440")
+
+    dataset = StormCastDataset(
+        data_root=data_config["root_dir"],
+        years=data_config["train_years"],
+        era5_variables=data_config["era5_variables"],
+        state_variables=data_config["state_variables"],
+        invariant_variables=data_config["invariant_variables"],
+        image_size=data_config["image_size"],
+        input_steps=data_config["input_steps"],
+        output_steps=data_config["output_steps"],
+        normalize=data_config["normalize"],
+    )
+    sampler = (
+        DistributedSampler(
+            dataset,
+            num_replicas=context.world_size,
+            rank=context.rank,
+            shuffle=True,
+        )
+        if context.distributed
+        else None
+    )
+    loader = DataLoader(
+        dataset,
+        batch_size=loader_config["batch_size"],
+        shuffle=sampler is None,
+        sampler=sampler,
+        num_workers=loader_config["num_workers"],
+        pin_memory=loader_config["pin_memory"],
+        drop_last=False,
+    )
+
+    regression, model = _build_stage_models(
+        config, stage, regression_weights, context.device
+    )
+    if initial_weights is not None and resume is None:
+        _load_initial_weights(model, initial_weights, stage)
+    model.to(context.device)
+    optimizer = torch.optim.Adam(
+        model.parameters(),
+        lr=training_config["learning_rate"],
+        betas=tuple(training_config["betas"]),
+        weight_decay=training_config["weight_decay"],
+    )
+
+    start_epoch = 0
+    start_batch = 0
+    global_step = 0
+    if resume is not None:
+        start_epoch, start_batch, global_step = load_training_checkpoint(
+            resume, model, optimizer, stage, context.device
+        )
+    if context.distributed:
+        model = DistributedDataParallel(
+            model,
+            device_ids=[context.local_rank],
+            output_device=context.local_rank,
+        )
+
+    checkpoint_dir = Path(training_config["checkpoint_dir"]) / stage
+    if context.is_main:
+        checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        print(
+            f"stage={stage} world_size={context.world_size} "
+            f"parameters={sum(parameter.numel() for parameter in model.parameters())}"
+        )
+
+    stop = False
+    if max_steps is not None and global_step >= max_steps:
+        stop = True
+    for epoch in range(start_epoch, training_config["epochs"]):
+        if stop:
+            break
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+        model.train()
+        for batch_index, batch in enumerate(loader):
+            if epoch == start_epoch and batch_index < start_batch:
+                continue
+            loss = _training_step(
+                stage,
+                model,
+                regression,
+                batch,
+                context.device,
+                training_config,
+            )
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            clip_norm = training_config.get("gradient_clip_norm")
+            if clip_norm is not None:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), clip_norm)
+            optimizer.step()
+            global_step += 1
+
+            reduced_loss = _mean_across_ranks(loss.detach(), context.world_size)
+            if context.is_main and global_step % training_config["log_interval"] == 0:
+                print(
+                    f"stage={stage} epoch={epoch + 1} step={global_step} "
+                    f"loss={reduced_loss.item():.8f}"
+                )
+            if (
+                context.is_main
+                and global_step % training_config["checkpoint_interval"] == 0
+            ):
+                save_training_checkpoint(
+                    checkpoint_dir / "model_bak.pt",
+                    model,
+                    optimizer,
+                    stage,
+                    epoch,
+                    batch_index + 1,
+                    global_step,
+                    config,
+                )
+            if max_steps is not None and global_step >= max_steps:
+                stop = True
+                break
+        if stop:
+            break
+
+    if context.is_main:
+        checkpoint = checkpoint_dir / "model_bak.pt"
+        save_training_checkpoint(
+            checkpoint,
+            model,
+            optimizer,
+            stage,
+            epoch if "epoch" in locals() else start_epoch,
+            batch_index + 1 if "batch_index" in locals() else start_batch,
+            global_step,
+            config,
+        )
+        print(f"checkpoint={checkpoint} steps={global_step}")
+
+
+def _build_stage_models(
+    config: dict[str, Any],
+    stage: str,
+    regression_weights: Path | None,
+    device: torch.device,
+) -> tuple[nn.Module | None, nn.Module]:
+    data_config = config["data"]
+    model_config = config["model"]
+    common = {
+        "image_size": model_config["image_size"],
+        "state_channels": len(data_config["state_variables"]),
+        "invariant_channels": len(data_config["invariant_variables"]),
+        "model_channels": model_config["model_channels"],
+        "channel_mult": model_config["channel_mult"],
+        "num_blocks": model_config["num_blocks"],
+        "attn_resolutions": model_config["attention_resolutions"],
+    }
+    if stage == "regression":
+        model = build_regression_model(
+            **common,
+            background_channels=len(data_config["era5_variables"]),
+        )
+        return None, model
+
+    if regression_weights is None:
+        raise ValueError("Diffusion training requires --regression-weights")
+    regression = _load_model_weights(
+        build_regression_model(
+            **common,
+            background_channels=len(data_config["era5_variables"]),
+        ),
+        regression_weights,
+        "regression",
+    ).to(device)
+    regression.eval()
+    regression.requires_grad_(False)
+    return regression, build_diffusion_model(**common)
+
+
+def _training_step(
+    stage: str,
+    model: nn.Module,
+    regression: nn.Module | None,
+    batch: dict[str, Any],
+    device: torch.device,
+    training_config: dict[str, Any],
+) -> torch.Tensor:
+    background = batch["background"].to(device, dtype=torch.float32)
+    state, target = (
+        tensor.to(device, dtype=torch.float32) for tensor in batch["state"]
+    )
+    invariant = batch["invariant"].to(device, dtype=torch.float32)
+    if invariant.ndim == 3:
+        invariant = invariant.unsqueeze(0)
+    if invariant.shape[0] == 1 and state.shape[0] > 1:
+        invariant = invariant.expand(state.shape[0], -1, -1, -1)
+
+    if stage == "regression":
+        condition = torch.cat((state, background, invariant), dim=1)
+        return regression_loss(model, condition, target)
+
+    if regression is None:
+        raise RuntimeError("Regression model is required for diffusion training")
+    with torch.no_grad():
+        regression_condition = torch.cat((state, background, invariant), dim=1)
+        regression_prediction = regression(regression_condition)
+        residual = target - regression_prediction
+        condition = torch.cat((state, regression_prediction, invariant), dim=1)
+    return edm_residual_loss(
+        model,
+        residual,
+        condition,
+        sigma_data=training_config["sigma_data"],
+        p_mean=training_config["P_mean"],
+        p_std=training_config["P_std"],
+    )
+
+
+def save_training_checkpoint(
+    path: Path,
+    model: nn.Module,
+    optimizer: torch.optim.Optimizer,
+    stage: str,
+    epoch: int,
+    batch_in_epoch: int,
+    global_step: int,
+    config: dict[str, Any],
+) -> None:
+    model = model.module if isinstance(model, DistributedDataParallel) else model
+    temporary = path.with_suffix(path.suffix + ".tmp")
+    torch.save(
+        {
+            "stage": stage,
+            "epoch": epoch,
+            "batch_in_epoch": batch_in_epoch,
+            "global_step": global_step,
+            "model_state_dict": model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "config": config,
+        },
+        temporary,
+    )
+    temporary.replace(path)
+
+
+def load_training_checkpoint(
+    path: Path,
+    model: nn.Module,
+    optimizer: torch.optim.Optimizer,
+    stage: str,
+    device: torch.device,
+) -> tuple[int, int, int]:
+    checkpoint = torch.load(path, map_location=device, weights_only=False)
+    if checkpoint["stage"] != stage:
+        raise ValueError(
+            f"Checkpoint stage is {checkpoint['stage']}, requested stage is {stage}"
+        )
+    model.load_state_dict(checkpoint["model_state_dict"], strict=True)
+    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+    return (
+        int(checkpoint["epoch"]),
+        int(checkpoint.get("batch_in_epoch", 0)),
+        int(checkpoint["global_step"]),
+    )
+
+
+def initialize_distributed() -> DistributedContext:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    if not torch.cuda.is_available():
+        raise RuntimeError("StormCast training requires a CUDA/HIP device")
+    torch.cuda.set_device(local_rank)
+    if world_size > 1:
+        dist.init_process_group(backend="nccl", init_method="env://")
+        rank = dist.get_rank()
+    else:
+        rank = 0
+    return DistributedContext(
+        device=torch.device("cuda", local_rank),
+        rank=rank,
+        local_rank=local_rank,
+        world_size=world_size,
+    )
+
+
+def _load_initial_weights(model: nn.Module, path: Path, stage: str) -> None:
+    loaded = _load_model_weights(model, path, stage)
+    model.load_state_dict(loaded.state_dict(), strict=True)
+
+
+def _load_model_weights(model: nn.Module, path: Path, stage: str) -> nn.Module:
+    checkpoint = torch.load(path, map_location="cpu", weights_only=False)
+    state = checkpoint.get("model_state_dict", checkpoint)
+    model.load_state_dict(state, strict=True)
+    return model
+
+
+def _mean_across_ranks(value: torch.Tensor, world_size: int) -> torch.Tensor:
+    if world_size > 1:
+        dist.all_reduce(value, op=dist.ReduceOp.SUM)
+        value /= world_size
+    return value
+
+
+def _resolve_config_paths(config: dict[str, Any], project_root: Path) -> None:
+    for section, key in (
+        ("data", "root_dir"),
+        ("training", "checkpoint_dir"),
+    ):
+        path = Path(config[section][key])
+        if not path.is_absolute():
+            config[section][key] = str((project_root / path).resolve())
+    for key in ("regression_weights", "diffusion_weights"):
+        value = config["model"].get(key)
+        if value:
+            path = Path(value)
+            if not path.is_absolute():
+                config["model"][key] = str((project_root / path).resolve())
+    for key in ("initial_weights", "resume_checkpoint", "regression_weights"):
+        value = config["training"].get(key)
+        if value:
+            path = Path(value)
+            if not path.is_absolute():
+                config["training"][key] = str((project_root / path).resolve())
+
+
+def _optional_path(value: str | Path | None) -> Path | None:
+    return None if value is None else Path(value)
+
+
+def _seed_everything(seed: int, rank: int) -> None:
+    seed += rank
+    random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+
+
+if __name__ == "__main__":
+    main()