Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,10 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+results/figures/test_case_800_profiles.png filter=lfs diff=lfs merge=lfs -text
+results/figures/test_case_880_profiles.png filter=lfs diff=lfs merge=lfs -text
+results/figures/test_case_960_profiles.png filter=lfs diff=lfs merge=lfs -text
+results/figures/test_case_994_fields.png filter=lfs diff=lfs merge=lfs -text
+results/figures/test_case_994_profiles.png filter=lfs diff=lfs merge=lfs -text
+results/figures/test_error_summary.png filter=lfs diff=lfs merge=lfs -text
+results/figures/training_curve.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,37 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- fluid-dynamics
+library_name: pytorch
+---
+
+# RONOM
+
+## Model
+
+This package reproduces the one-dimensional viscous Burgers input generalization experiment. The trained checkpoint is `weight/best_model.pth`; model code, configuration, training, inference, and visualization scripts are included. The package was trained once; reported numbers are not guarantees for another run.
+
+## Data
+
+Download [OneScience/RONOM-Burgers](https://www.modelscope.cn/datasets/OneScience/RONOM-Burgers) with `modelscope download --dataset OneScience/RONOM-Burgers --local_dir ./data`. Burgers.mat contains 1,000 trajectories: input [1000,1025], output [1000,101,1025], tspan [1,101], float64. Split: 600 train, 200 validation, 200 test.
+
+Data provenance: [https://zenodo.org/records/21451819](https://zenodo.org/records/21451819).
+
+## Usage
+
+```bash
+hf download OneScience-Group/RONOM --local-dir ./RONOM
+cd RONOM
+pip install torch numpy scipy h5py pyyaml matplotlib
+python scripts/inference.py
+python scripts/result.py
+```
+
+For training, run `python scripts/train.py` after installing the appropriate accelerator-enabled PyTorch environment and downloading the data above. See `config/config.yaml` for experiment parameters. Preserve the included checkpoint before retraining.
+
+## Result and citation
+
+test trajectory RMSE: 0.00170540; paper Table 1(a): 0.00134. See `results/` for the recorded outputs and the [original paper or data page](https://zenodo.org/records/21451819) for the research context. The dataset license, where not separately stated by its author, must not be inferred from the code license.

config.json

@@ -0,0 +1,10 @@
+{
+  "model_type": "ronom_burgers",
+  "architectures": [
+    "RONOM"
+  ],
+  "framework": "pytorch",
+  "experiment": "one-dimensional viscous Burgers input generalization",
+  "checkpoint": "weight/best_model.pth",
+  "dataset": "OneScience/RONOM-Burgers"
+}

config/config.yaml

@@ -0,0 +1,42 @@
+seed: 0
+data:
+  directory: data
+  filename: Burgers.mat
+  train_samples: 600
+  validation_samples: 200
+  test_samples: 200
+  space_stride: 32
+  time_stride: 10
+model:
+  grid_size: 33
+  features: [64, 128, 128, 256]
+  latent_channels: 4
+  fourier_features: 32
+  velocity_widths: [64, 64, 64]
+  padding_size: 2
+  alpha: 1.0e-5
+  sobolev_reg: 1.0e-8
+  bc_reg: 10.0
+  step_size: 0.1
+training:
+  epochs: 5000
+  batch_size: 32
+  learning_rate: 0.001
+  weight_decay: 0.0
+  scheduler_frequency: 250
+  scheduler_gamma: 0.9
+  validation_frequency: 1000
+  checkpoint_frequency: 100
+  jit_flow: true
+  device: cuda
+  num_threads: 8
+paths:
+  results: results
+  checkpoint: weight/best_model.pth
+  last_checkpoint: weight/last_model.pth
+inference:
+  batch_size: 32
+  device: cuda
+visualization:
+  sample_indices: [0, 80, 160]
+  dpi: 200

model/__init__.py

@@ -0,0 +1 @@
+from .ronom import RONOM

model/ronom.py

@@ -0,0 +1,165 @@
+"""Independent implementation of RONOM, arXiv:2507.12814v2, Sec. 2 and SM3.
+
+The projection solves weighted regularized least squares in a Gaussian basis;
+it is not a learned Fourier operator. All geometric matrices are path independent.
+"""
+import math
+from typing import List
+
+import numpy as np
+from scipy.special import erf
+import torch
+from torch import nn
+
+
+def gaussian(x, centers, sigma):
+    return np.exp(-0.5 * ((x[:, None] - centers[None, :]) / sigma) ** 2)
+
+
+def laplacian_gram(centers, sigma):
+    """Closed-form integral of phi_i'' phi_j'' on [0,1], in float64.
+
+    Complete the square at the midpoint and integrate Gaussian moments 0,2,4.
+    The recurrence follows integration by parts, without numerical quadrature.
+    """
+    mid = (centers[:, None] + centers[None, :]) / 2
+    dist = centers[:, None] - centers[None, :]
+    lo, hi = -mid, 1 - mid
+    ea, eb = np.exp(-(lo / sigma)**2), np.exp(-(hi / sigma)**2)
+    j0 = sigma * math.sqrt(math.pi) / 2 * (erf(hi / sigma) - erf(lo / sigma))
+    j2 = sigma**2 / 2 * (j0 - hi * eb + lo * ea)
+    j4 = sigma**2 / 2 * (3 * j2 - hi**3 * eb + lo**3 * ea)
+    value = np.exp(-dist**2 / (4 * sigma**2)) / sigma**8 * (
+        j4 - (dist**2 / 2 + 2 * sigma**2) * j2
+        + (dist**2 / 4 - sigma**2)**2 * j0)
+    return (value + value.T) / 2
+
+
+def projection_matrix(x_in, x_measure, centers, sigma, alpha, sobolev_reg, bc_reg):
+    basis = gaussian(x_in, centers, sigma)
+    # Trapezoidal integral weights; all input-generalization points are uniform.
+    w = np.empty(len(x_in), dtype=np.float64)
+    w[0], w[-1] = (x_in[1] - x_in[0]) / 2, (x_in[-1] - x_in[-2]) / 2
+    w[1:-1] = (x_in[2:] - x_in[:-2]) / 2
+    edge = gaussian(np.array([0., 1.]), centers, sigma)
+    # Official residual convention scales the derivative residual by sigma.
+    deriv = -(np.array([0., 1.])[:, None] - centers) / sigma * edge
+    residual = np.stack([edge[0] - edge[1], deriv[0] - deriv[1]])
+    rhs = basis.T * w
+    normal = (rhs @ basis + sobolev_reg * laplacian_gram(centers, sigma)
+              + bc_reg * (residual.T @ residual) + alpha * np.eye(len(centers)))
+    return gaussian(x_measure, centers, sigma) @ np.linalg.solve(normal, rhs)
+
+
+class LatentFlow(nn.Module):
+    """Autonomous neural ODE with the explicit fourth order 3/8 tableau."""
+    def __init__(self, dimension, widths, step_size):
+        super().__init__()
+        layers = []
+        prev = dimension
+        for width in widths:
+            layers += [nn.Linear(prev, width), nn.GELU()]
+            prev = width
+        layers.append(nn.Linear(prev, dimension))
+        self.velocity = nn.Sequential(*layers)
+        self.step_size = float(step_size)
+
+    def forward(self, initial: torch.Tensor, times: torch.Tensor) -> torch.Tensor:
+        state = initial
+        states = torch.jit.annotate(List[torch.Tensor], [])
+        states.append(state)
+        for i in range(times.numel() - 1):
+            # Sampling times are 0,.1,...,1. Substeps also allow a coarser query.
+            interval = float(times[i + 1] - times[i])
+            count = max(1, int(math.ceil(interval / self.step_size - 1e-6)))
+            h = interval / count
+            for _ in range(count):
+                a = self.velocity(state)
+                b = self.velocity(state + (h / 3.) * a)
+                c = self.velocity(state + h * (b - a / 3.))
+                d = self.velocity(state + h * (a - b + c))
+                state = state + (h / 8.) * (a + 3. * b + 3. * c + d)
+            states.append(state)
+        return torch.stack(states, dim=1)
+
+
+class RONOM(nn.Module):
+    def __init__(self, grid_size=33, features=(64, 128, 128, 256),
+                 latent_channels=4, fourier_features=32, velocity_widths=(64, 64, 64),
+                 padding_size=2, alpha=1e-5, sobolev_reg=1e-8, bc_reg=10.,
+                 step_size=0.1):
+        super().__init__()
+        self.grid_size = int(grid_size)
+        self.latent_channels = int(latent_channels)
+        self.alpha, self.sobolev_reg, self.bc_reg = alpha, sobolev_reg, bc_reg
+        self.sigma = 1 / ((grid_size - 1) * math.sqrt(-2 * math.log(0.5)))
+        self.register_buffer('frequencies', torch.randn(fourier_features, 1))
+        f = list(features)
+        self.lift = nn.Sequential(nn.Conv1d(1 + 2 * fourier_features, 64, 3, padding=1),
+                                  nn.ReLU(), nn.Conv1d(64, f[0] // 2, 3, padding=1))
+        self.output_project = nn.Sequential(
+            nn.ConvTranspose1d(f[0] // 2, 64, 3, padding=1), nn.ReLU(),
+            nn.ConvTranspose1d(64, 1, 3, padding=1))
+        enc = []
+        prev = f[0] // 2
+        for width in f:
+            enc += [nn.Conv1d(prev, width, 3, padding=1), nn.ReLU()]
+            prev = width
+        self.encoder = nn.Sequential(*enc)
+        self.to_latent = nn.Conv1d(f[-1], latent_channels, 3, padding=1)
+        dec = []
+        back = list(reversed(f)) + [f[0] // 2]
+        for left, right in zip(back[:-1], back[1:]):
+            dec += [nn.ConvTranspose1d(left, right, 3, padding=1), nn.ReLU()]
+        self.decoder = nn.Sequential(*dec)
+        self.from_latent = nn.ConvTranspose1d(latent_channels, f[-1], 3, padding=1)
+        x = np.linspace(0., 1., grid_size)
+        centers = np.arange(-padding_size, grid_size + padding_size) / (grid_size - 1)
+        proj = projection_matrix(x, x, centers, self.sigma, alpha, sobolev_reg, bc_reg)
+        kernel = gaussian(x, x, self.sigma)
+        recovery = np.linalg.solve(kernel + alpha * np.eye(grid_size), kernel.T).T
+        self.register_buffer('grid', torch.tensor(x, dtype=torch.float32))
+        self.register_buffer('centers', torch.tensor(centers, dtype=torch.float64))
+        self.register_buffer('projection', torch.tensor(proj, dtype=torch.float32))
+        self.register_buffer('recovery', torch.tensor(recovery, dtype=torch.float32))
+        self.latent_flow = LatentFlow(grid_size * latent_channels, velocity_widths, step_size)
+
+    def encode(self, initial, x_input=None):
+        if initial.ndim == 3:
+            initial = initial.squeeze(-1)
+        if x_input is None:
+            projected = initial @ self.projection.T
+        else:
+            matrix = projection_matrix(x_input.detach().cpu().double().numpy(),
+                self.grid.detach().cpu().double().numpy(), self.centers.cpu().numpy(),
+                self.sigma, self.alpha, self.sobolev_reg, self.bc_reg)
+            projected = initial @ torch.as_tensor(matrix.T, device=initial.device, dtype=initial.dtype)
+        phase = 2 * math.pi * self.grid[:, None] @ self.frequencies.T
+        positional = torch.cat([phase.cos(), phase.sin()], dim=-1).T
+        lifted_input = torch.cat([projected[:, None, :],
+                                  positional[None].expand(initial.shape[0], -1, -1)], dim=1)
+        return self.to_latent(self.encoder(self.lift(lifted_input)))
+
+    def decode(self, latent, x_output=None):
+        b, nt = latent.shape[:2]
+        features = latent.reshape(b * nt, self.latent_channels, self.grid_size)
+        values = self.output_project(self.decoder(self.from_latent(features)))
+        values = values.reshape(b, nt, self.grid_size)
+        if x_output is None:
+            return values @ self.recovery.T
+        train_x = self.grid.detach().cpu().double().numpy()
+        target_x = x_output.detach().cpu().double().numpy()
+        kernel = gaussian(train_x, train_x, self.sigma)
+        matrix = np.linalg.solve(kernel + self.alpha * np.eye(self.grid_size),
+                                  gaussian(target_x, train_x, self.sigma).T).T
+        return values @ torch.as_tensor(matrix.T, device=values.device, dtype=values.dtype)
+
+    def forward(self, initial, times, return_latent=False):
+        # CPU time coordinates avoid a host/device synchronization at every RK stage.
+        z0 = self.encode(initial)
+        zt = self.latent_flow(z0.flatten(1), times)
+        zt = zt.reshape(initial.shape[0], times.numel(), self.latent_channels, self.grid_size)
+        prediction = self.decode(zt)
+        if return_latent:
+            return prediction, zt
+        return prediction

results/best_validation_metrics.json

@@ -0,0 +1,23 @@
+{
+  "epoch": 5000,
+  "split": "validation",
+  "samples": 200,
+  "mse": 5.125773963908943e-06,
+  "global_rmse": 0.002264017218112297,
+  "trajectory_rmse_mean": 0.0016242767846566335,
+  "trajectory_rmse_std_ddof1": 0.0015811384775408664,
+  "relative_l2_mean": 0.009477731318529135,
+  "rmse_by_time": [
+    0.0026190899833518078,
+    0.0022809309520447386,
+    0.002138092286151645,
+    0.002365773666200291,
+    0.0027355348428027863,
+    0.002744344484310206,
+    0.0023092583850567467,
+    0.0019961579303643937,
+    0.0018891513302002588,
+    0.001776062215740342,
+    0.0017600350993885276
+  ]
+}

results/data_provenance.json

@@ -0,0 +1,42 @@
+{
+  "file": "/public/home/liuyushuang/zhaomingsi/onescience-cfd-data/RONOM/Burgers.mat",
+  "sha256": "5c4f361d8318d4a00c61d4d57f3d49eeeb6643096b7ec356a059f0f95d014c71",
+  "metadata": {
+    "gamma": [
+      [
+        4
+      ]
+    ],
+    "sigma": [
+      [
+        62500
+      ]
+    ],
+    "tau": [
+      [
+        25
+      ]
+    ],
+    "visc": [
+      [
+        0.01
+      ]
+    ]
+  },
+  "provenance": "User-prepared original Zenodo record 21451819; no redownload",
+  "shape": {
+    "input": [
+      1000,
+      1025
+    ],
+    "output": [
+      1000,
+      101,
+      1025
+    ],
+    "tspan": [
+      1,
+      101
+    ]
+  }
+}

results/reproduction_summary.json

@@ -0,0 +1,95 @@
+{
+  "status": "completed",
+  "training": {
+    "status": "completed",
+    "completed_epochs": 5000,
+    "batches_per_epoch": 19,
+    "training_samples_per_epoch": 600,
+    "best_epoch": 5000,
+    "best_validation_mse": 5.125773963908943e-06,
+    "selection_split": "validation[600:800]",
+    "test_used_for_selection": false,
+    "checkpoint": "weight/best_model.pth",
+    "data_sha256": "5c4f361d8318d4a00c61d4d57f3d49eeeb6643096b7ec356a059f0f95d014c71",
+    "elapsed_seconds": 8878.240034841001,
+    "slurm_job_id": "1060058"
+  },
+  "independent_test": {
+    "samples": 200,
+    "mse": 4.673988774995294e-06,
+    "global_rmse": 0.002161940973985019,
+    "trajectory_rmse_mean": 0.001705397173261511,
+    "trajectory_rmse_std_ddof1": 0.0013320967174911497,
+    "relative_l2_mean": 0.00986708805065414,
+    "rmse_by_time": [
+      0.002461057978856684,
+      0.002075854227337168,
+      0.001999032913809955,
+      0.0020728556299004206,
+      0.0018694716441612323,
+      0.0022200476157488527,
+      0.0024226956088001946,
+      0.002467867249023815,
+      0.002239861862973637,
+      0.001978805748854005,
+      0.0018544607168863218
+    ],
+    "split": "independent_test",
+    "indices": [
+      800,
+      999
+    ],
+    "checkpoint_epoch": 5000,
+    "checkpoint_sha256": "fd6917b63e596ad919bf7ad8504794039d303a3910925acac75f3fccd4ac08a8",
+    "checkpoint_validation_mse": 5.125773963908943e-06,
+    "selection_split": "validation[600:800]",
+    "paper_rmse_mean": 0.00134,
+    "paper_rmse_std": 0.00162,
+    "elapsed_seconds": 9.619060612050816,
+    "slurm_job_id": "1060058"
+  },
+  "verification": {
+    "status": "passed",
+    "history_epochs": 5000,
+    "training_samples_per_epoch": 600,
+    "batches_per_epoch": 19,
+    "validation_epochs": [
+      1,
+      1001,
+      2001,
+      3001,
+      4001,
+      5000
+    ],
+    "best_epoch": 5000,
+    "best_validation_mse": 5.125773963908943e-06,
+    "test_indices": [
+      800,
+      999
+    ],
+    "test_samples": 200,
+    "test_shape": [
+      200,
+      11,
+      33
+    ],
+    "latent_shape": [
+      200,
+      11,
+      4,
+      33
+    ],
+    "recomputed_rmse_mean": 0.0017053971732615108,
+    "recomputed_rmse_std_ddof1": 0.0013320967174911497,
+    "global_rmse": 0.002161940973985019,
+    "checkpoint_sha256": "fd6917b63e596ad919bf7ad8504794039d303a3910925acac75f3fccd4ac08a8",
+    "figures_verified": 20,
+    "paper_rmse_mean": 0.00134,
+    "mean_rmse_gap_percent": 27.268445765784378,
+    "test_used_for_selection": false,
+    "source_data_sha256": "5c4f361d8318d4a00c61d4d57f3d49eeeb6643096b7ec356a059f0f95d014c71"
+  },
+  "figures": 20,
+  "paper": "https://arxiv.org/pdf/2507.12814v2",
+  "numeric_match_to_paper": false
+}

results/runtime_result.json

@@ -0,0 +1,28 @@
+{
+  "execution_channel": "local_slurm",
+  "config_source": "onescience.json",
+  "preflight_passed": true,
+  "submission_state": "submitted",
+  "execution_state": "COMPLETED",
+  "exit_code": "0:0",
+  "job_id": "1060058",
+  "smoke_job_id": "1059370",
+  "node": "a01r3n02",
+  "resources": {
+    "nodes": 1,
+    "dcu": 1,
+    "cpus": 8,
+    "memory": "30G"
+  },
+  "slurm_elapsed": "02:29:24",
+  "local_log_dir": "/public/home/liuyushuang/zhaomingsi/onescience-cfd-repro-packages/RONOM/logs",
+  "logs": [
+    "/public/home/liuyushuang/zhaomingsi/onescience-cfd-repro-packages/RONOM/results/train.log",
+    "/public/home/liuyushuang/zhaomingsi/onescience-cfd-repro-packages/RONOM/results/inference.log",
+    "/public/home/liuyushuang/zhaomingsi/onescience-cfd-repro-packages/RONOM/results/visualization.log"
+  ],
+  "completed_epochs": 5000,
+  "test_samples": 200,
+  "figures": 20,
+  "blocking_reason": null
+}

results/table1a_comparison.csv

@@ -0,0 +1,3 @@
+source,split,samples,rmse_mean,rmse_std_ddof1,units
+paper_Table1a_RONOM,paper_test,200,0.00134,0.00162,raw
+this_run,independent_test_800_999,200,0.001705397173261511,0.0013320967174911497,raw

results/test_metrics.json

@@ -0,0 +1,34 @@
+{
+  "samples": 200,
+  "mse": 4.673988774995294e-06,
+  "global_rmse": 0.002161940973985019,
+  "trajectory_rmse_mean": 0.001705397173261511,
+  "trajectory_rmse_std_ddof1": 0.0013320967174911497,
+  "relative_l2_mean": 0.00986708805065414,
+  "rmse_by_time": [
+    0.002461057978856684,
+    0.002075854227337168,
+    0.001999032913809955,
+    0.0020728556299004206,
+    0.0018694716441612323,
+    0.0022200476157488527,
+    0.0024226956088001946,
+    0.002467867249023815,
+    0.002239861862973637,
+    0.001978805748854005,
+    0.0018544607168863218
+  ],
+  "split": "independent_test",
+  "indices": [
+    800,
+    999
+  ],
+  "checkpoint_epoch": 5000,
+  "checkpoint_sha256": "fd6917b63e596ad919bf7ad8504794039d303a3910925acac75f3fccd4ac08a8",
+  "checkpoint_validation_mse": 5.125773963908943e-06,
+  "selection_split": "validation[600:800]",
+  "paper_rmse_mean": 0.00134,
+  "paper_rmse_std": 0.00162,
+  "elapsed_seconds": 9.619060612050816,
+  "slurm_job_id": "1060058"
+}

results/test_per_sample.csv

@@ -0,0 +1,201 @@
+sample_id,rmse,relative_l2
+800,0.0014441642533977318,0.00729123012996783
+801,0.002415821063563245,0.00960939631888287
+802,0.0011914687871388155,0.008451511382333349
+803,0.0012626003363246977,0.007894797343731797
+804,0.001998250550575508,0.008627655792169167
+805,0.0008154827869567817,0.010417249737812542
+806,0.0011013256291670644,0.00666085619823366
+807,0.0015790514604787243,0.008515506502624587
+808,0.003363841125022772,0.015478086942061958
+809,0.001860558779885756,0.010520708433998343
+810,0.0015722720141994162,0.009038333846308927
+811,0.0013418697131035065,0.008836094380108874
+812,0.0035894301379888758,0.01612087508704082
+813,0.0012646054778170076,0.010149752156164118
+814,0.002119074684764565,0.009492618292748448
+815,0.0007998953236767689,0.01666992334013468
+816,0.0011044745489972795,0.0067535312444658045
+817,0.0012831633204547967,0.0080612579188031
+818,0.000712199600738031,0.010070663533787555
+819,0.0011953804786144532,0.007086889618089626
+820,0.0011368105585113162,0.010047049930414158
+821,0.0014393966085775218,0.007766688394612785
+822,0.0028297305790365543,0.014422616433333901
+823,0.0008518209117680049,0.010362130183934495
+824,0.0015096009527889858,0.010858929407126289
+825,0.0009111668439799336,0.006464505910483904
+826,0.0011160574425272873,0.007794441801974531
+827,0.0011383052123295698,0.008206011598699653
+828,0.0008093319939719636,0.005389986818021365
+829,0.0015325703497033512,0.008833484364188529
+830,0.0011792855042533938,0.006164376190559706
+831,0.0013148246876960113,0.00604776096694483
+832,0.001455539736817489,0.007089218935786096
+833,0.001300337276554276,0.01375691214333825
+834,0.0018383855676432533,0.00802335054490402
+835,0.006325736936347732,0.03315410707726788
+836,0.0008055008257981556,0.014882097065569967
+837,0.0025166074208316557,0.011857703471058157
+838,0.0015345968253095863,0.010198831702372176
+839,0.0007412714997609998,0.007310470754295764
+840,0.0008701879555373449,0.005872218593710851
+841,0.0012929898835578905,0.008509449938652292
+842,0.0012409216815030717,0.0074217976140607276
+843,0.001135883691994004,0.009481252722396994
+844,0.0012506143762502397,0.005568246185989062
+845,0.0011502996472816083,0.009546127334200056
+846,0.0008205328328738069,0.013193053100485566
+847,0.0024833228468586752,0.010898839482961392
+848,0.0012799428387503683,0.007132848003360383
+849,0.0009348639236117146,0.008923647952120545
+850,0.007537085395672275,0.02642515778598448
+851,0.000959899735035825,0.008902070500208876
+852,0.000986228428168412,0.005933525267133642
+853,0.0009762482221984804,0.009761204137796731
+854,0.0013718022317925273,0.010185100781627281
+855,0.0010658512589603404,0.00919749785675796
+856,0.0017357300505692697,0.011288073938221118
+857,0.0014464055179929172,0.00831458797149399
+858,0.0010702700699737973,0.008537243183456014
+859,0.0010050917335225926,0.006474680376006104
+860,0.00306822761824498,0.014621531613193026
+861,0.0033928099370821195,0.0185655581334803
+862,0.0013527266189997836,0.006640214040641837
+863,0.0015559520463899752,0.010385949127214997
+864,0.001105718633584989,0.005764066935622933
+865,0.0010621524921265507,0.007696815553640195
+866,0.001261436300985645,0.00645613708700552
+867,0.0017576913752563894,0.010991553427473917
+868,0.0013099092578256833,0.009590739739171884
+869,0.0007613496377322205,0.006311841112983999
+870,0.0017574812372593408,0.009835706150111607
+871,0.004579392292460886,0.014371718283453802
+872,0.0015475964335534154,0.006230175452350933
+873,0.00148630639030871,0.009307520959803396
+874,0.0016762723633592241,0.006335640785602478
+875,0.0008748569414768277,0.0055368537345673935
+876,0.0017752065839310906,0.009570745165451207
+877,0.0009693422221720172,0.008563731232211389
+878,0.0007937653405507869,0.005892526467193541
+879,0.001039564762835509,0.00665225038016452
+880,0.001457881787085484,0.006899389637768518
+881,0.0018651977924340215,0.008489585871742623
+882,0.0010858070845518957,0.005558571478030946
+883,0.0012538390100052196,0.007565115823016464
+884,0.0008817052410440639,0.0046771205813646004
+885,0.0009310306394476048,0.006235420755048469
+886,0.0013793459972799656,0.007405686863358377
+887,0.0020846975641691766,0.013736941560302536
+888,0.0011948592085205228,0.0075800670034093754
+889,0.0011697477719223286,0.005507243210246042
+890,0.004551107853771669,0.020819895890820953
+891,0.003637871440803657,0.0163985772348673
+892,0.0007941509008577267,0.005470727708808298
+893,0.0012703569185660787,0.007230502272474798
+894,0.0030672970438073593,0.01151415953579286
+895,0.0011810650603748742,0.006550353576006067
+896,0.0012441613372044002,0.014603132185172781
+897,0.0029938960498603505,0.014752847710556592
+898,0.0008727711794136182,0.005515131326416716
+899,0.0011201524781288253,0.007810446284357195
+900,0.001233991416089354,0.007057580632594449
+901,0.0011650017881304974,0.007650767564492308
+902,0.0011996249695272127,0.006720864212309036
+903,0.0012370310054518714,0.0074587668715494675
+904,0.0009220205593254075,0.007096993832622519
+905,0.0009633853193826409,0.008508680908907603
+906,0.0007135869913291987,0.011259870770668952
+907,0.0010521268966999114,0.008066662355631086
+908,0.0009190429554308218,0.010030657046073343
+909,0.0013659371418535157,0.012432603051923527
+910,0.0009393835036039246,0.010479185025272513
+911,0.0009575362081858569,0.009120206511523378
+912,0.002702246207780349,0.010972459690995962
+913,0.0011266866886340261,0.0072892865526466
+914,0.0016682201969247475,0.008558695096768583
+915,0.0015066277308996962,0.008812798822178306
+916,0.0015375844064554301,0.008466948786086675
+917,0.0018666939150836967,0.008781869201336296
+918,0.004188875524084949,0.01215421534930819
+919,0.0031559402816445405,0.011026492850696725
+920,0.0009635181897632188,0.00826463296361614
+921,0.0035347368818962905,0.015725360016018052
+922,0.0014358204245286086,0.012940127641777835
+923,0.0009964901053022947,0.010119682500378042
+924,0.0009148418169323768,0.007392940254705161
+925,0.0016691999941050087,0.011411197426739681
+926,0.002387696487705466,0.012382249128682523
+927,0.0014521756816838005,0.00765643279091908
+928,0.0009902341380583805,0.01413048801942241
+929,0.002825700257807701,0.015226357917784709
+930,0.005160293489520457,0.016374975027971365
+931,0.0014003052210915922,0.006057548485411315
+932,0.0022625319342623688,0.009827343709087302
+933,0.0015242896197079703,0.00821868161145458
+934,0.001033083962603733,0.007102402629403483
+935,0.001420422999209446,0.0060663995460390014
+936,0.0009671579441707368,0.011546146828978348
+937,0.0012568418093653058,0.00954057253659006
+938,0.0019279405865385943,0.008715780498962775
+939,0.004221112616735485,0.018773587824957756
+940,0.0009647876962190149,0.007381329255001716
+941,0.0007270810859020719,0.008070871146743775
+942,0.0015694417270356253,0.01053232539653554
+943,0.0012707776046804089,0.01294442406153623
+944,0.004469268902409168,0.01791724149430663
+945,0.0009714431061166003,0.009614952188550357
+946,0.0014078634321375632,0.00938219391916491
+947,0.0009141998010265314,0.009410290158824807
+948,0.002605327225377241,0.01058111719026333
+949,0.0009009842271187858,0.006988902258785145
+950,0.0019418328237304457,0.01089337040143236
+951,0.001839604033349563,0.008286063786727929
+952,0.0012464553272033385,0.005839170559141925
+953,0.0016172507856935163,0.007244302616848377
+954,0.0013819677906637625,0.005714478430086377
+955,0.0013170377048367344,0.009037444084165236
+956,0.0010936816214972879,0.008607176858602072
+957,0.0012824458019790733,0.009049715234738824
+958,0.0007623849167848237,0.007785652530933292
+959,0.0012166669716038524,0.007187985134346414
+960,0.0015404606598713488,0.010387211792061974
+961,0.0009880663253630198,0.008116852182480316
+962,0.0012553970757991468,0.00845053892089981
+963,0.000702529926807744,0.00587484272041473
+964,0.0032501025434846158,0.011024587248372082
+965,0.0019096188518064033,0.012373153384138191
+966,0.0009021126397589328,0.009761469355660276
+967,0.0013885713963696115,0.006857308533615614
+968,0.0020933809454012015,0.00818737999326205
+969,0.0015117437447880135,0.012136279960406511
+970,0.0010510918757149453,0.004821037008264643
+971,0.0025011940432743177,0.015027695090694041
+972,0.003646446780417154,0.013519626547469004
+973,0.0011198295822278475,0.007990426050346788
+974,0.0014245268649068067,0.006864751006748974
+975,0.0012871312827115028,0.007565817128583534
+976,0.000886059705367791,0.005678057728736154
+977,0.000955530971755569,0.009408212237748423
+978,0.001069682752971588,0.0064699506263734
+979,0.0010199505489099584,0.005575725699605437
+980,0.0024029757571904214,0.009653973749049586
+981,0.0007437867902824578,0.01704335584185775
+982,0.004265423452123864,0.01586745992820094
+983,0.002167071679798764,0.012244394124590935
+984,0.000929962437772895,0.00893554650136493
+985,0.002525234246333844,0.013478586959083477
+986,0.0022838601217988913,0.011788063646633392
+987,0.003402670753787411,0.013849056177897001
+988,0.0010771493054250135,0.00789478518348717
+989,0.001083582156428052,0.006593746009821497
+990,0.0007881397517817494,0.011011966524056654
+991,0.006703697063184707,0.023413969354371778
+992,0.0012929017590327214,0.006987158099561493
+993,0.0011427798511789333,0.008311275758676951
+994,0.012773607533505092,0.035389758709408387
+995,0.0009239440952716762,0.0157285345242972
+996,0.0014329196350219976,0.012676394059059157
+997,0.001202171970288141,0.006952197076749028
+998,0.0017470733328887329,0.011172317814995052
+999,0.001339933014064422,0.006860201160684589

results/test_predictions.npz

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

results/training_summary.json

@@ -0,0 +1,14 @@
+{
+  "status": "completed",
+  "completed_epochs": 5000,
+  "batches_per_epoch": 19,
+  "training_samples_per_epoch": 600,
+  "best_epoch": 5000,
+  "best_validation_mse": 5.125773963908943e-06,
+  "selection_split": "validation[600:800]",
+  "test_used_for_selection": false,
+  "checkpoint": "weight/best_model.pth",
+  "data_sha256": "5c4f361d8318d4a00c61d4d57f3d49eeeb6643096b7ec356a059f0f95d014c71",
+  "elapsed_seconds": 8878.240034841001,
+  "slurm_job_id": "1060058"
+}

results/verification.json

@@ -0,0 +1,41 @@
+{
+  "status": "passed",
+  "history_epochs": 5000,
+  "training_samples_per_epoch": 600,
+  "batches_per_epoch": 19,
+  "validation_epochs": [
+    1,
+    1001,
+    2001,
+    3001,
+    4001,
+    5000
+  ],
+  "best_epoch": 5000,
+  "best_validation_mse": 5.125773963908943e-06,
+  "test_indices": [
+    800,
+    999
+  ],
+  "test_samples": 200,
+  "test_shape": [
+    200,
+    11,
+    33
+  ],
+  "latent_shape": [
+    200,
+    11,
+    4,
+    33
+  ],
+  "recomputed_rmse_mean": 0.0017053971732615108,
+  "recomputed_rmse_std_ddof1": 0.0013320967174911497,
+  "global_rmse": 0.002161940973985019,
+  "checkpoint_sha256": "fd6917b63e596ad919bf7ad8504794039d303a3910925acac75f3fccd4ac08a8",
+  "figures_verified": 20,
+  "paper_rmse_mean": 0.00134,
+  "mean_rmse_gap_percent": 27.268445765784378,
+  "test_used_for_selection": false,
+  "source_data_sha256": "5c4f361d8318d4a00c61d4d57f3d49eeeb6643096b7ec356a059f0f95d014c71"
+}

results/visualization_manifest.json

@@ -0,0 +1,42 @@
+{
+  "figures": [
+    "figures/test_case_800_fields.png",
+    "figures/test_case_800_fields.pdf",
+    "figures/test_case_800_profiles.png",
+    "figures/test_case_800_profiles.pdf",
+    "figures/test_case_880_fields.png",
+    "figures/test_case_880_fields.pdf",
+    "figures/test_case_880_profiles.png",
+    "figures/test_case_880_profiles.pdf",
+    "figures/test_case_960_fields.png",
+    "figures/test_case_960_fields.pdf",
+    "figures/test_case_960_profiles.png",
+    "figures/test_case_960_profiles.pdf",
+    "figures/test_case_994_fields.png",
+    "figures/test_case_994_fields.pdf",
+    "figures/test_case_994_profiles.png",
+    "figures/test_case_994_profiles.pdf",
+    "figures/training_curve.png",
+    "figures/training_curve.pdf",
+    "figures/test_error_summary.png",
+    "figures/test_error_summary.pdf"
+  ],
+  "sample_ids": [
+    800,
+    880,
+    960,
+    994
+  ],
+  "dpi": 200,
+  "test_shape": [
+    200,
+    11,
+    33
+  ],
+  "training_epochs": 5000,
+  "colormaps": {
+    "field": "RdBu_r",
+    "absolute_error": "magma"
+  },
+  "note": "Validation selected weights. Test figures and metrics are post-training only."
+}

scripts/common.py

@@ -0,0 +1,193 @@
+"""Shared configuration, raw Burgers data and metric contracts."""
+import argparse
+import hashlib
+import json
+import os
+from pathlib import Path
+import random
+import socket
+import sys
+
+import numpy as np
+from scipy.io import loadmat
+import torch
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+from model.ronom import RONOM
+
+
+def parser(description):
+    p = argparse.ArgumentParser(description=description)
+    p.add_argument('--config', default='config/config.yaml')
+    return p
+
+
+def load_config(filename):
+    with open(filename, encoding='utf-8') as handle:
+        cfg = yaml.safe_load(handle)
+    return cfg
+
+
+def resolve(path):
+    path = Path(path).expanduser()
+    return path if path.is_absolute() else ROOT / path
+
+
+def save_json(path, value):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_suffix(path.suffix + '.tmp')
+    tmp.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + '\n')
+    tmp.replace(path)
+
+
+def update_run_manifest(cfg, **changes):
+    path = resolve(cfg['paths']['results']) / 'run_manifest.json'
+    if path.exists():
+        value = json.loads(path.read_text())
+        value.update(changes)
+        save_json(path, value)
+
+
+def seed_everything(seed, threads=8):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.set_num_threads(threads)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.benchmark = False
+
+
+def device_for(name):
+    if name in ('cuda', 'dcu'):
+        if not os.environ.get('SLURM_JOB_ID'):
+            raise RuntimeError('Accelerator execution must run inside a Slurm allocation.')
+        if not torch.cuda.is_available():
+            raise RuntimeError('The allocated compute node has no accessible accelerator.')
+        if torch.cuda.device_count() != 1:
+            raise RuntimeError('Expected one Slurm-visible DCU; check allocation device visibility. '
+                               f'CUDA_VISIBLE_DEVICES={os.environ.get("CUDA_VISIBLE_DEVICES")} '
+                               f'HIP_VISIBLE_DEVICES={os.environ.get("HIP_VISIBLE_DEVICES")} '
+                               f'ROCR_VISIBLE_DEVICES={os.environ.get("ROCR_VISIBLE_DEVICES")} '
+                               f'SLURM_JOB_GPUS={os.environ.get("SLURM_JOB_GPUS")}')
+        print(f'Slurm job={os.environ["SLURM_JOB_ID"]} host={socket.gethostname()} '
+              f'accelerator={torch.cuda.get_device_name()} count={torch.cuda.device_count()}', flush=True)
+        return torch.device('cuda')
+    return torch.device(name)
+
+
+def load_splits(cfg, names):
+    """Read only the supplied .mat. No downloads, augmentation, or normalization.
+
+    scipy decompresses the raw MATLAB array; only requested split tensors survive
+    this function. train.py requests train/validation and never constructs test.
+    """
+    d = cfg['data']
+    path = resolve(d['directory']) / d['filename']
+    raw = loadmat(path, variable_names=['input', 'output', 'tspan', 'visc'])
+    if 'visc' not in raw or not np.isclose(float(raw['visc'].squeeze()), 0.01):
+        raise ValueError('Raw Burgers viscosity must equal 0.01')
+    initial, target = raw['input'], raw['output']
+    times = raw['tspan'].reshape(-1)
+    if initial.shape != (1000, 1025) or target.shape != (1000, 101, 1025):
+        raise ValueError(f'Unexpected raw Burgers shapes: {initial.shape}, {target.shape}')
+    if times.shape != (101,) or not np.allclose(times, np.linspace(0, 1, 101), atol=1e-12):
+        raise ValueError('Raw tspan does not match the paper interval [0,1].')
+    counts = (d['train_samples'], d['validation_samples'], d['test_samples'])
+    if counts != (600, 200, 200):
+        raise ValueError('Formal Burgers protocol requires exactly 600/200/200 samples.')
+    if (d['space_stride'], d['time_stride']) != (32, 10):
+        raise ValueError('This experiment requires strides 32 in space and 10 in time.')
+    spans = {'train': (0, 600), 'validation': (600, 800), 'test': (800, 1000)}
+    splits = {}
+    for name in names:
+        begin, end = spans[name]
+        u = np.ascontiguousarray(initial[begin:end, ::32], dtype=np.float32)
+        y = np.ascontiguousarray(target[begin:end, ::10, ::32], dtype=np.float32)
+        if not np.isfinite(u).all() or not np.isfinite(y).all():
+            raise ValueError(f'Nonfinite data in {name}')
+        mismatch = float(np.max(np.abs(u - y[:, 0])))
+        if mismatch > 1e-5:
+            raise ValueError(f'input and output(t=0) disagree in {name}: {mismatch}')
+        splits[name] = (torch.from_numpy(u), torch.from_numpy(y), np.arange(begin, end))
+        print(f'Data split={name} indices=[{begin},{end}) input={u.shape} '
+              f'target={y.shape} initial_consistency_max={mismatch:.3e}', flush=True)
+    return splits, torch.tensor(times[::10], dtype=torch.float32), path
+
+
+def digest(path):
+    sha = hashlib.sha256()
+    with open(path, 'rb') as stream:
+        for block in iter(lambda: stream.read(8 * 1024 * 1024), b''):
+            sha.update(block)
+    return sha.hexdigest()
+
+
+def build_model(cfg, device):
+    model = RONOM(**cfg['model']).to(device)
+    if cfg['training'].get('jit_flow', False):
+        model.latent_flow = torch.jit.script(model.latent_flow)
+    count = sum(p.numel() for p in model.parameters())
+    if count != 414217:
+        raise ValueError(f'Model parameter count {count} differs from paper Appendix B: 414217')
+    print(f'RONOM parameters={count}, latent=4x33=132', flush=True)
+    return model
+
+
+def checkpoint_save(path, value):
+    path = resolve(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_suffix(path.suffix + '.tmp')
+    torch.save(value, tmp)
+    tmp.replace(path)
+
+
+def load_best(cfg, device):
+    path = resolve(cfg['paths']['checkpoint'])
+    checkpoint = torch.load(path, map_location='cpu', weights_only=False)
+    # Only architecture is compared: data/results paths never gate loading.
+    if checkpoint['model_config'] != cfg['model']:
+        raise ValueError('Model architecture in config differs from checkpoint.')
+    model = build_model(cfg, device)
+    model.load_state_dict(checkpoint['model_state_dict'], strict=True)
+    model.eval()
+    print(f'Loaded best checkpoint epoch={checkpoint["epoch"]} '
+          f'validation_mse={checkpoint["validation_mse"]:.9e}', flush=True)
+    return model, checkpoint, path
+
+
+def metrics(prediction, target):
+    err = prediction.double() - target.double()
+    mse_i = err.square().flatten(1).mean(1)
+    rmse_i = mse_i.sqrt()
+    relative = err.flatten(1).norm(dim=1) / target.double().flatten(1).norm(dim=1).clamp_min(1e-30)
+    return {
+        'samples': len(rmse_i),
+        'mse': mse_i.mean().item(),
+        'global_rmse': mse_i.mean().sqrt().item(),
+        'trajectory_rmse_mean': rmse_i.mean().item(),
+        'trajectory_rmse_std_ddof1': rmse_i.std(correction=1).item() if len(rmse_i) > 1 else 0.,
+        'relative_l2_mean': relative.mean().item(),
+        'rmse_by_time': err.square().mean(dim=(0, 2)).sqrt().cpu().tolist(),
+    }, rmse_i.cpu().numpy(), relative.cpu().numpy()
+
+
+@torch.no_grad()
+def evaluate(model, initial, target, times, batch_size, collect_latent=False):
+    predictions, latents = [], []
+    model.eval()
+    for begin in range(0, len(initial), batch_size):
+        output = model(initial[begin:begin + batch_size], times, return_latent=collect_latent)
+        if collect_latent:
+            output, zt = output
+            latents.append(zt.cpu())
+        predictions.append(output.cpu())
+    predicted = torch.cat(predictions)
+    if not torch.isfinite(predicted).all():
+        raise RuntimeError('Evaluation produced nonfinite predictions')
+    summary, per_sample, relative = metrics(predicted, target.cpu())
+    return summary, predicted, per_sample, relative, torch.cat(latents) if latents else None

scripts/inference.py

@@ -0,0 +1,53 @@
+"""Load validation-selected weights and evaluate the last 200 independent cases."""
+import csv
+import os
+import time
+
+import numpy as np
+from common import (parser, load_config, seed_everything, device_for, resolve,
+                    load_splits, load_best, evaluate, save_json, digest, update_run_manifest)
+
+
+def main():
+    args = parser(__doc__).parse_args()
+    cfg = load_config(args.config)
+    seed_everything(cfg['seed'], cfg['training']['num_threads'])
+    device = device_for(cfg['inference']['device'])
+    model, checkpoint, checkpoint_path = load_best(cfg, device)
+    splits, times, data_path = load_splits(cfg, ['test'])
+    if checkpoint['data_sha256'] != digest(data_path):
+        raise ValueError('Raw test data checksum differs from the training source')
+    u, y, ids = splits['test']
+    start = time.monotonic()
+    summary, pred, rmse_i, relative_i, latent = evaluate(
+        model, u.to(device), y.to(device), times, cfg['inference']['batch_size'], True)
+    summary.update(split='independent_test', indices=[int(ids[0]), int(ids[-1])],
+                   checkpoint_epoch=checkpoint['epoch'], checkpoint_sha256=digest(checkpoint_path),
+                   checkpoint_validation_mse=checkpoint['validation_mse'],
+                   selection_split=checkpoint['selection_split'],
+                   paper_rmse_mean=0.00134, paper_rmse_std=0.00162,
+                   elapsed_seconds=time.monotonic() - start,
+                   slurm_job_id=os.environ.get('SLURM_JOB_ID'))
+    result_dir = resolve(cfg['paths']['results'])
+    result_dir.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(result_dir / 'test_predictions.npz', prediction=pred.numpy(),
+                        reference=y.numpy(), initial=u.numpy(), latent=latent.numpy(),
+                        sample_ids=ids, x=np.linspace(0, 1, 33), times=times.numpy())
+    with (result_dir / 'test_per_sample.csv').open('w', newline='') as file:
+        writer = csv.writer(file)
+        writer.writerow(['sample_id', 'rmse', 'relative_l2'])
+        writer.writerows(zip(ids, rmse_i, relative_i))
+    save_json(result_dir / 'test_metrics.json', summary)
+    update_run_manifest(cfg, status='inference_completed', test_samples=len(ids),
+                        test_indices=[int(ids[0]), int(ids[-1])],
+                        test_metrics_file='test_metrics.json')
+    print(f'INDEPENDENT TEST n={len(ids)} RMSE mean={summary["trajectory_rmse_mean"]:.9e} '
+          f'std(ddof1)={summary["trajectory_rmse_std_ddof1"]:.9e} '
+          f'global_RMSE={summary["global_rmse"]:.9e} '
+          f'relative_L2={summary["relative_l2_mean"]:.9e}', flush=True)
+    print(f'Table 1(a) units 1e-3: {summary["trajectory_rmse_mean"] * 1000:.4f} '
+          f'({summary["trajectory_rmse_std_ddof1"] * 1000:.4f}); paper 1.34 (1.62)', flush=True)
+
+
+if __name__ == '__main__':
+    main()

scripts/result.py

@@ -0,0 +1,102 @@
+"""Visualize independent Burgers predictions, errors and the actual training log."""
+import csv
+import json
+
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import numpy as np
+from common import parser, load_config, resolve, save_json, update_run_manifest
+
+
+def main():
+    args = parser(__doc__).parse_args()
+    cfg = load_config(args.config)
+    folder = resolve(cfg['paths']['results'])
+    data = np.load(folder / 'test_predictions.npz')
+    pred, truth, times, x = (data[k] for k in ['prediction', 'reference', 'times', 'x'])
+    if pred.shape != (200, 11, 33) or not np.isfinite(pred).all():
+        raise ValueError('Invalid independent test predictions')
+    summary = json.loads((folder / 'test_metrics.json').read_text())
+    rmse = np.sqrt(np.mean((pred - truth)**2, axis=(1, 2)))
+    plot_dir = folder / 'figures'
+    plot_dir.mkdir(exist_ok=True)
+    dpi = cfg['visualization']['dpi']
+    files = []
+
+    def save(fig, name):
+        for suffix in ['png', 'pdf']:
+            destination = plot_dir / f'{name}.{suffix}'
+            fig.savefig(destination, dpi=dpi, bbox_inches='tight')
+            files.append(str(destination.relative_to(folder)))
+        plt.close(fig)
+
+    selected = list(dict.fromkeys(cfg['visualization']['sample_indices'] + [int(rmse.argmax())]))
+    for i in selected:
+        uid = int(data['sample_ids'][i])
+        fig, axes = plt.subplots(1, 3, figsize=(12, 3.5), constrained_layout=True)
+        vlim = max(float(np.max(np.abs(truth[i]))), float(np.max(np.abs(pred[i]))))
+        for ax, field, title in zip(axes[:2], [truth[i], pred[i]], ['Reference', 'RONOM']):
+            im = ax.pcolormesh(times, x, field.T, shading='nearest', cmap='RdBu_r',
+                               vmin=-vlim, vmax=vlim)
+            ax.set(title=title, xlabel='t', ylabel='x')
+        fig.colorbar(im, ax=axes[:2], label='u(x,t)')
+        im = axes[2].pcolormesh(times, x, np.abs(pred[i] - truth[i]).T,
+                               shading='nearest', cmap='magma')
+        axes[2].set(title='Absolute error', xlabel='t', ylabel='x')
+        fig.colorbar(im, ax=axes[2], label='|prediction - reference|')
+        fig.suptitle(f'Independent test case {uid}, trajectory RMSE = {rmse[i]:.3e}')
+        save(fig, f'test_case_{uid}_fields')
+        fig, axes = plt.subplots(1, 3, figsize=(12, 3.2), constrained_layout=True)
+        for ax, j in zip(axes, [0, 5, 10]):
+            ax.plot(x, truth[i, j], color='black', label='Reference')
+            ax.plot(x, pred[i, j], '--o', ms=3, color='tab:blue', label='RONOM')
+            ax.set(title=f't = {times[j]:.1f}', xlabel='x', ylabel='u(x,t)')
+            ax.grid(alpha=.2)
+        axes[0].legend()
+        fig.suptitle(f'Independent test case {uid}')
+        save(fig, f'test_case_{uid}_profiles')
+
+    epochs = np.array([], dtype=int)
+    if (folder / 'training_history.csv').exists():
+        with (folder / 'training_history.csv').open() as f:
+            history = list(csv.DictReader(f))
+        epochs = np.array([int(r['epoch']) for r in history])
+        if len(epochs) != 5000 or not np.array_equal(epochs, np.arange(1, 5001)):
+            raise ValueError('Expected one complete formal training history, epochs 1 through 5000')
+        val_rows = [r for r in history if r['validation_mse']]
+        fig, ax = plt.subplots(figsize=(7, 4))
+        ax.semilogy(epochs, [float(r['train_mse']) for r in history], label='Training MSE')
+        ax.semilogy([int(r['epoch']) for r in val_rows],
+                    [float(r['validation_mse']) for r in val_rows], 'o-', label='Validation MSE')
+        ax.axvline(summary['checkpoint_epoch'], color='grey', ls=':', label='Selected epoch')
+        ax.set(xlabel='Epoch', ylabel='MSE', title='Formal 5000-epoch training')
+        ax.legend(); ax.grid(alpha=.2)
+        save(fig, 'training_curve')
+    else:
+        print('Training history is absent in this results directory; generating independent test figures.', flush=True)
+    fig, axes = plt.subplots(1, 2, figsize=(10, 3.6), constrained_layout=True)
+    axes[0].plot(times, summary['rmse_by_time'], 'o-')
+    axes[0].set(xlabel='t', ylabel='RMSE across independent test cases and x', title='Error by time')
+    axes[1].hist(rmse * 1000, bins=25, edgecolor='white')
+    axes[1].set(xlabel='Trajectory RMSE (× 10⁻³)', ylabel='Test cases', title='200 independent trajectories')
+    save(fig, 'test_error_summary')
+    comparison = folder / 'table1a_comparison.csv'
+    with comparison.open('w', newline='') as f:
+        writer = csv.writer(f)
+        writer.writerow(['source', 'split', 'samples', 'rmse_mean', 'rmse_std_ddof1', 'units'])
+        writer.writerow(['paper_Table1a_RONOM', 'paper_test', 200, .00134, .00162, 'raw'])
+        writer.writerow(['this_run', 'independent_test_800_999', 200,
+                         summary['trajectory_rmse_mean'], summary['trajectory_rmse_std_ddof1'], 'raw'])
+    save_json(folder / 'visualization_manifest.json', {
+        'figures': files, 'sample_ids': [int(data['sample_ids'][i]) for i in selected],
+        'dpi': dpi, 'test_shape': list(pred.shape), 'training_epochs': len(epochs) if len(epochs) else None,
+        'colormaps': {'field': 'RdBu_r', 'absolute_error': 'magma'},
+        'note': 'Validation selected weights. Test figures and metrics are post-training only.'})
+    update_run_manifest(cfg, status='completed', figure_files=len(files),
+                        visualization_manifest='visualization_manifest.json')
+    print(f'VISUALIZATION COMPLETE: {len(files)} figures in {plot_dir}', flush=True)
+
+
+if __name__ == '__main__':
+    main()

scripts/train.py

@@ -0,0 +1,152 @@
+"""Full 5000-epoch Burgers input-generalization experiment; validation selects best."""
+import csv
+import os
+from pathlib import Path
+import time
+
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+from common import (parser, load_config, seed_everything, device_for, load_splits,
+                    resolve, save_json, digest, build_model, checkpoint_save, evaluate,
+                    update_run_manifest)
+
+
+def main():
+    p = parser(__doc__)
+    p.add_argument('--resume', default=None, help='Resume a locally generated last checkpoint')
+    args = p.parse_args()
+    cfg = load_config(args.config)
+    tr = cfg['training']
+    if tr['epochs'] != 5000 or tr['batch_size'] != 32:
+        raise ValueError('Formal experiment requires 5000 epochs and batch size 32')
+    seed_everything(cfg['seed'], tr['num_threads'])
+    device = device_for(tr['device'])
+    result_dir = resolve(cfg['paths']['results'])
+    result_dir.mkdir(parents=True, exist_ok=True)
+    history_path = result_dir / 'training_history.csv'
+    if history_path.exists() and not args.resume:
+        raise FileExistsError('Training history already exists. Use --resume to continue this run.')
+    splits, times, data_path = load_splits(cfg, ['train', 'validation'])
+    train_u, train_y, train_ids = splits['train']
+    val_u, val_y, val_ids = splits['validation']
+    train_u, train_y = train_u.to(device), train_y.to(device)
+    val_u, val_y = val_u.to(device), val_y.to(device)
+    loader = DataLoader(TensorDataset(train_u, train_y), batch_size=32,
+                        shuffle=True, drop_last=False, num_workers=0)
+    model = build_model(cfg, device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=tr['learning_rate'],
+                                 weight_decay=tr['weight_decay'])
+    scheduler = torch.optim.lr_scheduler.MultiplicativeLR(
+        optimizer, lr_lambda=lambda _: tr['scheduler_gamma'])
+    start_epoch, best_mse, best_epoch = 0, float('inf'), None
+    if args.resume:
+        old = torch.load(resolve(args.resume), map_location='cpu', weights_only=False)
+        if old['model_config'] != cfg['model'] or old['data_sha256'] != digest(data_path):
+            raise ValueError('Resume checkpoint architecture or source data differs')
+        model.load_state_dict(old['model_state_dict'])
+        optimizer.load_state_dict(old['optimizer_state_dict'])
+        scheduler.load_state_dict(old['scheduler_state_dict'])
+        torch.set_rng_state(old['torch_rng'])
+        torch.cuda.set_rng_state_all(old['cuda_rng'])
+        start_epoch, best_mse, best_epoch = old['epoch'], old['best_validation_mse'], old['best_epoch']
+        # Logs may contain epochs beyond the last atomic checkpoint after an
+        # interrupted allocation. Replay those epochs from saved RNG/optimizer
+        # state, keeping exactly one history row for every completed epoch.
+        with history_path.open(newline='') as previous:
+            reader = csv.DictReader(previous)
+            fields = reader.fieldnames
+            retained = [r for r in reader if int(r['epoch']) <= start_epoch]
+        if [int(r['epoch']) for r in retained] != list(range(1, start_epoch + 1)):
+            raise ValueError('Resume history is incomplete before the checkpoint epoch')
+        temporary_history = history_path.with_suffix('.csv.tmp')
+        with temporary_history.open('w', newline='') as previous:
+            writer = csv.DictWriter(previous, fieldnames=fields)
+            writer.writeheader()
+            writer.writerows(retained)
+        temporary_history.replace(history_path)
+        print(f'Resuming after completed epoch {start_epoch}', flush=True)
+    data_sha = digest(data_path)
+    save_json(result_dir / 'run_manifest.json', {
+        'status': 'running', 'config': cfg, 'data_sha256': data_sha,
+        'slurm_job_id': os.environ.get('SLURM_JOB_ID'), 'torch_version': torch.__version__,
+        'train_indices': train_ids.tolist(), 'validation_indices': val_ids.tolist(),
+        'test_indices_reserved': [800, 1000], 'times': times.tolist(),
+        'model_parameters': sum(v.numel() for v in model.parameters()),
+        'test_used_for_selection': False, 'precision': 'float32'})
+    header = ['epoch', 'train_half_mse', 'train_mse', 'last_batch_half_mse', 'learning_rate',
+              'validation_mse', 'validation_rmse_mean', 'validation_rmse_std', 'elapsed_seconds']
+    started = time.monotonic()
+    with history_path.open('a' if args.resume else 'w', newline='', buffering=1) as handle:
+        writer = csv.DictWriter(handle, fieldnames=header)
+        if not args.resume:
+            writer.writeheader()
+        for epoch_index in range(start_epoch, tr['epochs']):
+            epoch = epoch_index + 1
+            model.train()
+            total, seen, last_loss = 0., 0, None
+            lr_used = optimizer.param_groups[0]['lr']
+            for u, target in loader:
+                optimizer.zero_grad(set_to_none=True)
+                predicted = model(u, times)
+                loss = 0.5 * (predicted - target).square().mean()
+                if not torch.isfinite(loss):
+                    raise RuntimeError(f'Nonfinite training loss at epoch {epoch}')
+                loss.backward()
+                optimizer.step()
+                last_loss = float(loss.detach())
+                total += last_loss * len(u)
+                seen += len(u)
+            if seen != 600:
+                raise RuntimeError(f'Only {seen} training samples processed')
+            if epoch_index > 0 and epoch_index % tr['scheduler_frequency'] == 0:
+                scheduler.step()
+            row = {'epoch': epoch, 'train_half_mse': total / seen,
+                   'train_mse': 2 * total / seen, 'last_batch_half_mse': last_loss,
+                   'learning_rate': lr_used, 'elapsed_seconds': time.monotonic() - started}
+            val_summary = None
+            if epoch_index % tr['validation_frequency'] == 0 or epoch == tr['epochs']:
+                val_summary, *_ = evaluate(model, val_u, val_y, times, tr['batch_size'])
+                row.update(validation_mse=val_summary['mse'],
+                           validation_rmse_mean=val_summary['trajectory_rmse_mean'],
+                           validation_rmse_std=val_summary['trajectory_rmse_std_ddof1'])
+                if val_summary['mse'] < best_mse:
+                    best_mse, best_epoch = val_summary['mse'], epoch
+                    checkpoint_save(cfg['paths']['checkpoint'], {
+                        'format_version': 1, 'model_config': cfg['model'],
+                        'model_state_dict': model.state_dict(), 'epoch': epoch,
+                        'validation_mse': best_mse, 'validation_metrics': val_summary,
+                        'selection_split': 'validation[600:800]',
+                        'data_sha256': data_sha, 'seed': cfg['seed']})
+                    save_json(result_dir / 'best_validation_metrics.json',
+                              {'epoch': epoch, 'split': 'validation', **val_summary})
+                print(f'VALIDATION epoch={epoch} mse={val_summary["mse"]:.9e} '
+                      f'RMSE_mean={val_summary["trajectory_rmse_mean"]:.9e} '
+                      f'RMSE_std={val_summary["trajectory_rmse_std_ddof1"]:.9e} '
+                      f'best_epoch={best_epoch}', flush=True)
+            writer.writerow(row)
+            print(f'epoch={epoch:04d}/5000 train_loss(half_MSE)={total / seen:.9e} '
+                  f'train_MSE={2 * total / seen:.9e} lr={lr_used:.6e} '
+                  f'samples={seen} elapsed={row["elapsed_seconds"]:.1f}s', flush=True)
+            if epoch % tr['checkpoint_frequency'] == 0 or epoch == tr['epochs']:
+                checkpoint_save(cfg['paths']['last_checkpoint'], {
+                    'model_config': cfg['model'], 'model_state_dict': model.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'scheduler_state_dict': scheduler.state_dict(), 'epoch': epoch,
+                    'best_validation_mse': best_mse, 'best_epoch': best_epoch,
+                    'torch_rng': torch.get_rng_state(), 'cuda_rng': torch.cuda.get_rng_state_all(),
+                    'data_sha256': data_sha})
+    summary = {'status': 'completed', 'completed_epochs': tr['epochs'],
+               'batches_per_epoch': len(loader), 'training_samples_per_epoch': 600,
+               'best_epoch': best_epoch, 'best_validation_mse': best_mse,
+               'selection_split': 'validation[600:800]', 'test_used_for_selection': False,
+               'checkpoint': str(cfg['paths']['checkpoint']), 'data_sha256': data_sha,
+               'elapsed_seconds': time.monotonic() - started,
+               'slurm_job_id': os.environ.get('SLURM_JOB_ID')}
+    save_json(result_dir / 'training_summary.json', summary)
+    update_run_manifest(cfg, status='training_completed', completed_epochs=tr['epochs'],
+                        best_epoch=best_epoch, best_validation_mse=best_mse)
+    print('FORMAL TRAINING COMPLETE', summary, flush=True)
+
+
+if __name__ == '__main__':
+    main()

weight/best_model.pth

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