Upload folder using huggingface_hub

README.md

@@ -0,0 +1,166 @@
+---
+frameworks: PyTorch
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Weather Forecast
+- Global Weather Forecast
+- ERA5
+- Neural ODE
+tasks: []
+datasets:
+- OneScience/ERA5
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">ClimODE</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+ClimODE is a weather forecasting model proposed in 2024 by researchers from Aalto University and collaborating institutions.
+
+Paper: ClimODE: Climate and Weather Forecasting with Physics-informed Neural ODEs
+
+https://arxiv.org/abs/2404.10024
+
+# Model Description
+
+ClimODE is a physics-informed neural ordinary differential equation model for global, monthly-scale, and regional climate and weather forecasting. It represents atmospheric evolution as a continuous-time dynamical system and incorporates a transport-based physical inductive bias into the neural ODE.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Global weather forecasting | Train or evaluate ClimODE on ERA5 data following this project's five-variable protocol. |
+| Local quick validation | Use synthetic ERA5 HDF5 data to check data loading, training, inference, evaluation, and 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**
+
+- Training and inference require a GPU or DCU recognized by PyTorch. CPU can be used to generate synthetic data and inspect configuration, but cannot run the current training and inference scripts.
+- Multi-GPU training uses the NCCL backend. Ensure that 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/ClimODE --local-dir ./ClimODE
+cd ClimODE
+```
+
+### 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 an ERA5 data slice for training. Download it and confirm that the paths in `conf/config.yaml` point to the downloaded data:
+
+```bash
+hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data
+```
+
+ClimODE reads the five variables `z`, `t`, `t2m`, `u10`, and `v10`, regridding the source `721x1440` fields to the model's `32x64` grid. The source variable mapping is defined in `conf/config.yaml`.
+
+### Generate Synthetic Data
+
+When real ERA5 data is unavailable, generate a full-resolution synthetic fixture for pipeline validation:
+
+```bash
+python scripts/fake_data.py
+```
+
+Synthetic data does not represent ERA5 and cannot reproduce the paper's metrics.
+
+### 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`.
+
+### Fine-tuning
+
+To fine-tune from an existing checkpoint, pass an explicit checkpoint and mode:
+
+```bash
+python scripts/train.py --mode finetune --checkpoint data/checkpoints/model_bak.pth
+```
+
+An official pretrained checkpoint can be selected explicitly with `--use-pretrained --pretrained-checkpoint <path>`.
+
+### Training Weights
+
+This repository provides a `weight/` directory for ClimODE checkpoints. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+Inference reads `data/checkpoints/model_bak.pth` by default. If it is unavailable, pass `--checkpoint` explicitly:
+
+```bash
+python scripts/inference.py
+```
+
+Predictions, uncertainty estimates, and targets are written to `result/output/`.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+The script computes latitude-weighted RMSE, ACC, and CRPS, and writes per-variable figures to `result/output/figures/`.
+
+# 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 OneScience adaptation of the ClimODE paper and is not the official Aalto-QuML release.

conf/config.yaml

@@ -0,0 +1,91 @@
+model:
+  name: ClimODE
+  default_checkpoint: ./data/checkpoints/model_bak.pth
+  pretrained_checkpoint: ./weight/ClimODE_global.pt
+  input_channels: 5
+  input_height: 32
+  input_width: 64
+  solver: euler
+  atol: 0.005
+  rtol: 0.005
+  use_attention: true
+  use_uncertainty: true
+  use_positional_encoder: false
+  learning_rate: 0.0005
+  weight_decay: 0.00001
+  epochs: 10
+  checkpoint_dir: ./data/checkpoints
+
+data:
+  data_dir: ./data
+  # ERA5Dataset receives this root and resolves data/<year>.h5 itself.
+  raw_data_dir: ./data/data
+  static_dir: ./data/static
+  output_dir: ./result/output
+  time_step_hours: 6
+  input_steps: 1
+  output_steps: 1
+  # ClimODE min-max normalization; ERA5Dataset's mean/std transform stays disabled.
+  normalize: true
+  # ClimODE uses the following order throughout the model and metrics.
+  variables: [z, t, t2m, u10, v10]
+  variable_sources:
+    z: geopotential_500
+    t: temperature_850
+    t2m: 2m_temperature
+    u10: 10m_u_component_of_wind
+    v10: 10m_v_component_of_wind
+  train_years: [2014, 2015]
+  val_years: [2016]
+  test_years: [2017]
+  raw_height: 721
+  raw_width: 1440
+  model_height: 32
+  model_width: 64
+  regrid_method: bilinear
+  stats_dir: ./data/static
+  static_file: ./data/static/constants.h5
+  dataloader:
+    batch_size: 8
+    num_workers: 0
+    pin_memory: false
+    drop_last: false
+  sequence_length: 8
+
+velocity:
+  cache_dir: ./data/checkpoints/velocity
+  optimizer: Adam
+  learning_rate: 2.0
+  epochs: 200
+  smoothing_alpha: 1.0e-7
+  kernel_sigma: 1.0
+
+training:
+  epochs: 10
+  finetune_epochs: 5
+  finetune_learning_rate: 0.00005
+  mode: scratch
+  seed: 42
+  ddp_backend: nccl
+  max_batches: null
+  log_file: ./result/train.jsonl
+
+fake_data:
+  # The default is intentionally small in time but preserves the real ERA5 grid.
+  timesteps: 12
+  years: [2014, 2015, 2016, 2017]
+  height: 721
+  width: 1440
+  seed: 42
+  dtype: float32
+
+runtime:
+  device: cuda
+  module: sghpc-mpi-gcc/26.3
+  conda_env: develop_base
+
+output:
+  checkpoint_name: model_bak.pth
+  metrics_file: ./result/metrics.json
+  prediction_file: ./result/output/predictions.npy
+  figure_dir: ./result/output/figures

config.json

@@ -0,0 +1,79 @@
+{
+  "model_name": "ClimODE",
+  "model_type": "climode",
+  "architectures": [
+    "ClimODE",
+    "ClimateEncoderFreeUncertain"
+  ],
+  "framework": "PyTorch",
+  "domain": "atmosphere",
+  "task": "global-weather-forecasting-with-neural-odes",
+  "implementation": {
+    "entry_point": "model/climode.py",
+    "scope": "OneScience adapter of the physics-informed neural transport ODE with optional attention and probabilistic output; it is not an AutoModel-compatible Transformers implementation"
+  },
+  "architecture": {
+    "family": "physics-informed neural ODE with advective transport and convolutional velocity/noise networks",
+    "input_format": "BCHW",
+    "input_grid_shape": [
+      32,
+      64
+    ],
+    "input_channels": 5,
+    "output_channels": 5,
+    "time_step_hours": 6,
+    "solver": "euler",
+    "absolute_tolerance": 0.005,
+    "relative_tolerance": 0.005,
+    "use_attention": true,
+    "use_uncertainty": true,
+    "use_positional_encoder": false,
+    "resnet_repetitions": [
+      5,
+      3,
+      2
+    ],
+    "resnet_hidden_channels": [
+      128,
+      64,
+      10
+    ],
+    "history_frames_for_context": 3
+  },
+  "data": {
+    "dataset": "ERA5",
+    "source_grid_shape": [
+      721,
+      1440
+    ],
+    "model_grid_shape": [
+      32,
+      64
+    ],
+    "regrid_method": "bilinear with periodic longitude handling",
+    "variables": [
+      "z",
+      "t",
+      "t2m",
+      "u10",
+      "v10"
+    ],
+    "variable_sources": {
+      "z": "geopotential_500",
+      "t": "temperature_850",
+      "t2m": "2m_temperature",
+      "u10": "10m_u_component_of_wind",
+      "v10": "10m_v_component_of_wind"
+    },
+    "input_steps": 1,
+    "output_steps": 1,
+    "normalization": "ClimODE min-max normalization; source statistics are loaded from the configured static directory",
+    "uncertainty_output": "mean and standard deviation fields when use_uncertainty is enabled"
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/climode.py",
+    "scripts/data_loader.py",
+    "scripts/train.py"
+  ]
+}

configuration.json

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

model/climode.py

