Upload folder using huggingface_hub

README.md

@@ -0,0 +1,142 @@
+---
+license: mit
+language:
+- en
+tags:
+- OneScience
+- fluid-mechanics
+- flow-field-prediction
+- convolutional-neural-operator
+frameworks: PyTorch
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">CNO</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+CNO (Convolutional Neural Operator) is a neural operator model for learning partial differential equation (PDE) operators, proposed by **Bogdan Raonić et al.** and published at **NeurIPS 2023**. CNO combines convolutional neural networks with operator learning in continuous function spaces. Its anti-aliased activations and filtered upsampling and downsampling operations reduce errors caused by discretization and changes in resolution. The model can directly predict PDE solutions from initial conditions, source terms, or parameter fields. Using the OneScience skill workflow, this project independently reproduces the CNO experiment that predicts the horizontal velocity component of the two-dimensional incompressible Navier–Stokes equations from \(t=0\) to \(T=1\).
+
+Paper: [Convolutional Neural Operators for Robust and Accurate Learning of PDEs](https://arxiv.org/abs/2302.01178)
+
+# Model Description
+
+CNO uses a **U-Net-like multiscale encoder-decoder architecture** that combines local convolutions, skip connections, and filtered upsampling and downsampling to learn function-to-function mappings for PDEs. Its key feature is an **anti-aliasing mechanism** in nonlinear activations and scale transformations. This reduces errors caused by discretization and resolution changes, better preserves the properties of continuous operators, and improves both operator-learning accuracy and cross-resolution robustness.
+
+## Use Cases
+
+| Use case | Description |
+| --- | --- |
+| PDE solution-operator learning | Learns mappings from initial conditions, source terms, or parameter fields to PDE solutions. |
+| Fluid-dynamics prediction | Provides fast surrogate predictions for problems such as Navier–Stokes and compressible Euler equations. |
+| Multiscale physical-field modeling | Handles complex PDE solutions containing both low- and high-frequency spatial structures. |
+| Cross-resolution prediction | Supports inference at different spatial discretization resolutions to evaluate resolution generalization. |
+| ID/OOD generalization | Evaluates model robustness under in-distribution and out-of-distribution PDE parameters or physical conditions. |
+
+# 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/CNO --local_dir ./CNO
+cd CNO
+```
+
+### 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 two-dimensional incompressible Navier–Stokes data from the RPB_CNO dataset. The training file is `NavierStokes_64x64_IN.h5`. Each sample contains:
+
+- `input`: the horizontal velocity component at \(t=0\), with shape `(64, 64)`;
+- `output`: the horizontal velocity component at \(T=1\), with shape `(64, 64)`.
+
+A channel dimension is added after loading, so both inputs and labels in each training batch have shape `(batch_size, 1, 64, 64)`.
+
+The current configuration uses 750 training samples, 128 validation samples, and 128 in-distribution test samples. The out-of-distribution test set contains 128 samples from `NavierStokes_128x128_OUT.h5`.
+
+Download the data with:
+
+```bash
+modelscope download --dataset OneScience/RPB_CNO --local_dir ./data
+```
+
+After downloading, set `paths.data_dir` in `config/config.yaml` to the directory containing the HDF5 files.
+
+### Training
+
+The default configuration corresponds to the two-dimensional incompressible Navier–Stokes experiment in the paper, with spectral viscosity approximately \(\nu=4\times10^{-4}\). The model learns the mapping from the initial horizontal velocity component at \(t=0\) to the velocity component at \(T=1\).
+
+```bash
+python scripts/train.py --config config/config.yaml --device auto
+```
+
+The complete training state with the lowest median relative L1 error in physical space on the validation set is saved to `weight/best_model.pth`. It includes the model, optimizer, learning-rate scheduler, and normalization parameters.
+
+### 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 model performs a single-step prediction from the initial horizontal velocity component at \(t=0\) to the velocity component at \(T=1\). It does not perform closed-loop multistep or trajectory rollout prediction. The default inference batch size is 16.
+
+```bash
+python scripts/inference.py --config config/config.yaml
+```
+
+### Evaluation and Visualization
+
+After training and inference, run:
+
+```bash
+python scripts/result.py --config config/config.yaml --sample-index 0
+```
+
+# 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
+
+- Original paper: [Convolutional Neural Operators for Robust and Accurate Learning of PDEs](https://arxiv.org/abs/2302.01178)
+- This project is an independent reproduction of the CNO paper. The official implementation is licensed under the MIT License. The project code, model weights, training data, and third-party dependencies remain subject to their respective license terms.

config/config.yaml

@@ -0,0 +1,81 @@
+# Paper-faithful configuration for arXiv:2302.01178, Table 12 and Appendix C.
+experiment:
+  name: cno_navier_stokes_2d
+  paper: https://arxiv.org/pdf/2302.01178
+  seed: 0                    # Paper-unspecified, explicit reproducibility choice.
+  deterministic: true
+
+paths:
+  data_dir: /public/share/sugonhpcapp01/onestore/onedatasets/RPB_CNO/data
+  train_file: NavierStokes_64x64_IN.h5
+  id_test_file: NavierStokes_64x64_IN.h5
+  ood_test_file: NavierStokes_128x128_OUT.h5
+  checkpoint: weight/best_model.pth
+  results_dir: results
+
+data:
+  input_key: input
+  output_key: output
+  train:
+    start: 0
+    stop: 750
+  validation:
+    start: 768
+    stop: 896
+  test_id:
+    start: 896
+    stop: 1024
+  test_ood:
+    start: 0
+    stop: 128
+
+# The paper requires [0,1] training normalization and reuse at test time but
+# does not publish its constants. These extrema were measured once from the
+# supplied 64x64 ID benchmark and are fixed for every split, including OOD.
+normalization:
+  source: supplied_64x64_id_benchmark
+  input_min: -1.4294605255126953
+  input_max: 1.4294605255126953
+  output_min: -2.0383081436157227
+  output_max: 2.0602376461029053
+  epsilon: 1.0e-12
+
+model:
+  in_channels: 1
+  out_channels: 1
+  base_width: 32             # d_e in Table 12.
+  levels: 3                  # M in Table 12.
+  bottleneck_residual_blocks: 8
+  intermediate_residual_blocks: 1
+  kernel_size: 3
+  latent_channels: 64        # Paper-unspecified lift/project internal width.
+  activation_upsampling_factor: 2  # N_sigma.
+  filter_taps: 12            # N_tap.
+  filter_half_width: 0.8     # c_h.
+  cutoff_denominator: 2.0001
+  leaky_relu_slope: 0.2      # Paper-unspecified, official supplemental fact.
+
+training:
+  epochs: 1000
+  batch_size: 32
+  num_workers: 4
+  optimizer: Adam
+  learning_rate: 0.001
+  weight_decay: 1.0e-10
+  scheduler: StepLR
+  scheduler_step_size: 1
+  scheduler_gamma: 0.98
+  early_stopping_patience: 50
+  log_interval: 5
+  device: auto
+
+inference:
+  batch_size: 16
+  num_workers: 2
+  device: auto
+  metric_epsilon: 1.0e-12
+
+paper_reference:
+  metric: relative_median_l1_percent
+  id: 2.76
+  ood: 7.04

configuration.json

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

models/FNO.py

@@ -0,0 +1,437 @@
+"""Convolutional Neural Operator used by the CNO Navier--Stokes experiment.
+
+The filename is fixed by the reproduction request.  This module implements a
+CNO, not a Fourier Neural Operator.  The implementation is written from the
+architecture and filter description in Sections 2 and C.1.4 of arXiv:2302.01178.
+No source code from the authors' repository is included here.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+
+def _design_windowed_sinc(
+    num_taps: int,
+    resample_factor: int,
+    cutoff_denominator: float,
+    half_width: float,
+) -> Tensor:
+    """Construct a finite, symmetric low-pass windowed-sinc filter.
+
+    The paper fixes ``N_tap=12``, ``c_h=0.8`` and a cutoff arbitrarily close
+    to the target Nyquist frequency, ``s/2.0001``.  At a high-rate grid used
+    for factor-r resampling, this corresponds to a normalized cutoff of
+    ``1 / (r * 2.0001)`` cycles per sample.  The exact finite window is not
+    specified in the paper; a Kaiser window is used as an explicit,
+    configurable approximation to ``scipy.signal.firwin``.
+    """
+    if num_taps < 2:
+        raise ValueError(f"num_taps must be at least 2, got {num_taps}")
+    if resample_factor < 1:
+        raise ValueError("resample_factor must be positive")
+    if cutoff_denominator <= 2.0:
+        raise ValueError("cutoff_denominator must be greater than 2")
+    if half_width <= 0:
+        raise ValueError("half_width must be positive")
+
+    dtype = torch.float64
+    positions = torch.arange(num_taps, dtype=dtype) - (num_taps - 1) / 2
+    cutoff = 1.0 / (resample_factor * cutoff_denominator)
+    ideal = 2.0 * cutoff * torch.sinc(2.0 * cutoff * positions)
+
+    # c_h=0.8 maps to a conventional beta=8.6 window.  Keeping the relation
+    # explicit makes the paper-unspecified window choice auditable.
+    beta = 8.6 * half_width / 0.8
+    window = torch.kaiser_window(num_taps, periodic=False, beta=beta, dtype=dtype)
+    kernel = ideal * window
+    kernel = kernel / kernel.sum()
+    return kernel.to(torch.float32)
+
+
+class FixedSincResample2d(nn.Module):
+    """Separable periodic 2-D windowed-sinc up/downsampling."""
+
+    def __init__(
+        self,
+        factor: int = 2,
+        num_taps: int = 12,
+        cutoff_denominator: float = 2.0001,
+        half_width: float = 0.8,
+    ) -> None:
+        super().__init__()
+        if factor < 1:
+            raise ValueError("factor must be positive")
+        self.factor = int(factor)
+        kernel = _design_windowed_sinc(
+            num_taps=num_taps,
+            resample_factor=factor,
+            cutoff_denominator=cutoff_denominator,
+            half_width=half_width,
+        )
+        self.register_buffer("kernel", kernel, persistent=True)
+
+    def _filter(self, x: Tensor, gain: float = 1.0) -> Tensor:
+        if x.ndim != 4:
+            raise ValueError(f"expected BCHW input, got shape {tuple(x.shape)}")
+        channels = x.shape[1]
+        kernel = self.kernel.to(device=x.device, dtype=x.dtype)
+        taps = int(kernel.numel())
+        pad_left = (taps - 1) // 2
+        pad_right = taps - 1 - pad_left
+
+        weight_x = (kernel * gain).view(1, 1, 1, taps).repeat(channels, 1, 1, 1)
+        x = F.pad(x, (pad_left, pad_right, 0, 0), mode="circular")
+        x = F.conv2d(x, weight_x, groups=channels)
+
+        weight_y = kernel.view(1, 1, taps, 1).repeat(channels, 1, 1, 1)
+        x = F.pad(x, (0, 0, pad_left, pad_right), mode="circular")
+        return F.conv2d(x, weight_y, groups=channels)
+
+    def upsample(self, x: Tensor) -> Tensor:
+        if self.factor == 1:
+            return x
+        batch, channels, height, width = x.shape
+        up = x.new_zeros(batch, channels, height * self.factor, width * self.factor)
+        up[..., :: self.factor, :: self.factor] = x
+        return self._filter(up, gain=float(self.factor * self.factor))
+
+    def downsample(self, x: Tensor) -> Tensor:
+        if self.factor == 1:
+            return x
+        if x.shape[-2] % self.factor or x.shape[-1] % self.factor:
+            raise ValueError(
+                f"spatial shape {tuple(x.shape[-2:])} is not divisible by {self.factor}"
+            )
+        return self._filter(x)[..., :: self.factor, :: self.factor]
+
+
+class BandlimitedActivation(nn.Module):
+    """Paper Eq. (2.6): upsample, activate, then low-pass/downsample."""
+
+    def __init__(
+        self,
+        upsampling_factor: int = 2,
+        num_taps: int = 12,
+        cutoff_denominator: float = 2.0001,
+        half_width: float = 0.8,
+        negative_slope: float = 0.2,
+    ) -> None:
+        super().__init__()
+        self.negative_slope = float(negative_slope)
+        self.resampler = FixedSincResample2d(
+            factor=upsampling_factor,
+            num_taps=num_taps,
+            cutoff_denominator=cutoff_denominator,
+            half_width=half_width,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.resampler.upsample(x)
+        x = F.leaky_relu(x, negative_slope=self.negative_slope)
+        return self.resampler.downsample(x)
+
+
+def _periodic_conv(in_channels: int, out_channels: int, kernel_size: int) -> nn.Conv2d:
+    if kernel_size % 2 != 1:
+        raise ValueError("CNO convolution kernel_size must be odd")
+    return nn.Conv2d(
+        in_channels,
+        out_channels,
+        kernel_size=kernel_size,
+        padding=kernel_size // 2,
+        padding_mode="circular",
+    )
+
+
+class CNOBlock(nn.Module):
+    """Physical-space convolution followed by bandlimited activation/resampling."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        resample: Literal["same", "up", "down"],
+        activation_kwargs: dict,
+        batch_norm: bool = True,
+    ) -> None:
+        super().__init__()
+        self.resample = resample
+        self.conv = _periodic_conv(in_channels, out_channels, kernel_size)
+        self.norm = nn.BatchNorm2d(out_channels) if batch_norm else nn.Identity()
+        self.activation = BandlimitedActivation(**activation_kwargs)
+        self.resampler = FixedSincResample2d(
+            factor=activation_kwargs["upsampling_factor"],
+            num_taps=activation_kwargs["num_taps"],
+            cutoff_denominator=activation_kwargs["cutoff_denominator"],
+            half_width=activation_kwargs["half_width"],
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.activation(self.norm(self.conv(x)))
+        if self.resample == "down":
+            return self.resampler.downsample(x)
+        if self.resample == "up":
+            return self.resampler.upsample(x)
+        return x
+
+
+class ResidualBlock(nn.Module):
+    """Paper Eq. (2.7): identity plus K o Sigma o K."""
+
+    def __init__(self, channels: int, kernel_size: int, activation_kwargs: dict) -> None:
+        super().__init__()
+        self.conv1 = _periodic_conv(channels, channels, kernel_size)
+        self.norm1 = nn.BatchNorm2d(channels)
+        self.activation = BandlimitedActivation(**activation_kwargs)
+        self.conv2 = _periodic_conv(channels, channels, kernel_size)
+        self.norm2 = nn.BatchNorm2d(channels)
+
+    def forward(self, x: Tensor) -> Tensor:
+        residual = self.norm1(self.conv1(x))
+        residual = self.activation(residual)
+        residual = self.norm2(self.conv2(residual))
+        return x + residual
+
+
+class InvariantBlock(nn.Module):
+    """Paper Eq. (2.8): Sigma o K at an unchanged representation."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        activation_kwargs: dict,
+    ) -> None:
+        super().__init__()
+        self.conv = _periodic_conv(in_channels, out_channels, kernel_size)
+        self.norm = nn.BatchNorm2d(out_channels)
+        self.activation = BandlimitedActivation(**activation_kwargs)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.activation(self.norm(self.conv(x)))
+
+
+class LiftProjectBlock(nn.Module):
+    """Two convolutions with no BatchNorm, as specified for lift/project."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        latent_channels: int,
+        kernel_size: int,
+        activation_kwargs: dict,
+    ) -> None:
+        super().__init__()
+        self.conv1 = _periodic_conv(in_channels, latent_channels, kernel_size)
+        self.activation = BandlimitedActivation(**activation_kwargs)
+        self.conv2 = _periodic_conv(latent_channels, out_channels, kernel_size)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv2(self.activation(self.conv1(x)))
+
+
+def _residual_stack(
+    channels: int,
+    count: int,
+    kernel_size: int,
+    activation_kwargs: dict,
+) -> nn.Module:
+    if count == 0:
+        return nn.Identity()
+    return nn.Sequential(
+        *[
+            ResidualBlock(channels, kernel_size, activation_kwargs)
+            for _ in range(count)
+        ]
+    )
+
+
+class CNO2d(nn.Module):
+    """Operator U-Net CNO for the paper's 2-D Navier--Stokes experiment."""
+
+    def __init__(
+        self,
+        in_channels: int = 1,
+        out_channels: int = 1,
+        base_width: int = 32,
+        levels: int = 3,
+        bottleneck_residual_blocks: int = 8,
+        intermediate_residual_blocks: int = 1,
+        kernel_size: int = 3,
+        latent_channels: int = 64,
+        activation_upsampling_factor: int = 2,
+        filter_taps: int = 12,
+        filter_half_width: float = 0.8,
+        cutoff_denominator: float = 2.0001,
+        leaky_relu_slope: float = 0.2,
+    ) -> None:
+        super().__init__()
+        if levels < 1:
+            raise ValueError("levels must be positive")
+        if base_width % 2:
+            raise ValueError("base_width must be even because lift width is d_e/2")
+
+        self.in_channels = int(in_channels)
+        self.out_channels = int(out_channels)
+        self.base_width = int(base_width)
+        self.levels = int(levels)
+        self.required_divisor = 2**levels
+        lift_width = base_width // 2
+        encoder_widths = [base_width * (2**index) for index in range(levels)]
+        activation_kwargs = {
+            "upsampling_factor": activation_upsampling_factor,
+            "num_taps": filter_taps,
+            "cutoff_denominator": cutoff_denominator,
+            "half_width": filter_half_width,
+            "negative_slope": leaky_relu_slope,
+        }
+
+        self.lift = LiftProjectBlock(
+            in_channels,
+            lift_width,
+            latent_channels,
+            kernel_size,
+            activation_kwargs,
+        )
+
+        down_blocks: list[nn.Module] = []
+        intermediate_blocks: list[nn.Module] = []
+        current_width = lift_width
+        for level, next_width in enumerate(encoder_widths):
+            down_blocks.append(
+                CNOBlock(
+                    current_width,
+                    next_width,
+                    kernel_size,
+                    "down",
+                    activation_kwargs,
+                )
+            )
+            # The bottleneck has its own N_res,b stack; N_res,i belongs to the
+            # genuinely intermediate resolutions only.
+            count = intermediate_residual_blocks if level < levels - 1 else 0
+            intermediate_blocks.append(
+                _residual_stack(next_width, count, kernel_size, activation_kwargs)
+            )
+            current_width = next_width
+        self.encoder = nn.ModuleList(down_blocks)
+        self.encoder_residuals = nn.ModuleList(intermediate_blocks)
+
+        self.bottleneck = _residual_stack(
+            encoder_widths[-1],
+            bottleneck_residual_blocks,
+            kernel_size,
+            activation_kwargs,
+        )
+
+        pre_patch: list[nn.Module] = []
+        post_patch: list[nn.Module] = []
+        up_blocks: list[nn.Module] = []
+        decoder_current = encoder_widths[-1]
+        output_widths = list(reversed([lift_width] + encoder_widths[:-1]))
+        for next_width in output_widths:
+            pre_patch.append(
+                InvariantBlock(
+                    decoder_current,
+                    decoder_current,
+                    kernel_size,
+                    activation_kwargs,
+                )
+            )
+            post_patch.append(
+                InvariantBlock(
+                    decoder_current * 2,
+                    decoder_current,
+                    kernel_size,
+                    activation_kwargs,
+                )
+            )
+            up_blocks.append(
+                CNOBlock(
+                    decoder_current,
+                    next_width,
+                    kernel_size,
+                    "up",
+                    activation_kwargs,
+                )
+            )
+            decoder_current = next_width
+        self.decoder_pre_patch = nn.ModuleList(pre_patch)
+        self.decoder_post_patch = nn.ModuleList(post_patch)
+        self.decoder = nn.ModuleList(up_blocks)
+
+        self.project = LiftProjectBlock(
+            lift_width * 2,
+            out_channels,
+            latent_channels,
+            kernel_size,
+            activation_kwargs,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        if x.ndim != 4:
+            raise ValueError(f"CNO2d expects BCHW input, got {tuple(x.shape)}")
+        if x.shape[1] != self.in_channels:
+            raise ValueError(
+                f"expected {self.in_channels} channels, got {x.shape[1]}"
+            )
+        height, width = x.shape[-2:]
+        if height % self.required_divisor or width % self.required_divisor:
+            raise ValueError(
+                f"spatial shape {(height, width)} must be divisible by "
+                f"2**levels={self.required_divisor}"
+            )
+
+        lifted = self.lift(x)
+        encoded = lifted
+        skips: list[Tensor] = []
+        for down, residuals in zip(self.encoder, self.encoder_residuals):
+            encoded = residuals(down(encoded))
+            skips.append(encoded)
+
+        decoded = self.bottleneck(encoded)
+        for pre, post, up, skip in zip(
+            self.decoder_pre_patch,
+            self.decoder_post_patch,
+            self.decoder,
+            reversed(skips),
+        ):
+            decoded = pre(decoded)
+            if decoded.shape[-2:] != skip.shape[-2:]:
+                raise RuntimeError(
+                    "decoder/skip spatial mismatch before patching: "
+                    f"{tuple(decoded.shape)} versus {tuple(skip.shape)}"
+                )
+            decoded = post(torch.cat((decoded, skip), dim=1))
+            decoded = up(decoded)
+
+        if decoded.shape[-2:] != lifted.shape[-2:]:
+            raise RuntimeError(
+                f"final decoder/lift mismatch: {decoded.shape} versus {lifted.shape}"
+            )
+        output = self.project(torch.cat((decoded, lifted), dim=1))
+        if output.shape[-2:] != (height, width):
+            raise RuntimeError(
+                f"CNO changed output grid from {(height, width)} to {output.shape[-2:]}"
+            )
+        return output
+
+
+def build_model(model_config: dict) -> CNO2d:
+    """Build a CNO2d from the ``model`` section of config.yaml."""
+    return CNO2d(**model_config)
+
+
+def count_trainable_parameters(model: nn.Module) -> int:
+    return sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)
+
+
+__all__ = ["CNO2d", "build_model", "count_trainable_parameters"]

models/__init__.py

@@ -0,0 +1,5 @@
+"""Model package for the CNO reproduction."""
+
+from .FNO import CNO2d, build_model, count_trainable_parameters
+
+__all__ = ["CNO2d", "build_model", "count_trainable_parameters"]

scripts/__init__.py

@@ -0,0 +1 @@
+"""Executable and shared utilities for the CNO reproduction."""

scripts/common.py

@@ -0,0 +1,269 @@
+"""Shared data, normalization, metric, and serialization utilities."""
+
+from __future__ import annotations
+
+import json
+import os
+import random
+import tempfile
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any, Sequence
+
+import h5py
+import numpy as np
+import torch
+from torch import Tensor
+from torch.utils.data import Dataset
+import yaml
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+
+def load_config(path: str | Path) -> dict[str, Any]:
+    config_path = Path(path).expanduser().resolve()
+    if not config_path.is_file():
+        raise FileNotFoundError(f"configuration file not found: {config_path}")
+    with config_path.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    if not isinstance(config, dict):
+        raise ValueError(f"configuration root must be a mapping: {config_path}")
+    for section in ("experiment", "paths", "data", "normalization", "model"):
+        if section not in config:
+            raise KeyError(f"missing required config section: {section}")
+    return config
+
+
+def project_path(path: str | Path) -> Path:
+    candidate = Path(path).expanduser()
+    return candidate.resolve() if candidate.is_absolute() else (PROJECT_ROOT / candidate).resolve()
+
+
+def data_file(config: dict[str, Any], filename_key: str) -> Path:
+    directory = Path(config["paths"]["data_dir"]).expanduser()
+    path = (directory / config["paths"][filename_key]).resolve()
+    if not path.is_file():
+        raise FileNotFoundError(f"data file not found: {path}")
+    return path
+
+
+def numeric_sample_ids(split: dict[str, int]) -> list[int]:
+    start, stop = int(split["start"]), int(split["stop"])
+    if start < 0 or stop <= start:
+        raise ValueError(f"invalid half-open sample range [{start}, {stop})")
+    return list(range(start, stop))
+
+
+def set_reproducibility(seed: int, deterministic: bool = True) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    if deterministic:
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+        try:
+            torch.use_deterministic_algorithms(True, warn_only=True)
+        except TypeError:
+            torch.use_deterministic_algorithms(True)
+
+
+def select_device(requested: str) -> torch.device:
+    requested = requested.lower()
+    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("CUDA was requested but torch.cuda.is_available() is false")
+    return device
+
+
+@dataclass(frozen=True)
+class MinMaxNormalizer:
+    input_min: float
+    input_max: float
+    output_min: float
+    output_max: float
+    epsilon: float = 1.0e-12
+    source: str = ""
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "MinMaxNormalizer":
+        values = config["normalization"]
+        result = cls(
+            input_min=float(values["input_min"]),
+            input_max=float(values["input_max"]),
+            output_min=float(values["output_min"]),
+            output_max=float(values["output_max"]),
+            epsilon=float(values.get("epsilon", 1.0e-12)),
+            source=str(values.get("source", "")),
+        )
+        result.validate()
+        return result
+
+    @classmethod
+    def from_state(cls, state: dict[str, Any]) -> "MinMaxNormalizer":
+        result = cls(**state)
+        result.validate()
+        return result
+
+    def validate(self) -> None:
+        values = (self.input_min, self.input_max, self.output_min, self.output_max)
+        if not all(np.isfinite(value) for value in values):
+            raise ValueError(f"normalization contains nonfinite values: {values}")
+        if self.input_max - self.input_min <= self.epsilon:
+            raise ValueError("input normalization range is zero or negative")
+        if self.output_max - self.output_min <= self.epsilon:
+            raise ValueError("output normalization range is zero or negative")
+
+    def normalize_input(self, value: Tensor) -> Tensor:
+        return (value - self.input_min) / (self.input_max - self.input_min)
+
+    def normalize_output(self, value: Tensor) -> Tensor:
+        return (value - self.output_min) / (self.output_max - self.output_min)
+
+    def denormalize_input(self, value: Tensor) -> Tensor:
+        return value * (self.input_max - self.input_min) + self.input_min
+
+    def denormalize_output(self, value: Tensor) -> Tensor:
+        return value * (self.output_max - self.output_min) + self.output_min
+
+    def state_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+class NavierStokesH5Dataset(Dataset[tuple[Tensor, Tensor, int]]):
+    """Lazy reader for the supplied ``Sample_i/{input,output}`` benchmark."""
+
+    def __init__(
+        self,
+        path: str | Path,
+        sample_ids: Sequence[int],
+        normalizer: MinMaxNormalizer,
+        input_key: str = "input",
+        output_key: str = "output",
+    ) -> None:
+        self.path = Path(path).expanduser().resolve()
+        if not self.path.is_file():
+            raise FileNotFoundError(f"HDF5 file not found: {self.path}")
+        self.sample_ids = [int(sample_id) for sample_id in sample_ids]
+        if not self.sample_ids:
+            raise ValueError("dataset sample_ids must not be empty")
+        self.normalizer = normalizer
+        self.input_key = input_key
+        self.output_key = output_key
+        self._handle: h5py.File | None = None
+        self._validate_contract()
+
+    def _validate_contract(self) -> None:
+        with h5py.File(self.path, "r") as handle:
+            for sample_id in (self.sample_ids[0], self.sample_ids[-1]):
+                group_name = f"Sample_{sample_id}"
+                if group_name not in handle:
+                    raise KeyError(f"missing group {group_name} in {self.path}")
+                group = handle[group_name]
+                if self.input_key not in group or self.output_key not in group:
+                    raise KeyError(
+                        f"{group_name} must contain {self.input_key!r} and {self.output_key!r}"
+                    )
+                input_shape = tuple(group[self.input_key].shape)
+                output_shape = tuple(group[self.output_key].shape)
+                if len(input_shape) != 2 or input_shape != output_shape:
+                    raise ValueError(
+                        f"invalid field shapes in {group_name}: {input_shape}, {output_shape}"
+                    )
+
+    def _file(self) -> h5py.File:
+        if self._handle is None:
+            self._handle = h5py.File(self.path, "r")
+        return self._handle
+
+    def __len__(self) -> int:
+        return len(self.sample_ids)
+
+    def __getitem__(self, index: int) -> tuple[Tensor, Tensor, int]:
+        sample_id = self.sample_ids[index]
+        group = self._file()[f"Sample_{sample_id}"]
+        input_array = np.asarray(group[self.input_key], dtype=np.float32)
+        output_array = np.asarray(group[self.output_key], dtype=np.float32)
+        if input_array.shape != output_array.shape or input_array.ndim != 2:
+            raise ValueError(f"invalid shapes for Sample_{sample_id}")
+        if not np.isfinite(input_array).all() or not np.isfinite(output_array).all():
+            raise ValueError(f"nonfinite field values in Sample_{sample_id}")
+        input_tensor = torch.from_numpy(input_array.copy()).unsqueeze(0)
+        output_tensor = torch.from_numpy(output_array.copy()).unsqueeze(0)
+        return (
+            self.normalizer.normalize_input(input_tensor),
+            self.normalizer.normalize_output(output_tensor),
+            sample_id,
+        )
+
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        state["_handle"] = None
+        return state
+
+    def close(self) -> None:
+        if self._handle is not None:
+            self._handle.close()
+            self._handle = None
+
+    def __del__(self) -> None:
+        # h5py modules may already be partially torn down during interpreter
+        # shutdown.  Explicit ``close`` remains available for normal control
+        # flow; finalization must never emit a spurious exception.
+        try:
+            self.close()
+        except Exception:
+            self._handle = None
+
+
+def relative_l1_per_sample(prediction: Tensor, target: Tensor, epsilon: float) -> Tensor:
+    if prediction.shape != target.shape:
+        raise ValueError(
+            f"prediction/target shape mismatch: {prediction.shape} versus {target.shape}"
+        )
+    reduce_dims = tuple(range(1, prediction.ndim))
+    numerator = torch.sum(torch.abs(prediction - target), dim=reduce_dims)
+    denominator = torch.sum(torch.abs(target), dim=reduce_dims).clamp_min(epsilon)
+    return numerator / denominator
+
+
+def atomic_json_dump(payload: Any, path: str | Path) -> None:
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.NamedTemporaryFile(
+        mode="w", encoding="utf-8", dir=destination.parent, suffix=".json", delete=False
+    ) as handle:
+        json.dump(payload, handle, indent=2, ensure_ascii=False)
+        handle.write("\n")
+        temporary = Path(handle.name)
+    os.replace(temporary, destination)
+
+
+def atomic_torch_save(payload: Any, path: str | Path) -> None:
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.NamedTemporaryFile(dir=destination.parent, suffix=".pth", delete=False) as handle:
+        temporary = Path(handle.name)
+    try:
+        torch.save(payload, temporary)
+        os.replace(temporary, destination)
+    finally:
+        if temporary.exists():
+            temporary.unlink()
+
+
+def atomic_npz_save(path: str | Path, **arrays: np.ndarray) -> None:
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.NamedTemporaryFile(dir=destination.parent, suffix=".npz", delete=False) as handle:
+        temporary = Path(handle.name)
+    try:
+        np.savez_compressed(temporary, **arrays)
+        os.replace(temporary, destination)
+    finally:
+        if temporary.exists():
+            temporary.unlink()

scripts/inference.py

@@ -0,0 +1,231 @@
+#!/usr/bin/env python3
+"""Evaluate a trained CNO on the paper's ID and OOD Navier--Stokes sets."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+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 models.FNO import build_model
+from scripts.common import (
+    MinMaxNormalizer,
+    NavierStokesH5Dataset,
+    atomic_json_dump,
+    atomic_npz_save,
+    data_file,
+    load_config,
+    numeric_sample_ids,
+    project_path,
+    relative_l1_per_sample,
+    select_device,
+    set_reproducibility,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--config", default=str(PROJECT_ROOT / "config" / "config.yaml")
+    )
+    parser.add_argument("--checkpoint", default=None)
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--batch-size", type=int, default=None)
+    parser.add_argument("--output-dir", default=None)
+    parser.add_argument(
+        "--splits",
+        nargs="+",
+        choices=("id", "ood"),
+        default=("id", "ood"),
+    )
+    return parser.parse_args()
+
+
+def load_trained_model(
+    config: dict[str, Any], checkpoint_path: Path, device: torch.device
+) -> tuple[torch.nn.Module, MinMaxNormalizer, dict[str, Any]]:
+    if not checkpoint_path.is_file():
+        raise FileNotFoundError(f"checkpoint not found: {checkpoint_path}")
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    required = {
+        "model_state_dict",
+        "normalization",
+        "epoch",
+        "best_val_relative_l1",
+    }
+    missing = sorted(required.difference(checkpoint))
+    if missing:
+        raise KeyError(f"checkpoint is missing required keys: {missing}")
+    checkpoint_config = checkpoint.get("config", {})
+    if checkpoint_config and checkpoint_config.get("model") != config["model"]:
+        raise ValueError("checkpoint model configuration differs from config.yaml")
+    model = build_model(config["model"]).to(device)
+    model.load_state_dict(checkpoint["model_state_dict"], strict=True)
+    model.eval()
+    normalizer = MinMaxNormalizer.from_state(checkpoint["normalization"])
+    return model, normalizer, checkpoint
+
+
+@torch.inference_mode()
+def evaluate_split(
+    split_name: str,
+    model: torch.nn.Module,
+    loader: DataLoader,
+    normalizer: MinMaxNormalizer,
+    device: torch.device,
+    epsilon: float,
+) -> tuple[dict[str, np.ndarray], dict[str, Any]]:
+    all_ids: list[np.ndarray] = []
+    all_inputs: list[np.ndarray] = []
+    all_targets: list[np.ndarray] = []
+    all_predictions: list[np.ndarray] = []
+    all_ratios: list[np.ndarray] = []
+    completed = 0
+    total = len(loader.dataset)
+
+    for batch_index, (inputs, targets, sample_ids) in enumerate(loader, start=1):
+        inputs = inputs.to(device, non_blocking=True)
+        targets = targets.to(device, non_blocking=True)
+        predictions = model(inputs)
+        physical_inputs = normalizer.denormalize_input(inputs)
+        physical_targets = normalizer.denormalize_output(targets)
+        physical_predictions = normalizer.denormalize_output(predictions)
+        ratios = relative_l1_per_sample(
+            physical_predictions, physical_targets, epsilon
+        )
+
+        all_ids.append(np.asarray(sample_ids, dtype=np.int64))
+        all_inputs.append(physical_inputs.cpu().numpy().astype(np.float32))
+        all_targets.append(physical_targets.cpu().numpy().astype(np.float32))
+        all_predictions.append(physical_predictions.cpu().numpy().astype(np.float32))
+        all_ratios.append(ratios.cpu().numpy().astype(np.float64))
+        completed += inputs.shape[0]
+        running = np.concatenate(all_ratios) * 100.0
+        print(
+            f"inference split={split_name} batch={batch_index}/{len(loader)} "
+            f"samples={completed}/{total} running_rel_l1_median={np.median(running):.6f}%",
+            flush=True,
+        )
+
+    arrays = {
+        "sample_ids": np.concatenate(all_ids),
+        "inputs": np.concatenate(all_inputs),
+        "targets": np.concatenate(all_targets),
+        "predictions": np.concatenate(all_predictions),
+        "relative_l1": np.concatenate(all_ratios),
+    }
+    percentages = arrays["relative_l1"] * 100.0
+    metrics = {
+        "sample_count": int(percentages.size),
+        "resolution": [int(arrays["inputs"].shape[-2]), int(arrays["inputs"].shape[-1])],
+        "relative_l1_median_percent": float(np.median(percentages)),
+        "relative_l1_mean_percent": float(np.mean(percentages)),
+        "relative_l1_std_percent": float(np.std(percentages)),
+        "relative_l1_min_percent": float(np.min(percentages)),
+        "relative_l1_max_percent": float(np.max(percentages)),
+    }
+    return arrays, metrics
+
+
+def main() -> None:
+    args = parse_args()
+    config = load_config(args.config)
+    inference_config = config["inference"]
+    device = select_device(args.device or str(inference_config["device"]))
+    batch_size = int(
+        args.batch_size if args.batch_size is not None else inference_config["batch_size"]
+    )
+    if batch_size < 1:
+        raise ValueError("batch_size must be positive")
+    seed = int(config["experiment"]["seed"])
+    set_reproducibility(seed, bool(config["experiment"].get("deterministic", True)))
+
+    checkpoint_path = (
+        Path(args.checkpoint).expanduser().resolve()
+        if args.checkpoint
+        else project_path(config["paths"]["checkpoint"])
+    )
+    output_dir = (
+        Path(args.output_dir).expanduser().resolve()
+        if args.output_dir
+        else project_path(config["paths"]["results_dir"])
+    )
+    output_dir.mkdir(parents=True, exist_ok=True)
+    model, normalizer, checkpoint = load_trained_model(config, checkpoint_path, device)
+
+    split_specs = {
+        "id": ("id_test_file", "test_id"),
+        "ood": ("ood_test_file", "test_ood"),
+    }
+    all_metrics: dict[str, Any] = {
+        "schema_version": "cno-navier-stokes-metrics-v1",
+        "checkpoint": str(checkpoint_path),
+        "checkpoint_epoch": int(checkpoint["epoch"]),
+        "best_validation_relative_l1_percent": float(
+            checkpoint["best_val_relative_l1"]
+        ),
+        "normalization": normalizer.state_dict(),
+        "paper_reference": config.get("paper_reference", {}),
+        "splits": {},
+    }
+    print(
+        f"inference device={device} checkpoint_epoch={checkpoint['epoch']} "
+        f"best_val={float(checkpoint['best_val_relative_l1']):.6f}%",
+        flush=True,
+    )
+
+    for split_name in args.splits:
+        filename_key, split_key = split_specs[split_name]
+        dataset = NavierStokesH5Dataset(
+            data_file(config, filename_key),
+            numeric_sample_ids(config["data"][split_key]),
+            normalizer,
+            str(config["data"]["input_key"]),
+            str(config["data"]["output_key"]),
+        )
+        loader = DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=int(inference_config["num_workers"]),
+            pin_memory=device.type == "cuda",
+            persistent_workers=int(inference_config["num_workers"]) > 0,
+        )
+        arrays, metrics = evaluate_split(
+            split_name,
+            model,
+            loader,
+            normalizer,
+            device,
+            float(inference_config["metric_epsilon"]),
+        )
+        artifact_path = output_dir / f"{split_name}_predictions.npz"
+        atomic_npz_save(artifact_path, **arrays)
+        metrics["predictions_file"] = str(artifact_path)
+        all_metrics["splits"][split_name] = metrics
+        print(
+            f"evaluation split={split_name} n={metrics['sample_count']} "
+            f"rel_l1_median={metrics['relative_l1_median_percent']:.6f}% "
+            f"mean={metrics['relative_l1_mean_percent']:.6f}% "
+            f"std={metrics['relative_l1_std_percent']:.6f}% "
+            f"saved={artifact_path}",
+            flush=True,
+        )
+
+    metrics_path = output_dir / "metrics.json"
+    atomic_json_dump(all_metrics, metrics_path)
+    print(f"metrics saved path={metrics_path}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+"""Visualize CNO fields and ID/OOD relative-L1 distributions."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--results-dir", default=str(PROJECT_ROOT / "results"))
+    parser.add_argument(
+        "--sample-index",
+        type=int,
+        default=None,
+        help="array index to plot; default is the sample nearest the split median error",
+    )
+    parser.add_argument("--dpi", type=int, default=180)
+    return parser.parse_args()
+
+
+def load_predictions(path: Path) -> dict[str, np.ndarray]:
+    if not path.is_file():
+        raise FileNotFoundError(f"prediction artifact not found: {path}")
+    with np.load(path) as payload:
+        required = {"sample_ids", "inputs", "targets", "predictions", "relative_l1"}
+        missing = sorted(required.difference(payload.files))
+        if missing:
+            raise KeyError(f"{path} is missing arrays: {missing}")
+        arrays = {key: np.asarray(payload[key]) for key in required}
+    count = arrays["sample_ids"].shape[0]
+    for key in ("inputs", "targets", "predictions", "relative_l1"):
+        if arrays[key].shape[0] != count:
+            raise ValueError(f"sample count mismatch for {key} in {path}")
+        if not np.isfinite(arrays[key]).all():
+            raise ValueError(f"nonfinite values in {key} from {path}")
+    if arrays["inputs"].ndim != 4 or arrays["inputs"].shape[1] != 1:
+        raise ValueError(f"expected N1HW fields in {path}, got {arrays['inputs'].shape}")
+    return arrays
+
+
+def representative_index(errors: np.ndarray, requested: int | None) -> int:
+    if requested is not None:
+        if requested < 0 or requested >= errors.size:
+            raise IndexError(f"sample-index {requested} outside [0,{errors.size})")
+        return requested
+    median = np.median(errors)
+    return int(np.argmin(np.abs(errors - median)))
+
+
+def plot_fields(
+    split: str,
+    arrays: dict[str, np.ndarray],
+    output_path: Path,
+    requested_index: int | None,
+    dpi: int,
+) -> None:
+    errors_percent = arrays["relative_l1"] * 100.0
+    index = representative_index(errors_percent, requested_index)
+    input_field = arrays["inputs"][index, 0]
+    target = arrays["targets"][index, 0]
+    prediction = arrays["predictions"][index, 0]
+    absolute_error = np.abs(prediction - target)
+    field_min = float(min(target.min(), prediction.min()))
+    field_max = float(max(target.max(), prediction.max()))
+
+    fig, axes = plt.subplots(1, 4, figsize=(16, 3.8), constrained_layout=True)
+    input_image = axes[0].imshow(input_field, origin="lower", cmap="RdBu_r")
+    fig.colorbar(input_image, ax=axes[0], shrink=0.78)
+    target_image = axes[1].imshow(
+        target, origin="lower", cmap="RdBu_r", vmin=field_min, vmax=field_max
+    )
+    prediction_image = axes[2].imshow(
+        prediction, origin="lower", cmap="RdBu_r", vmin=field_min, vmax=field_max
+    )
+    error_image = axes[3].imshow(absolute_error, origin="lower", cmap="magma")
+    fig.colorbar(target_image, ax=[axes[1], axes[2]], shrink=0.78)
+    fig.colorbar(error_image, ax=axes[3], shrink=0.78)
+    titles = ("Initial velocity", "Target at T=1", "CNO prediction", "Absolute error")
+    for axis, title in zip(axes, titles):
+        axis.set_title(title)
+        axis.set_xlabel("x index")
+        axis.set_ylabel("y index")
+    sample_id = int(arrays["sample_ids"][index])
+    fig.suptitle(
+        f"{split.upper()} Sample_{sample_id} — relative L1={errors_percent[index]:.3f}%"
+    )
+    fig.savefig(output_path, dpi=dpi)
+    plt.close(fig)
+
+
+def plot_distribution(
+    split_arrays: dict[str, dict[str, np.ndarray]],
+    paper_reference: dict,
+    output_path: Path,
+    dpi: int,
+) -> None:
+    fig, axes = plt.subplots(1, 2, figsize=(11, 4), constrained_layout=True)
+    colors = {"id": "#2166ac", "ood": "#b2182b"}
+    for split in ("id", "ood"):
+        values = split_arrays[split]["relative_l1"] * 100.0
+        axes[0].hist(values, bins=20, alpha=0.55, label=split.upper(), color=colors[split])
+        axes[0].axvline(
+            np.median(values), color=colors[split], linewidth=2, linestyle="-"
+        )
+        reference = paper_reference.get(split)
+        if reference is not None:
+            axes[0].axvline(
+                float(reference), color=colors[split], linewidth=1.5, linestyle="--"
+            )
+    axes[0].set_title("Per-sample relative L1")
+    axes[0].set_xlabel("Relative L1 (%)")
+    axes[0].set_ylabel("Count")
+    axes[0].legend(title="solid=reproduction\ndashed=paper")
+
+    values = [
+        split_arrays["id"]["relative_l1"] * 100.0,
+        split_arrays["ood"]["relative_l1"] * 100.0,
+    ]
+    box = axes[1].boxplot(values, tick_labels=["ID", "OOD"], patch_artist=True)
+    for patch, color in zip(box["boxes"], (colors["id"], colors["ood"])):
+        patch.set_facecolor(color)
+        patch.set_alpha(0.55)
+    axes[1].set_title("Error distribution summary")
+    axes[1].set_ylabel("Relative L1 (%)")
+    fig.savefig(output_path, dpi=dpi)
+    plt.close(fig)
+
+
+def main() -> None:
+    args = parse_args()
+    results_dir = Path(args.results_dir).expanduser().resolve()
+    metrics_path = results_dir / "metrics.json"
+    if not metrics_path.is_file():
+        raise FileNotFoundError(f"metrics file not found: {metrics_path}")
+    with metrics_path.open("r", encoding="utf-8") as handle:
+        metrics = json.load(handle)
+
+    split_arrays = {
+        split: load_predictions(results_dir / f"{split}_predictions.npz")
+        for split in ("id", "ood")
+    }
+    for split, arrays in split_arrays.items():
+        output_path = results_dir / f"{split}_fields.png"
+        plot_fields(split, arrays, output_path, args.sample_index, args.dpi)
+        values = arrays["relative_l1"] * 100.0
+        print(
+            f"visualization split={split} median={np.median(values):.6f}% "
+            f"mean={np.mean(values):.6f}% saved={output_path}",
+            flush=True,
+        )
+
+    distribution_path = results_dir / "error_distribution.png"
+    plot_distribution(
+        split_arrays,
+        metrics.get("paper_reference", {}),
+        distribution_path,
+        args.dpi,
+    )
+    print(f"visualization saved={distribution_path}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,331 @@
+#!/usr/bin/env python3
+"""Train the paper-configured CNO on the 2-D Navier--Stokes benchmark."""
+
+from __future__ import annotations
+
+import argparse
+import platform
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+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 models.FNO import build_model, count_trainable_parameters
+from scripts.common import (
+    MinMaxNormalizer,
+    NavierStokesH5Dataset,
+    atomic_json_dump,
+    atomic_torch_save,
+    data_file,
+    load_config,
+    numeric_sample_ids,
+    project_path,
+    relative_l1_per_sample,
+    select_device,
+    set_reproducibility,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--config",
+        default=str(PROJECT_ROOT / "config" / "config.yaml"),
+        help="experiment YAML configuration",
+    )
+    parser.add_argument("--device", default=None, help="override training.device")
+    parser.add_argument("--epochs", type=int, default=None, help="override training.epochs")
+    parser.add_argument(
+        "--resume",
+        default=None,
+        help="resume a complete training state; not a weight-only initialization",
+    )
+    return parser.parse_args()
+
+
+def _make_loader(
+    dataset: NavierStokesH5Dataset,
+    batch_size: int,
+    workers: int,
+    shuffle: bool,
+    seed: int,
+    device: torch.device,
+) -> DataLoader:
+    generator = torch.Generator()
+    generator.manual_seed(seed)
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=workers,
+        pin_memory=device.type == "cuda",
+        persistent_workers=workers > 0,
+        generator=generator,
+    )
+
+
+@torch.inference_mode()
+def validate(
+    model: torch.nn.Module,
+    loader: DataLoader,
+    normalizer: MinMaxNormalizer,
+    device: torch.device,
+    epsilon: float,
+) -> dict[str, float]:
+    model.eval()
+    ratios: list[torch.Tensor] = []
+    for inputs, targets, _ in loader:
+        inputs = inputs.to(device, non_blocking=True)
+        targets = targets.to(device, non_blocking=True)
+        predictions = model(inputs)
+        predictions = normalizer.denormalize_output(predictions)
+        targets = normalizer.denormalize_output(targets)
+        ratios.append(relative_l1_per_sample(predictions, targets, epsilon).cpu())
+    values = torch.cat(ratios).numpy() * 100.0
+    return {
+        "median_percent": float(np.median(values)),
+        "mean_percent": float(np.mean(values)),
+        "std_percent": float(np.std(values)),
+    }
+
+
+def _resume_training(
+    path: Path,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    scheduler: torch.optim.lr_scheduler.LRScheduler,
+    normalizer: MinMaxNormalizer,
+    device: torch.device,
+) -> tuple[int, float, int, list[dict[str, Any]]]:
+    if not path.is_file():
+        raise FileNotFoundError(f"resume checkpoint not found: {path}")
+    state = torch.load(path, map_location=device, weights_only=False)
+    required = {
+        "model_state_dict",
+        "optimizer_state_dict",
+        "scheduler_state_dict",
+        "epoch",
+        "best_val_relative_l1",
+        "normalization",
+    }
+    missing = sorted(required.difference(state))
+    if missing:
+        raise KeyError(f"resume checkpoint is missing keys: {missing}")
+    checkpoint_normalizer = MinMaxNormalizer.from_state(state["normalization"])
+    if checkpoint_normalizer != normalizer:
+        raise ValueError("resume checkpoint normalization differs from config")
+    model.load_state_dict(state["model_state_dict"], strict=True)
+    optimizer.load_state_dict(state["optimizer_state_dict"])
+    scheduler.load_state_dict(state["scheduler_state_dict"])
+    return (
+        int(state["epoch"]) + 1,
+        float(state["best_val_relative_l1"]),
+        int(state.get("bad_epochs", 0)),
+        list(state.get("history", [])),
+    )
+
+
+def main() -> None:
+    args = parse_args()
+    config = load_config(args.config)
+    train_config = config["training"]
+    epochs = int(args.epochs if args.epochs is not None else train_config["epochs"])
+    if epochs < 1:
+        raise ValueError("epochs must be positive")
+    device = select_device(args.device or str(train_config["device"]))
+    seed = int(config["experiment"]["seed"])
+    set_reproducibility(seed, bool(config["experiment"].get("deterministic", True)))
+    normalizer = MinMaxNormalizer.from_config(config)
+
+    source = data_file(config, "train_file")
+    train_dataset = NavierStokesH5Dataset(
+        source,
+        numeric_sample_ids(config["data"]["train"]),
+        normalizer,
+        str(config["data"]["input_key"]),
+        str(config["data"]["output_key"]),
+    )
+    validation_dataset = NavierStokesH5Dataset(
+        source,
+        numeric_sample_ids(config["data"]["validation"]),
+        normalizer,
+        str(config["data"]["input_key"]),
+        str(config["data"]["output_key"]),
+    )
+    batch_size = int(train_config["batch_size"])
+    workers = int(train_config["num_workers"])
+    train_loader = _make_loader(train_dataset, batch_size, workers, True, seed, device)
+    validation_loader = _make_loader(
+        validation_dataset, batch_size, workers, False, seed, device
+    )
+
+    model = build_model(config["model"]).to(device)
+    parameter_count = count_trainable_parameters(model)
+    optimizer_name = str(train_config["optimizer"])
+    if optimizer_name != "Adam":
+        raise ValueError(f"paper reproduction requires Adam, got {optimizer_name}")
+    optimizer = torch.optim.Adam(
+        model.parameters(),
+        lr=float(train_config["learning_rate"]),
+        weight_decay=float(train_config["weight_decay"]),
+    )
+    if str(train_config["scheduler"]) != "StepLR":
+        raise ValueError("paper reproduction requires StepLR")
+    scheduler = torch.optim.lr_scheduler.StepLR(
+        optimizer,
+        step_size=int(train_config["scheduler_step_size"]),
+        gamma=float(train_config["scheduler_gamma"]),
+    )
+
+    start_epoch = 1
+    best_validation = float("inf")
+    bad_epochs = 0
+    history: list[dict[str, Any]] = []
+    if args.resume:
+        start_epoch, best_validation, bad_epochs, history = _resume_training(
+            Path(args.resume).expanduser().resolve(),
+            model,
+            optimizer,
+            scheduler,
+            normalizer,
+            device,
+        )
+
+    checkpoint_path = project_path(config["paths"]["checkpoint"])
+    results_dir = project_path(config["paths"]["results_dir"])
+    history_path = results_dir / "training_history.json"
+    patience = int(train_config["early_stopping_patience"])
+    log_interval = max(1, int(train_config["log_interval"]))
+
+    print(
+        f"experiment={config['experiment']['name']} device={device} "
+        f"python={platform.python_version()} torch={torch.__version__}",
+        flush=True,
+    )
+    print(
+        f"train_samples={len(train_dataset)} val_samples={len(validation_dataset)} "
+        f"batch_size={batch_size} parameters={parameter_count:,} "
+        f"checkpoint={checkpoint_path}",
+        flush=True,
+    )
+
+    for epoch in range(start_epoch, epochs + 1):
+        model.train()
+        loss_sum = 0.0
+        sample_count = 0
+        learning_rate = float(optimizer.param_groups[0]["lr"])
+        for batch_index, (inputs, targets, _) in enumerate(train_loader, start=1):
+            inputs = inputs.to(device, non_blocking=True)
+            targets = targets.to(device, non_blocking=True)
+            optimizer.zero_grad(set_to_none=True)
+            predictions = model(inputs)
+            loss = F.l1_loss(predictions, targets)
+            if not torch.isfinite(loss):
+                raise FloatingPointError(
+                    f"nonfinite training loss at epoch={epoch}, batch={batch_index}: {loss}"
+                )
+            loss.backward()
+            optimizer.step()
+            batch_samples = inputs.shape[0]
+            loss_sum += float(loss.detach()) * batch_samples
+            sample_count += batch_samples
+            if batch_index % log_interval == 0 or batch_index == len(train_loader):
+                print(
+                    f"train epoch={epoch}/{epochs} batch={batch_index}/{len(train_loader)} "
+                    f"loss={float(loss.detach()):.8f} running_loss={loss_sum/sample_count:.8f}",
+                    flush=True,
+                )
+
+        train_loss = loss_sum / sample_count
+        validation = validate(
+            model,
+            validation_loader,
+            normalizer,
+            device,
+            float(config["normalization"]["epsilon"]),
+        )
+        scheduler.step()
+        improved = validation["median_percent"] < best_validation
+        if improved:
+            best_validation = validation["median_percent"]
+            bad_epochs = 0
+        else:
+            bad_epochs += 1
+
+        record = {
+            "epoch": epoch,
+            "learning_rate": learning_rate,
+            "train_l1": train_loss,
+            "validation_relative_l1_median_percent": validation["median_percent"],
+            "validation_relative_l1_mean_percent": validation["mean_percent"],
+            "validation_relative_l1_std_percent": validation["std_percent"],
+            "best_validation_percent": best_validation,
+        }
+        history.append(record)
+        print(
+            f"eval epoch={epoch}/{epochs} lr={learning_rate:.8g} "
+            f"train_l1={train_loss:.8f} "
+            f"val_rel_l1_median={validation['median_percent']:.6f}% "
+            f"val_rel_l1_mean={validation['mean_percent']:.6f}% "
+            f"best={best_validation:.6f}% bad_epochs={bad_epochs}/{patience}",
+            flush=True,
+        )
+
+        if improved:
+            checkpoint = {
+                "schema_version": "cno-navier-stokes-checkpoint-v1",
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "scheduler_state_dict": scheduler.state_dict(),
+                "epoch": epoch,
+                "best_val_relative_l1": best_validation,
+                "bad_epochs": bad_epochs,
+                "normalization": normalizer.state_dict(),
+                "config": config,
+                "seed": seed,
+                "parameter_count": parameter_count,
+                "history": history,
+            }
+            atomic_torch_save(checkpoint, checkpoint_path)
+            print(
+                f"checkpoint saved path={checkpoint_path} "
+                f"val_rel_l1_median={best_validation:.6f}%",
+                flush=True,
+            )
+
+        atomic_json_dump(
+            {
+                "experiment": config["experiment"],
+                "device": str(device),
+                "parameter_count": parameter_count,
+                "normalization": normalizer.state_dict(),
+                "best_validation_percent": best_validation,
+                "history": history,
+            },
+            history_path,
+        )
+        if bad_epochs >= patience:
+            print(
+                f"early stopping at epoch={epoch}; no improvement for {patience} epochs",
+                flush=True,
+            )
+            break
+
+    print(
+        f"training complete best_val_rel_l1_median={best_validation:.6f}% "
+        f"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:630ad36593abc8fcd55a79b5fd25add748236f087f66eca551576f21e8d552ea
+size 39430846