Upload folder using huggingface_hub

README.md

@@ -0,0 +1,50 @@
+---
+license: mit
+library_name: pytorch
+tags: [neural-operator, pde-solver, transolver, linear-attention, elasticity]
+metrics: [relative-l2]
+---
+
+# LinearNO — Elastic Stress Surrogate (Geo-FNO Elasticity)
+
+Predicts the per-node **von Mises stress** of a hyper-elastic unit cell with a central void on
+the Geo-FNO **Elasticity** benchmark. Attention block: **LinearNO** (asymmetric linear attention,
+`shared_qk`, `project_out=False`), a drop-in
+replacement for Transolver's Physics-Attention.
+
+- **Test relative L2:** 0.005968 (this checkpoint).
+- **Params:** 582275 (≤ the Transolver baseline, 713,665).
+- **Interactive demo:** https://huggingface.co/spaces/Efradeca/elastic-stress-surrogate
+
+## Results (Geo-FNO Elasticity test relative L2, mean ± std over 3 seeds)
+
+| Model | test rel-L2 | Params | Notes |
+|---|---|---|---|
+| Transolver baseline (reproduced) | 0.00678 ± 0.0012 | 713,665 | reproduces the published 0.0064 |
+| LinearNO | 0.00664 ± 0.0008 | 713,089 | ≤ baseline params; does **not** reach the paper's 0.0050 |
+| **LinearNO + equilibrium reg (this model)** | 0.00668 ± 0.0006 | 582,275 | accuracy preserved + **~360× lower ∇·σ residual** |
+
+The equilibrium regularizer reduces the discrete equilibrium residual ‖∇·σ‖² by ~360× (6.1e6 → 1.7e4)
+with no accuracy cost, and this physical consistency is **maintained out-of-distribution** (OOD residual
+1.7e4 vs in-dist 1.6e4). It does **not** improve OOD *accuracy* (degradation comparable to LinearNO) —
+reported honestly, not massaged.
+
+## Usage
+```python
+import torch, json
+from safetensors.torch import load_file
+from model import load_checkpoint  # bundled loader (or use stress_operator.models.build_model)
+```
+
+## Honesty caveats (read these)
+1. **Scope:** a *narrow* result on one benchmark — **not** a foundation model. The material is
+   hyper-elastic (rubber-like), **not steel**; the method transfers, a steel result needs steel data.
+2. **LinearNO** has no official public repo; it was reimplemented from the paper's equations
+   (arXiv:2511.06294). Only the attention block differs from Transolver (ICML 2024, arXiv:2402.02366, MIT).
+3. **Reproduction variance:** small dataset (1000 train) → run-to-run variance; numbers are
+   eager (deterministic), mean ± std over seeds {0,1,2}.
+4. The interactive demo uses a **synthetic** mesh (no FEM ground truth for arbitrary void shapes).
+
+## Citation
+Transolver (Wu et al., ICML 2024, arXiv:2402.02366) · LinearNO (Hu et al., AAAI 2026,
+arXiv:2511.06294) · Geo-FNO (Li et al., 2022, arXiv:2207.05209).

config.json

@@ -0,0 +1,38 @@
+{
+  "model": {
+    "attention": "linearno",
+    "linearno_variant": "shared_qk",
+    "linearno_project_out": false,
+    "space_dim": 2,
+    "fun_dim": 0,
+    "out_dim": 3,
+    "n_layers": 8,
+    "n_hidden": 128,
+    "n_heads": 8,
+    "dim_head": 16,
+    "slice_num": 64,
+    "mlp_ratio": 1,
+    "dropout": 0.0,
+    "unified_pos": false,
+    "ref": 8
+  },
+  "normalizer": {
+    "mean": 187.73873901367188,
+    "std": 127.08287811279297
+  },
+  "scale_S": 187.73873901367188,
+  "metrics": {
+    "test_rel_l2": 0.005968,
+    "best_test_rel_l2": 0.005968,
+    "test_residual": 16361.362044798194,
+    "test_residual_scaled": 0.46420697510242465,
+    "scale_S": 187.73873901367188,
+    "train_rel_l2": 0.004208,
+    "n_params": 582275,
+    "epochs": 500,
+    "approach": "A",
+    "lambda": 0.01
+  },
+  "task": "geo-fno-elasticity-stress",
+  "license": "mit"
+}

model.py

@@ -0,0 +1,80 @@
+"""CPU inference helpers for the demo (loads a trained checkpoint; handles 1- and 3-channel heads).
+
+Kept dependency-light so it also works inside a Hugging Face Space. The model architecture is
+imported from the installed ``stress_operator`` package; the HF export bundles a self-contained
+copy (see scripts/05_export_to_hf.sh).
+"""
+from __future__ import annotations
+
+import os
+import sys
+
+import numpy as np
+import torch
+
+# Make `stress_operator` importable in BOTH layouts: the local repo (package under ../src) and a
+# bundled Hugging Face Space (package is a sibling of this file).
+_HERE = os.path.dirname(os.path.abspath(__file__))
+for _p in (_HERE, os.path.join(_HERE, "..", "src")):
+    if _p not in sys.path:
+        sys.path.insert(0, _p)
+
+from stress_operator.models.transolver import build_model  # noqa: E402
+
+
+def load_checkpoint(ckpt_path: str, device: str = "cpu"):
+    """Load a trained model. Supports two formats:
+
+    - ``*.safetensors`` (HF deployment): weights from safetensors + ``config.json`` (sibling)
+      holding the model config, normalizer mean/std, and scale_S.
+    - ``*.pt`` (local training checkpoint): a dict with state_dict / normalizer / config.
+    """
+    if ckpt_path.endswith(".safetensors"):
+        import json
+
+        from safetensors.torch import load_file
+
+        cfg_path = os.path.join(os.path.dirname(ckpt_path) or ".", "config.json")
+        with open(cfg_path) as f:
+            cfg = json.load(f)
+        model_cfg = cfg["model"]
+        state_dict = load_file(ckpt_path, device=device)
+        mean = torch.tensor(float(cfg["normalizer"]["mean"]), device=device).reshape(1, 1, 1)
+        std = torch.tensor(float(cfg["normalizer"]["std"]), device=device).reshape(1, 1, 1)
+        scale_S = cfg.get("scale_S", None)
+        metrics = cfg.get("metrics", {})
+    else:
+        ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+        model_cfg = ckpt["config"]["model"]
+        state_dict = ckpt["state_dict"]
+        mean = ckpt["normalizer"]["mean"].to(device)
+        std = ckpt["normalizer"]["std"].to(device)
+        scale_S = ckpt.get("scale_S", None)
+        metrics = ckpt.get("metrics", {})
+
+    model = build_model(model_cfg).to(device)
+    model.load_state_dict(state_dict)
+    model.eval()
+    info = {
+        "out_dim": model_cfg.get("out_dim", 1),
+        "scale_S": scale_S,
+        "attention": model_cfg.get("attention", "physics"),
+        "metrics": metrics,
+    }
+    return model, (mean, std), info
+
+
+@torch.no_grad()
+def predict_stress(model, coords: np.ndarray, norm, info, device: str = "cpu") -> np.ndarray:
+    """coords: (N, 2) -> per-node von Mises stress (N,) in physical units."""
+    mean, std = norm
+    x = torch.as_tensor(coords, dtype=torch.float32, device=device).unsqueeze(0)  # (1,N,2)
+    out = model(x, None)[0]  # (N, out_dim)
+    if info["out_dim"] == 3:
+        from stress_operator.losses.equilibrium import von_mises
+
+        S = info["scale_S"] or 1.0
+        stress = von_mises(out) * S          # 3-channel tensor -> von Mises, undo target scale
+    else:
+        stress = out[:, 0] * std + mean      # de-normalize scalar prediction
+    return stress.detach().cpu().numpy()

model.safetensors

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

stress_operator/__init__.py

@@ -0,0 +1,7 @@
+"""Equilibrium-regularized linear neural operator for elastic stress fields.
+
+Pipeline: reproduce Transolver baseline -> replace attention with LinearNO ->
+add equilibrium-residual regularizer + OOD evaluation -> ship to Hugging Face.
+"""
+
+__version__ = "0.1.0"

stress_operator/data/__init__.py

@@ -0,0 +1,15 @@
+from .dataset import (
+    ElasticityDataset,
+    UnitTransformer,
+    load_raw_arrays,
+    build_splits,
+    build_splits_from_indices,
+)
+
+__all__ = [
+    "ElasticityDataset",
+    "UnitTransformer",
+    "load_raw_arrays",
+    "build_splits",
+    "build_splits_from_indices",
+]

stress_operator/data/dataset.py

