Upload folder using huggingface_hub

README.md

@@ -0,0 +1,138 @@
+---
+frameworks: PyTorch
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Precipitation Nowcasting
+- Weather Forecast
+- MRMS
+tasks: []
+datasets:
+- OneScience/MRMS
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">NowcastNet_Earth</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+NowcastNet is a large model for extreme-precipitation nowcasting proposed by a team from Tsinghua University. The research was published in the main edition of *Nature*.
+
+Paper: Skilful nowcasting of extreme precipitation with NowcastNet
+
+https://www.nature.com/articles/s41586-023-06184-4
+
+# Model Description
+
+NowcastNet combines data-driven deep learning with numerical methods based on physical equations in a unified framework. Two core networks work together to model precipitation processes at different spatial scales.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Short-term precipitation nowcasting training | Train NowcastNet with MRMS data. |
+| Local quick validation | Use synthetic data to check data loading, model training and inference, and visualization of inference results. |
+| ModelScope / OneCode execution | Download the standalone model package, install dependencies, and run the scripts directly. |
+| Multi-GPU training | Use `torchrun` for data-parallel training across multiple GPUs or accelerators on one host. |
+
+# 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**
+
+- A GPU or DCU is recommended.
+- CPU can be used for import and small-scale connectivity verification; full training and inference will be slow.
+- 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/NowcastNet --local-dir ./NowcastNet
+cd NowcastNet
+```
+
+### 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 check the data protocol and program flow; it does not represent real MRMS data or forecast quality:
+
+```bash
+python scripts/fake_data.py
+```
+
+### Training
+
+Single GPU:
+
+```bash
+python scripts/train.py
+```
+
+Multi-GPU:
+
+```bash
+torchrun --nproc_per_node=8 scripts/train.py
+```
+
+Training weights are saved to `data/checkpoints/` by default.
+
+### Training Weights
+
+This repository provides weights trained on MRMS data in the `weight/` folder. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+Inference reads the training weights from `data/checkpoints/` by default:
+
+```bash
+python scripts/inference.py
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+# Official OneScience Resources
+
+| 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 and License
+
+- This repository is a reproduction of the original NowcastNet paper.

conf/config.yaml

@@ -0,0 +1,41 @@
+model:
+  input_length: 9
+  total_length: 29
+  img_height: 512
+  img_width: 512
+  img_ch: 2
+  ngf: 32
+  device: auto
+  pretrained_model: ""
+  checkpoint_dir: "./data/checkpoint"
+  checkpoint_prefix: "model_bak"
+
+data:
+  data_dir: "./data/data"
+  split: train
+  image_height: 512
+  image_width: 512
+  total_length: 29
+  input_data_type: float32
+  scale_divisor: 10.0
+  offset: 3.0
+  clip_max: 128.0
+  batch_size: 1
+  num_workers: 0
+
+training:
+  epochs: 10
+  lr: 0.0003
+  weight_decay: 0.0
+  seed: 42
+  mode: generator
+  output_dir: "./results/training"
+  loss_history_file: "loss_history.json"
+
+inference:
+  output_dir: "./results/inference"
+  num_samples: 1
+  threshold: 16.0
+
+visualization:
+  output_dir: "./results/visualization"

config.json

@@ -0,0 +1,40 @@
+{
+  "model_name": "NowcastNet",
+  "model_type": "nowcastnet",
+  "architectures": [
+    "Net"
+  ],
+  "framework": "PyTorch",
+  "domain": "atmosphere",
+  "task": "extreme-precipitation-nowcasting",
+  "implementation": {
+    "entry_point": "model/nowcastnet.py",
+    "scope": "evolution and generative networks for radar precipitation nowcasting"
+  },
+  "architecture": {
+    "family": "physics-guided evolution network plus conditional generative network",
+    "input_length": 9,
+    "total_length": 29,
+    "prediction_length": 20,
+    "image_shape": [
+      512,
+      512
+    ],
+    "input_channels": 1,
+    "base_channels": 32,
+    "noise_channels": 32
+  },
+  "data": {
+    "dataset": "MRMS",
+    "protocol": "synthetic_radar_precipitation",
+    "scale_divisor": 10.0,
+    "offset": 3.0,
+    "clip_max": 128.0
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/nowcastnet.py",
+    "model/evolution_network.py",
+    "model/generative_network.py"
+  ]
+}

configuration.json

@@ -0,0 +1 @@
+{"framework":"Jax","task":"other"}

model/evolution_module.py

