Upload folder using huggingface_hub

README.md

@@ -0,0 +1,148 @@
+---
+license: mit
+language:
+- en
+tags:
+- OneScience
+- fluid-mechanics
+- flow-field-prediction
+- neural-operator
+- Navier-Stokes
+frameworks: PyTorch
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">FNO</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+FNO (Fourier Neural Operator) is a class of neural operators for parameterized partial differential equations. It directly learns mappings from input functions to solution functions by parameterizing integral kernels in Fourier space. Using the OneScience skill workflow, this project independently reproduces the FNO-2D experiment for predicting the vorticity of two-dimensional incompressible Navier–Stokes flows.
+
+Paper: [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)
+
+# Model Description
+
+This implementation takes 10 consecutive 64 × 64 vorticity fields as input and autoregressively predicts the next 10 frames. It first maps the historical fields and two-dimensional periodic coordinates into a latent space of width 32, then applies four Fourier layers. Each layer retains 12 Fourier modes along each spatial dimension, adds a 1 × 1 local convolution to the spectral convolution, and applies BatchNorm and ReLU. A 32 → 128 → 1 projection head generates the vorticity field for the next time step. Each prediction is appended to the input window for closed-loop rollout inference.
+
+## Use Cases
+
+| Use case | Description |
+| --- | --- |
+| Parameterized PDE operator learning | Learns mappings from PDE parameters, coefficient fields, or initial conditions to solution fields, especially when a PDE must be solved repeatedly for many parameter settings. |
+| Burgers' equation prediction | Predicts future states from initial conditions for the one-dimensional Burgers' equation, demonstrating operator learning for nonlinear evolution equations. |
+| Darcy flow prediction | Predicts steady-state solutions from two-dimensional diffusion or permeability coefficient fields for applications such as porous-media flow and groundwater seepage. |
+| Navier–Stokes flow prediction | Autoregressively predicts the evolution of two-dimensional incompressible flow from historical vorticity fields. |
+
+# Usage
+
+## 1. Using OneCode
+
+Try intelligent, one-click AI4S programming in the OneCode online environment:
+
+[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware requirements**
+
+- A GPU or DCU is recommended.
+- A CPU can be used for import checks and small-scale connectivity tests, but full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version for the current cluster, is recommended.
+
+### Download the Model Package
+
+```bash
+modelscope download --model OneScience/FNO --local_dir ./FNO
+cd FNO
+```
+
+### Set Up the Runtime Environment
+
+**DCU environment**
+
+```bash
+# Activate DTK and Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# Installation with uv is also supported
+pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+**GPU environment**
+
+```bash
+# 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
+# Installation with uv is also supported
+pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+### Training Data
+
+The experiment uses `NavierStokes_V1e-5_N1200_T20.mat`, whose shape is `[N,H,W,T]=[1200,64,64,20]`. For each trajectory, the first 10 frames are inputs and the final 10 frames are prediction targets. The first 1,000 trajectories are used for training and the remaining 200 for testing. There is no separate validation set, and no normalization is applied.
+
+Download the data with:
+
+```bash
+modelscope download --dataset OneScience/fno --local_dir ./data
+```
+
+After downloading, set `data.root` in `config/config.yaml` to the data directory and confirm that `data.file` matches the MAT filename above. The training script strictly validates the field name, shape, dtype, and finite values.
+
+### Training
+
+The default configuration corresponds to the FNO-2D experiment with `ν=1e-5` and `T=20` in the paper. It trains for 500 epochs with a batch size of 20, uses Adam with an initial learning rate of `1e-3`, and halves the learning rate every 100 epochs.
+
+```bash
+python scripts/train.py --config config/config.yaml --device auto
+```
+
+The weights with the best full-trajectory relative L2 error on the training set are saved to `weight/best_model.pth`. The latest complete training state from every epoch is saved to `weight/last_model.pth`, and the training history is written to `results/train_history.json`.
+
+### Trained Weights
+
+`weight/best_model.pth` contains the best weights from the full training run and can be used directly for inference.
+
+### Inference
+
+Before running inference, make sure the configured data path is valid and `weight/best_model.pth` exists. The standard inference run uses the 200 configured test trajectories and a 10-step closed-loop rollout, printing batch progress and the final relative L2 error in real time:
+
+```bash
+python scripts/inference.py --config config/config.yaml
+```
+
+Inference outputs are saved as:
+
+- `results/predictions.npz`: predictions, ground truth, sample indices, and prediction times;
+- `results/metrics.json`: overall and per-step metrics, including comparisons with the reference values from the paper;
+- `results/per_sample_metrics.csv`: per-sample relative L2 errors.
+
+### Evaluation and Visualization
+
+After training and inference, run:
+
+```bash
+python scripts/result.py --config config/config.yaml --sample-index 0
+```
+
+The script recomputes metrics from `predictions.npz`, cross-checks the JSON and CSV outputs, best epoch, and prediction shape, and then generates:
+
+- `results/training_curves.png`: training/test errors and training-objective curves;
+- `results/sample_000_rollout.png`: ground truth, predictions, and absolute errors at `t=11, 16, 20` for a representative sample;
+- `results/run_metadata.json`: configuration, runtime environment, assumptions, file hashes, and quality checks;
+- `results/summary.md`: a summary of experimental results.
+
+# Official OneScience Resources
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+- Paper: [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)
+- The code in this model package is licensed under the MIT License. Use of the model weights is also subject to the licenses and applicable terms of the training data and third-party dependencies.

config/config.yaml

@@ -0,0 +1,99 @@
+paper:
+  title: "Fourier Neural Operator for Parametric Partial Differential Equations"
+  arxiv: "2010.08895"
+  experiment: "FNO-2D Navier-Stokes, nu=1e-5, T=20"
+  viscosity: 1.0e-5
+  reference_relative_l2: 0.1556
+  reference_parameter_count: 414517
+  reference_epoch_seconds_v100: 127.80
+
+data:
+  root: "/public/share/sugonhpcapp01/onestore/onedatasets/FNO_data"
+  file: "NavierStokes_V1e-5_N1200_T20.mat"
+  key: "u"
+  layout: "N,H,W,T"
+  dtype: "float32"
+  expected_shape: [1200, 64, 64, 20]
+  resolution: [64, 64]
+  ntrain: 1000
+  ntest: 200
+  train_start: 0
+  test_start: 1000
+  history: 10
+  horizon: 10
+  recording_interval: 1.0
+  future_times: [11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
+  normalization: "none"
+
+model:
+  name: "FNO2d"
+  input_channels: 10
+  output_channels: 1
+  use_grid: true
+  grid_channels: 2
+  grid_include_endpoint: false
+  width: 32
+  modes1: 12
+  modes2: 12
+  num_layers: 4
+  projection_width: 128
+  activation: "relu"
+  normalization: "batch_norm"
+  block_order: "relu(batch_norm(spectral_plus_pointwise))"
+  fft_norm: "backward"
+  spectral_init: "scaled_uniform_complex"
+
+training:
+  epochs: 500
+  batch_size: 20
+  optimizer: "adam"
+  learning_rate: 0.001
+  weight_decay: 0.0
+  scheduler: "step_lr"
+  scheduler_step_size: 100
+  scheduler_gamma: 0.5
+  seed: 0
+  dtype: "float32"
+  amp: false
+  gradient_clipping: null
+  ema: false
+  distributed: "single"
+  num_workers: 0
+  pin_memory: true
+  deterministic: true
+  relative_l2_epsilon: 1.0e-12
+  train_rollout_steps: 10
+  evaluation_rollout_steps: 10
+  checkpoint_monitor: "train_full_relative_l2"
+  checkpoint_mode: "min"
+  evaluate_test_every_epoch: true
+
+inference:
+  batch_size: 20
+  seed: 0
+  dtype: "float32"
+  rollout_steps: 10
+
+paths:
+  checkpoint: "weight/best_model.pth"
+  results_dir: "results"
+  train_history: "results/train_history.json"
+  predictions: "results/predictions.npz"
+  metrics: "results/metrics.json"
+  per_sample_metrics: "results/per_sample_metrics.csv"
+  training_curves: "results/training_curves.png"
+  rollout_figure: "results/sample_000_rollout.png"
+  run_metadata: "results/run_metadata.json"
+  summary: "results/summary.md"
+
+assumptions:
+  - "The paper does not specify a validation split; the best checkpoint is selected using train full-trajectory relative L2, never test error."
+  - "The paper does not specify batch size or seed; batch_size=20 and seed=0 are explicit engineering assumptions."
+  - "The paper does not define the exact relative-L2 reduction; ratios are computed per sample and then averaged with epsilon=1e-12."
+  - "The paper does not specify the projection hidden width; projection_width=128 is configurable."
+  - "Coordinate-grid input is configurable and enabled; the periodic grid excludes the duplicated endpoint."
+  - "The block ordering is ReLU(BatchNorm(spectral + pointwise)); the paper states ReLU and batch normalization but not their exact order."
+  - "Adam uses weight_decay=0.0 because the paper does not state an additional weight-decay regularizer."
+
+conflicts:
+  - "The paper states d_v=32 but reports 414,517 parameters without enough connection details to reproduce both uniquely. Width 32 takes precedence and the actual parameter count must be reported."

configuration.json

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

models/FNO.py

@@ -0,0 +1,219 @@
+"""Paper-driven Fourier Neural Operator for 2-D Navier--Stokes rollout.
+
+This is an independent implementation of Equations (2), (4), and (5) in
+arXiv:2010.08895.  It does not copy the authors' repository implementation.
+The public interface is channel-last because the ten history frames are the
+input function features; Fourier blocks operate channel-first internally.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+from torch import Tensor, nn
+
+
+class SpectralConv2d(nn.Module):
+    """Truncated 2-D Fourier integral operator on a periodic grid."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        modes1: int,
+        modes2: int,
+        fft_norm: str = "backward",
+    ) -> None:
+        super().__init__()
+        if min(in_channels, out_channels, modes1, modes2) <= 0:
+            raise ValueError("channels and retained Fourier modes must be positive")
+        if fft_norm not in {"backward", "forward", "ortho"}:
+            raise ValueError(f"Unsupported FFT normalization: {fft_norm}")
+
+        self.in_channels = int(in_channels)
+        self.out_channels = int(out_channels)
+        self.modes1 = int(modes1)
+        self.modes2 = int(modes2)
+        self.fft_norm = fft_norm
+
+        # MISSING in the paper: exact complex-weight initialization.  The
+        # scale is explicit and seed-controlled by the caller's torch seed.
+        scale = 1.0 / (self.in_channels * self.out_channels)
+        shape = (self.in_channels, self.out_channels, self.modes1, self.modes2)
+        self.weight_positive = nn.Parameter(
+            scale * torch.complex(torch.rand(shape), torch.rand(shape))
+        )
+        self.weight_negative = nn.Parameter(
+            scale * torch.complex(torch.rand(shape), torch.rand(shape))
+        )
+
+    @staticmethod
+    def _multiply_modes(inputs: Tensor, weights: Tensor) -> Tensor:
+        return torch.einsum("bixy,ioxy->boxy", inputs, weights)
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        if inputs.ndim != 4:
+            raise ValueError(f"SpectralConv2d expects [B,C,H,W], got {inputs.shape}")
+        batch, channels, height, width = inputs.shape
+        if channels != self.in_channels:
+            raise ValueError(
+                f"Expected {self.in_channels} channels, received {channels}"
+            )
+        if 2 * self.modes1 > height:
+            raise ValueError(
+                f"modes1={self.modes1} overlaps positive/negative bands for H={height}"
+            )
+        if self.modes2 > width // 2 + 1:
+            raise ValueError(
+                f"modes2={self.modes2} exceeds rFFT width {width // 2 + 1}"
+            )
+
+        spectrum = torch.fft.rfft2(inputs, norm=self.fft_norm)
+        output_spectrum = torch.zeros(
+            batch,
+            self.out_channels,
+            height,
+            width // 2 + 1,
+            device=inputs.device,
+            dtype=spectrum.dtype,
+        )
+        positive_weight = self.weight_positive.to(dtype=spectrum.dtype)
+        negative_weight = self.weight_negative.to(dtype=spectrum.dtype)
+        output_spectrum[:, :, : self.modes1, : self.modes2] = self._multiply_modes(
+            spectrum[:, :, : self.modes1, : self.modes2], positive_weight
+        )
+        output_spectrum[:, :, -self.modes1 :, : self.modes2] = self._multiply_modes(
+            spectrum[:, :, -self.modes1 :, : self.modes2], negative_weight
+        )
+        return torch.fft.irfft2(
+            output_spectrum, s=(height, width), norm=self.fft_norm
+        )
+
+
+class FourierBlock2d(nn.Module):
+    """One paper Fourier layer with local W, batch norm, and ReLU."""
+
+    def __init__(self, width: int, modes1: int, modes2: int, fft_norm: str) -> None:
+        super().__init__()
+        self.spectral = SpectralConv2d(width, width, modes1, modes2, fft_norm)
+        self.pointwise = nn.Conv2d(width, width, kernel_size=1)
+        self.batch_norm = nn.BatchNorm2d(width)
+        self.activation = nn.ReLU()
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        return self.activation(
+            self.batch_norm(self.spectral(inputs) + self.pointwise(inputs))
+        )
+
+
+class FNO2d(nn.Module):
+    """Four-layer FNO-2D mapping ten vorticity frames to the next frame."""
+
+    def __init__(
+        self,
+        input_channels: int = 10,
+        output_channels: int = 1,
+        width: int = 32,
+        modes1: int = 12,
+        modes2: int = 12,
+        num_layers: int = 4,
+        projection_width: int = 128,
+        use_grid: bool = True,
+        grid_include_endpoint: bool = False,
+        expected_resolution: tuple[int, int] = (64, 64),
+        fft_norm: str = "backward",
+    ) -> None:
+        super().__init__()
+        if num_layers != 4:
+            raise ValueError(
+                f"The paper reproduction requires four Fourier layers, got {num_layers}"
+            )
+        if len(expected_resolution) != 2 or min(expected_resolution) <= 0:
+            raise ValueError("expected_resolution must contain two positive dimensions")
+        if min(input_channels, output_channels, width, projection_width) <= 0:
+            raise ValueError("model widths and channel counts must be positive")
+
+        self.input_channels = int(input_channels)
+        self.output_channels = int(output_channels)
+        self.width = int(width)
+        self.modes1 = int(modes1)
+        self.modes2 = int(modes2)
+        self.num_layers = int(num_layers)
+        self.projection_width = int(projection_width)
+        self.use_grid = bool(use_grid)
+        self.grid_include_endpoint = bool(grid_include_endpoint)
+        self.expected_resolution = tuple(int(value) for value in expected_resolution)
+
+        lifting_channels = self.input_channels + (2 if self.use_grid else 0)
+        self.lifting = nn.Linear(lifting_channels, self.width)
+        self.fourier_blocks = nn.ModuleList(
+            [
+                FourierBlock2d(self.width, self.modes1, self.modes2, fft_norm)
+                for _ in range(self.num_layers)
+            ]
+        )
+        self.projection_hidden = nn.Linear(self.width, self.projection_width)
+        self.projection_activation = nn.ReLU()
+        self.projection_output = nn.Linear(
+            self.projection_width, self.output_channels
+        )
+
+    def _grid(self, batch: int, height: int, width: int, inputs: Tensor) -> Tensor:
+        if self.grid_include_endpoint:
+            x = torch.linspace(0.0, 1.0, height, device=inputs.device, dtype=inputs.dtype)
+            y = torch.linspace(0.0, 1.0, width, device=inputs.device, dtype=inputs.dtype)
+        else:
+            x = torch.arange(height, device=inputs.device, dtype=inputs.dtype) / height
+            y = torch.arange(width, device=inputs.device, dtype=inputs.dtype) / width
+        grid_x, grid_y = torch.meshgrid(x, y, indexing="ij")
+        grid = torch.stack((grid_x, grid_y), dim=-1)
+        return grid.unsqueeze(0).expand(batch, -1, -1, -1)
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        if inputs.ndim != 4:
+            raise ValueError(f"FNO2d expects [B,H,W,T_history], got {inputs.shape}")
+        batch, height, width, features = inputs.shape
+        if features != self.input_channels:
+            raise ValueError(
+                f"Expected {self.input_channels} history channels, received {features}"
+            )
+        if (height, width) != self.expected_resolution:
+            raise ValueError(
+                f"Expected resolution {self.expected_resolution}, received {(height, width)}"
+            )
+        if not inputs.is_floating_point():
+            raise TypeError(f"FNO2d expects floating-point input, got {inputs.dtype}")
+
+        lifted_inputs = inputs
+        if self.use_grid:
+            lifted_inputs = torch.cat(
+                (inputs, self._grid(batch, height, width, inputs)), dim=-1
+            )
+        hidden = self.lifting(lifted_inputs).permute(0, 3, 1, 2).contiguous()
+        for block in self.fourier_blocks:
+            hidden = block(hidden)
+        hidden = hidden.permute(0, 2, 3, 1).contiguous()
+        hidden = self.projection_activation(self.projection_hidden(hidden))
+        return self.projection_output(hidden)
+
+
+def build_model_from_config(config: Mapping[str, Any]) -> FNO2d:
+    """Construct the exact paper-reproduction model from a parsed YAML mapping."""
+    model = config["model"]
+    data = config["data"]
+    resolution = tuple(int(value) for value in data["resolution"])
+    return FNO2d(
+        input_channels=int(model["input_channels"]),
+        output_channels=int(model["output_channels"]),
+        width=int(model["width"]),
+        modes1=int(model["modes1"]),
+        modes2=int(model["modes2"]),
+        num_layers=int(model["num_layers"]),
+        projection_width=int(model["projection_width"]),
+        use_grid=bool(model["use_grid"]),
+        grid_include_endpoint=bool(model.get("grid_include_endpoint", False)),
+        expected_resolution=resolution,
+        fft_norm=str(model.get("fft_norm", "backward")),
+    )

models/__init__.py

@@ -0,0 +1,5 @@
+"""FNO-2D model exports."""
+
+from .FNO import FNO2d, SpectralConv2d, build_model_from_config
+
+__all__ = ["FNO2d", "SpectralConv2d", "build_model_from_config"]

scripts/inference.py

@@ -0,0 +1,392 @@
+#!/usr/bin/env python3
+"""Run strict closed-loop inference for the trained FNO-2D checkpoint."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import hashlib
+import json
+import os
+import sys
+import time
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+if str(Path(__file__).resolve().parent) not in sys.path:
+    sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+from models import build_model_from_config  # noqa: E402
+from train import (  # noqa: E402
+    atomic_write_json,
+    build_datasets,
+    build_loader,
+    data_file_from_config,
+    environment_metadata,
+    load_config,
+    resolve_project_path,
+    seed_everything,
+    select_device,
+    synchronize,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Evaluate a trained FNO-2D checkpoint on the fixed test split."
+    )
+    parser.add_argument(
+        "--config", type=Path, default=PROJECT_ROOT / "config" / "config.yaml"
+    )
+    parser.add_argument("--checkpoint", type=Path, default=None)
+    parser.add_argument("--output-dir", type=Path, default=None)
+    parser.add_argument("--device", default="auto")
+    parser.add_argument("--batch-size", type=int, default=None)
+    parser.add_argument("--max-test-samples", type=int, default=None, help="Smoke only.")
+    parser.add_argument("--rollout-steps", type=int, default=None, help="Smoke only.")
+    return parser.parse_args()
+
+
+def load_checkpoint(path: Path, device: torch.device) -> dict[str, Any]:
+    if not path.is_file():
+        raise FileNotFoundError(f"Checkpoint does not exist: {path}")
+    try:
+        checkpoint = torch.load(path, map_location=device, weights_only=False)
+    except TypeError:
+        checkpoint = torch.load(path, map_location=device)
+    if not isinstance(checkpoint, dict):
+        raise TypeError("Checkpoint root must be a mapping")
+    required = {
+        "model_state_dict",
+        "config",
+        "epoch",
+        "parameter_count",
+        "monitor",
+        "test_selected",
+    }
+    missing = sorted(required.difference(checkpoint))
+    if missing:
+        raise KeyError(f"Checkpoint is missing required keys: {missing}")
+    if checkpoint["test_selected"] is not False:
+        raise ValueError("This reproduction forbids a test-selected checkpoint")
+    if checkpoint["monitor"] != "train_full_relative_l2":
+        raise ValueError(f"Unexpected checkpoint monitor: {checkpoint['monitor']}")
+    return checkpoint
+
+
+def nested_value(mapping: dict[str, Any], dotted_key: str) -> Any:
+    value: Any = mapping
+    for key in dotted_key.split("."):
+        value = value[key]
+    return value
+
+
+def validate_checkpoint_config(
+    current: dict[str, Any], checkpoint_config: dict[str, Any]
+) -> None:
+    keys = (
+        "data.key",
+        "data.layout",
+        "data.dtype",
+        "data.expected_shape",
+        "data.resolution",
+        "data.ntrain",
+        "data.ntest",
+        "data.test_start",
+        "data.history",
+        "data.horizon",
+        "data.normalization",
+        "model.input_channels",
+        "model.output_channels",
+        "model.use_grid",
+        "model.grid_include_endpoint",
+        "model.width",
+        "model.modes1",
+        "model.modes2",
+        "model.num_layers",
+        "model.projection_width",
+        "model.fft_norm",
+        "training.dtype",
+        "training.relative_l2_epsilon",
+    )
+    differences = []
+    for key in keys:
+        current_value = nested_value(current, key)
+        checkpoint_value = nested_value(checkpoint_config, key)
+        if current_value != checkpoint_value:
+            differences.append(f"{key}: current={current_value!r}, checkpoint={checkpoint_value!r}")
+    if differences:
+        raise ValueError("Checkpoint/config mismatch:\n" + "\n".join(differences))
+
+
+def sha256_file(path: Path, chunk_size: int = 1024 * 1024) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while chunk := handle.read(chunk_size):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def compute_metrics(
+    prediction: np.ndarray, target: np.ndarray, epsilon: float
+) -> tuple[np.ndarray, np.ndarray]:
+    if prediction.shape != target.shape:
+        raise ValueError(f"Prediction/target mismatch: {prediction.shape} vs {target.shape}")
+    if prediction.ndim != 4:
+        raise ValueError(f"Expected [N,H,W,T], received {prediction.shape}")
+    if not np.isfinite(prediction).all() or not np.isfinite(target).all():
+        raise FloatingPointError("Prediction or target contains NaN/Inf")
+
+    difference = prediction.astype(np.float64) - target.astype(np.float64)
+    target64 = target.astype(np.float64)
+    full_numerator = np.linalg.norm(difference.reshape(prediction.shape[0], -1), axis=1)
+    full_denominator = np.linalg.norm(target64.reshape(target.shape[0], -1), axis=1)
+    full = full_numerator / (full_denominator + epsilon)
+
+    difference_by_time = np.moveaxis(difference, -1, 1).reshape(
+        prediction.shape[0], prediction.shape[-1], -1
+    )
+    target_by_time = np.moveaxis(target64, -1, 1).reshape(
+        target.shape[0], target.shape[-1], -1
+    )
+    per_lead = np.linalg.norm(difference_by_time, axis=2) / (
+        np.linalg.norm(target_by_time, axis=2) + epsilon
+    )
+    return full, per_lead
+
+
+def atomic_save_npz(path: Path, **arrays: np.ndarray) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_suffix(path.suffix + ".tmp")
+    with temporary.open("wb") as handle:
+        np.savez_compressed(handle, **arrays)
+    os.replace(temporary, path)
+
+
+def atomic_write_csv(
+    path: Path,
+    sample_indices: np.ndarray,
+    full_metrics: np.ndarray,
+    lead_metrics: np.ndarray,
+    time_values: np.ndarray,
+) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_suffix(path.suffix + ".tmp")
+    header = ["sample_index", "relative_l2_full"] + [
+        f"relative_l2_t{int(value)}" for value in time_values
+    ]
+    with temporary.open("w", encoding="utf-8", newline="") as handle:
+        writer = csv.writer(handle)
+        writer.writerow(header)
+        for row, sample_index in enumerate(sample_indices):
+            writer.writerow(
+                [int(sample_index), f"{full_metrics[row]:.17g}"]
+                + [f"{value:.17g}" for value in lead_metrics[row]]
+            )
+    os.replace(temporary, path)
+
+
+def main() -> None:
+    args = parse_args()
+    config = load_config(args.config)
+    inference = config["inference"]
+    training = config["training"]
+    seed = int(inference.get("seed", training["seed"]))
+    seed_everything(seed, bool(training.get("deterministic", True)))
+    device = select_device(args.device)
+    checkpoint_path = (
+        resolve_project_path(config["paths"]["checkpoint"])
+        if args.checkpoint is None
+        else args.checkpoint.expanduser().resolve()
+    )
+    output_dir = (
+        resolve_project_path(config["paths"]["results_dir"])
+        if args.output_dir is None
+        else args.output_dir.expanduser().resolve()
+    )
+    output_dir.mkdir(parents=True, exist_ok=True)
+    checkpoint = load_checkpoint(checkpoint_path, device)
+    validate_checkpoint_config(config, checkpoint["config"])
+
+    checkpoint_run_type = str(checkpoint.get("run_type", "formal"))
+    if checkpoint_run_type == "formal" and (
+        args.max_test_samples is not None
+        or args.rollout_steps is not None
+        or args.output_dir is not None
+        or args.checkpoint is not None
+    ):
+        raise ValueError(
+            "Formal inference uses the exact configured checkpoint, test split, horizon, "
+            "and results path; overrides are only allowed for smoke checkpoints"
+        )
+    horizon = int(config["data"]["horizon"])
+    rollout_steps = horizon if args.rollout_steps is None else int(args.rollout_steps)
+    if not 1 <= rollout_steps <= horizon:
+        raise ValueError(f"rollout_steps must be in [1,{horizon}]")
+
+    _, test_dataset = build_datasets(
+        config,
+        max_train_samples=1,
+        max_test_samples=args.max_test_samples,
+        rollout_steps=rollout_steps,
+    )
+    batch_size = int(
+        inference["batch_size"] if args.batch_size is None else args.batch_size
+    )
+    test_loader = build_loader(
+        test_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=int(training.get("num_workers", 0)),
+        pin_memory=bool(training.get("pin_memory", True)) and device.type == "cuda",
+        seed=seed,
+    )
+
+    # Preserve complex64 spectral parameters while moving the model to device.
+    model = build_model_from_config(config).to(device=device)
+    incompatible = model.load_state_dict(checkpoint["model_state_dict"], strict=True)
+    if incompatible.missing_keys or incompatible.unexpected_keys:
+        raise RuntimeError(f"Strict state load failed: {incompatible}")
+    parameter_count = sum(parameter.numel() for parameter in model.parameters())
+    if parameter_count != int(checkpoint["parameter_count"]):
+        raise ValueError(
+            f"Parameter count mismatch: model={parameter_count}, "
+            f"checkpoint={checkpoint['parameter_count']}"
+        )
+    model.eval()
+
+    predictions: list[np.ndarray] = []
+    targets: list[np.ndarray] = []
+    synchronize(device)
+    started = time.perf_counter()
+    processed = 0
+    with torch.inference_mode():
+        for batch_number, (history, target) in enumerate(test_loader, start=1):
+            history = history.to(device=device, dtype=torch.float32, non_blocking=True)
+            target_device = target.to(device=device, dtype=torch.float32, non_blocking=True)
+            window = history
+            batch_prediction: list[torch.Tensor] = []
+            for step in range(rollout_steps):
+                prediction_step = model(window)
+                if not torch.isfinite(prediction_step).all():
+                    raise FloatingPointError(
+                        f"Non-finite prediction at batch {batch_number}, step {step + 1}"
+                    )
+                batch_prediction.append(prediction_step)
+                window = torch.cat((window[..., 1:], prediction_step), dim=-1)
+            prediction = torch.cat(batch_prediction, dim=-1)
+            predictions.append(prediction.cpu().numpy().astype(np.float32, copy=False))
+            targets.append(target_device.cpu().numpy().astype(np.float32, copy=False))
+            processed += int(history.shape[0])
+            print(
+                f"inference_batch={batch_number:03d}/{len(test_loader):03d} "
+                f"processed={processed}/{len(test_dataset)}",
+                flush=True,
+            )
+    synchronize(device)
+    duration = time.perf_counter() - started
+
+    prediction_array = np.concatenate(predictions, axis=0)
+    target_array = np.concatenate(targets, axis=0)
+    expected_shape = (
+        len(test_dataset),
+        int(config["data"]["resolution"][0]),
+        int(config["data"]["resolution"][1]),
+        rollout_steps,
+    )
+    if prediction_array.shape != expected_shape or target_array.shape != expected_shape:
+        raise ValueError(
+            f"Unexpected inference arrays: prediction={prediction_array.shape}, "
+            f"target={target_array.shape}, expected={expected_shape}"
+        )
+    epsilon = float(training["relative_l2_epsilon"])
+    full_metrics, lead_metrics = compute_metrics(prediction_array, target_array, epsilon)
+    test_start = int(config["data"]["test_start"])
+    sample_indices = np.arange(
+        test_start, test_start + len(test_dataset), dtype=np.int64
+    )
+    configured_times = np.asarray(config["data"]["future_times"], dtype=np.float32)
+    time_values = configured_times[:rollout_steps]
+
+    if args.output_dir is None:
+        predictions_path = resolve_project_path(config["paths"]["predictions"])
+        metrics_path = resolve_project_path(config["paths"]["metrics"])
+        csv_path = resolve_project_path(config["paths"]["per_sample_metrics"])
+    else:
+        predictions_path = output_dir / "predictions.npz"
+        metrics_path = output_dir / "metrics.json"
+        csv_path = output_dir / "per_sample_metrics.csv"
+
+    atomic_save_npz(
+        predictions_path,
+        prediction=prediction_array,
+        target=target_array,
+        sample_indices=sample_indices,
+        time_values=time_values,
+    )
+    atomic_write_csv(csv_path, sample_indices, full_metrics, lead_metrics, time_values)
+
+    paper_metric = float(config["paper"]["reference_relative_l2"])
+    mean_full = float(full_metrics.mean())
+    metrics_payload: dict[str, Any] = {
+        "schema_version": "fno-ns2d-metrics-v1",
+        "created_at": datetime.now(timezone.utc).isoformat(),
+        "run_type": checkpoint_run_type,
+        "config_path": str(args.config.expanduser().resolve()),
+        "checkpoint_path": str(checkpoint_path),
+        "checkpoint_sha256": sha256_file(checkpoint_path),
+        "checkpoint_epoch": int(checkpoint["epoch"]),
+        "checkpoint_monitor": checkpoint["monitor"],
+        "test_selected": bool(checkpoint["test_selected"]),
+        "data_path": str(data_file_from_config(config)),
+        "prediction_path": str(predictions_path),
+        "per_sample_metrics_path": str(csv_path),
+        "sample_count": int(len(test_dataset)),
+        "prediction_shape": list(prediction_array.shape),
+        "sample_indices": {"first": int(sample_indices[0]), "last": int(sample_indices[-1])},
+        "time_values": [float(value) for value in time_values],
+        "metric": {
+            "name": "samplewise_relative_l2",
+            "formula": "||prediction-target||_2/(||target||_2+epsilon), then arithmetic mean over samples",
+            "epsilon": epsilon,
+            "full_trajectory_mean": mean_full,
+            "full_trajectory_std": float(full_metrics.std(ddof=0)),
+            "mean_step_relative_l2": float(lead_metrics.mean()),
+            "per_lead_mean": [float(value) for value in lead_metrics.mean(axis=0)],
+            "per_lead_std": [float(value) for value in lead_metrics.std(axis=0, ddof=0)],
+        },
+        "paper_comparison": {
+            "paper_relative_l2": paper_metric,
+            "signed_difference": mean_full - paper_metric,
+            "absolute_difference": abs(mean_full - paper_metric),
+        },
+        "parameter_count": parameter_count,
+        "paper_parameter_count": int(config["paper"]["reference_parameter_count"]),
+        "parameter_count_difference": parameter_count
+        - int(config["paper"]["reference_parameter_count"]),
+        "runtime": {
+            **environment_metadata(device),
+            "duration_seconds": duration,
+            "batch_size": batch_size,
+        },
+        "assumptions": config.get("assumptions", []),
+    }
+    atomic_write_json(metrics_path, metrics_payload)
+    print(
+        f"inference_complete samples={len(test_dataset)} duration={duration:.3f}s "
+        f"full_relative_l2={mean_full:.8f} paper={paper_metric:.8f} "
+        f"absolute_difference={abs(mean_full-paper_metric):.8f}",
+        flush=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,505 @@
+#!/usr/bin/env python3
+"""Validate real FNO outputs and render paper-comparison figures."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import hashlib
+import json
+import os
+import sys
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
+import numpy as np  # noqa: E402
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+if str(Path(__file__).resolve().parent) not in sys.path:
+    sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+from inference import compute_metrics  # noqa: E402
+from train import atomic_write_json, load_config, resolve_project_path  # noqa: E402
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Validate FNO inference artifacts and generate scientific figures."
+    )
+    parser.add_argument(
+        "--config", type=Path, default=PROJECT_ROOT / "config" / "config.yaml"
+    )
+    parser.add_argument("--output-dir", type=Path, default=None)
+    parser.add_argument(
+        "--sample-index", type=int, default=0, help="Local test-set index to visualize."
+    )
+    return parser.parse_args()
+
+
+def read_json(path: Path) -> dict[str, Any]:
+    if not path.is_file():
+        raise FileNotFoundError(f"Required JSON artifact is missing: {path}")
+    with path.open("r", encoding="utf-8") as handle:
+        payload = json.load(handle)
+    if not isinstance(payload, dict):
+        raise TypeError(f"Expected a JSON mapping in {path}")
+    return payload
+
+
+def sha256_file(path: Path, chunk_size: int = 1024 * 1024) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while chunk := handle.read(chunk_size):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def atomic_save_figure(figure: plt.Figure, path: Path, dpi: int = 300) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_suffix(path.suffix + ".tmp")
+    figure.savefig(temporary, format="png", dpi=dpi, bbox_inches="tight")
+    plt.close(figure)
+    os.replace(temporary, path)
+    if path.stat().st_size == 0:
+        raise RuntimeError(f"Generated an empty figure: {path}")
+
+
+def atomic_write_text(path: Path, content: str) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_suffix(path.suffix + ".tmp")
+    with temporary.open("w", encoding="utf-8") as handle:
+        handle.write(content)
+    os.replace(temporary, path)
+
+
+def verify_csv(
+    path: Path,
+    sample_indices: np.ndarray,
+    full_metrics: np.ndarray,
+    lead_metrics: np.ndarray,
+    time_values: np.ndarray,
+) -> None:
+    if not path.is_file():
+        raise FileNotFoundError(f"Per-sample metrics CSV is missing: {path}")
+    expected_header = ["sample_index", "relative_l2_full"] + [
+        f"relative_l2_t{int(value)}" for value in time_values
+    ]
+    with path.open("r", encoding="utf-8", newline="") as handle:
+        rows = list(csv.reader(handle))
+    if not rows or rows[0] != expected_header:
+        raise ValueError(f"Unexpected CSV header in {path}: {rows[0] if rows else None}")
+    if len(rows) - 1 != len(sample_indices):
+        raise ValueError(f"Expected {len(sample_indices)} CSV rows, found {len(rows)-1}")
+    for row_index, row in enumerate(rows[1:]):
+        if int(row[0]) != int(sample_indices[row_index]):
+            raise ValueError(f"CSV sample order mismatch at row {row_index + 2}")
+        observed = np.asarray([float(value) for value in row[1:]], dtype=np.float64)
+        expected = np.concatenate(
+            ([full_metrics[row_index]], lead_metrics[row_index].astype(np.float64))
+        )
+        if not np.allclose(observed, expected, rtol=1e-12, atol=1e-12):
+            raise ValueError(f"CSV metric mismatch for sample {sample_indices[row_index]}")
+
+
+def validate_history(history_payload: dict[str, Any]) -> list[dict[str, Any]]:
+    records = history_payload.get("history")
+    if not isinstance(records, list) or not records:
+        raise ValueError("Training history contains no epoch records")
+    formal = history_payload.get("run_type") == "formal"
+    requested = int(history_payload.get("epochs_requested", len(records)))
+    if formal and (requested != 500 or len(records) != 500):
+        raise ValueError(
+            f"Formal paper reproduction requires 500 epochs, got requested={requested}, "
+            f"records={len(records)}"
+        )
+    required = (
+        "epoch",
+        "learning_rate",
+        "duration_seconds",
+        "train_step_loss_sum",
+        "train_mean_step_relative_l2",
+        "train_full_relative_l2",
+        "test_mean_step_relative_l2",
+        "test_full_relative_l2",
+        "best",
+    )
+    for position, record in enumerate(records, start=1):
+        missing = [key for key in required if key not in record]
+        if missing:
+            raise KeyError(f"Epoch record {position} is missing {missing}")
+        if int(record["epoch"]) != position:
+            raise ValueError(f"Epoch sequence is not contiguous at record {position}")
+        numeric = [float(record[key]) for key in required[1:-1]]
+        if not np.isfinite(numeric).all():
+            raise FloatingPointError(f"Non-finite training history at epoch {position}")
+    return records
+
+
+def make_training_figure(
+    records: list[dict[str, Any]], paper_metric: float, best_epoch: int
+) -> plt.Figure:
+    epochs = np.asarray([record["epoch"] for record in records], dtype=np.int64)
+    train_full = np.asarray(
+        [record["train_full_relative_l2"] for record in records], dtype=np.float64
+    )
+    test_full = np.asarray(
+        [record["test_full_relative_l2"] for record in records], dtype=np.float64
+    )
+    train_loss = np.asarray(
+        [record["train_step_loss_sum"] for record in records], dtype=np.float64
+    )
+    train_step = np.asarray(
+        [record["train_mean_step_relative_l2"] for record in records], dtype=np.float64
+    )
+    test_step = np.asarray(
+        [record["test_mean_step_relative_l2"] for record in records], dtype=np.float64
+    )
+
+    figure, axes = plt.subplots(1, 2, figsize=(12.5, 4.8), constrained_layout=True)
+    left = axes[0]
+    left.plot(epochs, train_full, label="Train full relative L2", linewidth=1.6)
+    left.plot(epochs, test_full, label="Test full relative L2", linewidth=1.6)
+    left.axhline(
+        paper_metric,
+        color="black",
+        linestyle="--",
+        linewidth=1.2,
+        label=f"Paper benchmark ({paper_metric:.4f})",
+    )
+    left.axvline(
+        best_epoch,
+        color="tab:green",
+        linestyle=":",
+        linewidth=1.2,
+        label=f"Best checkpoint epoch ({best_epoch})",
+    )
+    if np.all(train_full > 0) and np.all(test_full > 0):
+        left.set_yscale("log")
+    left.set_xlabel("Epoch")
+    left.set_ylabel("Full-trajectory relative L2")
+    left.set_title("FNO-2D rollout error")
+    left.grid(True, alpha=0.25)
+    left.legend(fontsize=8)
+
+    right = axes[1]
+    loss_line = right.plot(
+        epochs,
+        train_loss,
+        color="tab:blue",
+        label="Train 10-step loss sum",
+        linewidth=1.5,
+    )
+    right.set_xlabel("Epoch")
+    right.set_ylabel("Summed step relative L2", color="tab:blue")
+    right.tick_params(axis="y", labelcolor="tab:blue")
+    right.grid(True, alpha=0.25)
+    diagnostic = right.twinx()
+    train_line = diagnostic.plot(
+        epochs,
+        train_step,
+        color="tab:orange",
+        label="Train mean-step relative L2",
+        linewidth=1.3,
+    )
+    test_line = diagnostic.plot(
+        epochs,
+        test_step,
+        color="tab:red",
+        label="Test mean-step relative L2",
+        linewidth=1.3,
+    )
+    diagnostic.set_ylabel("Mean-step relative L2")
+    right.set_title("Training objective and step diagnostics")
+    lines = loss_line + train_line + test_line
+    right.legend(lines, [line.get_label() for line in lines], fontsize=8, loc="best")
+    return figure
+
+
+def representative_leads(number_of_steps: int) -> list[int]:
+    if number_of_steps <= 0:
+        raise ValueError("At least one rollout step is required")
+    return sorted({0, number_of_steps // 2, number_of_steps - 1})
+
+
+def make_rollout_figure(
+    prediction: np.ndarray,
+    target: np.ndarray,
+    sample_indices: np.ndarray,
+    time_values: np.ndarray,
+    lead_metrics: np.ndarray,
+    local_sample: int,
+) -> plt.Figure:
+    if not 0 <= local_sample < prediction.shape[0]:
+        raise IndexError(
+            f"sample-index {local_sample} is outside [0,{prediction.shape[0] - 1}]"
+        )
+    lead_indices = representative_leads(prediction.shape[-1])
+    figure, axes = plt.subplots(
+        len(lead_indices),
+        3,
+        figsize=(11.5, 3.25 * len(lead_indices)),
+        squeeze=False,
+        constrained_layout=True,
+    )
+    global_sample = int(sample_indices[local_sample])
+    for row, lead_index in enumerate(lead_indices):
+        truth = target[local_sample, :, :, lead_index]
+        estimate = prediction[local_sample, :, :, lead_index]
+        absolute_error = np.abs(estimate - truth)
+        shared_limit = max(float(np.max(np.abs(truth))), float(np.max(np.abs(estimate))), 1e-12)
+        error_limit = max(float(np.max(absolute_error)), 1e-12)
+        time_value = int(time_values[lead_index])
+        relative_error = float(lead_metrics[local_sample, lead_index])
+
+        fields = (truth, estimate, absolute_error)
+        titles = (
+            f"Target vorticity w\nsample={global_sample}, t={time_value}",
+            f"Predicted vorticity w\nrelative L2={relative_error:.5f}",
+            f"Absolute error |prediction-target|\nt={time_value}",
+        )
+        for column, (field, title) in enumerate(zip(fields, titles)):
+            axis = axes[row, column]
+            if column < 2:
+                image = axis.imshow(
+                    field,
+                    origin="lower",
+                    extent=(0.0, 1.0, 0.0, 1.0),
+                    interpolation="nearest",
+                    cmap="RdBu_r",
+                    vmin=-shared_limit,
+                    vmax=shared_limit,
+                )
+                color_label = "Vorticity w (unit not specified)"
+            else:
+                image = axis.imshow(
+                    field,
+                    origin="lower",
+                    extent=(0.0, 1.0, 0.0, 1.0),
+                    interpolation="nearest",
+                    cmap="magma",
+                    vmin=0.0,
+                    vmax=error_limit,
+                )
+                color_label = "Absolute error"
+            axis.set_aspect("equal")
+            axis.set_xlabel("x")
+            axis.set_ylabel("y")
+            axis.set_title(title, fontsize=9)
+            colorbar = figure.colorbar(image, ax=axis, shrink=0.82)
+            colorbar.set_label(color_label, fontsize=8)
+    return figure
+
+
+def main() -> None:
+    args = parse_args()
+    config = load_config(args.config)
+    output_dir = (
+        resolve_project_path(config["paths"]["results_dir"])
+        if args.output_dir is None
+        else args.output_dir.expanduser().resolve()
+    )
+    if args.output_dir is None:
+        history_path = resolve_project_path(config["paths"]["train_history"])
+        predictions_path = resolve_project_path(config["paths"]["predictions"])
+        metrics_path = resolve_project_path(config["paths"]["metrics"])
+        csv_path = resolve_project_path(config["paths"]["per_sample_metrics"])
+        training_figure_path = resolve_project_path(config["paths"]["training_curves"])
+        rollout_figure_path = resolve_project_path(config["paths"]["rollout_figure"])
+        metadata_path = resolve_project_path(config["paths"]["run_metadata"])
+        summary_path = resolve_project_path(config["paths"]["summary"])
+    else:
+        history_path = output_dir / "train_history.json"
+        predictions_path = output_dir / "predictions.npz"
+        metrics_path = output_dir / "metrics.json"
+        csv_path = output_dir / "per_sample_metrics.csv"
+        training_figure_path = output_dir / "training_curves.png"
+        rollout_figure_path = output_dir / "sample_000_rollout.png"
+        metadata_path = output_dir / "run_metadata.json"
+        summary_path = output_dir / "summary.md"
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    history_payload = read_json(history_path)
+    metrics_payload = read_json(metrics_path)
+    if metrics_payload.get("run_type") == "formal" and args.output_dir is not None:
+        raise ValueError("Formal result generation must use the configured results directory")
+    if history_payload.get("run_type") != metrics_payload.get("run_type"):
+        raise ValueError("Training history and inference metrics have different run types")
+    records = validate_history(history_payload)
+    if not predictions_path.is_file():
+        raise FileNotFoundError(f"Predictions artifact is missing: {predictions_path}")
+    with np.load(predictions_path, allow_pickle=False) as archive:
+        required_arrays = {"prediction", "target", "sample_indices", "time_values"}
+        missing_arrays = required_arrays.difference(archive.files)
+        if missing_arrays:
+            raise KeyError(f"Predictions NPZ is missing {sorted(missing_arrays)}")
+        prediction = archive["prediction"]
+        target = archive["target"]
+        sample_indices = archive["sample_indices"]
+        time_values = archive["time_values"]
+
+    if prediction.dtype != np.float32 or target.dtype != np.float32:
+        raise TypeError("Prediction and target arrays must be float32")
+    if prediction.shape != target.shape or prediction.ndim != 4:
+        raise ValueError(f"Invalid prediction/target shapes: {prediction.shape}, {target.shape}")
+    if sample_indices.shape != (prediction.shape[0],):
+        raise ValueError("sample_indices shape does not match predictions")
+    if time_values.shape != (prediction.shape[-1],):
+        raise ValueError("time_values shape does not match rollout horizon")
+    if not np.array_equal(sample_indices, np.arange(sample_indices[0], sample_indices[0] + len(sample_indices))):
+        raise ValueError("sample_indices must be unique, contiguous, and ordered")
+    if not np.all(np.diff(time_values.astype(np.float64)) > 0):
+        raise ValueError("time_values must be strictly increasing")
+
+    formal = metrics_payload.get("run_type") == "formal"
+    if formal:
+        expected_shape = (
+            int(config["data"]["ntest"]),
+            int(config["data"]["resolution"][0]),
+            int(config["data"]["resolution"][1]),
+            int(config["data"]["horizon"]),
+        )
+        if prediction.shape != expected_shape:
+            raise ValueError(f"Formal prediction shape must be {expected_shape}, got {prediction.shape}")
+        expected_indices = np.arange(
+            int(config["data"]["test_start"]),
+            int(config["data"]["test_start"]) + int(config["data"]["ntest"]),
+        )
+        if not np.array_equal(sample_indices, expected_indices):
+            raise ValueError("Formal sample indices do not match the fixed test split")
+
+    epsilon = float(config["training"]["relative_l2_epsilon"])
+    full_metrics, lead_metrics = compute_metrics(prediction, target, epsilon)
+    observed_mean = float(metrics_payload["metric"]["full_trajectory_mean"])
+    if not np.isclose(full_metrics.mean(), observed_mean, rtol=1e-8, atol=1e-8):
+        raise ValueError(
+            f"metrics.json full relative L2 mismatch: recomputed={full_metrics.mean()}, "
+            f"stored={observed_mean}"
+        )
+    stored_leads = np.asarray(metrics_payload["metric"]["per_lead_mean"], dtype=np.float64)
+    if not np.allclose(lead_metrics.mean(axis=0), stored_leads, rtol=1e-8, atol=1e-8):
+        raise ValueError("metrics.json per-lead values do not match predictions")
+    verify_csv(csv_path, sample_indices, full_metrics, lead_metrics, time_values)
+
+    best_epoch = int(history_payload["best_epoch"])
+    checkpoint_epoch = int(metrics_payload["checkpoint_epoch"])
+    if best_epoch != checkpoint_epoch:
+        raise ValueError(
+            f"History best epoch {best_epoch} does not match checkpoint epoch {checkpoint_epoch}"
+        )
+    paper_metric = float(config["paper"]["reference_relative_l2"])
+    training_figure = make_training_figure(records, paper_metric, best_epoch)
+    atomic_save_figure(training_figure, training_figure_path, dpi=300)
+    rollout_figure = make_rollout_figure(
+        prediction,
+        target,
+        sample_indices,
+        time_values,
+        lead_metrics,
+        args.sample_index,
+    )
+    atomic_save_figure(rollout_figure, rollout_figure_path, dpi=300)
+
+    artifact_paths = {
+        "train_history": history_path,
+        "predictions": predictions_path,
+        "metrics": metrics_path,
+        "per_sample_metrics": csv_path,
+        "training_curves": training_figure_path,
+        "rollout_figure": rollout_figure_path,
+    }
+    artifact_metadata = {
+        name: {
+            "path": str(path),
+            "size_bytes": path.stat().st_size,
+            "sha256": sha256_file(path),
+        }
+        for name, path in artifact_paths.items()
+    }
+    main_metric = float(full_metrics.mean())
+    run_metadata = {
+        "schema_version": "fno-ns2d-run-metadata-v1",
+        "created_at": datetime.now(timezone.utc).isoformat(),
+        "run_type": metrics_payload.get("run_type"),
+        "config_path": str(args.config.expanduser().resolve()),
+        "data_path": metrics_payload["data_path"],
+        "checkpoint_path": metrics_payload["checkpoint_path"],
+        "checkpoint_epoch": checkpoint_epoch,
+        "test_selected": metrics_payload["test_selected"],
+        "split": {
+            "train": int(history_payload.get("train_samples", 1000)),
+            "validation": 0,
+            "test": int(prediction.shape[0]),
+        },
+        "prediction_shape": list(prediction.shape),
+        "normalization": config["data"]["normalization"],
+        "metric_formula": metrics_payload["metric"]["formula"],
+        "full_trajectory_relative_l2": main_metric,
+        "paper_relative_l2": paper_metric,
+        "signed_difference": main_metric - paper_metric,
+        "absolute_difference": abs(main_metric - paper_metric),
+        "parameter_count": metrics_payload["parameter_count"],
+        "paper_parameter_count": metrics_payload["paper_parameter_count"],
+        "parameter_count_difference": metrics_payload["parameter_count_difference"],
+        "runtime": {**metrics_payload["runtime"], "matplotlib": matplotlib.__version__},
+        "assumptions": config.get("assumptions", []),
+        "conflicts": config.get("conflicts", []),
+        "artifacts": artifact_metadata,
+        "quality_checks": {
+            "all_values_finite": True,
+            "metrics_recomputed_from_npz": True,
+            "json_metrics_match": True,
+            "csv_metrics_match": True,
+            "best_epoch_matches_checkpoint": True,
+            "figures_nonempty": True,
+        },
+    }
+    atomic_write_json(metadata_path, run_metadata)
+
+    summary = f"""# FNO-2D Navier–Stokes reproduction result
+
+## Summary
+
+- Run type: `{metrics_payload.get('run_type')}`
+- Test trajectories: {prediction.shape[0]}
+- Forecast shape: `{list(prediction.shape)}`
+- Mean full-trajectory relative L2: **{main_metric:.8f}**
+- Paper FNO-2D reference (`ν=1e-5`, `T=20`, 1000 train): **{paper_metric:.4f}**
+- Absolute difference: **{abs(main_metric-paper_metric):.8f}**
+- Checkpoint epoch: {checkpoint_epoch}; selected by train full relative L2 (`test_selected=false`).
+
+## Data and method
+
+The model uses the fixed first 1000 trajectories for training and the final {prediction.shape[0]} trajectories for testing. Ten observed vorticity frames initialize a closed-loop rollout; every predicted frame updates the next input window. No target frame is used after initialization, and no data normalization, padding, augmentation, or PDE-residual loss is applied.
+
+The reported metric is computed per sample as `||prediction-target||₂/(||target||₂+1e-12)` over the full space-time forecast and then averaged. It was recomputed directly from `predictions.npz` and cross-checked against JSON and CSV outputs.
+
+## Reproducibility limitations
+
+The paper does not specify the exact relative-L2 reduction, batch size, random seed, projection hidden width, coordinate-input choice, block ordering, or checkpoint-selection protocol. These choices are recorded explicitly in `config/config.yaml` and `run_metadata.json`. The paper width 32 also cannot be uniquely reconciled with the reported 414,517 parameters from the published connection details; the actual parameter count is reported rather than hidden.
+
+## Artifacts
+
+- `{training_figure_path.name}`: train/test rollout errors and step-loss diagnostics.
+- `{rollout_figure_path.name}`: target, prediction, and absolute-error vorticity fields.
+- `{predictions_path.name}`: full test prediction and target arrays.
+- `{metrics_path.name}` and `{csv_path.name}`: aggregate and per-sample metrics.
+- `{metadata_path.name}`: provenance, software, assumptions, hashes, and quality checks.
+"""
+    atomic_write_text(summary_path, summary)
+    print(
+        f"result_complete full_relative_l2={main_metric:.8f} "
+        f"training_figure={training_figure_path} rollout_figure={rollout_figure_path}",
+        flush=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,763 @@
+#!/usr/bin/env python3
+"""Train the paper-specified recurrent FNO-2D Navier--Stokes model."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import platform
+import random
+import sys
+import time
+from copy import deepcopy
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import scipy
+import scipy.io
+import torch
+import yaml
+from torch import Tensor, nn
+from torch.utils.data import DataLoader, TensorDataset
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from models import build_model_from_config  # noqa: E402
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Train FNO-2D on the validated Navier-Stokes trajectory MAT file."
+    )
+    parser.add_argument(
+        "--config",
+        type=Path,
+        default=PROJECT_ROOT / "config" / "config.yaml",
+        help="YAML configuration path.",
+    )
+    parser.add_argument("--epochs", type=int, default=None, help="Override epochs (smoke only).")
+    parser.add_argument(
+        "--rollout-steps", type=int, default=None, help="Override rollout steps (smoke only)."
+    )
+    parser.add_argument(
+        "--max-train-samples", type=int, default=None, help="Limit train samples (smoke only)."
+    )
+    parser.add_argument(
+        "--max-test-samples", type=int, default=None, help="Limit test samples (smoke only)."
+    )
+    parser.add_argument("--batch-size", type=int, default=None, help="Override batch size.")
+    parser.add_argument(
+        "--device", default="auto", help="auto, cpu, cuda, or an explicit torch device."
+    )
+    parser.add_argument("--checkpoint", type=Path, default=None, help="Output checkpoint path.")
+    parser.add_argument(
+        "--resume",
+        type=Path,
+        default=None,
+        help="Resume from a formal training-state checkpoint.",
+    )
+    parser.add_argument("--output-dir", type=Path, default=None, help="Output directory.")
+    parser.add_argument(
+        "--run-type", choices=("formal", "smoke"), default="formal", help="Run provenance tag."
+    )
+    return parser.parse_args()
+
+
+def load_config(path: Path) -> dict[str, Any]:
+    config_path = path.expanduser().resolve()
+    if not config_path.is_file():
+        raise FileNotFoundError(f"Configuration file does not exist: {config_path}")
+    with config_path.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    if not isinstance(config, dict):
+        raise ValueError("Configuration root must be a mapping")
+    validate_config(config)
+    return config
+
+
+def validate_config(config: dict[str, Any]) -> None:
+    for section in ("paper", "data", "model", "training", "paths"):
+        if section not in config or not isinstance(config[section], dict):
+            raise ValueError(f"Missing configuration section: {section}")
+
+    data = config["data"]
+    model = config["model"]
+    training = config["training"]
+    required_data = {
+        "root",
+        "file",
+        "key",
+        "layout",
+        "dtype",
+        "expected_shape",
+        "resolution",
+        "ntrain",
+        "ntest",
+        "test_start",
+        "history",
+        "horizon",
+        "normalization",
+    }
+    required_model = {
+        "input_channels",
+        "output_channels",
+        "width",
+        "modes1",
+        "modes2",
+        "num_layers",
+        "projection_width",
+        "use_grid",
+    }
+    required_training = {
+        "epochs",
+        "batch_size",
+        "optimizer",
+        "learning_rate",
+        "weight_decay",
+        "scheduler",
+        "scheduler_step_size",
+        "scheduler_gamma",
+        "seed",
+        "dtype",
+        "relative_l2_epsilon",
+        "checkpoint_monitor",
+    }
+    for label, mapping, required in (
+        ("data", data, required_data),
+        ("model", model, required_model),
+        ("training", training, required_training),
+    ):
+        missing = sorted(required.difference(mapping))
+        if missing:
+            raise ValueError(f"Missing {label} configuration keys: {missing}")
+
+    if data["layout"] != "N,H,W,T":
+        raise ValueError("This reproduction requires data.layout=N,H,W,T")
+    if data["dtype"] != "float32" or training["dtype"] != "float32":
+        raise ValueError("The audited reproduction requires float32 data and training")
+    if data["normalization"] != "none":
+        raise ValueError("The paper-faithful default requires normalization=none")
+    if int(data["history"]) != 10 or int(model["input_channels"]) != 10:
+        raise ValueError("The FNO-2D experiment requires ten history channels")
+    if int(data["horizon"]) != 10 or int(model["output_channels"]) != 1:
+        raise ValueError("The experiment requires a ten-step rollout of one-step outputs")
+    if int(model["width"]) != 32 or int(model["num_layers"]) != 4:
+        raise ValueError("Strict paper settings require width=32 and num_layers=4")
+    if int(model["modes1"]) != 12 or int(model["modes2"]) != 12:
+        raise ValueError("Strict paper settings require 12 retained modes per axis")
+    if str(training["optimizer"]).lower() != "adam":
+        raise ValueError("The paper specifies Adam")
+    if str(training["scheduler"]).lower() != "step_lr":
+        raise ValueError("The paper schedule is represented by StepLR")
+    if str(training["checkpoint_monitor"]) != "train_full_relative_l2":
+        raise ValueError("Test metrics must not select the checkpoint")
+
+
+def resolve_project_path(value: str | Path) -> Path:
+    path = Path(value).expanduser()
+    return path.resolve() if path.is_absolute() else (PROJECT_ROOT / path).resolve()
+
+
+def data_file_from_config(config: dict[str, Any]) -> Path:
+    data = config["data"]
+    path = Path(str(data["root"])).expanduser() / str(data["file"])
+    path = path.resolve()
+    if not path.is_file():
+        raise FileNotFoundError(f"Navier-Stokes MAT file does not exist: {path}")
+    return path
+
+
+def load_trajectory_array(config: dict[str, Any]) -> np.ndarray:
+    data = config["data"]
+    path = data_file_from_config(config)
+    key = str(data["key"])
+    payload = scipy.io.loadmat(path, variable_names=[key])
+    if key not in payload:
+        raise KeyError(f"MAT field {key!r} is missing from {path}")
+    trajectories = payload[key]
+    expected_shape = tuple(int(value) for value in data["expected_shape"])
+    if trajectories.shape != expected_shape:
+        raise ValueError(
+            f"Expected {key} shape {expected_shape}, received {trajectories.shape}"
+        )
+    if trajectories.dtype != np.float32:
+        raise TypeError(f"Expected {key} dtype float32, received {trajectories.dtype}")
+    if not np.isfinite(trajectories).all():
+        raise ValueError("Trajectory array contains NaN or Inf")
+    return trajectories
+
+
+def build_datasets(
+    config: dict[str, Any],
+    max_train_samples: int | None = None,
+    max_test_samples: int | None = None,
+    rollout_steps: int | None = None,
+) -> tuple[TensorDataset, TensorDataset]:
+    data = config["data"]
+    trajectories = load_trajectory_array(config)
+    history = int(data["history"])
+    horizon = int(data["horizon"])
+    steps = horizon if rollout_steps is None else int(rollout_steps)
+    if not 1 <= steps <= horizon:
+        raise ValueError(f"rollout_steps must be in [1,{horizon}], got {steps}")
+
+    ntrain = int(data["ntrain"])
+    ntest = int(data["ntest"])
+    train_start = int(data.get("train_start", 0))
+    test_start = int(data["test_start"])
+    train_count = ntrain if max_train_samples is None else min(ntrain, max_train_samples)
+    test_count = ntest if max_test_samples is None else min(ntest, max_test_samples)
+    if train_count <= 0 or test_count <= 0:
+        raise ValueError("Training and test sample counts must be positive")
+
+    train = trajectories[train_start : train_start + train_count]
+    test = trajectories[test_start : test_start + test_count]
+    train_history = torch.from_numpy(train[..., :history])
+    train_target = torch.from_numpy(train[..., history : history + steps])
+    test_history = torch.from_numpy(test[..., :history])
+    test_target = torch.from_numpy(test[..., history : history + steps])
+    expected_hw = tuple(int(value) for value in data["resolution"])
+    expected_history_shape = (expected_hw[0], expected_hw[1], history)
+    if tuple(train_history.shape[1:]) != expected_history_shape:
+        raise ValueError(f"Invalid train history shape: {train_history.shape}")
+    if tuple(test_history.shape[1:]) != expected_history_shape:
+        raise ValueError(f"Invalid test history shape: {test_history.shape}")
+    return TensorDataset(train_history, train_target), TensorDataset(test_history, test_target)
+
+
+def seed_everything(seed: int, deterministic: bool) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    if hasattr(torch.backends, "cudnn"):
+        torch.backends.cudnn.deterministic = deterministic
+        torch.backends.cudnn.benchmark = not deterministic
+    if deterministic:
+        torch.use_deterministic_algorithms(True, warn_only=True)
+
+
+def select_device(requested: str) -> torch.device:
+    if requested == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = torch.device(requested)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError(f"Requested {requested}, but torch reports no CUDA/DCU device")
+    return device
+
+
+def relative_l2_per_sample(prediction: Tensor, target: Tensor, epsilon: float) -> Tensor:
+    if prediction.shape != target.shape:
+        raise ValueError(f"Relative L2 shape mismatch: {prediction.shape} vs {target.shape}")
+    if prediction.ndim < 2:
+        raise ValueError("Relative L2 inputs must include a batch and at least one feature axis")
+    difference = torch.linalg.vector_norm(
+        (prediction - target).reshape(prediction.shape[0], -1), dim=1
+    )
+    denominator = torch.linalg.vector_norm(target.reshape(target.shape[0], -1), dim=1)
+    return difference / (denominator + epsilon)
+
+
+def autoregressive_rollout(
+    model: nn.Module,
+    history: Tensor,
+    target: Tensor,
+    epsilon: float,
+) -> tuple[Tensor, Tensor, Tensor]:
+    if history.ndim != 4 or target.ndim != 4:
+        raise ValueError("history and target must be channel-last four-dimensional tensors")
+    window = history
+    predictions: list[Tensor] = []
+    step_ratios: list[Tensor] = []
+    for step in range(target.shape[-1]):
+        prediction = model(window)
+        expected_step_shape = (*history.shape[:-1], 1)
+        if tuple(prediction.shape) != expected_step_shape:
+            raise ValueError(
+                f"Model returned {tuple(prediction.shape)}, expected {expected_step_shape}"
+            )
+        target_step = target[..., step : step + 1]
+        predictions.append(prediction)
+        step_ratios.append(relative_l2_per_sample(prediction, target_step, epsilon))
+        window = torch.cat((window[..., 1:], prediction), dim=-1)
+    rollout = torch.cat(predictions, dim=-1)
+    ratios = torch.stack(step_ratios, dim=1)
+    backward_loss = ratios.mean(dim=0).sum()
+    return rollout, ratios, backward_loss
+
+
+def build_loader(
+    dataset: TensorDataset,
+    batch_size: int,
+    shuffle: bool,
+    num_workers: int,
+    pin_memory: bool,
+    seed: int,
+) -> DataLoader:
+    generator = torch.Generator()
+    generator.manual_seed(seed)
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        drop_last=False,
+        generator=generator,
+    )
+
+
+def run_epoch(
+    model: nn.Module,
+    loader: DataLoader,
+    device: torch.device,
+    epsilon: float,
+    optimizer: torch.optim.Optimizer | None,
+) -> dict[str, float]:
+    training = optimizer is not None
+    model.train(training)
+    total_samples = 0
+    total_step_ratio = 0.0
+    total_full_ratio = 0.0
+    steps_per_sample: int | None = None
+
+    context = torch.enable_grad() if training else torch.inference_mode()
+    with context:
+        for history, target in loader:
+            history = history.to(device=device, dtype=torch.float32, non_blocking=True)
+            target = target.to(device=device, dtype=torch.float32, non_blocking=True)
+            if training:
+                optimizer.zero_grad(set_to_none=True)
+            prediction, step_ratios, backward_loss = autoregressive_rollout(
+                model, history, target, epsilon
+            )
+            full_ratios = relative_l2_per_sample(prediction, target, epsilon)
+            if not torch.isfinite(backward_loss) or not torch.isfinite(full_ratios).all():
+                raise FloatingPointError("Non-finite training/evaluation loss encountered")
+            if training:
+                backward_loss.backward()
+                for name, parameter in model.named_parameters():
+                    if parameter.grad is not None and not torch.isfinite(parameter.grad).all():
+                        raise FloatingPointError(f"Non-finite gradient in parameter {name}")
+                optimizer.step()
+
+            batch_samples = int(history.shape[0])
+            total_samples += batch_samples
+            total_step_ratio += float(step_ratios.detach().sum().cpu())
+            total_full_ratio += float(full_ratios.detach().sum().cpu())
+            steps_per_sample = int(step_ratios.shape[1])
+
+    if total_samples == 0 or steps_per_sample is None:
+        raise RuntimeError("DataLoader produced no batches")
+    return {
+        "step_loss_sum": total_step_ratio / total_samples,
+        "mean_step_relative_l2": total_step_ratio / (total_samples * steps_per_sample),
+        "full_relative_l2": total_full_ratio / total_samples,
+        "samples": float(total_samples),
+        "rollout_steps": float(steps_per_sample),
+    }
+
+
+def synchronize(device: torch.device) -> None:
+    if device.type == "cuda" and torch.cuda.is_available():
+        torch.cuda.synchronize(device)
+
+
+def environment_metadata(device: torch.device) -> dict[str, Any]:
+    device_name = "cpu"
+    if device.type == "cuda" and torch.cuda.is_available():
+        device_name = torch.cuda.get_device_name(device)
+    return {
+        "python": platform.python_version(),
+        "pytorch": torch.__version__,
+        "numpy": np.__version__,
+        "scipy": scipy.__version__,
+        "pyyaml": yaml.__version__,
+        "device": str(device),
+        "device_name": device_name,
+        "torch_cuda_version": torch.version.cuda,
+        "hostname": platform.node(),
+    }
+
+
+def atomic_write_json(path: Path, payload: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_suffix(path.suffix + ".tmp")
+    with temporary.open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2, ensure_ascii=False)
+        handle.write("\n")
+    os.replace(temporary, path)
+
+
+def atomic_torch_save(path: Path, payload: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_suffix(path.suffix + ".tmp")
+    torch.save(payload, temporary)
+    os.replace(temporary, path)
+
+
+def load_training_checkpoint(path: Path, device: torch.device) -> dict[str, Any]:
+    checkpoint_path = path.expanduser().resolve()
+    if not checkpoint_path.is_file():
+        raise FileNotFoundError(f"Resume checkpoint does not exist: {checkpoint_path}")
+    try:
+        checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    except TypeError:
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+    if not isinstance(checkpoint, dict):
+        raise TypeError("Resume checkpoint root must be a mapping")
+    required = {
+        "epoch",
+        "model_state_dict",
+        "optimizer_state_dict",
+        "scheduler_state_dict",
+        "best_train_full_relative_l2",
+        "monitor",
+        "test_selected",
+        "config",
+        "run_type",
+    }
+    missing = sorted(required.difference(checkpoint))
+    if missing:
+        raise KeyError(f"Resume checkpoint is missing required keys: {missing}")
+    if checkpoint["run_type"] != "formal":
+        raise ValueError("Formal training can only resume a formal checkpoint")
+    if checkpoint["monitor"] != "train_full_relative_l2":
+        raise ValueError(f"Unexpected checkpoint monitor: {checkpoint['monitor']}")
+    if checkpoint["test_selected"] is not False:
+        raise ValueError("This reproduction forbids resuming a test-selected checkpoint")
+    return checkpoint
+
+
+def capture_rng_state(train_loader: DataLoader) -> dict[str, Any]:
+    state: dict[str, Any] = {
+        "python": random.getstate(),
+        "numpy": np.random.get_state(),
+        "torch": torch.get_rng_state(),
+        "train_loader_generator": train_loader.generator.get_state(),
+    }
+    if torch.cuda.is_available():
+        state["cuda"] = torch.cuda.get_rng_state_all()
+    return state
+
+
+def restore_rng_state(state: dict[str, Any], train_loader: DataLoader) -> None:
+    random.setstate(state["python"])
+    np.random.set_state(state["numpy"])
+    torch.set_rng_state(state["torch"].cpu())
+    train_loader.generator.set_state(state["train_loader_generator"].cpu())
+    if torch.cuda.is_available() and "cuda" in state:
+        torch.cuda.set_rng_state_all([value.cpu() for value in state["cuda"]])
+
+
+def main() -> None:
+    args = parse_args()
+    config = load_config(args.config)
+    training = config["training"]
+    seed = int(training["seed"])
+    seed_everything(seed, bool(training.get("deterministic", True)))
+    device = select_device(args.device)
+
+    epochs = int(training["epochs"] if args.epochs is None else args.epochs)
+    batch_size = int(training["batch_size"] if args.batch_size is None else args.batch_size)
+    horizon = int(config["data"]["horizon"])
+    rollout_steps = horizon if args.rollout_steps is None else int(args.rollout_steps)
+    if epochs <= 0 or batch_size <= 0:
+        raise ValueError("epochs and batch_size must be positive")
+    if args.run_type == "formal" and any(
+        value is not None
+        for value in (
+            args.epochs,
+            args.rollout_steps,
+            args.max_train_samples,
+            args.max_test_samples,
+            args.output_dir,
+            args.checkpoint,
+        )
+    ):
+        raise ValueError(
+            "Formal runs use the exact configured data, epochs, rollout, checkpoint, "
+            "and results paths; overrides require --run-type smoke"
+        )
+
+    output_dir = (
+        resolve_project_path(config["paths"]["results_dir"])
+        if args.output_dir is None
+        else args.output_dir.expanduser().resolve()
+    )
+    checkpoint_path = (
+        resolve_project_path(config["paths"]["checkpoint"])
+        if args.checkpoint is None
+        else args.checkpoint.expanduser().resolve()
+    )
+    latest_checkpoint_path = checkpoint_path.with_name("last_model.pth")
+    history_path = (
+        resolve_project_path(config["paths"]["train_history"])
+        if args.output_dir is None
+        else output_dir / "train_history.json"
+    )
+    output_dir.mkdir(parents=True, exist_ok=True)
+    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
+
+    train_dataset, test_dataset = build_datasets(
+        config,
+        max_train_samples=args.max_train_samples,
+        max_test_samples=args.max_test_samples,
+        rollout_steps=rollout_steps,
+    )
+    pin_memory = bool(training.get("pin_memory", True)) and device.type == "cuda"
+    train_loader = build_loader(
+        train_dataset,
+        batch_size,
+        True,
+        int(training.get("num_workers", 0)),
+        pin_memory,
+        seed,
+    )
+    test_loader = build_loader(
+        test_dataset,
+        batch_size,
+        False,
+        int(training.get("num_workers", 0)),
+        pin_memory,
+        seed,
+    )
+
+    # Move devices without forcing a global dtype conversion: the pointwise
+    # parameters are float32 while spectral weights must remain complex64.
+    model = build_model_from_config(config).to(device=device)
+    parameter_count = sum(parameter.numel() for parameter in model.parameters())
+    paper_parameter_count = int(config["paper"]["reference_parameter_count"])
+    optimizer = torch.optim.Adam(
+        model.parameters(),
+        lr=float(training["learning_rate"]),
+        weight_decay=float(training["weight_decay"]),
+    )
+    scheduler = torch.optim.lr_scheduler.StepLR(
+        optimizer,
+        step_size=int(training["scheduler_step_size"]),
+        gamma=float(training["scheduler_gamma"]),
+    )
+    epsilon = float(training["relative_l2_epsilon"])
+
+    run_started = datetime.now(timezone.utc).isoformat()
+    history_payload: dict[str, Any] = {
+        "schema_version": "fno-ns2d-train-history-v1",
+        "run_type": args.run_type,
+        "started_at": run_started,
+        "completed_at": None,
+        "config_path": str(args.config.expanduser().resolve()),
+        "data_path": str(data_file_from_config(config)),
+        "checkpoint_path": str(checkpoint_path),
+        "config": deepcopy(config),
+        "environment": environment_metadata(device),
+        "parameter_count": parameter_count,
+        "paper_parameter_count": paper_parameter_count,
+        "parameter_count_difference": parameter_count - paper_parameter_count,
+        "test_selected": False,
+        "epochs_requested": epochs,
+        "rollout_steps": rollout_steps,
+        "train_samples": len(train_dataset),
+        "test_samples": len(test_dataset),
+        "history": [],
+    }
+
+    best_metric = float("inf")
+    start_epoch = 1
+    resume_exact_rng = True
+    if args.resume is not None:
+        resume_path = args.resume.expanduser().resolve()
+        resume_checkpoint = load_training_checkpoint(resume_path, device)
+        if resume_checkpoint["config"] != config:
+            raise ValueError("Resume checkpoint configuration differs from the current YAML")
+        resume_epoch = int(resume_checkpoint["epoch"])
+        if not 1 <= resume_epoch < epochs:
+            raise ValueError(
+                f"Resume epoch must be in [1,{epochs - 1}], received {resume_epoch}"
+            )
+        if not history_path.is_file():
+            raise FileNotFoundError(
+                f"Training history required for audited resume is missing: {history_path}"
+            )
+        with history_path.open("r", encoding="utf-8") as handle:
+            previous_history = json.load(handle)
+        if not isinstance(previous_history, dict):
+            raise TypeError("Existing training history root must be a mapping")
+        records = previous_history.get("history")
+        if not isinstance(records, list) or len(records) < resume_epoch:
+            raise ValueError(
+                f"Existing history has {len(records) if isinstance(records, list) else 0} "
+                f"records, fewer than resume epoch {resume_epoch}"
+            )
+        retained_records = records[:resume_epoch]
+        if int(retained_records[-1]["epoch"]) != resume_epoch:
+            raise ValueError("Existing training history is not contiguous at the resume epoch")
+        checkpoint_metric = float(resume_checkpoint["best_train_full_relative_l2"])
+        history_best = min(float(record["train_full_relative_l2"]) for record in retained_records)
+        if not np.isclose(checkpoint_metric, history_best, rtol=1e-7, atol=1e-9):
+            raise ValueError(
+                f"Resume checkpoint best metric {checkpoint_metric} differs from retained "
+                f"history best {history_best}"
+            )
+
+        model.load_state_dict(resume_checkpoint["model_state_dict"], strict=True)
+        optimizer.load_state_dict(resume_checkpoint["optimizer_state_dict"])
+        scheduler.load_state_dict(resume_checkpoint["scheduler_state_dict"])
+        if not bool(resume_checkpoint.get("scheduler_step_applied", False)):
+            scheduler.step()
+        rng_state = resume_checkpoint.get("rng_state")
+        if isinstance(rng_state, dict):
+            restore_rng_state(rng_state, train_loader)
+        else:
+            resume_exact_rng = False
+
+        discarded_records = len(records) - resume_epoch
+        history_payload = previous_history
+        history_payload["history"] = retained_records
+        history_payload["completed_at"] = None
+        history_payload["completed"] = False
+        history_payload["epochs_requested"] = epochs
+        history_payload["resume_exact_rng"] = resume_exact_rng
+        resume_events = history_payload.setdefault("resume_events", [])
+        resume_events.append(
+            {
+                "resumed_at": run_started,
+                "checkpoint_path": str(resume_path),
+                "checkpoint_epoch": resume_epoch,
+                "discarded_history_records": discarded_records,
+                "exact_rng_state_restored": resume_exact_rng,
+            }
+        )
+        best_metric = checkpoint_metric
+        start_epoch = resume_epoch + 1
+        atomic_write_json(history_path, history_payload)
+        print(
+            f"resume checkpoint={resume_path} epoch={resume_epoch} "
+            f"discarded_history_records={discarded_records} "
+            f"exact_rng_state_restored={'yes' if resume_exact_rng else 'no'}",
+            flush=True,
+        )
+    print(
+        f"device={device} parameters={parameter_count} "
+        f"paper_parameters={paper_parameter_count} delta={parameter_count-paper_parameter_count}",
+        flush=True,
+    )
+
+    if device.type == "cuda" and torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats(device)
+    for epoch in range(start_epoch, epochs + 1):
+        synchronize(device)
+        epoch_start = time.perf_counter()
+        lr_used = float(optimizer.param_groups[0]["lr"])
+        train_metrics = run_epoch(model, train_loader, device, epsilon, optimizer)
+        test_metrics = run_epoch(model, test_loader, device, epsilon, optimizer=None)
+        synchronize(device)
+        duration = time.perf_counter() - epoch_start
+        peak_memory_bytes = (
+            int(torch.cuda.max_memory_allocated(device))
+            if device.type == "cuda" and torch.cuda.is_available()
+            else 0
+        )
+
+        monitor = float(train_metrics["full_relative_l2"])
+        is_best = monitor < best_metric
+        if is_best:
+            best_metric = monitor
+            checkpoint = {
+                "schema_version": "fno-ns2d-checkpoint-v1",
+                "epoch": epoch,
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "scheduler_state_dict": scheduler.state_dict(),
+                "best_train_full_relative_l2": best_metric,
+                "test_full_relative_l2_at_best_epoch": float(
+                    test_metrics["full_relative_l2"]
+                ),
+                "monitor": "train_full_relative_l2",
+                "test_selected": False,
+                "config": deepcopy(config),
+                "seed": seed,
+                "parameter_count": parameter_count,
+                "paper_parameter_count": paper_parameter_count,
+                "parameter_count_difference": parameter_count - paper_parameter_count,
+                "run_type": args.run_type,
+                "scheduler_step_applied": False,
+                "rng_state": capture_rng_state(train_loader),
+            }
+            atomic_torch_save(checkpoint_path, checkpoint)
+
+        epoch_record = {
+            "epoch": epoch,
+            "learning_rate": lr_used,
+            "duration_seconds": duration,
+            "train_step_loss_sum": float(train_metrics["step_loss_sum"]),
+            "train_mean_step_relative_l2": float(
+                train_metrics["mean_step_relative_l2"]
+            ),
+            "train_full_relative_l2": float(train_metrics["full_relative_l2"]),
+            "test_mean_step_relative_l2": float(test_metrics["mean_step_relative_l2"]),
+            "test_full_relative_l2": float(test_metrics["full_relative_l2"]),
+            "peak_accelerator_memory_bytes": peak_memory_bytes,
+            "best": is_best,
+        }
+        history_payload["history"].append(epoch_record)
+        history_payload["best_epoch"] = next(
+            record["epoch"]
+            for record in reversed(history_payload["history"])
+            if record["best"]
+        )
+        history_payload["best_train_full_relative_l2"] = best_metric
+        atomic_write_json(history_path, history_payload)
+        print(
+            f"epoch={epoch:04d}/{epochs:04d} time={duration:.3f}s lr={lr_used:.6g} "
+            f"train_step_loss_sum={train_metrics['step_loss_sum']:.8f} "
+            f"train_step_rel_l2={train_metrics['mean_step_relative_l2']:.8f} "
+            f"train_full_rel_l2={train_metrics['full_relative_l2']:.8f} "
+            f"test_step_rel_l2={test_metrics['mean_step_relative_l2']:.8f} "
+            f"test_full_rel_l2={test_metrics['full_relative_l2']:.8f} "
+            f"peak_mem_gib={peak_memory_bytes / (1024 ** 3):.3f} "
+            f"best={'yes' if is_best else 'no'}",
+            flush=True,
+        )
+        scheduler.step()
+        latest_checkpoint = {
+            "schema_version": "fno-ns2d-training-state-v1",
+            "epoch": epoch,
+            "model_state_dict": model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "scheduler_state_dict": scheduler.state_dict(),
+            "best_train_full_relative_l2": best_metric,
+            "latest_train_full_relative_l2": float(train_metrics["full_relative_l2"]),
+            "latest_test_full_relative_l2": float(test_metrics["full_relative_l2"]),
+            "monitor": "train_full_relative_l2",
+            "test_selected": False,
+            "config": deepcopy(config),
+            "seed": seed,
+            "parameter_count": parameter_count,
+            "paper_parameter_count": paper_parameter_count,
+            "parameter_count_difference": parameter_count - paper_parameter_count,
+            "run_type": args.run_type,
+            "scheduler_step_applied": True,
+            "rng_state": capture_rng_state(train_loader),
+        }
+        atomic_torch_save(latest_checkpoint_path, latest_checkpoint)
+
+    history_payload["completed_at"] = datetime.now(timezone.utc).isoformat()
+    history_payload["completed"] = True
+    atomic_write_json(history_path, history_payload)
+    print(
+        f"training_complete best_epoch={history_payload['best_epoch']} "
+        f"best_train_full_rel_l2={best_metric:.8f} checkpoint={checkpoint_path}",
+        flush=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

weight/best_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a99344d653c72dd75cda2625181964c9f83ce3b2b4ad132aa2373a0a86ba883
+size 28489073

weight/last_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f51c51bb0b78668aaff0a530ed5ed699a4562c0cdd2577c83a4770a6105f56d0
+size 28489137