@@ -0,0 +1,180 @@
+"""Geo-FNO Elasticity dataset.
+
+Reconciled against the authoritative reference files (see ``docs/RECONCILIATION.md``):
+  - Geo-FNO ``elasticity/elas_geofno.py``            (raw .npy axis layout, filenames)
+  - Transolver ``exp_elas.py`` + ``utils/normalizer.py``  (split, output normalizer, loss scale)
+
+Raw array layout (the **sample axis is LAST** in every raw ``.npy``):
+
+  Random_UnitCell_sigma_10.npy : (972, 2000)      -> per-node von Mises stress (target)
+  Random_UnitCell_XY_10.npy    : (972, 2, 2000)   -> node coordinates (input)
+  Random_UnitCell_rr_10.npy    : (42, 2000)       -> void-boundary radii (geometry params; OOD)
+
+Split (verbatim from ``exp_elas.py``): train = first ``ntrain`` samples, test = LAST ``ntest``,
+out of 2000 total. The middle samples are unused.
+"""
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+SIGMA_FILE = "Random_UnitCell_sigma_10.npy"
+XY_FILE = "Random_UnitCell_XY_10.npy"
+RR_FILE = "Random_UnitCell_rr_10.npy"
+
+N_NODES = 972
+N_TOTAL = 2000
+
+
+class UnitTransformer:
+    """Global scalar z-score normalizer for the stress target.
+
+    Verbatim behavior of Transolver ``utils/normalizer.py::UnitTransformer``:
+    mean/std are reduced over dims ``(0, 1)`` (samples and nodes), giving a single
+    scalar mean and std. ``decode`` is applied to predictions before the metric.
+    """
+
+    def __init__(self, x: torch.Tensor):
+        self.mean = x.mean(dim=(0, 1), keepdim=True)
+        self.std = x.std(dim=(0, 1), keepdim=True) + 1e-8
+
+    def to(self, device) -> "UnitTransformer":
+        self.mean = self.mean.to(device)
+        self.std = self.std.to(device)
+        return self
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        return (x - self.mean) / self.std
+
+    def decode(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.std + self.mean
+
+
+def _find_file(data_dir: str, name: str) -> str:
+    """Locate ``name`` under ``data_dir`` (gdown may nest files in a subfolder)."""
+    direct = os.path.join(data_dir, name)
+    if os.path.isfile(direct):
+        return direct
+    for root, _dirs, files in os.walk(data_dir):
+        if name in files:
+            return os.path.join(root, name)
+    raise FileNotFoundError(
+        f"Could not find {name} under {data_dir!r}. "
+        f"Run `make data` (or python -m stress_operator.data.download --out {data_dir})."
+    )
+
+
+def load_raw_arrays(data_dir: str):
+    """Load and reorient the raw arrays so the sample axis is first.
+
+    Returns float32 tensors:
+        coords : (N_TOTAL, 972, 2)
+        sigma  : (N_TOTAL, 972)        physical (de-normalized) stress
+        rr     : (N_TOTAL, 42)         geometry params (may be absent -> None)
+    """
+    sigma_np = np.load(_find_file(data_dir, SIGMA_FILE))      # (972, 2000)
+    xy_np = np.load(_find_file(data_dir, XY_FILE))            # (972, 2, 2000)
+
+    sigma = torch.tensor(sigma_np, dtype=torch.float).permute(1, 0).contiguous()        # (2000, 972)
+    coords = torch.tensor(xy_np, dtype=torch.float).permute(2, 0, 1).contiguous()       # (2000, 972, 2)
+
+    rr: Optional[torch.Tensor] = None
+    try:
+        rr_np = np.load(_find_file(data_dir, RR_FILE))        # (42, 2000)
+        rr = torch.tensor(rr_np, dtype=torch.float).permute(1, 0).contiguous()          # (2000, 42)
+    except FileNotFoundError:
+        rr = None
+
+    return coords, sigma, rr
+
+
+@dataclass
+class Splits:
+    train_coords: torch.Tensor
+    train_sigma: torch.Tensor          # physical
+    test_coords: torch.Tensor
+    test_sigma: torch.Tensor           # physical
+    normalizer: UnitTransformer
+    train_rr: Optional[torch.Tensor] = None
+    test_rr: Optional[torch.Tensor] = None
+
+
+def build_splits(data_dir: str, ntrain: int = 1000, ntest: int = 200) -> Splits:
+    """Build the train/test split exactly as ``exp_elas.py`` does.
+
+    train = first ``ntrain``; test = LAST ``ntest``. The normalizer is fit on the
+    (physical) train stress only.
+    """
+    coords, sigma, rr = load_raw_arrays(data_dir)
+
+    train_coords = coords[:ntrain]
+    train_sigma = sigma[:ntrain]
+    test_coords = coords[-ntest:]
+    test_sigma = sigma[-ntest:]
+
+    normalizer = UnitTransformer(train_sigma)
+
+    train_rr = rr[:ntrain] if rr is not None else None
+    test_rr = rr[-ntest:] if rr is not None else None
+
+    return Splits(
+        train_coords=train_coords,
+        train_sigma=train_sigma,
+        test_coords=test_coords,
+        test_sigma=test_sigma,
+        normalizer=normalizer,
+        train_rr=train_rr,
+        test_rr=test_rr,
+    )
+
+
+def build_splits_from_indices(data_dir: str, train_idx, test_idx) -> Splits:
+    """Build a split from explicit sample indices over all 2000 samples (used by the OOD eval).
+
+    The normalizer is fit on the (physical) train stress only.
+    """
+    coords, sigma, rr = load_raw_arrays(data_dir)
+    train_idx = torch.as_tensor(np.asarray(train_idx), dtype=torch.long)
+    test_idx = torch.as_tensor(np.asarray(test_idx), dtype=torch.long)
+
+    train_sigma = sigma[train_idx]
+    normalizer = UnitTransformer(train_sigma)
+    return Splits(
+        train_coords=coords[train_idx],
+        train_sigma=train_sigma,
+        test_coords=coords[test_idx],
+        test_sigma=sigma[test_idx],
+        normalizer=normalizer,
+        train_rr=(rr[train_idx] if rr is not None else None),
+        test_rr=(rr[test_idx] if rr is not None else None),
+    )
+
+
+class ElasticityDataset(Dataset):
+    """Yields ``(coords, sigma)`` per sample.
+
+    ``coords`` : (972, 2) node coordinates (model input)
+    ``sigma``  : (972, 1) **physical** stress target
+
+    Normalization is intentionally *not* applied here: the training loop predicts in
+    normalized space and decodes before the relative-L2 loss (identical to ``exp_elas.py``;
+    see ``docs/RECONCILIATION.md`` §4). Keeping physical targets in the dataset makes the
+    de-normalize-before-metric contract explicit and impossible to forget.
+    """
+
+    def __init__(self, coords: torch.Tensor, sigma: torch.Tensor):
+        assert coords.shape[0] == sigma.shape[0]
+        self.coords = coords
+        # ensure (n, 972, 1)
+        self.sigma = sigma if sigma.dim() == 3 else sigma.unsqueeze(-1)
+
+    def __len__(self) -> int:
+        return self.coords.shape[0]
+
+    def __getitem__(self, idx: int):
+        return self.coords[idx], self.sigma[idx]

stress_operator/data/download.py

@@ -0,0 +1,118 @@
+"""Download + verify the Geo-FNO Elasticity dataset.
+
+Source: the Geo-FNO Google Drive collection (license MIT, ``neuraloperator/Geo-FNO``).
+
+IMPORTANT: the canonical Drive *folder*
+``https://drive.google.com/drive/folders/1YBuaoTdOSr_qzaow-G-iwvbUI7fiUzu8`` is the **entire**
+neuraloperator dataset collection (~8 GB: airfoil, car-cfd, channel-shocks, elasticity, ...).
+For this project we only need three files (~47 MB total), so we fetch them **by file id**
+rather than downloading the whole folder. The ids below were resolved once from that folder
+(``gdown.download_folder(..., skip_download=True)``) — see ``docs/RECONCILIATION.md``.
+
+Usage:
+    python -m stress_operator.data.download --out data
+    python -m stress_operator.data.download --out data --report-only   # just verify + print shapes
+"""
+from __future__ import annotations
+
+import argparse
+import hashlib
+import os
+import sys
+
+import numpy as np
+
+# Resolved file ids for elasticity/Meshes/* inside the Geo-FNO Drive collection.
+DRIVE_FILE_IDS = {
+    "Random_UnitCell_sigma_10.npy": "1Ia5izgUum-IQLdO6PW70HO8AdAqA_IVb",
+    "Random_UnitCell_XY_10.npy": "1I-fO-RsFvD3nqBuFrg67R0yqTFdD_gpA",
+    "Random_UnitCell_rr_10.npy": "1Pjliqhxegoe5VpoLrpBa9n3P4gX9MfTt",
+}
+# Mirror Transolver's expected layout: <out>/elasticity/Meshes/<file>.
+SUBDIR = os.path.join("elasticity", "Meshes")
+
+REQUIRED_FILES = ["Random_UnitCell_sigma_10.npy", "Random_UnitCell_XY_10.npy"]
+
+# Authoritative raw shapes (sample axis LAST). See docs/RECONCILIATION.md.
+EXPECTED_RAW_SHAPES = {
+    "Random_UnitCell_sigma_10.npy": (972, 2000),
+    "Random_UnitCell_XY_10.npy": (972, 2, 2000),
+    "Random_UnitCell_rr_10.npy": (42, 2000),
+}
+
+
+def _find(data_dir: str, name: str):
+    """Locate ``name`` anywhere under ``data_dir`` (handles nested layouts)."""
+    direct = os.path.join(data_dir, name)
+    if os.path.isfile(direct):
+        return direct
+    for root, _dirs, files in os.walk(data_dir):
+        if name in files:
+            return os.path.join(root, name)
+    return None
+
+
+def have_all(data_dir: str) -> bool:
+    return all(_find(data_dir, f) is not None for f in REQUIRED_FILES)
+
+
+def download(out_dir: str, force: bool = False) -> None:
+    import gdown
+
+    dest_dir = os.path.join(out_dir, SUBDIR)
+    os.makedirs(dest_dir, exist_ok=True)
+    for name, file_id in DRIVE_FILE_IDS.items():
+        existing = _find(out_dir, name)
+        if existing is not None and not force:
+            print(f"[download] {name}: already present ({existing}); skipping.")
+            continue
+        out_path = os.path.join(dest_dir, name)
+        print(f"[download] fetching {name} (id={file_id}) -> {out_path}")
+        gdown.download(id=file_id, output=out_path, quiet=False)
+
+
+def _md5(path: str, chunk: int = 1 << 20) -> str:
+    h = hashlib.md5()
+    with open(path, "rb") as f:
+        for block in iter(lambda: f.read(chunk), b""):
+            h.update(block)
+    return h.hexdigest()
+
+
+def report(data_dir: str) -> int:
+    """Print shape/dtype/stats and check against expected raw shapes. Returns exit code."""
+    ok = True
+    print(f"[report] scanning {data_dir!r}")
+    for name in DRIVE_FILE_IDS:
+        path = _find(data_dir, name)
+        required = name in REQUIRED_FILES
+        if path is None:
+            print(f"  - {name}: {'MISSING (required)' if required else 'missing (optional)'}")
+            if required:
+                ok = False
+            continue
+        arr = np.load(path)
+        exp = EXPECTED_RAW_SHAPES.get(name)
+        flag = "" if (exp is None or tuple(arr.shape) == tuple(exp)) else f"  <-- WARNING: expected {exp}"
+        print(
+            f"  - {name}: shape={tuple(arr.shape)} dtype={arr.dtype} "
+            f"min={float(arr.min()):.4g} max={float(arr.max()):.4g} md5={_md5(path)[:8]}{flag}"
+        )
+    print("[report] OK" if ok else "[report] FAILED: required files missing")
+    return 0 if ok else 1
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description="Download/verify the Geo-FNO Elasticity dataset (3 files).")
+    ap.add_argument("--out", default="data", help="output directory")
+    ap.add_argument("--report-only", action="store_true", help="skip download, just verify + print shapes")
+    ap.add_argument("--force", action="store_true", help="re-download even if files exist")
+    args = ap.parse_args()
+
+    if not args.report_only:
+        download(args.out, force=args.force)
+    return report(args.out)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