@@ -0,0 +1,91 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import spectral_norm
+
+class DoubleConv(nn.Module):
+
+    def   __init__(self, in_channels, out_channels, kernel=3, mid_channels=None):
+        super().__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+
+        self.double_conv = nn.Sequential(
+            nn.BatchNorm2d(in_channels),
+            nn.ReLU(inplace=True),
+            spectral_norm(nn.Conv2d(in_channels, mid_channels, kernel_size=kernel, padding=kernel//2)),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            spectral_norm(nn.Conv2d(mid_channels, out_channels, kernel_size=kernel, padding=kernel//2)),
+        )
+        self.single_conv = nn.Sequential(
+            nn.BatchNorm2d(in_channels),
+            spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=kernel, padding=kernel // 2))
+        )
+
+    def forward(self, x):
+        shortcut = self.single_conv(x)
+        x = self.double_conv(x)
+        x = x + shortcut
+        return x
+
+class Down(nn.Module):
+
+    def __init__(self, in_channels, out_channels, kernel=3):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            DoubleConv(in_channels, out_channels, kernel)
+        )
+
+    def forward(self, x):
+        x = self.maxpool_conv(x)
+        return x
+
+class Up(nn.Module):
+
+    def __init__(self, in_channels, out_channels, bilinear=True, kernel=3):
+        super().__init__()
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, kernel=kernel, mid_channels=in_channels // 2)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels, out_channels, kernel)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        # input is CHW
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+
+        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
+                        diffY // 2, diffY - diffY // 2])
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+
+class Up_S(nn.Module):
+
+    def __init__(self, in_channels, out_channels, bilinear=True, kernel=3):
+        super().__init__()
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, kernel=kernel, mid_channels=in_channels)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels, in_channels, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels, out_channels, kernel)
+
+    def forward(self, x):
+        x = self.up(x)
+        return self.conv(x)
+
+
+class OutConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(OutConv, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        return self.conv(x)
+

model/evolution_network.py

@@ -0,0 +1,48 @@
+import torch.nn.functional as F
+from .evolution_module import *
+
+class Evolution_Network(nn.Module):
+    def __init__(self, n_channels, n_classes, base_c=64, bilinear=True):
+        super(Evolution_Network, self).__init__()
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+        self.bilinear = bilinear
+        base_c = base_c
+        self.inc = DoubleConv(n_channels, base_c)
+        self.down1 = Down(base_c * 1, base_c * 2)
+        self.down2 = Down(base_c * 2, base_c * 4)
+        self.down3 = Down(base_c * 4, base_c * 8)
+        factor = 2 if bilinear else 1
+        self.down4 = Down(base_c * 8, base_c * 16 // factor)
+
+        self.up1 = Up(base_c * 16, base_c * 8 // factor, bilinear)
+        self.up2 = Up(base_c * 8, base_c * 4 // factor, bilinear)
+        self.up3 = Up(base_c * 4, base_c * 2 // factor, bilinear)
+        self.up4 = Up(base_c * 2, base_c * 1, bilinear)
+        self.outc = OutConv(base_c * 1, n_classes)
+        self.gamma = nn.Parameter(torch.zeros(1, n_classes, 1, 1), requires_grad=True)
+
+        self.up1_v = Up(base_c * 16, base_c * 8 // factor, bilinear)
+        self.up2_v = Up(base_c * 8, base_c * 4 // factor, bilinear)
+        self.up3_v = Up(base_c * 4, base_c * 2 // factor, bilinear)
+        self.up4_v = Up(base_c * 2, base_c * 1, bilinear)
+        self.outc_v = OutConv(base_c * 1, n_classes * 2)
+
+    def forward(self, x):
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x = self.up1(x5, x4)
+        x = self.up2(x, x3)
+        x = self.up3(x, x2)
+        x = self.up4(x, x1)
+        x = self.outc(x) * self.gamma
+
+        v = self.up1_v(x5, x4)
+        v = self.up2_v(v, x3)
+        v = self.up3_v(v, x2)
+        v = self.up4_v(v, x1)
+        v = self.outc_v(v)
+        return x, v

model/generation_module.py

@@ -0,0 +1,94 @@
+# SPADE Module and Block are adapted from Nvidia SPADE project (https://github.com/NVlabs/SPADE).
+
+import re
+import sys
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.utils.spectral_norm as spectral_norm
+
+class GenBlock(nn.Module):
+    def __init__(self, fin, fout, opt, use_se=False, dilation=1, double_conv=False):
+        super().__init__()
+        self.learned_shortcut = (fin != fout)
+        fmiddle = min(fin, fout)
+        self.opt = opt
+        self.double_conv = double_conv
+
+        self.pad = nn.ReflectionPad2d(dilation)
+        self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=0, dilation=dilation)
+        self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=0, dilation=dilation)
+
+        if self.learned_shortcut:
+            self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False)
+
+        self.conv_0 = spectral_norm(self.conv_0)
+        self.conv_1 = spectral_norm(self.conv_1)
+        if self.learned_shortcut:
+            self.conv_s = spectral_norm(self.conv_s)
+
+        ic = opt.evo_ic
+        
+        self.norm_0 = SPADE(fin, ic)
+        self.norm_1 = SPADE(fmiddle, ic)
+        if self.learned_shortcut:
+            self.norm_s = SPADE(fin, ic)
+
+    def forward(self, x, evo):
+        x_s = self.shortcut(x, evo)
+        dx = self.conv_0(self.pad(self.actvn(self.norm_0(x, evo))))
+        if self.double_conv:
+            dx = self.conv_1(self.pad(self.actvn(self.norm_1(dx, evo))))
+
+        out = x_s + dx
+
+        return out
+
+    def shortcut(self, x, evo):
+        if self.learned_shortcut:
+            x_s = self.conv_s(self.norm_s(x, evo))
+        else:
+            x_s = x
+        return x_s
+
+    def actvn(self, x):
+        return F.leaky_relu(x, 2e-1)
+
+
+class SPADE(nn.Module):
+    def __init__(self, norm_nc, label_nc):
+        super().__init__()
+
+        ks = 3
+
+        self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)
+        nhidden = 64
+        ks = 3
+        pw = ks // 2
+        self.mlp_shared = nn.Sequential(
+            nn.ReflectionPad2d(pw),
+            nn.Conv2d(label_nc, nhidden, kernel_size=ks, padding=0),
+            nn.ReLU()
+        )
+        self.pad = nn.ReflectionPad2d(pw)
+        self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=0)
+        self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=0)
+
+    def forward(self, x, evo):
+
+        normalized = self.param_free_norm(x)
+        evo = F.adaptive_avg_pool2d(evo, output_size=x.size()[2:])
+
+        actv = self.mlp_shared(evo)
+
+        gamma = self.mlp_gamma(self.pad(actv))
+        beta = self.mlp_beta(self.pad(actv))
+
+        out = normalized * (1 + gamma) + beta
+
+        return out
+
+
+
+

model/generative_network.py

@@ -0,0 +1,61 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Function
+from .generation_module import GenBlock
+from .evolution_module import *
+
+class Generative_Encoder(nn.Module):
+    def __init__(self, n_channels, base_c=64):
+        super(Generative_Encoder, self).__init__()
+        base_c = base_c
+        self.inc = DoubleConv(n_channels, base_c, kernel=3)
+        self.down1 = Down(base_c * 1, base_c * 2, 3)
+        self.down2 = Down(base_c * 2, base_c * 4, 3)
+        self.down3 = Down(base_c * 4, base_c * 8, 3)
+
+    def forward(self, x):
+        x = self.inc(x)
+        x = self.down1(x)
+        x = self.down2(x)
+        x = self.down3(x)
+        return x
+
+class Generative_Decoder(nn.Module):
+    def __init__(self, opt):
+        super().__init__()
+        self.opt = opt
+        nf = opt.ngf
+
+        ic = opt.ic_feature
+        self.fc = nn.Conv2d(ic, 8 * nf, 3, padding=1)
+
+        self.head_0 = GenBlock(8 * nf, 8 * nf, opt)
+
+        self.G_middle_0 = GenBlock(8 * nf, 4 * nf, opt, double_conv=True)
+        self.G_middle_1 = GenBlock(4 * nf, 4 * nf, opt, double_conv=True)
+
+        self.up_0 = GenBlock(4 * nf, 2 * nf, opt)
+
+        self.up_1 = GenBlock(2 * nf, 1 * nf, opt, double_conv=True)
+        self.up_2 = GenBlock(1 * nf, 1 * nf, opt, double_conv=True)
+
+        final_nc = nf * 1
+
+        self.conv_img = nn.Conv2d(final_nc, self.opt.gen_oc, 3, padding=1)
+        self.up = nn.Upsample(scale_factor=2)
+    
+    def forward(self, x, evo):
+        x = self.fc(x)
+        x = self.head_0(x, evo)
+        x = self.up(x)
+        x = self.G_middle_0(x, evo)
+        x = self.G_middle_1(x, evo)
+        x = self.up(x)
+        x = self.up_0(x, evo)
+        x = self.up(x)
+        x = self.up_1(x, evo)
+        x = self.up_2(x, evo)
+        x = self.conv_img(F.leaky_relu(x, 2e-1))
+        return x

model/layers_utils.py

@@ -0,0 +1,85 @@
+import torch.nn.functional as F
+import torch
+import torch.nn as nn
+
+
+def make_grid(input):
+    B, C, H, W = input.size()
+    # mesh grid
+    device = input.device
+    xx = torch.arange(0, W, device=device).view(1, -1).repeat(H, 1)
+    yy = torch.arange(0, H, device=device).view(-1, 1).repeat(1, W)
+    xx = xx.view(1, 1, H, W).repeat(B, 1, 1, 1)
+    yy = yy.view(1, 1, H, W).repeat(B, 1, 1, 1)
+    grid = torch.cat((xx, yy), 1).float()
+
+    return grid
+
+def warp(input, flow, grid, mode="bilinear", padding_mode="zeros"):
+
+    B, C, H, W = input.size()
+    vgrid = grid + flow
+
+    vgrid[:, 0, :, :] = 2.0 * vgrid[:, 0, :, :].clone() / max(W - 1, 1) - 1.0
+    vgrid[:, 1, :, :] = 2.0 * vgrid[:, 1, :, :].clone() / max(H - 1, 1) - 1.0
+    vgrid = vgrid.permute(0, 2, 3, 1)
+    output = torch.nn.functional.grid_sample(input, vgrid, padding_mode=padding_mode, mode=mode, align_corners=True)
+    return output
+
+def l2normalize(v, eps=1e-12):
+    return v / (v.norm() + eps)
+
+
+class spectral_norm(nn.Module):
+    def __init__(self, module, name='weight', power_iterations=1):
+        super(spectral_norm, self).__init__()
+        self.module = module
+        self.name = name
+        self.power_iterations = power_iterations
+        if not self._made_params():
+            self._make_params()
+
+    def _update_u_v(self):
+        u = getattr(self.module, self.name + "_u")
+        v = getattr(self.module, self.name + "_v")
+        w = getattr(self.module, self.name + "_bar")
+
+        height = w.data.shape[0]
+        for _ in range(self.power_iterations):
+            v.data = l2normalize(torch.mv(torch.t(w.view(height,-1).data), u.data))
+            u.data = l2normalize(torch.mv(w.view(height,-1).data, v.data))
+
+        sigma = u.dot(w.view(height, -1).mv(v))
+        setattr(self.module, self.name, w / sigma.expand_as(w))
+
+    def _made_params(self):
+        try:
+            u = getattr(self.module, self.name + "_u")
+            v = getattr(self.module, self.name + "_v")
+            w = getattr(self.module, self.name + "_bar")
+            return True
+        except AttributeError:
+            return False
+
+    def _make_params(self):
+        w = getattr(self.module, self.name)
+
+        height = w.data.shape[0]
+        width = w.view(height, -1).data.shape[1]
+
+        u = nn.Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)
+        v = nn.Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)
+        u.data = l2normalize(u.data)
+        v.data = l2normalize(v.data)
+        w_bar = nn.Parameter(w.data)
+
+        del self.module._parameters[self.name]
+
+        self.module.register_parameter(self.name + "_u", u)
+        self.module.register_parameter(self.name + "_v", v)
+        self.module.register_parameter(self.name + "_bar", w_bar)
+
+
+    def forward(self, *args):
+        self._update_u_v()
+        return self.module.forward(*args)

