Upload folder using huggingface_hub

README.md

@@ -0,0 +1,197 @@
+---
+license: mit
+language:
+- en
+tags:
+- OneScience
+- fluid-mechanics
+- flow-field-prediction
+- neural-operator
+frameworks: PyTorch
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">DeepONet</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+DeepONet is a deep neural network for operator learning proposed by a research team affiliated with Brown University and published in *Nature Machine Intelligence* in 2021. Unlike conventional neural networks, which primarily learn mappings between finite-dimensional vectors, DeepONet learns nonlinear operator mappings directly between function spaces, establishing a relationship between an input function and its corresponding output function. It can approximate solution operators defined by ordinary differential equations, partial differential equations, and other physical systems, providing an efficient data-driven modeling approach for complex dynamical systems and scientific computing.
+
+Using the OneScience skill workflow, this project independently reproduces experiments related to the DeepONet paper.
+
+Paper: [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://arxiv.org/abs/1910.03193)
+
+# Model Description
+
+DeepONet uses a dual-network architecture consisting of a Branch Net and a Trunk Net. The Branch Net encodes discrete samples of an input function at fixed sensor locations, while the Trunk Net encodes the spatial or spatiotemporal coordinates at which predictions are requested. The output of the target operator at a specified location is obtained from the inner product of the two feature vectors plus a bias term.
+
+## Use Cases
+
+| Use case | Description |
+| --- | --- |
+| Operator learning | Learns mappings between function spaces, directly predicting an output function \(G(u)\) from an input function \(u\). This is useful for function-to-function mappings that conventional neural networks cannot easily handle. |
+| Spatiotemporal field prediction | Provides fast surrogate predictions for fluid problems such as Navier–Stokes and compressible Euler equations. |
+| Multiscale physical-field modeling | The Trunk Net can directly accept multidimensional coordinates such as \((x,t)\), making it suitable for predicting temperature, concentration, diffusion, and other fields that vary in space and time. |
+| Multiple query-point prediction | For a fixed input function, solutions at different spatial or temporal locations can be predicted by changing only the query coordinates supplied to the Trunk Net. |
+
+# 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/DeepONet --local_dir ./DeepONet
+cd DeepONet
+```
+
+### 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
+
+This project does not depend on an external dataset. The E3 data is generated on the fly by `models/dataset.py` according to `config/config.yaml` and corresponds to the following parameterized one-dimensional PDE:
+
+$$
+\frac{\partial u}{\partial t}
++\alpha\frac{\partial(u^2)}{\partial x}
+-\beta\frac{\partial^2u}{\partial x^2}
++\gamma\frac{\partial^3u}{\partial x^3}
+=\delta(t,x),
+$$
+
+where $x\in[0,16)$ and $t\in[0,4]$, with periodic boundary conditions in the spatial dimension. The equation parameters are sampled independently:
+
+$$
+\alpha\sim\mathcal U(0,3),\quad
+\beta\sim\mathcal U(0,0.4),\quad
+\gamma\sim\mathcal U(0,1).
+$$
+
+The forcing term and initial condition are defined as:
+
+$$
+\delta(t,x)=\sum_{j=1}^{5}
+A_j\sin\left(
+\omega_jt+\frac{2\pi k_jx}{16}+\phi_j
+\right),
+\qquad
+u(0,x)=\delta(0,x),
+$$
+
+where:
+
+- $A_j\sim\mathcal U(-0.5,0.5)$;
+- $\omega_j=-0.4$;
+- $k_j\in\{1,2,3\}$;
+- $\phi_j\sim\mathcal U(0,2\pi)$.
+
+Reference solutions are first generated on 200 spatial grid points and then downsampled to 100 points. The nonlinear flux is discretized using a fifth-order WENO scheme, and time integration uses a fourth-order Runge–Kutta method. Each trajectory contains 250 time points and has shape:
+
+$$
+u\in\mathbb R^{250\times100}.
+$$
+
+The model uses a history window of length $K=25$ to predict the next 25 time steps. One supervised sample can be written as:
+
+$$
+\left(
+u_{i:i+K-1},\,x,\,t,\,(\alpha,\beta,\gamma);
+\ u_{i+K:i+2K-1}
+\right).
+$$
+
+### Training
+
+The default configuration reproduces four operator-learning experiments from the DeepONet paper: antiderivative, nonlinear ODE, forced pendulum, and diffusion–reaction equation. By default, the training script runs the antiderivative experiment, with input functions and reference solutions generated on the fly by `models/dataset.py`.
+
+Run the default antiderivative experiment:
+
+```bash
+python scripts/train.py \
+    --config config/config.yaml \
+    --experiment antiderivative \
+    --device auto
+```
+
+Set `--experiment` to `all` to run all four main experiments from the paper sequentially.
+
+During training, the training loss, test MSE, and relative L2 error are printed at the intervals configured in `config/config.yaml`. When the test MSE improves, the model weights, optimizer state, current iteration, evaluation metrics, and effective runtime configuration are saved to `weight/best_model.pth`.
+
+### Trained Weights
+
+`weight/best_model.pth` contains the best weights from the full antiderivative experiment and can be used directly for inference or fine-tuning.
+
+### Inference
+
+Before running inference, make sure the configured data path is valid and `weight/best_model.pth` exists.
+
+```bash
+python scripts/inference.py \
+    --config config/config.yaml \
+    --experiment <experiment-name> \
+    --variant unstacked_bias \
+    --mode <inference-mode> \
+    --device auto
+```
+
+Available inference modes are:
+
+- `random_test`: random test set;
+- `ood`: out-of-distribution ODE inputs;
+- `pde_grid`: two-dimensional spatiotemporal field for the diffusion–reaction equation.
+
+The default batch size is 8,192 and can be changed with `--batch-size`. Predictions and evaluation metrics are saved to the `results` directory.
+
+### Evaluation and Visualization
+
+After training and inference, summarize existing experiment results and generate training curves, prediction comparisons, and an evaluation report with:
+
+```bash
+python scripts/result.py --config config/config.yaml
+```
+
+# 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: [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://arxiv.org/abs/1910.03193)
+- This project is an independent reproduction of the DeepONet 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,202 @@
+project:
+  name: DeepONet
+  paper: "DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators"
+  arxiv: "1910.03193v3"
+  paper_url: "https://arxiv.org/abs/1910.03193"
+  implementation_policy: "paper_first_independent_reimplementation"
+  paper_scale: true
+
+paths:
+  checkpoint: weight/best_model.pth
+  results: results
+  cache: results/cache
+
+runtime:
+  seed: 0                    # ASSUMPTION: the paper does not report seeds.
+  device: auto
+  dtype: float32             # ASSUMPTION: precision is not reported.
+  num_workers: 0
+
+function_space:
+  type: grf
+  length_scale: 0.2          # PAPER: default RBF-GRF length scale.
+  grf_grid_size: 1000        # OFFICIAL_REPO_REFERENCE: paper omits this discretization.
+  jitter: 1.0e-13            # OFFICIAL_REPO_REFERENCE: numerical Cholesky stabilizer.
+  interpolation: cubic       # OFFICIAL_REPO_REFERENCE.
+  generation_chunk_size: 128
+  chebyshev:
+    degree: 10
+    coefficient_bound: 1.0
+
+model_defaults:
+  architecture: deeponet
+  activation: relu           # OFFICIAL_REPO_REFERENCE: paper does not state activation.
+  initializer: xavier_normal # OFFICIAL_REPO_REFERENCE: PyTorch equivalent of Glorot normal.
+  dense_bias: true
+  branch_depth: 2            # ASSUMPTION: number of Dense layers after the input.
+  trunk_depth: 3             # ASSUMPTION: number of Dense layers after the input.
+  width: 40
+  latent_dim: 40
+
+variants:
+  unstacked_bias:
+    architecture: deeponet
+    stacked: false
+    branch_output_bias: true
+    global_bias: true
+  unstacked_no_bias:
+    architecture: deeponet
+    stacked: false
+    branch_output_bias: false
+    global_bias: false
+  stacked_bias:
+    architecture: deeponet
+    stacked: true
+    branch_output_bias: true
+    global_bias: true
+  stacked_no_bias:
+    architecture: deeponet
+    stacked: true
+    branch_output_bias: false
+    global_bias: false
+  fnn:
+    architecture: fnn
+    depth: 3
+    width: 40
+    output_bias: true
+
+training_defaults:
+  optimizer: adam
+  learning_rate: 1.0e-3      # PAPER.
+  batch_size: null           # OFFICIAL_REPO_REFERENCE: null means full batch.
+  evaluation_batch_size: 8192
+  weight_decay: 0.0          # No paper evidence; disabled.
+  print_every: 1000
+  evaluate_every: 1000
+  checkpoint_metric: test_mse
+  resume: false
+
+solver_defaults:
+  method: RK45               # PAPER: Runge-Kutta (4,5); method name from official reference.
+  rtol: 1.0e-7               # ASSUMPTION: solver tolerances are not reported.
+  atol: 1.0e-9
+  pde_newton_tolerance: 1.0e-10
+  pde_newton_max_iterations: 20
+
+experiments:
+  antiderivative:
+    equation: "ds/dx = u(x), s(0) = 0"
+    sensor_points: 100
+    trunk_dim: 1
+    domain_end: 1.0
+    train_size: 10000
+    test_size: 100000
+    iterations: 50000
+    default_variant: unstacked_bias
+    sweeps:
+      architectures: [fnn, unstacked_no_bias, unstacked_bias, stacked_no_bias, stacked_bias]
+      fnn_depths: [2, 3, 4]
+      fnn_widths: [20, 40, 100]
+      learning_rates: [0.0001, 0.001, 0.01]
+
+  nonlinear_ode:
+    equation: "ds/dx = -s^2 + u(x), s(0) = 0"
+    sensor_points: 100
+    trunk_dim: 1
+    domain_end: 1.0
+    train_size: 10000
+    test_size: 100000
+    iterations: 100000
+    default_variant: unstacked_bias
+    trim_fraction: 0.001
+    ood_functions: [linear, sin_pi, sin_2pi]
+    sweeps:
+      architectures: [stacked_no_bias, stacked_bias, unstacked_no_bias, unstacked_bias]
+
+  pendulum:
+    equation: "ds1/dt = s2; ds2/dt = -k*sin(s1) + u(t); s(0) = (0,0)"
+    sensor_points: 100
+    trunk_dim: 1
+    domain_end: 1.0
+    k: 1.0
+    train_size: 10000
+    test_size: 100000
+    iterations: 100000
+    default_variant: unstacked_bias
+    convergence_iterations: 500000
+    convergence_test_size: 1000000
+    sweeps:
+      sensor_points: [10, 20, 50, 100, 200]
+      domain_end: [1.0, 2.0, 3.0, 4.0]
+      length_scale: [0.1, 0.2, 0.5, 1.0]
+      k: [0.5, 1.0, 2.0]
+      width: [20, 40, 100]
+      train_size: [1000, 5000, 10000]
+      chebyshev_degree: [5, 10, 20]
+
+  diffusion_reaction:
+    equation: "s_t = D*s_xx + k*s^2 + u(x)"
+    sensor_points: 100
+    trunk_dim: 2
+    domain_end: 1.0
+    train_functions: 100
+    points_per_function: 1000
+    test_functions: 1000
+    test_points_per_function: 1000
+    test_size: 1000000
+    iterations: 500000
+    default_variant: unstacked_bias
+    diffusion: 0.01
+    reaction: 0.01
+    space_points: 100
+    time_points: 100
+    branch_depth: 2
+    trunk_depth: 3
+    width: 100
+    latent_dim: 100
+    sweeps:
+      train_functions: [50, 100, 200, 400]
+      points_per_function: [10, 100, 1000, 10000]
+
+inference:
+  batch_size: 8192
+  ood_query_points: 100
+  save_inputs: true
+
+smoke_test:
+  project:
+    paper_scale: false
+  function_space:
+    grf_grid_size: 32
+    generation_chunk_size: 16
+  training_defaults:
+    batch_size: 16
+    evaluation_batch_size: 32
+    print_every: 1
+    evaluate_every: 1
+  experiments:
+    antiderivative:
+      sensor_points: 16
+      train_size: 32
+      test_size: 24
+      iterations: 2
+    nonlinear_ode:
+      sensor_points: 16
+      train_size: 24
+      test_size: 16
+      iterations: 2
+    pendulum:
+      sensor_points: 16
+      train_size: 24
+      test_size: 16
+      iterations: 2
+    diffusion_reaction:
+      sensor_points: 16
+      train_functions: 4
+      points_per_function: 8
+      test_functions: 2
+      test_points_per_function: 8
+      test_size: 16
+      space_points: 12
+      time_points: 12
+      iterations: 2

models/DeepONet.py

@@ -0,0 +1,269 @@
+"""Independent PyTorch implementation of the DeepONet in arXiv:1910.03193.
+
+The paper is the architectural authority.  No implementation from the official
+repository is imported or copied.  ReLU and Xavier-normal initialization are
+configurable details used only because the paper leaves them unspecified.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Mapping
+
+import torch
+from torch import Tensor, nn
+
+
+def _activation(name: str) -> nn.Module:
+    choices = {
+        "relu": nn.ReLU,
+        "tanh": nn.Tanh,
+        "gelu": nn.GELU,
+        "silu": nn.SiLU,
+    }
+    try:
+        return choices[name.lower()]()
+    except KeyError as exc:
+        raise ValueError(f"Unsupported activation {name!r}; choose {sorted(choices)}") from exc
+
+
+class DenseNetwork(nn.Module):
+    """A dense network where ``depth`` counts all Linear layers."""
+
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        depth: int,
+        width: int,
+        activation: str,
+        *,
+        activate_output: bool,
+        dense_bias: bool = True,
+        output_bias: bool = True,
+    ) -> None:
+        super().__init__()
+        if depth < 1:
+            raise ValueError("depth must be at least one")
+        if min(input_dim, output_dim, width) < 1:
+            raise ValueError("input_dim, output_dim and width must be positive")
+
+        layers = []
+        current_dim = input_dim
+        for layer_index in range(depth):
+            is_output = layer_index == depth - 1
+            next_dim = output_dim if is_output else width
+            layers.append(
+                nn.Linear(
+                    current_dim,
+                    next_dim,
+                    bias=output_bias if is_output else dense_bias,
+                )
+            )
+            if not is_output or activate_output:
+                layers.append(_activation(activation))
+            current_dim = next_dim
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        return self.layers(inputs)
+
+
+class DeepONet(nn.Module):
+    """Stacked or unstacked DeepONet with the paper's branch/trunk fusion."""
+
+    def __init__(
+        self,
+        branch_input_dim: int,
+        trunk_input_dim: int,
+        latent_dim: int,
+        *,
+        branch_depth: int = 2,
+        trunk_depth: int = 3,
+        width: int = 40,
+        activation: str = "relu",
+        stacked: bool = False,
+        dense_bias: bool = True,
+        branch_output_bias: bool = True,
+        global_bias: bool = True,
+        initializer: str = "xavier_normal",
+    ) -> None:
+        super().__init__()
+        self.branch_input_dim = int(branch_input_dim)
+        self.trunk_input_dim = int(trunk_input_dim)
+        self.latent_dim = int(latent_dim)
+        self.stacked = bool(stacked)
+
+        branch_kwargs = dict(
+            input_dim=self.branch_input_dim,
+            output_dim=1 if self.stacked else self.latent_dim,
+            depth=branch_depth,
+            width=width,
+            activation=activation,
+            activate_output=False,
+            dense_bias=dense_bias,
+            output_bias=branch_output_bias,
+        )
+        if self.stacked:
+            self.branch = nn.ModuleList(
+                DenseNetwork(**branch_kwargs) for _ in range(self.latent_dim)
+            )
+        else:
+            self.branch = DenseNetwork(**branch_kwargs)
+
+        self.trunk = DenseNetwork(
+            input_dim=self.trunk_input_dim,
+            output_dim=self.latent_dim,
+            depth=trunk_depth,
+            width=width,
+            activation=activation,
+            activate_output=True,
+            dense_bias=dense_bias,
+            output_bias=dense_bias,
+        )
+        if global_bias:
+            self.output_bias = nn.Parameter(torch.zeros(1))
+        else:
+            self.register_parameter("output_bias", None)
+        self.reset_parameters(initializer)
+
+    def reset_parameters(self, initializer: str = "xavier_normal") -> None:
+        for module in self.modules():
+            if not isinstance(module, nn.Linear):
+                continue
+            if initializer == "xavier_normal":
+                nn.init.xavier_normal_(module.weight)
+            elif initializer == "xavier_uniform":
+                nn.init.xavier_uniform_(module.weight)
+            else:
+                raise ValueError(f"Unsupported initializer {initializer!r}")
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+    def encode_branch(self, branch_inputs: Tensor) -> Tensor:
+        if branch_inputs.ndim != 2 or branch_inputs.shape[1] != self.branch_input_dim:
+            raise ValueError(
+                f"branch input must have shape [N,{self.branch_input_dim}], "
+                f"got {tuple(branch_inputs.shape)}"
+            )
+        if self.stacked:
+            return torch.cat([head(branch_inputs) for head in self.branch], dim=-1)
+        return self.branch(branch_inputs)
+
+    def forward(self, branch_inputs: Tensor, trunk_inputs: Tensor) -> Tensor:
+        if trunk_inputs.ndim != 2 or trunk_inputs.shape[1] != self.trunk_input_dim:
+            raise ValueError(
+                f"trunk input must have shape [N,{self.trunk_input_dim}], "
+                f"got {tuple(trunk_inputs.shape)}"
+            )
+        if branch_inputs.shape[0] != trunk_inputs.shape[0]:
+            raise ValueError("branch and trunk batches must contain the same number of rows")
+        branch_features = self.encode_branch(branch_inputs)
+        trunk_features = self.trunk(trunk_inputs)
+        prediction = torch.sum(branch_features * trunk_features, dim=-1, keepdim=True)
+        if self.output_bias is not None:
+            prediction = prediction + self.output_bias
+        return prediction
+
+
+class FNNBaseline(nn.Module):
+    """Paper baseline that concatenates sensor values and the query coordinate."""
+
+    def __init__(
+        self,
+        branch_input_dim: int,
+        trunk_input_dim: int,
+        *,
+        depth: int = 3,
+        width: int = 40,
+        activation: str = "relu",
+        output_bias: bool = True,
+        initializer: str = "xavier_normal",
+    ) -> None:
+        super().__init__()
+        self.branch_input_dim = int(branch_input_dim)
+        self.trunk_input_dim = int(trunk_input_dim)
+        self.network = DenseNetwork(
+            input_dim=self.branch_input_dim + self.trunk_input_dim,
+            output_dim=1,
+            depth=depth,
+            width=width,
+            activation=activation,
+            activate_output=False,
+            output_bias=output_bias,
+        )
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                if initializer == "xavier_normal":
+                    nn.init.xavier_normal_(module.weight)
+                elif initializer == "xavier_uniform":
+                    nn.init.xavier_uniform_(module.weight)
+                else:
+                    raise ValueError(f"Unsupported initializer {initializer!r}")
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+
+    def forward(self, branch_inputs: Tensor, trunk_inputs: Tensor) -> Tensor:
+        if branch_inputs.ndim != 2 or branch_inputs.shape[1] != self.branch_input_dim:
+            raise ValueError("invalid branch input shape")
+        if trunk_inputs.ndim != 2 or trunk_inputs.shape[1] != self.trunk_input_dim:
+            raise ValueError("invalid trunk input shape")
+        return self.network(torch.cat((branch_inputs, trunk_inputs), dim=-1))
+
+
+def _merged_model_config(config: Mapping[str, Any], experiment: str) -> Dict[str, Any]:
+    if experiment not in config.get("experiments", {}):
+        raise KeyError(f"Unknown experiment {experiment!r}")
+    merged = dict(config.get("model_defaults", {}))
+    for key in ("branch_depth", "trunk_depth", "width", "latent_dim"):
+        if key in config["experiments"][experiment]:
+            merged[key] = config["experiments"][experiment][key]
+    return merged
+
+
+def build_model(
+    config: Mapping[str, Any],
+    experiment: str,
+    variant: str | None = None,
+) -> nn.Module:
+    """Build a model from the YAML-compatible configuration mapping."""
+
+    experiment_config = config["experiments"][experiment]
+    variant_name = variant or experiment_config["default_variant"]
+    try:
+        variant_config = config["variants"][variant_name]
+    except KeyError as exc:
+        raise KeyError(f"Unknown model variant {variant_name!r}") from exc
+    model_config = _merged_model_config(config, experiment)
+    common = dict(
+        branch_input_dim=int(experiment_config["sensor_points"]),
+        trunk_input_dim=int(experiment_config["trunk_dim"]),
+        activation=str(model_config["activation"]),
+        initializer=str(model_config["initializer"]),
+    )
+    if variant_config["architecture"] == "fnn":
+        return FNNBaseline(
+            **common,
+            depth=int(variant_config["depth"]),
+            width=int(variant_config["width"]),
+            output_bias=bool(variant_config.get("output_bias", True)),
+        )
+    if variant_config["architecture"] != "deeponet":
+        raise ValueError(f"Unsupported architecture {variant_config['architecture']!r}")
+    return DeepONet(
+        **common,
+        latent_dim=int(model_config["latent_dim"]),
+        branch_depth=int(model_config["branch_depth"]),
+        trunk_depth=int(model_config["trunk_depth"]),
+        width=int(model_config["width"]),
+        stacked=bool(variant_config["stacked"]),
+        dense_bias=bool(model_config.get("dense_bias", True)),
+        branch_output_bias=bool(variant_config["branch_output_bias"]),
+        global_bias=bool(variant_config["global_bias"]),
+    )
+
+
+def count_parameters(model: nn.Module) -> int:
+    return sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)
+
+
+__all__ = ["DeepONet", "FNNBaseline", "build_model", "count_parameters"]

models/dataset.py

@@ -0,0 +1,637 @@
+"""Synthetic datasets for the four main experiments in arXiv:1910.03193.
+
+All equations and split rules follow the paper.  The numerical choices that the
+paper omits are configurable and documented in ``config/config.yaml``.
+"""
+
+from __future__ import annotations
+
+import copy
+import hashlib
+import json
+import os
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, Mapping, Sequence, Tuple
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+try:
+    from scipy.integrate import cumulative_trapezoid, solve_ivp
+    from scipy.interpolate import CubicSpline
+except ImportError:  # Deferred error keeps model-only imports usable.
+    cumulative_trapezoid = None
+    solve_ivp = None
+    CubicSpline = None
+
+
+def deep_update(base: Dict[str, Any], override: Mapping[str, Any]) -> Dict[str, Any]:
+    """Recursively update ``base`` without mutating the caller's mapping."""
+
+    for key, value in override.items():
+        if isinstance(value, Mapping) and isinstance(base.get(key), Mapping):
+            base[key] = deep_update(dict(base[key]), value)
+        else:
+            base[key] = copy.deepcopy(value)
+    return base
+
+
+def resolve_config(config: Mapping[str, Any], smoke_test: bool = False) -> Dict[str, Any]:
+    resolved = copy.deepcopy(dict(config))
+    smoke_override = resolved.pop("smoke_test", {})
+    if smoke_test:
+        resolved = deep_update(resolved, smoke_override)
+    resolved.setdefault("project", {})["paper_scale"] = not smoke_test
+    return resolved
+
+
+class OperatorDataset(Dataset):
+    """Triplets with compact storage for repeated PDE branch functions."""
+
+    def __init__(
+        self,
+        branch_functions: np.ndarray,
+        trunk: np.ndarray,
+        target: np.ndarray,
+        function_index: np.ndarray | None = None,
+        metadata: Mapping[str, Any] | None = None,
+    ) -> None:
+        branch_functions = np.asarray(branch_functions, dtype=np.float32)
+        trunk = np.asarray(trunk, dtype=np.float32)
+        target = np.asarray(target, dtype=np.float32)
+        if function_index is None:
+            function_index = np.arange(len(trunk), dtype=np.int64)
+        function_index = np.asarray(function_index, dtype=np.int64)
+        if branch_functions.ndim != 2 or trunk.ndim != 2 or target.ndim != 2:
+            raise ValueError("branch, trunk and target arrays must all have rank two")
+        if target.shape[1] != 1 or len(trunk) != len(target) or len(trunk) != len(function_index):
+            raise ValueError("trunk, target and function_index lengths must agree")
+        if len(function_index) and (
+            function_index.min() < 0 or function_index.max() >= len(branch_functions)
+        ):
+            raise ValueError("function_index refers outside branch_functions")
+        self.branch_functions = branch_functions
+        self.trunk = trunk
+        self.target = target
+        self.function_index = function_index
+        self.metadata = dict(metadata or {})
+
+    def __len__(self) -> int:
+        return len(self.trunk)
+
+    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        function_id = self.function_index[index]
+        return (
+            torch.from_numpy(self.branch_functions[function_id]),
+            torch.from_numpy(self.trunk[index]),
+            torch.from_numpy(self.target[index]),
+        )
+
+    def expanded_branch(self) -> np.ndarray:
+        return self.branch_functions[self.function_index]
+
+
+class FunctionSpaceSampler:
+    """GRF or Chebyshev function sampler on a reusable fine grid."""
+
+    def __init__(self, config: Mapping[str, Any], domain_end: float) -> None:
+        self.config = dict(config)
+        self.domain_end = float(domain_end)
+        self.grid_size = int(self.config["grf_grid_size"])
+        self.grid = np.linspace(0.0, self.domain_end, self.grid_size, dtype=np.float64)
+        self.kind = str(self.config.get("type", "grf")).lower()
+        self._cholesky: np.ndarray | None = None
+        if self.kind == "grf":
+            length_scale = float(self.config["length_scale"])
+            distances = self.grid[:, None] - self.grid[None, :]
+            covariance = np.exp(-(distances**2) / (2.0 * length_scale**2))
+            jitter = float(self.config.get("jitter", 1.0e-13))
+            identity = np.eye(self.grid_size, dtype=np.float64)
+            for attempt in range(6):
+                try:
+                    self._cholesky = np.linalg.cholesky(covariance + jitter * identity)
+                    break
+                except np.linalg.LinAlgError:
+                    jitter *= 10.0
+            if self._cholesky is None:
+                raise np.linalg.LinAlgError("GRF covariance Cholesky failed after jitter fallback")
+        elif self.kind != "chebyshev":
+            raise ValueError(f"Unsupported function space {self.kind!r}")
+
+    def sample(self, count: int, rng: np.random.Generator) -> np.ndarray:
+        if self.kind == "grf":
+            standard_normal = rng.standard_normal((self.grid_size, count))
+            return (self._cholesky @ standard_normal).T
+        cheb = self.config.get("chebyshev", {})
+        degree = int(cheb.get("degree", 10))
+        bound = float(cheb.get("coefficient_bound", 1.0))
+        coefficients = rng.uniform(-bound, bound, size=(count, degree + 1))
+        mapped_grid = 2.0 * self.grid / self.domain_end - 1.0
+        return np.stack(
+            [np.polynomial.chebyshev.chebval(mapped_grid, row) for row in coefficients], axis=0
+        )
+
+    def interpolate(self, values: np.ndarray, points: np.ndarray) -> np.ndarray:
+        values = np.asarray(values, dtype=np.float64)
+        points = np.asarray(points, dtype=np.float64)
+        method = str(self.config.get("interpolation", "cubic")).lower()
+        if method == "cubic":
+            _require_scipy("cubic GRF interpolation")
+            return np.asarray(CubicSpline(self.grid, values, axis=-1)(points))
+        if method != "linear":
+            raise ValueError(f"Unsupported interpolation method {method!r}")
+        if values.ndim == 1:
+            return np.interp(points, self.grid, values)
+        return np.stack([np.interp(points, self.grid, row) for row in values], axis=0)
+
+
+def _require_scipy(operation: str) -> None:
+    if solve_ivp is None or CubicSpline is None or cumulative_trapezoid is None:
+        raise ImportError(f"SciPy is required for {operation}; install it in the execution environment")
+
+
+def _rowwise_linear_interpolation(
+    grid: np.ndarray, values: np.ndarray, points: np.ndarray
+) -> np.ndarray:
+    points = np.clip(np.asarray(points), grid[0], grid[-1])
+    right = np.searchsorted(grid, points, side="right")
+    right = np.clip(right, 1, len(grid) - 1)
+    left = right - 1
+    fraction = (points - grid[left]) / (grid[right] - grid[left])
+    rows = np.arange(len(points))
+    return values[rows, left] * (1.0 - fraction) + values[rows, right] * fraction
+
+
+def solve_antiderivative(
+    input_grid: np.ndarray, input_values: np.ndarray, query_points: np.ndarray
+) -> np.ndarray:
+    _require_scipy("antiderivative reference generation")
+    integral = cumulative_trapezoid(input_values, input_grid, axis=-1, initial=0.0)
+    return _rowwise_linear_interpolation(input_grid, integral, query_points)
+
+
+def solve_nonlinear_ode(
+    input_grid: np.ndarray,
+    input_values: np.ndarray,
+    query_points: np.ndarray,
+    solver_config: Mapping[str, Any],
+) -> np.ndarray:
+    _require_scipy("nonlinear ODE reference generation")
+    query_points = np.asarray(query_points, dtype=np.float64)
+    if query_points.ndim != 1:
+        raise ValueError("query_points must be one-dimensional")
+    if not len(query_points):
+        return np.empty(0, dtype=np.float64)
+    interpolant = CubicSpline(input_grid, input_values)
+    maximum = float(np.max(query_points))
+    if maximum == 0.0:
+        return np.zeros_like(query_points)
+    solution = solve_ivp(
+        lambda x, state: -state**2 + interpolant(x),
+        (0.0, maximum),
+        np.zeros(1, dtype=np.float64),
+        method=str(solver_config.get("method", "RK45")),
+        rtol=float(solver_config.get("rtol", 1.0e-7)),
+        atol=float(solver_config.get("atol", 1.0e-9)),
+        dense_output=True,
+    )
+    if not solution.success:
+        raise RuntimeError(f"nonlinear ODE solve failed: {solution.message}")
+    return np.asarray(solution.sol(query_points)[0])
+
+
+def solve_pendulum(
+    input_grid: np.ndarray,
+    input_values: np.ndarray,
+    query_points: np.ndarray,
+    k: float,
+    solver_config: Mapping[str, Any],
+) -> np.ndarray:
+    _require_scipy("pendulum reference generation")
+    query_points = np.asarray(query_points, dtype=np.float64)
+    if not len(query_points):
+        return np.empty(0, dtype=np.float64)
+    interpolant = CubicSpline(input_grid, input_values)
+    maximum = float(np.max(query_points))
+    if maximum == 0.0:
+        return np.zeros_like(query_points)
+
+    def right_hand_side(time: float, state: np.ndarray) -> np.ndarray:
+        return np.asarray((state[1], -k * np.sin(state[0]) + interpolant(time)))
+
+    solution = solve_ivp(
+        right_hand_side,
+        (0.0, maximum),
+        np.zeros(2, dtype=np.float64),
+        method=str(solver_config.get("method", "RK45")),
+        rtol=float(solver_config.get("rtol", 1.0e-7)),
+        atol=float(solver_config.get("atol", 1.0e-9)),
+        dense_output=True,
+    )
+    if not solution.success:
+        raise RuntimeError(f"pendulum solve failed: {solution.message}")
+    return np.asarray(solution.sol(query_points)[0])
+
+
+def _solve_tridiagonal(
+    lower: np.ndarray, diagonal: np.ndarray, upper: np.ndarray, rhs: np.ndarray
+) -> np.ndarray:
+    """Thomas algorithm for a nonsingular tridiagonal system."""
+
+    lower = np.asarray(lower, dtype=np.float64).copy()
+    diagonal = np.asarray(diagonal, dtype=np.float64).copy()
+    upper = np.asarray(upper, dtype=np.float64).copy()
+    rhs = np.asarray(rhs, dtype=np.float64).copy()
+    for index in range(1, len(diagonal)):
+        if abs(diagonal[index - 1]) < np.finfo(np.float64).eps:
+            raise np.linalg.LinAlgError("zero pivot in tridiagonal solve")
+        multiplier = lower[index - 1] / diagonal[index - 1]
+        diagonal[index] -= multiplier * upper[index - 1]
+        rhs[index] -= multiplier * rhs[index - 1]
+    output = np.empty_like(rhs)
+    output[-1] = rhs[-1] / diagonal[-1]
+    for index in range(len(diagonal) - 2, -1, -1):
+        output[index] = (rhs[index] - upper[index] * output[index + 1]) / diagonal[index]
+    return output
+
+
+def solve_diffusion_reaction(
+    spatial_input: np.ndarray,
+    *,
+    diffusion: float,
+    reaction: float,
+    space_points: int,
+    time_points: int,
+    solver_config: Mapping[str, Any],
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Fully implicit time stepping with second-order centered spatial differences."""
+
+    x_grid = np.linspace(0.0, 1.0, int(space_points), dtype=np.float64)
+    t_grid = np.linspace(0.0, 1.0, int(time_points), dtype=np.float64)
+    source = np.asarray(spatial_input, dtype=np.float64)
+    if source.shape != x_grid.shape:
+        raise ValueError(f"spatial_input must have shape {(len(x_grid),)}, got {source.shape}")
+    field = np.zeros((len(t_grid), len(x_grid)), dtype=np.float64)
+    if len(x_grid) < 3 or len(t_grid) < 2:
+        raise ValueError("PDE grid requires at least 3 spatial and 2 temporal points")
+    dx = x_grid[1] - x_grid[0]
+    dt = t_grid[1] - t_grid[0]
+    ratio = float(diffusion) * dt / (dx * dx)
+    tolerance = float(solver_config.get("pde_newton_tolerance", 1.0e-10))
+    max_iterations = int(solver_config.get("pde_newton_max_iterations", 20))
+    interior_source = source[1:-1]
+    interior_size = len(interior_source)
+    off_diagonal = np.full(interior_size - 1, -ratio, dtype=np.float64)
+
+    for time_index in range(1, len(t_grid)):
+        old = field[time_index - 1, 1:-1]
+        estimate = old.copy()
+        for _ in range(max_iterations):
+            padded = np.pad(estimate, (1, 1), mode="constant")
+            laplacian_term = padded[:-2] - 2.0 * estimate + padded[2:]
+            residual = (
+                estimate
+                - old
+                - ratio * laplacian_term
+                - dt * float(reaction) * estimate**2
+                - dt * interior_source
+            )
+            diagonal = 1.0 + 2.0 * ratio - 2.0 * dt * float(reaction) * estimate
+            update = _solve_tridiagonal(
+                off_diagonal, diagonal, off_diagonal, -residual
+            )
+            estimate += update
+            if np.max(np.abs(update)) <= tolerance:
+                break
+        else:
+            raise RuntimeError(
+                f"PDE Newton solve did not converge at time index {time_index}"
+            )
+        field[time_index, 1:-1] = estimate
+    if not np.isfinite(field).all():
+        raise FloatingPointError("PDE solver produced NaN or infinity")
+    return x_grid, t_grid, field
+
+
+def _generate_ode_like(
+    config: Mapping[str, Any],
+    experiment: str,
+    count: int,
+    seed: int,
+) -> OperatorDataset:
+    experiment_config = config["experiments"][experiment]
+    function_config = config["function_space"]
+    solver_config = config["solver_defaults"]
+    rng = np.random.default_rng(seed)
+    sampler = FunctionSpaceSampler(function_config, float(experiment_config["domain_end"]))
+    sensors = np.linspace(
+        0.0,
+        float(experiment_config["domain_end"]),
+        int(experiment_config["sensor_points"]),
+        dtype=np.float64,
+    )
+    branch = np.empty((count, len(sensors)), dtype=np.float32)
+    trunk = rng.uniform(0.0, float(experiment_config["domain_end"]), size=(count, 1))
+    target = np.empty((count, 1), dtype=np.float32)
+    chunk_size = int(function_config.get("generation_chunk_size", 128))
+    for start in range(0, count, chunk_size):
+        stop = min(count, start + chunk_size)
+        fine_values = sampler.sample(stop - start, rng)
+        branch[start:stop] = sampler.interpolate(fine_values, sensors).astype(np.float32)
+        local_queries = trunk[start:stop, 0]
+        if experiment == "antiderivative":
+            target[start:stop, 0] = solve_antiderivative(
+                sampler.grid, fine_values, local_queries
+            ).astype(np.float32)
+            continue
+        for local_index, values in enumerate(fine_values):
+            query = np.asarray([local_queries[local_index]])
+            if experiment == "nonlinear_ode":
+                answer = solve_nonlinear_ode(sampler.grid, values, query, solver_config)
+            elif experiment == "pendulum":
+                answer = solve_pendulum(
+                    sampler.grid,
+                    values,
+                    query,
+                    float(experiment_config["k"]),
+                    solver_config,
+                )
+            else:
+                raise ValueError(f"Unsupported ODE-like experiment {experiment!r}")
+            target[start + local_index, 0] = answer[0]
+    return OperatorDataset(
+        branch,
+        trunk.astype(np.float32),
+        target,
+        metadata={"experiment": experiment, "seed": seed, "function_count": count},
+    )
+
+
+def _generate_pde(
+    config: Mapping[str, Any],
+    split: str,
+    seed: int,
+) -> OperatorDataset:
+    experiment_config = config["experiments"]["diffusion_reaction"]
+    if split == "train":
+        function_count = int(experiment_config["train_functions"])
+        points_per_function = int(experiment_config["points_per_function"])
+    else:
+        function_count = int(experiment_config["test_functions"])
+        points_per_function = int(experiment_config["test_points_per_function"])
+    rng = np.random.default_rng(seed)
+    sampler = FunctionSpaceSampler(config["function_space"], 1.0)
+    sensors = np.linspace(0.0, 1.0, int(experiment_config["sensor_points"]))
+    branch = np.empty((function_count, len(sensors)), dtype=np.float32)
+    total_points = function_count * points_per_function
+    trunk = np.empty((total_points, 2), dtype=np.float32)
+    target = np.empty((total_points, 1), dtype=np.float32)
+    function_index = np.repeat(np.arange(function_count, dtype=np.int64), points_per_function)
+    chunk_size = int(config["function_space"].get("generation_chunk_size", 128))
+    cursor = 0
+    for start in range(0, function_count, chunk_size):
+        stop = min(function_count, start + chunk_size)
+        fine_batch = sampler.sample(stop - start, rng)
+        branch[start:stop] = sampler.interpolate(fine_batch, sensors).astype(np.float32)
+        for values in fine_batch:
+            spatial_grid = np.linspace(0.0, 1.0, int(experiment_config["space_points"]))
+            spatial_input = sampler.interpolate(values, spatial_grid)
+            x_grid, t_grid, field = solve_diffusion_reaction(
+                spatial_input,
+                diffusion=float(experiment_config["diffusion"]),
+                reaction=float(experiment_config["reaction"]),
+                space_points=int(experiment_config["space_points"]),
+                time_points=int(experiment_config["time_points"]),
+                solver_config=config["solver_defaults"],
+            )
+            grid_size = len(x_grid) * len(t_grid)
+            flat_indices = rng.choice(
+                grid_size,
+                size=points_per_function,
+                replace=points_per_function > grid_size,
+            )
+            time_indices, space_indices = np.divmod(flat_indices, len(x_grid))
+            next_cursor = cursor + points_per_function
+            trunk[cursor:next_cursor, 0] = x_grid[space_indices]
+            trunk[cursor:next_cursor, 1] = t_grid[time_indices]
+            target[cursor:next_cursor, 0] = field[time_indices, space_indices]
+            cursor = next_cursor
+    return OperatorDataset(
+        branch,
+        trunk,
+        target,
+        function_index,
+        metadata={
+            "experiment": "diffusion_reaction",
+            "split": split,
+            "seed": seed,
+            "function_count": function_count,
+            "points_per_function": points_per_function,
+            "group_isolated": True,
+        },
+    )
+
+
+def _fingerprint(config: Mapping[str, Any], experiment: str, split: str, seed: int) -> str:
+    payload = json.dumps(
+        {"config": config, "experiment": experiment, "split": split, "seed": seed},
+        sort_keys=True,
+        separators=(",", ":"),
+    ).encode("utf-8")
+    return hashlib.sha256(payload).hexdigest()
+
+
+def _cache_path(
+    config: Mapping[str, Any], project_root: Path, experiment: str, split: str, seed: int
+) -> Tuple[Path, str]:
+    fingerprint = _fingerprint(config, experiment, split, seed)
+    relative_root = Path(config["paths"]["cache"])
+    root = relative_root if relative_root.is_absolute() else project_root / relative_root
+    return root / f"{experiment}_{split}_{fingerprint[:16]}.npz", fingerprint
+
+
+def _save_cache(path: Path, dataset: OperatorDataset, fingerprint: str) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    metadata = dict(dataset.metadata)
+    metadata["fingerprint"] = fingerprint
+    with tempfile.NamedTemporaryFile(dir=path.parent, suffix=".npz", delete=False) as handle:
+        temporary_path = Path(handle.name)
+    try:
+        np.savez_compressed(
+            temporary_path,
+            branch_functions=dataset.branch_functions,
+            trunk=dataset.trunk,
+            target=dataset.target,
+            function_index=dataset.function_index,
+            metadata=np.asarray(json.dumps(metadata, sort_keys=True)),
+        )
+        os.replace(temporary_path, path)
+    finally:
+        if temporary_path.exists():
+            temporary_path.unlink()
+
+
+def _load_cache(path: Path, fingerprint: str) -> OperatorDataset:
+    with np.load(path, allow_pickle=False) as payload:
+        metadata = json.loads(str(payload["metadata"].item()))
+        if metadata.get("fingerprint") != fingerprint:
+            raise ValueError(f"cache fingerprint mismatch for {path}")
+        return OperatorDataset(
+            payload["branch_functions"],
+            payload["trunk"],
+            payload["target"],
+            payload["function_index"],
+            metadata,
+        )
+
+
+def build_split(
+    config: Mapping[str, Any],
+    experiment: str,
+    split: str,
+    project_root: str | Path,
+    *,
+    use_cache: bool = True,
+) -> OperatorDataset:
+    """Build or load one independent train/test split."""
+
+    if split not in {"train", "test"}:
+        raise ValueError("split must be 'train' or 'test'")
+    if experiment not in config.get("experiments", {}):
+        raise KeyError(f"Unknown experiment {experiment!r}")
+    root = Path(project_root).resolve()
+    base_seed = int(config["runtime"]["seed"])
+    seed = base_seed + (0 if split == "train" else 100_000)
+    path, fingerprint = _cache_path(config, root, experiment, split, seed)
+    if use_cache and path.exists():
+        return _load_cache(path, fingerprint)
+    if experiment == "diffusion_reaction":
+        dataset = _generate_pde(config, split, seed)
+    else:
+        size_key = "train_size" if split == "train" else "test_size"
+        dataset = _generate_ode_like(
+            config, experiment, int(config["experiments"][experiment][size_key]), seed
+        )
+    dataset.metadata.update(
+        {
+            "split": split,
+            "paper_scale": bool(config["project"]["paper_scale"]),
+            "fingerprint": fingerprint,
+        }
+    )
+    if use_cache:
+        _save_cache(path, dataset, fingerprint)
+    return dataset
+
+
+def build_datasets(
+    config: Mapping[str, Any], experiment: str, project_root: str | Path
+) -> Tuple[OperatorDataset, OperatorDataset]:
+    train = build_split(config, experiment, "train", project_root)
+    test = build_split(config, experiment, "test", project_root)
+    return train, test
+
+
+def analytic_input(name: str, coordinates: np.ndarray) -> np.ndarray:
+    if name == "linear":
+        return coordinates
+    if name == "sin_pi":
+        return np.sin(np.pi * coordinates)
+    if name == "sin_2pi":
+        return np.sin(2.0 * np.pi * coordinates)
+    if name == "x_sin_2pi":
+        return coordinates * np.sin(2.0 * np.pi * coordinates)
+    raise KeyError(f"Unknown analytic input {name!r}")
+
+
+def generate_ood_data(
+    config: Mapping[str, Any], experiment: str, query_points: int | None = None
+) -> Dict[str, np.ndarray]:
+    if experiment == "diffusion_reaction":
+        raise ValueError("Use generate_pde_grid_case for the PDE")
+    experiment_config = config["experiments"][experiment]
+    count = int(query_points or config["inference"]["ood_query_points"])
+    domain_end = float(experiment_config["domain_end"])
+    sensors = np.linspace(0.0, domain_end, int(experiment_config["sensor_points"]))
+    queries = np.linspace(0.0, domain_end, count)
+    fine_grid = np.linspace(0.0, domain_end, max(1000, count))
+    names = experiment_config.get("ood_functions", ["linear", "sin_pi", "sin_2pi"])
+    all_branch, all_trunk, all_target, all_labels = [], [], [], []
+    for name in names:
+        fine_values = analytic_input(str(name), fine_grid)
+        branch = analytic_input(str(name), sensors)
+        if experiment == "antiderivative":
+            _require_scipy("antiderivative OOD reference")
+            integral = cumulative_trapezoid(fine_values, fine_grid, initial=0.0)
+            target = np.interp(queries, fine_grid, integral)
+        elif experiment == "nonlinear_ode":
+            target = solve_nonlinear_ode(
+                fine_grid, fine_values, queries, config["solver_defaults"]
+            )
+        elif experiment == "pendulum":
+            target = solve_pendulum(
+                fine_grid,
+                fine_values,
+                queries,
+                float(experiment_config["k"]),
+                config["solver_defaults"],
+            )
+        else:
+            raise ValueError(f"Unsupported experiment {experiment!r}")
+        all_branch.append(np.repeat(branch[None, :], count, axis=0))
+        all_trunk.append(queries[:, None])
+        all_target.append(target[:, None])
+        all_labels.extend([str(name)] * count)
+    return {
+        "branch": np.concatenate(all_branch).astype(np.float32),
+        "trunk": np.concatenate(all_trunk).astype(np.float32),
+        "target": np.concatenate(all_target).astype(np.float32),
+        "labels": np.asarray(all_labels),
+    }
+
+
+def generate_pde_grid_case(config: Mapping[str, Any], seed: int) -> Dict[str, np.ndarray]:
+    experiment_config = config["experiments"]["diffusion_reaction"]
+    sampler = FunctionSpaceSampler(config["function_space"], 1.0)
+    rng = np.random.default_rng(seed)
+    fine_values = sampler.sample(1, rng)[0]
+    sensors = np.linspace(0.0, 1.0, int(experiment_config["sensor_points"]))
+    branch_vector = sampler.interpolate(fine_values, sensors).astype(np.float32)
+    spatial_grid = np.linspace(0.0, 1.0, int(experiment_config["space_points"]))
+    source = sampler.interpolate(fine_values, spatial_grid)
+    x_grid, t_grid, field = solve_diffusion_reaction(
+        source,
+        diffusion=float(experiment_config["diffusion"]),
+        reaction=float(experiment_config["reaction"]),
+        space_points=int(experiment_config["space_points"]),
+        time_points=int(experiment_config["time_points"]),
+        solver_config=config["solver_defaults"],
+    )
+    x_mesh, t_mesh = np.meshgrid(x_grid, t_grid)
+    trunk = np.column_stack((x_mesh.ravel(), t_mesh.ravel())).astype(np.float32)
+    return {
+        "branch": np.repeat(branch_vector[None, :], len(trunk), axis=0),
+        "trunk": trunk,
+        "target": field.reshape(-1, 1).astype(np.float32),
+        "source": source.astype(np.float32),
+        "x": x_grid.astype(np.float32),
+        "t": t_grid.astype(np.float32),
+        "grid_shape": np.asarray(field.shape, dtype=np.int64),
+    }
+
+
+__all__ = [
+    "OperatorDataset",
+    "FunctionSpaceSampler",
+    "resolve_config",
+    "build_split",
+    "build_datasets",
+    "generate_ood_data",
+    "generate_pde_grid_case",
+    "solve_antiderivative",
+    "solve_nonlinear_ode",
+    "solve_pendulum",
+    "solve_diffusion_reaction",
+]

scripts/inference.py

@@ -0,0 +1,269 @@
+#!/usr/bin/env python3
+"""Inference for random-test, OOD-curve, and PDE-grid DeepONet cases."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import os
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, Mapping, Tuple
+
+import numpy as np
+import torch
+import yaml
+from torch.utils.data import DataLoader
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from models.DeepONet import build_model  # noqa: E402
+from models.dataset import (  # noqa: E402
+    OperatorDataset,
+    build_split,
+    generate_ood_data,
+    generate_pde_grid_case,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "config/config.yaml")
+    parser.add_argument("--checkpoint", type=Path, default=None)
+    parser.add_argument("--experiment", default="antiderivative")
+    parser.add_argument("--variant", default="unstacked_bias")
+    parser.add_argument("--mode", choices=("random_test", "ood", "pde_grid"), default="random_test")
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--batch-size", type=int, default=None)
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--output-dir", type=Path, default=None)
+    parser.add_argument("--smoke-test", action="store_true")
+    return parser.parse_args()
+
+
+def load_yaml(path: Path) -> Dict[str, Any]:
+    with path.expanduser().resolve().open("r", encoding="utf-8") as handle:
+        value = yaml.safe_load(handle)
+    if not isinstance(value, dict):
+        raise ValueError(f"Configuration {path} must contain a mapping")
+    return value
+
+
+def resolve_path(value: str | Path) -> Path:
+    path = Path(value).expanduser()
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def torch_load(path: Path, map_location: str | torch.device = "cpu") -> Any:
+    try:
+        return torch.load(path, map_location=map_location, weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location=map_location)
+
+
+def load_entry(checkpoint_path: Path, experiment: str, variant: str) -> Mapping[str, Any]:
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(
+            f"Checkpoint {checkpoint_path} does not exist. Train the requested entry first; "
+            "inference never evaluates random weights."
+        )
+    bundle = torch_load(checkpoint_path)
+    if not isinstance(bundle, dict) or not isinstance(bundle.get("entries"), dict):
+        raise ValueError(f"Checkpoint {checkpoint_path} is not a DeepONet indexed bundle")
+    key = f"{experiment}/{variant}"
+    if key not in bundle["entries"]:
+        available = sorted(bundle["entries"])
+        raise KeyError(f"Checkpoint entry {key!r} is missing; available entries: {available}")
+    entry = bundle["entries"][key]
+    if not isinstance(entry.get("run_config"), dict):
+        raise ValueError(f"Checkpoint entry {key!r} has no resolved run_config")
+    return entry
+
+
+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"CUDA device {requested!r} requested but CUDA is unavailable")
+    return device
+
+
+def metrics(prediction: np.ndarray, target: np.ndarray, trim_fraction: float) -> Dict[str, float]:
+    error = prediction.astype(np.float64).reshape(-1) - target.astype(np.float64).reshape(-1)
+    squared = error**2
+    output = {
+        "test_mse": float(np.mean(squared)),
+        "relative_l2": float(
+            np.linalg.norm(error)
+            / max(np.linalg.norm(target.astype(np.float64).reshape(-1)), np.finfo(np.float64).eps)
+        ),
+    }
+    if trim_fraction > 0.0:
+        remove_count = min(len(squared) - 1, int(np.ceil(len(squared) * trim_fraction)))
+        kept_count = len(squared) - remove_count
+        output["trimmed_test_mse"] = float(np.mean(np.partition(squared, kept_count - 1)[:kept_count]))
+    return output
+
+
+@torch.inference_mode()
+def predict_dataset(
+    model: torch.nn.Module,
+    dataset: OperatorDataset,
+    device: torch.device,
+    batch_size: int,
+) -> Tuple[np.ndarray, np.ndarray]:
+    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)
+    predictions, targets = [], []
+    model.eval()
+    for branch, trunk, target in loader:
+        prediction = model(branch.to(device), trunk.to(device))
+        predictions.append(prediction.detach().cpu().numpy())
+        targets.append(target.numpy())
+    prediction_array = np.concatenate(predictions).astype(np.float32)
+    target_array = np.concatenate(targets).astype(np.float32)
+    if prediction_array.shape != target_array.shape or prediction_array.ndim != 2:
+        raise ValueError(
+            f"prediction/target shape mismatch: {prediction_array.shape} versus {target_array.shape}"
+        )
+    if not np.isfinite(prediction_array).all():
+        raise FloatingPointError("model prediction contains NaN or infinity")
+    return prediction_array, target_array
+
+
+def _safe_output_dir(
+    requested: Path | None,
+    results_root: Path,
+    experiment: str,
+    variant: str,
+    mode: str,
+) -> Path:
+    output = results_root / experiment / variant / mode if requested is None else resolve_path(requested)
+    results_root = results_root.resolve()
+    output = output.resolve()
+    if os.path.commonpath((str(results_root), str(output))) != str(results_root):
+        raise ValueError(f"Inference outputs must remain below {results_root}, got {output}")
+    output.mkdir(parents=True, exist_ok=True)
+    return output
+
+
+def _atomic_json(path: Path, value: Mapping[str, Any]) -> None:
+    with tempfile.NamedTemporaryFile(
+        mode="w", encoding="utf-8", dir=path.parent, suffix=".json", delete=False
+    ) as handle:
+        json.dump(value, handle, indent=2, sort_keys=True)
+        handle.write("\n")
+        temporary_path = Path(handle.name)
+    try:
+        os.replace(temporary_path, path)
+    finally:
+        if temporary_path.exists():
+            temporary_path.unlink()
+
+
+def _atomic_npz(path: Path, arrays: Mapping[str, np.ndarray]) -> None:
+    with tempfile.NamedTemporaryFile(dir=path.parent, suffix=".npz", delete=False) as handle:
+        temporary_path = Path(handle.name)
+    try:
+        np.savez_compressed(temporary_path, **arrays)
+        os.replace(temporary_path, path)
+    finally:
+        if temporary_path.exists():
+            temporary_path.unlink()
+
+
+def final_train_loss(results_root: Path, experiment: str, variant: str) -> float | None:
+    history_path = results_root / experiment / variant / "history.csv"
+    if not history_path.exists():
+        return None
+    with history_path.open("r", encoding="utf-8", newline="") as handle:
+        rows = list(csv.DictReader(handle))
+    return float(rows[-1]["train_loss"]) if rows else None
+
+
+def prepare_data(
+    config: Mapping[str, Any], experiment: str, mode: str, seed: int
+) -> Tuple[OperatorDataset, Dict[str, np.ndarray]]:
+    if mode == "random_test":
+        dataset = build_split(config, experiment, "test", PROJECT_ROOT)
+        return dataset, {
+            "branch_functions": dataset.branch_functions,
+            "function_index": dataset.function_index,
+            "trunk": dataset.trunk,
+        }
+    if mode == "ood":
+        payload = generate_ood_data(config, experiment)
+        dataset = OperatorDataset(payload["branch"], payload["trunk"], payload["target"])
+        return dataset, {"branch": payload["branch"], "trunk": payload["trunk"], "labels": payload["labels"]}
+    if experiment != "diffusion_reaction":
+        raise ValueError("pde_grid mode is only valid for diffusion_reaction")
+    payload = generate_pde_grid_case(config, seed)
+    dataset = OperatorDataset(payload["branch"], payload["trunk"], payload["target"])
+    extras = {key: value for key, value in payload.items() if key != "target"}
+    return dataset, extras
+
+
+def main() -> None:
+    args = parse_args()
+    file_config = load_yaml(args.config)
+    checkpoint_path = resolve_path(args.checkpoint or file_config["paths"]["checkpoint"])
+    entry = load_entry(checkpoint_path, args.experiment, args.variant)
+    config = entry["run_config"]
+    stored_paper_scale = bool(entry.get("paper_scale", config["project"].get("paper_scale", True)))
+    if args.smoke_test and stored_paper_scale:
+        raise ValueError("--smoke-test was requested, but the selected checkpoint is paper-scale")
+    device_name = args.device or str(config["runtime"].get("device", "auto"))
+    device = select_device(device_name)
+    model = build_model(config, args.experiment, args.variant).to(device)
+    model.load_state_dict(entry["model_state"], strict=True)
+    seed = int(args.seed if args.seed is not None else config["runtime"]["seed"] + 200_000)
+    dataset, extra_arrays = prepare_data(config, args.experiment, args.mode, seed)
+    batch_size = int(args.batch_size or config["inference"]["batch_size"])
+    prediction, target = predict_dataset(model, dataset, device, batch_size)
+    trim_fraction = float(config["experiments"][args.experiment].get("trim_fraction", 0.0))
+    result_metrics = metrics(prediction, target, trim_fraction)
+
+    results_root = resolve_path(config["paths"]["results"])
+    output_dir = _safe_output_dir(
+        args.output_dir, results_root, args.experiment, args.variant, args.mode
+    )
+    train_loss = final_train_loss(results_root, args.experiment, args.variant)
+    if train_loss is not None and args.mode == "random_test":
+        result_metrics["generalization_error"] = result_metrics["test_mse"] - train_loss
+    metadata = {
+        "experiment": args.experiment,
+        "variant": args.variant,
+        "mode": args.mode,
+        "checkpoint": str(checkpoint_path),
+        "checkpoint_iteration": int(entry["iteration"]),
+        "checkpoint_best_metric": float(entry["best_metric"]),
+        "seed": seed,
+        "paper_scale": stored_paper_scale,
+        "sample_count": len(dataset),
+    }
+    arrays: Dict[str, np.ndarray] = {
+        "prediction": prediction,
+        "target": target,
+        "metadata": np.asarray(json.dumps(metadata, sort_keys=True)),
+    }
+    arrays.update(extra_arrays)
+    _atomic_npz(output_dir / "predictions.npz", arrays)
+    _atomic_json(output_dir / "metrics.json", {**metadata, **result_metrics})
+    metric_text = " ".join(f"{name}={value:.8e}" for name, value in result_metrics.items())
+    print(
+        f"INFERENCE experiment={args.experiment} variant={args.variant} mode={args.mode} "
+        f"paper_scale={stored_paper_scale} {metric_text}",
+        flush=True,
+    )
+    print(
+        f"SAVED metrics={output_dir / 'metrics.json'} predictions={output_dir / 'predictions.npz'}",
+        flush=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,236 @@
+#!/usr/bin/env python3
+"""Create figures and a summary from existing DeepONet result artifacts."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Mapping, Tuple
+
+import numpy as np
+import yaml
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "config/config.yaml")
+    parser.add_argument("--experiment", default=None)
+    parser.add_argument("--variant", default=None)
+    return parser.parse_args()
+
+
+def load_config(path: Path) -> Dict[str, Any]:
+    with path.expanduser().resolve().open("r", encoding="utf-8") as handle:
+        value = yaml.safe_load(handle)
+    if not isinstance(value, dict):
+        raise ValueError(f"Configuration {path} must contain a mapping")
+    return value
+
+
+def resolve_path(value: str | Path) -> Path:
+    path = Path(value).expanduser()
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def candidate_runs(results_root: Path, experiment: str | None, variant: str | None) -> Iterable[Tuple[str, str, Path]]:
+    experiment_dirs = [results_root / experiment] if experiment else sorted(results_root.iterdir())
+    for experiment_dir in experiment_dirs:
+        if not experiment_dir.is_dir() or experiment_dir.name in {"cache", "figures"}:
+            continue
+        variant_dirs = [experiment_dir / variant] if variant else sorted(experiment_dir.iterdir())
+        for variant_dir in variant_dirs:
+            if variant_dir.is_dir():
+                yield experiment_dir.name, variant_dir.name, variant_dir
+
+
+def read_history(path: Path) -> Dict[str, np.ndarray]:
+    with path.open("r", encoding="utf-8", newline="") as handle:
+        rows = list(csv.DictReader(handle))
+    if not rows:
+        raise ValueError(f"History {path} is empty")
+    columns: Dict[str, np.ndarray] = {}
+    for key in rows[0]:
+        if key == "paper_scale":
+            continue
+        try:
+            columns[key] = np.asarray([float(row[key]) for row in rows if row.get(key, "") != ""])
+        except ValueError:
+            continue
+    return columns
+
+
+def plot_history(history_path: Path, output_path: Path, title: str) -> None:
+    columns = read_history(history_path)
+    iterations = columns["iteration"]
+    fig, axis = plt.subplots(figsize=(7.2, 4.5))
+    axis.plot(iterations, columns["train_loss"], label="train MSE", linewidth=1.8)
+    axis.plot(iterations, columns["test_mse"], label="test MSE", linewidth=1.8)
+    if "trimmed_test_mse" in columns:
+        axis.plot(iterations, columns["trimmed_test_mse"], label="trimmed test MSE", linewidth=1.5)
+    axis.set_yscale("log")
+    axis.set_xlabel("Iteration")
+    axis.set_ylabel("Mean squared error")
+    axis.set_title(title)
+    axis.grid(True, which="both", alpha=0.25)
+    axis.legend()
+    fig.tight_layout()
+    fig.savefig(output_path, dpi=180)
+    plt.close(fig)
+
+
+def _metadata(payload: Mapping[str, np.ndarray]) -> Dict[str, Any]:
+    return json.loads(str(payload["metadata"].item())) if "metadata" in payload else {}
+
+
+def plot_ood(payload: Mapping[str, np.ndarray], output_path: Path, title: str) -> None:
+    labels = payload["labels"].astype(str)
+    trunk = payload["trunk"][:, 0]
+    target = payload["target"][:, 0]
+    prediction = payload["prediction"][:, 0]
+    unique_labels = list(dict.fromkeys(labels.tolist()))
+    fig, axes = plt.subplots(len(unique_labels), 1, figsize=(7.2, 3.0 * len(unique_labels)), squeeze=False)
+    for axis, label in zip(axes[:, 0], unique_labels):
+        selected = labels == label
+        order = np.argsort(trunk[selected])
+        axis.plot(trunk[selected][order], target[selected][order], label="reference", linewidth=2.0)
+        axis.plot(trunk[selected][order], prediction[selected][order], "--", label="DeepONet", linewidth=1.8)
+        axis.set_title(label)
+        axis.set_xlabel("Query coordinate")
+        axis.set_ylabel("Solution")
+        axis.grid(True, alpha=0.25)
+        axis.legend()
+    fig.suptitle(title)
+    fig.tight_layout()
+    fig.savefig(output_path, dpi=180)
+    plt.close(fig)
+
+
+def plot_pde(payload: Mapping[str, np.ndarray], output_path: Path, title: str) -> None:
+    shape = tuple(int(value) for value in payload["grid_shape"])
+    target = payload["target"].reshape(shape)
+    prediction = payload["prediction"].reshape(shape)
+    absolute_error = np.abs(prediction - target)
+    x_grid = payload["x"]
+    t_grid = payload["t"]
+    extent = (float(x_grid.min()), float(x_grid.max()), float(t_grid.min()), float(t_grid.max()))
+    shared_min = float(min(target.min(), prediction.min()))
+    shared_max = float(max(target.max(), prediction.max()))
+    fig, axes = plt.subplots(1, 3, figsize=(13.5, 3.8), constrained_layout=True)
+    first = axes[0].imshow(target, origin="lower", aspect="auto", extent=extent, vmin=shared_min, vmax=shared_max)
+    axes[0].set_title("Reference")
+    axes[1].imshow(prediction, origin="lower", aspect="auto", extent=extent, vmin=shared_min, vmax=shared_max)
+    axes[1].set_title("Prediction")
+    error_image = axes[2].imshow(absolute_error, origin="lower", aspect="auto", extent=extent)
+    axes[2].set_title("Absolute error")
+    for axis in axes:
+        axis.set_xlabel("x")
+        axis.set_ylabel("t")
+    fig.colorbar(first, ax=axes[:2], shrink=0.82, label="s(x,t)")
+    fig.colorbar(error_image, ax=axes[2], shrink=0.82, label="absolute error")
+    fig.suptitle(title)
+    fig.savefig(output_path, dpi=180)
+    plt.close(fig)
+
+
+def plot_generic(payload: Mapping[str, np.ndarray], output_path: Path, title: str) -> None:
+    prediction = payload["prediction"].reshape(-1)
+    target = payload["target"].reshape(-1)
+    fig, axes = plt.subplots(1, 2, figsize=(10.5, 4.2))
+    limit_min = float(min(prediction.min(), target.min()))
+    limit_max = float(max(prediction.max(), target.max()))
+    axes[0].scatter(target, prediction, s=7, alpha=0.45)
+    axes[0].plot((limit_min, limit_max), (limit_min, limit_max), "k--", linewidth=1)
+    axes[0].set_xlabel("Reference")
+    axes[0].set_ylabel("Prediction")
+    axes[0].set_title("Prediction parity")
+    axes[1].hist(np.abs(prediction - target), bins=40)
+    axes[1].set_xlabel("Absolute error")
+    axes[1].set_ylabel("Count")
+    axes[1].set_title("Error distribution")
+    fig.suptitle(title)
+    fig.tight_layout()
+    fig.savefig(output_path, dpi=180)
+    plt.close(fig)
+
+
+def write_summary(path: Path, records: List[Mapping[str, Any]]) -> None:
+    with path.open("w", encoding="utf-8") as handle:
+        json.dump({"runs": records}, handle, indent=2, sort_keys=True)
+        handle.write("\n")
+
+
+def main() -> None:
+    args = parse_args()
+    config = load_config(args.config)
+    results_root = resolve_path(config["paths"]["results"])
+    if not results_root.exists():
+        raise FileNotFoundError(f"Results directory {results_root} does not exist")
+    figures_root = results_root / "figures"
+    figures_root.mkdir(parents=True, exist_ok=True)
+    records: List[Mapping[str, Any]] = []
+    artifact_count = 0
+    for experiment, variant, run_dir in candidate_runs(results_root, args.experiment, args.variant):
+        title = f"{experiment} / {variant}"
+        history_path = run_dir / "history.csv"
+        if history_path.exists():
+            output = figures_root / f"{experiment}_{variant}_loss.png"
+            plot_history(history_path, output, title)
+            print(f"FIGURE {output}", flush=True)
+            artifact_count += 1
+        mode_dirs = sorted(
+            path
+            for path in run_dir.iterdir()
+            if path.is_dir() and ((path / "predictions.npz").exists() or (path / "metrics.json").exists())
+        )
+        # Backward-compatible fallback for results created before mode directories were introduced.
+        artifact_dirs = mode_dirs or [run_dir]
+        for artifact_dir in artifact_dirs:
+            prediction_path = artifact_dir / "predictions.npz"
+            metadata: Dict[str, Any] = {}
+            if prediction_path.exists():
+                with np.load(prediction_path, allow_pickle=False) as payload_file:
+                    payload = {key: payload_file[key] for key in payload_file.files}
+                metadata = _metadata(payload)
+                suffix = str(metadata.get("mode", artifact_dir.name))
+                output = figures_root / f"{experiment}_{variant}_{suffix}.png"
+                mode_title = f"{title} / {suffix}"
+                if "labels" in payload:
+                    plot_ood(payload, output, mode_title)
+                elif "grid_shape" in payload:
+                    plot_pde(payload, output, mode_title)
+                else:
+                    plot_generic(payload, output, mode_title)
+                print(f"FIGURE {output}", flush=True)
+                artifact_count += 1
+            metrics_path = artifact_dir / "metrics.json"
+            if metrics_path.exists():
+                with metrics_path.open("r", encoding="utf-8") as handle:
+                    record = json.load(handle)
+                records.append(record)
+                metric_names = ("test_mse", "trimmed_test_mse", "relative_l2", "generalization_error")
+                metric_text = " ".join(
+                    f"{name}={float(record[name]):.8e}" for name in metric_names if name in record
+                )
+                print(
+                    f"METRICS experiment={experiment} variant={variant} "
+                    f"mode={record.get('mode', metadata.get('mode', artifact_dir.name))} {metric_text}",
+                    flush=True,
+                )
+    if artifact_count == 0 and not records:
+        raise FileNotFoundError("No history.csv, predictions.npz, or metrics.json matched the request")
+    summary_path = figures_root / "summary.json"
+    write_summary(summary_path, records)
+    print(f"SUMMARY {summary_path} runs={len(records)} figures={artifact_count}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,369 @@
+#!/usr/bin/env python3
+"""Train the paper-scale or explicitly reduced DeepONet experiments."""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import csv
+import os
+import random
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, Iterable, Mapping, Tuple
+
+import numpy as np
+import torch
+import yaml
+from torch import Tensor, nn
+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.DeepONet import build_model, count_parameters  # noqa: E402
+from models.dataset import OperatorDataset, build_datasets, resolve_config  # noqa: E402
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "config/config.yaml")
+    parser.add_argument(
+        "--experiment",
+        default="antiderivative",
+        help="One experiment name or 'all' for all four paper experiments.",
+    )
+    parser.add_argument(
+        "--variant",
+        default=None,
+        help="Variant name, 'all', or omit to use each experiment's paper default.",
+    )
+    parser.add_argument("--smoke-test", action="store_true", help="Run reduced non-paper settings.")
+    parser.add_argument("--device", default=None, help="cpu, cuda, cuda:N, or auto")
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--resume", action="store_true")
+    parser.add_argument("--no-cache", action="store_true")
+    return parser.parse_args()
+
+
+def load_config(path: Path, smoke_test: bool) -> Dict[str, Any]:
+    with path.expanduser().resolve().open("r", encoding="utf-8") as handle:
+        raw = yaml.safe_load(handle)
+    if not isinstance(raw, Mapping):
+        raise ValueError(f"Configuration {path} must contain a mapping")
+    return resolve_config(raw, smoke_test=smoke_test)
+
+
+def resolve_path(path_value: str, root: Path = PROJECT_ROOT) -> Path:
+    path = Path(path_value).expanduser()
+    return path if path.is_absolute() else root / path
+
+
+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"CUDA device {requested!r} requested but CUDA is unavailable")
+    return device
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def make_loader(
+    dataset: OperatorDataset,
+    batch_size: int | None,
+    *,
+    shuffle: bool,
+    num_workers: int,
+    pin_memory: bool,
+) -> DataLoader:
+    effective_batch = len(dataset) if batch_size is None else int(batch_size)
+    if effective_batch < 1:
+        raise ValueError("batch size must be positive")
+    return DataLoader(
+        dataset,
+        batch_size=effective_batch,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        drop_last=False,
+    )
+
+
+def metric_values(prediction: np.ndarray, target: np.ndarray, trim_fraction: float = 0.0) -> Dict[str, float]:
+    squared_error = np.square(prediction.astype(np.float64) - target.astype(np.float64)).reshape(-1)
+    mse = float(np.mean(squared_error))
+    denominator = float(np.linalg.norm(target.astype(np.float64).reshape(-1)))
+    relative_l2 = float(
+        np.linalg.norm(prediction.astype(np.float64).reshape(-1) - target.astype(np.float64).reshape(-1))
+        / max(denominator, np.finfo(np.float64).eps)
+    )
+    metrics = {"test_mse": mse, "relative_l2": relative_l2}
+    if trim_fraction > 0.0:
+        remove_count = min(len(squared_error) - 1, int(np.ceil(len(squared_error) * trim_fraction)))
+        kept = np.partition(squared_error, len(squared_error) - remove_count - 1)[
+            : len(squared_error) - remove_count
+        ]
+        metrics["trimmed_test_mse"] = float(np.mean(kept))
+    return metrics
+
+
+@torch.inference_mode()
+def evaluate(
+    model: nn.Module,
+    dataset: OperatorDataset,
+    device: torch.device,
+    batch_size: int,
+    trim_fraction: float,
+    num_workers: int,
+) -> Dict[str, float]:
+    loader = make_loader(
+        dataset,
+        batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=device.type == "cuda",
+    )
+    predictions, targets = [], []
+    model.eval()
+    for branch, trunk, target in loader:
+        output = model(branch.to(device), trunk.to(device))
+        predictions.append(output.detach().cpu().numpy())
+        targets.append(target.numpy())
+    return metric_values(np.concatenate(predictions), np.concatenate(targets), trim_fraction)
+
+
+def _cpu_copy(value: Any) -> Any:
+    if isinstance(value, Tensor):
+        return value.detach().cpu()
+    if isinstance(value, dict):
+        return {key: _cpu_copy(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_cpu_copy(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(_cpu_copy(item) for item in value)
+    return value
+
+
+def torch_load(path: Path, map_location: str | torch.device = "cpu") -> Any:
+    try:
+        return torch.load(path, map_location=map_location, weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location=map_location)
+
+
+def load_bundle(path: Path) -> Dict[str, Any]:
+    if not path.exists():
+        return {"format_version": 1, "entries": {}}
+    bundle = torch_load(path)
+    if not isinstance(bundle, dict) or not isinstance(bundle.get("entries"), dict):
+        raise ValueError(f"Checkpoint {path} is not a DeepONet indexed bundle")
+    return bundle
+
+
+def save_bundle_entry(path: Path, key: str, entry: Mapping[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    bundle = load_bundle(path)
+    bundle["format_version"] = 1
+    bundle["entries"][key] = _cpu_copy(dict(entry))
+    with tempfile.NamedTemporaryFile(dir=path.parent, suffix=".pth", delete=False) as handle:
+        temporary_path = Path(handle.name)
+    try:
+        torch.save(bundle, temporary_path)
+        os.replace(temporary_path, path)
+    finally:
+        if temporary_path.exists():
+            temporary_path.unlink()
+
+
+def append_history(path: Path, row: Mapping[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    exists = path.exists()
+    with path.open("a", newline="", encoding="utf-8") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(row.keys()))
+        if not exists:
+            writer.writeheader()
+        writer.writerow(row)
+        handle.flush()
+
+
+def variants_for(config: Mapping[str, Any], experiment: str, requested: str | None) -> Iterable[str]:
+    if requested is None:
+        return [str(config["experiments"][experiment]["default_variant"])]
+    if requested == "all":
+        return list(config["variants"])
+    if requested not in config["variants"]:
+        raise KeyError(f"Unknown variant {requested!r}; choose from {list(config['variants'])}")
+    return [requested]
+
+
+def run_training(
+    config: Dict[str, Any],
+    experiment: str,
+    variant: str,
+    *,
+    no_cache: bool,
+    resume: bool,
+) -> None:
+    experiment_config = config["experiments"][experiment]
+    training_config = config["training_defaults"]
+    device = select_device(str(config["runtime"]["device"]))
+    seed = int(config["runtime"]["seed"])
+    set_seed(seed)
+
+    if no_cache:
+        from models.dataset import build_split
+
+        train_dataset = build_split(config, experiment, "train", PROJECT_ROOT, use_cache=False)
+        test_dataset = build_split(config, experiment, "test", PROJECT_ROOT, use_cache=False)
+    else:
+        train_dataset, test_dataset = build_datasets(config, experiment, PROJECT_ROOT)
+
+    model = build_model(config, experiment, variant).to(device)
+    optimizer = torch.optim.Adam(
+        model.parameters(),
+        lr=float(training_config["learning_rate"]),
+        weight_decay=float(training_config.get("weight_decay", 0.0)),
+    )
+    criterion = nn.MSELoss(reduction="mean")
+    checkpoint_path = resolve_path(str(config["paths"]["checkpoint"]))
+    entry_key = f"{experiment}/{variant}"
+    start_iteration = 0
+    best_metric = float("inf")
+    if resume and checkpoint_path.exists():
+        entry = load_bundle(checkpoint_path)["entries"].get(entry_key)
+        if entry is None:
+            raise KeyError(f"Cannot resume: {entry_key!r} is absent from {checkpoint_path}")
+        model.load_state_dict(entry["model_state"], strict=True)
+        optimizer.load_state_dict(entry["optimizer_state"])
+        start_iteration = int(entry["iteration"])
+        best_metric = float(entry["best_metric"])
+
+    batch_size = training_config.get("batch_size")
+    train_loader = make_loader(
+        train_dataset,
+        None if batch_size is None else int(batch_size),
+        shuffle=True,
+        num_workers=int(config["runtime"].get("num_workers", 0)),
+        pin_memory=device.type == "cuda",
+    )
+    train_iterator = iter(train_loader)
+    iterations = int(experiment_config["iterations"])
+    print_every = int(training_config["print_every"])
+    evaluate_every = int(training_config["evaluate_every"])
+    eval_batch_size = int(training_config["evaluation_batch_size"])
+    trim_fraction = float(experiment_config.get("trim_fraction", 0.0))
+    history_path = resolve_path(str(config["paths"]["results"])) / experiment / variant / "history.csv"
+    paper_scale = bool(config["project"]["paper_scale"])
+    if not resume and history_path.exists():
+        history_path.unlink()
+
+    print(
+        f"START experiment={experiment} variant={variant} device={device} "
+        f"parameters={count_parameters(model)} train_points={len(train_dataset)} "
+        f"test_points={len(test_dataset)} iterations={iterations} paper_scale={paper_scale}",
+        flush=True,
+    )
+    for iteration in range(start_iteration + 1, iterations + 1):
+        try:
+            branch, trunk, target = next(train_iterator)
+        except StopIteration:
+            train_iterator = iter(train_loader)
+            branch, trunk, target = next(train_iterator)
+        model.train()
+        optimizer.zero_grad(set_to_none=True)
+        prediction = model(branch.to(device), trunk.to(device))
+        loss = criterion(prediction, target.to(device))
+        if not torch.isfinite(loss):
+            raise FloatingPointError(f"non-finite training loss at iteration {iteration}")
+        loss.backward()
+        optimizer.step()
+        train_loss = float(loss.detach().cpu())
+
+        should_evaluate = iteration == 1 or iteration % evaluate_every == 0 or iteration == iterations
+        should_print = iteration == 1 or iteration % print_every == 0 or should_evaluate
+        metrics: Dict[str, float] = {}
+        if should_evaluate:
+            metrics = evaluate(
+                model,
+                test_dataset,
+                device,
+                eval_batch_size,
+                trim_fraction,
+                int(config["runtime"].get("num_workers", 0)),
+            )
+            row: Dict[str, Any] = {
+                "iteration": iteration,
+                "train_loss": train_loss,
+                "test_mse": metrics["test_mse"],
+                "relative_l2": metrics["relative_l2"],
+                "generalization_error": metrics["test_mse"] - train_loss,
+                "paper_scale": paper_scale,
+            }
+            if "trimmed_test_mse" in metrics:
+                row["trimmed_test_mse"] = metrics["trimmed_test_mse"]
+            append_history(history_path, row)
+            if metrics["test_mse"] < best_metric:
+                best_metric = metrics["test_mse"]
+                save_bundle_entry(
+                    checkpoint_path,
+                    entry_key,
+                    {
+                        "experiment": experiment,
+                        "variant": variant,
+                        "iteration": iteration,
+                        "best_metric": best_metric,
+                        "metric_name": "test_mse",
+                        "model_state": model.state_dict(),
+                        "optimizer_state": optimizer.state_dict(),
+                        "run_config": copy.deepcopy(config),
+                        "paper_scale": paper_scale,
+                    },
+                )
+        if should_print:
+            metric_text = " ".join(f"{name}={value:.8e}" for name, value in metrics.items())
+            print(
+                f"TRAIN experiment={experiment} variant={variant} iteration={iteration}/{iterations} "
+                f"train_loss={train_loss:.8e} {metric_text}".rstrip(),
+                flush=True,
+            )
+    print(
+        f"DONE experiment={experiment} variant={variant} best_test_mse={best_metric:.8e} "
+        f"checkpoint={checkpoint_path} history={history_path}",
+        flush=True,
+    )
+
+
+def main() -> None:
+    args = parse_args()
+    config = load_config(args.config, args.smoke_test)
+    if args.device is not None:
+        config["runtime"]["device"] = args.device
+    if args.seed is not None:
+        config["runtime"]["seed"] = args.seed
+    experiments = list(config["experiments"]) if args.experiment == "all" else [args.experiment]
+    unknown = [name for name in experiments if name not in config["experiments"]]
+    if unknown:
+        raise KeyError(f"Unknown experiments: {unknown}; choose from {list(config['experiments'])}")
+    for experiment in experiments:
+        for variant in variants_for(config, experiment, args.variant):
+            run_training(
+                config,
+                experiment,
+                variant,
+                no_cache=args.no_cache,
+                resume=args.resume or bool(config["training_defaults"].get("resume", False)),
+            )
+
+
+if __name__ == "__main__":
+    main()

weight/best_model.pth

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