Upload folder using huggingface_hub

README.md

@@ -0,0 +1,150 @@
+---
+frameworks: PyTorch
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Ocean Simulation
+- Global Ocean Forecasting
+- OM4
+- ConvNeXt
+tasks: []
+datasets:
+- M2LInES/Samudra-OM4
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">Samudra</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+Samudra is a global ocean emulator developed by the M2LInES team.
+
+Paper: Samudra: An AI Global Ocean Emulator for Climate
+
+https://doi.org/10.1029/2024GL114318
+
+# Model Description
+
+Samudra predicts global ocean states on an approximately one-degree grid with a five-day time step. It is designed to emulate the evolution of the OM4 ocean circulation model with a deep neural network.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Global ocean simulation | Train the model on OM4 data following the 77-state-channel and 4-forcing-channel Samudra protocol. |
+| Local quick validation | Use synthetic NPZ data to check training, inference, and ocean-field visualization. |
+| ModelScope / OneCode execution | Download the standalone model package, install dependencies, and run the scripts directly. |
+| Multi-GPU training | Launch PyTorch DistributedDataParallel with `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**
+
+- 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/Samudra --local-dir ./Samudra
+cd Samudra
+```
+
+### 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 official training data is generated by NOAA/GFDL OM4. The M2LInES project provides the dataset and documentation:
+
+https://huggingface.co/datasets/M2LInES/Samudra-OM4
+
+The dataset must be converted to the native NPZ layout expected by this repository. When real OM4 data is unavailable, generate a synthetic fixture for pipeline validation:
+
+```bash
+python scripts/fake_data.py
+```
+
+The synthetic fixture contains 77 prognostic state channels and 4 boundary-forcing channels and is not suitable for scientific evaluation.
+
+### Training
+
+Single GPU:
+
+```bash
+python scripts/train.py
+```
+
+Multi-GPU:
+
+```bash
+torchrun --nproc_per_node=8 scripts/train.py
+```
+
+The default checkpoint is saved to `data/checkpoints/model_bak.pth`.
+
+### Training Weights
+
+This repository provides a `weight/` directory for Samudra checkpoints. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+Inference performs an autoregressive rollout from `data/test.npz` and reads `data/checkpoints/model_bak.pth` by default:
+
+```bash
+python scripts/inference.py
+```
+
+Predictions are written to `result/output/prediction.npz`.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+The default outputs are `result/forecast_maps.png` and `result/temperature_profile.png`.
+
+# 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
+
+- Paper: https://doi.org/10.1029/2024GL114318
+- This repository is an independent reproduction of the original Samudra paper.

conf/config.yaml

@@ -0,0 +1,53 @@
+project:
+  name: Samudra
+  default_model_name: samudra
+  seed: 1
+
+fake_data:
+  time: 12
+  height: 32
+  width: 64
+
+data:
+  source_dir: ./data
+  output_dir: ./data
+  format: npz
+  grid:
+    height: 180
+    width: 360
+  depths_m: [2.5, 10.0, 22.5, 40.0, 65.0, 105.0, 165.0, 250.0, 375.0, 550.0, 775.0, 1050.0, 1400.0, 1850.0, 2400.0, 3100.0, 4000.0, 5000.0, 6000.0]
+  time_step_days: 5
+  input_state_steps: 2
+  output_state_steps: 2
+  recurrent_passes: 4
+
+model:
+  variant: thermo_dynamic
+  input_channels: 158
+  output_channels: 154
+  state_channels: 77
+  widths: [200, 250, 300, 400]
+  dilations: [1, 2, 4, 8]
+  bottleneck_width: 400
+  bottleneck_dilation: 8
+  checkpoint: null
+  load_official_weights: false
+
+training:
+  batch_size: 2
+  epochs: 10
+  learning_rate: 0.0006
+  weight_decay: 0.0
+  num_workers: 0
+  scheduler: cosine
+  save_frequency: 5
+  resume_checkpoint: null
+  output_dir: ./data/checkpoints
+
+inference:
+  checkpoint: null
+  rollout_steps: 6
+  output_dir: ./result/output
+
+visualization:
+  output_dir: ./result

config.json