model/model_factory.py

@@ -0,0 +1,32 @@
+"""Factory and checkpoint helpers for the project-local NowcastNet model."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Mapping
+
+import torch
+
+from .nowcastnet import Net
+
+
+def build_model(config, device: torch.device | str = "cpu") -> Net:
+    """Build a model from a mapping or namespace without requiring OneScience."""
+    if isinstance(config, Mapping):
+        config = SimpleNamespace(**config)
+    config.device = torch.device(device)
+    config.evo_ic = config.total_length - config.input_length
+    config.gen_oc = config.total_length - config.input_length
+    config.ic_feature = config.ngf * 10
+    return Net(config).to(config.device)
+
+
+def load_checkpoint(model: torch.nn.Module, path: str, device: torch.device | str = "cpu") -> Mapping:
+    try:
+        state = torch.load(path, map_location=device, weights_only=True)
+    except TypeError:  # torch < 2.0
+        state = torch.load(path, map_location=device)
+    checkpoint = state if isinstance(state, Mapping) else {}
+    model_state = checkpoint["state_dict"] if "state_dict" in checkpoint else state
+    model.load_state_dict(model_state, strict=True)
+    return checkpoint

model/noise_projector.py

@@ -0,0 +1,39 @@
+import torch
+import torch.nn as nn
+from .layers_utils import spectral_norm
+
+class Noise_Projector(nn.Module):
+    def __init__(self, input_length, configs):
+        super(Noise_Projector, self).__init__()
+        self.input_length = input_length
+        self.conv_first = spectral_norm(nn.Conv2d(self.input_length, self.input_length * 2, kernel_size=3, padding=1))
+        self.L1 = ProjBlock(self.input_length * 2, self.input_length * 4)
+        self.L2 = ProjBlock(self.input_length * 4, self.input_length * 8)
+        self.L3 = ProjBlock(self.input_length * 8, self.input_length * 16)
+        self.L4 = ProjBlock(self.input_length * 16, self.input_length * 32)
+
+    def forward(self, x):
+        x = self.conv_first(x)
+        x = self.L1(x)
+        x = self.L2(x)
+        x = self.L3(x)
+        x = self.L4(x)
+
+        return x
+
+
+class ProjBlock(nn.Module):
+    def __init__(self, in_channel, out_channel):
+        super(ProjBlock, self).__init__()
+        self.one_conv = spectral_norm(nn.Conv2d(in_channel, out_channel-in_channel, kernel_size=1, padding=0))
+        self.double_conv = nn.Sequential(
+            spectral_norm(nn.Conv2d(in_channel, out_channel, kernel_size=3, padding=1)),
+            nn.ReLU(),
+            spectral_norm(nn.Conv2d(out_channel, out_channel, kernel_size=3, padding=1))
+        )
+
+    def forward(self, x):
+        x1 = torch.cat([x, self.one_conv(x)], dim=1)
+        x2 = self.double_conv(x)
+        output = x1 + x2
+        return output

model/nowcastnet.py

@@ -0,0 +1,74 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from .layers_utils import warp, make_grid
+from .generative_network import Generative_Encoder, Generative_Decoder
+from .evolution_network import Evolution_Network
+from .noise_projector import Noise_Projector
+
+class Net(nn.Module):
+    def __init__(self, configs):
+        super(Net, self).__init__()
+        self.configs = configs
+        self.pred_length = self.configs.total_length - self.configs.input_length
+
+        self.evo_net = Evolution_Network(self.configs.input_length, self.pred_length, base_c=32)
+        self.gen_enc = Generative_Encoder(self.configs.total_length, base_c=self.configs.ngf)
+        self.gen_dec = Generative_Decoder(self.configs)
+        self.proj = Noise_Projector(self.configs.ngf, configs)
+
+        sample_tensor = torch.zeros(1, 1, self.configs.img_height, self.configs.img_width)
+        self.register_buffer("grid", make_grid(sample_tensor), persistent=False)
+
+    def forward(self, all_frames):
+        if all_frames.ndim == 4:
+            all_frames = all_frames.unsqueeze(-1)
+        if all_frames.ndim != 5 or all_frames.shape[-1] < 1:
+            raise ValueError("Expected frames with shape [B,T,H,W,C], C>=1")
+        all_frames = all_frames[:, :, :, :, :1]
+
+        frames = all_frames.permute(0, 1, 4, 2, 3)
+        batch = frames.shape[0]
+        height = frames.shape[3]
+        width = frames.shape[4]
+
+        # Input Frames
+        input_frames = frames[:, :self.configs.input_length]
+        input_frames = input_frames.reshape(batch, self.configs.input_length, height, width)
+
+        # Evolution Network
+        intensity, motion = self.evo_net(input_frames)
+        motion_ = motion.reshape(batch, self.pred_length, 2, height, width)
+        intensity_ = intensity.reshape(batch, self.pred_length, 1, height, width)
+        series = []
+        last_frames = all_frames[:, (self.configs.input_length - 1):self.configs.input_length, :, :, 0]
+        grid = self.grid.to(frames.device).repeat(batch, 1, 1, 1)
+        for i in range(self.pred_length):
+            last_frames = warp(last_frames, motion_[:, i], grid, mode="nearest", padding_mode="border")
+            last_frames = last_frames + intensity_[:, i]
+            series.append(last_frames)
+        evo_result = torch.cat(series, dim=1)
+
+        evo_result = evo_result/128
+        
+        # Generative Network
+        evo_feature = self.gen_enc(torch.cat([input_frames, evo_result], dim=1))
+
+        noise = torch.randn(batch, self.configs.ngf, max(1, height // 32), max(1, width // 32), device=frames.device)
+        projected = self.proj(noise)
+        # The official reshape/permute is a 4x pixel shuffle: 32*ngf
+        # projected channels become 2*ngf channels at H/8 x W/8. Expressing
+        # it directly preserves the checkpoint layout and also works at the
+        # official 512x512 resolution (where the released literal reshape is
+        # dimensionally inconsistent).
+        if projected.shape[2] * 4 == height // 8 and projected.shape[3] * 4 == width // 8:
+            noise_feature = F.pixel_shuffle(projected, 4)
+        else:
+            target_hw = (max(1, height // 8), max(1, width // 8))
+            noise_feature = F.interpolate(projected[:, : 2 * self.configs.ngf], size=target_hw, mode="nearest")
+
+        feature = torch.cat([evo_feature, noise_feature], dim=1)
+        gen_result = self.gen_dec(feature, evo_result)
+
+        return gen_result.unsqueeze(-1)

scripts/data_loader.py

@@ -0,0 +1,101 @@
+"""MRMS data loading utilities matching the official NowcastNet contract."""
+
+from __future__ import annotations
+
+import struct
+import zlib
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+
+def _read_png_gray16(path: Path) -> np.ndarray:
+    """Read the 16-bit grayscale PNG subset used by MRMS."""
+    raw = path.read_bytes()
+    if raw[:8] != b"\x89PNG\r\n\x1a\n":
+        raise ValueError(f"Not a PNG file: {path}")
+    pos = 8
+    idat: list[bytes] = []
+    width = height = bit_depth = color_type = None
+    while pos < len(raw):
+        length = struct.unpack(">I", raw[pos:pos + 4])[0]
+        kind = raw[pos + 4:pos + 8]
+        payload = raw[pos + 8:pos + 8 + length]
+        pos += length + 12
+        if kind == b"IHDR":
+            width, height, bit_depth, color_type = struct.unpack(">IIBB", payload[:10])
+        elif kind == b"IDAT":
+            idat.append(payload)
+        elif kind == b"IEND":
+            break
+    if width is None or height is None or bit_depth != 16 or color_type != 0:
+        raise ValueError(f"Expected 16-bit grayscale PNG: {path}")
+
+    decoded = zlib.decompress(b"".join(idat))
+    row_bytes = width * 2
+    previous = np.zeros(row_bytes, dtype=np.uint8)
+    rows: list[np.ndarray] = []
+    cursor = 0
+    for _ in range(height):
+        filter_type = decoded[cursor]
+        row = np.frombuffer(decoded[cursor + 1:cursor + 1 + row_bytes], dtype=np.uint8).copy()
+        cursor += row_bytes + 1
+        for index in range(row_bytes):
+            left = row[index - 2] if index >= 2 else 0
+            up = previous[index]
+            upper_left = previous[index - 2] if index >= 2 else 0
+            if filter_type == 1:
+                value = left
+            elif filter_type == 2:
+                value = up
+            elif filter_type == 3:
+                value = (int(left) + int(up)) // 2
+            elif filter_type == 4:
+                predictor = int(left) + int(up) - int(upper_left)
+                distances = (abs(predictor - int(left)), abs(predictor - int(up)), abs(predictor - int(upper_left)))
+                value = (left, up, upper_left)[int(np.argmin(distances))]
+            elif filter_type == 0:
+                value = 0
+            else:
+                raise ValueError(f"Unsupported PNG filter {filter_type}")
+            row[index] = (int(row[index]) + int(value)) & 255
+        rows.append(row)
+        previous = row
+    return np.frombuffer(b"".join(row.tobytes() for row in rows), dtype=">u2").astype(np.uint16).reshape(height, width)
+
+
+def decode_mrms_event(event_dir: str | Path, image_height: int = 512, image_width: int = 512, total_length: int = 29) -> np.ndarray:
+    event_dir = Path(event_dir)
+    event_name = event_dir.name
+    frames = [_read_png_gray16(event_dir / f"{event_name}-{index:02d}.png") for index in range(total_length)]
+    data = np.stack(frames).astype(np.float32) / 10.0 - 3.0
+    if data.shape[1:] != (image_height, image_width):
+        raise ValueError(f"MRMS frame shape {data.shape[1:]} != {(image_height, image_width)}")
+    mask = np.ones_like(data, dtype=np.float32)
+    mask[data < 0] = 0
+    data[data < 0] = 0
+    data = np.clip(data, 0, 128)
+    return np.stack([data, mask], axis=-1)
+
+
+class MRMSDataset(Dataset):
+    def __init__(self, data_dir: str | Path, image_height: int = 512, image_width: int = 512, total_length: int = 29, split: str = "train"):
+        self.data_dir = Path(data_dir)
+        self.events = sorted(path for path in self.data_dir.iterdir() if path.is_dir())
+        if split == "train":
+            self.events = [path for path in self.events if len(list(path.glob("*.png"))) == total_length]
+        elif split != "test":
+            raise ValueError("split must be train or test")
+        self.image_height = image_height
+        self.image_width = image_width
+        self.total_length = total_length
+
+    def __len__(self) -> int:
+        return len(self.events)
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor | str]:
+        event = self.events[index]
+        frames = decode_mrms_event(event, self.image_height, self.image_width, self.total_length)
+        return {"radar_frames": torch.from_numpy(frames), "event": event.name}

scripts/fake_data.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+"""Generate MRMS-shaped synthetic events for contract and smoke checks."""
+from __future__ import annotations
+
+import argparse
+import struct
+import zlib
+from pathlib import Path
+
+import numpy as np
+
+
+def write_png_gray16(path: Path, array: np.ndarray) -> None:
+    array = np.asarray(array, dtype=">u2")
+    raw = b"".join(b"\x00" + row.tobytes() for row in array)
+    def chunk(kind: bytes, payload: bytes) -> bytes:
+        return struct.pack(">I", len(payload)) + kind + payload + struct.pack(">I", zlib.crc32(kind + payload) & 0xffffffff)
+    header = struct.pack(">IIBBBBB", array.shape[1], array.shape[0], 16, 0, 0, 0, 0)
+    path.write_bytes(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", header) + chunk(b"IDAT", zlib.compress(raw, 1)) + chunk(b"IEND", b""))
+
+
+def generate(output_dir: str | Path, events: int = 2, height: int = 512, width: int = 512, total_length: int = 29, seed: int = 42) -> None:
+    root = Path(output_dir)
+    rng = np.random.default_rng(seed)
+    yy, xx = np.mgrid[:height, :width]
+    for event_index in range(events):
+        event = root / f"synthetic_{event_index:04d}"
+        event.mkdir(parents=True, exist_ok=True)
+        cx, cy = width * (0.25 + 0.2 * event_index), height * 0.45
+        for frame in range(total_length):
+            center_x = cx + frame * 1.5
+            rain = 12.0 * np.exp(-((xx - center_x) ** 2 + (yy - cy) ** 2) / (2 * (max(height, width) * 0.12) ** 2))
+            rain += rng.normal(0, 0.08, size=(height, width))
+            # Inverse official encoding: decoded = uint16 / 10 - 3.
+            encoded = np.clip(np.rint((np.maximum(rain, 0) + 3.0) * 10), 0, 65535).astype(np.uint16)
+            write_png_gray16(event / f"{event.name}-{frame:02d}.png", encoded)
+    print(f"generated {events} MRMS events at {root} with shape ({total_length},{height},{width})")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output-dir", type=str, default="data/data")
+    parser.add_argument("--events", type=int, default=2)
+    parser.add_argument("--height", type=int, default=512)
+    parser.add_argument("--width", type=int, default=512)
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+    generate(args.output_dir, args.events, args.height, args.width, seed=args.seed)

scripts/inference.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+import argparse
+from pathlib import Path
+import sys
+import numpy as np
+import torch
+import yaml
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+from scripts.data_loader import MRMSDataset
+from model.model_factory import build_model, load_checkpoint
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="NowcastNet inference")
+    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
+    parser.add_argument("--data-dir", help="MRMS event directory; defaults to data.data_dir in config.yaml")
+    parser.add_argument("--checkpoint")
+    parser.add_argument("--output-dir")
+    parser.add_argument("--device", default="auto")
+    parser.add_argument("--height", type=int)
+    parser.add_argument("--width", type=int)
+    parser.add_argument("--ngf", type=int)
+    args = parser.parse_args()
+    cfg = yaml.safe_load(Path(args.config).read_text())
+    mc, dc, ic = cfg["model"], cfg["data"], cfg["inference"]
+    height, width = args.height or dc["image_height"], args.width or dc["image_width"]
+    if args.ngf: mc["ngf"] = args.ngf
+    mc["img_height"], mc["img_width"] = height, width
+    device = torch.device("cuda" if args.device == "auto" and torch.cuda.is_available() else ("cpu" if args.device == "auto" else args.device))
+    checkpoint = Path(args.checkpoint) if args.checkpoint else PROJECT_ROOT / mc["checkpoint_dir"] / f"{mc.get('checkpoint_prefix', 'model_bak')}.pth"
+    if not checkpoint.is_file():
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint}. Run scripts/train.py first or pass --checkpoint.")
+    model = build_model(mc, device); load_checkpoint(model, checkpoint, device); model.eval()
+    data_dir = Path(args.data_dir) if args.data_dir else PROJECT_ROOT / dc["data_dir"]
+    ds = MRMSDataset(data_dir, height, width, dc["total_length"], "test")
+    out = Path(args.output_dir) if args.output_dir else PROJECT_ROOT / ic["output_dir"]
+    out.mkdir(parents=True, exist_ok=True)
+    with torch.no_grad():
+        for item in ds:
+            frames = item["radar_frames"].unsqueeze(0).to(device=device, dtype=torch.float32)
+            pred = model(frames).squeeze(0).squeeze(-1).cpu().numpy()
+            target = frames[0, mc["input_length"]:, :, :, 0].cpu().numpy()
+            np.save(out / f"{item['event']}_pred.npy", pred)
+            np.save(out / f"{item['event']}_input.npy", frames[0, :mc["input_length"], :, :, 0].cpu().numpy())
+            np.save(out / f"{item['event']}_target.npy", target)
+            print(item["event"], pred.shape, float(pred.min()), float(pred.max()))
+
+if __name__ == "__main__": main()

scripts/result.py

@@ -0,0 +1,160 @@
+#!/usr/bin/env python3
+"""Render NowcastNet predictions and truth comparisons as RGB PNG files."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+import struct
+import zlib
+
+import numpy as np
+import yaml
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+RAIN_THRESHOLDS = np.asarray([0.1, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0], dtype=np.float32)
+RAIN_COLORS = np.asarray(
+    [
+        [0, 0, 0],
+        [70, 70, 70],
+        [0, 110, 255],
+        [0, 205, 255],
+        [0, 190, 80],
+        [255, 230, 0],
+        [255, 145, 0],
+        [235, 35, 30],
+        [205, 0, 180],
+    ],
+    dtype=np.uint8,
+)
+ERROR_THRESHOLDS = np.asarray([0.1, 0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0], dtype=np.float32)
+ERROR_COLORS = np.asarray(
+    [
+        [0, 0, 0],
+        [40, 40, 40],
+        [35, 80, 170],
+        [30, 165, 215],
+        [80, 200, 120],
+        [245, 225, 65],
+        [245, 145, 45],
+        [220, 55, 40],
+        [245, 245, 245],
+    ],
+    dtype=np.uint8,
+)
+
+
+def _png_chunk(kind: bytes, payload: bytes) -> bytes:
+    checksum = zlib.crc32(kind + payload) & 0xFFFFFFFF
+    return struct.pack(">I", len(payload)) + kind + payload + struct.pack(">I", checksum)
+
+
+def write_png(path: Path, image: np.ndarray) -> None:
+    """Write an H x W x 3 uint8 array as a standards-compliant RGB PNG."""
+    image = np.asarray(image, dtype=np.uint8)
+    if image.ndim != 3 or image.shape[2] != 3:
+        raise ValueError(f"Expected RGB image [H,W,3], got {image.shape}")
+    raw = b"".join(b"\x00" + row.tobytes() for row in image)
+    header = struct.pack(">IIBBBBB", image.shape[1], image.shape[0], 8, 2, 0, 0, 0)
+    path.write_bytes(
+        b"\x89PNG\r\n\x1a\n"
+        + _png_chunk(b"IHDR", header)
+        + _png_chunk(b"IDAT", zlib.compress(raw, 1))
+        + _png_chunk(b"IEND", b"")
+    )
+
+
+def colorize(image: np.ndarray, thresholds: np.ndarray, colors: np.ndarray) -> np.ndarray:
+    values = np.nan_to_num(np.asarray(image, dtype=np.float32), nan=0.0, posinf=128.0, neginf=0.0)
+    return colors[np.searchsorted(thresholds, np.maximum(values, 0.0), side="right")]
+
+
+def comparison_image(truth: np.ndarray, prediction: np.ndarray) -> np.ndarray:
+    truth_rgb = colorize(truth, RAIN_THRESHOLDS, RAIN_COLORS)
+    prediction_rgb = colorize(prediction, RAIN_THRESHOLDS, RAIN_COLORS)
+    error_rgb = colorize(np.abs(prediction - truth), ERROR_THRESHOLDS, ERROR_COLORS)
+    separator = np.full((truth.shape[0], 4, 3), 255, dtype=np.uint8)
+    return np.concatenate([truth_rgb, separator, prediction_rgb, separator, error_rgb], axis=1)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Render NowcastNet inference results as PNG images")
+    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
+    parser.add_argument("--input-dir", help="directory containing *_pred.npy and *_target.npy")
+    parser.add_argument("--output-dir")
+    parser.add_argument("--threshold", type=float)
+    args = parser.parse_args()
+
+    cfg = yaml.safe_load(Path(args.config).read_text())
+    src = Path(args.input_dir) if args.input_dir else PROJECT_ROOT / cfg["inference"]["output_dir"]
+    out = Path(args.output_dir) if args.output_dir else PROJECT_ROOT / cfg["visualization"]["output_dir"]
+    threshold = args.threshold if args.threshold is not None else float(cfg["inference"]["threshold"])
+    expected_frames = int(cfg["model"]["total_length"]) - int(cfg["model"]["input_length"])
+    prediction_dir = out / "predictions"
+    comparison_dir = out / "comparison"
+    prediction_dir.mkdir(parents=True, exist_ok=True)
+    comparison_dir.mkdir(parents=True, exist_ok=True)
+
+    summary: dict[str, dict[str, object]] = {}
+    pred_paths = sorted(src.glob("*_pred.npy"))
+    if not pred_paths:
+        raise FileNotFoundError(f"No *_pred.npy inference results found under {src}")
+
+    for pred_path in pred_paths:
+        event = pred_path.name.removesuffix("_pred.npy")
+        target_path = src / f"{event}_target.npy"
+        if not target_path.is_file():
+            raise FileNotFoundError(
+                f"Truth file not found: {target_path}. Rerun scripts/inference.py to export targets."
+            )
+        prediction = np.load(pred_path)
+        truth = np.load(target_path)
+        if prediction.shape != truth.shape:
+            raise ValueError(f"Prediction shape {prediction.shape} != truth shape {truth.shape} for {event}")
+        if prediction.ndim != 3 or prediction.shape[0] != expected_frames:
+            raise ValueError(
+                f"Expected {expected_frames} frames [T,H,W] for {event}, got {prediction.shape}"
+            )
+
+        absolute_error = np.abs(prediction - truth)
+        mae_by_lead = absolute_error.mean(axis=(1, 2))
+        rmse_by_lead = np.sqrt(np.square(prediction - truth).mean(axis=(1, 2)))
+        for index in range(expected_frames):
+            filename = f"{event}_t{index + 1:02d}.png"
+            write_png(
+                prediction_dir / filename,
+                colorize(prediction[index], RAIN_THRESHOLDS, RAIN_COLORS),
+            )
+            write_png(
+                comparison_dir / filename,
+                comparison_image(truth[index], prediction[index]),
+            )
+
+        summary[event] = {
+            "shape": list(prediction.shape),
+            "prediction_png_count": expected_frames,
+            "comparison_png_count": expected_frames,
+            "comparison_layout": ["truth", "prediction", "absolute_error"],
+            "prediction_min": float(prediction.min()),
+            "prediction_max": float(prediction.max()),
+            "prediction_mean": float(prediction.mean()),
+            "threshold": threshold,
+            "threshold_fraction": float((prediction >= threshold).mean()),
+            "mae": float(absolute_error.mean()),
+            "rmse": float(np.sqrt(np.square(prediction - truth).mean())),
+            "mae_by_lead": [float(value) for value in mae_by_lead],
+            "rmse_by_lead": [float(value) for value in rmse_by_lead],
+        }
+
+    summary_path = out / "summary.json"
+    summary_path.write_text(json.dumps(summary, indent=2) + "\n")
+    print(f"prediction_png_dir={prediction_dir}")
+    print(f"comparison_png_dir={comparison_dir}")
+    print(f"summary={summary_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,150 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import random
+from pathlib import Path
+import sys
+
+import numpy as np
+import torch
+import yaml
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+# ``model`` and ``scripts`` are namespace packages; no __init__.py is required.
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+from scripts.data_loader import MRMSDataset
+from model.model_factory import build_model, load_checkpoint
+
+
+def clone_state_dict(model: torch.nn.Module) -> dict[str, torch.Tensor]:
+    """Copy the current parameters to CPU so later epochs cannot mutate them."""
+    network = model.module if hasattr(model, "module") else model
+    return {name: tensor.detach().cpu().clone() for name, tensor in network.state_dict().items()}
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="NowcastNet single-card training")
+    parser.add_argument("--config", default=str(Path(__file__).parents[1] / "conf/config.yaml"))
+    parser.add_argument("--data-dir", help="MRMS event directory; defaults to data.data_dir in config.yaml")
+    parser.add_argument("--device", default="auto")
+    parser.add_argument("--epochs", type=int)
+    parser.add_argument("--checkpoint")
+    parser.add_argument("--output-dir")
+    parser.add_argument("--log-dir", help="training log directory; defaults to training.output_dir in config.yaml")
+    parser.add_argument("--checkpoint-prefix")
+    parser.add_argument("--height", type=int)
+    parser.add_argument("--width", type=int)
+    parser.add_argument("--ngf", type=int)
+    parser.add_argument("--distributed", action="store_true", help="use torchrun for single-node or multi-node DDP")
+    args = parser.parse_args()
+    cfg = yaml.safe_load(Path(args.config).read_text())
+    mc, dc, tc = cfg["model"], cfg["data"], cfg["training"]
+    if args.height: mc["img_height"] = dc["image_height"] = args.height
+    if args.width: mc["img_width"] = dc["image_width"] = args.width
+    if args.ngf: mc["ngf"] = args.ngf
+    data_dir = Path(args.data_dir) if args.data_dir else PROJECT_ROOT / dc["data_dir"]
+    distributed = args.distributed or int(__import__("os").environ.get("WORLD_SIZE", "1")) > 1
+    rank = 0
+    if distributed:
+        import os, torch.distributed as dist
+        dist.init_process_group(backend="nccl" if torch.cuda.is_available() else "gloo", init_method="env://")
+        rank = dist.get_rank()
+        if args.device == "auto" and torch.cuda.is_available():
+            torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
+    if distributed and torch.cuda.is_available():
+        device = torch.device("cuda", int(__import__("os").environ.get("LOCAL_RANK", 0)))
+    else:
+        device = torch.device("cuda" if args.device == "auto" and torch.cuda.is_available() else ("cpu" if args.device == "auto" else args.device))
+    seed = int(tc["seed"]); random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
+    model = build_model(mc, device)
+    best_loss = float("inf")
+    best_epoch = 0
+    best_state = None
+    if args.checkpoint:
+        checkpoint = load_checkpoint(model, args.checkpoint, device)
+        if "best_loss" in checkpoint:
+            best_loss = float(checkpoint["best_loss"])
+            best_epoch = int(checkpoint.get("epoch", 0))
+            best_state = clone_state_dict(model)
+    dataset = MRMSDataset(data_dir, dc["image_height"], dc["image_width"], dc["total_length"], "train")
+    if not dataset:
+        raise RuntimeError(f"No complete MRMS events found under {data_dir}")
+    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
+    loader = DataLoader(dataset, batch_size=int(dc["batch_size"]), shuffle=sampler is None, sampler=sampler, num_workers=int(dc["num_workers"]), drop_last=False)
+    if distributed:
+        from torch.nn.parallel import DistributedDataParallel as DDP
+        # The official generator keeps architectural branches whose
+        # parameters are not touched by the L1 smoke-training path.
+        model = DDP(
+            model,
+            device_ids=[device.index] if device.type == "cuda" else None,
+            find_unused_parameters=True,
+        )
+    optimizer = torch.optim.Adam(model.parameters(), lr=float(tc["lr"]), weight_decay=float(tc["weight_decay"]))
+    epochs = args.epochs or int(tc["epochs"])
+    checkpoint_dir = Path(args.output_dir) if args.output_dir else PROJECT_ROOT / mc["checkpoint_dir"]
+    checkpoint_prefix = args.checkpoint_prefix or mc.get("checkpoint_prefix", "model_bak")
+    log_dir = Path(args.log_dir) if args.log_dir else PROJECT_ROOT / tc["output_dir"]
+    loss_history_path = log_dir / tc.get("loss_history_file", "loss_history.json")
+    loss_history: list[dict[str, float | int]] = []
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    if rank == 0:
+        log_dir.mkdir(parents=True, exist_ok=True)
+    for epoch in range(epochs):
+        if sampler is not None: sampler.set_epoch(epoch)
+        model.train()
+        loss_sum = torch.zeros(1, device=device)
+        sample_count = torch.zeros(1, device=device)
+        for batch in loader:
+            frames = batch["radar_frames"].to(device=device, dtype=torch.float32)
+            pred = model(frames)
+            target = frames[:, int(mc["input_length"]):, :, :, :1].squeeze(-1)
+            loss = torch.mean(torch.abs(pred.squeeze(-1) - target))
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+            loss_sum += loss.detach() * frames.shape[0]
+            sample_count += frames.shape[0]
+        if distributed:
+            torch.distributed.all_reduce(loss_sum)
+            torch.distributed.all_reduce(sample_count)
+        epoch_loss = (loss_sum / sample_count.clamp_min(1)).item()
+        if rank == 0:
+            print(f"epoch={epoch + 1}/{epochs} loss={epoch_loss:.6f}")
+            if epoch_loss < best_loss:
+                best_loss = epoch_loss
+                best_epoch = epoch + 1
+                best_state = clone_state_dict(model)
+            loss_history.append({"epoch": epoch + 1, "loss": epoch_loss})
+            loss_history_path.write_text(
+                json.dumps(
+                    {
+                        "epochs": loss_history,
+                        "best_epoch": best_epoch,
+                        "best_loss": best_loss,
+                    },
+                    indent=2,
+                )
+                + "\n"
+            )
+    if rank == 0:
+        if best_state is None:
+            raise RuntimeError("Training finished without producing a model state")
+        checkpoint_path = checkpoint_dir / f"{checkpoint_prefix}.pth"
+        torch.save(
+            {"state_dict": best_state, "epoch": best_epoch, "best_loss": best_loss},
+            checkpoint_path,
+        )
+        print(f"checkpoint={checkpoint_path}")
+        print(f"loss_history={loss_history_path}")
+    if distributed:
+        torch.distributed.barrier()
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()