@@ -0,0 +1,483 @@
+"""ClimODE's official global neural transport model.
+
+The equations and layer layout follow Aalto-QuML/ClimODE.  The small wrapper
+at the bottom adds configuration-friendly construction and checkpoint loading;
+the core ``ClimateEncoderFreeUncertain`` class intentionally keeps the
+official state and tensor layout.
+"""
+
+from __future__ import annotations
+
+import importlib
+from pathlib import Path
+import sys
+from typing import Any, Sequence
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from torchdiffeq import odeint as _torchdiffeq_odeint
+except ImportError:  # pragma: no cover - exercised when optional dependency is absent
+    _torchdiffeq_odeint = None
+
+
+def _euler_odeint(func, y0: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    """Equivalent fixed-step Euler solver used only when torchdiffeq is absent."""
+
+    states = [y0]
+    for index in range(1, len(t)):
+        previous = states[-1]
+        dt = t[index] - t[index - 1]
+        states.append(previous + dt * func(t[index - 1], previous))
+    return torch.stack(states, dim=0)
+
+
+def odeint(func, y0: torch.Tensor, t: torch.Tensor, method: str = "euler", **kwargs):
+    if _torchdiffeq_odeint is not None:
+        return _torchdiffeq_odeint(func, y0, t, method=method, **kwargs)
+    if method != "euler":
+        raise ImportError(
+            "torchdiffeq is required for solver=%r; install the official dependency."
+            % method
+        )
+    return _euler_odeint(func, y0, t)
+
+
+class OptimVelocity(nn.Module):
+    """Learn the initial per-channel velocity used to start the ODE system."""
+
+    def __init__(self, num_years: int, height: int, width: int, out_channels: int = 5):
+        super().__init__()
+        self.out_channels = out_channels
+        self.v_x = nn.Parameter(
+            torch.randn(num_years, 1, out_channels, height, width)
+        )
+        self.v_y = nn.Parameter(
+            torch.randn(num_years, 1, out_channels, height, width)
+        )
+
+    def forward(self, data: torch.Tensor):
+        u_y = torch.gradient(data, dim=3)[0]
+        u_x = torch.gradient(data, dim=4)[0]
+        divergence = torch.gradient(self.v_y, dim=3)[0] + torch.gradient(
+            self.v_x, dim=4
+        )[0]
+        adv = self.v_x * u_x + self.v_y * u_y + data * divergence
+        return adv, self.v_x, self.v_y
+
+
+class BoundaryPad(nn.Module):
+    """Reflect at the poles and wrap around the longitude seam."""
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        return F.pad(F.pad(value, (0, 0, 1, 1), "reflect"), (1, 1, 0, 0), "circular")
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.activation = nn.LeakyReLU(0.3)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=0)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        self.drop = nn.Dropout(p=0.1)
+        self.shortcut = (
+            nn.Conv2d(in_channels, out_channels, kernel_size=1)
+            if in_channels != out_channels
+            else nn.Identity()
+        )
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        value_padded = F.pad(
+            F.pad(value, (0, 0, 1, 1), "reflect"), (1, 1, 0, 0), "circular"
+        )
+        hidden = self.activation(self.bn1(self.conv1(value_padded)))
+        hidden = F.pad(
+            F.pad(hidden, (0, 0, 1, 1), "reflect"), (1, 1, 0, 0), "circular"
+        )
+        hidden = self.activation(self.bn2(self.conv2(hidden)))
+        return self.drop(hidden) + self.shortcut(value)
+
+
+class ClimateResNet2D(nn.Module):
+    def __init__(
+        self,
+        num_channels: int,
+        layers: Sequence[int],
+        hidden_size: Sequence[int],
+    ):
+        super().__init__()
+        if len(layers) != len(hidden_size):
+            raise ValueError("layers and hidden_size must have equal lengths")
+        self.layer_cnn = nn.ModuleList(
+            [nn.Sequential(*blocks_for_layer) for blocks_for_layer in self._split_blocks(layers, hidden_size, num_channels)]
+        )
+
+    @staticmethod
+    def _split_blocks(
+        layers: Sequence[int], hidden_size: Sequence[int], num_channels: int
+    ) -> list[list[nn.Module]]:
+        groups: list[list[nn.Module]] = []
+        in_channels = num_channels
+        for repetitions, out_channels in zip(layers, hidden_size):
+            group = [ResidualBlock(in_channels, out_channels)]
+            group.extend(
+                ResidualBlock(out_channels, out_channels)
+                for _ in range(1, int(repetitions))
+            )
+            groups.append(group)
+            in_channels = out_channels
+        return groups
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        output = value.float()
+        for layer in self.layer_cnn:
+            output = layer(output)
+        return output
+
+
+class SelfAttnConv(nn.Module):
+    """Official key-query-value attention convolution."""
+
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.query = self._conv(in_channels, in_channels // 8, stride=1)
+        self.key = self._key_conv(in_channels, in_channels // 8, stride=2)
+        self.value = self._key_conv(in_channels, out_channels, stride=2)
+        self.post_map = nn.Sequential(
+            nn.Conv2d(out_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        )
+        self.out_ch = out_channels
+
+    @staticmethod
+    def _conv(n_in: int, n_out: int, stride: int) -> nn.Sequential:
+        return nn.Sequential(
+            BoundaryPad(),
+            nn.Conv2d(n_in, n_in // 2, kernel_size=3, stride=stride, padding=0),
+            nn.LeakyReLU(0.3),
+            BoundaryPad(),
+            nn.Conv2d(n_in // 2, n_out, kernel_size=3, stride=stride, padding=0),
+            nn.LeakyReLU(0.3),
+            BoundaryPad(),
+            nn.Conv2d(n_out, n_out, kernel_size=3, stride=stride, padding=0),
+        )
+
+    @staticmethod
+    def _key_conv(n_in: int, n_out: int, stride: int) -> nn.Sequential:
+        return nn.Sequential(
+            nn.Conv2d(n_in, n_in // 2, kernel_size=3, stride=stride, padding=0),
+            nn.LeakyReLU(0.3),
+            nn.Conv2d(n_in // 2, n_out, kernel_size=3, stride=stride, padding=0),
+            nn.LeakyReLU(0.3),
+            nn.Conv2d(n_out, n_out, kernel_size=3, stride=1, padding=0),
+        )
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        size = value.size()
+        value = value.float()
+        query = self.query(value).flatten(-2, -1)
+        key = self.key(value).flatten(-2, -1)
+        val = self.value(value).flatten(-2, -1)
+        beta = F.softmax(torch.bmm(query.transpose(1, 2), key), dim=1)
+        output = torch.bmm(val, beta.transpose(1, 2))
+        output = output.view(-1, self.out_ch, size[-2], size[-1]).contiguous()
+        return self.post_map(output)
+
+
+class ClimateEncoderFreeUncertain(nn.Module):
+    """Official global ``Climate_encoder_free_uncertain`` implementation."""
+
+    def __init__(
+        self,
+        num_channels: int = 5,
+        const_channels: int = 2,
+        out_types: int = 5,
+        method: str = "euler",
+        use_att: bool = True,
+        use_err: bool = True,
+        use_pos: bool = False,
+    ):
+        super().__init__()
+        self.layers = [5, 3, 2]
+        self.hidden = [128, 64, 2 * out_types]
+        input_channels = 30 + out_types * int(use_pos) + 34 * (1 - int(use_pos))
+        self.vel_f = ClimateResNet2D(input_channels, self.layers, self.hidden)
+        if use_att:
+            self.vel_att = SelfAttnConv(input_channels, 10)
+            self.gamma = nn.Parameter(torch.tensor([0.1]))
+
+        self.scales = num_channels
+        self.const_channel = const_channels
+        self.out_ch = out_types
+        self.past_samples: torch.Tensor | int = 0
+        self.const_info: torch.Tensor | int = 0
+        self.lat_map: torch.Tensor | int = 0
+        self.lon_map: torch.Tensor | int = 0
+        self.method = method
+        err_in = 9 + out_types * int(use_pos) + 34 * (1 - int(use_pos))
+        if use_err:
+            self.noise_net = ClimateResNet2D(err_in, [3, 2, 2], [128, 64, 2 * out_types])
+        if use_pos:
+            self.pos_enc = ClimateResNet2D(4, [2, 1, 1], [32, 16, out_types])
+        self.att = use_att
+        self.err = use_err
+        self.pos = use_pos
+        self.pos_feat: torch.Tensor | int = 0
+        self.lsm: torch.Tensor | int = 0
+        self.oro: torch.Tensor | int = 0
+
+    def update_param(self, params: Sequence[torch.Tensor]) -> None:
+        if len(params) != 4:
+            raise ValueError("update_param expects past_samples, constants, latitude, longitude")
+        self.past_samples, self.const_info, self.lat_map, self.lon_map = params
+
+    def _time_features(
+        self, t: torch.Tensor, height: int, width: int, batch_size: int
+    ) -> tuple[torch.Tensor, ...]:
+        t_emb = ((t * 100) % 24).view(1, 1, 1, 1).expand(batch_size, 1, height, width)
+        sin_t = torch.sin(torch.pi * t_emb / 12 - torch.pi / 2)
+        cos_t = torch.cos(torch.pi * t_emb / 12 - torch.pi / 2)
+        sin_season = torch.sin(torch.pi * t_emb / (12 * 365) - torch.pi / 2)
+        cos_season = torch.cos(torch.pi * t_emb / (12 * 365) - torch.pi / 2)
+        return t_emb, torch.cat([sin_t, cos_t], dim=1), torch.cat(
+            [sin_season, cos_season], dim=1
+        )
+
+    def pde(self, t: torch.Tensor, state: torch.Tensor) -> torch.Tensor:
+        height, width = state.shape[-2:]
+        ds = state[:, -self.out_ch :].view(-1, self.out_ch, height, width).float()
+        velocity = state[:, : 2 * self.out_ch].view(
+            -1, 2 * self.out_ch, height, width
+        ).float()
+        t_emb, day_emb, season_emb = self._time_features(
+            t, height, width, ds.shape[0]
+        )
+        ds_grad_x = torch.gradient(ds, dim=3)[0]
+        ds_grad_y = torch.gradient(ds, dim=2)[0]
+        nabla_u = torch.cat([ds_grad_x, ds_grad_y], dim=1)
+
+        if self.pos:
+            combined = torch.cat(
+                [t_emb / 24, day_emb, season_emb, nabla_u, velocity, ds, self.pos_feat],
+                dim=1,
+            )
+        else:
+            cos_lat, sin_lat = torch.cos(self.new_lat_map), torch.sin(self.new_lat_map)
+            cos_lon, sin_lon = torch.cos(self.new_lon_map), torch.sin(self.new_lon_map)
+            time_cyclic = torch.cat([day_emb, season_emb], dim=1)
+            pos_feats = torch.cat(
+                [
+                    cos_lat,
+                    cos_lon,
+                    sin_lat,
+                    sin_lon,
+                    sin_lat * cos_lon,
+                    sin_lat * sin_lon,
+                ],
+                dim=1,
+            )
+            pos_time = self.get_time_pos_embedding(time_cyclic, pos_feats)
+            combined = torch.cat(
+                [
+                    t_emb / 24,
+                    day_emb,
+                    season_emb,
+                    nabla_u,
+                    velocity,
+                    ds,
+                    self.new_lat_map,
+                    self.new_lon_map,
+                    self.lsm,
+                    self.oro,
+                    pos_feats,
+                    pos_time,
+                ],
+                dim=1,
+            )
+
+        dv = self.vel_f(combined)
+        if self.att:
+            dv = dv + self.gamma * self.vel_att(combined)
+        v_x = velocity[:, : self.out_ch]
+        v_y = velocity[:, self.out_ch :]
+        advection = v_x * ds_grad_x + v_y * ds_grad_y
+        advection = advection + ds * (
+            torch.gradient(v_x, dim=3)[0] + torch.gradient(v_y, dim=2)[0]
+        )
+        return torch.cat([dv, advection], dim=1)
+
+    @staticmethod
+    def get_time_pos_embedding(
+        time_features: torch.Tensor, position_features: torch.Tensor
+    ) -> torch.Tensor:
+        outputs = [feature.unsqueeze(1) * position_features for feature in time_features.unbind(1)]
+        return torch.cat(outputs, dim=1)
+
+    def noise_net_contrib(
+        self,
+        time: torch.Tensor,
+        pos_enc: torch.Tensor,
+        s_final: torch.Tensor,
+        height: int,
+        width: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        t_emb = (time % 24).view(-1, 1, 1, 1, 1)
+        sin_t = torch.sin(torch.pi * t_emb / 12 - torch.pi / 2).expand(
+            len(s_final), s_final.shape[1], 1, height, width
+        )
+        cos_t = torch.cos(torch.pi * t_emb / 12 - torch.pi / 2).expand(
+            len(s_final), s_final.shape[1], 1, height, width
+        )
+        sin_season = torch.sin(torch.pi * t_emb / (12 * 365) - torch.pi / 2).expand(
+            len(s_final), s_final.shape[1], 1, height, width
+        )
+        cos_season = torch.cos(torch.pi * t_emb / (12 * 365) - torch.pi / 2).expand(
+            len(s_final), s_final.shape[1], 1, height, width
+        )
+        pos_rep = pos_enc.expand(len(s_final), s_final.shape[1], -1, height, width)
+        pos_rep = pos_rep.flatten(start_dim=0, end_dim=1)
+        time_cyclic = torch.cat([sin_t, cos_t, sin_season, cos_season], dim=2)
+        time_cyclic = time_cyclic.flatten(start_dim=0, end_dim=1)
+        pos_time = self.get_time_pos_embedding(time_cyclic, pos_rep[:, 2:-2])
+        combined = torch.cat(
+            [time_cyclic, s_final.flatten(start_dim=0, end_dim=1), pos_rep, pos_time],
+            dim=1,
+        )
+        final_out = self.noise_net(combined).view(
+            len(time), -1, 2 * self.out_ch, height, width
+        )
+        mean = s_final + final_out[:, :, : self.out_ch]
+        std = F.softplus(final_out[:, :, self.out_ch :])
+        return mean, std
+
+    def forward(
+        self,
+        time_steps: torch.Tensor,
+        data: torch.Tensor,
+        atol: float = 0.1,
+        rtol: float = 0.1,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if not isinstance(self.past_samples, torch.Tensor):
+            raise RuntimeError("Call update_param before forward")
+        height, width = self.past_samples.shape[-2:]
+        values = data.float().view(-1, self.out_ch, height, width)
+        final_data = torch.cat([self.past_samples, values], dim=1)
+        init_time = time_steps[0].item() * 6
+        final_time = time_steps[-1].item() * 6
+        steps_val = final_time - init_time
+
+        if self.pos:
+            lat_map = self.lat_map.unsqueeze(0) * torch.pi / 180
+            lon_map = self.lon_map.unsqueeze(0) * torch.pi / 180
+            pos_rep = torch.cat([lat_map.unsqueeze(0), lon_map.unsqueeze(0), self.const_info], dim=1)
+            self.pos_feat = self.pos_enc(pos_rep).expand(
+                values.shape[0], -1, values.shape[-2], values.shape[-1]
+            )
+            final_pos_enc = self.pos_feat
+        else:
+            self.oro = self.const_info[0, 0]
+            self.lsm = self.const_info[0, 1]
+            self.lsm = self.lsm.unsqueeze(0).expand(values.shape[0], -1, height, width)
+            self.oro = F.normalize(self.oro).unsqueeze(0).expand(values.shape[0], -1, height, width)
+            self.new_lat_map = self.lat_map.expand(values.shape[0], 1, height, width) * torch.pi / 180
+            self.new_lon_map = self.lon_map.expand(values.shape[0], 1, height, width) * torch.pi / 180
+            cos_lat, sin_lat = torch.cos(self.new_lat_map), torch.sin(self.new_lat_map)
+            cos_lon, sin_lon = torch.cos(self.new_lon_map), torch.sin(self.new_lon_map)
+            pos_feats = torch.cat(
+                [cos_lat, cos_lon, sin_lat, sin_lon, sin_lat * cos_lon, sin_lat * sin_lon],
+                dim=1,
+            )
+            final_pos_enc = torch.cat(
+                [self.new_lat_map, self.new_lon_map, pos_feats, self.lsm, self.oro], dim=1
+            )
+
+        integration_steps = max(int(steps_val) + 1, 1)
+        new_time_steps = torch.linspace(
+            init_time, final_time, steps=integration_steps, device=values.device
+        )
+        ode_time = 0.01 * new_time_steps.float()
+        final_result = odeint(
+            self.pde,
+            final_data,
+            ode_time,
+            method=self.method,
+            atol=atol,
+            rtol=rtol,
+        )
+        s_final = final_result[:, :, -self.out_ch :].view(
+            len(ode_time), -1, self.out_ch, height, width
+        )
+        sampled = s_final[0 : len(s_final) : 6]
+        if self.err:
+            mean, std = self.noise_net_contrib(
+                time_steps, final_pos_enc, sampled, height, width
+            )
+            return mean, std, sampled
+        return sampled, torch.zeros_like(sampled), sampled
+
+
+class ClimODE(ClimateEncoderFreeUncertain):
+    """Configuration-friendly public model name."""
+
+    def __init__(
+        self,
+        num_channels: int = 5,
+        const_channels: int = 2,
+        out_types: int = 5,
+        method: str = "euler",
+        use_attention: bool = True,
+        use_uncertainty: bool = True,
+        use_positional_encoder: bool = False,
+    ):
+        super().__init__(
+            num_channels=num_channels,
+            const_channels=const_channels,
+            out_types=out_types,
+            method=method,
+            use_att=use_attention,
+            use_err=use_uncertainty,
+            use_pos=use_positional_encoder,
+        )
+
+
+def _register_checkpoint_compat_modules() -> None:
+    """Expose legacy top-level module names used by official pickle files."""
+
+    for legacy_name, current_name in (
+        ("model_function", "model.model_function"),
+        ("model_utils", "model.model_utils"),
+    ):
+        module = importlib.import_module(current_name)
+        sys.modules.setdefault(legacy_name, module)
+
+
+def load_checkpoint(path: str | Path, map_location: str | torch.device = "cpu") -> nn.Module:
+    """Load an official full-object checkpoint or a state-dict checkpoint."""
+
+    checkpoint_path = Path(path)
+    if not checkpoint_path.is_file():
+        raise FileNotFoundError(checkpoint_path)
+    _register_checkpoint_compat_modules()
+    try:
+        checkpoint = torch.load(checkpoint_path, map_location=map_location, weights_only=False)
+    except TypeError:  # Older PyTorch does not expose weights_only.
+        checkpoint = torch.load(checkpoint_path, map_location=map_location)
+    if isinstance(checkpoint, nn.Module):
+        return checkpoint
+    model = ClimODE()
+    if isinstance(checkpoint, dict):
+        state_dict = checkpoint.get("state_dict", checkpoint.get("model", checkpoint))
+    else:
+        state_dict = checkpoint
+    model.load_state_dict(state_dict)
+    return model
+
+
+# Names kept for the official checkpoint's pickle module/class references.
+Climate_encoder_free_uncertain = ClimateEncoderFreeUncertain
+Climate_ResNet_2D = ClimateResNet2D
+Self_attn_conv = SelfAttnConv
+boundarypad = BoundaryPad

model/model_function.py

@@ -0,0 +1,21 @@
+"""Compatibility exports for the official ClimODE checkpoint pickle."""
+
+from model.climode import (
+    ClimateEncoderFreeUncertain,
+    ClimateResNet2D,
+    Climate_encoder_free_uncertain,
+    Climate_ResNet_2D,
+    OptimVelocity,
+    SelfAttnConv,
+    Self_attn_conv,
+)
+
+__all__ = [
+    "Climate_encoder_free_uncertain",
+    "Climate_ResNet_2D",
+    "OptimVelocity",
+    "Self_attn_conv",
+    "ClimateEncoderFreeUncertain",
+    "ClimateResNet2D",
+    "SelfAttnConv",
+]

model/model_utils.py

@@ -0,0 +1,8 @@
+"""Compatibility exports for the official ClimODE checkpoint pickle."""
+
+from model.climode import BoundaryPad, ResidualBlock, SelfAttnConv
+
+boundarypad = BoundaryPad
+Self_attn_conv = SelfAttnConv
+
+__all__ = ["boundarypad", "ResidualBlock", "Self_attn_conv"]

scripts/data_loader.py

@@ -0,0 +1,314 @@
+"""OneScience ERA5Dataset adapter for ClimODE's 32x64 global grid."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Iterable, Sequence
+
+import h5py
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+
+try:
+    from onescience.datapipes.climate.era5 import ERA5Dataset
+except ImportError as exc:  # pragma: no cover - exercised only without OneScience
+    ERA5Dataset = None
+    _ERA5_IMPORT_ERROR = exc
+else:
+    _ERA5_IMPORT_ERROR = None
+
+
+OFFICIAL_VARIABLES = ("z", "t", "t2m", "u10", "v10")
+
+
+def _require_era5dataset() -> None:
+    if ERA5Dataset is None:
+        raise ImportError(
+            "ClimODE data loading requires OneScience ERA5Dataset; "
+            "activate an environment containing onescience before running."
+        ) from _ERA5_IMPORT_ERROR
+
+
+def _as_channel_vector(values: np.ndarray | torch.Tensor) -> torch.Tensor:
+    tensor = torch.as_tensor(values, dtype=torch.float32)
+    return tensor.reshape(-1)
+
+
+def _regrid_periodic(frame: torch.Tensor, target_size: tuple[int, int]) -> torch.Tensor:
+    """Bilinearly sample [C,H,W] on WeatherBench cell centers."""
+
+    if frame.ndim != 3:
+        raise ValueError(f"Expected [C,H,W], got {tuple(frame.shape)}")
+    target_height, target_width = target_size
+    if frame.shape[-2:] == target_size:
+        return frame
+    periodic = torch.cat([frame, frame[..., :1]], dim=-1).unsqueeze(0)
+    latitude = torch.linspace(
+        90.0 - 90.0 / target_height,
+        -90.0 + 90.0 / target_height,
+        target_height,
+        device=frame.device,
+        dtype=frame.dtype,
+    )
+    longitude = (
+        torch.arange(target_width, device=frame.device, dtype=frame.dtype)
+        * (360.0 / target_width)
+    )
+    lat2d, lon2d = torch.meshgrid(latitude, longitude, indexing="ij")
+    grid = torch.stack([lon2d / 180.0 - 1.0, -lat2d / 90.0], dim=-1)
+    return F.grid_sample(
+        periodic,
+        grid.unsqueeze(0),
+        mode="bilinear",
+        padding_mode="border",
+        align_corners=True,
+    )[0]
+
+
+def _load_stats(stats_dir: str | Path, channels: int) -> tuple[torch.Tensor, torch.Tensor]:
+    stats_path = Path(stats_dir)
+    minimum = _as_channel_vector(np.load(stats_path / "min_values.npy"))
+    maximum = _as_channel_vector(np.load(stats_path / "max_values.npy"))
+    if minimum.numel() != channels or maximum.numel() != channels:
+        raise ValueError(
+            f"Expected {channels} channel statistics, got {minimum.numel()} and {maximum.numel()}"
+        )
+    if torch.any(maximum <= minimum):
+        raise ValueError("All max_values must be greater than min_values")
+    return minimum, maximum
+
+
+def _normalize(frame: torch.Tensor, minimum: torch.Tensor, maximum: torch.Tensor) -> torch.Tensor:
+    scale = (maximum - minimum).clamp_min(torch.finfo(frame.dtype).eps)
+    return (frame - minimum[:, None, None]) / scale[:, None, None]
+
+
+class ClimODEDataset(Dataset):
+    """Return three history frames and the following target frame.
+
+    The underlying annual files are always read by OneScience ``ERA5Dataset``.
+    No direct HDF5 field access is used here, which keeps the model adapter
+    compatible with the OneScience ERA5 contract.
+    """
+
+    def __init__(
+        self,
+        data_dir: str | Path,
+        years: Sequence[int],
+        used_variables: Sequence[str] = OFFICIAL_VARIABLES,
+        stats_dir: str | Path | None = None,
+        model_size: tuple[int, int] = (32, 64),
+        normalize: bool = True,
+        input_steps: int = 1,
+        output_steps: int = 1,
+    ) -> None:
+        _require_era5dataset()
+        if tuple(used_variables) != OFFICIAL_VARIABLES:
+            raise ValueError(
+                "ClimODE requires the exact channel order ['z','t','t2m','u10','v10']"
+            )
+        if input_steps != 1 or output_steps != 1:
+            raise ValueError("The ClimODE adapter currently uses one input and one target step")
+        if len(years) == 0:
+            raise ValueError("At least one year is required")
+        # ``data_dir`` is the OneScience dataset root containing data/*.h5,
+        # rather than the nested data/ directory itself.
+        self.data_dir = Path(data_dir)
+        self.years = [int(year) for year in years]
+        self.variables = tuple(used_variables)
+        self.model_size = tuple(model_size)
+        self.normalize = normalize
+        self.era5 = ERA5Dataset(
+            dataset_dir=str(self.data_dir),
+            used_years=self.years,
+            used_variables=list(self.variables),
+            input_steps=1,
+            output_steps=1,
+            normalize=False,
+        )
+        if (self.era5.H, self.era5.W) != (721, 1440):
+            raise ValueError(
+                "ClimODE raw-data adapter expects (721,1440), "
+                f"got ({self.era5.H},{self.era5.W})"
+            )
+        self.samples_per_year = self.era5.samples_per_year
+        if self.samples_per_year < 3:
+            raise ValueError("Each year needs at least four frames for a three-frame history and target")
+        self.samples_per_year_with_history = self.samples_per_year - 2
+        self.minimum, self.maximum = (
+            _load_stats(stats_dir or self.data_dir / "static", len(self.variables))
+            if normalize
+            else (torch.zeros(len(self.variables)), torch.ones(len(self.variables)))
+        )
+
+    def __len__(self) -> int:
+        return len(self.years) * self.samples_per_year_with_history
+
+    def _frame(self, sample_index: int, target: bool = False) -> torch.Tensor:
+        invar, outvar, _, _, _ = self.era5[sample_index]
+        frame = outvar if target else invar
+        frame = _regrid_periodic(torch.as_tensor(frame, dtype=torch.float32), self.model_size)
+        if self.normalize:
+            frame = _normalize(frame, self.minimum, self.maximum)
+        return frame
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor | int | str]:
+        if index < 0:
+            index += len(self)
+        if index < 0 or index >= len(self):
+            raise IndexError(index)
+        year_index = index // self.samples_per_year_with_history
+        local_index = index % self.samples_per_year_with_history
+        base = year_index * self.samples_per_year + local_index + 2
+        history = torch.stack([self._frame(base - 2), self._frame(base - 1), self._frame(base)])
+        target = self._frame(base, target=True)
+        return {
+            "history": history,
+            "input": history[-1],
+            "target": target,
+            "year": self.years[year_index],
+            "step_index": local_index + 2,
+        }
+
+
+class ClimODESeriesDataset(Dataset):
+    """Official-style sequence batches with years as the inner batch axis.
+
+    Each item contains a contiguous sequence for every requested year.  The
+    outer DataLoader should use ``batch_size=1``; the sequence length plays the
+    role of the official training batch of time points.
+    """
+
+    def __init__(
+        self,
+        data_dir: str | Path,
+        years: Sequence[int],
+        used_variables: Sequence[str] = OFFICIAL_VARIABLES,
+        stats_dir: str | Path | None = None,
+        model_size: tuple[int, int] = (32, 64),
+        sequence_length: int = 8,
+        normalize: bool = True,
+    ) -> None:
+        _require_era5dataset()
+        if tuple(used_variables) != OFFICIAL_VARIABLES:
+            raise ValueError("ClimODE requires the exact channel order ['z','t','t2m','u10','v10']")
+        if sequence_length < 1:
+            raise ValueError("sequence_length must be positive")
+        self.data_dir = Path(data_dir)
+        self.years = [int(year) for year in years]
+        self.variables = tuple(used_variables)
+        self.model_size = tuple(model_size)
+        self.sequence_length = int(sequence_length)
+        self.normalize = normalize
+        self.era5 = ERA5Dataset(
+            dataset_dir=str(self.data_dir),
+            used_years=self.years,
+            used_variables=list(self.variables),
+            input_steps=1,
+            output_steps=1,
+            normalize=False,
+        )
+        if (self.era5.H, self.era5.W) != (721, 1440):
+            raise ValueError(
+                "ClimODE raw-data adapter expects (721,1440), "
+                f"got ({self.era5.H},{self.era5.W})"
+            )
+        self.samples_per_year = self.era5.samples_per_year
+        self.frames_per_year = self.era5.T
+        first_start = 2
+        # The official DataLoader keeps its final, possibly shorter batch.
+        self.starts = list(range(first_start, self.frames_per_year, self.sequence_length))
+        if not self.starts:
+            raise ValueError(
+                f"Not enough frames ({self.era5.T}) for sequence_length={sequence_length} "
+                "and a three-frame history"
+            )
+        self.minimum, self.maximum = (
+            _load_stats(stats_dir or self.data_dir / "static", len(self.variables))
+            if normalize
+            else (torch.zeros(len(self.variables)), torch.ones(len(self.variables)))
+        )
+
+    def __len__(self) -> int:
+        return len(self.starts)
+
+    def _frame(self, year_index: int, frame_index: int) -> torch.Tensor:
+        if frame_index < 0 or frame_index >= self.frames_per_year:
+            raise IndexError(frame_index)
+        sample_index = year_index * self.samples_per_year + min(
+            frame_index, self.samples_per_year - 1
+        )
+        invar, outvar, _, _, _ = self.era5[sample_index]
+        # ERA5Dataset's final input index is T-2; its paired target is frame T-1.
+        frame = outvar if frame_index == self.frames_per_year - 1 else invar
+        frame = _regrid_periodic(torch.as_tensor(frame, dtype=torch.float32), self.model_size)
+        if self.normalize:
+            frame = _normalize(frame, self.minimum, self.maximum)
+        return frame
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor | int]:
+        start = self.starts[index]
+        history_per_year = []
+        sequence_per_year = []
+        for year_index in range(len(self.years)):
+            history_per_year.append(
+                torch.stack(
+                    [
+                        self._frame(year_index, start - 2),
+                        self._frame(year_index, start - 1),
+                        self._frame(year_index, start),
+                    ]
+                )
+            )
+            stop = min(start + self.sequence_length, self.frames_per_year)
+            sequence_per_year.append(
+                torch.stack(
+                    [self._frame(year_index, step) for step in range(start, stop)]
+                )
+            )
+        stop = min(start + self.sequence_length, self.frames_per_year)
+        return {
+            "history": torch.stack(history_per_year, dim=0),
+            "observations": torch.stack(sequence_per_year, dim=1),
+            "time_steps": torch.arange(start, stop, dtype=torch.float32),
+            "sequence_index": index,
+        }
+
+
+def load_constants(
+    static_file: str | Path,
+    expected_size: tuple[int, int] = (32, 64),
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Load [orography, lsm], latitude and longitude from constants.h5."""
+
+    with h5py.File(static_file, "r") as handle:
+        constants = torch.stack(
+            [
+                torch.as_tensor(handle["orography"][:], dtype=torch.float32),
+                torch.as_tensor(handle["lsm"][:], dtype=torch.float32),
+            ]
+        ).unsqueeze(0)
+        lat2d = torch.as_tensor(handle["lat2d"][:], dtype=torch.float32)
+        lon2d = torch.as_tensor(handle["lon2d"][:], dtype=torch.float32)
+    if tuple(constants.shape[-2:]) != expected_size:
+        raise ValueError(f"Static constants have shape {tuple(constants.shape[-2:])}")
+    return constants, lat2d, lon2d
+
+
+def make_dataloader(
+    dataset: Dataset,
+    batch_size: int,
+    shuffle: bool,
+    num_workers: int = 0,
+    pin_memory: bool = False,
+) -> DataLoader:
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        drop_last=False,
+    )

scripts/fake_data.py

@@ -0,0 +1,221 @@
+"""Generate small, deterministic ERA5-shaped HDF5 files for workflow checks."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Iterable
+
+import h5py
+import numpy as np
+import yaml
+
+
+VARIABLES = ("z", "t", "t2m", "u10", "v10")
+ERA5_HEIGHT = 721
+ERA5_WIDTH = 1440
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+
+def _parse_years(value: str | Iterable[int]) -> list[int]:
+    if isinstance(value, str):
+        return [int(item.strip()) for item in value.split(",") if item.strip()]
+    return [int(item) for item in value]
+
+
+def _synthetic_frame(
+    step: int,
+    year: int,
+    lat: np.ndarray,
+    lon: np.ndarray,
+) -> np.ndarray:
+    """Create smooth fields with distinct scales for the five official channels."""
+
+    lat_rad = np.deg2rad(lat)[:, None]
+    lon_rad = np.deg2rad(lon)[None, :]
+    phase = 2.0 * np.pi * (step + (year % 100)) / 1460.0
+    spatial = np.sin(lat_rad) + 0.35 * np.cos(lon_rad) + 0.15 * np.sin(
+        2.0 * lon_rad + phase
+    )
+    seasonal = np.cos(lat_rad) * np.sin(phase)
+
+    channels = np.stack(
+        [
+            5000.0 + 300.0 * spatial + 20.0 * seasonal,
+            260.0 + 12.0 * spatial + 2.0 * seasonal,
+            280.0 + 18.0 * spatial + 3.0 * seasonal,
+            4.0 * np.cos(lon_rad + phase) + 0.5 * spatial,
+            3.0 * np.sin(lon_rad - phase) - 0.5 * spatial,
+        ],
+        axis=0,
+    )
+    return channels.astype(np.float32, copy=False)
+
+
+def _write_static(static_dir: Path, height: int, width: int) -> None:
+    static_dir.mkdir(parents=True, exist_ok=True)
+    lat = np.linspace(
+        90.0 - 90.0 / height,
+        -90.0 + 90.0 / height,
+        height,
+        dtype=np.float32,
+    )
+    lon = np.linspace(0.0, 360.0, width, endpoint=False, dtype=np.float32)
+    lat2d, lon2d = np.meshgrid(lat, lon, indexing="ij")
+    orography = (1200.0 * np.maximum(np.cos(np.deg2rad(lat2d)), 0.0)).astype(
+        np.float32
+    )
+    lsm = (np.cos(np.deg2rad(lat2d)) > 0.25).astype(np.float32)
+
+    with h5py.File(static_dir / "constants.h5", "w") as handle:
+        handle.create_dataset("orography", data=orography)
+        handle.create_dataset("lsm", data=lsm)
+        handle.create_dataset("lat2d", data=lat2d)
+        handle.create_dataset("lon2d", data=lon2d)
+        handle.attrs["variables"] = np.asarray(
+            ["orography", "lsm"], dtype=h5py.string_dtype()
+        )
+
+
+def generate_data(
+    output_dir: str | Path,
+    years: Iterable[int],
+    timesteps: int,
+    height: int = ERA5_HEIGHT,
+    width: int = ERA5_WIDTH,
+    seed: int = 42,
+    overwrite: bool = False,
+    write_static: bool = True,
+) -> dict[str, list[float]]:
+    """Generate annual files and return per-channel global statistics."""
+
+    if timesteps < 4:
+        raise ValueError("timesteps must be at least 4 for a three-frame history")
+    if (height, width) != (ERA5_HEIGHT, ERA5_WIDTH):
+        raise ValueError(
+            "ClimODE virtual ERA5 data must use the raw shape "
+            f"({ERA5_HEIGHT}, {ERA5_WIDTH})"
+        )
+
+    root = Path(output_dir)
+    data_dir = root / "data"
+    static_dir = root / "static"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    years = _parse_years(years)
+    lat = np.linspace(90.0, -90.0, height, dtype=np.float32)
+    lon = np.linspace(0.0, 360.0, width, endpoint=False, dtype=np.float32)
+    rng = np.random.default_rng(seed)
+    minimum = np.full(len(VARIABLES), np.inf, dtype=np.float64)
+    maximum = np.full(len(VARIABLES), -np.inf, dtype=np.float64)
+    total = np.zeros(len(VARIABLES), dtype=np.float64)
+    total_sq = np.zeros(len(VARIABLES), dtype=np.float64)
+    total_count = 0
+
+    for year in years:
+        path = data_dir / f"{year}.h5"
+        if path.exists() and not overwrite:
+            raise FileExistsError(f"Refusing to overwrite existing file: {path}")
+        with h5py.File(path, "w") as handle:
+            fields = handle.create_dataset(
+                "fields",
+                shape=(timesteps, len(VARIABLES), height, width),
+                dtype=np.float32,
+                chunks=(1, len(VARIABLES), height, width),
+            )
+            fields.attrs["variables"] = np.asarray(
+                VARIABLES, dtype=h5py.string_dtype()
+            )
+            fields.attrs["time_step"] = 6
+            for step in range(timesteps):
+                frame = _synthetic_frame(step, year, lat, lon)
+                # A tiny deterministic per-frame perturbation keeps years distinct
+                # without materializing another 20 MB random tensor per frame.
+                frame += np.float32(rng.normal(0.0, 1.0e-3))
+                fields[step] = frame
+                flat = frame.reshape(len(VARIABLES), -1).astype(np.float64)
+                minimum = np.minimum(minimum, flat.min(axis=1))
+                maximum = np.maximum(maximum, flat.max(axis=1))
+                total += flat.sum(axis=1)
+                total_sq += np.square(flat).sum(axis=1)
+                total_count += flat.shape[1]
+
+            # Placeholders are replaced with statistics over every requested year.
+            handle.create_dataset("global_means", shape=(1, len(VARIABLES), 1, 1), dtype=np.float32)
+            handle.create_dataset("global_stds", shape=(1, len(VARIABLES), 1, 1), dtype=np.float32)
+
+    means = (total / total_count).astype(np.float32)
+    variances = np.maximum(
+        total_sq / total_count - means.astype(np.float64) ** 2, 1.0e-12
+    )
+    stds = np.sqrt(variances).astype(np.float32)
+    for year in years:
+        with h5py.File(data_dir / f"{year}.h5", "r+") as handle:
+            handle["global_means"][:] = means.reshape(1, -1, 1, 1)
+            handle["global_stds"][:] = stds.reshape(1, -1, 1, 1)
+
+    static_height, static_width = 32, 64
+    if write_static:
+        _write_static(static_dir, static_height, static_width)
+    np.save(static_dir / "min_values.npy", minimum.astype(np.float32))
+    np.save(static_dir / "max_values.npy", maximum.astype(np.float32))
+    return {"min": minimum.tolist(), "max": maximum.tolist()}
+
+
+def _load_config(path: Path) -> dict:
+    with path.open("r", encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def _resolve(path: str | Path) -> Path:
+    value = Path(path)
+    return value if value.is_absolute() else PROJECT_ROOT / value
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--config", type=Path, default=PROJECT_ROOT / "conf/config.yaml"
+    )
+    parser.add_argument("--output-dir", type=Path, default=None)
+    parser.add_argument("--years", type=str, default=None, help="Comma-separated years")
+    parser.add_argument("--timesteps", 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)
+    parser.add_argument("--overwrite", action="store_true")
+    args = parser.parse_args()
+    config_path = _resolve(args.config)
+    config = _load_config(config_path) if config_path.exists() else {}
+    fake = config.get("fake_data", {})
+    output_dir = _resolve(
+        args.output_dir or config.get("data", {}).get("data_dir", "./data")
+    )
+    years = _parse_years(args.years) if args.years else fake.get("years", [2006, 2016, 2017])
+    stats = generate_data(
+        output_dir=output_dir,
+        years=years,
+        timesteps=(
+            args.timesteps
+            if args.timesteps is not None
+            else fake.get("timesteps", 8)
+        ),
+        height=(
+            args.height
+            if args.height is not None
+            else fake.get("height", ERA5_HEIGHT)
+        ),
+        width=(
+            args.width
+            if args.width is not None
+            else fake.get("width", ERA5_WIDTH)
+        ),
+        seed=args.seed if args.seed is not None else fake.get("seed", 42),
+        overwrite=args.overwrite,
+    )
+    print(f"Generated years={years} under {Path(output_dir).resolve()}")
+    print(f"min={stats['min']}")
+    print(f"max={stats['max']}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,206 @@
+"""Run ClimODE global forecasts and save machine-readable outputs."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+from torch.utils.data import DataLoader
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from model.climode import load_checkpoint
+from scripts.data_loader import ClimODESeriesDataset, load_constants
+from scripts.metrics import evaluate, save_metrics
+from scripts.velocity import fit_velocity_cache, load_velocity_cache
+
+
+def _load_config(path: Path) -> dict:
+    with path.open("r", encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def _parse_years(value: str | None, fallback: list[int]) -> list[int]:
+    if value is None:
+        return list(fallback)
+    years = [int(item.strip()) for item in value.split(",") if item.strip()]
+    if not years:
+        raise ValueError("year override must contain at least one integer")
+    return years
+
+
+def _device(value: str | None) -> torch.device:
+    if value:
+        return torch.device(value)
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def _resolve(path: str | Path) -> Path:
+    value = Path(path)
+    return value if value.is_absolute() else PROJECT_ROOT / value
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "conf/config.yaml")
+    parser.add_argument("--checkpoint", type=Path, default=None)
+    parser.add_argument("--device", type=str, default=None)
+    parser.add_argument("--test-years", type=str, default=None)
+    parser.add_argument("--sequence-length", type=int, default=None)
+    parser.add_argument("--max-samples", type=int, default=None)
+    parser.add_argument("--velocity-epochs", type=int, default=None)
+    parser.add_argument("--velocity-cache", type=Path, default=None)
+    parser.add_argument("--data-dir", type=Path, default=None, help="Override data.data_dir")
+    parser.add_argument("--stats-dir", type=Path, default=None, help="Override data.stats_dir")
+    parser.add_argument("--static-file", type=Path, default=None, help="Override data.static_file")
+    parser.add_argument("--output-dir", type=Path, default=None)
+    args = parser.parse_args()
+    args.config = _resolve(args.config)
+    args.checkpoint = _resolve(args.checkpoint) if args.checkpoint is not None else None
+    args.velocity_cache = (
+        _resolve(args.velocity_cache) if args.velocity_cache is not None else None
+    )
+    args.data_dir = _resolve(args.data_dir) if args.data_dir is not None else None
+    args.stats_dir = _resolve(args.stats_dir) if args.stats_dir is not None else None
+    args.static_file = _resolve(args.static_file) if args.static_file is not None else None
+    args.output_dir = _resolve(args.output_dir) if args.output_dir is not None else None
+    config = _load_config(args.config)
+    data_cfg, model_cfg, vel_cfg = config["data"], config["model"], config["velocity"]
+    root = _resolve(args.data_dir or data_cfg["data_dir"])
+    stats_dir = _resolve(args.stats_dir or data_cfg.get("stats_dir", root / "static"))
+    test_years = _parse_years(args.test_years, data_cfg["test_years"])
+    sequence_length = args.sequence_length or data_cfg.get("sequence_length", 8)
+    dataset = ClimODESeriesDataset(
+        root,
+        test_years,
+        stats_dir=stats_dir,
+        model_size=(data_cfg["model_height"], data_cfg["model_width"]),
+        sequence_length=sequence_length,
+        normalize=data_cfg.get("normalize", True),
+    )
+    loader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)
+    static_file = _resolve(args.static_file or data_cfg["static_file"])
+    constants, lat, lon = load_constants(
+        static_file, (data_cfg["model_height"], data_cfg["model_width"])
+    )
+    device = _device(args.device)
+    constants = constants.to(device)
+    lat_device, lon_device = lat.unsqueeze(0).to(device), lon.unsqueeze(0).to(device)
+
+    velocity_root = _resolve(vel_cfg["cache_dir"])
+    velocity_path = args.velocity_cache or (velocity_root / "test.pt")
+    if velocity_path.is_file():
+        velocity = load_velocity_cache(velocity_path, len(dataset))
+    else:
+        velocity = fit_velocity_cache(
+            dataset,
+            constants,
+            lat,
+            lon,
+            velocity_path,
+            epochs=args.velocity_epochs if args.velocity_epochs is not None else vel_cfg["epochs"],
+            learning_rate=vel_cfg["learning_rate"],
+            smoothing_alpha=vel_cfg["smoothing_alpha"],
+            kernel_sigma=vel_cfg["kernel_sigma"],
+        )
+
+    checkpoint_path = args.checkpoint
+    if checkpoint_path is None:
+        checkpoint_path = _resolve(model_cfg["default_checkpoint"])
+        if not checkpoint_path.is_file():
+            pretrained = _resolve(model_cfg["pretrained_checkpoint"])
+            if pretrained.is_file():
+                checkpoint_path = pretrained
+    if checkpoint_path is None or not checkpoint_path.is_file():
+        raise FileNotFoundError(
+            "No checkpoint found; pass --checkpoint or provide model.default_checkpoint"
+        )
+    model = load_checkpoint(checkpoint_path, map_location="cpu").to(device).eval()
+    predictions, uncertainties, targets = [], [], []
+    with torch.no_grad():
+        for sample_index, batch in enumerate(loader):
+            if args.max_samples is not None and sample_index >= args.max_samples:
+                break
+            observations = batch["observations"].squeeze(0).to(device)
+            time_steps = batch["time_steps"].squeeze(0).to(device)
+            initial = observations[0].unsqueeze(1)
+            model.update_param([velocity[sample_index].to(device), constants, lat_device, lon_device])
+            mean, std, _ = model(
+                time_steps,
+                initial,
+                atol=model_cfg["atol"],
+                rtol=model_cfg["rtol"],
+            )
+            # Index 0 is the analysis state used to initialize the ODE. Official
+            # evaluation starts at index 1, corresponding to a six-hour lead.
+            if mean.shape[0] > 1:
+                predictions.append(mean[1:].detach().cpu().numpy())
+                uncertainties.append(std[1:].detach().cpu().numpy())
+                targets.append(observations[1:].detach().cpu().numpy())
+    if not predictions:
+        raise RuntimeError("No test samples were processed")
+
+    valid_lengths = np.asarray([item.shape[0] for item in predictions], dtype=np.int64)
+    max_lead = int(valid_lengths.max())
+
+    def _pad(items: list[np.ndarray]) -> np.ndarray:
+        shape = (len(items), max_lead) + tuple(items[0].shape[1:])
+        padded = np.full(shape, np.nan, dtype=np.float32)
+        for index, item in enumerate(items):
+            padded[index, : item.shape[0]] = item
+        return padded
+
+    pred_array = _pad(predictions)
+    std_array = _pad(uncertainties)
+    target_array = _pad(targets)
+    scale = (dataset.maximum - dataset.minimum).numpy().reshape(1, 1, 1, 5, 1, 1)
+    offset = dataset.minimum.numpy().reshape(1, 1, 1, 5, 1, 1)
+    pred_physical = pred_array * scale + offset
+    target_physical = target_array * scale + offset
+    std_physical = std_array * scale
+    output_dir = args.output_dir or _resolve(data_cfg["output_dir"])
+    output_dir.mkdir(parents=True, exist_ok=True)
+    np.save(output_dir / "predictions.npy", pred_array)
+    np.save(output_dir / "std.npy", std_array)
+    np.save(output_dir / "targets.npy", target_array)
+    np.save(output_dir / "valid_lengths.npy", valid_lengths)
+    metrics = evaluate(
+        pred_physical,
+        target_physical,
+        lat.numpy(),
+        std_physical,
+        crps_predictions=pred_array,
+        crps_targets=target_array,
+        crps_std=std_array,
+        valid_lengths=valid_lengths,
+    )
+    metrics["checkpoint"] = str(checkpoint_path)
+    metrics["outputs_normalized"] = True
+    metrics_path = _resolve(config["output"]["metrics_file"])
+    if args.output_dir is not None:
+        metrics_path = output_dir.parent / "metrics.json"
+    save_metrics(metrics, metrics_path)
+    manifest = {
+        "checkpoint": str(checkpoint_path),
+        "samples": int(pred_array.shape[0]),
+        "shape": list(pred_array.shape),
+        "valid_lengths": valid_lengths.tolist(),
+        "variables": ["z", "t", "t2m", "u10", "v10"],
+        "output_dir": str(output_dir),
+        "metrics": str(metrics_path),
+    }
+    (output_dir / "inference_manifest.json").write_text(
+        json.dumps(manifest, indent=2), encoding="utf-8"
+    )
+    print(json.dumps(manifest))
+
+
+if __name__ == "__main__":
+    main()

scripts/metrics.py

@@ -0,0 +1,179 @@
+"""ClimODE evaluation metrics and output serialization helpers."""
+
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+from typing import Sequence
+
+import numpy as np
+
+try:
+    from scipy.special import erf as _erf
+except ImportError:  # pragma: no cover - only used in minimal environments
+    _erf = np.vectorize(math.erf)
+
+
+VARIABLES = ("z", "t", "t2m", "u10", "v10")
+
+
+def latitude_weights(lat2d: np.ndarray) -> np.ndarray:
+    lat = np.asarray(lat2d, dtype=np.float64)
+    if lat.ndim == 2:
+        lat = lat[:, 0]
+    weights = np.cos(np.deg2rad(lat))
+    weights = weights / np.mean(weights)
+    return weights[:, None]
+
+
+def _check_arrays(
+    predictions: np.ndarray,
+    targets: np.ndarray,
+    std: np.ndarray | None,
+    valid_lengths: Sequence[int] | None = None,
+) -> None:
+    if predictions.shape != targets.shape:
+        raise ValueError(f"predictions {predictions.shape} != targets {targets.shape}")
+    if predictions.ndim != 6:
+        raise ValueError("Expected [samples, lead, years, channels, height, width]")
+    if predictions.shape[3] != len(VARIABLES):
+        raise ValueError(f"Expected {len(VARIABLES)} channels, got {predictions.shape[3]}")
+    if std is not None and std.shape != predictions.shape:
+        raise ValueError(f"std {std.shape} != predictions {predictions.shape}")
+    if valid_lengths is not None:
+        lengths = np.asarray(valid_lengths, dtype=np.int64)
+        if lengths.shape != (predictions.shape[0],):
+            raise ValueError(f"valid_lengths {lengths.shape} != ({predictions.shape[0]},)")
+        if np.any(lengths < 1) or np.any(lengths > predictions.shape[1]):
+            raise ValueError("valid_lengths must be within the lead dimension")
+
+
+def _lead_mask(
+    predictions: np.ndarray,
+    valid_lengths: Sequence[int] | None,
+) -> np.ndarray:
+    lengths = (
+        np.full(predictions.shape[0], predictions.shape[1], dtype=np.int64)
+        if valid_lengths is None
+        else np.asarray(valid_lengths, dtype=np.int64)
+    )
+    return (np.arange(predictions.shape[1])[None, :] < lengths[:, None]).reshape(
+        predictions.shape[0], predictions.shape[1], 1, 1, 1, 1
+    )
+
+
+def _weighted_mean(values: np.ndarray, weights: np.ndarray) -> np.ndarray:
+    # values: [N,L,Y,K,H,W], weights: [H,1]
+    weighted = values * weights[None, None, None, None, :, :]
+    return weighted.mean(axis=(-1, -2))
+
+
+def latitude_weighted_rmse(
+    predictions: np.ndarray,
+    targets: np.ndarray,
+    lat2d: np.ndarray,
+    valid_lengths: Sequence[int] | None = None,
+) -> np.ndarray:
+    weights = latitude_weights(lat2d)
+    error = np.square(np.nan_to_num(predictions - targets, nan=0.0))
+    per_field = np.sqrt(_weighted_mean(error, weights))
+    valid = _lead_mask(predictions, valid_lengths)[..., 0, 0, 0, 0]
+    valid_fields = np.broadcast_to(valid[:, :, None, None], per_field.shape)
+    return (per_field * valid_fields).sum(axis=(0, 2)) / np.maximum(
+        valid_fields.sum(axis=(0, 2)), 1.0
+    )
+
+
+def anomaly_correlation(
+    predictions: np.ndarray,
+    targets: np.ndarray,
+    lat2d: np.ndarray,
+    valid_lengths: Sequence[int] | None = None,
+) -> np.ndarray:
+    weights = latitude_weights(lat2d)
+    valid = _lead_mask(predictions, valid_lengths)
+    valid_broadcast = np.broadcast_to(valid, targets.shape)
+    target_clean = np.nan_to_num(targets, nan=0.0)
+    valid_count = valid_broadcast.sum(axis=(0, 1))
+    # Official evaluation uses one test-set climatology for each year/channel/grid.
+    climatology = (target_clean * valid_broadcast).sum(axis=(0, 1)) / np.maximum(
+        valid_count, 1.0
+    )
+    pred_anomaly = np.nan_to_num(predictions, nan=0.0) - climatology[None, None]
+    target_anomaly = target_clean - climatology[None, None]
+    pred_anomaly -= pred_anomaly.mean(axis=(-1, -2), keepdims=True)
+    target_anomaly -= target_anomaly.mean(axis=(-1, -2), keepdims=True)
+    weighted_mask = weights[None, None, None, None] * valid
+    numerator = (pred_anomaly * target_anomaly * weighted_mask).sum(axis=(-1, -2))
+    pred_norm = np.sqrt((np.square(pred_anomaly) * weighted_mask).sum(axis=(-1, -2)))
+    target_norm = np.sqrt((np.square(target_anomaly) * weighted_mask).sum(axis=(-1, -2)))
+    per_field = numerator / np.maximum(pred_norm * target_norm, 1.0e-12)
+    valid_fields = np.broadcast_to(valid[..., 0, 0], per_field.shape)
+    return (per_field * valid_fields).sum(axis=(0, 2)) / np.maximum(
+        valid_fields.sum(axis=(0, 2)), 1.0
+    )
+
+
+def _normal_crps(
+    observations: np.ndarray,
+    means: np.ndarray,
+    scales: np.ndarray,
+) -> np.ndarray:
+    """Closed-form CRPS for a Gaussian predictive distribution."""
+
+    scales = np.maximum(np.asarray(scales, dtype=np.float64), 1.0e-6)
+    z = (np.asarray(observations, dtype=np.float64) - means) / scales
+    phi = np.exp(-0.5 * np.square(z)) / math.sqrt(2.0 * math.pi)
+    cdf = 0.5 * (1.0 + _erf(z / math.sqrt(2.0)))
+    return scales * (z * (2.0 * cdf - 1.0) + 2.0 * phi - 1.0 / math.sqrt(math.pi))
+
+
+def gaussian_crps(
+    predictions: np.ndarray,
+    targets: np.ndarray,
+    std: np.ndarray,
+    valid_lengths: Sequence[int] | None = None,
+) -> np.ndarray:
+    values = np.nan_to_num(_normal_crps(targets, predictions, std), nan=0.0)
+    mask = np.broadcast_to(_lead_mask(predictions, valid_lengths), values.shape)
+    return (values * mask).sum(axis=(0, 2, 4, 5)) / np.maximum(
+        mask.sum(axis=(0, 2, 4, 5)), 1.0
+    )
+
+
+def evaluate(
+    predictions: np.ndarray,
+    targets: np.ndarray,
+    lat2d: np.ndarray,
+    std: np.ndarray | None = None,
+    crps_predictions: np.ndarray | None = None,
+    crps_targets: np.ndarray | None = None,
+    crps_std: np.ndarray | None = None,
+    valid_lengths: Sequence[int] | None = None,
+) -> dict:
+    _check_arrays(predictions, targets, std, valid_lengths)
+    result = {
+        "variables": list(VARIABLES),
+        "lead_times_hours": [6 * (index + 1) for index in range(predictions.shape[1])],
+        "rmse": latitude_weighted_rmse(predictions, targets, lat2d, valid_lengths).tolist(),
+        "acc": anomaly_correlation(predictions, targets, lat2d, valid_lengths).tolist(),
+        "rmse_space": "physical",
+        "acc_space": "physical",
+    }
+    if std is not None:
+        result["crps"] = gaussian_crps(
+            crps_predictions if crps_predictions is not None else predictions,
+            crps_targets if crps_targets is not None else targets,
+            crps_std if crps_std is not None else std,
+            valid_lengths,
+        ).tolist()
+        result["crps_space"] = "normalized" if crps_predictions is not None else "physical"
+        result["crps_implementation"] = "closed_form_gaussian"
+    return result
+
+
+def save_metrics(metrics: dict, path: str | Path) -> None:
+    output = Path(path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    output.write_text(json.dumps(metrics, indent=2), encoding="utf-8")

scripts/result.py

@@ -0,0 +1,114 @@
+"""Compute metrics and render ClimODE forecast maps from saved outputs."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from scripts.metrics import evaluate, save_metrics
+
+
+def _resolve(path: str | Path) -> Path:
+    value = Path(path)
+    return value if value.is_absolute() else PROJECT_ROOT / value
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "conf/config.yaml")
+    parser.add_argument("--predictions", type=Path, default=None)
+    parser.add_argument("--targets", type=Path, default=None)
+    parser.add_argument("--std", type=Path, default=None)
+    parser.add_argument("--output-dir", type=Path, default=None)
+    parser.add_argument("--stats-dir", type=Path, default=None)
+    parser.add_argument("--static-file", type=Path, default=None)
+    parser.add_argument("--sample", type=int, default=0)
+    parser.add_argument("--lead", type=int, default=0)
+    args = parser.parse_args()
+    with args.config.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    output_dir = args.output_dir or _resolve(config["data"]["output_dir"])
+    predictions = np.load(args.predictions or output_dir / "predictions.npy")
+    targets = np.load(args.targets or output_dir / "targets.npy")
+    std_path = args.std or output_dir / "std.npy"
+    std = np.load(std_path) if std_path.is_file() else None
+    lengths_path = output_dir / "valid_lengths.npy"
+    valid_lengths = np.load(lengths_path) if lengths_path.is_file() else None
+    static_file = _resolve(args.static_file or config["data"]["static_file"])
+    import h5py
+
+    with h5py.File(static_file, "r") as handle:
+        lat2d = handle["lat2d"][:]
+    stats_dir = _resolve(
+        args.stats_dir
+        or config["data"].get("stats_dir", Path(config["data"]["data_dir"]) / "static")
+    )
+    minimum = np.load(stats_dir / "min_values.npy").reshape(1, 1, 1, 5, 1, 1)
+    maximum = np.load(stats_dir / "max_values.npy").reshape(1, 1, 1, 5, 1, 1)
+    scale = maximum - minimum
+    metrics = evaluate(
+        predictions * scale + minimum,
+        targets * scale + minimum,
+        lat2d,
+        std * scale if std is not None else None,
+        crps_predictions=predictions if std is not None else None,
+        crps_targets=targets if std is not None else None,
+        crps_std=std,
+        valid_lengths=valid_lengths,
+    )
+    metrics_path = output_dir.parent / "metrics.json"
+    save_metrics(metrics, metrics_path)
+
+    figure_dir = output_dir / "figures"
+    figure_dir.mkdir(parents=True, exist_ok=True)
+    try:
+        import matplotlib
+
+        matplotlib.use("Agg")
+        import matplotlib.pyplot as plt
+    except ImportError as exc:
+        raise RuntimeError("Visualization requires matplotlib in the active environment") from exc
+
+    if not 0 <= args.sample < predictions.shape[0]:
+        raise IndexError(f"sample must be in [0,{predictions.shape[0] - 1}]")
+    if not 0 <= args.lead < predictions.shape[1]:
+        raise IndexError(f"lead must be in [0,{predictions.shape[1] - 1}]")
+    if valid_lengths is not None and args.lead >= int(valid_lengths[args.sample]):
+        raise IndexError(
+            f"lead {args.lead} is padding for sample {args.sample}; "
+            f"valid length is {int(valid_lengths[args.sample])}"
+        )
+    names = ["z", "t", "t2m", "u10", "v10"]
+    for channel, name in enumerate(names):
+        prediction = predictions[args.sample, args.lead, 0, channel]
+        target = targets[args.sample, args.lead, 0, channel]
+        difference = prediction - target
+        figure, axes = plt.subplots(1, 3, figsize=(12, 3.4), constrained_layout=True)
+        for axis, image, title in zip(
+            axes,
+            (prediction, target, difference),
+            ("prediction", "target", "difference"),
+        ):
+            cmap = "RdBu_r" if title == "difference" else "viridis"
+            plot = axis.imshow(image, cmap=cmap, origin="upper", aspect="auto")
+            axis.set_title(title)
+            axis.set_xlabel("longitude index")
+            axis.set_ylabel("latitude index")
+            figure.colorbar(plot, ax=axis, shrink=0.8)
+        figure.suptitle(f"ClimODE {name}, lead={(args.lead + 1) * 6} h")
+        figure.savefig(figure_dir / f"{name}_lead_{(args.lead + 1) * 6:03d}h.png", dpi=150)
+        plt.close(figure)
+    print(json.dumps({"metrics": str(metrics_path), "figures": str(figure_dir)}))
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,403 @@
+"""Train or fine-tune ClimODE with OneScience ERA5 data."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import yaml
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+# Allow ``python scripts/train.py`` to resolve project-local packages.
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from model.climode import ClimODE, load_checkpoint
+from scripts.data_loader import ClimODESeriesDataset, load_constants
+from scripts.velocity import fit_velocity_cache, load_velocity_cache
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def _device(value: str | None) -> torch.device:
+    if value:
+        return torch.device(value)
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def _init_distributed(backend: str) -> tuple[bool, int, int]:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    if world_size == 1:
+        return False, 0, 1
+    if not dist.is_initialized():
+        dist.init_process_group(backend=backend)
+    return True, dist.get_rank(), world_size
+
+
+def _nll(mean: torch.Tensor, std: torch.Tensor, truth: torch.Tensor, var_coeff: float) -> torch.Tensor:
+    distribution = torch.distributions.Normal(mean, 1.0e-3 + std)
+    return (-distribution.log_prob(truth)).mean() + var_coeff * (std.square()).sum()
+
+
+def _load_yaml(path: Path) -> dict:
+    with path.open("r", encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def _resolve(path: str | Path) -> Path:
+    value = Path(path)
+    return value if value.is_absolute() else PROJECT_ROOT / value
+
+
+def _parse_years(value: str | None, fallback: list[int]) -> list[int]:
+    if value is None:
+        return list(fallback)
+    years = [int(item.strip()) for item in value.split(",") if item.strip()]
+    if not years:
+        raise ValueError("year override must contain at least one integer")
+    return years
+
+
+def _model_from_args(config: dict, args: argparse.Namespace, device: torch.device) -> nn.Module:
+    model_cfg = config["model"]
+    use_pretrained = bool(getattr(args, "use_pretrained", False))
+    pretrained_checkpoint = getattr(args, "pretrained_checkpoint", None)
+    if args.mode == "resume":
+        checkpoint = args.checkpoint or _resolve(model_cfg["default_checkpoint"])
+        model = load_checkpoint(checkpoint, map_location="cpu")
+    elif args.mode == "finetune" or use_pretrained:
+        checkpoint = args.checkpoint
+        if checkpoint is None and (use_pretrained or pretrained_checkpoint is not None):
+            checkpoint = pretrained_checkpoint or model_cfg.get("pretrained_checkpoint")
+        if checkpoint is None:
+            raise ValueError(
+                "finetune requires --checkpoint, or explicitly pass "
+                "--use-pretrained [--pretrained-checkpoint PATH]"
+            )
+        checkpoint = _resolve(checkpoint)
+        if not checkpoint.is_file():
+            raise FileNotFoundError(f"Checkpoint not found: {checkpoint}")
+        model = load_checkpoint(checkpoint, map_location="cpu")
+    else:
+        model = ClimODE(
+            num_channels=5,
+            const_channels=2,
+            out_types=5,
+            method=args.solver or model_cfg.get("solver", "euler"),
+            use_attention=model_cfg.get("use_attention", True),
+            use_uncertainty=model_cfg.get("use_uncertainty", True),
+            use_positional_encoder=model_cfg.get("use_positional_encoder", False),
+        )
+    return model.to(device)
+
+
+def _run_epoch(
+    model,
+    loader,
+    velocity,
+    constants,
+    lat,
+    lon,
+    device,
+    optimizer,
+    var_coeff,
+    max_batches,
+    atol,
+    rtol,
+):
+    training = optimizer is not None
+    model.train(training)
+    total = 0.0
+    count = 0
+    for batch_index, batch in enumerate(loader):
+        if max_batches is not None and batch_index >= max_batches:
+            break
+        observations = batch["observations"].squeeze(0).to(device)
+        time_steps = batch["time_steps"].squeeze(0).to(device)
+        sequence_index = int(batch["sequence_index"].item())
+        past_velocity = velocity[sequence_index].to(device)
+        target = observations
+        initial = observations[0].unsqueeze(1)
+        model_core = model.module if isinstance(model, DistributedDataParallel) else model
+        model_core.update_param([past_velocity, constants, lat, lon])
+        if training:
+            optimizer.zero_grad(set_to_none=True)
+        with torch.set_grad_enabled(training):
+            mean, std, _ = model(time_steps, initial, atol=atol, rtol=rtol)
+            loss = _nll(mean, std, target, var_coeff)
+            loss = loss + 0.001 * sum(parameter.square().sum() for parameter in model.parameters())
+            if training:
+                loss.backward()
+                optimizer.step()
+        total += float(loss.detach())
+        count += 1
+    if dist.is_initialized():
+        totals = torch.tensor([total, float(count)], dtype=torch.float64, device=device)
+        dist.all_reduce(totals, op=dist.ReduceOp.SUM)
+        total, count = float(totals[0].item()), int(totals[1].item())
+    return total / max(count, 1), count
+
+
+def _prepare_velocity(
+    dataset,
+    constants: torch.Tensor,
+    lat: torch.Tensor,
+    lon: torch.Tensor,
+    path: Path,
+    epochs: int,
+    learning_rate: float,
+    smoothing_alpha: float,
+    kernel_sigma: float,
+    distributed: bool,
+    rank: int,
+) -> torch.Tensor:
+    """Build a split cache once, then let every DDP rank read the same result."""
+
+    if path.is_file():
+        return load_velocity_cache(path, len(dataset))
+    if distributed:
+        if rank == 0:
+            fit_velocity_cache(
+                dataset,
+                constants,
+                lat,
+                lon,
+                path,
+                epochs=epochs,
+                learning_rate=learning_rate,
+                smoothing_alpha=smoothing_alpha,
+                kernel_sigma=kernel_sigma,
+            )
+        dist.barrier()
+        return load_velocity_cache(path, len(dataset))
+    return fit_velocity_cache(
+        dataset,
+        constants,
+        lat,
+        lon,
+        path,
+        epochs=epochs,
+        learning_rate=learning_rate,
+        smoothing_alpha=smoothing_alpha,
+        kernel_sigma=kernel_sigma,
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--config", type=Path, default=PROJECT_ROOT / "conf/config.yaml"
+    )
+    parser.add_argument("--mode", choices=["scratch", "finetune", "resume"], default=None)
+    parser.add_argument("--checkpoint", type=Path, default=None)
+    parser.add_argument(
+        "--use-pretrained",
+        action="store_true",
+        help="Explicitly initialize from the official pretrained checkpoint",
+    )
+    parser.add_argument(
+        "--pretrained-checkpoint",
+        type=Path,
+        default=None,
+        help="Override model.pretrained_checkpoint when --use-pretrained is set",
+    )
+    parser.add_argument("--solver", choices=["euler", "rk4", "dopri5", "dopri8", "midpoint"], default=None)
+    parser.add_argument("--epochs", type=int, default=None)
+    parser.add_argument("--sequence-length", type=int, default=None)
+    parser.add_argument("--velocity-epochs", type=int, default=None)
+    parser.add_argument("--velocity-cache", type=Path, default=None)
+    parser.add_argument("--data-dir", type=Path, default=None, help="Override data.data_dir")
+    parser.add_argument("--stats-dir", type=Path, default=None, help="Override data.stats_dir")
+    parser.add_argument("--static-file", type=Path, default=None, help="Override data.static_file")
+    parser.add_argument("--checkpoint-dir", type=Path, default=None)
+    parser.add_argument("--log-file", type=Path, default=None)
+    parser.add_argument("--device", type=str, default=None)
+    parser.add_argument("--max-batches", type=int, default=None)
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--train-years", type=str, default=None, help="Comma-separated year override")
+    parser.add_argument("--val-years", type=str, default=None, help="Comma-separated year override")
+    args = parser.parse_args()
+    args.config = _resolve(args.config)
+    args.checkpoint = _resolve(args.checkpoint) if args.checkpoint is not None else None
+    args.pretrained_checkpoint = (
+        _resolve(args.pretrained_checkpoint)
+        if args.pretrained_checkpoint is not None
+        else None
+    )
+    args.data_dir = _resolve(args.data_dir) if args.data_dir is not None else None
+    args.stats_dir = _resolve(args.stats_dir) if args.stats_dir is not None else None
+    args.static_file = _resolve(args.static_file) if args.static_file is not None else None
+    args.velocity_cache = (
+        _resolve(args.velocity_cache) if args.velocity_cache is not None else None
+    )
+    args.checkpoint_dir = (
+        _resolve(args.checkpoint_dir) if args.checkpoint_dir is not None else None
+    )
+    args.log_file = _resolve(args.log_file) if args.log_file is not None else None
+    config = _load_yaml(args.config)
+    model_cfg, data_cfg, vel_cfg, train_cfg = config["model"], config["data"], config["velocity"], config["training"]
+    args.mode = args.mode or train_cfg.get("mode", "scratch")
+    if args.mode == "resume" and args.use_pretrained:
+        raise ValueError("--use-pretrained cannot be combined with --mode resume")
+    if (
+        args.pretrained_checkpoint is not None
+        and args.mode not in {"finetune"}
+        and not args.use_pretrained
+    ):
+        raise ValueError(
+            "--pretrained-checkpoint requires --use-pretrained or "
+            "--mode finetune"
+        )
+    if args.mode == "scratch" and args.checkpoint is not None:
+        raise ValueError(
+            "--checkpoint is ignored in scratch mode; use --mode resume or "
+            "--mode finetune explicitly"
+        )
+    if args.use_pretrained and args.mode == "scratch":
+        args.mode = "finetune"
+    args.solver = args.solver or model_cfg.get("solver", "euler")
+    args.sequence_length = args.sequence_length or data_cfg.get("sequence_length", 8)
+    args.velocity_epochs = args.velocity_epochs if args.velocity_epochs is not None else vel_cfg.get("epochs", 200)
+    args.max_batches = args.max_batches if args.max_batches is not None else train_cfg.get("max_batches")
+    set_seed(args.seed if args.seed is not None else train_cfg.get("seed", 42))
+    distributed, rank, world_size = _init_distributed(train_cfg.get("ddp_backend", "nccl"))
+    device = _device(args.device)
+    if distributed and device.type == "cuda":
+        device = torch.device("cuda", int(os.environ.get("LOCAL_RANK", "0")))
+    if device.type == "cuda":
+        if device.index is None:
+            device = torch.device("cuda", 0)
+        torch.cuda.set_device(device)
+
+    root = _resolve(args.data_dir or data_cfg["data_dir"])
+    stats_dir = _resolve(args.stats_dir or data_cfg.get("stats_dir", root / "static"))
+    train_set = ClimODESeriesDataset(
+        root,
+        _parse_years(args.train_years, data_cfg["train_years"]),
+        stats_dir=stats_dir,
+        model_size=(data_cfg["model_height"], data_cfg["model_width"]),
+        sequence_length=args.sequence_length,
+        normalize=data_cfg.get("normalize", True),
+    )
+    val_set = ClimODESeriesDataset(
+        root,
+        _parse_years(args.val_years, data_cfg["val_years"]),
+        stats_dir=stats_dir,
+        model_size=(data_cfg["model_height"], data_cfg["model_width"]),
+        sequence_length=args.sequence_length,
+        normalize=data_cfg.get("normalize", True),
+    )
+    train_sampler = DistributedSampler(train_set, shuffle=True) if distributed else None
+    val_sampler = DistributedSampler(val_set, shuffle=False) if distributed else None
+    train_loader = DataLoader(train_set, batch_size=1, sampler=train_sampler, shuffle=train_sampler is None, num_workers=data_cfg["dataloader"]["num_workers"])
+    val_loader = DataLoader(val_set, batch_size=1, sampler=val_sampler, shuffle=False, num_workers=data_cfg["dataloader"]["num_workers"])
+    static_file = _resolve(args.static_file or data_cfg["static_file"])
+    constants, lat, lon = load_constants(static_file, (data_cfg["model_height"], data_cfg["model_width"]))
+    constants, lat, lon = constants.to(device), lat.unsqueeze(0).to(device), lon.unsqueeze(0).to(device)
+
+    # Relative paths follow the project working directory, matching the
+    # config/checkpoint conventions used by the reference earth projects.
+    velocity_root = _resolve(vel_cfg["cache_dir"])
+    velocity_path = args.velocity_cache or (velocity_root / "train.pt")
+    train_velocity = _prepare_velocity(
+        train_set,
+        constants,
+        lat.squeeze(0).cpu(),
+        lon.squeeze(0).cpu(),
+        velocity_path,
+        epochs=args.velocity_epochs,
+        learning_rate=vel_cfg["learning_rate"],
+        smoothing_alpha=vel_cfg["smoothing_alpha"],
+        kernel_sigma=vel_cfg["kernel_sigma"],
+        distributed=distributed,
+        rank=rank,
+    )
+    val_velocity_path = velocity_path.with_name("val.pt")
+    val_velocity = _prepare_velocity(
+        val_set,
+        constants,
+        lat.squeeze(0).cpu(),
+        lon.squeeze(0).cpu(),
+        val_velocity_path,
+        epochs=args.velocity_epochs,
+        learning_rate=vel_cfg["learning_rate"],
+        smoothing_alpha=vel_cfg["smoothing_alpha"],
+        kernel_sigma=vel_cfg["kernel_sigma"],
+        distributed=distributed,
+        rank=rank,
+    )
+
+    model = _model_from_args(config, args, device)
+    if distributed:
+        model = DistributedDataParallel(model, device_ids=[device.index] if device.type == "cuda" else None)
+    lr = train_cfg.get("finetune_learning_rate", 5.0e-5) if args.mode == "finetune" else model_cfg.get("learning_rate", 5.0e-4)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=model_cfg.get("weight_decay", 1.0e-5))
+    epochs = args.epochs or (train_cfg.get("finetune_epochs", 40) if args.mode == "finetune" else train_cfg.get("epochs", 300))
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs)
+    start_epoch = 0
+    if args.mode == "resume":
+        resume_path = args.checkpoint or _resolve(model_cfg["default_checkpoint"])
+        try:
+            resume = torch.load(resume_path, map_location="cpu", weights_only=True)
+        except TypeError:
+            resume = torch.load(resume_path, map_location="cpu")
+        state = resume.get("model", resume.get("state_dict"))
+        (model.module if isinstance(model, DistributedDataParallel) else model).load_state_dict(state)
+        if "optimizer" in resume:
+            optimizer.load_state_dict(resume["optimizer"])
+        if "scheduler" in resume:
+            scheduler.load_state_dict(resume["scheduler"])
+        start_epoch = int(resume.get("epoch", -1)) + 1
+
+    checkpoint_dir = _resolve(args.checkpoint_dir or model_cfg["checkpoint_dir"])
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    log_path = _resolve(args.log_file or train_cfg.get("log_file", "./result/train.jsonl"))
+    log_path.parent.mkdir(parents=True, exist_ok=True)
+    best_val = float("inf")
+    for epoch in range(start_epoch, epochs):
+        if train_sampler is not None:
+            train_sampler.set_epoch(epoch)
+        var_coeff = 1.0e-3 if epoch == 0 else 2.0 * scheduler.get_last_lr()[0]
+        train_loss, train_count = _run_epoch(
+            model, train_loader, train_velocity, constants, lat, lon, device,
+            optimizer, var_coeff, args.max_batches, model_cfg["atol"], model_cfg["rtol"]
+        )
+        with torch.no_grad():
+            val_loss, val_count = _run_epoch(
+                model, val_loader, val_velocity, constants, lat, lon, device,
+                None, var_coeff, args.max_batches, model_cfg["atol"], model_cfg["rtol"]
+            )
+        scheduler.step()
+        record = {"epoch": epoch, "train_loss": train_loss, "val_loss": val_loss, "train_batches": train_count, "val_batches": val_count, "lr": scheduler.get_last_lr()[0]}
+        if rank == 0:
+            with log_path.open("a", encoding="utf-8") as handle:
+                handle.write(json.dumps(record) + "\n")
+            if val_loss < best_val:
+                best_val = val_loss
+                torch.save({"model": (model.module if isinstance(model, DistributedDataParallel) else model).state_dict(), "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(), "epoch": epoch}, checkpoint_dir / "model_bak.pth")
+            print(json.dumps(record))
+    if distributed:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

scripts/velocity.py

@@ -0,0 +1,136 @@
+"""Initial velocity fitting and cache management for ClimODE training."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from model.climode import OptimVelocity
+
+try:
+    from torchcubicspline import NaturalCubicSpline, natural_cubic_spline_coeffs
+except ImportError:  # pragma: no cover - dependency is optional for tiny smoke tests
+    NaturalCubicSpline = None
+    natural_cubic_spline_coeffs = None
+
+
+def _time_derivative(history: torch.Tensor, interval_hours: float = 6.0) -> torch.Tensor:
+    """Estimate the derivative at the final history point.
+
+    The cubic-spline path is identical to the official implementation.  The
+    finite-difference fallback is only for environments without the optional
+    package and is explicitly reported to the caller.
+    """
+
+    if history.ndim != 5:
+        raise ValueError(f"Expected history [N,3,K,H,W], got {tuple(history.shape)}")
+    if natural_cubic_spline_coeffs is not None:
+        times = torch.arange(3, device=history.device, dtype=history.dtype) * interval_hours
+        values = history.permute(1, 0, 2, 3, 4)
+        coeffs = natural_cubic_spline_coeffs(times, values)
+        spline = NaturalCubicSpline(coeffs)
+        return spline.derivative(times[-1])
+    return (3.0 * history[:, 2] - 4.0 * history[:, 1] + history[:, 0]) / (2.0 * interval_hours)
+
+
+def build_rbf_kernel(
+    lat2d: torch.Tensor,
+    lon2d: torch.Tensor,
+    sigma: float = 1.0,
+) -> torch.Tensor:
+    coords = torch.stack([lat2d.reshape(-1), lon2d.reshape(-1)], dim=1).float()
+    distances = torch.cdist(coords, coords).square()
+    kernel = torch.exp(-distances / (2.0 * sigma * sigma))
+    return torch.linalg.inv(kernel)
+
+
+def optimize_velocity(
+    history: torch.Tensor,
+    current: torch.Tensor,
+    kernel_inv: torch.Tensor,
+    epochs: int = 200,
+    learning_rate: float = 2.0,
+    smoothing_alpha: float = 1.0e-7,
+) -> torch.Tensor:
+    """Fit [N,2K,H,W] velocities using the official penalized objective."""
+
+    if current.ndim != 4:
+        raise ValueError(f"Expected current [N,K,H,W], got {tuple(current.shape)}")
+    num_years, channels, height, width = current.shape
+    model = OptimVelocity(num_years, height, width, channels).to(current.device)
+    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+    delta_u = _time_derivative(history)
+    best_loss = float("inf")
+    best_velocity = None
+    for _ in range(max(int(epochs), 1)):
+        optimizer.zero_grad(set_to_none=True)
+        out, vx, vy = model(current.unsqueeze(1))
+        vx_flat = vx.view(num_years, channels, -1, 1)
+        vy_flat = vy.view(num_years, channels, -1, 1)
+        kernel = kernel_inv.to(current.device).expand(num_years, channels, -1, -1)
+        smooth_x = torch.matmul(torch.matmul(vx_flat.transpose(2, 3), kernel), vx_flat).mean()
+        smooth_y = torch.matmul(torch.matmul(vy_flat.transpose(2, 3), kernel), vy_flat).mean()
+        loss = nn.functional.mse_loss(out.squeeze(1), delta_u) + smoothing_alpha * (smooth_x + smooth_y)
+        loss.backward()
+        optimizer.step()
+        if float(loss.detach()) < best_loss:
+            best_loss = float(loss.detach())
+            best_velocity = torch.cat([vx.detach(), vy.detach()], dim=2).squeeze(1).clone()
+    if best_velocity is None:
+        raise RuntimeError("Velocity optimization produced no result")
+    return best_velocity
+
+
+def fit_velocity_cache(
+    dataset,
+    constants: torch.Tensor,
+    lat2d: torch.Tensor,
+    lon2d: torch.Tensor,
+    output_path: str | Path,
+    epochs: int = 200,
+    learning_rate: float = 2.0,
+    smoothing_alpha: float = 1.0e-7,
+    kernel_sigma: float = 1.0,
+) -> torch.Tensor:
+    del constants  # Kept in the signature to make the training handoff explicit.
+    kernel_inv = build_rbf_kernel(lat2d, lon2d, kernel_sigma)
+    velocities = []
+    for index in range(len(dataset)):
+        item = dataset[index]
+        history = item["history"].float()
+        current = item["observations"][0].float()
+        velocities.append(
+            optimize_velocity(
+                history,
+                current,
+                kernel_inv,
+                epochs=epochs,
+                learning_rate=learning_rate,
+                smoothing_alpha=smoothing_alpha,
+            )
+        )
+    result = torch.stack(velocities)
+    path = Path(output_path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save({"velocity": result, "starts": dataset.starts, "years": dataset.years}, path)
+    return result
+
+
+def load_velocity_cache(path: str | Path, expected_length: int | None = None) -> torch.Tensor:
+    try:
+        checkpoint = torch.load(path, map_location="cpu", weights_only=True)
+    except TypeError:  # PyTorch before weights_only support.
+        checkpoint = torch.load(path, map_location="cpu")
+    velocity = checkpoint["velocity"] if isinstance(checkpoint, dict) else checkpoint
+    if expected_length is not None and len(velocity) != expected_length:
+        raise ValueError(f"Velocity cache length {len(velocity)} != dataset length {expected_length}")
+    return velocity.float()

weight/.gitkeep

@@ -0,0 +1 @@
+placeholder for external ClimODE checkpoints