@@ -0,0 +1,70 @@
+{
+  "model_name": "Samudra",
+  "model_type": "samudra",
+  "architectures": [
+    "Samudra"
+  ],
+  "framework": "PyTorch",
+  "domain": "ocean",
+  "task": "global-ocean-emulation",
+  "implementation": {
+    "entry_point": "model/samudra.py",
+    "scope": "thermo-dynamic Samudra v1 ConvNeXt U-Net with periodic longitude handling"
+  },
+  "architecture": {
+    "family": "dilated ConvNeXt U-Net",
+    "variant": "thermo_dynamic",
+    "input_channels": 158,
+    "output_channels": 154,
+    "state_channels": 77,
+    "widths": [
+      200,
+      250,
+      300,
+      400
+    ],
+    "dilations": [
+      1,
+      2,
+      4,
+      8
+    ],
+    "bottleneck_width": 400,
+    "activation": "capped GELU"
+  },
+  "data": {
+    "dataset": "NOAA/GFDL OM4",
+    "spatial_resolution_degrees": 1.0,
+    "time_step_days": 5,
+    "input_state_steps": 2,
+    "output_state_steps": 2,
+    "recurrent_passes": 4,
+    "depth_levels_m": [
+      2.5,
+      10.0,
+      22.5,
+      40.0,
+      65.0,
+      105.0,
+      165.0,
+      250.0,
+      375.0,
+      550.0,
+      775.0,
+      1050.0,
+      1400.0,
+      1850.0,
+      2400.0,
+      3100.0,
+      4000.0,
+      5000.0,
+      6000.0
+    ],
+    "protocol": "samudra_v1_om4"
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/samudra.py",
+    "scripts/train.py"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "ocean_emulation",
+  "model": "Samudra",
+  "input_format": "BCHW",
+  "protocol": "samudra_v1_om4",
+  "default_config": "conf/config.yaml",
+  "train": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/samudra.py

@@ -0,0 +1,198 @@
+"""Paper-version Samudra ConvNeXt U-Net implementation.
+
+The layer layout follows the official ``samudra_om4_v1/model.yaml`` and the
+official v1 ``blocks.py`` implementation. Inputs are channel-stacked tensors
+with shape ``[batch, 158, lat, lon]`` for the full thermo-dynamic variant.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Iterable
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def circular_pad_width(x: torch.Tensor, padding: int) -> torch.Tensor:
+    """Apply periodic longitude padding, including widths smaller than padding."""
+    if padding == 0:
+        return x
+    width = x.shape[-1]
+    repeats = (padding + width - 1) // width
+    tiled = x.repeat(1, 1, 1, repeats)
+    left = tiled[..., -padding:]
+    right = tiled[..., :padding]
+    return torch.cat((left, x, right), dim=-1)
+
+
+def globe_pad(x: torch.Tensor, padding: int) -> torch.Tensor:
+    x = circular_pad_width(x, padding)
+    return F.pad(x, (0, 0, padding, padding), mode="constant")
+
+
+class CappedGELU(nn.Module):
+    def __init__(self, cap_value: float = 10.0):
+        super().__init__()
+        self.gelu = nn.GELU()
+        # Keep this constant out of DDP buffers. DDP synchronizes buffers at
+        # every forward, which would mutate the autograd version counter during
+        # a recurrent multi-step rollout.
+        self.cap = float(cap_value)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.clamp(self.gelu(x), max=self.cap)
+
+
+class ConvNeXtBlock(nn.Module):
+    """Official v1 residual block with dilated 3x3 convolutions."""
+
+    def __init__(self, in_channels: int, out_channels: int, dilation: int):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.dilation = dilation
+        self.padding = dilation
+        self.skip = nn.Identity() if in_channels == out_channels else nn.Conv2d(in_channels, out_channels, 1)
+        hidden = in_channels * 4
+        self.layers = nn.ModuleList(
+            [
+                nn.Conv2d(in_channels, hidden, 3, dilation=dilation),
+                nn.BatchNorm2d(hidden),
+                CappedGELU(),
+                nn.Conv2d(hidden, hidden, 3, dilation=dilation),
+                nn.BatchNorm2d(hidden),
+                CappedGELU(),
+                nn.Conv2d(hidden, out_channels, 1),
+            ]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        skip = self.skip(x)
+        for layer in self.layers:
+            if isinstance(layer, nn.Conv2d) and layer.kernel_size == (3, 3):
+                x = globe_pad(x, self.padding)
+            x = layer(x)
+        return skip + x
+
+
+class PeriodicBilinearUpsample(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        width = x.shape[-1]
+        padded = F.pad(x, (1, 1, 0, 0), mode="circular")
+        result = F.interpolate(padded, scale_factor=2, mode="bilinear", align_corners=False)
+        return result[..., 2 : 2 + width * 2]
+
+
+class SamudraUNet(nn.Module):
+    def __init__(self, in_channels: int, widths: Iterable[int], dilations: Iterable[int]):
+        super().__init__()
+        widths = list(widths)
+        dilations = list(dilations)
+        if len(widths) != 4 or len(dilations) != 4:
+            raise ValueError("paper v1 requires four widths and four dilation rates")
+        channels = [in_channels, *widths]
+        self.down_blocks = nn.ModuleList(
+            ConvNeXtBlock(channels[i], channels[i + 1], dilations[i]) for i in range(4)
+        )
+        self.pools = nn.ModuleList(nn.AvgPool2d(2) for _ in range(4))
+        self.middle = ConvNeXtBlock(widths[-1], widths[-1], dilations[-1])
+        self.first_up = PeriodicBilinearUpsample()
+        reversed_widths = list(reversed(widths))
+        reversed_dilations = list(reversed(dilations))
+        self.up_blocks = nn.ModuleList()
+        self.upsamples = nn.ModuleList()
+        current = widths[-1]
+        for index in range(3):
+            target = reversed_widths[index + 1]
+            self.up_blocks.append(ConvNeXtBlock(current, target, reversed_dilations[index]))
+            self.upsamples.append(PeriodicBilinearUpsample())
+            current = target
+        self.final = ConvNeXtBlock(current, widths[0], reversed_dilations[-1])
+        self.out_channels = widths[0]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        skips = []
+        for block, pool in zip(self.down_blocks, self.pools):
+            x = block(x)
+            skips.append(x)
+            x = pool(x)
+        x = self.middle(x)
+        x = self.first_up(x)
+        x = self._merge(x, skips[-1])
+        for block, up, skip in zip(self.up_blocks, self.upsamples, reversed(skips[:-1])):
+            x = block(x)
+            x = up(x)
+            x = self._merge(x, skip)
+        return self.final(x)
+
+    @staticmethod
+    def _merge(x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:
+        height = min(x.shape[-2], skip.shape[-2])
+        width = min(x.shape[-1], skip.shape[-1])
+        x = x[..., :height, :width]
+        skip = skip[..., :height, :width]
+        if x.shape[1] != skip.shape[1]:
+            raise RuntimeError(f"skip channel mismatch: {x.shape[1]} != {skip.shape[1]}")
+        return x + skip
+
+
+class Samudra(nn.Module):
+    """Samudra v1 single-scale emulator.
+
+    ``variant='thermo_dynamic'`` uses the paper's 158-to-154 interface.
+    ``variant='thermo'`` uses 78 input channels and 154 output channels.
+    """
+
+    def __init__(
+        self,
+        variant: str = "thermo_dynamic",
+        input_channels: int | None = None,
+        output_channels: int | None = None,
+        widths: Iterable[int] = (200, 250, 300, 400),
+        dilations: Iterable[int] = (1, 2, 4, 8),
+    ):
+        super().__init__()
+        if variant not in {"thermo_dynamic", "thermo"}:
+            raise ValueError("variant must be thermo_dynamic or thermo")
+        default_input = 158 if variant == "thermo_dynamic" else 78
+        default_output = 154
+        self.variant = variant
+        self.input_channels = default_input if input_channels is None else input_channels
+        self.output_channels = default_output if output_channels is None else output_channels
+        self.unet = SamudraUNet(self.input_channels, widths, dilations)
+        self.decoder = nn.Conv2d(self.unet.out_channels, self.output_channels, 3)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.ndim != 4 or x.shape[1] != self.input_channels:
+            raise ValueError(f"expected [batch, {self.input_channels}, lat, lon], got {tuple(x.shape)}")
+        features = self.unet(x)
+        features = globe_pad(features, 1)
+        return self.decoder(features)
+
+    def load_official_checkpoint(self, path: str | Path, strict: bool = True) -> None:
+        checkpoint = torch.load(path, map_location="cpu")
+        if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
+            checkpoint = checkpoint["state_dict"]
+        if not isinstance(checkpoint, dict):
+            raise TypeError("official checkpoint must contain a state dictionary")
+        # Older local checkpoints stored the fixed GELU cap as a buffer.
+        checkpoint = {key: value for key, value in checkpoint.items() if not key.endswith(".cap")}
+        self.load_state_dict(checkpoint, strict=strict)
+
+
+def build_model(config: dict) -> Samudra:
+    """Build a model from the project's YAML model section."""
+    model_config = config.get("model", config)
+    model = Samudra(
+        variant=model_config.get("variant", "thermo_dynamic"),
+        input_channels=model_config.get("input_channels"),
+        output_channels=model_config.get("output_channels"),
+        widths=model_config.get("widths", (200, 250, 300, 400)),
+        dilations=model_config.get("dilations", (1, 2, 4, 8)),
+    )
+    checkpoint = model_config.get("checkpoint")
+    if checkpoint:
+        model.load_official_checkpoint(checkpoint)
+    return model

scripts/_bootstrap.py

@@ -0,0 +1,11 @@
+"""Make the project root importable when an entry point is run by path."""
+
+from pathlib import Path
+import sys
+
+ROOT = Path(__file__).resolve().parents[1]
+SCRIPT_DIR = str(ROOT / "scripts")
+if SCRIPT_DIR in sys.path:
+    sys.path.remove(SCRIPT_DIR)
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))