stress_operator/data/ood_split.py

@@ -0,0 +1,87 @@
+"""Out-of-distribution (covariate-shift) split from void-geometry statistics (Gate D).
+
+No new FEM is needed: we re-use the existing 2000 samples but split them by a per-sample
+geometric statistic of the void boundary radii ``rr`` (shape ``(n, 42)``). The model trains on
+the lower percentile (typical voids) and is tested on the upper percentile (large/extreme voids
+it never saw) — a genuine covariate shift in geometry.
+
+Split rule (documented exactly, master plan §2.6):
+  stat_i = STAT(rr_i)                          # default STAT = max radius ("void size")
+  threshold = quantile(stat, train_frac)        # default train_frac = 0.80
+  in-distribution (train+id-test): stat_i <= threshold   (lower ~80%)
+  out-of-distribution (ood-test) : stat_i >  threshold   (upper ~20%)
+The in-distribution pool is further split into train / id-test by ``id_test_frac`` (held-out,
+seeded) so in-dist and OOD errors are both measured on data unseen during training.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict
+
+import numpy as np
+import torch
+
+STATS = {
+    "max": lambda rr: rr.max(axis=1),       # largest radius -> void size (primary)
+    "mean": lambda rr: rr.mean(axis=1),
+    "std": lambda rr: rr.std(axis=1),        # radius spread -> "lobiness"
+    "range": lambda rr: rr.max(axis=1) - rr.min(axis=1),
+}
+
+
+def geometry_stat(rr: torch.Tensor, kind: str = "max") -> np.ndarray:
+    if kind not in STATS:
+        raise ValueError(f"unknown stat {kind!r}; choose from {list(STATS)}")
+    return STATS[kind](rr.detach().cpu().numpy())
+
+
+@dataclass
+class OODSplit:
+    train_idx: np.ndarray       # in-distribution, used for training
+    id_test_idx: np.ndarray     # in-distribution, held out for testing
+    ood_test_idx: np.ndarray    # out-of-distribution (extreme geometry), held out
+    threshold: float
+    stat_kind: str
+    info: Dict
+
+
+def make_ood_split(
+    rr: torch.Tensor,
+    stat_kind: str = "max",
+    train_frac: float = 0.80,
+    id_test_frac: float = 0.20,
+    seed: int = 0,
+) -> OODSplit:
+    """Build a covariate-shift split from ``rr`` statistics. Deterministic given ``seed``."""
+    stat = geometry_stat(rr, stat_kind)            # (n,)
+    n = stat.shape[0]
+    threshold = float(np.quantile(stat, train_frac))
+
+    in_dist = np.where(stat <= threshold)[0]
+    ood = np.where(stat > threshold)[0]
+
+    rng = np.random.default_rng(seed)
+    perm = rng.permutation(in_dist)
+    n_id_test = int(round(len(in_dist) * id_test_frac))
+    id_test_idx = np.sort(perm[:n_id_test])
+    train_idx = np.sort(perm[n_id_test:])
+
+    info = {
+        "n_total": int(n),
+        "n_train": int(train_idx.size),
+        "n_id_test": int(id_test_idx.size),
+        "n_ood_test": int(ood.size),
+        "threshold": threshold,
+        "stat_kind": stat_kind,
+        "stat_train_max": float(stat[train_idx].max()),
+        "stat_ood_min": float(stat[ood].min()),
+        "stat_ood_max": float(stat[ood].max()),
+    }
+    return OODSplit(
+        train_idx=train_idx,
+        id_test_idx=id_test_idx,
+        ood_test_idx=np.sort(ood),
+        threshold=threshold,
+        stat_kind=stat_kind,
+        info=info,
+    )

stress_operator/eval.py

@@ -0,0 +1,27 @@
+"""Evaluation: mean test relative-L2 on a split (de-normalized).
+
+Mirrors the eval loop in Transolver ``exp_elas.py``: predictions are decoded with the
+(train-fitted) normalizer before the metric; targets are physical.
+"""
+from __future__ import annotations
+
+import torch
+from torch.utils.data import DataLoader
+
+from .losses.relative_l2 import relative_l2
+
+
+@torch.no_grad()
+def evaluate(model, loader: DataLoader, normalizer, device) -> float:
+    """Return mean relative-L2 over the loader (physical units)."""
+    model.eval()
+    total = 0.0
+    n = 0
+    for coords, sigma in loader:
+        coords = coords.to(device)
+        sigma = sigma.to(device)  # physical (B, N, 1)
+        out = model(coords, None)  # (B, N, 1) normalized
+        out = normalizer.decode(out)  # -> physical
+        total += relative_l2(out, sigma, reduction="sum").item()
+        n += coords.shape[0]
+    return total / max(n, 1)

stress_operator/losses/__init__.py

@@ -0,0 +1,3 @@
+from .relative_l2 import relative_l2, RelativeL2Loss
+
+__all__ = ["relative_l2", "RelativeL2Loss"]

stress_operator/losses/equilibrium.py

@@ -0,0 +1,140 @@
+"""Discrete equilibrium-residual operators on the unstructured point cloud (Gate C).
+
+Static mechanical equilibrium with no body force requires the stress tensor to be
+divergence-free in the interior:  div(sigma) = 0, i.e.
+    d sigma_xx/dx + d sigma_xy/dy = 0
+    d sigma_xy/dx + d sigma_yy/dy = 0
+This holds in the interior regardless of boundary conditions, so it is a valid soft constraint.
+
+Two operators (master plan §2.6):
+  - Approach A (principled): per-node spatial gradient via a weighted moving-least-squares (MLS)
+    fit over the kNN neighborhood. Exact for linear fields -> validatable analytically. The model
+    outputs 3 channels (sigma_xx, sigma_yy, sigma_xy); the equilibrium residual is
+    L_eq = mean ||div(sigma)||^2 over interior nodes.
+  - Approach B (fallback): graph-Laplacian smoothness prior L_smooth = mean (L f)^2 — a
+    plausibility prior *motivated by* (not equal to) equilibrium.
+
+MANDATORY: validate against an analytic field (tests/test_equilibrium.py) before using in the loss.
+"""
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+
+
+def knn_indices(coords: np.ndarray, k: int) -> Tuple[np.ndarray, np.ndarray]:
+    """Return (dist, idx) of the k nearest neighbors (excluding self) for each node."""
+    try:
+        from scipy.spatial import cKDTree
+
+        tree = cKDTree(coords)
+        dist, idx = tree.query(coords, k=k + 1)
+    except Exception:  # pragma: no cover - scipy fallback
+        diff = coords[:, None, :] - coords[None, :, :]
+        d2 = (diff ** 2).sum(-1)
+        idx = np.argsort(d2, axis=1)[:, : k + 1]
+        dist = np.sqrt(np.take_along_axis(d2, idx, axis=1))
+    return dist[:, 1:], idx[:, 1:]  # drop self (column 0)
+
+
+def build_mls_gradient_operators(
+    coords: torch.Tensor, k: int = 12, eps: float = 1e-9
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Build dense (N, N) operators Gx, Gy with (Gx @ f) ~ df/dx and (Gy @ f) ~ df/dy.
+
+    MLS: per node i, minimize sum_j w_ij (f_j - f_i - g . (x_j - x_i))^2 for g = grad f_i.
+    The solution is linear in f, so the per-node stencil rows form constant matrices Gx, Gy
+    (depend only on the mesh). Exact for linear f (reproduces affine functions).
+    """
+    cpu = coords.detach().cpu().numpy().astype(np.float64)
+    n = cpu.shape[0]
+    dist, idx = knn_indices(cpu, k)
+    sigma2 = float((dist.mean()) ** 2) + eps  # Gaussian bandwidth ~ mean neighbor distance
+
+    Gx = np.zeros((n, n), dtype=np.float64)
+    Gy = np.zeros((n, n), dtype=np.float64)
+    for i in range(n):
+        nbr = idx[i]
+        d = cpu[nbr] - cpu[i]                      # (k, 2)
+        w = np.exp(-(d ** 2).sum(1) / sigma2)      # (k,)
+        A = np.einsum("k,ka,kb->ab", w, d, d)      # (2, 2)  weighted moment matrix
+        Ainv = np.linalg.pinv(A)                   # robust to near-degenerate neighborhoods
+        C = (w[:, None] * d) @ Ainv                # (k, 2)  coefficient per neighbor
+        Gx[i, nbr] += C[:, 0]
+        Gx[i, i] -= C[:, 0].sum()
+        Gy[i, nbr] += C[:, 1]
+        Gy[i, i] -= C[:, 1].sum()
+    return (
+        torch.tensor(Gx, dtype=torch.float32, device=coords.device),
+        torch.tensor(Gy, dtype=torch.float32, device=coords.device),
+    )
+
+
+def build_graph_laplacian(coords: torch.Tensor, k: int = 12, eps: float = 1e-9) -> torch.Tensor:
+    """Unnormalized graph Laplacian L = D - W with Gaussian kNN weights (Approach B)."""
+    cpu = coords.detach().cpu().numpy().astype(np.float64)
+    n = cpu.shape[0]
+    dist, idx = knn_indices(cpu, k)
+    sigma2 = float((dist.mean()) ** 2) + eps
+    W = np.zeros((n, n), dtype=np.float64)
+    for i in range(n):
+        nbr = idx[i]
+        w = np.exp(-(dist[i] ** 2) / sigma2)
+        W[i, nbr] = w
+    W = 0.5 * (W + W.T)  # symmetrize
+    L = np.diag(W.sum(1)) - W
+    return torch.tensor(L, dtype=torch.float32, device=coords.device)
+
+
+def interior_mask(coords: torch.Tensor, tol: float = 0.03) -> torch.Tensor:
+    """Boolean mask of interior nodes: inside the [0,1]^2 outer box by ``tol`` on every side.
+
+    The outer-box boundary carries the applied traction / clamp, where div(sigma)=0 need not
+    hold; interior nodes (including those around the void) are valid for the constraint.
+    """
+    x, y = coords[..., 0], coords[..., 1]
+    return (x > tol) & (x < 1 - tol) & (y > tol) & (y < 1 - tol)
+
+
+def divergence(stress3: torch.Tensor, Gx: torch.Tensor, Gy: torch.Tensor) -> torch.Tensor:
+    """Divergence of the symmetric 2x2 stress tensor field.
+
+    stress3: (..., N, 3) = (sigma_xx, sigma_yy, sigma_xy). Returns (..., N, 2) = (div_x, div_y).
+    """
+    sxx = stress3[..., 0]
+    syy = stress3[..., 1]
+    sxy = stress3[..., 2]
+    div_x = torch.einsum("ij,...j->...i", Gx, sxx) + torch.einsum("ij,...j->...i", Gy, sxy)
+    div_y = torch.einsum("ij,...j->...i", Gx, sxy) + torch.einsum("ij,...j->...i", Gy, syy)
+    return torch.stack([div_x, div_y], dim=-1)
+
+
+def equilibrium_residual(
+    stress3: torch.Tensor,
+    Gx: torch.Tensor,
+    Gy: torch.Tensor,
+    mask: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """L_eq = mean ||div(sigma)||^2 over interior nodes. stress3: (..., N, 3)."""
+    div = divergence(stress3, Gx, Gy)        # (..., N, 2)
+    sq = (div ** 2).sum(-1)                   # (..., N)
+    if mask is not None:
+        sq = sq[..., mask]
+    return sq.mean()
+
+
+def laplacian_smoothness(field: torch.Tensor, L: torch.Tensor) -> torch.Tensor:
+    """Approach B: L_smooth = mean (L f)^2. field: (..., N) or (..., N, 1)."""
+    f = field[..., 0] if field.dim() >= 2 and field.shape[-1] == 1 else field
+    Lf = torch.einsum("ij,...j->...i", L, f)
+    return (Lf ** 2).mean()
+
+
+def von_mises(stress3: torch.Tensor) -> torch.Tensor:
+    """von Mises stress from (sigma_xx, sigma_yy, sigma_xy): sqrt(sxx^2 - sxx*syy + syy^2 + 3 sxy^2)."""
+    sxx = stress3[..., 0]
+    syy = stress3[..., 1]
+    sxy = stress3[..., 2]
+    return torch.sqrt(torch.clamp(sxx ** 2 - sxx * syy + syy ** 2 + 3 * sxy ** 2, min=0.0) + 1e-12)

stress_operator/losses/relative_l2.py

@@ -0,0 +1,50 @@
+"""Relative-L2 metric and data loss.
+
+Faithful to Transolver ``utils/testloss.py::TestLoss.rel`` (the number compared to the
+published 0.0064 / 0.0050). Computed **after de-normalizing** predictions (see
+``docs/RECONCILIATION.md`` §4-5). No denominator epsilon is used by default, matching upstream.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+def relative_l2(
+    pred: torch.Tensor,
+    true: torch.Tensor,
+    reduction: str = "mean",
+    eps: float = 0.0,
+) -> torch.Tensor:
+    """Per-sample relative L2: ``||pred - true||_2 / ||true||_2``.
+
+    Args:
+        pred, true: tensors shaped ``(B, ...)``; everything after dim 0 is flattened
+            per sample. Must be **de-normalized** (physical units).
+        reduction: ``"mean"`` (default), ``"sum"`` (matches ``size_average=False`` upstream
+            before the caller divides by ``n``), or ``"none"`` for the per-sample vector.
+        eps: optional denominator epsilon (default 0.0 to match upstream exactly).
+    """
+    b = pred.shape[0]
+    diff = torch.linalg.vector_norm(pred.reshape(b, -1) - true.reshape(b, -1), ord=2, dim=1)
+    den = torch.linalg.vector_norm(true.reshape(b, -1), ord=2, dim=1)
+    rel = diff / (den + eps)
+    if reduction == "mean":
+        return rel.mean()
+    if reduction == "sum":
+        return rel.sum()
+    if reduction == "none":
+        return rel
+    raise ValueError(f"unknown reduction {reduction!r}")
+
+
+class RelativeL2Loss(nn.Module):
+    """Module wrapper around :func:`relative_l2`."""
+
+    def __init__(self, reduction: str = "mean", eps: float = 0.0):
+        super().__init__()
+        self.reduction = reduction
+        self.eps = eps
+
+    def forward(self, pred: torch.Tensor, true: torch.Tensor) -> torch.Tensor:
+        return relative_l2(pred, true, reduction=self.reduction, eps=self.eps)

stress_operator/models/__init__.py

@@ -0,0 +1,3 @@
+from .transolver import StressOperator, build_model, count_parameters, make_attention
+
+__all__ = ["StressOperator", "build_model", "count_parameters", "make_attention"]

stress_operator/models/blocks.py

@@ -0,0 +1,77 @@
+"""Building blocks: MLP and the pre-norm Transolver block.
+
+Faithful to Transolver ``model/Transolver_Irregular_Mesh.py`` (MIT). The only structural
+change is that the attention module is injected (so Stage 2 can swap Physics-Attention for
+LinearNO without touching anything else).
+"""
+import torch.nn as nn
+
+ACTIVATION = {
+    "gelu": nn.GELU,
+    "tanh": nn.Tanh,
+    "sigmoid": nn.Sigmoid,
+    "relu": nn.ReLU,
+    "leaky_relu": lambda: nn.LeakyReLU(0.1),
+    "softplus": nn.Softplus,
+    "ELU": nn.ELU,
+    "silu": nn.SiLU,
+}
+
+
+class MLP(nn.Module):
+    """Transolver MLP: Linear->act (->[Linear->act]^n_layers) ->Linear, optional residual."""
+
+    def __init__(self, n_input, n_hidden, n_output, n_layers=1, act="gelu", res=True):
+        super().__init__()
+        if act not in ACTIVATION:
+            raise NotImplementedError(act)
+        act_cls = ACTIVATION[act]
+        self.n_layers = n_layers
+        self.res = res
+        self.linear_pre = nn.Sequential(nn.Linear(n_input, n_hidden), act_cls())
+        self.linear_post = nn.Linear(n_hidden, n_output)
+        self.linears = nn.ModuleList(
+            [nn.Sequential(nn.Linear(n_hidden, n_hidden), act_cls()) for _ in range(n_layers)]
+        )
+
+    def forward(self, x):
+        x = self.linear_pre(x)
+        for layer in self.linears:
+            x = layer(x) + x if self.res else layer(x)
+        return self.linear_post(x)
+
+
+class TransolverBlock(nn.Module):
+    """Pre-norm transformer block; the last block also carries the decoder head.
+
+        fx = fx + Attn(LayerNorm(fx))
+        fx = fx + MLP(LayerNorm(fx))
+        (last block) return Linear(LayerNorm(fx))  -> out_dim
+    """
+
+    def __init__(
+        self,
+        attention: nn.Module,
+        hidden_dim: int,
+        dropout: float = 0.0,
+        act: str = "gelu",
+        mlp_ratio: int = 1,
+        last_layer: bool = False,
+        out_dim: int = 1,
+    ):
+        super().__init__()
+        self.last_layer = last_layer
+        self.ln_1 = nn.LayerNorm(hidden_dim)
+        self.Attn = attention
+        self.ln_2 = nn.LayerNorm(hidden_dim)
+        self.mlp = MLP(hidden_dim, hidden_dim * mlp_ratio, hidden_dim, n_layers=0, res=False, act=act)
+        if last_layer:
+            self.ln_3 = nn.LayerNorm(hidden_dim)
+            self.mlp2 = nn.Linear(hidden_dim, out_dim)
+
+    def forward(self, fx):
+        fx = self.Attn(self.ln_1(fx)) + fx
+        fx = self.mlp(self.ln_2(fx)) + fx
+        if self.last_layer:
+            return self.mlp2(self.ln_3(fx))
+        return fx

stress_operator/models/linear_no.py

@@ -0,0 +1,101 @@
+"""LinearNO attention block (Stage 2 / Gate B).
+
+Reimplemented from the equations of:
+    "Transolver is a Linear Transformer: Revisiting Physics-Attention through the Lens of
+     Linear Attention" — Hu, Liu, Qiao, Sun, Dou (NUDT), AAAI 2026, arXiv:2511.06294.
+There is no official public repo; only the attention block differs from Transolver.
+
+Core idea: Physics-Attention is the special case of linear attention
+``Attention(Q,K,V) ~ phi(Q) (psi^T(K) V)`` in which (a) phi and psi come from the SAME linear
+layer (differing only by normalization) and (b) there is an extra slice self-attention step.
+LinearNO removes BOTH constraints:
+  1. learn Q and K projections independently (break weight sharing), and
+  2. drop the slice self-attention (identity).
+
+  LinearNO(H) = phi(Q) @ ( psi(K)^T @ V )
+  Q = H Wq ; K = H Wk ; V = Linear_V(H)
+  phi(Q) = softmax_over_M( Linear_Q(Q) )   # (N, M)  rows sum to 1  <- softmax along M
+  psi(K) = softmax_over_N( Linear_K(K) )   # (N, M)  cols sum to 1  <- softmax along N
+
+### THE make-or-break detail (paper Table 6, Elasticity):
+    phi softmax over M (slices), psi softmax over N (points) -> 0.0050  (correct)
+    swapping these dims -> 0.0081..0.0112 (wrong). If LinearNO lands at 0.008-0.011, the
+    softmax dimensions are almost certainly swapped.
+
+Associativity: compute ``psi^T V`` first (M x d), then ``phi @ (...)`` -> O(N*M*d), linear in N.
+
+Two variants (see docs/RECONCILIATION.md for the parameter-count tension):
+  - "independent" (paper skeleton, default): separate Wq, Wk, Wv as full dim->inner
+    projections. Most literal to the plan's reference code. ~0.85M params (> baseline).
+  - "shared_qk": share one dim->inner base for the slice projections (phi, psi), keep a
+    separate dim->inner V; breaks slice-weight sharing via two separate small slice layers.
+    ~0.72M params (~= baseline). Use this to satisfy the Gate-B "<= baseline params" constraint.
+An optional ``project_out=False`` drops the output projection (folds the per-head concat
+directly), reaching ~0.59M (the plan's stated LinearNO size).
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+
+class LinearNO(nn.Module):
+    def __init__(
+        self,
+        dim,
+        heads=8,
+        dim_head=16,
+        slice_num=64,
+        dropout=0.0,
+        variant: str = "independent",
+        project_out: bool = True,
+    ):
+        super().__init__()
+        inner = heads * dim_head
+        self.h, self.m, self.dh = heads, slice_num, dim_head
+        self.variant = variant
+        self.project_out = project_out
+
+        if variant == "independent":
+            # Independent Q, K, V projections (paper modification 1, literal).
+            self.to_q = nn.Linear(dim, inner, bias=False)
+            self.to_k = nn.Linear(dim, inner, bias=False)
+            self.to_v = nn.Linear(dim, inner, bias=False)
+        elif variant == "shared_qk":
+            # Share one base projection for the slice space (Q==K base), keep V separate.
+            self.to_qk = nn.Linear(dim, inner, bias=False)
+            self.to_v = nn.Linear(dim, inner, bias=False)
+        else:
+            raise ValueError(f"unknown variant {variant!r} (expected 'independent' or 'shared_qk')")
+
+        # Asymmetric slice projections: query->slices and key->slices are SEPARATE weights.
+        self.lin_q = nn.Linear(dim_head, slice_num)
+        self.lin_k = nn.Linear(dim_head, slice_num)
+
+        if project_out:
+            self.to_out = nn.Sequential(nn.Linear(inner, dim), nn.Dropout(dropout))
+        else:
+            self.to_out = nn.Dropout(dropout)
+
+    def _heads(self, t, B, N):
+        return t.reshape(B, N, self.h, self.dh).permute(0, 2, 1, 3).contiguous()  # (B,H,N,dh)
+
+    def forward(self, x):  # x: (B, N, C)
+        B, N, C = x.shape
+        if self.variant == "independent":
+            q = self._heads(self.to_q(x), B, N)  # (B,H,N,dh)
+            k = self._heads(self.to_k(x), B, N)
+            v = self._heads(self.to_v(x), B, N)
+        else:  # shared_qk
+            base = self._heads(self.to_qk(x), B, N)
+            q = base
+            k = base
+            v = self._heads(self.to_v(x), B, N)
+
+        phi = self.lin_q(q).softmax(dim=-1)  # (B,H,N,M)  softmax OVER M (slices)  <- rows sum to 1
+        psi = self.lin_k(k).softmax(dim=-2)  # (B,H,N,M)  softmax OVER N (points)  <- cols sum to 1
+        kv = torch.einsum("bhnm,bhnd->bhmd", psi, v)  # (B,H,M,dh)  cheap inner product first
+        out = torch.einsum("bhnm,bhmd->bhnd", phi, kv)  # (B,H,N,dh)  linear in N
+        out = rearrange(out, "b h n d -> b n (h d)")
+        return self.to_out(out)

stress_operator/models/physics_attention.py

@@ -0,0 +1,78 @@
+"""Physics-Attention for irregular meshes / point clouds.
+
+This module is reused **verbatim** from the Transolver project and is licensed under the MIT
+License. Attribution preserved per the license terms.
+
+    Transolver: A Fast Transformer Solver for PDEs on General Geometries
+    Wu, Luo, Wang, Wang, Long (THUML / Tsinghua University), ICML 2024 Spotlight.
+    Source: https://github.com/thuml/Transolver  (MIT License)
+    Paper:  arXiv:2402.02366
+
+Key constants preserved exactly: temperature init 0.5, orthogonal init of the slice
+projection, slice-norm epsilon 1e-5, attention scale dim_head**-0.5.
+
+Note: in the assembled model (see transolver.py), the operator's global ``trunc_normal_``
+weight initialization runs *after* this block is constructed and overwrites the orthogonal
+init of ``in_project_slice`` — this matches the upstream assembly order and is what produced
+the published 0.0064. We keep the orthogonal init here to stay verbatim with upstream.
+"""
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+
+class Physics_Attention_Irregular_Mesh(nn.Module):
+    ## for irregular meshes in 1D, 2D or 3D space
+    # NOTE: the dim_head=64 default is upstream's generic default. For the Elasticity config the
+    # assembler (transolver.py) always passes dim_head = n_hidden // n_heads = 16; the default here
+    # is never used in this repo. See docs/RECONCILIATION.md §3.
+    def __init__(self, dim, heads=8, dim_head=64, dropout=0., slice_num=64):
+        super().__init__()
+        inner_dim = dim_head * heads
+        self.dim_head = dim_head
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        self.softmax = nn.Softmax(dim=-1)
+        self.dropout = nn.Dropout(dropout)
+        self.temperature = nn.Parameter(torch.ones([1, heads, 1, 1]) * 0.5)
+
+        self.in_project_x = nn.Linear(dim, inner_dim)
+        self.in_project_fx = nn.Linear(dim, inner_dim)
+        self.in_project_slice = nn.Linear(dim_head, slice_num)
+        for l in [self.in_project_slice]:
+            torch.nn.init.orthogonal_(l.weight)  # use a principled initialization
+        self.to_q = nn.Linear(dim_head, dim_head, bias=False)
+        self.to_k = nn.Linear(dim_head, dim_head, bias=False)
+        self.to_v = nn.Linear(dim_head, dim_head, bias=False)
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        # B N C
+        B, N, C = x.shape
+
+        ### (1) Slice
+        fx_mid = self.in_project_fx(x).reshape(B, N, self.heads, self.dim_head) \
+            .permute(0, 2, 1, 3).contiguous()  # B H N C
+        x_mid = self.in_project_x(x).reshape(B, N, self.heads, self.dim_head) \
+            .permute(0, 2, 1, 3).contiguous()  # B H N C
+        slice_weights = self.softmax(self.in_project_slice(x_mid) / self.temperature)  # B H N G
+        slice_norm = slice_weights.sum(2)  # B H G
+        slice_token = torch.einsum("bhnc,bhng->bhgc", fx_mid, slice_weights)
+        slice_token = slice_token / ((slice_norm + 1e-5)[:, :, :, None].repeat(1, 1, 1, self.dim_head))
+
+        ### (2) Attention among slice tokens
+        q_slice_token = self.to_q(slice_token)
+        k_slice_token = self.to_k(slice_token)
+        v_slice_token = self.to_v(slice_token)
+        dots = torch.matmul(q_slice_token, k_slice_token.transpose(-1, -2)) * self.scale
+        attn = self.softmax(dots)
+        attn = self.dropout(attn)
+        out_slice_token = torch.matmul(attn, v_slice_token)  # B H G D
+
+        ### (3) Deslice
+        out_x = torch.einsum("bhgc,bhng->bhnc", out_slice_token, slice_weights)
+        out_x = rearrange(out_x, 'b h n d -> b n (h d)')
+        return self.to_out(out_x)

stress_operator/models/transolver.py

@@ -0,0 +1,159 @@
+"""Assemble the stress operator: encoder -> N pre-norm blocks -> decoder head.
+
+Faithful reproduction of Transolver ``model/Transolver_Irregular_Mesh.py::Model`` (MIT). The
+attention type is configurable (``physics`` for Stage 1, ``linearno`` for Stage 2) and is the
+*only* thing that changes between gates — a controlled comparison.
+
+Reproduction note: ``initialize_weights()`` runs after the blocks are built, so the global
+``trunc_normal_(std=0.02)`` init is applied to every Linear, **overwriting** the orthogonal
+init of each attention's ``in_project_slice``. This matches the upstream assembly order.
+"""
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from .blocks import MLP, TransolverBlock
+from .physics_attention import Physics_Attention_Irregular_Mesh
+
+
+def make_attention(
+    kind: str,
+    dim,
+    heads,
+    dim_head,
+    dropout,
+    slice_num,
+    linearno_variant: str = "shared_qk",
+    linearno_project_out: bool = False,
+) -> nn.Module:
+    if kind == "physics":
+        return Physics_Attention_Irregular_Mesh(
+            dim, heads=heads, dim_head=dim_head, dropout=dropout, slice_num=slice_num
+        )
+    if kind == "linearno":
+        from .linear_no import LinearNO  # lazy: only needed at Stage 2
+
+        return LinearNO(
+            dim,
+            heads=heads,
+            dim_head=dim_head,
+            slice_num=slice_num,
+            dropout=dropout,
+            variant=linearno_variant,
+            project_out=linearno_project_out,
+        )
+    raise ValueError(f"unknown attention kind {kind!r} (expected 'physics' or 'linearno')")
+
+
+class StressOperator(nn.Module):
+    def __init__(
+        self,
+        attention: str = "physics",
+        space_dim: int = 2,
+        n_layers: int = 8,
+        n_hidden: int = 128,
+        dropout: float = 0.0,
+        n_heads: int = 8,
+        dim_head: int | None = None,
+        mlp_ratio: int = 1,
+        fun_dim: int = 0,
+        out_dim: int = 1,
+        slice_num: int = 64,
+        unified_pos: bool = False,
+        ref: int = 8,
+        act: str = "gelu",
+        linearno_variant: str = "shared_qk",
+        linearno_project_out: bool = False,
+    ):
+        super().__init__()
+        if dim_head is None:
+            dim_head = n_hidden // n_heads  # = 16 for the Elasticity config (repo value)
+        self.attention_kind = attention
+        self.unified_pos = unified_pos
+        self.ref = ref
+        self.n_hidden = n_hidden
+
+        in_dim = (fun_dim + ref * ref) if unified_pos else (fun_dim + space_dim)
+        self.preprocess = MLP(in_dim, n_hidden * 2, n_hidden, n_layers=0, res=False, act=act)
+
+        self.blocks = nn.ModuleList(
+            [
+                TransolverBlock(
+                    attention=make_attention(
+                        attention, n_hidden, n_heads, dim_head, dropout, slice_num,
+                        linearno_variant=linearno_variant,
+                        linearno_project_out=linearno_project_out,
+                    ),
+                    hidden_dim=n_hidden,
+                    dropout=dropout,
+                    act=act,
+                    mlp_ratio=mlp_ratio,
+                    last_layer=(i == n_layers - 1),
+                    out_dim=out_dim,
+                )
+                for i in range(n_layers)
+            ]
+        )
+        self.initialize_weights()
+        self.placeholder = nn.Parameter((1 / n_hidden) * torch.rand(n_hidden, dtype=torch.float))
+
+    def initialize_weights(self):
+        self.apply(self._init_weights)
+
+    @staticmethod
+    def _init_weights(m):
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d)):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def get_grid(self, x):
+        """Unified positional grid (only used when ``unified_pos`` is True). Device-agnostic."""
+        b = x.shape[0]
+        device = x.device
+        gx = torch.linspace(0, 1, self.ref, device=device).reshape(1, self.ref, 1, 1).repeat(b, 1, self.ref, 1)
+        gy = torch.linspace(0, 1, self.ref, device=device).reshape(1, 1, self.ref, 1).repeat(b, self.ref, 1, 1)
+        grid_ref = torch.cat((gx, gy), dim=-1).reshape(b, self.ref * self.ref, 2)
+        pos = torch.sqrt(((x[:, :, None, :] - grid_ref[:, None, :, :]) ** 2).sum(-1))
+        return pos.reshape(b, x.shape[1], self.ref * self.ref).contiguous()
+
+    def forward(self, x, fx=None):
+        # x: (B, N, space_dim) node coordinates; fx: optional extra input function
+        if self.unified_pos:
+            x = self.get_grid(x)
+        fx = self.preprocess(x if fx is None else torch.cat((x, fx), dim=-1))
+        fx = fx + self.placeholder[None, None, :]
+        for block in self.blocks:
+            fx = block(fx)
+        return fx  # (B, N, out_dim)
+
+
+def build_model(model_cfg: dict) -> StressOperator:
+    """Instantiate :class:`StressOperator` from a config dict (configs/*.yaml ``model`` block)."""
+    return StressOperator(
+        attention=model_cfg.get("attention", "physics"),
+        space_dim=model_cfg.get("space_dim", 2),
+        n_layers=model_cfg.get("n_layers", 8),
+        n_hidden=model_cfg.get("n_hidden", 128),
+        dropout=model_cfg.get("dropout", 0.0),
+        n_heads=model_cfg.get("n_heads", 8),
+        dim_head=model_cfg.get("dim_head", None),
+        mlp_ratio=model_cfg.get("mlp_ratio", 1),
+        fun_dim=model_cfg.get("fun_dim", 0),
+        out_dim=model_cfg.get("out_dim", 1),
+        slice_num=model_cfg.get("slice_num", 64),
+        unified_pos=model_cfg.get("unified_pos", False),
+        ref=model_cfg.get("ref", 8),
+        act=model_cfg.get("act", "gelu"),
+        linearno_variant=model_cfg.get("linearno_variant", "shared_qk"),
+        linearno_project_out=model_cfg.get("linearno_project_out", False),
+    )
+
+
+def count_parameters(model: nn.Module) -> int:
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)

