Upload folder using huggingface_hub

README.md

@@ -0,0 +1,156 @@
+---
+license: mit
+language:
+- en
+tags:
+- OneScience
+- fluid-mechanics
+- flow-field-prediction
+- graph-neural-network
+- partial-differential-equations
+frameworks: PyTorch
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">MP_PDE</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+MP-PDE (Message Passing Neural PDE Solver) is a neural PDE solver based on message passing, proposed by researchers from the University of Amsterdam and other institutions and published at ICLR 2022. It targets the spatiotemporal prediction of parameterized partial differential equations (PDEs). By learning local interactions between spatial nodes, it predicts numerical PDE solutions and adapts well to different equation parameters, spatial resolutions, and mesh structures.
+
+Using the OneScience skill workflow, this project independently reproduces the E3 parameterized PDE prediction experiment from the MP-PDE paper to evaluate spatiotemporal evolution predictions under different PDE parameter settings.
+
+Paper: [Message Passing Neural PDE Solvers](https://arxiv.org/abs/2202.03376)
+
+# Model Description
+
+MP-PDE uses an encoder-processor-decoder architecture consisting of an MLP encoder, a multilayer message-passing graph neural network (GNN processor), and a one-dimensional CNN decoder. It encodes PDE states from a historical time window together with spatial coordinates and equation parameters as node features. Multiple message-passing layers learn local spatial relationships, while temporal bundling predicts several future time steps at once for autoregressive spatiotemporal PDE forecasting.
+
+## Use Cases
+
+| Use case | Description |
+| --- | --- |
+| Cross-resolution PDE prediction | Graph structures describe relationships between spatial nodes with less dependence on fixed convolutional grids, supporting training and generalization across spatial discretization resolutions. |
+| Irregular-mesh computation | Message passing uses relative spatial positions between nodes and can handle nonuniform or irregular spatial discretizations, as demonstrated in the paper's wave-equation experiment. |
+| Long-horizon autoregressive prediction | Techniques such as temporal bundling and the pushforward trick reduce error accumulation during autoregressive prediction and improve the stability of long PDE rollouts. |
+
+# 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/MP_PDE --local_dir ./MP_PDE
+cd MP_PDE
+```
+
+### 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. Training data is synthesized on the fly by `models/dataset.py` according to `config.yaml`. The default input function $u$ is sampled from a zero-mean Gaussian random field (GRF) with covariance kernel:
+
+$$
+k_l(x_1,x_2)
+=
+\exp\left(
+-\frac{\lVert x_1-x_2\rVert^2}{2l^2}
+\right),
+$$
+
+where the correlation length is $l=0.2$. Input functions are first generated on a fine grid of 1,000 points, then sampled at 100 equally spaced sensor locations using cubic interpolation. These discrete samples are used as the Branch Net input in MP_PDE.
+
+Each supervised-learning sample is represented as:
+
+$$
+\left(u,\ y,\ G(u)(y)\right),
+$$
+
+where:
+
+- $u$: discrete values of the input function at 100 sensor locations;
+- $y$: the query location to predict—one-dimensional coordinates for ODE experiments and two-dimensional spatiotemporal coordinates $(x,t)$ for PDE experiments;
+- $G(u)(y)$: the reference solution computed by the corresponding differential-equation solver, representing the output of operator $G$ applied to input function $u$ at location $y$.
+
+### Training
+
+The default configuration reproduces the E3 parameterized PDE experiment from the MP-PDE paper. Training data is generated on the fly by `models/dataset.py` according to `config/config.yaml`, and the model uses a historical window of 25 time steps to predict the next 25 steps.
+
+Run E3 training with:
+
+```bash
+python scripts/train.py \
+    --config config/config.yaml \
+    --generate-data \
+    --device auto
+```
+
+### Trained Weights
+
+`weight/best_model.pth` contains the best weights from the full training run and can be used directly for inference or fine-tuning.
+
+### Inference
+
+Before running inference, make sure the E3 dataset `data/e3.h5` and trained weights `weight/best_model.pth` exist.
+
+```bash
+python scripts/inference.py \
+    --config config/config.yaml \
+    --device auto
+```
+
+### Evaluation and Visualization
+
+After training and inference, generate prediction visualizations 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: [Message Passing Neural PDE Solvers](https://arxiv.org/abs/2202.03376)
+- This project is an independent reproduction of the MP_PDE paper. The project code, model weights, training data, and third-party dependencies remain subject to their respective license terms.

config/config.yaml

@@ -0,0 +1,114 @@
+experiment:
+  name: E3
+  model: MP-PDE
+  paper: https://arxiv.org/abs/2202.03376
+  evidence_priority: paper_then_official_repo_for_missing_details
+  ambiguity_policy: paper_strict  # paper_strict: omega=[-0.4,-0.4]; official_consistency: [-0.4,0.4]
+
+paths:
+  data: data/e3.h5
+  checkpoint: weight/best_model.pth
+  results: results
+  train_history: results/train_history.json
+  predictions: results/predictions.npz
+  metrics: results/metrics.json
+
+data:
+  domain_length: 16.0
+  final_time: 4.0
+  num_time_points: 250  # E3 Table 1; nearby prose says 200
+  high_resolution_nx: 200
+  resolution: 100
+  supported_resolutions: [100, 50, 40]
+  sample_fraction: 0.3333333333333333  # user-authorized reduced run, 784/2352 trajectories exactly
+  full_experiment_samples:
+    train: 2096  # paper; official README says 2048
+    valid: 128   # paper missing; official README reference
+    test: 128    # paper missing; official README reference
+  train_samples: 698
+  valid_samples: 43
+  test_samples: 43
+  seed: 0
+  equation:
+    alpha_range: [0.0, 3.0]
+    beta_range: [0.0, 0.4]
+    gamma_range: [0.0, 1.0]
+  forcing:
+    terms: 5
+    amplitude_range: [-0.5, 0.5]
+    paper_omega_range: [-0.4, -0.4]
+    official_consistency_omega_range: [-0.4, 0.4]
+    modes: [1, 2, 3]
+    phase_range: [0.0, 6.283185307179586]
+  generation:
+    cfl: 0.2
+    min_dt: 1.0e-8
+    max_substeps: 10000
+    weno_epsilon: 1.0e-6
+    compression: gzip
+    compression_level: 4
+  parallel_generation:
+    workers: 8
+    max_in_flight: 16
+    flush_every: 8
+    resume_partial: true
+
+model:
+  time_window: 25  # paper missing; official training command reference
+  hidden_dim: 164
+  message_passing_layers: 6
+  neighbor_offsets: [-3, -2, -1, 1, 2, 3]
+  aggregation: sum
+  instance_norm_affine: false
+  decoder:
+    middle_channels: 8
+    kernel_sizes: [16, 26]
+    strides: [3, 1]
+  scaling:
+    coordinates: true
+    parameters: true
+    solution: none
+
+training:
+  epochs: 20
+  batch_size: 16
+  num_workers: 0
+  optimizer: AdamW
+  learning_rate: 1.0e-4
+  weight_decay: 1.0e-8
+  loss: rmse
+  loss_epsilon: 1.0e-12
+  max_no_grad_pushes: 1
+  scheduler:
+    enabled: true
+    milestones: [1, 5, 10, 15]
+    gamma: 0.4
+  precision: float32
+  amp: false
+  gradient_clip: null
+  seed: 0
+  deterministic: true
+  device: auto
+  print_interval: 10
+  validation_rollout_samples: null
+  test_rollout_samples: null
+
+inference:
+  batch_size: 16
+  split: test
+  max_samples: null
+  device: auto
+  forecast_start_index: 25
+
+visualization:
+  sample_index: 0
+  time_indices: [25, 75, 150, 249]
+  dpi: 160
+
+paper_reference:
+  metric: forecast_accumulated_mse
+  values:
+    100: 4.26
+    50: 3.74
+    40: 3.70
+  note: provenance_only_not_an_assertion

models/__init__.py

@@ -0,0 +1,6 @@
+"""Independent MP-PDE E3 reproduction components."""
+
+from .dataset import E3Dataset, generate_e3_hdf5
+from .pde import MPPDESolver, periodic_neighbor_indices
+
+__all__ = ["E3Dataset", "MPPDESolver", "generate_e3_hdf5", "periodic_neighbor_indices"]

models/dataset.py

@@ -0,0 +1,572 @@
+"""E3 trajectory generation and HDF5 loading for MP-PDE.
+
+This is an independent implementation from the equations and numerical-method
+description in arXiv:2202.03376.  No official repository source is used.
+"""
+
+from __future__ import annotations
+
+import argparse
+from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait
+import json
+import multiprocessing
+import os
+from pathlib import Path
+from typing import Any, Dict, Iterable, Mapping, Optional, Tuple
+
+import h5py
+import numpy as np
+import torch
+import yaml
+from torch.utils.data import Dataset
+
+
+SPLITS = ("train", "valid", "test")
+PARAMETER_ORDER = ("alpha", "beta", "gamma")
+PARALLEL_SCHEMA_VERSION = "mp_pde_e3_parallel_resume_v1"
+_WORKER_X: Optional[np.ndarray] = None
+_WORKER_TIMES: Optional[np.ndarray] = None
+_WORKER_CONFIG: Optional[Dict[str, Any]] = None
+_WORKER_STRIDE: Optional[int] = None
+
+
+def _shift(values: np.ndarray, offset: int) -> np.ndarray:
+    """Return values[i + offset] under periodic indexing."""
+    return np.roll(values, -offset, axis=-1)
+
+
+def _weno5_left(values: np.ndarray, epsilon: float) -> np.ndarray:
+    """Fifth-order WENO left state at every i+1/2 interface."""
+    um2, um1, u0 = _shift(values, -2), _shift(values, -1), values
+    up1, up2 = _shift(values, 1), _shift(values, 2)
+    p0 = (2.0 * um2 - 7.0 * um1 + 11.0 * u0) / 6.0
+    p1 = (-um1 + 5.0 * u0 + 2.0 * up1) / 6.0
+    p2 = (2.0 * u0 + 5.0 * up1 - up2) / 6.0
+    b0 = (13.0 / 12.0) * (um2 - 2.0 * um1 + u0) ** 2 + 0.25 * (um2 - 4.0 * um1 + 3.0 * u0) ** 2
+    b1 = (13.0 / 12.0) * (um1 - 2.0 * u0 + up1) ** 2 + 0.25 * (um1 - up1) ** 2
+    b2 = (13.0 / 12.0) * (u0 - 2.0 * up1 + up2) ** 2 + 0.25 * (3.0 * u0 - 4.0 * up1 + up2) ** 2
+    alpha = np.stack((0.1 / (epsilon + b0) ** 2, 0.6 / (epsilon + b1) ** 2, 0.3 / (epsilon + b2) ** 2))
+    weights = alpha / np.sum(alpha, axis=0, keepdims=True)
+    return weights[0] * p0 + weights[1] * p1 + weights[2] * p2
+
+
+def _weno5_right(values: np.ndarray, epsilon: float) -> np.ndarray:
+    """Fifth-order WENO right state at every i+1/2 interface."""
+    um1, u0 = _shift(values, -1), values
+    up1, up2, up3 = _shift(values, 1), _shift(values, 2), _shift(values, 3)
+    p0 = (2.0 * up3 - 7.0 * up2 + 11.0 * up1) / 6.0
+    p1 = (-up2 + 5.0 * up1 + 2.0 * u0) / 6.0
+    p2 = (2.0 * up1 + 5.0 * u0 - um1) / 6.0
+    b0 = (13.0 / 12.0) * (up1 - 2.0 * up2 + up3) ** 2 + 0.25 * (3.0 * up1 - 4.0 * up2 + up3) ** 2
+    b1 = (13.0 / 12.0) * (u0 - 2.0 * up1 + up2) ** 2 + 0.25 * (u0 - up2) ** 2
+    b2 = (13.0 / 12.0) * (um1 - 2.0 * u0 + up1) ** 2 + 0.25 * (um1 - 4.0 * u0 + 3.0 * up1) ** 2
+    alpha = np.stack((0.1 / (epsilon + b0) ** 2, 0.6 / (epsilon + b1) ** 2, 0.3 / (epsilon + b2) ** 2))
+    weights = alpha / np.sum(alpha, axis=0, keepdims=True)
+    return weights[0] * p0 + weights[1] * p1 + weights[2] * p2
+
+
+def godunov_quadratic_flux(left: np.ndarray, right: np.ndarray) -> np.ndarray:
+    """Godunov flux for the convex scalar flux f(u)=u**2."""
+    left_flux, right_flux = left * left, right * right
+    rarefaction = left <= right
+    rare_flux = np.where((left <= 0.0) & (right >= 0.0), 0.0, np.minimum(left_flux, right_flux))
+    shock_flux = np.maximum(left_flux, right_flux)
+    return np.where(rarefaction, rare_flux, shock_flux)
+
+
+def weno5_flux_derivative(values: np.ndarray, dx: float, epsilon: float = 1.0e-6) -> np.ndarray:
+    """Conservative derivative d_x(u**2) on a periodic uniform grid."""
+    interface_flux = godunov_quadratic_flux(_weno5_left(values, epsilon), _weno5_right(values, epsilon))
+    return (interface_flux - np.roll(interface_flux, 1, axis=-1)) / dx
+
+
+def fourth_order_second_derivative(values: np.ndarray, dx: float) -> np.ndarray:
+    return (-_shift(values, 2) + 16.0 * _shift(values, 1) - 30.0 * values + 16.0 * _shift(values, -1) - _shift(values, -2)) / (12.0 * dx**2)
+
+
+def fourth_order_third_derivative(values: np.ndarray, dx: float) -> np.ndarray:
+    return (_shift(values, -3) - 8.0 * _shift(values, -2) + 13.0 * _shift(values, -1) - 13.0 * _shift(values, 1) + 8.0 * _shift(values, 2) - _shift(values, 3)) / (8.0 * dx**3)
+
+
+def sample_e3_parameters(rng: np.random.Generator, cfg: Mapping[str, Any]) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
+    equation, forcing = cfg["equation"], cfg["forcing"]
+    params = np.array(
+        [rng.uniform(*equation["alpha_range"]), rng.uniform(*equation["beta_range"]), rng.uniform(*equation["gamma_range"])],
+        dtype=np.float64,
+    )
+    policy = cfg["ambiguity_policy"]
+    if policy == "paper_strict":
+        omega_range = forcing["paper_omega_range"]
+    elif policy == "official_consistency":
+        omega_range = forcing["official_consistency_omega_range"]
+    else:
+        raise ValueError(f"Unknown ambiguity_policy={policy!r}")
+    terms = int(forcing["terms"])
+    provenance = {
+        "amplitude": rng.uniform(*forcing["amplitude_range"], size=terms),
+        "omega": rng.uniform(*omega_range, size=terms),
+        "mode": rng.choice(np.asarray(forcing["modes"], dtype=np.int64), size=terms),
+        "phase": rng.uniform(*forcing["phase_range"], size=terms),
+    }
+    return params, provenance
+
+
+def evaluate_forcing(time: float, x: np.ndarray, forcing: Mapping[str, np.ndarray], domain_length: float) -> np.ndarray:
+    phase = (
+        forcing["omega"][:, None] * time
+        + 2.0 * np.pi * forcing["mode"][:, None] * x[None, :] / domain_length
+        + forcing["phase"][:, None]
+    )
+    return np.sum(forcing["amplitude"][:, None] * np.sin(phase), axis=0)
+
+
+def e3_rhs(
+    time: float,
+    state: np.ndarray,
+    x: np.ndarray,
+    params: np.ndarray,
+    forcing: Mapping[str, np.ndarray],
+    domain_length: float,
+    weno_epsilon: float,
+) -> np.ndarray:
+    alpha, beta, gamma = params
+    dx = domain_length / state.shape[-1]
+    return (
+        evaluate_forcing(time, x, forcing, domain_length)
+        - alpha * weno5_flux_derivative(state, dx, weno_epsilon)
+        + beta * fourth_order_second_derivative(state, dx)
+        - gamma * fourth_order_third_derivative(state, dx)
+    )
+
+
+def _stable_step(state: np.ndarray, params: np.ndarray, dx: float, cfl: float) -> float:
+    alpha, beta, gamma = np.abs(params)
+    limits = []
+    wave_speed = 2.0 * alpha * float(np.max(np.abs(state)))
+    if wave_speed > 1.0e-14:
+        limits.append(dx / wave_speed)
+    if beta > 1.0e-14:
+        limits.append(dx**2 / (2.0 * beta))
+    if gamma > 1.0e-14:
+        limits.append(dx**3 / (6.0 * gamma))
+    return cfl * min(limits) if limits else np.inf
+
+
+def generate_trajectory(
+    x: np.ndarray,
+    save_times: np.ndarray,
+    params: np.ndarray,
+    forcing: Mapping[str, np.ndarray],
+    cfg: Mapping[str, Any],
+) -> np.ndarray:
+    """Integrate one trajectory using RK4 and stability-limited substeps."""
+    domain_length = float(cfg["domain_length"])
+    generation = cfg["generation"]
+    dx = domain_length / x.size
+    state = evaluate_forcing(float(save_times[0]), x, forcing, domain_length).astype(np.float64)
+    trajectory = np.empty((save_times.size, x.size), dtype=np.float32)
+    trajectory[0] = state
+    current_time = float(save_times[0])
+    for output_index, target_time in enumerate(save_times[1:], start=1):
+        substeps = 0
+        while current_time < float(target_time) - 1.0e-14:
+            stable = _stable_step(state, params, dx, float(generation["cfl"]))
+            remaining = float(target_time) - current_time
+            step = min(stable, remaining)
+            if step < float(generation["min_dt"]) and remaining > float(generation["min_dt"]):
+                raise RuntimeError(f"Stable RK4 step {step:.3e} fell below min_dt at t={current_time:.6g}")
+            step = remaining if remaining <= float(generation["min_dt"]) else step
+            rhs_args = (x, params, forcing, domain_length, float(generation["weno_epsilon"]))
+            k1 = e3_rhs(current_time, state, *rhs_args)
+            k2 = e3_rhs(current_time + 0.5 * step, state + 0.5 * step * k1, *rhs_args)
+            k3 = e3_rhs(current_time + 0.5 * step, state + 0.5 * step * k2, *rhs_args)
+            k4 = e3_rhs(current_time + step, state + step * k3, *rhs_args)
+            state = state + (step / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
+            current_time += step
+            substeps += 1
+            if substeps > int(generation["max_substeps"]):
+                raise RuntimeError(f"max_substeps exceeded while advancing to t={target_time:.6g}")
+            if not np.all(np.isfinite(state)):
+                raise FloatingPointError(f"Non-finite E3 state at t={current_time:.6g}")
+        current_time = float(target_time)
+        trajectory[output_index] = state
+    return trajectory
+
+
+def _generation_config(config: Mapping[str, Any]) -> Dict[str, Any]:
+    data = config["data"]
+    return {
+        "ambiguity_policy": config["experiment"]["ambiguity_policy"],
+        "domain_length": data["domain_length"],
+        "equation": data["equation"],
+        "forcing": data["forcing"],
+        "generation": data["generation"],
+    }
+
+
+def _initialize_generation_worker(
+    x_high: np.ndarray, times: np.ndarray, generation_cfg: Mapping[str, Any], stride: int
+) -> None:
+    """Initialize immutable state used by one trajectory worker process."""
+    global _WORKER_X, _WORKER_TIMES, _WORKER_CONFIG, _WORKER_STRIDE
+    _WORKER_X = x_high
+    _WORKER_TIMES = times
+    _WORKER_CONFIG = dict(generation_cfg)
+    _WORKER_STRIDE = int(stride)
+
+
+def _generate_sample_worker(payload: Tuple[int, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]) -> Tuple[int, np.ndarray]:
+    """Generate one independent trajectory; HDF5 remains owned by the parent."""
+    if _WORKER_X is None or _WORKER_TIMES is None or _WORKER_CONFIG is None or _WORKER_STRIDE is None:
+        raise RuntimeError("Parallel E3 worker was not initialized")
+    sample_index, params, amplitude, omega, mode, phase = payload
+    forcing = {"amplitude": amplitude, "omega": omega, "mode": mode, "phase": phase}
+    trajectory = generate_trajectory(_WORKER_X, _WORKER_TIMES, params, forcing, _WORKER_CONFIG)
+    return sample_index, trajectory[:, ::_WORKER_STRIDE]
+
+
+def _sample_payloads(group: h5py.Group, indices: Iterable[int]) -> Iterable[Tuple[int, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]]:
+    for sample_index in indices:
+        yield (
+            sample_index,
+            np.asarray(group["params"][sample_index], dtype=np.float64),
+            np.asarray(group["forcing_amplitude"][sample_index], dtype=np.float64),
+            np.asarray(group["forcing_omega"][sample_index], dtype=np.float64),
+            np.asarray(group["forcing_mode"][sample_index], dtype=np.int64),
+            np.asarray(group["forcing_phase"][sample_index], dtype=np.float64),
+        )
+
+
+def _parallel_results(
+    payloads: Iterable[Tuple[int, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]],
+    workers: int,
+    max_in_flight: int,
+    x_high: np.ndarray,
+    times: np.ndarray,
+    generation_cfg: Mapping[str, Any],
+    stride: int,
+) -> Iterable[Tuple[int, np.ndarray]]:
+    """Yield completed trajectories while keeping the process queue bounded."""
+    if workers == 1:
+        _initialize_generation_worker(x_high, times, generation_cfg, stride)
+        for payload in payloads:
+            yield _generate_sample_worker(payload)
+        return
+
+    thread_variables = ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS")
+    previous_environment = {name: os.environ.get(name) for name in thread_variables}
+    for name in thread_variables:
+        os.environ[name] = "1"
+    executor = ProcessPoolExecutor(
+        max_workers=workers,
+        mp_context=multiprocessing.get_context("spawn"),
+        initializer=_initialize_generation_worker,
+        initargs=(x_high, times, dict(generation_cfg), stride),
+    )
+    iterator = iter(payloads)
+    futures = set()
+    try:
+        for _ in range(max_in_flight):
+            try:
+                futures.add(executor.submit(_generate_sample_worker, next(iterator)))
+            except StopIteration:
+                break
+        while futures:
+            completed, futures = wait(futures, return_when=FIRST_COMPLETED)
+            for future in completed:
+                yield future.result()
+                try:
+                    futures.add(executor.submit(_generate_sample_worker, next(iterator)))
+                except StopIteration:
+                    pass
+    finally:
+        executor.shutdown(wait=True, cancel_futures=True)
+        for name, value in previous_environment.items():
+            if value is None:
+                os.environ.pop(name, None)
+            else:
+                os.environ[name] = value
+
+
+def _generation_signature(
+    config: Mapping[str, Any], counts: Mapping[str, int], nt: int, high_nx: int, target_nx: int
+) -> str:
+    payload = {
+        "schema_version": PARALLEL_SCHEMA_VERSION,
+        "paper": str(config["experiment"]["paper"]),
+        "ambiguity_policy": str(config["experiment"]["ambiguity_policy"]),
+        "generation_config": _generation_config(config),
+        "counts": {name: int(counts[name]) for name in SPLITS},
+        "nt": int(nt),
+        "high_resolution_nx": int(high_nx),
+        "resolution": int(target_nx),
+        "seed": int(config["data"]["seed"]),
+    }
+    return json.dumps(payload, sort_keys=True, separators=(",", ":"))
+
+
+def _initialize_partial_file(
+    partial_path: Path,
+    config: Mapping[str, Any],
+    counts: Mapping[str, int],
+    nt: int,
+    high_nx: int,
+    target_nx: int,
+    x: np.ndarray,
+    times: np.ndarray,
+    generation_cfg: Mapping[str, Any],
+    compression: Optional[str],
+    compression_opts: Optional[int],
+    signature: str,
+) -> None:
+    data = config["data"]
+    base_seed = int(data["seed"])
+    with h5py.File(partial_path, "w") as handle:
+        handle.attrs["schema_version"] = PARALLEL_SCHEMA_VERSION
+        handle.attrs["generation_signature"] = signature
+        handle.attrs["experiment"] = "MP-PDE E3"
+        handle.attrs["paper"] = str(config["experiment"]["paper"])
+        handle.attrs["ambiguity_policy"] = str(config["experiment"]["ambiguity_policy"])
+        handle.attrs["parameter_order"] = json.dumps(PARAMETER_ORDER)
+        handle.attrs["generation_config"] = json.dumps(generation_cfg, sort_keys=True)
+        handle.attrs["seed"] = base_seed
+        handle.attrs["high_resolution_nx"] = high_nx
+        handle.attrs["resolution"] = target_nx
+        for split_index, split in enumerate(SPLITS):
+            count = int(counts[split])
+            rng = np.random.default_rng(np.random.SeedSequence([base_seed, split_index]))
+            group = handle.create_group(split)
+            group.attrs["seed_derivation"] = json.dumps([base_seed, split_index])
+            group.create_dataset("x", data=x)
+            group.create_dataset("t", data=times.astype(np.float32))
+            group.create_dataset(
+                "u", shape=(count, nt, target_nx), dtype="f4", chunks=(1, min(nt, 32), target_nx),
+                compression=compression, compression_opts=compression_opts,
+            )
+            params_ds = group.create_dataset("params", shape=(count, 3), dtype="f4")
+            terms = int(data["forcing"]["terms"])
+            amplitude_ds = group.create_dataset("forcing_amplitude", shape=(count, terms), dtype="f4")
+            omega_ds = group.create_dataset("forcing_omega", shape=(count, terms), dtype="f4")
+            mode_ds = group.create_dataset("forcing_mode", shape=(count, terms), dtype="i8")
+            phase_ds = group.create_dataset("forcing_phase", shape=(count, terms), dtype="f4")
+            group.create_dataset("completed", shape=(count,), dtype="bool", data=np.zeros(count, dtype=bool))
+            for sample_index in range(count):
+                params, forcing = sample_e3_parameters(rng, generation_cfg)
+                params_ds[sample_index] = params
+                amplitude_ds[sample_index] = forcing["amplitude"]
+                omega_ds[sample_index] = forcing["omega"]
+                mode_ds[sample_index] = forcing["mode"]
+                phase_ds[sample_index] = forcing["phase"]
+        handle.flush()
+
+
+def generate_e3_hdf5(
+    config: Mapping[str, Any],
+    output_path: Path | str,
+    *,
+    sample_counts: Optional[Mapping[str, int]] = None,
+    nt: Optional[int] = None,
+    high_resolution_nx: Optional[int] = None,
+    resolution: Optional[int] = None,
+    workers: Optional[int] = None,
+    max_in_flight: Optional[int] = None,
+    flush_every: Optional[int] = None,
+    resume_partial: Optional[bool] = None,
+    overwrite: bool = False,
+) -> Path:
+    """Generate E3 splits with process workers and a single resumable HDF5 writer."""
+    output_path = Path(output_path)
+    if output_path.exists() and not overwrite:
+        raise FileExistsError(f"Refusing to overwrite existing dataset: {output_path}")
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    partial_path = output_path.with_suffix(output_path.suffix + ".partial")
+
+    data = config["data"]
+    parallel = data.get("parallel_generation", {})
+    workers = int(workers if workers is not None else parallel.get("workers", 1))
+    max_in_flight = int(max_in_flight if max_in_flight is not None else parallel.get("max_in_flight", 2 * workers))
+    flush_every = int(flush_every if flush_every is not None else parallel.get("flush_every", workers))
+    resume_partial = bool(resume_partial if resume_partial is not None else parallel.get("resume_partial", True))
+    if workers < 1 or max_in_flight < workers or flush_every < 1:
+        raise ValueError("workers>=1, max_in_flight>=workers, and flush_every>=1 are required")
+    nt = int(nt or data["num_time_points"])
+    high_nx = int(high_resolution_nx or data["high_resolution_nx"])
+    target_nx = int(resolution or data["resolution"])
+    if high_nx % target_nx != 0:
+        raise ValueError(f"high_resolution_nx={high_nx} must be divisible by resolution={target_nx}")
+    if high_nx < 7 or nt < 2:
+        raise ValueError("Generation requires high_resolution_nx>=7 and nt>=2")
+    counts = dict(sample_counts or {name: int(data[f"{name}_samples"]) for name in SPLITS})
+    if set(counts) != set(SPLITS) or any(int(counts[name]) <= 0 for name in SPLITS):
+        raise ValueError(f"sample_counts must provide positive counts for {SPLITS}")
+
+    domain_length, final_time = float(data["domain_length"]), float(data["final_time"])
+    x_high = np.linspace(0.0, domain_length, high_nx, endpoint=False, dtype=np.float64)
+    stride = high_nx // target_nx
+    x = x_high[::stride].astype(np.float32)
+    times = np.linspace(0.0, final_time, nt, dtype=np.float64)
+    generation_cfg = _generation_config(config)
+    compression = data["generation"].get("compression")
+    compression_opts = int(data["generation"].get("compression_level", 4)) if compression == "gzip" else None
+    signature = _generation_signature(config, counts, nt, high_nx, target_nx)
+    if partial_path.exists() and overwrite:
+        partial_path.unlink()
+    if partial_path.exists() and not resume_partial:
+        raise FileExistsError(f"Partial dataset exists; enable resume_partial or use --overwrite: {partial_path}")
+    if not partial_path.exists():
+        _initialize_partial_file(
+            partial_path, config, counts, nt, high_nx, target_nx, x, times, generation_cfg,
+            compression, compression_opts, signature,
+        )
+
+    try:
+        with h5py.File(partial_path, "r+") as handle:
+            stored_signature = str(handle.attrs.get("generation_signature", ""))
+            if stored_signature != signature:
+                raise ValueError("Partial dataset configuration does not match this run; archive it or use --overwrite")
+            print(
+                f"[generate] workers={workers} max_in_flight={max_in_flight} flush_every={flush_every} "
+                f"resume_partial={resume_partial}", flush=True,
+            )
+            for split in SPLITS:
+                count = int(counts[split])
+                group = handle[split]
+                completed_ds = group["completed"]
+                pending_indices = np.flatnonzero(~np.asarray(completed_ds[:], dtype=bool)).tolist()
+                completed_count = count - len(pending_indices)
+                if pending_indices:
+                    print(f"[generate] split={split} resume_completed={completed_count}/{count}", flush=True)
+                payloads = _sample_payloads(group, pending_indices)
+                for sample_index, trajectory in _parallel_results(
+                    payloads, workers, max_in_flight, x_high, times, generation_cfg, stride
+                ):
+                    group["u"][sample_index] = trajectory
+                    completed_ds[sample_index] = True
+                    completed_count += 1
+                    if completed_count % flush_every == 0 or completed_count == count:
+                        handle.flush()
+                    print(
+                        f"[generate] split={split} completed={completed_count}/{count} "
+                        f"sample_index={sample_index} workers={workers}", flush=True,
+                    )
+                if not np.all(np.asarray(completed_ds[:], dtype=bool)):
+                    raise RuntimeError(f"Split {split} is incomplete after generation")
+                handle.flush()
+        os.replace(partial_path, output_path)
+    except Exception:
+        print(f"Generation failed; partial file retained at {partial_path}", flush=True)
+        raise
+    return output_path
+
+
+class E3Dataset(Dataset):
+    """Lazy, process-safe reader for one E3 HDF5 split."""
+
+    def __init__(self, path: Path | str, split: str, expected_nt: Optional[int] = None, expected_nx: Optional[int] = None):
+        self.path = Path(path)
+        self.split = split
+        self._handle: Optional[h5py.File] = None
+        if split not in SPLITS:
+            raise ValueError(f"split must be one of {SPLITS}, got {split!r}")
+        if not self.path.is_file():
+            raise FileNotFoundError(f"E3 dataset not found: {self.path}")
+        with h5py.File(self.path, "r") as handle:
+            if split not in handle:
+                raise KeyError(f"Missing HDF5 group {split!r}")
+            group = handle[split]
+            required = {"u", "x", "t", "params", "forcing_amplitude", "forcing_omega", "forcing_mode", "forcing_phase"}
+            missing = required.difference(group.keys())
+            if missing:
+                raise KeyError(f"Missing HDF5 fields in {split}: {sorted(missing)}")
+            shape = group["u"].shape
+            if len(shape) != 3 or group["params"].shape != (shape[0], 3) or group["x"].shape != (shape[2],) or group["t"].shape != (shape[1],):
+                raise ValueError(f"Inconsistent E3 schema in split={split}: u={shape}")
+            if expected_nt is not None and shape[1] != expected_nt:
+                raise ValueError(f"Expected nt={expected_nt}, found {shape[1]}")
+            if expected_nx is not None and shape[2] != expected_nx:
+                raise ValueError(f"Expected nx={expected_nx}, found {shape[2]}")
+            self.length, self.nt, self.nx = shape
+            x, t = group["x"][:], group["t"][:]
+            if not np.all(np.isfinite(x)) or not np.all(np.isfinite(t)) or np.any(np.diff(t) <= 0.0):
+                raise ValueError("Grid/time metadata are non-finite or non-monotone")
+            if x.size > 1 and not np.allclose(np.diff(x), np.diff(x)[0], rtol=1e-5, atol=1e-7):
+                raise ValueError("E3 x grid must be uniform")
+            if t.size > 1 and not np.allclose(np.diff(t), np.diff(t)[0], rtol=1e-5, atol=1e-7):
+                raise ValueError("E3 saved times must be uniform")
+
+    def _group(self) -> h5py.Group:
+        if self._handle is None:
+            self._handle = h5py.File(self.path, "r")
+        return self._handle[self.split]
+
+    def __len__(self) -> int:
+        return self.length
+
+    def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:
+        group = self._group()
+        trajectory = np.asarray(group["u"][index], dtype=np.float32)
+        params = np.asarray(group["params"][index], dtype=np.float32)
+        if not np.all(np.isfinite(trajectory)) or not np.all(np.isfinite(params)):
+            raise FloatingPointError(f"Non-finite data at split={self.split}, sample={index}")
+        return {
+            "u": torch.from_numpy(trajectory),
+            "x": torch.from_numpy(np.asarray(group["x"][:], dtype=np.float32)),
+            "t": torch.from_numpy(np.asarray(group["t"][:], dtype=np.float32)),
+            "params": torch.from_numpy(params),
+            "index": torch.tensor(index, dtype=torch.long),
+        }
+
+    def close(self) -> None:
+        if self._handle is not None:
+            self._handle.close()
+            self._handle = None
+
+    def __del__(self) -> None:
+        self.close()
+
+
+def _load_yaml(path: Path) -> Dict[str, Any]:
+    with path.open("r", encoding="utf-8") as stream:
+        config = yaml.safe_load(stream)
+    if not isinstance(config, dict):
+        raise ValueError(f"Config must contain a mapping: {path}")
+    return config
+
+
+def main() -> None:
+    project_root = Path(__file__).resolve().parents[1]
+    parser = argparse.ArgumentParser(description="Generate the MP-PDE E3 HDF5 dataset")
+    parser.add_argument("--config", type=Path, default=project_root / "config/config.yaml")
+    parser.add_argument("--output", type=Path, default=None)
+    parser.add_argument("--overwrite", action="store_true")
+    parser.add_argument("--train-samples", type=int)
+    parser.add_argument("--valid-samples", type=int)
+    parser.add_argument("--test-samples", type=int)
+    parser.add_argument("--nt", type=int)
+    parser.add_argument("--high-resolution-nx", type=int)
+    parser.add_argument("--resolution", type=int)
+    parser.add_argument("--workers", type=int)
+    parser.add_argument("--max-in-flight", type=int)
+    parser.add_argument("--flush-every", type=int)
+    parser.add_argument("--no-resume-partial", action="store_true")
+    args = parser.parse_args()
+    config = _load_yaml(args.config.resolve())
+    configured = Path(config["paths"]["data"])
+    output = args.output or (configured if configured.is_absolute() else project_root / configured)
+    default_counts = {name: int(config["data"][f"{name}_samples"]) for name in SPLITS}
+    counts = {
+        "train": args.train_samples or default_counts["train"],
+        "valid": args.valid_samples or default_counts["valid"],
+        "test": args.test_samples or default_counts["test"],
+    }
+    generated = generate_e3_hdf5(
+        config, output, sample_counts=counts, nt=args.nt, high_resolution_nx=args.high_resolution_nx,
+        resolution=args.resolution, workers=args.workers, max_in_flight=args.max_in_flight,
+        flush_every=args.flush_every, resume_partial=False if args.no_resume_partial else None,
+        overwrite=args.overwrite,
+    )
+    print(f"Generated E3 dataset: {generated}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

models/pde.py

@@ -0,0 +1,211 @@
+"""Paper-priority MP-PDE model for experiment E3."""
+
+from __future__ import annotations
+
+from typing import Iterable, Sequence, Tuple
+
+import torch
+from torch import Tensor, nn
+
+
+class Swish(nn.Module):
+    def forward(self, values: Tensor) -> Tensor:
+        return values * torch.sigmoid(values)
+
+
+class TwoLayerMLP(nn.Module):
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__()
+        self.network = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            Swish(),
+            nn.Linear(hidden_dim, output_dim),
+            Swish(),
+        )
+
+    def forward(self, values: Tensor) -> Tensor:
+        return self.network(values)
+
+
+def periodic_neighbor_indices(num_nodes: int, offsets: Iterable[int], device: torch.device | None = None) -> Tensor:
+    """Return source indices [num_nodes, num_neighbors] for each target node."""
+    offsets_tensor = torch.as_tensor(tuple(offsets), dtype=torch.long, device=device)
+    if num_nodes < 7:
+        raise ValueError(f"The six-neighbor periodic graph requires num_nodes>=7, found {num_nodes}")
+    if offsets_tensor.numel() != 6 or offsets_tensor.unique().numel() != 6 or torch.any(offsets_tensor == 0):
+        raise ValueError("MP-PDE E3 requires exactly six distinct non-zero neighbor offsets")
+    target = torch.arange(num_nodes, dtype=torch.long, device=device)[:, None]
+    return torch.remainder(target + offsets_tensor[None, :], num_nodes)
+
+
+class MessagePassingLayer(nn.Module):
+    """Equation (8)--(9) processor layer with sum aggregation."""
+
+    def __init__(self, hidden_dim: int, history_dim: int, parameter_dim: int, affine_norm: bool):
+        super().__init__()
+        edge_dim = 2 * hidden_dim + history_dim + 1 + parameter_dim
+        node_dim = 2 * hidden_dim + parameter_dim
+        self.edge_mlp = TwoLayerMLP(edge_dim, hidden_dim, hidden_dim)
+        self.node_mlp = TwoLayerMLP(node_dim, hidden_dim, hidden_dim)
+        self.norm = nn.InstanceNorm1d(hidden_dim, affine=affine_norm, track_running_stats=False)
+
+    def forward(
+        self,
+        hidden: Tensor,
+        history: Tensor,
+        x: Tensor,
+        parameters: Tensor,
+        neighbors: Tensor,
+        domain_length: float,
+    ) -> Tensor:
+        batch, num_nodes, hidden_dim = hidden.shape
+        num_neighbors = neighbors.shape[1]
+        source_hidden = hidden[:, neighbors, :]
+        target_hidden = hidden[:, :, None, :].expand(-1, -1, num_neighbors, -1)
+        source_history = history[:, neighbors, :]
+        history_difference = history[:, :, None, :] - source_history
+        source_x = x[:, neighbors]
+        displacement = x[:, :, None] - source_x
+        displacement = torch.remainder(displacement + 0.5 * domain_length, domain_length) - 0.5 * domain_length
+        theta_edges = parameters[:, None, None, :].expand(-1, num_nodes, num_neighbors, -1)
+        edge_input = torch.cat(
+            (target_hidden, source_hidden, history_difference, displacement[..., None], theta_edges), dim=-1
+        )
+        messages = self.edge_mlp(edge_input)
+        aggregated = messages.sum(dim=2)
+        theta_nodes = parameters[:, None, :].expand(-1, num_nodes, -1)
+        node_update = self.node_mlp(torch.cat((hidden, aggregated, theta_nodes), dim=-1))
+        return self.norm((hidden + node_update).transpose(1, 2)).transpose(1, 2)
+
+
+class TemporalDecoder(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int,
+        time_window: int,
+        middle_channels: int = 8,
+        kernels: Sequence[int] = (16, 26),
+        strides: Sequence[int] = (3, 1),
+    ):
+        super().__init__()
+        if len(kernels) != 2 or len(strides) != 2:
+            raise ValueError("The paper-priority decoder requires exactly two convolutions")
+        length_after_first = (hidden_dim - int(kernels[0])) // int(strides[0]) + 1
+        output_length = (length_after_first - int(kernels[1])) // int(strides[1]) + 1
+        if output_length != time_window:
+            raise ValueError(
+                f"Decoder does not close hidden={hidden_dim} to K={time_window}: output length={output_length}"
+            )
+        self.network = nn.Sequential(
+            nn.Conv1d(1, middle_channels, kernel_size=int(kernels[0]), stride=int(strides[0])),
+            Swish(),
+            nn.Conv1d(middle_channels, 1, kernel_size=int(kernels[1]), stride=int(strides[1])),
+        )
+
+    def forward(self, hidden: Tensor) -> Tensor:
+        batch, num_nodes, hidden_dim = hidden.shape
+        decoded = self.network(hidden.reshape(batch * num_nodes, 1, hidden_dim))
+        return decoded.reshape(batch, num_nodes, decoded.shape[-1])
+
+
+class MPPDESolver(nn.Module):
+    """Message-passing neural solver mapping K E3 states to the next K states."""
+
+    def __init__(
+        self,
+        time_window: int = 25,
+        hidden_dim: int = 164,
+        message_passing_layers: int = 6,
+        neighbor_offsets: Sequence[int] = (-3, -2, -1, 1, 2, 3),
+        domain_length: float = 16.0,
+        final_time: float = 4.0,
+        parameter_maxima: Sequence[float] = (3.0, 0.4, 1.0),
+        scale_coordinates: bool = True,
+        scale_parameters: bool = True,
+        instance_norm_affine: bool = False,
+        decoder_middle_channels: int = 8,
+        decoder_kernels: Sequence[int] = (16, 26),
+        decoder_strides: Sequence[int] = (3, 1),
+    ):
+        super().__init__()
+        if time_window <= 0 or hidden_dim <= 0 or message_passing_layers <= 0:
+            raise ValueError("time_window, hidden_dim, and message_passing_layers must be positive")
+        self.time_window = int(time_window)
+        self.hidden_dim = int(hidden_dim)
+        self.neighbor_offsets: Tuple[int, ...] = tuple(int(value) for value in neighbor_offsets)
+        self.domain_length = float(domain_length)
+        self.final_time = float(final_time)
+        self.scale_coordinates = bool(scale_coordinates)
+        self.scale_parameters = bool(scale_parameters)
+        maxima = torch.tensor(tuple(float(value) for value in parameter_maxima), dtype=torch.float32)
+        if maxima.shape != (3,) or torch.any(maxima <= 0.0):
+            raise ValueError("parameter_maxima must contain three positive values")
+        self.register_buffer("parameter_maxima", maxima, persistent=False)
+        self.encoder = TwoLayerMLP(self.time_window + 1 + 1 + 3, self.hidden_dim, self.hidden_dim)
+        self.processor = nn.ModuleList(
+            MessagePassingLayer(self.hidden_dim, self.time_window, 3, instance_norm_affine)
+            for _ in range(int(message_passing_layers))
+        )
+        self.decoder = TemporalDecoder(
+            self.hidden_dim, self.time_window, decoder_middle_channels, decoder_kernels, decoder_strides
+        )
+
+    def _canonicalize_inputs(
+        self, history: Tensor, x: Tensor, current_time: Tensor | float, parameters: Tensor
+    ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+        if history.ndim != 3:
+            raise ValueError(f"history must have shape [B,N,K], found {tuple(history.shape)}")
+        batch, num_nodes, window = history.shape
+        if window != self.time_window:
+            raise ValueError(f"Expected history K={self.time_window}, found {window}")
+        if num_nodes < 7:
+            raise ValueError(f"MP-PDE requires N>=7, found {num_nodes}")
+        if parameters.shape != (batch, 3):
+            raise ValueError(f"parameters must have shape [{batch},3], found {tuple(parameters.shape)}")
+        if x.ndim == 1:
+            if x.shape[0] != num_nodes:
+                raise ValueError(f"x length {x.shape[0]} does not match N={num_nodes}")
+            x = x[None, :].expand(batch, -1)
+        elif x.shape != (batch, num_nodes):
+            raise ValueError(f"x must have shape [N] or [B,N], found {tuple(x.shape)}")
+        time = torch.as_tensor(current_time, dtype=history.dtype, device=history.device)
+        if time.ndim == 0:
+            time = time.expand(batch)
+        elif time.shape == (batch, 1):
+            time = time[:, 0]
+        elif time.shape != (batch,):
+            raise ValueError(f"current_time must be scalar or shape [B], found {tuple(time.shape)}")
+        return history, x.to(history), time, parameters.to(history)
+
+    def forward(
+        self,
+        history: Tensor,
+        x: Tensor,
+        current_time: Tensor | float,
+        parameters: Tensor,
+        dt: Tensor | float,
+        *,
+        return_derivative: bool = False,
+    ) -> Tensor | Tuple[Tensor, Tensor]:
+        history, x, current_time, parameters = self._canonicalize_inputs(history, x, current_time, parameters)
+        batch, num_nodes, _ = history.shape
+        scaled_x = x / self.domain_length if self.scale_coordinates else x
+        scaled_time = current_time / self.final_time if self.scale_coordinates else current_time
+        scaled_parameters = parameters / self.parameter_maxima.to(parameters) if self.scale_parameters else parameters
+        time_feature = scaled_time[:, None, None].expand(-1, num_nodes, 1)
+        parameter_features = scaled_parameters[:, None, :].expand(-1, num_nodes, -1)
+        encoded = torch.cat((history, scaled_x[..., None], time_feature, parameter_features), dim=-1)
+        hidden = self.encoder(encoded)
+        neighbors = periodic_neighbor_indices(num_nodes, self.neighbor_offsets, history.device)
+        for layer in self.processor:
+            hidden = layer(hidden, history, x, scaled_parameters, neighbors, self.domain_length)
+        derivative = self.decoder(hidden)
+        step = torch.as_tensor(dt, dtype=history.dtype, device=history.device)
+        if step.ndim == 0:
+            step = step.expand(batch)
+        elif step.shape != (batch,):
+            raise ValueError(f"dt must be scalar or shape [B], found {tuple(step.shape)}")
+        offsets = torch.arange(1, self.time_window + 1, dtype=history.dtype, device=history.device)
+        delta_times = step[:, None, None] * offsets[None, None, :]
+        prediction = history[:, :, -1:] + delta_times * derivative
+        return (prediction, derivative) if return_derivative else prediction

scripts/__init__.py

@@ -0,0 +1 @@
+"""Command-line entrypoints for the MP-PDE E3 reproduction."""

scripts/inference.py

@@ -0,0 +1,216 @@
+"""Autoregressive MP-PDE E3 rollout and evaluation."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, Iterable, Mapping
+
+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.dataset import E3Dataset  # noqa: E402
+from scripts.train import build_model, choose_device, load_config, project_path, rollout_batch, set_seed  # noqa: E402
+
+
+CRITICAL_CONFIG_KEYS = (
+    ("model", "time_window"),
+    ("model", "hidden_dim"),
+    ("model", "message_passing_layers"),
+    ("model", "neighbor_offsets"),
+    ("model", "aggregation"),
+    ("model", "decoder", "kernel_sizes"),
+    ("model", "decoder", "strides"),
+    ("model", "scaling", "coordinates"),
+    ("model", "scaling", "parameters"),
+    ("data", "equation", "alpha_range"),
+    ("data", "equation", "beta_range"),
+    ("data", "equation", "gamma_range"),
+)
+
+
+def _nested(mapping: Mapping[str, Any], key: Iterable[str]) -> Any:
+    value: Any = mapping
+    for component in key:
+        value = value[component]
+    return value
+
+
+def verify_checkpoint_config(runtime: Mapping[str, Any], stored: Mapping[str, Any]) -> None:
+    mismatches = []
+    for key in CRITICAL_CONFIG_KEYS:
+        runtime_value, stored_value = _nested(runtime, key), _nested(stored, key)
+        if runtime_value != stored_value:
+            mismatches.append(f"{'.'.join(key)}: runtime={runtime_value!r}, checkpoint={stored_value!r}")
+    if mismatches:
+        raise ValueError("Checkpoint/config architecture mismatch:\n  " + "\n  ".join(mismatches))
+
+
+def compute_metrics(prediction: np.ndarray, target: np.ndarray, forecast_start: int) -> Dict[str, Any]:
+    if prediction.shape != target.shape or prediction.ndim != 3:
+        raise ValueError(f"prediction/target must share [S,T,N] shape, found {prediction.shape}/{target.shape}")
+    error = prediction[:, forecast_start:] - target[:, forecast_start:]
+    if error.size == 0 or not np.all(np.isfinite(error)):
+        raise FloatingPointError("Forecast error is empty or non-finite")
+    squared = error.astype(np.float64) ** 2
+    per_time_mse = np.mean(squared, axis=(0, 2))
+    numerator = np.linalg.norm(error.reshape(error.shape[0], -1), axis=1)
+    denominator = np.linalg.norm(target[:, forecast_start:].reshape(error.shape[0], -1), axis=1) + 1.0e-12
+    return {
+        "accumulated_mse": float(np.sum(per_time_mse)),
+        "rmse": float(np.sqrt(np.mean(squared))),
+        "mae": float(np.mean(np.abs(error))),
+        "relative_l2": float(np.mean(numerator / denominator)),
+        "per_time_mse": per_time_mse,
+    }
+
+
+def atomic_npz(destination: Path, **arrays: Any) -> None:
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.NamedTemporaryFile(dir=destination.parent, prefix=destination.stem + ".", suffix=".npz", delete=False) as stream:
+        temporary = Path(stream.name)
+    try:
+        np.savez_compressed(temporary, **arrays)
+        os.replace(temporary, destination)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def atomic_json(destination: Path, payload: Mapping[str, Any]) -> None:
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.NamedTemporaryFile("w", dir=destination.parent, prefix=destination.name + ".", suffix=".tmp", delete=False, encoding="utf-8") as stream:
+        json.dump(payload, stream, indent=2, sort_keys=True)
+        stream.write("\n")
+        temporary = Path(stream.name)
+    os.replace(temporary, destination)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Run MP-PDE E3 autoregressive inference")
+    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "config/config.yaml")
+    parser.add_argument("--checkpoint", type=Path)
+    parser.add_argument("--data", type=Path)
+    parser.add_argument("--split", type=str)
+    parser.add_argument("--output-dir", type=Path)
+    parser.add_argument("--device", type=str)
+    parser.add_argument("--batch-size", type=int)
+    parser.add_argument("--max-samples", type=int)
+    parser.add_argument("--seed", type=int)
+    args = parser.parse_args()
+
+    config = load_config(args.config.resolve())
+    inference_cfg = config["inference"]
+    checkpoint_path = args.checkpoint or project_path(config["paths"]["checkpoint"])
+    data_path = args.data or project_path(config["paths"]["data"])
+    output_dir = args.output_dir or project_path(config["paths"]["results"])
+    if not checkpoint_path.is_file():
+        raise FileNotFoundError(f"Real trained checkpoint not found: {checkpoint_path}. Run scripts/train.py first.")
+    if not data_path.is_file():
+        raise FileNotFoundError(f"E3 HDF5 data not found: {data_path}. Generate data before inference.")
+    seed = int(args.seed if args.seed is not None else config["training"]["seed"])
+    set_seed(seed, bool(config["training"]["deterministic"]))
+    device = choose_device(args.device or str(inference_cfg["device"]))
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    if "model_state" not in checkpoint or "resolved_config" not in checkpoint:
+        raise KeyError("Checkpoint must contain model_state and resolved_config")
+    verify_checkpoint_config(config, checkpoint["resolved_config"])
+    model = build_model(config).to(device)
+    model.load_state_dict(checkpoint["model_state"], strict=True)
+    model.eval()
+
+    split = args.split or str(inference_cfg["split"])
+    dataset = E3Dataset(
+        data_path, split, expected_nt=int(config["data"]["num_time_points"]), expected_nx=int(config["data"]["resolution"])
+    )
+    batch_size = int(args.batch_size or inference_cfg["batch_size"])
+    configured_max = inference_cfg.get("max_samples")
+    max_samples = args.max_samples if args.max_samples is not None else configured_max
+    if max_samples is not None and int(max_samples) <= 0:
+        raise ValueError("max_samples must be positive")
+    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)
+    window = int(config["model"]["time_window"])
+    predictions, targets, parameters, sample_indices = [], [], [], []
+    first_x, first_t = None, None
+    processed = 0
+    with torch.inference_mode():
+        for batch_index, batch in enumerate(loader, start=1):
+            if max_samples is not None:
+                remaining = int(max_samples) - processed
+                if remaining <= 0:
+                    break
+                if batch["u"].shape[0] > remaining:
+                    batch = {key: value[:remaining] for key, value in batch.items()}
+            prediction = rollout_batch(model, batch, device, window).cpu().numpy()
+            predictions.append(prediction)
+            targets.append(batch["u"].numpy())
+            parameters.append(batch["params"].numpy())
+            sample_indices.append(batch["index"].numpy())
+            if first_x is None:
+                first_x, first_t = batch["x"][0].numpy(), batch["t"][0].numpy()
+            processed += prediction.shape[0]
+            print(f"inference batch={batch_index} processed={processed}", flush=True)
+    if not predictions or first_x is None or first_t is None:
+        raise RuntimeError("Inference produced no samples")
+    prediction_array = np.concatenate(predictions)
+    target_array = np.concatenate(targets)
+    parameter_array = np.concatenate(parameters)
+    index_array = np.concatenate(sample_indices)
+    forecast_start = int(inference_cfg["forecast_start_index"])
+    if forecast_start != window:
+        raise ValueError("E3 inference forecast_start_index must equal model.time_window")
+    if not np.array_equal(prediction_array[:, :forecast_start], target_array[:, :forecast_start]):
+        raise AssertionError("Rollout seed differs from the first K ground-truth states")
+    metrics = compute_metrics(prediction_array, target_array, forecast_start)
+    per_time_mse = metrics.pop("per_time_mse")
+    reference_values = config["paper_reference"]["values"]
+    nx = int(prediction_array.shape[2])
+    paper_reference = reference_values.get(nx, reference_values.get(str(nx)))
+    metrics_payload: Dict[str, Any] = {
+        **metrics,
+        "samples": int(prediction_array.shape[0]),
+        "split": split,
+        "nx": nx,
+        "nt": int(prediction_array.shape[1]),
+        "forecast_start_index": forecast_start,
+        "forecast_end_index_inclusive": int(prediction_array.shape[1] - 1),
+        "metric_scope": "forecast_only; accumulated_mse=sum_time(mean_sample_space(error^2))",
+        "checkpoint": str(checkpoint_path),
+        "checkpoint_epoch": int(checkpoint.get("epoch", -1)),
+        "checkpoint_validation_accumulated_mse": float(checkpoint.get("best_validation_accumulated_mse", float("nan"))),
+        "ambiguity_policy": config["experiment"]["ambiguity_policy"],
+        "paper_reference": {"value": paper_reference, "provenance_only_not_assertion": True},
+    }
+    predictions_path = output_dir / Path(config["paths"]["predictions"]).name
+    metrics_path = output_dir / Path(config["paths"]["metrics"]).name
+    atomic_npz(
+        predictions_path,
+        prediction=prediction_array,
+        target=target_array,
+        x=first_x,
+        t=first_t,
+        params=parameter_array,
+        sample_indices=index_array,
+        forecast_start_index=np.asarray(forecast_start, dtype=np.int64),
+        per_time_mse=per_time_mse,
+    )
+    atomic_json(metrics_path, metrics_payload)
+    print(
+        f"metrics samples={metrics_payload['samples']} accumulated_mse={metrics['accumulated_mse']:.8e} "
+        f"rmse={metrics['rmse']:.8e} mae={metrics['mae']:.8e} relative_l2={metrics['relative_l2']:.8e}",
+        flush=True,
+    )
+    print(f"Saved predictions: {predictions_path}", flush=True)
+    print(f"Saved metrics: {metrics_path}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,144 @@
+"""Create MP-PDE E3 visualizations from real inference/training artifacts."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any, Dict
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import yaml
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+
+def load_config(path: Path) -> Dict[str, Any]:
+    with path.open("r", encoding="utf-8") as stream:
+        return yaml.safe_load(stream)
+
+
+def project_path(value: str | Path) -> Path:
+    path = Path(value)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Plot real MP-PDE E3 rollout results")
+    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "config/config.yaml")
+    parser.add_argument("--predictions", type=Path)
+    parser.add_argument("--metrics", type=Path)
+    parser.add_argument("--history", type=Path)
+    parser.add_argument("--output-dir", type=Path)
+    parser.add_argument("--sample-index", type=int)
+    args = parser.parse_args()
+    config = load_config(args.config.resolve())
+    predictions_path = args.predictions or project_path(config["paths"]["predictions"])
+    metrics_path = args.metrics or project_path(config["paths"]["metrics"])
+    history_path = args.history or project_path(config["paths"]["train_history"])
+    output_dir = args.output_dir or project_path(config["paths"]["results"])
+    if not predictions_path.is_file() or not metrics_path.is_file():
+        raise FileNotFoundError(
+            f"Real inference artifacts are required: {predictions_path} and {metrics_path}. "
+            "Run scripts/inference.py after training; placeholder data will not be generated."
+        )
+    with np.load(predictions_path, allow_pickle=False) as archive:
+        required = {"prediction", "target", "x", "t", "params", "sample_indices", "forecast_start_index", "per_time_mse"}
+        missing = required.difference(archive.files)
+        if missing:
+            raise KeyError(f"predictions.npz is missing fields: {sorted(missing)}")
+        prediction, target = archive["prediction"], archive["target"]
+        x, t, per_time_mse = archive["x"], archive["t"], archive["per_time_mse"]
+        forecast_start = int(archive["forecast_start_index"])
+    with metrics_path.open("r", encoding="utf-8") as stream:
+        metrics = json.load(stream)
+    if prediction.shape != target.shape or prediction.ndim != 3:
+        raise ValueError(f"Expected matching [S,T,N] arrays, found {prediction.shape}/{target.shape}")
+    if prediction.shape[1:] != (t.size, x.size) or not np.all(np.isfinite(prediction)) or not np.all(np.isfinite(target)):
+        raise ValueError("Prediction axes do not match x/t or contain non-finite values")
+    recomputed = np.mean((prediction[:, forecast_start:] - target[:, forecast_start:]) ** 2, axis=(0, 2))
+    if per_time_mse.shape != recomputed.shape or not np.allclose(per_time_mse, recomputed, rtol=2e-5, atol=1e-8):
+        raise ValueError("Stored per_time_mse is inconsistent with prediction and target")
+    if not np.isclose(float(metrics["accumulated_mse"]), float(np.sum(recomputed)), rtol=2e-5, atol=1e-8):
+        raise ValueError("metrics.json accumulated_mse is inconsistent with predictions.npz")
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    sample = int(args.sample_index if args.sample_index is not None else config["visualization"]["sample_index"])
+    if sample < 0 or sample >= prediction.shape[0]:
+        raise IndexError(f"sample_index={sample} outside [0,{prediction.shape[0]})")
+    dpi = int(config["visualization"]["dpi"])
+    time_indices = [int(index) for index in config["visualization"]["time_indices"]]
+    if any(index < 0 or index >= t.size for index in time_indices):
+        raise IndexError(f"Configured time_indices exceed nt={t.size}")
+
+    figure, axes = plt.subplots(len(time_indices), 1, figsize=(9, 2.4 * len(time_indices)), sharex=True)
+    axes = np.atleast_1d(axes)
+    for axis, time_index in zip(axes, time_indices):
+        axis.plot(x, target[sample, time_index], color="black", linewidth=1.5, label="target")
+        axis.plot(x, prediction[sample, time_index], color="tab:blue", linewidth=1.2, linestyle="--", label="MP-PDE")
+        axis.set_ylabel("u")
+        axis.set_title(f"t={t[time_index]:.4f}, index={time_index}")
+        axis.grid(alpha=0.2)
+    axes[0].legend(loc="best")
+    axes[-1].set_xlabel("x")
+    figure.tight_layout()
+    rollout_path = output_dir / "e3_rollout.png"
+    figure.savefig(rollout_path, dpi=dpi)
+    plt.close(figure)
+
+    absolute_error = np.abs(prediction[sample] - target[sample])
+    figure, axes = plt.subplots(2, 1, figsize=(10, 7), gridspec_kw={"height_ratios": [2.2, 1.0]})
+    image = axes[0].imshow(
+        absolute_error.T, origin="lower", aspect="auto", extent=(float(t[0]), float(t[-1]), float(x[0]), float(x[-1])), cmap="magma"
+    )
+    axes[0].axvline(float(t[forecast_start]), color="white", linestyle="--", linewidth=1.0, label="forecast start")
+    axes[0].set_ylabel("x")
+    axes[0].set_title("Absolute rollout error")
+    axes[0].legend(loc="upper right")
+    figure.colorbar(image, ax=axes[0], label="|prediction-target|")
+    axes[1].plot(t[forecast_start:], per_time_mse, color="tab:red")
+    axes[1].set_xlabel("t")
+    axes[1].set_ylabel("MSE")
+    axes[1].set_title(f"Per-time MSE; accumulated={metrics['accumulated_mse']:.6g}")
+    axes[1].grid(alpha=0.2)
+    figure.tight_layout()
+    error_path = output_dir / "e3_error.png"
+    figure.savefig(error_path, dpi=dpi)
+    plt.close(figure)
+
+    created = [rollout_path, error_path]
+    if history_path.is_file():
+        with history_path.open("r", encoding="utf-8") as stream:
+            history = json.load(stream)
+        if not isinstance(history, list) or not history:
+            raise ValueError(f"Training history is empty or malformed: {history_path}")
+        epochs = [int(item["epoch"]) + 1 for item in history]
+        figure, axes = plt.subplots(1, 2, figsize=(10, 4))
+        axes[0].plot(epochs, [item["train_rmse"] for item in history], label="train bundle RMSE")
+        axes[0].plot(epochs, [item["validation_bundle_rmse"] for item in history], label="validation bundle RMSE")
+        axes[0].set_yscale("log")
+        axes[0].set_xlabel("epoch")
+        axes[0].set_ylabel("RMSE")
+        axes[0].legend()
+        axes[0].grid(alpha=0.2)
+        axes[1].plot(epochs, [item["validation_accumulated_mse"] for item in history], color="tab:purple")
+        axes[1].set_yscale("log")
+        axes[1].set_xlabel("epoch")
+        axes[1].set_ylabel("validation accumulated MSE")
+        axes[1].grid(alpha=0.2)
+        figure.tight_layout()
+        training_path = output_dir / "training_curve.png"
+        figure.savefig(training_path, dpi=dpi)
+        plt.close(figure)
+        created.append(training_path)
+    else:
+        print(f"Training history not found; skipped training curve: {history_path}", flush=True)
+    for path in created:
+        print(f"Saved figure: {path}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

scripts/smoke_test.py

@@ -0,0 +1,103 @@
+"""Bounded integration smoke test; creates artifacts only in a temporary directory."""
+
+from __future__ import annotations
+
+import copy
+import json
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+from torch.utils.data import DataLoader
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from models.dataset import E3Dataset, generate_e3_hdf5  # noqa: E402
+from models.pde import periodic_neighbor_indices  # noqa: E402
+from scripts.inference import atomic_json, atomic_npz, compute_metrics  # noqa: E402
+from scripts.train import build_model, load_config, rmse_loss, rollout_batch  # noqa: E402
+
+
+def main() -> None:
+    canonical = load_config(PROJECT_ROOT / "config/config.yaml")
+    assert canonical["model"]["hidden_dim"] == 164
+    assert canonical["model"]["time_window"] == 25
+    neighbors = periodic_neighbor_indices(40, canonical["model"]["neighbor_offsets"])
+    assert neighbors.shape == (40, 6)
+    assert torch.all(torch.tensor([row.unique().numel() == 6 for row in neighbors]))
+
+    with tempfile.TemporaryDirectory(prefix="mp_pde_smoke_") as temporary_directory:
+        root = Path(temporary_directory)
+        config = copy.deepcopy(canonical)
+        config["data"]["num_time_points"] = 50
+        config["data"]["high_resolution_nx"] = 40
+        config["data"]["resolution"] = 40
+        config["data"]["train_samples"] = 1
+        config["data"]["valid_samples"] = 1
+        config["data"]["test_samples"] = 1
+        config["data"]["parallel_generation"].update({"workers": 2, "max_in_flight": 2, "flush_every": 1})
+        config["visualization"]["time_indices"] = [25, 30, 40, 49]
+        data_path = root / "e3_smoke.h5"
+        generate_e3_hdf5(config, data_path, sample_counts={"train": 1, "valid": 1, "test": 1})
+        partial_path = data_path.with_suffix(data_path.suffix + ".partial")
+        data_path.replace(partial_path)
+        generate_e3_hdf5(config, data_path, sample_counts={"train": 1, "valid": 1, "test": 1})
+        dataset = E3Dataset(data_path, "test", expected_nt=50, expected_nx=40)
+        batch = next(iter(DataLoader(dataset, batch_size=1, shuffle=False)))
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = build_model(config).to(device)
+        window = 25
+        history = batch["u"][:, :window].transpose(1, 2).to(device)
+        target = batch["u"][:, window : 2 * window].transpose(1, 2).to(device)
+        times = batch["t"]
+        prediction = model(
+            history, batch["x"].to(device), times[:, window - 1].to(device), batch["params"].to(device),
+            (times[:, 1] - times[:, 0]).to(device),
+        )
+        assert prediction.shape == (1, 40, 25)
+        loss = rmse_loss(prediction, target, 1.0e-12)
+        loss.backward()
+        assert np.isfinite(float(loss.detach().cpu()))
+        checkpoint_path = root / "smoke_checkpoint.pth"
+        torch.save({"model_state": model.state_dict(), "resolved_config": config}, checkpoint_path)
+        reloaded = build_model(config).to(device)
+        reloaded.load_state_dict(torch.load(checkpoint_path, map_location=device, weights_only=False)["model_state"], strict=True)
+        reloaded.eval()
+        rollout = rollout_batch(reloaded, batch, device, window).cpu().numpy()
+        target_array = batch["u"].numpy()
+        metrics = compute_metrics(rollout, target_array, window)
+        assert np.isfinite(metrics["accumulated_mse"])
+
+        results = root / "results"
+        per_time_mse = metrics.pop("per_time_mse")
+        atomic_npz(
+            results / "predictions.npz", prediction=rollout, target=target_array, x=batch["x"][0].numpy(),
+            t=batch["t"][0].numpy(), params=batch["params"].numpy(), sample_indices=batch["index"].numpy(),
+            forecast_start_index=np.asarray(window, dtype=np.int64), per_time_mse=per_time_mse,
+        )
+        atomic_json(results / "metrics.json", {**metrics, "samples": 1})
+        with (results / "train_history.json").open("w", encoding="utf-8") as stream:
+            json.dump([{"epoch": 0, "train_rmse": float(loss.detach().cpu()), "validation_bundle_rmse": float(loss.detach().cpu()), "validation_accumulated_mse": metrics["accumulated_mse"]}], stream)
+        config["paths"].update(
+            {"predictions": str(results / "predictions.npz"), "metrics": str(results / "metrics.json"), "train_history": str(results / "train_history.json"), "results": str(results)}
+        )
+        smoke_config = root / "config.yaml"
+        with smoke_config.open("w", encoding="utf-8") as stream:
+            yaml.safe_dump(config, stream, sort_keys=False)
+        subprocess.run([sys.executable, str(PROJECT_ROOT / "scripts/result.py"), "--config", str(smoke_config)], check=True)
+        for name in ("e3_rollout.png", "e3_error.png", "training_curve.png"):
+            assert (results / name).is_file() and (results / name).stat().st_size > 0
+        print(
+            f"smoke_ok device={device} loss={float(loss.detach().cpu()):.8e} "
+            f"accumulated_mse={metrics['accumulated_mse']:.8e}", flush=True,
+        )
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,395 @@
+"""Train the paper-priority MP-PDE E3 solver."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import random
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, Iterable, Mapping, Optional, Tuple
+
+import numpy as np
+import torch
+import yaml
+from torch import Tensor
+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.dataset import E3Dataset, generate_e3_hdf5  # noqa: E402
+from models.pde import MPPDESolver  # noqa: E402
+
+
+def load_config(path: Path) -> Dict[str, Any]:
+    with path.open("r", encoding="utf-8") as stream:
+        config = yaml.safe_load(stream)
+    if not isinstance(config, dict):
+        raise ValueError(f"Config must contain a mapping: {path}")
+    return config
+
+
+def project_path(value: str | Path) -> Path:
+    path = Path(value)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def choose_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 but CUDA is unavailable: {requested}")
+    return device
+
+
+def set_seed(seed: int, deterministic: bool) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    if deterministic:
+        torch.use_deterministic_algorithms(True, warn_only=True)
+        if torch.backends.cudnn.is_available():
+            torch.backends.cudnn.benchmark = False
+            torch.backends.cudnn.deterministic = True
+
+
+def build_model(config: Mapping[str, Any]) -> MPPDESolver:
+    model_cfg, data_cfg = config["model"], config["data"]
+    decoder, scaling = model_cfg["decoder"], model_cfg["scaling"]
+    parameter_maxima = [data_cfg["equation"][name][1] for name in ("alpha_range", "beta_range", "gamma_range")]
+    if model_cfg["aggregation"] != "sum":
+        raise ValueError("Paper Eq. (9) requires model.aggregation=sum")
+    if scaling["solution"] != "none":
+        raise ValueError("The paper does not specify solution normalization; model.scaling.solution must be none")
+    return MPPDESolver(
+        time_window=int(model_cfg["time_window"]),
+        hidden_dim=int(model_cfg["hidden_dim"]),
+        message_passing_layers=int(model_cfg["message_passing_layers"]),
+        neighbor_offsets=model_cfg["neighbor_offsets"],
+        domain_length=float(data_cfg["domain_length"]),
+        final_time=float(data_cfg["final_time"]),
+        parameter_maxima=parameter_maxima,
+        scale_coordinates=bool(scaling["coordinates"]),
+        scale_parameters=bool(scaling["parameters"]),
+        instance_norm_affine=bool(model_cfg["instance_norm_affine"]),
+        decoder_middle_channels=int(decoder["middle_channels"]),
+        decoder_kernels=decoder["kernel_sizes"],
+        decoder_strides=decoder["strides"],
+    )
+
+
+def rmse_loss(prediction: Tensor, target: Tensor, epsilon: float) -> Tensor:
+    return torch.sqrt(torch.mean((prediction - target) ** 2) + epsilon)
+
+
+def _extract_windows(trajectories: Tensor, starts: Tensor, window: int, no_grad_pushes: int) -> Tuple[Tensor, Tensor]:
+    histories, targets = [], []
+    target_offset = window * (no_grad_pushes + 1)
+    for sample, start_tensor in zip(trajectories, starts):
+        start = int(start_tensor)
+        histories.append(sample[start : start + window].transpose(0, 1))
+        targets.append(sample[start + target_offset : start + target_offset + window].transpose(0, 1))
+    return torch.stack(histories), torch.stack(targets)
+
+
+def train_step(
+    batch: Mapping[str, Tensor],
+    model: MPPDESolver,
+    optimizer: torch.optim.Optimizer,
+    device: torch.device,
+    epoch: int,
+    config: Mapping[str, Any],
+    window_generator: torch.Generator,
+) -> Tuple[float, int]:
+    training = config["training"]
+    window = int(config["model"]["time_window"])
+    max_pushes = int(training["max_no_grad_pushes"])
+    no_grad_pushes = 0 if epoch == 0 else int(torch.randint(0, max_pushes + 1, (1,), generator=window_generator))
+    trajectories = batch["u"]
+    num_times = trajectories.shape[1]
+    max_start = num_times - window * (no_grad_pushes + 2)
+    if max_start < 0:
+        raise ValueError(f"nt={num_times} is too short for K={window} and pushes={no_grad_pushes}")
+    starts = torch.randint(0, max_start + 1, (trajectories.shape[0],), generator=window_generator)
+    history, target = _extract_windows(trajectories, starts, window, no_grad_pushes)
+    history, target = history.to(device), target.to(device)
+    x, parameters = batch["x"].to(device), batch["params"].to(device)
+    times = batch["t"]
+    dt = (times[:, 1] - times[:, 0]).to(device)
+    current_time = torch.stack([times[index, int(start) + window - 1] for index, start in enumerate(starts)]).to(device)
+    for _ in range(no_grad_pushes):
+        with torch.no_grad():
+            history = model(history, x, current_time, parameters, dt).detach()
+        current_time = current_time + window * dt
+    optimizer.zero_grad(set_to_none=True)
+    prediction = model(history, x, current_time, parameters, dt)
+    loss = rmse_loss(prediction, target, float(training["loss_epsilon"]))
+    loss.backward()
+    clip = training.get("gradient_clip")
+    if clip is not None:
+        torch.nn.utils.clip_grad_norm_(model.parameters(), float(clip))
+    optimizer.step()
+    return float(loss.detach().cpu()), no_grad_pushes
+
+
+@torch.inference_mode()
+def evaluate_bundle_rmse(
+    model: MPPDESolver, loader: Iterable[Mapping[str, Tensor]], device: torch.device, window: int, epsilon: float
+) -> float:
+    model.eval()
+    squared_error, count = 0.0, 0
+    for batch in loader:
+        history = batch["u"][:, :window].transpose(1, 2).to(device)
+        target = batch["u"][:, window : 2 * window].transpose(1, 2).to(device)
+        x, times, parameters = batch["x"].to(device), batch["t"], batch["params"].to(device)
+        dt = (times[:, 1] - times[:, 0]).to(device)
+        current_time = times[:, window - 1].to(device)
+        prediction = model(history, x, current_time, parameters, dt)
+        squared_error += float(torch.sum((prediction - target) ** 2).cpu())
+        count += prediction.numel()
+    return float(np.sqrt(squared_error / max(count, 1) + epsilon))
+
+
+@torch.inference_mode()
+def rollout_batch(model: MPPDESolver, batch: Mapping[str, Tensor], device: torch.device, window: int) -> Tensor:
+    target = batch["u"].to(device)
+    x, parameters = batch["x"].to(device), batch["params"].to(device)
+    times = batch["t"].to(device)
+    dt = times[:, 1] - times[:, 0]
+    total_times = target.shape[1]
+    prediction = torch.empty_like(target)
+    prediction[:, :window] = target[:, :window]
+    for start in range(window, total_times, window):
+        history = prediction[:, start - window : start].transpose(1, 2)
+        current_time = times[:, start - 1]
+        bundle = model(history, x, current_time, parameters, dt).transpose(1, 2)
+        length = min(window, total_times - start)
+        prediction[:, start : start + length] = bundle[:, :length]
+        if not torch.all(torch.isfinite(bundle)):
+            raise FloatingPointError(f"Non-finite rollout output at forecast index {start}")
+    return prediction
+
+
+@torch.inference_mode()
+def evaluate_rollout(
+    model: MPPDESolver,
+    loader: Iterable[Mapping[str, Tensor]],
+    device: torch.device,
+    window: int,
+    max_samples: Optional[int] = None,
+) -> Dict[str, float]:
+    model.eval()
+    per_time_sse: Optional[Tensor] = None
+    element_count_per_time = 0
+    total_samples = 0
+    for batch in loader:
+        if max_samples is not None:
+            remaining = int(max_samples) - total_samples
+            if remaining <= 0:
+                break
+            if batch["u"].shape[0] > remaining:
+                batch = {key: value[:remaining] for key, value in batch.items()}
+        prediction = rollout_batch(model, batch, device, window)
+        target = batch["u"].to(device)
+        error = prediction[:, window:] - target[:, window:]
+        current = torch.sum(error.double() ** 2, dim=(0, 2)).cpu()
+        per_time_sse = current if per_time_sse is None else per_time_sse + current
+        element_count_per_time += error.shape[0] * error.shape[2]
+        total_samples += error.shape[0]
+    if per_time_sse is None or total_samples == 0:
+        raise RuntimeError("No samples were available for rollout evaluation")
+    per_time_mse = per_time_sse / element_count_per_time
+    return {
+        "accumulated_mse": float(per_time_mse.sum()),
+        "rmse": float(torch.sqrt(per_time_mse.mean())),
+        "samples": float(total_samples),
+    }
+
+
+def _rng_state() -> Dict[str, Any]:
+    return {
+        "python": random.getstate(),
+        "numpy": np.random.get_state(),
+        "torch": torch.get_rng_state(),
+        "cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
+    }
+
+
+def _restore_rng_state(state: Mapping[str, Any]) -> None:
+    random.setstate(state["python"])
+    np.random.set_state(state["numpy"])
+    torch.set_rng_state(state["torch"])
+    if state.get("cuda") is not None and torch.cuda.is_available():
+        torch.cuda.set_rng_state_all(state["cuda"])
+
+
+def atomic_torch_save(payload: Mapping[str, Any], destination: Path) -> None:
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.NamedTemporaryFile(dir=destination.parent, prefix=destination.name + ".", suffix=".tmp", delete=False) as stream:
+        temporary = Path(stream.name)
+    try:
+        torch.save(dict(payload), temporary)
+        os.replace(temporary, destination)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def atomic_json_dump(payload: Any, destination: Path) -> None:
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.NamedTemporaryFile("w", dir=destination.parent, prefix=destination.name + ".", suffix=".tmp", delete=False, encoding="utf-8") as stream:
+        json.dump(payload, stream, indent=2, sort_keys=True)
+        stream.write("\n")
+        temporary = Path(stream.name)
+    os.replace(temporary, destination)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Train the MP-PDE E3 model")
+    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "config/config.yaml")
+    parser.add_argument("--generate-data", action="store_true", help="Generate configured data first when it is absent")
+    parser.add_argument("--overwrite-data", action="store_true")
+    parser.add_argument("--resume", type=Path)
+    parser.add_argument("--init-checkpoint", type=Path)
+    parser.add_argument("--device", type=str)
+    parser.add_argument("--epochs", type=int, help="Explicit run override; does not modify config.yaml")
+    args = parser.parse_args()
+    if args.resume is not None and args.init_checkpoint is not None:
+        parser.error("--resume and --init-checkpoint are mutually exclusive")
+
+    config = load_config(args.config.resolve())
+    training = config["training"]
+    if training["optimizer"] != "AdamW" or training["loss"] != "rmse":
+        raise ValueError("Paper-priority E3 requires AdamW and RMSE")
+    data_path = project_path(config["paths"]["data"])
+    if args.generate_data and (not data_path.exists() or args.overwrite_data):
+        generate_e3_hdf5(config, data_path, overwrite=args.overwrite_data)
+    if not data_path.is_file():
+        raise FileNotFoundError(f"Dataset missing: {data_path}. Run with --generate-data or invoke models/dataset.py")
+
+    seed = int(training["seed"])
+    set_seed(seed, bool(training["deterministic"]))
+    device = choose_device(args.device or str(training["device"]))
+    window = int(config["model"]["time_window"])
+    expected_nt, expected_nx = int(config["data"]["num_time_points"]), int(config["data"]["resolution"])
+    datasets = {split: E3Dataset(data_path, split, expected_nt, expected_nx) for split in ("train", "valid", "test")}
+    loader_generator = torch.Generator().manual_seed(seed)
+    loader_kwargs = dict(batch_size=int(training["batch_size"]), num_workers=int(training["num_workers"]), pin_memory=device.type == "cuda")
+    train_loader = DataLoader(datasets["train"], shuffle=True, generator=loader_generator, **loader_kwargs)
+    valid_loader = DataLoader(datasets["valid"], shuffle=False, **loader_kwargs)
+    test_loader = DataLoader(datasets["test"], shuffle=False, **loader_kwargs)
+    window_generator = torch.Generator().manual_seed(seed + 104729)
+
+    model = build_model(config).to(device)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=float(training["learning_rate"]), weight_decay=float(training["weight_decay"]))
+    scheduler_cfg = training["scheduler"]
+    scheduler = torch.optim.lr_scheduler.MultiStepLR(
+        optimizer, milestones=[int(value) for value in scheduler_cfg["milestones"]], gamma=float(scheduler_cfg["gamma"])
+    ) if scheduler_cfg["enabled"] else None
+    start_epoch, best_metric, history = 0, float("inf"), []
+
+    if args.init_checkpoint is not None:
+        checkpoint = torch.load(args.init_checkpoint, map_location=device, weights_only=False)
+        state = checkpoint.get("model_state", checkpoint)
+        model.load_state_dict(state, strict=True)
+        print(f"Initialized model weights from {args.init_checkpoint}; optimizer and scheduler were reset", flush=True)
+    if args.resume is not None:
+        checkpoint = torch.load(args.resume, map_location=device, weights_only=False)
+        required = {
+            "model_state", "optimizer_state", "epoch", "best_validation_accumulated_mse", "rng_state",
+            "loader_generator_state", "window_generator_state",
+        }
+        missing = required.difference(checkpoint)
+        if missing:
+            raise KeyError(f"Resume checkpoint lacks training state: {sorted(missing)}")
+        model.load_state_dict(checkpoint["model_state"], strict=True)
+        optimizer.load_state_dict(checkpoint["optimizer_state"])
+        if scheduler is not None:
+            if checkpoint.get("scheduler_state") is None:
+                raise KeyError("Resume checkpoint lacks scheduler_state")
+            scheduler.load_state_dict(checkpoint["scheduler_state"])
+        start_epoch = int(checkpoint["epoch"]) + 1
+        best_metric = float(checkpoint["best_validation_accumulated_mse"])
+        history = list(checkpoint.get("history", []))
+        _restore_rng_state(checkpoint["rng_state"])
+        loader_generator.set_state(checkpoint["loader_generator_state"])
+        window_generator.set_state(checkpoint["window_generator_state"])
+        print(f"Resumed training from {args.resume} at epoch {start_epoch}", flush=True)
+
+    epochs = int(args.epochs or training["epochs"])
+    if start_epoch >= epochs:
+        raise ValueError(f"Resume epoch {start_epoch} is not below requested epochs={epochs}")
+    checkpoint_path = project_path(config["paths"]["checkpoint"])
+    history_path = project_path(config["paths"]["train_history"])
+    print(f"Training MP-PDE E3 on device={device}, parameters={sum(p.numel() for p in model.parameters())}", flush=True)
+
+    for epoch in range(start_epoch, epochs):
+        model.train()
+        step_losses = []
+        for step, batch in enumerate(train_loader, start=1):
+            loss, pushes = train_step(batch, model, optimizer, device, epoch, config, window_generator)
+            step_losses.append(loss)
+            if step == 1 or step % int(training["print_interval"]) == 0 or step == len(train_loader):
+                learning_rate = optimizer.param_groups[0]["lr"]
+                print(
+                    f"epoch={epoch + 1}/{epochs} step={step}/{len(train_loader)} pushes={pushes} "
+                    f"loss_rmse={loss:.8e} lr={learning_rate:.3e}", flush=True,
+                )
+        validation_rmse = evaluate_bundle_rmse(model, valid_loader, device, window, float(training["loss_epsilon"]))
+        validation_rollout = evaluate_rollout(model, valid_loader, device, window, training.get("validation_rollout_samples"))
+        train_rmse = float(np.mean(step_losses))
+        learning_rate = float(optimizer.param_groups[0]["lr"])
+        epoch_record: Dict[str, Any] = {
+            "epoch": epoch,
+            "train_rmse": train_rmse,
+            "validation_bundle_rmse": validation_rmse,
+            "validation_accumulated_mse": validation_rollout["accumulated_mse"],
+            "learning_rate": learning_rate,
+        }
+        print(
+            f"epoch={epoch + 1}/{epochs} train_rmse={train_rmse:.8e} val_rmse={validation_rmse:.8e} "
+            f"val_accumulated_mse={validation_rollout['accumulated_mse']:.8e} lr={learning_rate:.3e}", flush=True,
+        )
+        improved = validation_rollout["accumulated_mse"] < best_metric
+        if scheduler is not None:
+            scheduler.step()
+        if improved:
+            best_metric = validation_rollout["accumulated_mse"]
+            test_rollout = evaluate_rollout(model, test_loader, device, window, training.get("test_rollout_samples"))
+            epoch_record["test_accumulated_mse_on_new_best"] = test_rollout["accumulated_mse"]
+            epoch_record["test_rmse_on_new_best"] = test_rollout["rmse"]
+            print(
+                f"new_best epoch={epoch + 1} val_accumulated_mse={best_metric:.8e} "
+                f"test_accumulated_mse={test_rollout['accumulated_mse']:.8e} test_rmse={test_rollout['rmse']:.8e}",
+                flush=True,
+            )
+        history.append(epoch_record)
+        atomic_json_dump(history, history_path)
+        if improved:
+            atomic_torch_save(
+                {
+                    "model_state": model.state_dict(),
+                    "optimizer_state": optimizer.state_dict(),
+                    "scheduler_state": scheduler.state_dict() if scheduler is not None else None,
+                    "epoch": epoch,
+                    "best_validation_accumulated_mse": best_metric,
+                    "resolved_config": config,
+                    "rng_state": _rng_state(),
+                    "loader_generator_state": loader_generator.get_state(),
+                    "window_generator_state": window_generator.get_state(),
+                    "history": history,
+                },
+                checkpoint_path,
+            )
+            print(f"Saved best checkpoint: {checkpoint_path}", flush=True)
+    print(f"Training complete; best_validation_accumulated_mse={best_metric:.8e}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

weight/best_model.pth

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