scripts/fake_data.py

@@ -0,0 +1,52 @@
+"""Generate virtual data in the native Samudra NPZ layout."""
+
+try:
+    from ._bootstrap import ROOT
+except ImportError:
+    from _bootstrap import ROOT
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+STATE_CHANNELS = 77
+BOUNDARY_CHANNELS = 4
+
+
+def generate(path: str | Path, time: int, height: int, width: int, seed: int) -> None:
+    if time < 10:
+        raise ValueError("time must be at least 10 for four-pass recurrent training")
+    rng = np.random.default_rng(seed)
+    prognostic = rng.standard_normal((time, STATE_CHANNELS, height, width), dtype=np.float32)
+    boundary = rng.standard_normal((time, BOUNDARY_CHANNELS, height, width), dtype=np.float32)
+    output = Path(path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, prognostic=prognostic, boundary=boundary)
+    print(f"saved native data: {output} prognostic={prognostic.shape} boundary={boundary.shape}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="./conf/config.yaml")
+    parser.add_argument("--train-output", default="./data/train.npz")
+    parser.add_argument("--test-output", default="./data/test.npz")
+    parser.add_argument("--time", type=int, default=None)
+    parser.add_argument("--height", type=int, default=None)
+    parser.add_argument("--width", type=int, default=None)
+    parser.add_argument("--seed", type=int, default=None)
+    args = parser.parse_args()
+    with open(args.config, encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    fake = config.get("fake_data", {})
+    time = args.time or int(fake.get("time", 12))
+    height = args.height or int(fake.get("height", 32))
+    width = args.width or int(fake.get("width", 64))
+    seed = args.seed if args.seed is not None else int(config["project"].get("seed", 1))
+    generate(args.train_output, time, height, width, seed)
+    generate(args.test_output, time, height, width, seed + 1)
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,109 @@
+"""Autoregressive inference for the paper-version Samudra model."""
+
+try:
+    from ._bootstrap import ROOT
+except ImportError:
+    from _bootstrap import ROOT
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+
+from model.samudra import build_model
+
+STATE_CHANNELS = 77
+BOUNDARY_CHANNELS = 4
+
+
+def load_data(path: str | Path) -> tuple[np.ndarray, np.ndarray]:
+    """Load and validate native Samudra time-major arrays."""
+    with np.load(path) as data:
+        prognostic = np.asarray(data["prognostic"], dtype=np.float32)
+        boundary = np.asarray(data["boundary"], dtype=np.float32)
+    if prognostic.ndim != 4 or prognostic.shape[1] != STATE_CHANNELS:
+        raise ValueError("prognostic must have shape [time, 77, lat, lon]")
+    if boundary.ndim != 4 or boundary.shape[1] != BOUNDARY_CHANNELS:
+        raise ValueError("boundary must have shape [time, 4, lat, lon]")
+    if prognostic.shape[0] != boundary.shape[0] or prognostic.shape[2:] != boundary.shape[2:]:
+        raise ValueError("prognostic and boundary time/grid dimensions must match")
+    return prognostic, boundary
+
+
+def load_checkpoint(model: torch.nn.Module, path: str, device: torch.device) -> None:
+    try:
+        checkpoint = torch.load(path, map_location=device, weights_only=True)
+    except TypeError:
+        checkpoint = torch.load(path, map_location=device)
+    state = checkpoint.get("model", checkpoint.get("state_dict", checkpoint))
+    model.load_state_dict(state)
+
+
+@torch.no_grad()
+def rollout(model, prognostic, boundary, steps, device):
+    # Training uses one consecutive boundary frame for each recurrent pass.
+    available_steps = boundary.shape[0] - 1
+    if prognostic.shape[0] < 2 or available_steps < 1:
+        raise ValueError(
+            "rollout requires at least two states and two boundary time samples"
+        )
+    if steps > available_steps:
+        raise ValueError(
+            f"requested {steps} model steps, but input data provides only {available_steps}; "
+            f"use --steps <= {available_steps} or provide a longer data file"
+        )
+    previous = torch.from_numpy(prognostic[0]).to(device)
+    current = torch.from_numpy(prognostic[1]).to(device)
+    predictions = []
+    for step in range(steps):
+        forcing = torch.from_numpy(boundary[step + 1]).to(device)
+        inputs = torch.cat((previous, current, forcing), dim=0).unsqueeze(0)
+        prediction = model(inputs).squeeze(0)
+        previous, current = prediction[:STATE_CHANNELS], prediction[STATE_CHANNELS:]
+        predictions.extend((previous.cpu().numpy(), current.cpu().numpy()))
+    return np.stack(predictions)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="./conf/config.yaml")
+    parser.add_argument("--data", default="./data/test.npz")
+    parser.add_argument("--checkpoint", default=None)
+    parser.add_argument("--output", default=None)
+    parser.add_argument("--steps", type=int, default=None)
+    parser.add_argument("--device", default=None)
+    args = parser.parse_args()
+    if args.steps is not None and args.steps < 1:
+        raise ValueError("--steps must be positive")
+    with open(args.config, encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu"))
+    model = build_model(config).to(device)
+    checkpoint = args.checkpoint or config["inference"].get("checkpoint") or "./data/checkpoints/model_bak.pth"
+    output_path = args.output or config["inference"].get("output_dir", "./result/output")
+    if Path(output_path).suffix != ".npz":
+        output_path = str(Path(output_path) / "prediction.npz")
+    load_checkpoint(model, checkpoint, device)
+    model.eval()
+    prognostic, boundary = load_data(args.data)
+    configured_steps = int(config["inference"].get("rollout_steps", 1))
+    if configured_steps < 1:
+        raise ValueError("inference.rollout_steps must be positive")
+    available_steps = boundary.shape[0] - 1
+    steps = args.steps if args.steps is not None else min(configured_steps, available_steps)
+    if args.steps is None and steps < configured_steps:
+        print(
+            f"input data supports {available_steps} model steps; "
+            f"using rollout_steps={steps}"
+        )
+    predictions = rollout(model, prognostic, boundary, steps, device)
+    output = Path(output_path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, predictions=predictions, forcing=boundary[1 : steps + 1])
+    print(f"saved predictions: {output} shape={predictions.shape}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,54 @@
+"""Create diagnostic plots from Samudra NPZ rollout outputs."""
+
+try:
+    from ._bootstrap import ROOT
+except ImportError:
+    from _bootstrap import ROOT
+
+import argparse
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+DEPTHS = (2.5, 10.0, 22.5, 40.0, 65.0, 105.0, 165.0, 250.0, 375.0, 550.0, 775.0, 1050.0, 1400.0, 1850.0, 2400.0, 3100.0, 4000.0, 5000.0, 6000.0)
+PROGNOSTIC_CHANNELS = tuple(
+    f"{name}_{level}"
+    for name in ("uo", "vo", "thetao", "so")
+    for level in range(len(DEPTHS))
+) + ("zos",)
+
+
+def plot_rollout(prediction_path: str, output_dir: str) -> None:
+    with np.load(prediction_path) as data:
+        predictions = np.asarray(data["predictions"])
+    output = Path(output_dir)
+    output.mkdir(parents=True, exist_ok=True)
+    fig, axes = plt.subplots(1, 2, figsize=(12, 4), constrained_layout=True)
+    axes[0].imshow(predictions[0, PROGNOSTIC_CHANNELS.index("zos")], cmap="coolwarm", aspect="auto")
+    axes[0].set_title("SSH forecast")
+    axes[1].imshow(predictions[0, PROGNOSTIC_CHANNELS.index("thetao_0")], cmap="turbo", aspect="auto")
+    axes[1].set_title("Surface potential temperature")
+    fig.savefig(output / "forecast_maps.png", dpi=150)
+    plt.close(fig)
+    indices = [PROGNOSTIC_CHANNELS.index(f"thetao_{i}") for i in range(len(DEPTHS))]
+    profile = predictions[:, indices].mean(axis=(0, 2, 3))
+    fig, ax = plt.subplots(figsize=(5, 5), constrained_layout=True)
+    ax.plot(profile, DEPTHS, marker="o")
+    ax.invert_yaxis()
+    ax.set(xlabel="potential temperature", ylabel="depth (m)", title="Mean temperature profile")
+    fig.savefig(output / "temperature_profile.png", dpi=150)
+    plt.close(fig)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prediction", default="./result/output/prediction.npz")
+    parser.add_argument("--output-dir", default="./result")
+    args = parser.parse_args()
+    plot_rollout(args.prediction, args.output_dir)
+
+
+if __name__ == "__main__":
+    main()
+    print("save to ./result/")

scripts/train.py

@@ -0,0 +1,197 @@
+"""Single-device and torchrun-based distributed training for Samudra v1."""
+
+try:
+    from ._bootstrap import ROOT
+except ImportError:
+    from _bootstrap import ROOT
+
+import argparse
+import os
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+from model.samudra import build_model
+
+STATE_CHANNELS = 77
+BOUNDARY_CHANNELS = 4
+
+
+def load_data(path: str | Path) -> tuple[np.ndarray, np.ndarray]:
+    """Load and validate native Samudra time-major arrays."""
+    with np.load(path) as data:
+        prognostic = np.asarray(data["prognostic"], dtype=np.float32)
+        boundary = np.asarray(data["boundary"], dtype=np.float32)
+    if prognostic.ndim != 4 or prognostic.shape[1] != STATE_CHANNELS:
+        raise ValueError("prognostic must have shape [time, 77, lat, lon]")
+    if boundary.ndim != 4 or boundary.shape[1] != BOUNDARY_CHANNELS:
+        raise ValueError("boundary must have shape [time, 4, lat, lon]")
+    if prognostic.shape[0] != boundary.shape[0] or prognostic.shape[2:] != boundary.shape[2:]:
+        raise ValueError("prognostic and boundary time/grid dimensions must match")
+    return prognostic, boundary
+
+
+class SamudraDataset(Dataset):
+    """Build recurrent training windows from native Samudra arrays."""
+
+    def __init__(self, path: str | Path, recurrent_passes: int):
+        self.prognostic, self.boundary = load_data(path)
+        self.recurrent_passes = recurrent_passes
+        self.end = self.prognostic.shape[0] - 2 * recurrent_passes
+        if self.end <= 1:
+            raise ValueError("dataset does not contain enough samples for recurrent training")
+
+    def __len__(self) -> int:
+        return self.end - 1
+
+    def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        t = index + 1
+        history = np.stack((self.prognostic[t - 1], self.prognostic[t]))
+        forcing = self.boundary[t : t + self.recurrent_passes]
+        labels = np.stack(
+            [
+                np.concatenate(
+                    (self.prognostic[t + 2 * step + 1], self.prognostic[t + 2 * step + 2])
+                )
+                for step in range(self.recurrent_passes)
+            ]
+        )
+        return torch.from_numpy(history), torch.from_numpy(forcing), torch.from_numpy(labels)
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def load_config(path: str) -> dict:
+    with open(path, encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def distributed_setup(device: torch.device) -> tuple[int, int, bool]:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    rank = int(os.environ.get("RANK", "0"))
+    distributed = world_size > 1
+    if distributed:
+        backend = "nccl" if device.type == "cuda" else "gloo"
+        torch.distributed.init_process_group(backend=backend)
+    return rank, world_size, distributed
+
+
+def save_checkpoint(path: Path, model: nn.Module, optimizer: torch.optim.Optimizer, scheduler, epoch: int, loss: float) -> None:
+    state = model.module.state_dict() if isinstance(model, DistributedDataParallel) else model.state_dict()
+    torch.save({"epoch": epoch, "loss": loss, "model": state, "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict()}, path)
+
+
+def freeze_batch_norm_stats(module: nn.Module) -> None:
+    """Prevent recurrent forwards from mutating BatchNorm buffers in one graph."""
+    for child in module.modules():
+        if isinstance(child, nn.modules.batchnorm._BatchNorm):
+            child.eval()
+
+
+def train(
+    config_path: str,
+    data_path: str,
+    device_name: str | None = None,
+    epochs_override: int | None = None,
+    output_dir_override: str | None = None,
+) -> None:
+    config = load_config(config_path)
+    rank = int(os.environ.get("RANK", "0"))
+    local_rank = int(os.environ.get("LOCAL_RANK", rank))
+    if device_name:
+        device = torch.device(device_name)
+    elif torch.cuda.is_available():
+        device = torch.device(f"cuda:{local_rank}")
+    else:
+        device = torch.device("cpu")
+    rank, world_size, distributed = distributed_setup(device)
+    set_seed(int(config["project"].get("seed", 1)) + rank)
+    if device.type == "cuda":
+        torch.cuda.set_device(device)
+
+    recurrent_passes = int(config["data"].get("recurrent_passes", 1))
+    dataset = SamudraDataset(data_path, recurrent_passes)
+    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True) if distributed else None
+    loader = DataLoader(dataset, batch_size=config["training"]["batch_size"], shuffle=sampler is None, sampler=sampler, num_workers=config["training"].get("num_workers", 0), pin_memory=device.type == "cuda")
+    model = build_model(config).to(device)
+    if distributed:
+        model = DistributedDataParallel(
+            model,
+            device_ids=[device.index] if device.type == "cuda" else None,
+            broadcast_buffers=False,
+        )
+    optimizer = torch.optim.Adam(model.parameters(), lr=config["training"]["learning_rate"], weight_decay=config["training"].get("weight_decay", 0.0))
+    epochs = epochs_override if epochs_override is not None else int(config["training"]["epochs"])
+    if epochs < 1:
+        raise ValueError("epochs must be positive")
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
+    resume = config["training"].get("resume_checkpoint")
+    start_epoch = 0
+    if resume:
+        try:
+            checkpoint = torch.load(resume, map_location=device, weights_only=True)
+        except TypeError:
+            checkpoint = torch.load(resume, map_location=device)
+        target = model.module if isinstance(model, DistributedDataParallel) else model
+        state = {key: value for key, value in checkpoint["model"].items() if not key.endswith(".cap")}
+        target.load_state_dict(state)
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        scheduler.load_state_dict(checkpoint["scheduler"])
+        start_epoch = int(checkpoint["epoch"]) + 1
+
+    output_dir = Path(output_dir_override or config["training"]["output_dir"])
+    output_dir.mkdir(parents=True, exist_ok=True)
+    for epoch in range(start_epoch, epochs):
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+        model.train()
+        freeze_batch_norm_stats(model)
+        total_loss = 0.0
+        for batch in loader:
+            optimizer.zero_grad(set_to_none=True)
+            history, forcing, labels = (item.to(device, non_blocking=True) for item in batch)
+            previous, current = history[:, 0], history[:, 1]
+            losses = []
+            for step in range(recurrent_passes):
+                prediction = model(torch.cat((previous, current, forcing[:, step]), dim=1))
+                losses.append(nn.functional.mse_loss(prediction, labels[:, step]))
+                previous, current = prediction[:, :STATE_CHANNELS], prediction[:, STATE_CHANNELS:]
+            loss = torch.stack(losses).mean()
+            loss.backward()
+            optimizer.step()
+            total_loss += float(loss.detach())
+        scheduler.step()
+        mean_loss = total_loss / max(1, len(loader))
+        if rank == 0:
+            print(f"epoch={epoch + 1} loss={mean_loss:.6e} lr={scheduler.get_last_lr()[0]:.6e}")
+            frequency = int(config["training"].get("save_frequency", 5))
+            if (epoch + 1) % frequency == 0 or epoch + 1 == epochs:
+                save_checkpoint(output_dir / f"epoch_{epoch + 1:04d}.pt", model, optimizer, scheduler, epoch, mean_loss)
+                save_checkpoint(output_dir / "model_bak.pth", model, optimizer, scheduler, epoch, mean_loss)
+    if distributed:
+        torch.distributed.destroy_process_group()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="./conf/config.yaml")
+    parser.add_argument("--data", default="./data/train.npz")
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--epochs", type=int, default=None)
+    parser.add_argument("--output-dir", default=None)
+    args = parser.parse_args()
+    train(args.config, args.data, args.device, args.epochs, args.output_dir)
+
+
+if __name__ == "__main__":
+    main()