stress_operator/seeds.py

@@ -0,0 +1,31 @@
+"""Deterministic seeding across python / numpy / torch."""
+from __future__ import annotations
+
+import os
+import random
+
+
+def set_seed(seed: int, deterministic: bool = True) -> None:
+    """Seed all RNGs used in training/eval.
+
+    Mirrors the upstream Transolver/Geo-FNO seeding (``torch.manual_seed``,
+    ``np.random.seed``, cudnn deterministic) and extends it to python's ``random``
+    and the ``PYTHONHASHSEED`` env var so runs are reproducible across seeds {0,1,2}.
+    """
+    import numpy as np
+    import torch
+
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    if deterministic:
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        # NOTE: torch.use_deterministic_algorithms(True) was tried but it markedly DEGRADED
+        # convergence of the Physics-Attention baseline (slice self-attention) on GPU (baseline
+        # eager landed ~0.0090 vs ~0.0068 without it), while leaving LinearNO unaffected. We
+        # therefore do not force it; reproducibility is handled by averaging seeds and disclosing
+        # run-to-run variance (README + model card). See docs/RECONCILIATION.md / PART 6 caveat 6.

stress_operator/train.py

@@ -0,0 +1,225 @@
+"""Single training entrypoint; ``--config`` selects the stage.
+
+Faithful to Transolver ``exp_elas.py``:
+  - AdamW (lr, weight_decay), CosineAnnealingLR(T_max=epochs)
+  - batch_size 1, gradient clipping at max_grad_norm (0.1)
+  - loss = relative-L2 in physical units: predictions are de-normalized before the loss;
+    targets are physical (decode(encode(s)) == s, so storing physical targets is equivalent).
+
+Writes a run-log JSON to ``results/`` per master plan §0.2.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import time
+from typing import Any, Dict, Optional
+
+import torch
+import yaml
+
+from .data.dataset import build_splits
+from .losses.relative_l2 import relative_l2
+from .models.transolver import build_model, count_parameters
+from .seeds import set_seed
+from .utils.logging import MODAL_RATES_PER_SEC, write_run_log
+
+
+def load_config(path: str) -> Dict[str, Any]:
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+def run_training(
+    config: Dict[str, Any],
+    seed: int,
+    data_dir: str,
+    device: Optional[str] = None,
+    gpu_name: str = "CPU",
+    results_path: Optional[str] = None,
+    ckpt_path: Optional[str] = None,
+    log_every: int = 50,
+    max_epochs: Optional[int] = None,
+    ntrain_override: Optional[int] = None,
+    splits=None,
+) -> Dict[str, Any]:
+    """Train one model for one seed; return final metrics and write a run-log JSON.
+
+    If ``ckpt_path`` is given, also save ``{state_dict, normalizer{mean,std}, config, seed,
+    metrics}`` (the normalizer stats are required to de-normalize predictions at inference).
+    If ``splits`` (a ``Splits`` from ``build_splits_from_indices``) is given, it overrides the
+    default first-1000/last-200 split (used by the OOD evaluation).
+    """
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+    set_seed(seed)
+
+    data_cfg = config["data"]
+    train_cfg = config["train"]
+    model_cfg = config["model"]
+
+    if splits is None:
+        ntrain = ntrain_override or data_cfg.get("ntrain", 1000)
+        ntest = data_cfg.get("ntest", 200)
+        splits = build_splits(data_dir, ntrain=ntrain, ntest=ntest)
+    ntest = splits.test_coords.shape[0]
+    normalizer = splits.normalizer.to(device)
+
+    # GPU-resident dataset: the whole thing is tiny (~10 MB), so we keep it on-device and
+    # batch by index. This removes DataLoader + per-iteration host->device + per-iteration
+    # .item() sync overhead, which dominates wall-clock at batch_size 1. The math is identical
+    # to the DataLoader path (same batch_size, same loss, same seeded shuffle order).
+    batch_size = train_cfg.get("batch_size", 1)
+    eval_every = int(train_cfg.get("eval_every", 1))
+
+    def _3d(t):
+        return (t if t.dim() == 3 else t.unsqueeze(-1)).to(device)
+
+    train_coords = splits.train_coords.to(device)   # (ntrain, 972, 2)
+    train_sigma = _3d(splits.train_sigma)            # (ntrain, 972, 1) physical
+    test_coords = splits.test_coords.to(device)
+    test_sigma = _3d(splits.test_sigma)
+    ntrain_eff = train_coords.shape[0]
+
+    base_model = build_model(model_cfg).to(device)
+    n_params = count_parameters(base_model)
+
+    # Optional torch.compile (CUDA graphs) to cut per-iteration kernel-launch overhead, which
+    # dominates wall-clock at batch_size 1. Same math, static input shape (1, 972, 2). The
+    # checkpoint is saved from base_model so its state_dict keys stay clean (no _orig_mod prefix).
+    model = base_model
+    if bool(train_cfg.get("compile", False)) and device == "cuda":
+        try:
+            model = torch.compile(base_model, mode="reduce-overhead")
+            print(f"[seed {seed}] torch.compile enabled (reduce-overhead)", flush=True)
+        except Exception as e:  # pragma: no cover
+            print(f"[seed {seed}] torch.compile failed ({e}); falling back to eager", flush=True)
+            model = base_model
+
+    lr = float(train_cfg.get("lr", 1e-3))
+    wd = float(train_cfg.get("weight_decay", 1e-5))
+    betas = tuple(train_cfg.get("betas", (0.9, 0.999)))
+    epochs = max_epochs or int(train_cfg.get("epochs", 500))
+    max_grad_norm = train_cfg.get("max_grad_norm", None)
+
+    optimizer = torch.optim.AdamW(base_model.parameters(), lr=lr, weight_decay=wd, betas=betas)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
+
+    shuffle_gen = torch.Generator().manual_seed(seed)  # reproducible per-epoch shuffle
+
+    @torch.no_grad()
+    def eval_test() -> float:
+        model.eval()
+        total = 0.0
+        for i in range(0, ntest, batch_size):
+            out = normalizer.decode(model(test_coords[i:i + batch_size], None))
+            total += relative_l2(out, test_sigma[i:i + batch_size], reduction="sum").item()
+        return total / ntest
+
+    t0 = time.time()
+    best_rel = float("inf")
+    test_rel = float("nan")
+    history = []
+    for ep in range(epochs):
+        model.train()
+        perm = torch.randperm(ntrain_eff, generator=shuffle_gen).to(device)
+        running = torch.zeros((), device=device)
+        for s in range(0, ntrain_eff, batch_size):
+            idx = perm[s:s + batch_size]
+            optimizer.zero_grad()
+            out = normalizer.decode(model(train_coords[idx], None))  # -> physical
+            loss = relative_l2(out, train_sigma[idx], reduction="sum")
+            loss.backward()
+            if max_grad_norm is not None:
+                torch.nn.utils.clip_grad_norm_(base_model.parameters(), max_grad_norm)
+            optimizer.step()
+            running += loss.detach()
+        scheduler.step()
+        train_rel = (running / ntrain_eff).item()
+
+        if (ep % eval_every == 0) or (ep >= epochs - 5):
+            test_rel = eval_test()
+            best_rel = min(best_rel, test_rel)
+        history.append({"epoch": ep, "train_rel": train_rel, "test_rel": test_rel})
+        if ep % log_every == 0 or ep == epochs - 1:
+            print(
+                f"[seed {seed}] epoch {ep:4d}  train_rel={train_rel:.5f}  test_rel={test_rel:.5f}",
+                flush=True,
+            )
+
+    wall = time.time() - t0
+    rate = MODAL_RATES_PER_SEC.get(gpu_name, 0.0)
+    est_cost = wall * rate
+
+    final_metrics = {
+        "test_rel_l2": round(test_rel, 6),
+        "best_test_rel_l2": round(best_rel, 6),
+        "train_rel_l2": round(train_rel, 6),
+        "n_params": n_params,
+        "epochs": epochs,
+    }
+
+    if ckpt_path is not None:
+        os.makedirs(os.path.dirname(ckpt_path) or ".", exist_ok=True)
+        torch.save(
+            {
+                "state_dict": base_model.state_dict(),
+                "normalizer": {
+                    "mean": normalizer.mean.detach().cpu(),
+                    "std": normalizer.std.detach().cpu(),
+                },
+                "config": config,
+                "seed": seed,
+                "metrics": final_metrics,
+            },
+            ckpt_path,
+        )
+        print(f"[seed {seed}] saved checkpoint -> {ckpt_path}")
+
+    if results_path is None:
+        os.makedirs("results", exist_ok=True)
+        results_path = os.path.join("results", f"{config.get('name','run')}_seed{seed}.json")
+    write_run_log(
+        path=results_path,
+        config=config,
+        seed=seed,
+        final_metrics=final_metrics,
+        wall_clock_sec=wall,
+        gpu=gpu_name,
+        est_cost_usd=est_cost,
+        extra={"history_tail": history[-5:]},
+    )
+    print(
+        f"[seed {seed}] DONE  test_rel_l2={test_rel:.6f}  best={best_rel:.6f}  "
+        f"params={n_params}  wall={wall:.0f}s  gpu={gpu_name}  est_cost=${est_cost:.4f}"
+    )
+    return final_metrics
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description="Train the stress operator (one seed).")
+    ap.add_argument("--config", required=True)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--data-dir", default="data")
+    ap.add_argument("--device", default=None)
+    ap.add_argument("--gpu-name", default="CPU", help="for cost accounting (A10 / A100-40GB / CPU)")
+    ap.add_argument("--results-path", default=None)
+    ap.add_argument("--max-epochs", type=int, default=None, help="override epochs (local smoke runs)")
+    ap.add_argument("--ntrain", type=int, default=None, help="override ntrain (local smoke runs)")
+    args = ap.parse_args()
+
+    config = load_config(args.config)
+    run_training(
+        config=config,
+        seed=args.seed,
+        data_dir=args.data_dir,
+        device=args.device,
+        gpu_name=args.gpu_name,
+        results_path=args.results_path,
+        max_epochs=args.max_epochs,
+        ntrain_override=args.ntrain,
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

stress_operator/train_eqreg.py

@@ -0,0 +1,234 @@
+"""Training with the equilibrium-residual regularizer (Stage 3 / Gate C).
+
+Two approaches (master plan §2.6), selected by ``config['equilibrium']['approach']``:
+
+  A (principled): model outputs 3 channels (sigma_xx, sigma_yy, sigma_xy). The data loss compares
+    the derived von Mises stress to the scalar target; the physics loss penalizes the discrete
+    divergence residual ||div(sigma)||^2 on interior nodes (operators validated analytically).
+    Target is scaled by a constant S so outputs stay O(1); relative-L2 is scale-invariant, so the
+    reported number equals the physical relative-L2.
+
+  B (fallback): model keeps the 1-channel scalar output (identical accuracy path to LinearNO);
+    the physics loss is a graph-Laplacian smoothness prior — a plausibility prior *motivated by*
+    (not equal to) equilibrium. Honest framing required.
+
+The per-sample discrete operators depend only on the mesh, so they are precomputed once (sparse)
+and reused every epoch. Total loss: L = data_loss + lambda * physics_loss.
+"""
+from __future__ import annotations
+
+import os
+import time
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+
+from .data.dataset import build_splits
+from .losses.equilibrium import (
+    build_graph_laplacian,
+    build_mls_gradient_operators,
+    interior_mask,
+    von_mises,
+)
+from .losses.relative_l2 import relative_l2
+from .models.transolver import build_model, count_parameters
+from .seeds import set_seed
+from .utils.logging import MODAL_RATES_PER_SEC, write_run_log
+
+
+def _precompute_grad_ops(
+    coords: torch.Tensor, k: int, tol: float, device
+) -> List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+    """Per-sample (Gx_sparse, Gy_sparse, interior_mask) for Approach A. Built once (CPU build)."""
+    ops = []
+    for s in range(coords.shape[0]):
+        c = coords[s]
+        Gx, Gy = build_mls_gradient_operators(c, k=k)
+        mask = interior_mask(c, tol).to(device)
+        ops.append((Gx.to_sparse().to(device), Gy.to_sparse().to(device), mask))
+    return ops
+
+
+def _precompute_laplacians(coords: torch.Tensor, k: int, device) -> List[torch.Tensor]:
+    """Per-sample sparse graph Laplacian for Approach B."""
+    return [build_graph_laplacian(coords[s], k=k).to_sparse().to(device) for s in range(coords.shape[0])]
+
+
+def _sparse_div_residual(stress3: torch.Tensor, Gx_s, Gy_s, mask) -> torch.Tensor:
+    """L_eq = mean_interior ||div(sigma)||^2 for one sample. stress3: (N, 3)."""
+    sxx = stress3[:, 0:1]
+    syy = stress3[:, 1:2]
+    sxy = stress3[:, 2:3]
+    div_x = torch.sparse.mm(Gx_s, sxx) + torch.sparse.mm(Gy_s, sxy)   # (N,1)
+    div_y = torch.sparse.mm(Gx_s, sxy) + torch.sparse.mm(Gy_s, syy)
+    sq = (div_x.squeeze(-1) ** 2 + div_y.squeeze(-1) ** 2)            # (N,)
+    return sq[mask].mean()
+
+
+def run_training_eqreg(
+    config: Dict[str, Any],
+    seed: int,
+    data_dir: str,
+    device: Optional[str] = None,
+    gpu_name: str = "CPU",
+    results_path: Optional[str] = None,
+    ckpt_path: Optional[str] = None,
+    log_every: int = 50,
+    max_epochs: Optional[int] = None,
+    splits=None,
+    lambda_override: Optional[float] = None,
+) -> Dict[str, Any]:
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+    set_seed(seed)
+
+    data_cfg = config["data"]
+    train_cfg = config["train"]
+    model_cfg = config["model"]
+    eq_cfg = config["equilibrium"]
+    approach = eq_cfg.get("approach", "A").upper()
+    lam = float(lambda_override) if lambda_override is not None else float(eq_cfg.get("lambda", 0.05))
+    knn_k = int(eq_cfg.get("knn_k", 12))
+    tol = float(eq_cfg.get("interior_tol", 0.03))
+
+    if splits is None:
+        ntrain = data_cfg.get("ntrain", 1000)
+        ntest = data_cfg.get("ntest", 200)
+        splits = build_splits(data_dir, ntrain=ntrain, ntest=ntest)
+    ntest = splits.test_coords.shape[0]
+    normalizer = splits.normalizer.to(device)
+
+    train_coords = splits.train_coords.to(device)
+    test_coords = splits.test_coords.to(device)
+    train_sigma = (splits.train_sigma if splits.train_sigma.dim() == 3 else splits.train_sigma.unsqueeze(-1)).to(device)
+    test_sigma = (splits.test_sigma if splits.test_sigma.dim() == 3 else splits.test_sigma.unsqueeze(-1)).to(device)
+    n_train = train_coords.shape[0]
+
+    # Scale for Approach A: keep von Mises target O(1). relative-L2 is scale-invariant, so the
+    # reported metric equals the physical relative-L2 regardless of S.
+    S = float(splits.train_sigma.mean()) if approach == "A" else 1.0
+
+    print(f"[eqreg seed {seed}] approach={approach} lambda={lam} k={knn_k} S={S:.2f}", flush=True)
+    t_build = time.time()
+    if approach == "A":
+        train_ops = _precompute_grad_ops(splits.train_coords, knn_k, tol, device)
+        test_ops = _precompute_grad_ops(splits.test_coords, knn_k, tol, device)
+    else:
+        train_ops = _precompute_laplacians(splits.train_coords, knn_k, device)
+        test_ops = _precompute_laplacians(splits.test_coords, knn_k, device)
+    print(f"[eqreg seed {seed}] precomputed operators in {time.time()-t_build:.0f}s", flush=True)
+
+    model = build_model(model_cfg).to(device)
+    n_params = count_parameters(model)
+
+    lr = float(train_cfg.get("lr", 1e-3))
+    wd = float(train_cfg.get("weight_decay", 1e-5))
+    betas = tuple(train_cfg.get("betas", (0.9, 0.999)))
+    epochs = max_epochs or int(train_cfg.get("epochs", 500))
+    max_grad_norm = train_cfg.get("max_grad_norm", None)
+    eval_every = int(train_cfg.get("eval_every", 10))
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd, betas=betas)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
+    gen = torch.Generator().manual_seed(seed)
+
+    def data_and_phys(out, sigma_phys, ops_i):
+        """Return (data_loss, phys_residual) for one sample. out: (N, C).
+
+        Approach A scale convention (IMPORTANT — do NOT 'decode' ``out``):
+        ``out`` is the RAW 3-channel model output (no normalizer applied). It learns the physical
+        stress tensor divided by the constant ``S`` (a pure scale, NOT the affine z-score). Because
+        von Mises is homogeneous of degree 1 and relative-L2 is jointly scale-invariant,
+            relative_l2(von_mises(out), sigma_phys/S) == relative_l2(S*von_mises(out), sigma_phys)
+        i.e. the reported data loss EQUALS the physical relative-L2 of the prediction ``S*von_mises(out)``
+        against the physical target (verified numerically). Applying the scalar normalizer's
+        ``decode`` here would be WRONG: it would add the von-Mises mean (~187) to every tensor
+        component. The divergence residual below is therefore in scaled units (= physical/S^2);
+        it is converted to physical units only for reporting (see final_metrics).
+        """
+        if approach == "A":
+            vm = von_mises(out)                      # (N,)  von Mises of the scaled tensor
+            data_loss = relative_l2(vm.unsqueeze(0).unsqueeze(-1), (sigma_phys / S).unsqueeze(0), reduction="mean")
+            Gx_s, Gy_s, mask = ops_i
+            phys = _sparse_div_residual(out, Gx_s, Gy_s, mask)  # scaled units (physical/S^2)
+        else:  # B: scalar output + Laplacian smoothness
+            pred = normalizer.decode(out)            # (N,1) physical
+            data_loss = relative_l2(pred.unsqueeze(0), sigma_phys.unsqueeze(0), reduction="mean")
+            L_s = ops_i
+            Lf = torch.sparse.mm(L_s, out[:, :1])    # operate on the (normalized) scalar field
+            phys = (Lf ** 2).mean()
+        return data_loss, phys
+
+    @torch.no_grad()
+    def evaluate() -> Tuple[float, float]:
+        model.eval()
+        tot_data, tot_phys = 0.0, 0.0
+        for i in range(ntest):
+            out = model(test_coords[i:i + 1], None)[0]   # (N, C)
+            dloss, phys = data_and_phys(out, test_sigma[i], test_ops[i])
+            tot_data += dloss.item()
+            tot_phys += phys.item()
+        return tot_data / ntest, tot_phys / ntest
+
+    t0 = time.time()
+    best_rel = float("inf")
+    test_rel = float("nan")
+    test_phys = float("nan")
+    history = []
+    for ep in range(epochs):
+        model.train()
+        perm = torch.randperm(n_train, generator=gen).tolist()
+        run_data = 0.0
+        for i in perm:
+            optimizer.zero_grad()
+            out = model(train_coords[i:i + 1], None)[0]   # (N, C)
+            dloss, phys = data_and_phys(out, train_sigma[i], train_ops[i])
+            loss = dloss + lam * phys
+            loss.backward()
+            if max_grad_norm is not None:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
+            optimizer.step()
+            run_data += dloss.item()
+        scheduler.step()
+        train_rel = run_data / n_train
+
+        if (ep % eval_every == 0) or (ep >= epochs - 5):
+            test_rel, test_phys = evaluate()
+            best_rel = min(best_rel, test_rel)
+        history.append({"epoch": ep, "train_rel": train_rel, "test_rel": test_rel, "test_phys": test_phys})
+        if ep % log_every == 0 or ep == epochs - 1:
+            print(f"[eqreg seed {seed}] epoch {ep:4d} train_rel={train_rel:.5f} "
+                  f"test_rel={test_rel:.5f} test_resid={test_phys:.4e}", flush=True)
+
+    wall = time.time() - t0
+    rate = MODAL_RATES_PER_SEC.get(gpu_name, 0.0)
+    # Report the residual in physical units. Approach A computes div on the S-scaled tensor, so the
+    # physical residual is test_phys * S^2; Approach B's Laplacian residual is already in the
+    # (normalized) field's units (S == 1).
+    test_residual_phys = test_phys * (S ** 2) if approach == "A" else test_phys
+    final_metrics = {
+        "test_rel_l2": round(test_rel, 6),
+        "best_test_rel_l2": round(best_rel, 6),
+        "test_residual": test_residual_phys,
+        "test_residual_scaled": test_phys,
+        "scale_S": S,
+        "train_rel_l2": round(train_rel, 6),
+        "n_params": n_params,
+        "epochs": epochs,
+        "approach": approach,
+        "lambda": lam,
+    }
+    if ckpt_path is not None:
+        os.makedirs(os.path.dirname(ckpt_path) or ".", exist_ok=True)
+        torch.save(
+            {"state_dict": model.state_dict(),
+             "normalizer": {"mean": normalizer.mean.detach().cpu(), "std": normalizer.std.detach().cpu()},
+             "scale_S": S, "config": config, "seed": seed, "metrics": final_metrics},
+            ckpt_path,
+        )
+    if results_path is None:
+        results_path = os.path.join("results", f"{config.get('name','eqreg')}_seed{seed}.json")
+    write_run_log(results_path, config, seed, final_metrics, wall, gpu_name, wall * rate,
+                  extra={"history_tail": history[-5:]})
+    print(f"[eqreg seed {seed}] DONE test_rel={test_rel:.6f} resid={test_phys:.4e} "
+          f"wall={wall:.0f}s est_cost=${wall*rate:.4f}", flush=True)
+    return final_metrics

stress_operator/utils/logging.py

@@ -0,0 +1,63 @@
+"""Run logging: every run writes a JSON ``{config, git_commit, seed, final_metrics,
+wall_clock, gpu, est_cost}`` to ``results/`` (master plan §0.2 principle 3)."""
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+from typing import Any, Dict, Optional
+
+
+def git_commit(default: str = "unknown") -> str:
+    try:
+        out = subprocess.check_output(
+            ["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL
+        )
+        return out.decode().strip()
+    except Exception:
+        return default
+
+
+def write_run_log(
+    path: str,
+    config: Dict[str, Any],
+    seed: int,
+    final_metrics: Dict[str, Any],
+    wall_clock_sec: float,
+    gpu: str,
+    est_cost_usd: float,
+    extra: Optional[Dict[str, Any]] = None,
+) -> str:
+    """Write a single run record as JSON and return the path."""
+    record = {
+        "config": config,
+        "git_commit": git_commit(),
+        "seed": seed,
+        "final_metrics": final_metrics,
+        "wall_clock_sec": round(float(wall_clock_sec), 2),
+        "gpu": gpu,
+        "est_cost_usd": round(float(est_cost_usd), 4),
+    }
+    if extra:
+        record.update(extra)
+    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+    with open(path, "w") as f:
+        json.dump(record, f, indent=2)
+    return path
+
+
+def estimate_cost(wall_clock_sec: float, rate_per_sec: float) -> float:
+    """GPU cost estimate = wall-clock seconds x per-second rate (master plan §2.7)."""
+    return wall_clock_sec * rate_per_sec
+
+
+# Modal per-second GPU rates (USD), from master plan §2.7 (modal.com/pricing, June 2026).
+MODAL_RATES_PER_SEC = {
+    "A10": 0.000306,
+    "A100-40GB": 0.000583,
+    "A100-80GB": 0.000694,
+    "H100": 0.001097,
+    "L4": 0.000222,
+    "T4": 0.000164,
+    "CPU": 0.0,
+}

stress_operator/utils/viz.py

@@ -0,0 +1,35 @@
+"""Shared stress-field renderer (used by the Gate-0 sanity plot and the Gradio demo)."""
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+
+
+def render_stress(
+    coords: np.ndarray,
+    sigma: np.ndarray,
+    title: str = "von Mises stress",
+    ax=None,
+    vmin: Optional[float] = None,
+    vmax: Optional[float] = None,
+    cmap: str = "coolwarm",
+):
+    """Scatter the per-node stress over the point cloud. Returns the matplotlib Figure.
+
+    coords: (N, 2) node coordinates; sigma: (N,) per-node stress.
+    """
+    import matplotlib.pyplot as plt
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(5, 5))
+    else:
+        fig = ax.figure
+    sc = ax.scatter(
+        coords[:, 0], coords[:, 1], c=sigma, s=18, cmap=cmap,
+        vmin=vmin, vmax=vmax, edgecolor="w", lw=0.1,
+    )
+    ax.set_aspect("equal")
+    ax.set_title(title)
+    fig.colorbar(sc, ax=ax, shrink=0.8, label="sigma")
+    return fig