Inference-only code: remove training, evaluation and experimental code

.gitignore

@@ -1,15 +1,10 @@
-# Raw and derived data are large. Keep only the small public property tables.
+# Data stays out of the repository except the property tables.
 data/*
 !data/properties/
-!data/properties/*.csv
-# Training artifacts, checkpoints, logs
+!data/properties/*
+# Prediction output and the mesh cache
 runs/
-# Model checkpoints anywhere
-*.pt
-# SDD scratch (progress ledger, briefs, review packages)
-.superpowers/
 # Python
 __pycache__/
 *.pyc
-.ipynb_checkpoints/
 .pytest_cache/

README.md

@@ -166,8 +166,9 @@ Each node gets 15 features computed from the mesh and its surface tags:
 - part volume, surface area and the fraction of surface that cools;
 - local modulus, the volume-to-surface ratio inside geodesic balls of 5, 15 and 40 mm, which is
   a local version of Chvorinov's modulus;
-- two boundary-condition channels from training, which are zero in this release because the
-  whole surface cools into the mold.
+- two boundary-condition channels: the mold effusivity on surface nodes (zero inside the
+  part), and a boundary heat-input channel, which is zero because a sand-mold surface adds no
+  heat.
 
 The alloy and mold enter as five dimensionless numbers: Stefan number, normalized freezing
 range, thermal diffusivity ratio, normalized superheat and normalized mold effusivity. They
@@ -216,13 +217,15 @@ spread is not monotone in `a`, so the solver brackets the smallest root explicit
 | `weights/ood_ranges.json` | <1 KB | Training envelope (part volume, size, distance to cooling) for out-of-distribution warnings |
 | `weights/SHA256SUMS` | <1 KB | Checksums |
 | `src/model_api.py` | | `CastSolidModel`, the mesh-in, field-out interface shown above |
+| `src/features.py`, `src/model_gnn.py`, `src/gbm_heads.py`, `src/hybrid.py` | | The model: node features, graph network, gradient-boosted heads, recomposition |
 | `data/properties/` | | Alloy and mold tables ([README](data/properties/README.md)) |
 
 None of the weight files can execute code when loaded. The GNN is stored as
 [safetensors](https://github.com/huggingface/safetensors) and the heads as
 [skops](https://skops.readthedocs.io/); before loading a skops file, `src/release_io.py` checks
-every type it names against an allowlist of scikit-learn and NumPy types. The rest of `src/`,
-`tools/` and `tests/` is the supporting code, described on GitHub.
+every type it names against an allowlist of scikit-learn and NumPy types. The rest of `src/`
+(meshing a STEP file, shrinkage screening, result export, a browser viewer) and `tests/` are the
+supporting pipeline, described on GitHub. Training code is not part of this release.
 
 ## Training data
 
@@ -251,6 +254,13 @@ Training meshes have roughly 7,500 to 10,000 nodes. The target at each node is t
 temperature drops through the solidus. Material properties come from handbook data. The raw
 simulation corpus (about 100 GB) is not distributed.
 
+One data issue shaped this project more than any model choice. The solver stored temperatures
+in dolfinx degree-of-freedom order, while coordinates and geometric features were in gmsh node
+order, so about 60% of the nodal targets were silently attached to the wrong positions. Every
+early model failure traced back to it. The training set was rebuilt with a recovered
+permutation between the two orderings. If you train on solver output, check node identity across
+every solver and mesh interface before tuning a model.
+
 ## Training procedure
 
 v5.1 came out of a chain of fine-tunes, not a single training run.
@@ -327,7 +337,8 @@ every result, and send high-consequence cases to a calibrated simulation and a p
 
 ## Reproducibility
 
-The published files were converted from the internal checkpoints by `tools/export_release.py`:
+The published files were converted from the internal training checkpoints and checked against
+them:
 
 | Release file | Internal artifact | Check |
 |---|---|---|
@@ -335,8 +346,9 @@ The published files were converted from the internal checkpoints by `tools/expor
 | `gbm_total_v3.skops` | `hgb_v3.joblib` | Identical predictions on 10,002 probe rows spanning each feature's training split range |
 | `gbm_std_v2.skops` | `hgb_std_v2.joblib` | Identical predictions on 10,002 probe rows spanning each feature's training split range |
 
-On a real test mesh, the published files and the internal checkpoints give the same field. The
-heads were trained with scikit-learn 1.5.2, and other versions may refuse to load them or change
+The code in this repository was also run side by side with the internal pipeline on five
+production parts (two of them with both alloys), plus a fixed-size mesh with a custom property
+set: every output field matched exactly. The heads were trained with scikit-learn 1.5.2, and other versions may refuse to load them or change
 their output. Other versions used: Python 3.11, PyTorch 2.3.0, PyTorch Geometric 2.7.0,
 NumPy 2.x.
 

requirements.txt

@@ -1,11 +1,10 @@
-# Install PyTorch 2.3 for your platform first (see README), plus the matching
-# torch-cluster wheel, e.g. for CPU:
+# Install PyTorch 2.3 for your platform first, plus the matching torch-cluster
+# wheel, for example on CPU:
 #   pip install "torch==2.3.0" --index-url https://download.pytorch.org/whl/cpu
 #   pip install torch-cluster -f https://data.pyg.org/whl/torch-2.3.0+cpu.html
 numpy>=2.0
 scipy>=1.11
 scikit-learn==1.5.2
-joblib>=1.3
 skops>=0.10
 safetensors>=0.4
 torch-geometric==2.7.0

src/baseline.py

@@ -1,63 +0,0 @@
-"""Per-node regression baselines: gradient boosting (primary) and a small MLP.
-
-Both predict standardized log1p(solidification_time) from standardized
-physics features. predict_seconds inverts the transform back to seconds.
-"""
-import numpy as np
-import torch
-import torch.nn as nn
-from sklearn.ensemble import HistGradientBoostingRegressor
-
-
-def train_hgb(X_std, y_log, max_iter: int = 300):
-    model = HistGradientBoostingRegressor(
-        max_iter=max_iter, learning_rate=0.1, max_leaf_nodes=63,
-        l2_regularization=1.0, early_stopping=True, validation_fraction=0.1,
-        random_state=0)
-    model.fit(X_std, y_log)          # trained directly on log1p(y)
-    return model
-
-
-class _MLP(nn.Module):
-    def __init__(self, in_dim):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(in_dim, 128), nn.ReLU(),
-            nn.Linear(128, 64), nn.ReLU(),
-            nn.Linear(64, 1))
-
-    def forward(self, x):
-        return self.net(x).squeeze(-1)
-
-
-def train_mlp(X_std, y_log, epochs: int = 30, device: str = "cpu", batch: int = 8192):
-    torch.manual_seed(0)
-    X = torch.tensor(np.asarray(X_std), dtype=torch.float32, device=device)
-    y = torch.tensor(np.asarray(y_log), dtype=torch.float32, device=device)
-    model = _MLP(X.shape[1]).to(device)
-    opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-5)
-    loss_fn = nn.SmoothL1Loss()
-    n = X.shape[0]
-    for _ in range(epochs):
-        perm = torch.randperm(n, device=device)
-        for i in range(0, n, batch):
-            idx = perm[i:i + batch]
-            opt.zero_grad()
-            loss = loss_fn(model(X[idx]), y[idx])
-            loss.backward()
-            opt.step()
-    model.eval()
-    return model
-
-
-def predict_seconds(model, X_std, norm) -> np.ndarray:
-    if isinstance(model, nn.Module):
-        dev = next(model.parameters()).device
-        with torch.no_grad():
-            pred_log = model(
-                torch.tensor(np.asarray(X_std), dtype=torch.float32, device=dev)
-            ).cpu().numpy()
-    else:
-        pred_log = model.predict(X_std)
-    # Both models predict raw log1p(seconds); invert uniformly.
-    return np.clip(np.expm1(pred_log), 0.0, None)

src/build_dataset.py

@@ -1,1199 +0,0 @@
-"""Reprocess raw .pt + paired .msh into cached per-node feature files.
-
-Writes one .npz per casting under <data_root>/features/graphs/ plus a global
-meta.npz and norm_stats.npz. Parallelised across CPU cores.
-"""
-import os
-# CPU-hygiene: cap every BLAS backend to 1 thread PER WORKER PROCESS, set
-# before numpy (or anything that transitively imports it) is first imported
-# in this process. This must run before `import numpy as np` below to take
-# effect for numpy's own BLAS init; ProcessPoolExecutor workers either
-# re-import this module from scratch ("spawn", re-running this) or fork from
-# a parent that already applied it ("fork", inherited via copy-on-write) --
-# mirrors tools/doe/doe.py's identical top-of-file pattern.
-for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS",
-           "MKL_NUM_THREADS", "NUMEXPR_NUM_THREADS"):
-    os.environ.setdefault(_v, "1")
-
-import csv
-import itertools
-import json
-import os.path as osp
-import re
-import shutil
-import sys
-from glob import glob
-
-import meshio
-import numpy as np
-import torch
-from scipy.spatial import cKDTree
-
-sys.path.insert(0, osp.dirname(osp.dirname(__file__)))
-from src.msh_parse import parse_msh
-from src.features import (compute_features, FEATURE_NAMES, FEATURE_NAMES_V3,
-                          FEATURE_NAMES_V4, FEATURE_NAMES_V5, G_VECTOR_NAMES,
-                          _node_features, alloy_g_vector, chvorinov_t_ref)
-from src.geometry import mesh_volume, total_area
-
-_FACE_RE = re.compile(r'_[+-][XYZ]')
-_MSH_SETS = ("msh_none", "msh_half", "msh_all")
-
-# Canonical raw temperature backup (dolfinx-DOF-order snapshots_nodal), on the
-# large Media drive. Used when data/pt/*.pt has been deleted to free space.
-# The .pt temps are byte-identical to these npz (verified). Override if moved.
-RAW_NPZ_DIRS = ("data/raw",)
-
-
-def find_msh(stem: str, data_root: str = "data"):
-    for s in _MSH_SETS:
-        hits = glob(osp.join(data_root, s, "*", f"{stem}.msh"))
-        if hits:
-            return hits[0]
-    return None
-
-
-def _load_raw_temps(stem: str, data_root: str = "data"):
-    """Return (temps_dof (T,N) float32, times (T,) float32) or None.
-
-    Prefers data/pt/<stem>.pt; falls back to the canonical raw npz backup so the
-    pipeline still rebuilds after data/pt is deleted."""
-    ptp = osp.join(data_root, "pt", f"{stem}.pt")
-    if osp.exists(ptp):
-        raw = torch.load(ptp, weights_only=False)
-        return raw["temps"].numpy().astype(np.float32), raw["times"].numpy().astype(np.float32)
-    for d in RAW_NPZ_DIRS:
-        hits = glob(osp.join(d, "*", f"{stem}.npz"))
-        if hits:
-            z = np.load(hits[0])
-            return z["snapshots_nodal"].astype(np.float32), z["times"].astype(np.float32)
-    return None
-
-
-def process_one(stem: str, data_root: str = "data", feature_version: int = 2):
-    mpath = find_msh(stem, data_root)
-    if mpath is None:
-        print(f"  WARN missing msh for {stem}")
-        return None
-    src = _load_raw_temps(stem, data_root)
-    if src is None:
-        print(f"  WARN no raw temps for {stem} (.pt or npz backup)")
-        return None
-    temps_dof, times = src
-    mesh = parse_msh(mpath)
-    n_pt = temps_dof.shape[1]
-    if mesh.node_xyz.shape[0] != n_pt:
-        print(f"  WARN node count mismatch {stem}: msh={mesh.node_xyz.shape[0]} raw={n_pt}")
-        return None
-
-    # Apply the node-ordering recovery permutation: the raw temps are in dolfinx
-    # DOF order, which is scrambled relative to the gmsh node order used by pos/
-    # mesh labels. perm g maps DOF->gmsh node: temps_gmsh[:, g] = temps_dof.
-    # Files without a validated perm are excluded (their per-node target is
-    # unrecoverable). See docs spec §1.5.
-    perm_path = osp.join(data_root, "features", "perms", f"{stem}.npy")
-    if not osp.exists(perm_path):
-        print(f"  WARN no recovery perm for {stem} (excluded)")
-        return None
-    g = np.load(perm_path)
-    if g.shape[0] != n_pt:
-        print(f"  WARN perm/temps size mismatch {stem}: perm={g.shape[0]} temps_N={n_pt}")
-        return None
-    temps = np.empty_like(temps_dof)
-    temps[:, g] = temps_dof                 # reorder into gmsh node order
-
-    feat = compute_features(mesh, temps, times, version=feature_version)
-
-    n_adiab = len(_FACE_RE.findall(stem))
-    prefix = stem.split("_")[0]
-    return {
-        "stem": stem,
-        "X": feat["X"],
-        "y": feat["y"],
-        "pos": feat["pos"],
-        "mesh_edges": feat["mesh_edges"],
-        "y_global": feat["y_global"],
-        "prefix": prefix,
-        "n_adiab": n_adiab,
-        "is_fully_cooled": n_adiab == 0,
-        "n_nodes": n_pt,
-    }
-
-
-import argparse
-from concurrent.futures import ProcessPoolExecutor
-
-
-def _init_worker():
-    """ProcessPoolExecutor initializer: lower scheduling priority + force
-    every BLAS backend down to 1 thread at RUNTIME via threadpoolctl. The
-    module-top env vars already cover a fresh/"spawn" worker's own numpy
-    import; this is the belt-and-braces fallback for a "fork" worker that
-    inherited an ALREADY-initialized BLAS library from the parent (in which
-    case env vars set post-fork would be too late) -- mirrors
-    tools/doe/doe.py's `_init_worker`."""
-    try:
-        os.nice(10)
-    except Exception:
-        pass
-    try:
-        import threadpoolctl
-        threadpoolctl.threadpool_limits(1)
-    except Exception:
-        pass
-
-
-def write_record(rec, out_dir):
-    gdir = osp.join(out_dir, "graphs")
-    os.makedirs(gdir, exist_ok=True)
-    np.savez(osp.join(gdir, f"{rec['stem']}.npz"), X=rec["X"], y=rec["y"], pos=rec["pos"], mesh_edges=rec["mesh_edges"])
-
-
-def _all_stems(data_root):
-    """Enumerate casting stems. Prefer data/pt; after it is deleted, fall back to
-    the recovery perms (the exact recoverable set), then the raw npz backup."""
-    stems = [osp.splitext(osp.basename(p))[0]
-             for p in glob(osp.join(data_root, "pt", "*.pt"))]
-    if not stems:
-        stems = [osp.splitext(osp.basename(p))[0]
-                 for p in glob(osp.join(data_root, "features", "perms", "*.npy"))]
-    if not stems:
-        for d in RAW_NPZ_DIRS:
-            stems = [osp.splitext(osp.basename(p))[0]
-                     for p in glob(osp.join(d, "*", "*.npz"))]
-            if stems:
-                break
-    return sorted(stems)
-
-
-# ============================================================================
-# feature_version == 4: join raw FEniCSx solves (`data/raw_v4/<run_id>.npz`,
-# {t,gradT,dTdt} already in gmsh node order -- see tools/doe/adapters.py's
-# map_to_gmsh) with the DOE manifest (per-run spec: material, mold/chill/
-# sleeve BC, hold_out). One graph per manifest row; `run_id` is the cache key
-# ("stem").
-#
-# IMPORTANT interface note: `src.features.compute_features(mesh, temps, times,
-# version=4, spec=...)` (Task 3) still expects raw (T,N) temperature snapshots
-# + times, and internally derives y via `solidification_time` -- that is the
-# OLD raw-pipeline contract (v2/v3). The FT2 raw solves are NOT snapshot
-# histories: `fenicsx_solver.py` already outputs the per-node crossing time
-# `t` directly (see its module docstring). So this v4 path does NOT call
-# `compute_features`; it composes the same lower-level building blocks
-# (`_node_features`, `alloy_g_vector`, `chvorinov_t_ref`) directly, using the
-# solver-provided `t` array as `y` with no re-derivation.
-# ============================================================================
-def _read_manifest_rows(manifest_path: str) -> list:
-    with open(manifest_path, newline="") as f:
-        return list(csv.DictReader(f))
-
-
-def _archive_mesh_work_dir(data_root: str) -> str:
-    d = osp.join(data_root, "mesh_work_v4")
-    os.makedirs(d, exist_ok=True)
-    return d
-
-
-def _find_msh_v4(row: dict, data_root: str = "data"):
-    """Resolve a manifest row's geometry to a .msh path.
-    - "path": geometry_id verbatim (a full .msh path).
-    - "primitive"/"scaled": reuse the existing 42-primitive lookup
-      (`find_msh`) by geometry_id.
-    - "archive": geometry_id is a STEP file path; mesh it via
-      `src.step_mesh.mesh_step` (same tool `tools.doe.adapters.mesh()` uses
-      for dataset D) into `<data_root>/mesh_work_v4/<run_id>.msh`, cached
-      across calls (the pre-pass Chvorinov-C fit and the main pass would
-      otherwise mesh the same archive part twice). Returns None (with a
-      printed WARN) if the STEP import/meshing itself fails -- a bad archive
-      part must WARN+skip its row, not crash the batch.
-    `scale_factor` (if != 1.0) is applied to the resolved mesh's node
-    coordinates by the caller, not here.
-
-    ft2a feature-BC fix: a non-empty `mesh_variant` column (written by
-    `tools/doe/augment_manifest_bc.py` for chill/sleeve rows) is an explicit
-    face-tagged msh_half path that OVERRIDES every other resolution branch
-    below -- mirrors `tools.doe.adapters.mesh()`'s identical check, which the
-    SOLVER side already uses. msh_none (what `find_msh` would otherwise
-    return) has no "Adiabatic" (BC-patch) facet tag, so without this override
-    the chill/sleeve/exo BC-dependent feature channels (is_adiabatic,
-    boundary_extraction, exothermic_input) would be computed as if the row
-    were plain sand, even though the row's SOLVER target used the tagged
-    msh_half mesh -- silently dropping the BC input signal from X. msh_none
-    and its msh_half variant share the same node order/coordinates (verified
-    -- only the facet tagging differs), so this override changes only the
-    tag-derived feature channels, not the geometric features or the
-    per-node alignment with the (msh_half-based) solver targets.
-
-    Task 3 (ft3-assembly-finetune): dataset=="E" rows (assembly/integration)
-    have mesh_variant as a literal direct path. For consistency with FT2a's
-    batch-resilience pattern, missing mesh_variant is WARN+skip (row-level
-    failure, not batch-killing), same as non-E rows."""
-    mesh_variant = (row.get("mesh_variant") or "").strip()
-    if mesh_variant:
-        if not osp.exists(mesh_variant):
-            print(f"  WARN mesh_variant {mesh_variant!r} (run_id="
-                  f"{row.get('run_id')!r}) does not exist -- stale path?")
-            return None
-        return mesh_variant
-    source = row.get("geometry_source") or "primitive"
-    gid = row["geometry_id"]
-    if source == "path":
-        return gid if osp.exists(gid) else None
-    if source == "archive":
-        out_msh = osp.join(_archive_mesh_work_dir(data_root),
-                          f"{row.get('run_id', 'archive')}.msh")
-        if osp.exists(out_msh):
-            return out_msh
-        try:
-            from src.step_mesh import mesh_step
-            mesh_step(gid, out_msh)
-        except Exception as e:  # noqa
-            print(f"  WARN mesh_step failed for run_id={row.get('run_id')} "
-                  f"(STEP={gid!r}): {e}")
-            return None
-        return out_msh if osp.exists(out_msh) else None
-    return find_msh(gid, data_root)
-
-
-def _resolve_material_props_v4(material_id: str, manifest_dir: str) -> dict:
-    """material_id -> {k, rho, cp, L_kJ_kg, TL, TS, Tpour}, mirroring
-    `tools.doe.adapters._resolve_material`'s 3-step fallback: the
-    owner-verified CSV, then `materials_synthetic.csv` next to the manifest
-    (SYN_* ids), then the 12 named anchor grades."""
-    from tools.doe.load_properties import load_materials
-    for m in load_materials():
-        if m.material_id == material_id:
-            return dict(k=m.k, rho=m.rho, cp=m.cp, L_kJ_kg=m.L,
-                       TL=m.TL, TS=m.TS, Tpour=m.Tpour)
-
-    synth_csv = osp.join(manifest_dir, "materials_synthetic.csv")
-    if osp.exists(synth_csv):
-        with open(synth_csv, newline="") as f:
-            for row in csv.DictReader(f):
-                if row.get("material_id") == material_id:
-                    return dict(
-                        k=float(row["k_W_mK"]), rho=float(row["rho_kg_m3"]),
-                        cp=float(row["cp_J_kgK"]), L_kJ_kg=float(row["L_kJ_kg"]),
-                        TL=float(row["TL_C"]), TS=float(row["TS_C"]),
-                        Tpour=float(row["Tpour_C"]))
-
-    from tools.doe.doe import _anchor_materials
-    for m in _anchor_materials():
-        if m.material_id == material_id:
-            return dict(k=m.k, rho=m.rho, cp=m.cp, L_kJ_kg=m.L,
-                       TL=m.TL, TS=m.TS, Tpour=m.Tpour)
-
-    raise KeyError(
-        f"material_id {material_id!r} not resolvable via load_properties, "
-        f"{synth_csv}, or tools.doe.doe._anchor_materials()")
-
-
-def _resolve_chill_effusivity_v4(row: dict) -> float:
-    """0.0 unless bc_config=='chills', in which case
-    `tools.doe.doe.get_chill_effusivities()[chill_material]` -- raises
-    ValueError (never a silent 0.0) if chill_material is empty/unknown."""
-    if (row.get("bc_config") or "none") != "chills":
-        return 0.0
-    from tools.doe.doe import get_chill_effusivities
-    chill_material = (row.get("chill_material") or "").strip()
-    effs = get_chill_effusivities()
-    if not chill_material or chill_material not in effs:
-        raise ValueError(
-            f"run_id={row.get('run_id')!r}: bc_config=='chills' but "
-            f"chill_material={chill_material!r} not in get_chill_effusivities() "
-            f"(known: {sorted(effs)}) -- refusing to silently use 0.0")
-    return float(effs[chill_material])
-
-
-def _resolve_exo_q0_v4(row: dict,
-                       exo_calibration_csv: str = "data/properties/exo_calibration.csv",
-                       tol: float = 0.05) -> float:
-    """0.0 unless bc_config=='exo_sleeve', in which case the nearest-MEF
-    `q0_W_m2` row of `exo_calibration_csv` (the sweep output of
-    `tools/doe/calibrate_exo.py`), resolved via the shared
-    `tools.doe.load_properties.resolve_exo_q0_by_mef` lookup (also used by
-    `src.assembly.validate_bodies` for Part B) -- raises ValueError (never a
-    silent 0.0) if the calibration table is missing or has no row within
-    `tol` of the row's sleeve_MEF, mirroring
-    `tools.doe.adapters._resolve_exo_params`. The ValueError message is
-    re-wrapped here with this row's run_id for context."""
-    if (row.get("bc_config") or "none") != "exo_sleeve":
-        return 0.0
-    from tools.doe.load_properties import resolve_exo_q0_by_mef
-    mef = float(row.get("sleeve_MEF") or 0.0)
-    try:
-        return resolve_exo_q0_by_mef(mef, exo_calibration_csv, tol)
-    except ValueError as e:
-        raise ValueError(f"run_id={row.get('run_id')!r}: {e}") from e
-
-
-# --------------------------------------------------------------------- #
-# I1 (ft2a-final-review): archive/DF feature-mesh coordinate correspondence.
-#
-# For archive rows the FEATURE builder (`_find_msh_v4`, above) re-meshes the
-# row's STEP file via `mesh_step` INDEPENDENTLY of the SOLVER side, which
-# also calls `mesh_step` (via `tools.doe.adapters.mesh()`) but into a
-# DIFFERENT work directory (`<doe_dir>/mesh_work/` vs this module's
-# `<data_root>/mesh_work_v4/`) -- two separate meshing invocations of the
-# same STEP file. Until now only the resulting node COUNT was checked to
-# match the raw npz's array length (see the `t.shape[0] != n_pt` check in
-# `_process_row_stage1`); node ORDER/coordinates were never verified. The
-# solver's `t`/`gradT`/`dTdt`/`t_liquidus` raw_v4 arrays are in the node
-# order of THAT solve-time mesh (`fenicsx_solver._dof_to_gmsh_perm`
-# coordinate-matches its dolfinx DOFs onto it, at ~1e-9 m tolerance); if a
-# re-mesh of the identical STEP ever produced the same node count in a
-# DIFFERENT order (a gmsh non-determinism this repo has not observed but
-# cannot rule out -- different gmsh version, different meshing-algorithm
-# code path, etc.), targets and features would silently misalign for that
-# row with no other symptom. This mirrors the strict bijection
-# `tools.doe.adapters.map_to_gmsh` already enforces on the solver side, on
-# the feature side.
-# --------------------------------------------------------------------- #
-_ARCHIVE_COORD_TOL_MM = 1e-6
-
-
-def _archive_reference_mesh_path(row: dict, manifest_dir: str) -> str:
-    """The ORIGINAL solve-time .msh for an archive/DF row, as cached by
-    `tools.doe.adapters.mesh()` under `<doe_dir>/mesh_work/<run_id>[_scaled]
-    .msh` (`adapters._work_dir()`'s naming convention). `manifest_dir`
-    (where manifest.csv lives) IS `<doe_dir>` in the standard layout (both
-    default to "doe_out"), so no separate config knob is needed."""
-    run_id = row.get("run_id", "<unknown>")
-    scale = float(row.get("scale_factor") or 1.0)
-    suffix = "_scaled" if scale != 1.0 else ""
-    return osp.join(manifest_dir, "mesh_work", f"{run_id}{suffix}.msh")
-
-
-def _assert_archive_mesh_coord_correspondence(mesh, row: dict, manifest_dir: str) -> None:
-    """No-op unless `row["geometry_source"] == "archive"` (msh_none/msh_half
-    primitive rows already share node order with the raw npz by construction
-    -- verified -- so this cost is deliberately not paid there).
-
-    For archive rows, raises a loud `ValueError` naming the run_id --
-    NEVER silently proceeds -- when:
-      - the reference solve-time mesh can't be found (an archive row with a
-        raw_v4 solve implies `tools.doe.adapters.mesh()` ran and cached one;
-        its absence means this check literally cannot be performed, which is
-        treated as a failure, not a silent pass);
-      - node counts differ; or
-      - node coordinates fail to correspond AT THE SAME INDEX, within
-        `_ARCHIVE_COORD_TOL_MM`.
-    The last check is deliberately NOT a plain nearest-neighbour SET
-    membership test (which a pure reordering of the identical point set
-    would pass): it uses a KDTree to find, for every feature-mesh node, its
-    nearest reference-mesh node (mirroring `fenicsx_solver.
-    _dof_to_gmsh_perm`'s own coordinate-matching approach) but then requires
-    BOTH that the match be within tolerance AND that the matched index equal
-    the feature node's own index (identity correspondence) -- since the two
-    meshes are expected to already share node order (same deterministic
-    `mesh_step` call on the same STEP), a same-count/same-position-set but
-    DIFFERENT-order mesh (exactly the silent-misalignment risk I1 flags)
-    must fail this check, not pass it."""
-    if (row.get("geometry_source") or "primitive") != "archive":
-        return
-    run_id = row.get("run_id", "<unknown>")
-    ref_path = _archive_reference_mesh_path(row, manifest_dir)
-    if not osp.exists(ref_path):
-        raise ValueError(
-            f"run_id={run_id!r}: archive-row coordinate-correspondence check "
-            f"cannot find the original solve-time mesh at {ref_path!r} (the "
-            f"one tools.doe.adapters.mesh() caches for the FEniCSx solve) -- "
-            "refusing to trust the feature builder's independently-remeshed "
-            "STEP without verifying it against it")
-
-    ref_pts = meshio.read(ref_path).points[:, :3].astype(np.float64)
-    feat_pts = np.asarray(mesh.node_xyz, dtype=np.float64)[:, :3]
-    if ref_pts.shape[0] != feat_pts.shape[0]:
-        raise ValueError(
-            f"run_id={run_id!r}: archive feature-mesh node count "
-            f"({feat_pts.shape[0]}) != solve-time mesh node count "
-            f"({ref_pts.shape[0]}) at {ref_path!r} -- a gmsh remesh of the "
-            "same STEP produced a DIFFERENT mesh; targets<->features cannot "
-            "be trusted to align for this row")
-
-    dist, idx = cKDTree(ref_pts).query(feat_pts)
-    max_dist = float(dist.max())
-    if max_dist > _ARCHIVE_COORD_TOL_MM:
-        raise ValueError(
-            f"run_id={run_id!r}: archive feature-mesh node coordinates do "
-            f"NOT correspond to the solve-time mesh at {ref_path!r} within "
-            f"{_ARCHIVE_COORD_TOL_MM} mm (max nearest-neighbour distance = "
-            f"{max_dist:.6g} mm) -- a gmsh remesh of the same STEP produced "
-            "a mesh whose node positions don't match; targets<->features "
-            "would silently misalign for this row")
-    if not np.array_equal(idx, np.arange(feat_pts.shape[0])):
-        n_reordered = int((idx != np.arange(feat_pts.shape[0])).sum())
-        raise ValueError(
-            f"run_id={run_id!r}: archive feature-mesh node ORDER does not "
-            f"match the solve-time mesh at {ref_path!r} -- {n_reordered} "
-            "node(s) matched a DIFFERENT index by nearest-neighbour "
-            "coordinate (same node positions, different order) -- a gmsh "
-            "remesh of the same STEP produced the same node count/positions "
-            "but a DIFFERENT order; targets<->features would silently "
-            "misalign for this row")
-
-
-def _load_mesh_for_row(row: dict, data_root: str):
-    """Resolve + parse a manifest row's mesh, applying scale_factor (if any)
-    to node coordinates. Returns None (with a printed warning) if the .msh
-    can't be found."""
-    mpath = _find_msh_v4(row, data_root)
-    if mpath is None:
-        print(f"  WARN missing msh for run_id={row.get('run_id')} "
-              f"(geometry_id={row.get('geometry_id')!r})")
-        return None
-    mesh = parse_msh(mpath)
-    scale = float(row.get("scale_factor") or 1.0)
-    if scale != 1.0:
-        mesh.node_xyz = mesh.node_xyz * scale
-    return mesh
-
-
-def _fit_chvorinov_c_per_family(rows: list, raw_dir: str, data_root: str) -> dict:
-    """Per alloy_family, a fixed-exponent (Chvorinov exponent==2) Chvorinov
-    constant C fit from the generated (t_total, V/A) pairs: log(C) =
-    mean(log(t_total) - 2*log(M)) -- the geometric mean of the per-row ratio
-    t_total/M^2. This is the "underlying fit" `gate_b.fit_chvorinov` performs
-    a free two-parameter (slope+intercept) regression of, specialized to a
-    FIXED slope of 2 (the assumed t=C*M^2 model): more robust than the free
-    regression when a family has few/degenerate (M, t) points (e.g. a single
-    geometry solved once -- `gate_b.fit_chvorinov` requires >=3 non-degenerate
-    points and would raise), which is expected for early/small manifests."""
-    # M-4 (ft2a-branch-review): dataset=="S" (sanity cross-check pairs, always
-    # a duplicate of an existing A/B/C row) rows are EXCLUDED here -- they
-    # must not double-weight their base row's (M, t_total) point in the
-    # per-family geometric-mean fit, mirroring `_build_splits_v4`'s existing
-    # S exclusion from train/test.
-    fam_points = {}
-    for row in rows:
-        if row.get("dataset") == "S":
-            continue
-        run_id = row["run_id"]
-        mesh = _load_mesh_for_row(row, data_root)
-        if mesh is None:
-            continue
-        raw_path = osp.join(raw_dir, f"{run_id}.npz")
-        if not osp.exists(raw_path):
-            continue
-        with np.load(raw_path) as z:
-            if "t" not in z.files:
-                continue
-            t_total = float(np.max(z["t"]))
-        cool_tri = mesh.tris.get("Cooling", np.empty((0, 3), np.int64))
-        adiab_tri = mesh.tris.get("Adiabatic", np.empty((0, 3), np.int64))
-        area = total_area(mesh.node_xyz, cool_tri) + total_area(mesh.node_xyz, adiab_tri)
-        vol = mesh_volume(mesh.node_xyz, mesh.tets)
-        if area <= 0 or vol <= 0 or t_total <= 0:
-            continue
-        M = vol / area
-        fam = row.get("alloy_family") or "unknown"
-        fam_points.setdefault(fam, []).append((M, t_total))
-
-    return _fit_c_from_points(fam_points)
-
-
-def _fit_c_from_points(fam_points: dict) -> dict:
-    """Shared core of the per-family Chvorinov-C geometric-mean fit: given
-    {alloy_family: [(M, t_total), ...]} (M = volume/area), returns
-    {alloy_family: C}. Factored out so `_fit_chvorinov_c_per_family`
-    (mesh-parsing, standalone/tested directly) and `build_v4`'s parallel
-    stage-1 pipeline (which already has volume/area/t_total from stage 1 and
-    must NOT re-parse the mesh to get them -- see M-3) can never silently
-    diverge in the actual fit math."""
-    C_by_family = {}
-    for fam, pts in fam_points.items():
-        M = np.array([p[0] for p in pts], dtype=np.float64)
-        T = np.array([p[1] for p in pts], dtype=np.float64)
-        C_by_family[fam] = float(np.exp(np.mean(np.log(T) - 2.0 * np.log(M))))
-    return C_by_family
-
-
-# ----------------------------------------------------------------------- #
-# M-3 (ft2a-branch-review): build_v4 used to parse each row's .msh TWICE --
-# once inside `_fit_chvorinov_c_per_family`'s pre-pass (for volume/area, to
-# compute M=V/A) and again inside `process_one_v4`'s main pass (for X/pos/
-# mesh_edges) -- and did both passes SERIALLY, row by row, which is the slow
-# path for the ~4725-row FT2 manifest. Fixed by splitting the per-row work
-# into two stages:
-#   stage 1 (`_process_row_stage1`, embarrassingly parallel, one mesh parse
-#     per row): resolve+parse the mesh ONCE, load the raw solve, resolve the
-#     BC/material spec, and compute X/pos/mesh_edges/volume/area/t_total.
-#     None of this needs the cross-row Chvorinov C fit -- `_node_features`
-#     never reads `spec["C_material"]` (only `chvorinov_t_ref` does, in stage
-#     2), so X can be computed here without a fitted C yet.
-#   (serial aggregate, cheap, no mesh reparse): fit C_by_family straight from
-#     stage 1's (volume, area, t_total) triples via `_fit_c_from_points` --
-#     the same fit math `_fit_chvorinov_c_per_family` uses, so the two paths
-#     can never silently diverge.
-#   stage 2 (`_finalize_row_stage2`, pure arithmetic, no mesh/IO): plugs the
-#     fitted C into t_ref/y_nondim/g and assembles the final record.
-# `build_v4` calls the two stages directly (stage 1 parallelised across the
-# whole manifest, once) instead of going through `process_one_v4`, which
-# remains as a back-compat single-row wrapper composing both stages in one
-# call for anyone calling it directly.
-# ----------------------------------------------------------------------- #
-def _process_row_stage1(row: dict, manifest_dir: str, raw_dir: str, data_root: str,
-                        exo_calibration_csv: str = "data/properties/exo_calibration.csv",
-                        feature_version: int = 4):
-    """Parallel-safe, picklable: parses the mesh exactly once. Returns a dict
-    of picklable (numpy/scalar) results for stage 2, or None (+ printed WARN)
-    if this row must be skipped -- every skip condition here is byte-for-byte
-    the same check `process_one_v4` used to make (missing mesh/raw/t, node-
-    count mismatch, unknown material_id, unresolvable chill/exo BC props).
-
-    `feature_version` (Task 2, bc-reach): threaded straight into
-    `_node_features` -- 4 (default, v4's 15 cols) or 5 (v4's 15 cols +
-    Task 1's 5 bc-reach cols, `build_v5`'s 20-wide X). Nothing else in this
-    function depends on the feature version; the BC-spec resolution above
-    is shared verbatim by both."""
-    run_id = row["run_id"]
-    mesh = _load_mesh_for_row(row, data_root)
-    if mesh is None:
-        return None
-
-    raw_path = osp.join(raw_dir, f"{run_id}.npz")
-    if not osp.exists(raw_path):
-        print(f"  WARN no raw_v4 npz for run_id={run_id}: {raw_path}")
-        return None
-    with np.load(raw_path) as z:
-        if "t" not in z.files:
-            print(f"  WARN raw_v4 npz for {run_id} has no 't' field (keys={z.files})")
-            return None
-        t = z["t"].astype(np.float32)
-        gradT = z["gradT"].astype(np.float32) if "gradT" in z.files else np.zeros_like(t)
-        dTdt = z["dTdt"].astype(np.float32) if "dTdt" in z.files else np.zeros_like(t)
-        # t_liquidus (FT2a: liquidus-crossing time, enables t - t_liquidus =
-        # local solidification time without re-simulating) is OPTIONAL --
-        # older raw_v4 npz files predate this field. Tolerate its absence
-        # (or a malformed length) by simply not carrying it forward; never
-        # crash and never fabricate a value (unlike gradT/dTdt, a missing
-        # t_liquidus has no safe zero-fallback -- see _finalize_row_stage2).
-        t_liquidus = None
-        if "t_liquidus" in z.files:
-            _tl = z["t_liquidus"].astype(np.float32)
-            if _tl.shape[0] == t.shape[0]:
-                t_liquidus = _tl
-            else:
-                print(f"  WARN t_liquidus length mismatch run_id={run_id}: "
-                      f"t={t.shape[0]} t_liquidus={_tl.shape[0]} -- skipping "
-                      "t_liquidus/local_solidification_time for this row")
-
-    n_pt = mesh.node_xyz.shape[0]
-    if t.shape[0] != n_pt:
-        print(f"  WARN node count mismatch run_id={run_id}: msh={n_pt} raw={t.shape[0]}")
-        return None
-
-    # I1 (ft2a-final-review): archive/DF rows only -- see the function's
-    # docstring above. Deliberately NOT wrapped in try/except: a mismatch (or
-    # a missing reference mesh) here is a hard, batch-aborting failure, not a
-    # row-level WARN+skip like every other check in this function -- it
-    # signals a potential systemic feature<->target misalignment, not a bad
-    # individual row.
-    _assert_archive_mesh_coord_correspondence(mesh, row, manifest_dir)
-
-    try:
-        spec = _resolve_material_props_v4(row["material_id"], manifest_dir)
-    except KeyError as e:
-        print(f"  WARN {e}")
-        return None
-
-    bc_config = row.get("bc_config") or "none"
-    spec["mold_effusivity"] = float(row["mold_effusivity"])
-    spec["bc_config"] = bc_config
-    spec["sleeve_MEF"] = float(row.get("sleeve_MEF") or 0.0)
-    # Row-level BC-property resolution failures (unknown chill_material, an
-    # exo_sleeve with no calibrated MEF / no exo_calibration.csv yet) must
-    # WARN+skip this one row, exactly like every other row-level failure
-    # above (missing mesh/raw/t, node-count mismatch, unknown material_id) --
-    # NOT abort the whole build_v4 batch. _resolve_chill_effusivity_v4 /
-    # _resolve_exo_q0_v4 raise ValueError by design (never a silent 0.0); this
-    # is where that "loud" behavior is turned into a "skip THIS row" outcome
-    # instead of letting it propagate out of build_v4 mid-manifest.
-    try:
-        spec["chill_effusivity"] = _resolve_chill_effusivity_v4(row)
-        spec["exo_q0"] = _resolve_exo_q0_v4(row, exo_calibration_csv)
-    except ValueError as e:
-        print(f"  WARN {e}")
-        return None
-
-    # `spec` here deliberately has NO "C_material" key yet -- `_node_features`
-    # (version=4/5) never reads it (only mold_effusivity/bc_config/
-    # chill_effusivity/sleeve_MEF/exo_q0), so X/pos/mesh_edges are fully
-    # computable before the cross-row Chvorinov fit exists.
-    X, pos, mesh_edges = _node_features(mesh, version=feature_version, spec=spec)
-
-    volume = float(X[0, FEATURE_NAMES.index("volume_mm3")])
-    area = float(X[0, FEATURE_NAMES.index("surface_area_mm2")])
-    fam = row.get("alloy_family") or "unknown"
-
-    return {
-        "run_id": run_id, "dataset": row.get("dataset", ""),
-        "alloy_family": fam, "bc_config": bc_config,
-        "hold_out": row.get("hold_out", ""),
-        "X": X, "pos": pos, "mesh_edges": mesh_edges,
-        "spec": spec, "volume": volume, "area": area,
-        "t": t, "gradT": gradT, "dTdt": dTdt, "t_liquidus": t_liquidus,
-        "t_total": float(np.max(t)), "n_nodes": n_pt,
-    }
-
-
-def _finalize_row_stage2(s1: dict, C_by_family: dict):
-    """Pure arithmetic + dict assembly, no mesh/IO: plugs the fitted
-    per-family Chvorinov C into stage 1's output. Returns None (+ printed
-    WARN) if no C was fitted for this row's family (e.g. every row in that
-    family failed stage 1)."""
-    fam = s1["alloy_family"]
-    run_id = s1["run_id"]
-    if fam not in C_by_family:
-        print(f"  WARN no Chvorinov C fitted for alloy_family={fam!r} "
-              f"(run_id={run_id}) -- skipped")
-        return None
-
-    spec = dict(s1["spec"])
-    spec["C_material"] = C_by_family[fam]
-    t = s1["t"]
-    t_ref = chvorinov_t_ref(s1["volume"], s1["area"], spec["C_material"])
-    g = alloy_g_vector(spec)
-    y_nondim = (t.astype(np.float64) / max(t_ref, 1e-9)).astype(np.float32)
-
-    rec = {
-        "stem": run_id, "run_id": run_id, "dataset": s1["dataset"],
-        "alloy_family": fam, "bc_config": s1["bc_config"],
-        "hold_out": s1["hold_out"],
-        "X": s1["X"], "pos": s1["pos"], "mesh_edges": s1["mesh_edges"],
-        "g": g, "y": t, "y_nondim": y_nondim, "t_ref": t_ref,
-        "gradT": s1["gradT"], "dTdt": s1["dTdt"],
-        "y_global": float(np.max(t)), "n_nodes": s1["n_nodes"],
-    }
-    t_liquidus = s1["t_liquidus"]
-    if t_liquidus is not None:
-        # local solidification time = t_solidus - t_liquidus, clipped at 0
-        # (the two independent threshold-crossing interpolations can, in
-        # principle, disagree by a tiny numerical amount right at a node
-        # that crosses both thresholds within the same accepted step).
-        rec["t_liquidus"] = t_liquidus
-        rec["local_solidification_time"] = np.clip(t - t_liquidus, 0.0, None).astype(np.float32)
-    return rec
-
-
-def process_one_v4(row: dict, manifest_dir: str, raw_dir: str, data_root: str,
-                   C_by_family: dict,
-                   exo_calibration_csv: str = "data/properties/exo_calibration.csv"):
-    """Back-compat single-row entrypoint (still parses the mesh exactly
-    once). `build_v4` does NOT call this -- it calls stage 1 across the
-    whole manifest first (parallelised) and stage 2 after the C fit, so a
-    row's mesh is parsed once per BUILD, not once per stage."""
-    s1 = _process_row_stage1(row, manifest_dir, raw_dir, data_root, exo_calibration_csv)
-    if s1 is None:
-        return None
-    return _finalize_row_stage2(s1, C_by_family)
-
-
-def write_record_v4(rec, out_dir):
-    gdir = osp.join(out_dir, "graphs")
-    os.makedirs(gdir, exist_ok=True)
-    extra = {}
-    if "t_liquidus" in rec:
-        extra["t_liquidus"] = rec["t_liquidus"]
-        extra["local_solidification_time"] = rec["local_solidification_time"]
-    np.savez(osp.join(gdir, f"{rec['stem']}.npz"),
-             X=rec["X"], pos=rec["pos"], mesh_edges=rec["mesh_edges"],
-             g=rec["g"], y=rec["y"], y_raw=rec["y"], y_nondim=rec["y_nondim"],
-             t_ref=np.float64(rec["t_ref"]),
-             gradT=rec["gradT"], dTdt=rec["dTdt"], **extra)
-
-
-def _build_splits_v4(meta: dict) -> dict:
-    """{"interp": {"train","test"}, "extrap": {"train","test"}} from the
-    manifest's hold_out column. dataset=="S" (sanity cross-check pairs) rows
-    are EXCLUDED from every train (and test) list here, regardless of their
-    hold_out value -- they are not part of the GNN train/val protocol."""
-    stems = list(meta["stem"]); dataset = list(meta["dataset"]); hold_out = list(meta["hold_out"])
-    def _stems_where(pred):
-        return sorted(s for s, d, h in zip(stems, dataset, hold_out)
-                     if d != "S" and pred(h))
-    train = _stems_where(lambda h: h == "train")
-    val_interp = _stems_where(lambda h: h == "val_interp")
-    val_extrap = _stems_where(lambda h: h == "val_extrap")
-    return {
-        "interp": {"train": train, "test": val_interp},
-        "extrap": {"train": train, "test": val_extrap},
-    }
-
-
-def _init_norm_accum(feature_version: int = 4, n_feat: int = None):
-    """Fresh accumulator state for the norm-stats running sums, shared by
-    `build_v4`/`build_v5`'s Phase 2 fold and `recompute_norm_stats_v4` (so
-    these can never silently diverge in the accumulation math).
-
-    `n_feat`, if given, overrides the width implied by `feature_version`
-    (used by `recompute_norm_stats_v4` to size the accumulator from an
-    on-disk graph's ACTUAL X width -- 15 for a v4 cache, 20 for a v5 one --
-    instead of assuming a version from the caller)."""
-    if n_feat is None:
-        n_feat = len(FEATURE_NAMES_V5) if feature_version >= 5 else len(FEATURE_NAMES_V4)
-    n_g = len(G_VECTOR_NAMES)
-    return {
-        "fsum": np.zeros(n_feat), "fsq": np.zeros(n_feat), "fcnt": 0,
-        "gsum": np.zeros(n_g), "gsq": np.zeros(n_g), "gcnt": 0,
-        "ylsum": 0.0, "ylsq": 0.0, "ycnt": 0,
-    }
-
-
-def _accumulate_norm_stats(acc: dict, X: np.ndarray, g: np.ndarray, y: np.ndarray) -> None:
-    """Fold one row's (X, g, y) into `acc` in place."""
-    X = X.astype(np.float64)
-    acc["fsum"] += X.sum(0); acc["fsq"] += (X ** 2).sum(0); acc["fcnt"] += X.shape[0]
-    g = g.astype(np.float64)
-    acc["gsum"] += g; acc["gsq"] += g ** 2; acc["gcnt"] += 1
-    yl = np.log1p(y.astype(np.float64))
-    acc["ylsum"] += yl.sum(); acc["ylsq"] += (yl ** 2).sum(); acc["ycnt"] += yl.size
-
-
-def _finalize_norm_stats(acc: dict) -> dict:
-    """Turn the running sums into the final {feat,g,y_log}_{mean,std} dict --
-    exactly the arithmetic `build_v4` used to do inline."""
-    feat_mean = acc["fsum"] / max(acc["fcnt"], 1)
-    feat_std = np.sqrt(np.clip(acc["fsq"] / max(acc["fcnt"], 1) - feat_mean ** 2, 1e-12, None))
-    g_mean = acc["gsum"] / max(acc["gcnt"], 1)
-    g_std = np.sqrt(np.clip(acc["gsq"] / max(acc["gcnt"], 1) - g_mean ** 2, 1e-12, None))
-    y_log_mean = acc["ylsum"] / max(acc["ycnt"], 1)
-    y_log_std = float(np.sqrt(max(acc["ylsq"] / max(acc["ycnt"], 1) - y_log_mean ** 2, 1e-12)))
-    return {
-        "feat_mean": feat_mean, "feat_std": feat_std,
-        "g_mean": g_mean, "g_std": g_std,
-        "y_log_mean": np.float64(y_log_mean), "y_log_std": np.float64(y_log_std),
-    }
-
-
-def recompute_norm_stats_v4(features_dir: str) -> dict:
-    """Recompute `norm_stats.npz` from EVERY graph `.npz` currently on disk
-    under `<features_dir>/graphs/`, using exactly the same accumulation
-    conventions as `build_v4`'s Phase 2 (`_accumulate_norm_stats` /
-    `_finalize_norm_stats` are the SAME helpers build_v4 folds each row
-    through -- this cannot silently diverge from a full build's math).
-
-    dataset=="S" rows are excluded, mirroring build_v4's own exclusion --
-    looked up via `<features_dir>/meta.npz`'s stem/dataset columns (still
-    the full, untouched manifest-wide meta written by the last full build;
-    a selective `only_stems` rebuild deliberately never rewrites it).
-
-    Overwrites `<features_dir>/norm_stats.npz` in place and returns the new
-    stats dict. Intended for use AFTER a selective (`only_stems`) rebuild,
-    to fold the rebuilt rows' (possibly changed) features back into the
-    global normalization without re-running the whole manifest.
-
-    Version-agnostic (v4 15-wide or v5 20-wide cache): the accumulator's
-    feature width is read off the FIRST graph `.npz` actually found on disk
-    (`X.shape[1]`), not assumed from a hardcoded FEATURE_NAMES_V4 -- so this
-    one function serves both a v4 and a v5 `features_dir` without a
-    parameter, and can never disagree with what's actually cached there."""
-    meta_path = osp.join(features_dir, "meta.npz")
-    dataset_by_stem = {}
-    if osp.exists(meta_path):
-        with np.load(meta_path, allow_pickle=True) as mz:
-            for stem, ds in zip(mz["stem"], mz["dataset"]):
-                dataset_by_stem[str(stem)] = str(ds)
-
-    gpaths = sorted(glob(osp.join(features_dir, "graphs", "*.npz")))
-    n_feat = None
-    if gpaths:
-        with np.load(gpaths[0]) as z0:
-            n_feat = int(z0["X"].shape[1])
-    acc = _init_norm_accum(n_feat=n_feat)
-    for gp in gpaths:
-        stem = osp.splitext(osp.basename(gp))[0]
-        if dataset_by_stem.get(stem) == "S":
-            continue
-        with np.load(gp) as z:
-            _accumulate_norm_stats(acc, z["X"], z["g"], z["y"])
-
-    stats = _finalize_norm_stats(acc)
-    np.savez(osp.join(features_dir, "norm_stats.npz"), **stats)
-    return stats
-
-
-def _recover_c_by_family_from_cache(rows: list, out_dir: str) -> dict:
-    """For a SELECTIVE (`only_stems`) rebuild: recover each needed
-    alloy_family's Chvorinov C by inverting `t_ref = C * (V/A)**2` from the
-    row's OWN currently-cached graph `.npz` (read BEFORE this rebuild
-    overwrites it). Volume/area are geometry-only and unchanged by Task 1's
-    exothermic_input re-encoding, so this exactly reproduces the ORIGINAL
-    full-manifest fit's C for that family -- to floating-point precision,
-    not a re-fit from the (possibly much smaller/skewed) `only_stems`
-    subset, which would silently shift t_ref/y_nondim for every row in that
-    family (see the Task 2 brief's Chvorinov-C correctness point).
-
-    One point per family suffices (algebraic inversion of a constant, not a
-    regression); if more than one row of a family is present among `rows`,
-    their recovered C values are averaged purely as a sanity cross-check
-    (they should already agree to float precision)."""
-    v_idx = FEATURE_NAMES.index("volume_mm3")
-    a_idx = FEATURE_NAMES.index("surface_area_mm2")
-    fam_c_vals = {}
-    for row in rows:
-        run_id = row["run_id"]
-        fam = row.get("alloy_family") or "unknown"
-        gpath = osp.join(out_dir, "graphs", f"{run_id}.npz")
-        if not osp.exists(gpath):
-            continue
-        with np.load(gpath) as z:
-            if "t_ref" not in z.files or "X" not in z.files:
-                continue
-            X = z["X"]
-            t_ref_old = float(z["t_ref"])
-            V = float(X[0, v_idx]); A = float(X[0, a_idx])
-        if V <= 0 or A <= 0 or t_ref_old <= 0:
-            continue
-        M = V / A
-        fam_c_vals.setdefault(fam, []).append(t_ref_old / (M ** 2))
-    return {fam: float(np.mean(vals)) for fam, vals in fam_c_vals.items()}
-
-
-def _run_stage1_v4(rows: list, manifest_dir: str, raw_dir: str, data_root: str,
-                   exo_calibration_csv: str, n_workers, feature_version: int = 4) -> list:
-    """Stage 1 across the whole manifest. Returns a list of `dict | None`,
-    ONE PER ROW, in the SAME ORDER as `rows` (required for the serial
-    aggregation phase below to be byte-identical to the pre-parallel serial
-    build, which folded rows into norm_stats/meta in manifest order).
-    `n_workers in (None, 1)` runs serially in-process (no pool overhead,
-    easiest to debug); otherwise a bounded `ProcessPoolExecutor` -- BLAS
-    capped (module-top env vars + `_init_worker`'s threadpoolctl fallback)
-    and niced, mirroring tools/doe/doe.py's execute() pool. `ex.map` (not
-    submit+as_completed) is used specifically because it yields results in
-    submission order regardless of completion order.
-
-    `feature_version` (Task 2, bc-reach): forwarded to every
-    `_process_row_stage1` call (serial or pooled) -- `build_v4` passes 4
-    (the default), `build_v5` passes 5."""
-    if n_workers is None or n_workers <= 1:
-        return [_process_row_stage1(r, manifest_dir, raw_dir, data_root,
-                                    exo_calibration_csv, feature_version)
-                for r in rows]
-    with ProcessPoolExecutor(max_workers=n_workers, initializer=_init_worker) as ex:
-        return list(ex.map(
-            _process_row_stage1, rows,
-            itertools.repeat(manifest_dir), itertools.repeat(raw_dir),
-            itertools.repeat(data_root), itertools.repeat(exo_calibration_csv),
-            itertools.repeat(feature_version)))
-
-
-def _build_v4_family(manifest_path: str, feature_version: int, raw_dir: str,
-                     data_root: str, out_dir: str, limit: int,
-                     exo_calibration_csv: str, n_workers, only_stems,
-                     write_splits) -> None:
-    """Shared implementation behind `build_v4` (feature_version=4, 15-wide
-    X) and `build_v5` (feature_version=5, 20-wide X, Task 2/bc-reach).
-    Everything here is feature-version-agnostic: X's width is entirely
-    determined by `_process_row_stage1`/`_node_features(version=
-    feature_version)` in stage 1 -- the per-family Chvorinov-C fit (or, for
-    a selective `only_stems` rebuild, its cache-recovery + same-family
-    fallback), stage-2 finalize + write, and norm-stats/meta accumulation
-    never read a hardcoded column count.
-
-    `write_splits(meta, out_dir)` is the ONE place the two callers diverge:
-    called only for a FULL (non-selective) build, after meta.npz/
-    norm_stats.npz are already written, it decides how splits.json is
-    produced -- `build_v4` fits it from this build's own meta
-    (`_build_splits_v4`); `build_v5` instead copies an existing v4 cache's
-    splits.json VERBATIM (same manifest/rows, so the split membership is
-    identical by construction, never independently re-derived and risking
-    divergence). A selective (`only_stems`) rebuild never calls
-    `write_splits` at all -- splits.json is left untouched, same as
-    meta.npz/norm_stats.npz.
-
-    `only_stems` (optional set/iterable of `run_id`): a SELECTIVE rebuild --
-    process only the manifest rows whose `run_id` is in `only_stems`, write
-    (only) their graph `.npz`, and skip `meta.npz`/`norm_stats.npz`/
-    `splits.json` entirely (those describe the WHOLE manifest and a partial
-    rebuild must never truncate/corrupt them). The per-family Chvorinov C
-    each selected row needs is NOT re-fit from the (possibly tiny/skewed)
-    subset -- it is recovered from the already-cached graph .npz on disk
-    (`_recover_c_by_family_from_cache`), so t_ref/y_nondim for unaffected
-    rows never drifts. Use `recompute_norm_stats_v4` afterwards to fold the
-    rebuilt rows back into `norm_stats.npz`."""
-    rows = _read_manifest_rows(manifest_path)
-    if limit:
-        rows = rows[:limit]
-    rows_full = rows
-    if only_stems is not None:
-        only_stems = set(only_stems)
-        rows = [r for r in rows if r["run_id"] in only_stems]
-    manifest_dir = osp.dirname(osp.abspath(manifest_path))
-    os.makedirs(out_dir, exist_ok=True)
-
-    # Phase 1 (parallel, embarrassingly-parallel per row -- M-3): resolve +
-    # parse each row's mesh EXACTLY ONCE (previously parsed twice: once in
-    # `_fit_chvorinov_c_per_family`'s pre-pass, once in `process_one_v4`'s
-    # main pass), load its raw solve, resolve its BC/material spec, and
-    # compute X/pos/mesh_edges/volume/area/t_total -- everything that does
-    # NOT depend on the cross-row Chvorinov C fit.
-    stage1 = _run_stage1_v4(rows, manifest_dir, raw_dir, data_root,
-                            exo_calibration_csv, n_workers,
-                            feature_version=feature_version)
-
-    # Phase 1.5 (serial aggregate, cheap -- no mesh reparse, just a fold over
-    # already-computed scalars): fit the per-family Chvorinov C, EXCLUDING
-    # dataset=="S" rows (M-4, mirrored from `_fit_chvorinov_c_per_family`)
-    # from the fit, using the exact same fit math (`_fit_c_from_points`) so
-    # this can never silently diverge from the standalone/tested function.
-    # Selective (only_stems) rebuild: NEVER re-fit C from this row subset --
-    # recover the ORIGINAL full-manifest C per family from the already-cached
-    # graph .npz on disk instead (see build_v4's docstring + the Task 2
-    # brief's Chvorinov-C correctness point).
-    if only_stems is not None:
-        C_by_family = _recover_c_by_family_from_cache(rows, out_dir)
-        # NEW stems (no graph .npz of their own cached yet) contribute no
-        # point to the lookup above, so a family made up ENTIRELY of new
-        # stems comes back with no C at all -- fall back to recovering that
-        # family's C from any OTHER (already-cached, from a prior full or
-        # selective build) row of the same family across the FULL manifest,
-        # excluding dataset=="S" rows (mirrors the fit's own exclusion, see
-        # `_fit_chvorinov_c_per_family`). This is still the same algebraic
-        # inversion of a per-family constant, just sourced from a wider row
-        # set -- never a re-fit from the (possibly tiny/skewed) only_stems
-        # subset.
-        needed = {r.get("alloy_family") or "unknown" for r in rows}
-        missing = needed - set(C_by_family)
-        if missing:
-            fallback = [r for r in rows_full
-                        if (r.get("alloy_family") or "unknown") in missing
-                        and r.get("dataset") != "S"]
-            C_by_family.update(_recover_c_by_family_from_cache(fallback, out_dir))
-        # If a needed family STILL has no C after the full-manifest fallback,
-        # every one of its rows would otherwise be silently skipped by
-        # `_finalize_row_stage2`'s warn-and-skip -- a selective rebuild must
-        # fail loudly instead of quietly writing zero graphs.
-        still_missing = needed - set(C_by_family)
-        if still_missing:
-            raise RuntimeError(
-                f"build_v{feature_version}(only_stems=...): no cached graph "
-                f"found anywhere in {out_dir!r} to recover Chvorinov C for "
-                f"alloy_family(s) {sorted(still_missing)!r} -- run a full "
-                f"(non-selective) build_v{feature_version} first so at least "
-                f"one graph of each family is cached.")
-    else:
-        fam_points = {}
-        for s1 in stage1:
-            if s1 is None or s1["dataset"] == "S":
-                continue
-            fam_points.setdefault(s1["alloy_family"], []).append(
-                (s1["volume"] / s1["area"], s1["t_total"]))
-        C_by_family = _fit_c_from_points(fam_points)
-
-    # Phase 2 (serial -- pure arithmetic + file writes, no mesh parsing or
-    # other heavy work, so parallelising it would only add IPC/pickling
-    # overhead): finalize each row now that C_by_family is known, write its
-    # graph .npz, and fold it into the running norm-stats/meta accumulators
-    # in the SAME `rows` order the old serial build used -- byte-identical
-    # accumulation order (floating-point sums are order-sensitive in
-    # principle, though not in practice at this row count/precision).
-    names = FEATURE_NAMES_V5 if feature_version >= 5 else FEATURE_NAMES_V4
-    acc = _init_norm_accum(feature_version=feature_version)
-    meta = {k: [] for k in ("stem", "run_id", "dataset", "alloy_family",
-                            "bc_config", "hold_out", "y_global", "n_nodes")}
-    n_written = 0
-
-    for s1 in stage1:
-        if s1 is None:
-            continue
-        rec = _finalize_row_stage2(s1, C_by_family)
-        if rec is None:
-            continue
-        write_record_v4(rec, out_dir)
-        n_written += 1
-        if only_stems is not None:
-            # Selective rebuild: meta/norm_stats/splits describe the WHOLE
-            # manifest and are deliberately left untouched (see
-            # `_build_v4_family`'s docstring) -- use `recompute_norm_stats_v4`
-            # afterwards instead.
-            continue
-        # M-4 (ft2a-branch-review): dataset=="S" rows still get a cached
-        # graph .npz (and a meta.npz row, so tooling that looks them up by
-        # run_id keeps working) but must NOT fold into norm_stats -- they
-        # duplicate an existing A/B/C geometry+material and would otherwise
-        # mildly double-weight that point in feature/g/y normalization,
-        # mirroring their existing exclusion from the Chvorinov-C fit above
-        # and from train/test splits (`_build_splits_v4`).
-        if rec["dataset"] != "S":
-            _accumulate_norm_stats(acc, rec["X"], rec["g"], rec["y"])
-        for k in meta:
-            meta[k].append(rec[k])
-
-    if only_stems is not None:
-        print(f"Rebuilt {n_written} v{feature_version} graphs (of {len(rows)} "
-              f"selected stems) -> {out_dir}/graphs (meta.npz/norm_stats.npz/"
-              "splits.json untouched -- run recompute_norm_stats_v4 separately)")
-        return
-
-    stats = _finalize_norm_stats(acc)
-
-    np.savez(osp.join(out_dir, "meta.npz"),
-             stem=np.array(meta["stem"]), run_id=np.array(meta["run_id"]),
-             dataset=np.array(meta["dataset"]),
-             alloy_family=np.array(meta["alloy_family"]),
-             bc_config=np.array(meta["bc_config"]),
-             hold_out=np.array(meta["hold_out"]),
-             y_global=np.array(meta["y_global"], dtype=np.float32),
-             n_nodes=np.array(meta["n_nodes"]),
-             feature_names=np.array(names))
-    np.savez(osp.join(out_dir, "norm_stats.npz"), **stats)
-
-    write_splits(meta, out_dir)
-
-    print(f"Built {len(meta['stem'])} v{feature_version} graphs "
-          f"(of {len(rows)} manifest rows) -> {out_dir}")
-
-
-def build_v4(manifest_path: str, raw_dir: str = "data/raw_v4", data_root: str = "data",
-            out_dir: str = "data/features_v4", limit: int = None,
-            exo_calibration_csv: str = "data/properties/exo_calibration.csv",
-            n_workers=1, only_stems=None):
-    """v4 (15-wide X) feature-cache build. See `_build_v4_family`'s
-    docstring for the full contract (stage 1/2 split, only_stems selective
-    rebuild + Chvorinov-C cache-recovery). splits.json is fit from this
-    build's own meta via `_build_splits_v4` -- the v4-specific case of
-    `_build_v4_family`'s `write_splits` hook."""
-    def _write_splits_v4(meta: dict, out_dir: str) -> None:
-        splits = _build_splits_v4(meta)
-        with open(osp.join(out_dir, "splits.json"), "w") as f:
-            json.dump(splits, f, indent=2)
-
-    return _build_v4_family(manifest_path, 4, raw_dir, data_root, out_dir,
-                            limit, exo_calibration_csv, n_workers, only_stems,
-                            _write_splits_v4)
-
-
-def build_v5(manifest_path: str, raw_dir: str = "data/raw_v4", data_root: str = "data",
-            out_dir: str = "data/features_v5", limit: int = None,
-            exo_calibration_csv: str = "data/properties/exo_calibration.csv",
-            n_workers=1, only_stems=None,
-            v4_splits_path: str = "data/features_v4/splits.json"):
-    """Task 2 (bc-reach): the 20-channel feature-cache build -- `build_v4`
-    parameterized to feature_version=5 via the shared `_build_v4_family`
-    implementation (see its docstring for the full stage 1/2 + only_stems
-    contract, unchanged here). X is 20-wide: v4's 15 columns UNCHANGED
-    (byte-identical to a v4 build of the same row -- Task 1's
-    `compute_bc_reach_channels`/`FEATURE_NAMES_V5` only APPEND 5 new
-    columns) + the 5 bc-reach columns from Task 1.
-
-    Two deliberate differences from `build_v4`:
-      - norm_stats.npz is FRESH over the full 20-wide X -- NOT v4's stats
-        reused or zero-padded. `_init_norm_accum(feature_version=5)` sizes
-        the accumulator to `len(FEATURE_NAMES_V5) == 20` and every row's
-        real 20-column X is folded into it (`_accumulate_norm_stats`,
-        unchanged -- it never assumes a column count).
-      - splits.json is NOT computed from this build's own meta -- v5 uses
-        the IDENTICAL manifest/rows/hold_out protocol as the v4 cache it's
-        built alongside, so `splits.json` is copied VERBATIM from
-        `v4_splits_path` (default `data/features_v4/splits.json`, the
-        production v4 cache) instead of re-deriving it a second time --
-        guaranteeing byte-for-byte the same split membership (incl. e_val)
-        with zero risk of the two derivations silently diverging. Raises
-        `FileNotFoundError` if `v4_splits_path` doesn't exist -- a v5 build
-        requires the v4 cache to already exist; this never silently skips
-        or fabricates a split. Selective (`only_stems`) rebuilds never call
-        this at all (same as `build_v4`), so `v4_splits_path` need not even
-        be reachable for those."""
-    def _copy_splits_from_v4(meta: dict, out_dir: str) -> None:
-        if not osp.exists(v4_splits_path):
-            raise FileNotFoundError(
-                f"build_v5: v4_splits_path {v4_splits_path!r} does not exist "
-                "-- build_v5 copies splits.json VERBATIM from an existing v4 "
-                "cache (same manifest/rows) instead of re-deriving it; run "
-                "build_v4 on this manifest first, or pass the correct "
-                "v4_splits_path")
-        shutil.copyfile(v4_splits_path, osp.join(out_dir, "splits.json"))
-
-    return _build_v4_family(manifest_path, 5, raw_dir, data_root, out_dir,
-                            limit, exo_calibration_csv, n_workers, only_stems,
-                            _copy_splits_from_v4)
-
-
-def build(data_root="data", out_dir="data/features", n_workers=None, limit=None,
-         feature_version=2, manifest_path=None, raw_dir="data/raw_v4",
-         exo_calibration_csv="data/properties/exo_calibration.csv"):
-    if feature_version >= 4:
-        if manifest_path is None:
-            raise ValueError("build(feature_version=4) requires manifest_path")
-        return build_v4(manifest_path=manifest_path, raw_dir=raw_dir,
-                        data_root=data_root, out_dir=out_dir, limit=limit,
-                        exo_calibration_csv=exo_calibration_csv,
-                        n_workers=(n_workers if n_workers is not None else 1))
-    stems = _all_stems(data_root)
-    if limit:
-        stems = stems[:limit]
-    os.makedirs(out_dir, exist_ok=True)
-
-    # Determine feature names based on version
-    names = FEATURE_NAMES_V3 if feature_version >= 3 else FEATURE_NAMES
-
-    # accumulators for global normalization (Welford-free: sum, sumsq, count)
-    fsum = np.zeros(len(names)); fsq = np.zeros(len(names)); fcnt = 0
-    ylsum = ylsq = ycnt = 0.0
-    meta = {k: [] for k in ("stem", "prefix", "n_adiab", "is_fully_cooled",
-                            "y_global", "n_nodes")}
-
-    def _accumulate(rec):
-        nonlocal fsum, fsq, fcnt, ylsum, ylsq, ycnt
-        write_record(rec, out_dir)
-        X = rec["X"].astype(np.float64)
-        fsum += X.sum(0); fsq += (X ** 2).sum(0); fcnt += X.shape[0]
-        yl = np.log1p(rec["y"].astype(np.float64))
-        ylsum += yl.sum(); ylsq += (yl ** 2).sum(); ycnt += yl.size
-        for k in meta:
-            meta[k].append(rec[k])
-
-    worker = process_one
-    if n_workers == 1:
-        for i, s in enumerate(stems):
-            rec = worker(s, data_root, feature_version)
-            if rec is not None:
-                _accumulate(rec)
-            if (i + 1) % 500 == 0:
-                print(f"  {i+1}/{len(stems)}")
-    else:
-        with ProcessPoolExecutor(max_workers=n_workers) as ex:
-            futs = {ex.submit(worker, s, data_root, feature_version): s for s in stems}
-            done = 0
-            for fut in futs:
-                rec = fut.result()
-                done += 1
-                if rec is not None:
-                    _accumulate(rec)
-                if done % 500 == 0:
-                    print(f"  {done}/{len(stems)}")
-
-    feat_mean = fsum / max(fcnt, 1)
-    feat_std = np.sqrt(np.clip(fsq / max(fcnt, 1) - feat_mean ** 2, 1e-12, None))
-    y_log_mean = ylsum / max(ycnt, 1)
-    y_log_std = float(np.sqrt(max(ylsq / max(ycnt, 1) - y_log_mean ** 2, 1e-12)))
-
-    np.savez(osp.join(out_dir, "meta.npz"),
-             stem=np.array(meta["stem"]),
-             prefix=np.array(meta["prefix"]),
-             n_adiab=np.array(meta["n_adiab"]),
-             is_fully_cooled=np.array(meta["is_fully_cooled"]),
-             y_global=np.array(meta["y_global"], dtype=np.float32),
-             n_nodes=np.array(meta["n_nodes"]),
-             feature_names=np.array(names))
-    np.savez(osp.join(out_dir, "norm_stats.npz"),
-             feat_mean=feat_mean, feat_std=feat_std,
-             y_log_mean=np.float64(y_log_mean), y_log_std=np.float64(y_log_std))
-    print(f"Built {len(meta['stem'])} graphs → {out_dir}")
-
-
-if __name__ == "__main__":
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--data-root", default="data")
-    ap.add_argument("--out-dir", default="data/features")
-    ap.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 2) - 2))
-    ap.add_argument("--limit", type=int, default=None)
-    ap.add_argument("--feature-version", type=int, default=2)
-    ap.add_argument("--manifest", default=None,
-                   help="DOE manifest.csv (required for --feature-version 4)")
-    ap.add_argument("--raw-dir", default="data/raw_v4",
-                   help="raw FEniCSx solve npz dir ({run_id}.npz -> t/gradT/dTdt), v4 only")
-    a = ap.parse_args()
-    build(a.data_root, a.out_dir, a.workers, a.limit, a.feature_version,
-         manifest_path=a.manifest, raw_dir=a.raw_dir)

src/compare_d_parts.py

@@ -1,214 +0,0 @@
-"""FT2b v4-inference Task 5: D-parts v4-vs-SP1 comparison.
-
-Before promoting the v4 hybrid (FiLM GNN + multi-alloy GBM, `src.evaluate_ft2b`)
-over the SP1 production hybrid (`runs/sp1_run3_final/best_gnn.pt` +
-`runs/baseline_v2/hgb_A.joblib`), both stacks are compared head-to-head on
-the real-part "dataset D" geometries (`tools/doe/prep_parts.py`'s cleaned
-STEP parts) that have a FEniCSx-ground-truth v4 feature cache. This is the
-*same* geometry fed through *both* stacks -- exactly the `_no_regression_check`
-technique in `src.evaluate_ft2b` (v4_X[:, :13]/[:9] through SP1's own norm
-stats), just scoped to dataset=="D" instead of dataset=="A", and reported
-per-stem as well as aggregated.
-
-Reused, not reimplemented: `load_ft2b_gnn`/`predict_graph_seconds` (GNN
-loading+inference, film-aware), `_evaluate_stems_v4` (v4 quantile-anchor hybrid
-per stem), `_sp1_hybrid_predict_stem` (SP1 mean-anchor hybrid on the v4
-cache's geometry), `hotspot_overlap` -- all from `src.evaluate_ft2b`, the
-single source of truth for both compositions (see that module's docstring
-for the quantile-anchor-vs-mean-anchor rationale).
-"""
-from __future__ import annotations
-
-import json
-import os
-import os.path as osp
-
-import joblib
-import numpy as np
-import torch
-
-from src import evaluate_ft2b as ft2b
-from src import metrics
-from src.dataload import load_norm
-from src.gbm_ft2b import _load_manifest_rows
-
-
-def _d_stems(manifest_rows: dict, features_dir: str) -> tuple[list[str], int]:
-    """Sorted `run_id`s with `dataset=="D"` AND a v4 graph cache on disk, plus
-    the raw count of `dataset=="D"` manifest rows (cache present or not) --
-    the latter lets callers tell "no D rows in the manifest at all" apart
-    from "D rows exist but none have been built into the v4 cache yet"."""
-    d_rows = [stem for stem, row in manifest_rows.items() if row.get("dataset") == "D"]
-    out = [s for s in d_rows if osp.exists(osp.join(features_dir, "graphs", f"{s}.npz"))]
-    return sorted(out), len(d_rows)
-
-
-def _stem_metrics(y_true: np.ndarray, y_pred: np.ndarray, xyz: np.ndarray) -> dict:
-    gid = np.zeros(y_true.shape[0], dtype=np.int64)
-    return {
-        "mae": metrics.mae(y_true, y_pred),
-        "hotspot": ft2b.hotspot_overlap(y_true, y_pred, xyz, gid),
-        "std_ratio": metrics.pred_std_ratio(y_true, y_pred, gid),
-    }
-
-
-def _aggregate_metrics(records: list[dict]) -> dict:
-    """Node-weighted aggregate over `{"y_true", "y_hybrid", "xyz"}` records.
-
-    Delegates to `ft2b._overall_metrics` (single source of truth for the
-    node-weighted concat/gid glue -- don't re-derive it here and risk it
-    drifting out of sync), just renaming `hotspot_overlap` -> `hotspot` to
-    match this module's own key convention."""
-    agg = ft2b._overall_metrics(records)
-    return {"mae": agg["mae"], "hotspot": agg["hotspot_overlap"],
-            "std_ratio": agg["std_ratio"]}
-
-
-def compare_d_parts(v4_ckpt, v4_gbm,
-                    sp1_ckpt=ft2b.DEFAULT_SP1_CKPT, sp1_hgb=ft2b.DEFAULT_SP1_HGB,
-                    features_dir="data/features_v4", manifest="doe_out/manifest.csv",
-                    out_json="runs/v4_inference/d_parts_comparison.json",
-                    baseline_features_dir=ft2b.DEFAULT_BASELINE_FEATURES_DIR,
-                    hgb_std=None,
-                    lambda_shrink: float = 1.0,
-                    w_total: float = 1.0,
-                    w_std: float = 1.0,
-                    device=None) -> dict:
-    """Head-to-head v4-hybrid vs SP1-hybrid comparison on dataset D stems.
-
-    Args:
-        v4_ckpt: path to the v4 GNN checkpoint, or an already-loaded
-            checkpoint dict (`ft2b.load_ft2b_gnn`'s own convention).
-        v4_gbm: path to the refit multi-alloy GBM (`src.gbm_ft2b.refit_gbm`),
-            or an already-loaded estimator.
-        sp1_ckpt / sp1_hgb: the SP1 production hybrid stack.
-        features_dir: v4 feature cache directory (`graphs/<stem>.npz`).
-        manifest: DOE manifest CSV (`run_id`==stem; needs `dataset`/
-            `geometry_id`/`alloy_family`/`bc_config` columns).
-        out_json: where to write the result as JSON (falsy skips the write;
-            the parent directory is created when writing).
-        baseline_features_dir: SP1 HGB norm stats (`norm_stats.npz`, 9 v2
-            features) -- same default as `evaluate_ft2b`.
-        hgb_std (variance-recomposition Task 3): an optional path to (or
-            already-loaded) GBM std-head model
-            (`src.gbm_ft2b.refit_gbm_std`'s output), forwarded ONLY to the
-            v4 side's `ft2b._evaluate_stems_v4` (the SP1 side has no
-            std-head analog and is unaffected). `None` (the default) is
-            byte-identical to before this argument existed (the v4 side
-            stays on the quantile-anchor). The returned dict gains a
-            top-level `"recomposition": "soft"|"anchor"` field (soft-anchor
-            Task 3: was `"affine"|"anchor"` -- `ft2b._evaluate_stems_v4` now
-            recomposes via `src.hybrid.hybrid_field_soft` instead of
-            `hybrid_field_affine` whenever `hgb_std` is given).
-        lambda_shrink (variance-recomposition amendment 2): forwarded
-            straight to the v4 side's `ft2b._evaluate_stems_v4` (ignored
-            when `hgb_std` is `None`). `1.0` (the default) is byte-identical
-            to before this argument existed.
-        w_total / w_std (soft-anchor Task 3): forwarded straight to the v4
-            side's `ft2b._evaluate_stems_v4`/`hybrid_field_soft`'s own
-            weights (ignored when `hgb_std` is `None`). `1.0`/`1.0` (the
-            defaults) reproduce `hybrid_field_soft`'s own default weighting.
-        device: "cuda"/"cpu"; default picks cuda if available.
-
-    Returns: `{"rows": [{"stem", "geometry_id", "alloy_family",
-        "v4": {"mae", "hotspot", "std_ratio"}, "sp1": {...}}, ...],
-        "v4": {...aggregate...}, "sp1": {...aggregate...}, "n_stems",
-        "n_d_stems_in_manifest", "verdict": {"mae_ok", "hotspot_ok"[, "note"]},
-        "recomposition"}`
-        where `mae_ok = v4.mae <= sp1.mae` and `hotspot_ok = v4.hotspot >=
-        sp1.hotspot` (both computed on the node-weighted aggregates,
-        mirroring `evaluate_ft2b`'s `_overall_metrics`) -- EXCEPT when
-        `n_stems == 0`, which forces BOTH verdict fields to `False` (an
-        empty comparison is not evidence of no-regression: `0.0 <= 0.0` and
-        `0.0 >= 0.0` would otherwise vacuously pass) and adds a `note`
-        explaining why, mentioning `n_stems`/`n_d_stems_in_manifest`.
-        `n_d_stems_in_manifest` is the raw count of `dataset=="D"` manifest
-        rows (cache present or not), for skip-visibility -- `n_stems` (<=
-        it) is how many of those actually had a `features_dir` graph cache
-        and were compared.
-    """
-    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-    manifest_dir = osp.dirname(osp.abspath(manifest)) or "."
-    manifest_rows = _load_manifest_rows(manifest)
-    stems, n_d_in_manifest = _d_stems(manifest_rows, features_dir)
-
-    # -- v4 hybrid (GNN shape, quantile-anchor GBM scale -- or, when `hgb_std`
-    # is given, the soft-anchor recomposition) -- exactly `evaluate_ft2b`'s
-    # own per-stem composition, reused verbatim.
-    v4_model, v4_cfg, v4_norm = ft2b.load_ft2b_gnn(v4_ckpt, device)
-    v4_gbm_model = joblib.load(v4_gbm) if isinstance(v4_gbm, (str, os.PathLike)) else v4_gbm
-    hgb_std_model = joblib.load(hgb_std) if isinstance(hgb_std, (str, os.PathLike)) else hgb_std
-    v4_records = ft2b._evaluate_stems_v4(v4_model, v4_cfg, v4_norm, v4_gbm_model, stems,
-                                        manifest_rows, manifest_dir, features_dir, device,
-                                        hgb_std=hgb_std_model, lambda_shrink=lambda_shrink,
-                                        w_total=w_total, w_std=w_std)
-    v4_by_stem = {r["stem"]: r for r in v4_records}
-
-    # -- SP1 hybrid (GNN shape, mean-anchor HGB scale) on the SAME v4-cache
-    # geometry, fed through SP1's own norm stats -- `_sp1_hybrid_predict_stem`'s
-    # own contract (see `evaluate_ft2b`'s `_no_regression_check` docstring for
-    # the full rationale/caveats of this "plug the v4 cache into the SP1
-    # stack" technique).
-    sp1_model, sp1_cfg, sp1_norm = ft2b.load_ft2b_gnn(sp1_ckpt, device)
-    hgb = joblib.load(sp1_hgb) if isinstance(sp1_hgb, (str, os.PathLike)) else sp1_hgb
-    hgb_norm = load_norm(baseline_features_dir)
-    sp1_by_stem = {}
-    for stem in stems:
-        y_true, hyb, pos_np = ft2b._sp1_hybrid_predict_stem(
-            stem, sp1_model, sp1_cfg, sp1_norm, hgb, hgb_norm, features_dir, device)
-        sp1_by_stem[stem] = {"y_true": y_true, "y_hybrid": hyb, "xyz": pos_np}
-
-    rows = []
-    for stem in stems:
-        v4r, sp1r = v4_by_stem[stem], sp1_by_stem[stem]
-        row = manifest_rows.get(stem, {})
-        rows.append({
-            "stem": stem,
-            "geometry_id": row.get("geometry_id") or "",
-            "alloy_family": v4r.get("family") or "unknown",
-            "v4": _stem_metrics(v4r["y_true"], v4r["y_hybrid"], v4r["xyz"]),
-            "sp1": _stem_metrics(sp1r["y_true"], sp1r["y_hybrid"], sp1r["xyz"]),
-        })
-
-    v4_agg = _aggregate_metrics([v4_by_stem[s] for s in stems])
-    sp1_agg = _aggregate_metrics([sp1_by_stem[s] for s in stems])
-
-    if stems:
-        verdict = {
-            "mae_ok": bool(v4_agg["mae"] <= sp1_agg["mae"]),
-            "hotspot_ok": bool(v4_agg["hotspot"] >= sp1_agg["hotspot"]),
-        }
-    else:
-        # An empty comparison must never read as a pass: with no records,
-        # both aggregates are all-zero, so a naive `0.0 <= 0.0` /
-        # `0.0 >= 0.0` would vacuously satisfy both gates. Force an
-        # explicit fail instead, with a note pointing at WHY (0 of
-        # n_d_in_manifest dataset=="D" manifest rows had a cached graph).
-        verdict = {
-            "mae_ok": False,
-            "hotspot_ok": False,
-            "note": (f"no D stems compared (n_stems=0 of "
-                     f"n_d_stems_in_manifest={n_d_in_manifest} dataset==\"D\" "
-                     "manifest rows have a features_dir graph cache) -- "
-                     "vacuous run, verdict forced to fail rather than "
-                     "trivially passing on empty aggregates"),
-        }
-
-    result = {
-        "rows": rows,
-        "v4": v4_agg,
-        "sp1": sp1_agg,
-        "n_stems": len(stems),
-        "n_d_stems_in_manifest": n_d_in_manifest,
-        "verdict": verdict,
-        "recomposition": "soft" if hgb_std_model is not None else "anchor",
-    }
-
-    if out_json:
-        out_dir = osp.dirname(out_json)
-        if out_dir:
-            os.makedirs(out_dir, exist_ok=True)
-        with open(out_json, "w") as f:
-            json.dump(result, f, indent=2)
-
-    return result

src/dataload.py

@@ -1,28 +0,0 @@
-"""Load cached per-node features for a set of stems into flat matrices."""
-import os.path as osp
-import numpy as np
-
-
-def load_norm(features_dir: str) -> dict:
-    d = np.load(osp.join(features_dir, "norm_stats.npz"))
-    return {"feat_mean": d["feat_mean"], "feat_std": d["feat_std"],
-            "y_log_mean": float(d["y_log_mean"]), "y_log_std": float(d["y_log_std"])}
-
-
-def load_matrix(features_dir: str, stems, max_nodes_per_graph=None, seed=0):
-    rng = np.random.default_rng(seed)
-    Xs, ys, gids = [], [], []
-    for gi, stem in enumerate(stems):
-        d = np.load(osp.join(features_dir, "graphs", f"{stem}.npz"))
-        X, y = d["X"], d["y"]
-        if max_nodes_per_graph is not None and X.shape[0] > max_nodes_per_graph:
-            idx = rng.choice(X.shape[0], max_nodes_per_graph, replace=False)
-            X, y = X[idx], y[idx]
-        Xs.append(X)
-        ys.append(y)
-        gids.append(np.full(X.shape[0], gi, dtype=np.int64))
-    return (np.concatenate(Xs), np.concatenate(ys), np.concatenate(gids))
-
-
-def standardize(X, norm) -> np.ndarray:
-    return (X - norm["feat_mean"]) / norm["feat_std"]

src/evaluate_ft2b.py

@@ -1,630 +0,0 @@
-"""FT2b Task 8: evaluation of the multi-alloy/multi-BC hybrid predictor.
-
-Consumes a trained v4 checkpoint (FiLM, 15 node features, `g`-conditioned,
-`node_in=15`) + the refit multi-alloy GBM scale-head (`src.gbm_ft2b`) and
-reports, on `interp.test` AND `extrap.test`: per-`alloy_family` and per-
-`bc_config` hybrid MAE (physical seconds), hot-spot top-10% overlap, and
-`std_ratio = std(pred)/std(true)`; a no-regression comparison against the
-SP1 hybrid on the all-sand subset; a Chvorinov-scaling sanity check; and a
-variance-collapse flag.
-
-=== y / y_raw / t_ref inversion chain (read before touching this file) ===
-`src/build_dataset.py` (`_finalize_row_stage2`, `write_record_v4`) writes,
-per v4 graph cache:
-    y       = t              (PHYSICAL seconds, per-node solidification time)
-    y_raw   = t               (same array, duplicate key for API clarity)
-    y_nondim = t / t_ref       (t_ref = C_family * (V/A)^2, Chvorinov reference)
-    t_ref   = C_family * (V/A)^2   (scalar, physical seconds)
-`build_v4`'s norm_stats.npz computes `y_log_mean`/`y_log_std` from
-`log1p(rec["y"])` i.e. `log1p(t)` -- log1p of the PHYSICAL target, exactly
-the same convention `features_v2`/`features_v3` use. `SolidGraphDataset`
-(`src/gnn_dataset.py`) trains on `standardize(log1p(d["y"]))` i.e.
-standardized log1p(PHYSICAL seconds) -- NOT the non-dim `y_nondim` field.
-Therefore the GNN's raw output, once unstandardized and expm1'd, is ALREADY
-physical seconds -- the exact same inversion `src.evaluate_gnn.
-predict_seconds_gnn` performs for SP1. No `t_ref` multiplication is needed
-to get physical seconds out of the v4 GNN. `y_nondim`/`t_ref`/`g` are real,
-separately-stored cache fields used elsewhere (the Chvorinov sanity check
-below reads `t_ref`; `g` feeds FiLM) -- they are NOT an alternate target
-space the model was trained on.
-
-=== v4 hybrid recomposition: QUANTILE-anchor (q=0.995), not MEAN-anchor ===
-`src.hybrid.hybrid_field(gnn, baseline)` (SP1's convention) re-anchors to the
-baseline's MEAN: `pred = gnn - mean(gnn) + mean(baseline)`. That is correct
-for SP1 because SP1's HGB baseline predicts a genuine PER-NODE value, so
-`mean(baseline)` is a real per-graph mean estimate. The FT2b GBM
-(`src.gbm_ft2b`) is different by design (Task 6): it predicts a single
-per-graph SCALAR, `log1p(total/max solidification time)` -- see its module
-docstring, "Target: per-graph `log1p(y_raw.max())`". Broadcasting that
-scalar to every node and then MEAN-anchoring would force
-`mean(pred) == gbm_total`, inflating every node by
-`(true_max - true_mean)` -- confirmed in review to blow up MAE by ~60x on an
-otherwise near-perfect GNN. `_max_anchor_hybrid` (an alias of
-`src.hybrid.hybrid_field_max_anchor`, the single source -- see that
-function's docstring) anchors to a high QUANTILE of the GNN field instead
-(default `q=0.995`; `q=1.0` reproduces the original exact-MAX behavior
-bit-for-bit): `pred_i = gnn_i - anchor + gbm_total` where
-`anchor = np.quantile(gnn, q)`, which makes `max(pred) ~= gbm_total` (the
-quantity the GBM actually predicts) while still preserving the GNN's field
-shape (its deviations from that anchor). The anchor moved off the exact
-max (Task 7) because it is fragile to a handful of outlier nodes -- a
-promotion-gate smoke test found a few inflated GNN log-space predictions
-on a real casting dragging ~90% of the field below the zero clip, since the
-single max node alone set the whole graph's additive offset; anchoring to
-the 99.5th percentile instead is robust to that. This function is used for
-EVERY v4 hybrid composition in this module (`_evaluate_stems_v4`, and
-therefore every per-family/per-bc metric, the v4 side of `no_regression`,
-and `chvorinov`'s `pred_total`). The SP1 side of `no_regression`
-(`_sp1_hybrid_predict_stem`) deliberately keeps `src.hybrid.hybrid_field`'s
-MEAN-anchor unchanged -- it is SP1's own genuine per-node HGB baseline, not
-a broadcast scalar, so the MEAN-anchor is correct there.
-"""
-from __future__ import annotations
-
-import json
-import os
-import os.path as osp
-from collections import defaultdict
-
-import joblib
-import numpy as np
-import torch
-
-from src import metrics
-from src.baseline import predict_seconds
-from src.dataload import load_norm, standardize
-from src.evaluate_gnn import predict_seconds_gnn
-from src.gbm_ft2b import (_load_manifest_rows, _row_features, predict_std_seconds,
-                          std_feature_row)
-from src.graph_build import (edge_attr, knn_edge_index, union_edge_attr,
-                             union_edge_index)
-from src.hybrid import hybrid_field, hybrid_field_soft
-from src.hybrid import hybrid_field_max_anchor as _max_anchor_hybrid
-from src.model_gnn import SolidGNN
-
-# The HGB v2/SP1 baseline was trained on the 9 v2 features = first 9 columns
-# of any newer feature version (ordering contract, SP1 spec §3). The SP1
-# production GNN was trained on the 13 v3 features = first 13 columns of the
-# v4 cache's X (`src/features.py`: "v4: v3's 13 columns UNCHANGED").
-N_V2_FEATURES = 9
-N_V3_FEATURES = 13
-
-DEFAULT_OUT_PATH = "runs/ft2b_stage3/eval.json"
-DEFAULT_SP1_CKPT = "runs/sp1_run3_final/best_gnn.pt"
-DEFAULT_SP1_HGB = "runs/baseline_v2/hgb_A.joblib"
-DEFAULT_BASELINE_FEATURES_DIR = "data/features_v2"
-
-
-def hotspot_overlap(y_true, y_pred, xyz, graph_id, top_frac: float = 0.10,
-                    tol_mm: float = 5.0) -> float:
-    """Top-`top_frac` latest-solidifying-node overlap, within `tol_mm`.
-
-    Thin, deliberately-not-reimplemented wrapper over
-    `src.metrics.hotspot_accuracy` -- the README-documented hot-spot metric
-    ("overlap of predicted vs. true top-10% latest-solidifying nodes, within
-    a tolerance"). Exposed under this module's own name so callers of
-    `evaluate_ft2b` don't need to reach into `src.metrics` directly.
-    """
-    return metrics.hotspot_accuracy(y_true, y_pred, xyz, graph_id,
-                                    top_frac=top_frac, tol_mm=tol_mm)
-
-
-# --------------------------------------------------------------------------- #
-# GNN loading / prediction (v4-aware: film + g)
-# --------------------------------------------------------------------------- #
-def load_ft2b_gnn(ckpt, device):
-    """Load a `{model, cfg, norm}` checkpoint into a `SolidGNN` (film-aware).
-
-    `ckpt` may be a path (loaded via `torch.load`) or an already-loaded
-    checkpoint dict -- convenient for tests that build a ckpt in-process.
-
-    Learned-heads Task 3: `cfg.get("deriv_heads")` truthy builds the model
-    with `deriv_heads=True` (the `SolidGNN` gains `grad_head`/`dtdt_head`,
-    see `src/model_gnn.py`) before the strict `load_state_dict` below -- a
-    heads checkpoint's saved `state["model"]` already carries the two heads'
-    weights (`src/train_heads.py` saves the full `model.state_dict()` of a
-    `deriv_heads=True` model), so this is still a plain strict load, not
-    `attach_deriv_heads` (that helper is only for loading a HEADLESS
-    state_dict into a heads-shaped model, e.g. inside the trainer). A
-    headless checkpoint (`cfg.get("deriv_heads")` absent/falsy, the default)
-    builds `deriv_heads=False` exactly as before -- byte-identical path.
-    """
-    state = torch.load(ckpt, map_location=device, weights_only=False) \
-        if isinstance(ckpt, (str, os.PathLike)) else ckpt
-    cfg, norm = state["cfg"], state["norm"]
-    model = SolidGNN(node_in=cfg.get("node_in", 15), edge_in=cfg.get("edge_in", 4),
-                     hidden=cfg["hidden"], n_blocks=cfg["n_blocks"],
-                     two_head=cfg.get("two_head", False),
-                     film=cfg.get("film", False), g_dim=cfg.get("g_dim", 5),
-                     deriv_heads=cfg.get("deriv_heads", False)).to(device)
-    model.load_state_dict(state["model"])
-    model.eval()
-    return model, cfg, norm
-
-
-def _build_edges(pos, cfg, mesh_edges=None):
-    if cfg.get("graph", "knn") == "union":
-        if mesh_edges is None:
-            raise ValueError("checkpoint expects a union graph but mesh_edges is missing")
-        ei, is_mesh = union_edge_index(pos, mesh_edges, k=cfg["k"])
-        return ei, union_edge_attr(pos, ei, is_mesh)
-    ei = knn_edge_index(pos, k=cfg["k"])
-    return ei, edge_attr(pos, ei)
-
-
-def _run_model(model, x, ei, ea, zeros, g, device):
-    """Shared forward-pass plumbing for `predict_graph_seconds`/
-    `predict_graph_derivs` (learned-heads Task 3): moves the already-built
-    tensors to `device` and calls the model once. Returns whatever tuple the
-    model itself returns -- 2-tuple `(y_nodal, y_global)` for a headless
-    model, 3-tuple `(y_nodal, y_global, deriv)` for a `deriv_heads=True` one
-    -- so callers must index (`out[0]`), never unpack-by-position, to stay
-    correct for both."""
-    with torch.no_grad():
-        return model(x.to(device), ei.to(device), ea.to(device), zeros, g=g)
-
-
-def _prepare_graph_inputs(cfg, norm, X, pos_np, mesh_edges, g_vec, device):
-    """Shared tensor-building for `predict_graph_seconds`/`predict_graph_derivs`
-    (learned-heads Task 3): standardizes `X`, builds the edge set per `cfg`,
-    and the FiLM `g` vector (if applicable). Returns `(x, ei, ea, zeros, g)`,
-    all ready to feed straight into `_run_model`."""
-    fmean = torch.tensor(norm["feat_mean"], dtype=torch.float32)
-    fstd = torch.tensor(norm["feat_std"], dtype=torch.float32)
-    x = (torch.tensor(X, dtype=torch.float32) - fmean) / fstd
-    pos = torch.tensor(pos_np, dtype=torch.float32)
-    me = torch.tensor(mesh_edges, dtype=torch.long) if mesh_edges is not None else None
-    ei, ea = _build_edges(pos, cfg, me)
-    zeros = torch.zeros(x.shape[0], dtype=torch.long, device=device)
-    g = None
-    if cfg.get("film", False) and g_vec is not None:
-        g = torch.tensor(np.asarray(g_vec), dtype=torch.float32).reshape(1, -1).to(device)
-    return x, ei, ea, zeros, g
-
-
-def predict_graph_seconds(model, cfg, norm, X, pos_np, mesh_edges=None, g_vec=None,
-                          device="cpu") -> np.ndarray:
-    """One graph's per-node hybrid-shape input -> physical-seconds GNN prediction.
-
-    `g_vec`: the cache's raw (5,) alloy/mould vector, forwarded un-standardized
-    (mirrors `SolidGraphDataset`, which does not standardize `g`) -- only used
-    when `cfg["film"]` is set.
-
-    Learned-heads Task 3: `model` may be a `deriv_heads=True` `SolidGNN`
-    (forward returns a 3-tuple `(y_nodal, y_global, deriv)` instead of the
-    headless 2-tuple) -- this function's OWN return contract is unchanged
-    either way (always just the `(N,)` physical-seconds array from
-    `y_nodal`); the extra `deriv` output (if any) is simply ignored here.
-    Use the sibling `predict_graph_derivs` to get the derivative heads'
-    predictions.
-    """
-    x, ei, ea, zeros, g = _prepare_graph_inputs(cfg, norm, X, pos_np, mesh_edges, g_vec, device)
-    out = _run_model(model, x, ei, ea, zeros, g, device)
-    yn = out[0]
-    return predict_seconds_gnn(yn.cpu().numpy(), norm)
-
-
-def predict_graph_derivs(model, cfg, norm, X, pos_np, mesh_edges=None, g_vec=None,
-                         device="cpu"):
-    """One graph's per-node hybrid-shape input -> physical-units learned
-    gradT (K/m) / dTdt (K/s), the two `deriv_heads` outputs (learned-heads
-    Task 3). Sibling of `predict_graph_seconds`, sharing its exact tensor
-    construction (standardized `X`, `cfg`-selected edge set, FiLM `g`) --
-    only the output side differs.
-
-    `model` MUST be a `deriv_heads=True` SolidGNN (forward returns a 3-tuple
-    `(y_nodal, y_global, deriv)`, `deriv` shape `(N, 2)` = standardized
-    `[log1p(gradT), log1p(|dTdt|)]`) -- a headless (2-tuple) model raises a
-    loud `ValueError` naming the mismatch rather than crashing on a tuple-
-    index error deeper in the call, or silently returning garbage.
-
-    `norm` must carry `deriv_log_mean`/`deriv_log_std` (each shape `(2,)`,
-    written by `src.train_heads.train_deriv_heads` into the checkpoint's
-    `norm` dict) -- the per-head log-space standardization stats the model
-    was trained against.
-
-    Inverse transform (pinned exactly, mirrors `predict_seconds_gnn`'s own
-    `pred_log = pred_std_log * std + mean; expm1`, applied per-head):
-    `deriv_std * z + deriv_mean -> expm1 -> clip(min=0)`. Both outputs are
-    non-negative by construction: `gradT` is a gradient MAGNITUDE (never
-    signed), and `dTdt` is returned as a positive magnitude in K/s -- this
-    matches the cache's own `dTdt` convention (`tools/doe/crossing_capture.py`
-    stores `dTdt_at = (T_prev - T_now) / dt`, i.e. a positive cooling rate by
-    construction on a monotonically-cooling casting; `src.train_heads`'s
-    `DerivGraphDataset` additionally takes `.abs()` defensively before
-    log-standardizing) -- and `src.shrinkage.niyama_exact` itself also takes
-    `np.abs(dTdt_K_per_s)` internally, so this function's non-negative
-    `dTdt` is compatible either way.
-
-    Returns `(gradT_K_per_m, dTdt_K_per_s)`, each an `(N,)` float64 array.
-    """
-    x, ei, ea, zeros, g = _prepare_graph_inputs(cfg, norm, X, pos_np, mesh_edges, g_vec, device)
-    out = _run_model(model, x, ei, ea, zeros, g, device)
-    if len(out) != 3:
-        raise ValueError(
-            "predict_graph_derivs: model did not return a deriv-heads output "
-            f"(expected a 3-tuple (y_nodal, y_global, deriv), got a "
-            f"{len(out)}-tuple) -- pass a deriv_heads=True SolidGNN / a heads "
-            "checkpoint (cfg['deriv_heads'] truthy, loaded via load_ft2b_gnn)")
-    deriv = out[2].cpu().numpy().astype(np.float64)
-    mean = np.asarray(norm["deriv_log_mean"], dtype=np.float64)
-    std = np.asarray(norm["deriv_log_std"], dtype=np.float64)
-    log_vals = deriv * std + mean
-    vals = np.clip(np.expm1(log_vals), 0.0, None)
-    return vals[:, 0], vals[:, 1]
-
-
-# --------------------------------------------------------------------------- #
-# Per-stem v4 hybrid prediction (GNN shape x GBM total-time scale)
-# --------------------------------------------------------------------------- #
-def _gbm_seconds_for_stem(gbm, stem, manifest_rows, manifest_dir, features_dir) -> float:
-    x_row, _y_log_true = _row_features(stem, manifest_rows, manifest_dir, features_dir)
-    log_pred = float(gbm.predict(x_row.reshape(1, -1))[0])
-    return float(np.clip(np.expm1(log_pred), 0.0, None))
-
-
-def _evaluate_stems_v4(model, cfg, norm, gbm, stems, manifest_rows, manifest_dir,
-                       features_dir, device, hgb_std=None, lambda_shrink: float = 1.0,
-                       w_total: float = 1.0, w_std: float = 1.0
-                       ) -> list[dict]:
-    """One record per stem: per-node y_true/y_hybrid/xyz + per-stem metadata
-    (alloy_family, bc_config, t_ref, predicted total time).
-
-    `hgb_std` (variance-recomposition Task 3): an optional loaded GBM
-    std-head model (`src.gbm_ft2b.refit_gbm_std`'s output). `None` (the
-    default) is byte-identical to before this argument existed --
-    `_max_anchor_hybrid` (quantile-anchor). When given, each stem's row
-    features are rebuilt via `src.gbm_ft2b.std_feature_row` (soft-anchor
-    Task 3: width-dispatched off `hgb_std`'s own `n_features_in_`, 16 or 17
-    -- a SEPARATE call from the one `_gbm_seconds_for_stem` makes internally
-    for the total, kept separate, rather than refactored to share one call,
-    so existing tests that monkeypatch `_gbm_seconds_for_stem` wholesale
-    keep working unmodified) and recomposition uses
-    `src.hybrid.hybrid_field_soft` (soft-anchor Task 3; was
-    `hybrid_field_affine`) instead; each such record also carries
-    `recomposition_info` (that function's `info` dict, now carrying
-    `"m"`/`"target_unreachable"`/... in place of affine's `"a_clamped"`).
-
-    `lambda_shrink` (variance-recomposition amendment 2): forwarded straight
-    to `hybrid_field_soft` (ignored when `hgb_std` is `None`, same as every
-    other soft-only knob here). `1.0` (the default) is byte-identical to
-    before this argument existed.
-
-    `w_total`/`w_std` (soft-anchor Task 3): forwarded straight to
-    `hybrid_field_soft`'s own weights (ignored when `hgb_std` is `None`).
-    `1.0`/`1.0` (the defaults) reproduce `hybrid_field_soft`'s own default
-    weighting."""
-    records = []
-    for stem in stems:
-        d = np.load(osp.join(features_dir, "graphs", f"{stem}.npz"))
-        X, pos_np = d["X"], d["pos"]
-        y_true = d["y_raw"] if "y_raw" in d.files else d["y"]
-        mesh_edges = d["mesh_edges"] if "mesh_edges" in d.files else None
-        g_vec = d["g"] if "g" in d.files else None
-        gnn_seconds = predict_graph_seconds(model, cfg, norm, X, pos_np, mesh_edges,
-                                            g_vec, device=device)
-        gbm_seconds = _gbm_seconds_for_stem(gbm, stem, manifest_rows, manifest_dir,
-                                            features_dir)
-        recomp_info = None
-        if hgb_std is not None:
-            n_std_features = int(getattr(hgb_std, "n_features_in_", 16))
-            x_row = std_feature_row(stem, manifest_rows, manifest_dir, features_dir,
-                                    n_std_features)
-            std_seconds = predict_std_seconds(hgb_std, x_row)
-            hyb, recomp_info = hybrid_field_soft(gnn_seconds, gbm_seconds, std_seconds,
-                                                 lambda_shrink=lambda_shrink,
-                                                 w_total=w_total, w_std=w_std)
-        else:
-            hyb = _max_anchor_hybrid(gnn_seconds, gbm_seconds)
-        row = manifest_rows.get(stem, {})
-        t_ref = float(d["t_ref"]) if "t_ref" in d.files else None
-        rec = {
-            "stem": stem,
-            "family": row.get("alloy_family") or "unknown",
-            "bc": row.get("bc_config") or "none",
-            "dataset": row.get("dataset") or "",
-            "y_true": y_true, "y_hybrid": hyb, "xyz": pos_np,
-            "gbm_seconds": gbm_seconds, "t_ref": t_ref,
-            "pred_total": float(np.max(hyb)),
-        }
-        if recomp_info is not None:
-            rec["recomposition_info"] = recomp_info
-        records.append(rec)
-    return records
-
-
-def _group_metrics(records, key: str) -> dict:
-    groups = defaultdict(list)
-    for r in records:
-        groups[r[key]].append(r)
-    out = {}
-    for gval, recs in groups.items():
-        y = np.concatenate([r["y_true"] for r in recs])
-        p = np.concatenate([r["y_hybrid"] for r in recs])
-        xyz = np.concatenate([r["xyz"] for r in recs])
-        gid = np.concatenate([np.full(r["y_true"].shape[0], i, dtype=np.int64)
-                              for i, r in enumerate(recs)])
-        out[str(gval)] = {
-            "n_stems": len(recs),
-            "n_nodes": int(y.shape[0]),
-            "mae": metrics.mae(y, p),
-            "hotspot_overlap": hotspot_overlap(y, p, xyz, gid),
-            "std_ratio": metrics.pred_std_ratio(y, p, gid),
-        }
-    return out
-
-
-def _overall_metrics(records) -> dict:
-    if not records:
-        return {"n_stems": 0, "n_nodes": 0, "mae": 0.0, "hotspot_overlap": 0.0,
-                "std_ratio": 0.0}
-    y = np.concatenate([r["y_true"] for r in records])
-    p = np.concatenate([r["y_hybrid"] for r in records])
-    xyz = np.concatenate([r["xyz"] for r in records])
-    gid = np.concatenate([np.full(r["y_true"].shape[0], i, dtype=np.int64)
-                          for i, r in enumerate(records)])
-    return {
-        "n_stems": len(records), "n_nodes": int(y.shape[0]),
-        "mae": metrics.mae(y, p),
-        "hotspot_overlap": hotspot_overlap(y, p, xyz, gid),
-        "std_ratio": metrics.pred_std_ratio(y, p, gid),
-    }
-
-
-# --------------------------------------------------------------------------- #
-# No-regression check vs the SP1 hybrid
-# --------------------------------------------------------------------------- #
-# DESIGN CHOICE (documented per Task 8 brief): the no-regression subset is
-# manifest dataset=="A" (all-sand, no BC variation -- the closest analog to
-# the OLD single-alloy regime SP1 was trained/validated on) intersected with
-# the union of interp.test/extrap.test. SP1 itself cannot be run on this
-# subset "as originally trained" (SP1 has no v4 cache -- its own
-# `data/features_v3` cache is a DIFFERENT, single-alloy dataset with
-# different geometries/splits). Instead we run the SP1 GNN + SP1 HGB THE SAME
-# WAY they are normally run (`src.evaluate_gnn.evaluate`): standardize with
-# SP1's OWN norm stats and feed it whatever 13/9-column feature vector it's
-# given. Since `FEATURE_NAMES_V4 == FEATURE_NAMES_V3 + [2 BC cols]` and
-# `FEATURE_NAMES_V3 == FEATURE_NAMES + [4 local-modulus/geodesic cols]`
-# (`src/features.py`), `v4_X[:, :13]` IS the v3 feature vector and
-# `v4_X[:, :9]` IS the v2 feature vector for the SAME geometry -- so this is
-# a legitimate "plug the v4 cache's geometry into the SP1 stack" comparison,
-# not an apples-to-oranges one. CAVEATS: (1) the v4 A-subset's alloy
-# properties vary (synthetic LHS-sampled alloys, not the single fixed alloy
-# SP1's target distribution assumes) -- SP1 has no way to know this, so its
-# scale prediction (HGB) is evaluated fully out-of-its-training-distribution
-# here; a "regression" in this comparison could reflect the alloy variety
-# alone, not the FT2b GNN's shape quality. (2) mesh/graph identity is
-# guaranteed (same X/pos/mesh_edges array), so any MAE delta is NOT a
-# geometry-mismatch artifact. (3) small subset sizes (a handful of A-rows in
-# the actual test splits) make this a noisy signal, not a hard gate by
-# itself -- read alongside per-family/per-bc numbers above.
-def _no_regression_subset(splits_d: dict, manifest_rows: dict) -> list[str]:
-    test_stems = set(splits_d["interp"]["test"]) | set(splits_d["extrap"]["test"])
-    return sorted(s for s in test_stems if manifest_rows.get(s, {}).get("dataset") == "A")
-
-
-def _sp1_hybrid_predict_stem(stem, sp1_model, sp1_cfg, sp1_norm, hgb, hgb_norm,
-                             features_dir, device):
-    d = np.load(osp.join(features_dir, "graphs", f"{stem}.npz"))
-    X, pos_np = d["X"], d["pos"]
-    y_true = d["y_raw"] if "y_raw" in d.files else d["y"]
-    mesh_edges = d["mesh_edges"] if "mesh_edges" in d.files else None
-    X13 = X[:, :N_V3_FEATURES]
-    gnn_seconds = predict_graph_seconds(sp1_model, sp1_cfg, sp1_norm, X13, pos_np,
-                                        mesh_edges, g_vec=None, device=device)
-    Xb = standardize(X[:, :N_V2_FEATURES], hgb_norm)
-    baseline_seconds = predict_seconds(hgb, Xb, hgb_norm)
-    hyb = hybrid_field(gnn_seconds, baseline_seconds)
-    return y_true, hyb, pos_np
-
-
-def _no_regression_check(splits_d, manifest_rows, features_dir, v4_records_by_stem,
-                         sp1_ckpt, sp1_hgb, baseline_features_dir, device,
-                         tolerance: float = 1.15) -> dict:
-    subset = _no_regression_subset(splits_d, manifest_rows)
-    note = ("subset = manifest dataset=='A' (all-sand) stems present in "
-            "interp.test | extrap.test; SP1 hybrid run on the SAME v4-cache "
-            "geometry via v4_X[:, :13] (v3 features) / v4_X[:, :9] (v2 "
-            "features) fed through SP1's own norm stats -- see this module's "
-            "`_no_regression_subset`/`_no_regression_check` docstring for the "
-            "full design rationale + caveats (alloy variety is OOD for SP1's "
-            "HGB; small-subset noise).")
-    if not subset:
-        zero = {"n_stems": 0, "mae": 0.0, "hotspot_overlap": 0.0, "std_ratio": 0.0}
-        return {"n_stems": 0, "v4_hybrid": zero, "sp1_hybrid": dict(zero),
-                "mae_ratio": 0.0, "regression_flag": False,
-                "note": note + " EMPTY SUBSET: no dataset=='A' test stems found."}
-
-    v4_recs = [v4_records_by_stem[s] for s in subset if s in v4_records_by_stem]
-    v4_metrics = _overall_metrics(v4_recs)
-
-    sp1_model, sp1_cfg, sp1_norm = load_ft2b_gnn(sp1_ckpt, device)
-    hgb = joblib.load(sp1_hgb) if isinstance(sp1_hgb, (str, os.PathLike)) else sp1_hgb
-    hgb_norm = load_norm(baseline_features_dir)
-
-    ys, ps, xyzs = [], [], []
-    for stem in subset:
-        y_true, hyb, pos_np = _sp1_hybrid_predict_stem(
-            stem, sp1_model, sp1_cfg, sp1_norm, hgb, hgb_norm, features_dir, device)
-        ys.append(y_true); ps.append(hyb); xyzs.append(pos_np)
-    y = np.concatenate(ys); p = np.concatenate(ps); xyz = np.concatenate(xyzs)
-    gid = np.concatenate([np.full(ys[i].shape[0], i, dtype=np.int64) for i in range(len(ys))])
-    sp1_metrics = {
-        "n_stems": len(subset), "n_nodes": int(y.shape[0]),
-        "mae": metrics.mae(y, p),
-        "hotspot_overlap": hotspot_overlap(y, p, xyz, gid),
-        "std_ratio": metrics.pred_std_ratio(y, p, gid),
-    }
-
-    mae_ratio = (v4_metrics["mae"] / sp1_metrics["mae"]) if sp1_metrics["mae"] > 1e-9 else 0.0
-    return {
-        "n_stems": len(subset),
-        "v4_hybrid": v4_metrics,
-        "sp1_hybrid": sp1_metrics,
-        "mae_ratio": float(mae_ratio),
-        "regression_flag": bool(v4_metrics["mae"] > sp1_metrics["mae"] * tolerance),
-        "note": note,
-    }
-
-
-# --------------------------------------------------------------------------- #
-# Chvorinov-scaling sanity check
-# --------------------------------------------------------------------------- #
-def _linfit(x: np.ndarray, y: np.ndarray) -> dict:
-    """Least-squares y = slope*x + intercept + R^2. Always finite: falls
-    back to a degenerate-but-finite {0,0,0} fit when there are fewer than 2
-    points or `x` is (numerically) constant, rather than raising or NaN'ing."""
-    n = len(x)
-    if n < 2 or np.ptp(x) < 1e-9:
-        return {"n": int(n), "slope": 0.0, "intercept": 0.0, "r2": 0.0}
-    slope, intercept = np.polyfit(x, y, 1)
-    y_hat = slope * x + intercept
-    ss_res = float(np.sum((y - y_hat) ** 2))
-    ss_tot = float(np.sum((y - np.mean(y)) ** 2))
-    r2 = 1.0 - ss_res / ss_tot if ss_tot > 1e-12 else 0.0
-    return {"n": int(n), "slope": float(slope), "intercept": float(intercept), "r2": float(r2)}
-
-
-def _chvorinov_check(records: list[dict]) -> dict:
-    """Regress the model's predicted total solidification time (`pred_total`
-    = max of the per-node hybrid field, per stem) against each stem's cached
-    analytic Chvorinov reference `t_ref = C_family * (V/A)^2`
-    (`src/build_dataset.py::_finalize_row_stage2`). A clean per-family fit
-    (slope near 1, high R^2) means the scale head still tracks Chvorinov
-    scaling after fine-tuning; reported per family AND overall."""
-    usable = [r for r in records if r["t_ref"] is not None and np.isfinite(r["t_ref"])]
-    groups = defaultdict(list)
-    for r in usable:
-        groups[r["family"]].append(r)
-    out = {}
-    for fam, recs in groups.items():
-        x = np.array([r["t_ref"] for r in recs], dtype=np.float64)
-        y = np.array([r["pred_total"] for r in recs], dtype=np.float64)
-        out[str(fam)] = _linfit(x, y)
-    x_all = np.array([r["t_ref"] for r in usable], dtype=np.float64)
-    y_all = np.array([r["pred_total"] for r in usable], dtype=np.float64)
-    out["overall"] = _linfit(x_all, y_all)
-    return out
-
-
-# --------------------------------------------------------------------------- #
-# Top-level entry point
-# --------------------------------------------------------------------------- #
-def _load_splits(splits) -> dict:
-    if isinstance(splits, (str, os.PathLike)):
-        with open(splits) as f:
-            return json.load(f)
-    return splits
-
-
-def evaluate_ft2b(ckpt, gbm, features_dir, manifest, splits,
-                  sp1_ckpt=DEFAULT_SP1_CKPT, sp1_hgb=DEFAULT_SP1_HGB,
-                  baseline_features_dir=DEFAULT_BASELINE_FEATURES_DIR,
-                  out_path=DEFAULT_OUT_PATH,
-                  variance_collapse_threshold: float = 0.9,
-                  no_regression_tolerance: float = 1.15,
-                  hgb_std=None,
-                  lambda_shrink: float = 1.0,
-                  w_total: float = 1.0,
-                  w_std: float = 1.0,
-                  device=None) -> dict:
-    """Evaluate a trained FT2b v4 checkpoint + refit GBM hybrid predictor.
-
-    Args:
-        ckpt: path to the v4 GNN checkpoint (`{model, cfg, norm}`), or an
-            already-loaded checkpoint dict.
-        gbm: path to the refit multi-alloy GBM (`src.gbm_ft2b.refit_gbm`'s
-            `out_path`), or an already-loaded estimator.
-        features_dir: v4 feature cache directory (`graphs/<stem>.npz`).
-        manifest: path to the DOE manifest CSV (`run_id`==stem; needs
-            `dataset`/`alloy_family`/`bc_config`/`material_id`/
-            `mold_effusivity` columns).
-        splits: path to `features_dir/splits.json`-shaped JSON, or an
-            already-loaded `{"interp": {"train","test"}, "extrap":
-            {"train","test"}}` dict.
-        sp1_ckpt / sp1_hgb / baseline_features_dir: the SP1 production
-            hybrid stack, for the no-regression comparison.
-        out_path: where to write the result as JSON (set `None`/falsy to
-            skip writing).
-        variance_collapse_threshold: any per-family/per-bc `std_ratio` below
-            this trips `variance_collapse_flag`.
-        no_regression_tolerance: the v4 hybrid may be at most this factor
-            worse (MAE) than the SP1 hybrid on the all-sand subset before
-            `no_regression.regression_flag` trips.
-        hgb_std (variance-recomposition Task 3): an optional path to (or
-            already-loaded) GBM std-head model
-            (`src.gbm_ft2b.refit_gbm_std`'s output). `None` (the default)
-            keeps every v4 hybrid composition in this module on the
-            quantile-anchor (`src.hybrid.hybrid_field_max_anchor`) --
-            byte-identical to before this argument existed. When given,
-            every v4 composition (`_evaluate_stems_v4`, and therefore every
-            per-family/per-bc metric, the v4 side of `no_regression`, and
-            `chvorinov`'s `pred_total`) uses the soft-anchor (soft-anchor
-            Task 3; was affine) recomposition instead
-            (`src.hybrid.hybrid_field_soft`); the SP1 side of
-            `no_regression` is unaffected (SP1 has no std-head analog).
-        lambda_shrink (variance-recomposition amendment 2): forwarded
-            straight to `_evaluate_stems_v4`/`hybrid_field_soft` (ignored
-            when `hgb_std` is `None`). `1.0` (the default) is byte-identical
-            to before this argument existed.
-        w_total / w_std (soft-anchor Task 3): forwarded straight to
-            `_evaluate_stems_v4`/`hybrid_field_soft`'s own weights (ignored
-            when `hgb_std` is `None`). `1.0`/`1.0` (the defaults) reproduce
-            `hybrid_field_soft`'s own default weighting.
-        device: "cuda"/"cpu"; default picks cuda if available.
-
-    Returns: nested dict with `interp`/`extrap` (each `per_alloy_family`,
-        `per_bc_config`, `overall`), `no_regression`, `chvorinov`,
-        `variance_collapse_flag`, `recomposition` ("soft" iff `hgb_std`
-        was given, else "anchor").
-    """
-    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-    model, cfg, norm = load_ft2b_gnn(ckpt, device)
-    gbm_model = joblib.load(gbm) if isinstance(gbm, (str, os.PathLike)) else gbm
-    hgb_std_model = joblib.load(hgb_std) if isinstance(hgb_std, (str, os.PathLike)) else hgb_std
-    manifest_dir = osp.dirname(osp.abspath(manifest)) or "."
-    manifest_rows = _load_manifest_rows(manifest)
-    splits_d = _load_splits(splits)
-
-    result: dict = {}
-    records_by_stem: dict = {}
-    for split_name in ("interp", "extrap"):
-        test_stems = splits_d[split_name]["test"]
-        records = _evaluate_stems_v4(model, cfg, norm, gbm_model, test_stems,
-                                     manifest_rows, manifest_dir, features_dir, device,
-                                     hgb_std=hgb_std_model, lambda_shrink=lambda_shrink,
-                                     w_total=w_total, w_std=w_std)
-        for r in records:
-            records_by_stem[r["stem"]] = r
-        result[split_name] = {
-            "n_stems": len(records),
-            "per_alloy_family": _group_metrics(records, "family"),
-            "per_bc_config": _group_metrics(records, "bc"),
-            "overall": _overall_metrics(records),
-        }
-
-    variance_collapse = False
-    for split_res in result.values():
-        for grouping in ("per_alloy_family", "per_bc_config"):
-            for m in split_res[grouping].values():
-                if m["std_ratio"] < variance_collapse_threshold:
-                    variance_collapse = True
-    result["variance_collapse_flag"] = variance_collapse
-
-    result["no_regression"] = _no_regression_check(
-        splits_d, manifest_rows, features_dir, records_by_stem,
-        sp1_ckpt, sp1_hgb, baseline_features_dir, device,
-        tolerance=no_regression_tolerance)
-
-    result["chvorinov"] = _chvorinov_check(list(records_by_stem.values()))
-
-    result["recomposition"] = "soft" if hgb_std_model is not None else "anchor"
-
-    if out_path:
-        out_dir = osp.dirname(out_path)
-        if out_dir:
-            os.makedirs(out_dir, exist_ok=True)
-        with open(out_path, "w") as f:
-            json.dump(result, f, indent=2)
-
-    return result

src/evaluate_gnn.py

@@ -1,150 +0,0 @@
-"""Evaluate a trained GNN on cache stems: GNN-only and (optionally) hybrid,
-optional rotation TTA, per-split comparison table. No augmentation leakage:
-TTA rotations are seeded and averaged in seconds-space."""
-import argparse
-import json
-import os.path as osp
-
-import numpy as np
-import torch
-
-from src.graph_build import (knn_edge_index, edge_attr, random_rotation_matrix, rotate,
-                             union_edge_index, union_edge_attr)
-from src.model_gnn import SolidGNN
-from src.dataload import load_norm, standardize
-from src.splits import load_splits
-from src.baseline import predict_seconds
-from src.hybrid import hybrid_field
-from src import metrics
-
-# The HGB baseline was trained on the 9 v2 features = first 9 columns of any
-# newer feature version (ordering contract, SP1 spec §3).
-N_BASELINE_FEATURES = 9
-SPLIT_NAMES = ("interpolation", "family", "fully_cooled")
-
-
-def predict_seconds_gnn(pred_std_log, norm) -> np.ndarray:
-    pred_log = pred_std_log * norm["y_log_std"] + norm["y_log_mean"]
-    return np.clip(np.expm1(pred_log), 0.0, None)
-
-
-def load_gnn(ckpt, device):
-    state = torch.load(ckpt, map_location=device, weights_only=False)
-    cfg, norm = state["cfg"], state["norm"]
-    model = SolidGNN(node_in=cfg.get("node_in", 9), edge_in=cfg.get("edge_in", 4),
-                     hidden=cfg["hidden"], n_blocks=cfg["n_blocks"],
-                     two_head=cfg.get("two_head", False)).to(device)
-    model.load_state_dict(state["model"])
-    model.eval()
-    return model, cfg, norm
-
-
-def _build_edges(pos, cfg, mesh_edges=None):
-    if cfg.get("graph", "knn") == "union":
-        if mesh_edges is None:
-            raise ValueError("checkpoint expects a union graph but mesh_edges is missing")
-        ei, is_mesh = union_edge_index(pos, mesh_edges, k=cfg["k"])
-        return ei, union_edge_attr(pos, ei, is_mesh)
-    ei = knn_edge_index(pos, k=cfg["k"])
-    return ei, edge_attr(pos, ei)
-
-
-def predict_graph_seconds(model, cfg, norm, x, pos, mesh_edges=None, device="cpu", tta=0):
-    """Mean GNN prediction (seconds) over identity + `tta` seeded random rotations."""
-    outs = []
-    rots = [None] + [random_rotation_matrix(seed=s) for s in range(tta)]
-    zeros = torch.zeros(x.shape[0], dtype=torch.long, device=device)
-    for R in rots:
-        p = pos if R is None else rotate(pos, R)
-        ei, ea = _build_edges(p, cfg, mesh_edges)
-        with torch.no_grad():
-            yn, _ = model(x.to(device), ei.to(device), ea.to(device), zeros)
-        outs.append(predict_seconds_gnn(yn.cpu().numpy(), norm))
-    return np.mean(outs, axis=0)
-
-
-def _metric_dict(y, p, xyz, gid) -> dict:
-    return {
-        "mae": metrics.mae(y, p),
-        "rmse": metrics.rmse(y, p),
-        "median_rel_error": metrics.median_rel_error(y, p),
-        "pred_std_ratio": metrics.pred_std_ratio(y, p, gid),
-        "hotspot_accuracy": metrics.hotspot_accuracy(y, p, xyz, gid),
-    }
-
-
-def evaluate(ckpt, features_dir, test_stems, device=None, tta=0,
-             hgb_path=None, baseline_features_dir="data/features_v2") -> dict:
-    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-    model, cfg, norm = load_gnn(ckpt, device)
-    fmean = torch.tensor(norm["feat_mean"], dtype=torch.float32)
-    fstd = torch.tensor(norm["feat_std"], dtype=torch.float32)
-    hgb = hgb_norm = None
-    if hgb_path is not None:
-        import joblib
-        hgb = joblib.load(hgb_path)
-        hgb_norm = load_norm(baseline_features_dir)
-
-    ys, preds, hyb, xyz, gid = [], [], [], [], []
-    for i, stem in enumerate(test_stems):
-        d = np.load(osp.join(features_dir, "graphs", f"{stem}.npz"))
-        X, y, pos_np = d["X"], d["y"], d["pos"]
-        x = (torch.tensor(X, dtype=torch.float32) - fmean) / fstd
-        pos = torch.tensor(pos_np, dtype=torch.float32)
-        me = (torch.tensor(d["mesh_edges"], dtype=torch.long)
-              if "mesh_edges" in d.files else None)
-        p = predict_graph_seconds(model, cfg, norm, x, pos, me, device=device, tta=tta)
-        preds.append(p); ys.append(y); xyz.append(pos_np)
-        gid.append(np.full(y.shape[0], i, dtype=np.int64))
-        if hgb is not None:
-            Xb = standardize(X[:, :N_BASELINE_FEATURES], hgb_norm)
-            hyb.append(hybrid_field(p, predict_seconds(hgb, Xb, hgb_norm)))
-
-    y = np.concatenate(ys); p = np.concatenate(preds)
-    xyz = np.concatenate(xyz); gid = np.concatenate(gid)
-    out = _metric_dict(y, p, xyz, gid)
-    if hgb is not None:
-        out["hybrid"] = _metric_dict(y, np.concatenate(hyb), xyz, gid)
-    return out
-
-
-def format_table(results) -> str:
-    lines = [f"{'split':<15}{'model':<8}{'MAE':>8}{'RMSE':>8}{'hot-spot':>10}{'std-ratio':>11}"]
-    for split, res in results.items():
-        rows = [("gnn", res)] + ([("hybrid", res["hybrid"])] if "hybrid" in res else [])
-        for label, m in rows:
-            lines.append(f"{split:<15}{label:<8}{m['mae']:>8.2f}{m['rmse']:>8.2f}"
-                         f"{m['hotspot_accuracy']:>10.3f}{m['pred_std_ratio']:>11.3f}")
-    return "\n".join(lines)
-
-
-def evaluate_all_splits(ckpt, features_dir, splits_path, hgb_path=None, tta=0,
-                        device=None, out_json=None,
-                        baseline_features_dir="data/features_v2") -> dict:
-    splits = load_splits(splits_path)
-    results = {}
-    for name in SPLIT_NAMES:
-        results[name] = evaluate(ckpt, features_dir, splits[name]["test"], device=device,
-                                 tta=tta, hgb_path=hgb_path,
-                                 baseline_features_dir=baseline_features_dir)
-    if out_json:
-        with open(out_json, "w") as f:
-            json.dump(results, f, indent=2)
-    print(format_table(results))
-    return results
-
-
-if __name__ == "__main__":
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--ckpt", required=True)
-    ap.add_argument("--features-dir", default="data/features_v2")
-    ap.add_argument("--splits", default=None, help="default: <features-dir>/splits.json")
-    ap.add_argument("--hgb", default=None)
-    ap.add_argument("--baseline-features-dir", default="data/features_v2")
-    ap.add_argument("--tta", type=int, default=0)
-    ap.add_argument("--out", default=None)
-    a = ap.parse_args()
-    evaluate_all_splits(a.ckpt, a.features_dir,
-                        a.splits or osp.join(a.features_dir, "splits.json"),
-                        hgb_path=a.hgb, tta=a.tta, out_json=a.out,
-                        baseline_features_dir=a.baseline_features_dir)

src/features.py

@@ -1,88 +1,33 @@
-"""Assemble per-node physics features (v2) + target for one casting.
+"""Node features and the alloy/mold conditioning vector.
 
-v2 vs v1: distance-to-cooling is GEODESIC (through the material); the degenerate
-engineered ratios (d_cool_sq, d_cool_over_dsurf) are dropped; node positions are
-returned for downstream graph construction. All features mm, invariant, from mesh
-+ facet labels only.
+Every feature comes from the mesh and its surface tags ("Cooling", optionally
+"Adiabatic"), in millimetres. Nothing depends on node numbering, file names or
+absolute position.
 """
-import logging
-
 import numpy as np
 
-from src.geometry import (triangle_areas, total_area, mesh_volume,
-                          nearest_distance, boundary_node_mask,
-                          edges_from_tets, geodesic_distance, local_modulus,
-                          geodesic_diameter_two_sweep)
-from src.targets import solidification_time, global_time
-
-logger = logging.getLogger(__name__)
+from src.geometry import (total_area, mesh_volume, nearest_distance, boundary_node_mask,
+                          edges_from_tets, geodesic_distance, local_modulus)
 
 FEATURE_NAMES = [
     "geodesic_d_cool", "d_surf", "d_adiab",
     "is_cooling", "is_adiabatic", "is_interior",
     "volume_mm3", "surface_area_mm2", "cooling_area_frac",
-]
-
-FEATURE_NAMES_V3 = FEATURE_NAMES + [
     "geodesic_d_adiab",
     "local_modulus_r5", "local_modulus_r15", "local_modulus_r40",
+    "boundary_extraction", "boundary_heat_input",
 ]
 LOCAL_MODULUS_RADII = (5.0, 15.0, 40.0)
 
-# v4: v3's 13 columns UNCHANGED (X[:, :9] == v2, X[:, :13] == v3 -- the HGB
-# baseline + backward-compat contract), with two continuous BC channels
-# APPENDED at the end (replacing is_cooling/is_adiabatic/is_interior as MODEL
-# input signal, though those binary tags stay in the first 9 columns as-is).
-FEATURE_NAMES_V4 = FEATURE_NAMES_V3 + ["boundary_extraction", "exothermic_input"]
-
-# v5: v4's 15 columns UNCHANGED, with five multi-scale part-normalized BC
-# influence channels APPENDED -- exo_mid/exo_long extend the v4
-# exothermic_input ("exo_short") reach to 0.15*D_geo/0.40*D_geo, and
-# chill_short/chill_mid/chill_long give chills the same three-scale
-# treatment (chills previously only had the surface-local
-# boundary_extraction). See BC_REACH_MID_FRAC/BC_REACH_LONG_FRAC above.
-FEATURE_NAMES_V5 = FEATURE_NAMES_V4 + [
-    "exo_mid", "exo_long", "chill_short", "chill_mid", "chill_long",
-]
-
-# Geodesic-diffusion length scale (mm) for the exothermic_input channel: the
-# sleeve's calibrated heat flux q0 decays as exp(-d_geo/lambda) away from the
-# adiab_only (BC-patch) source nodes, instead of being painted only on those
-# 1-2% of nodes as a sparse indicator (see docs/superpowers exo-fix design).
-# This is also reused, unchanged, as the "short" reach scale for the v5
-# chill_short channel below (bc-reach spec: same short-range mesh-scale
-# decay length, not part-size-normalized).
-EXO_DIFFUSION_LAMBDA_MM = 15.0
-
-# v5 "mid"/"long" BC-influence reach fractions -- PART-NORMALIZED (of
-# D_geo = the part's GEODESIC DIAMETER: a two-sweep-Dijkstra estimate over
-# the mesh-edge graph -- `geodesic_diameter_two_sweep` -- patch-INDEPENDENT
-# and through-metal), unlike the fixed EXO_DIFFUSION_LAMBDA_MM above.
-# CORRECTION (2026-08-02, Task-1 review round 1, F1/F3): D_geo was
-# originally (wrongly) `max(geodesic_d_cool)` -- that is max WALL DEPTH
-# (~35-94mm on real parts, and BC-PATCH-DEPENDENT since it depends on which
-# facets are tagged Cooling), not part size; it made L_long 6-17x too
-# short (shorter even than the v4 fixed 15mm scale on some rows), inverting
-# the intended short<mid<long ladder. The geodesic diameter is a single
-# per-part constant, identical across every BC row of the same geometry.
-# Motivation (bc-reach investigation): the fixed 15mm short scale decays to
-# ~0 by ~140mm, but a real exo sleeve's influence on a 600mm-class casting
-# extends past 400mm -- normalizing by the part's own geodesic diameter
-# lets the encoding transfer across part sizes instead of a single global
-# mm constant under/over-reaching.
-BC_REACH_MID_FRAC = 0.15
-BC_REACH_LONG_FRAC = 0.40
-
-# Global dimensionless alloy+mould vector (broadcast to every node by the
-# cache builder / FT2b; computed here RAW -- not standardized).
+# Alloy/mold conditioning vector (raw, not standardized).
 G_VECTOR_NAMES = ["Ste", "dTf_norm", "alpha_ratio", "superheat_norm", "mold_eff_norm"]
-T_REF = 1000.0          # deg C, normalizes temperature-difference channels
-ALPHA_REF = 1e-5        # m^2/s, ballpark alloy thermal diffusivity
-MOLD_EFF_REF = 1250.0   # W*s^0.5/(m^2*K), green-sand mold effusivity baseline
+T_REF = 1000.0          # deg C, normalizes temperature differences
+ALPHA_REF = 1e-5        # m^2/s, typical alloy thermal diffusivity
+MOLD_EFF_REF = 1250.0   # W*s^0.5/(m^2*K), reference mold effusivity
 
 
 def alloy_g_vector(props: dict) -> np.ndarray:
-    """Global dimensionless alloy+mould vector `g` (raw, not standardized).
+    """Five dimensionless numbers describing alloy and mold.
 
     `props` keys: k, rho, cp, L_kJ_kg, TL, TS, Tpour, mold_effusivity."""
     k = float(props["k"]); rho = float(props["rho"]); cp = float(props["cp"])
@@ -95,119 +40,12 @@ def alloy_g_vector(props: dict) -> np.ndarray:
     superheat_norm = (Tpour - TL) / T_REF
     mold_eff_norm = b_mold / MOLD_EFF_REF
     return np.array([Ste, dTf_norm, alpha_ratio, superheat_norm, mold_eff_norm],
-                     dtype=np.float64)
-
-
-def chvorinov_t_ref(volume_mm3, area_mm2, C_material) -> float:
-    """Chvorinov total-solidification-time estimate: C * (V/A)^2."""
-    return float(C_material) * (float(volume_mm3) / float(area_mm2)) ** 2
-
-
-def geodesic_diffuse(source_values: np.ndarray, length_mm: float, N: int,
-                     edges: np.ndarray, edge_len: np.ndarray) -> np.ndarray:
-    """Multi-source geodesic exponential-decay diffusion, factored out of
-    the original inline `exothermic_input` construction.
-
-    For every node: the max over distinct nonzero source magnitudes of
-    (magnitude * exp(-geodesic distance to the nearest node carrying that
-    magnitude / length_mm)). Distinct magnitudes are grouped and Dijkstra'd
-    independently, then combined by elementwise max -- for the single
-    uniform-magnitude source group this codebase actually uses (every BC
-    source node carries the same scalar peak), this reduces EXACTLY to
-    `peak * exp(-d_geo/length_mm)` via one ordinary multi-source Dijkstra
-    call, byte-identical to the original `exothermic_input` formula --
-    PROVIDED that peak is non-negative. The `out = np.maximum(out, ...)`
-    combine is an implicit >= 0 clip: a genuinely negative source magnitude
-    would get clamped toward 0 instead of decaying to a negative value (a
-    node's own zero-init `out` beats any negative contribution). No caller
-    in this codebase produces a negative peak today (`exo_q0/1e6` and
-    `(chill_eff-mold_eff)/mold_eff` are both >= 0 in every resolved spec),
-    so this is currently unreachable, not a live bug -- documented so a
-    future negative-magnitude source doesn't rediscover it the hard way.
-
-    Returns all-zeros WITHOUT invoking the Dijkstra solver when there are
-    no nonzero sources or `length_mm <= 0` (cheap all-sand / degenerate
-    path -- `np.exp(-inf) == 0.0` would otherwise handle disconnected mesh
-    components fine, but a zero/negative length is not physically
-    meaningful and must not divide-by-zero).
-    """
-    source_values = np.asarray(source_values, dtype=np.float64)
-    src_idx = np.flatnonzero(source_values != 0.0)
-    if src_idx.size == 0 or length_mm <= 0:
-        return np.zeros(N, dtype=np.float32)
-    out = np.zeros(N, dtype=np.float64)
-    for val in np.unique(source_values[src_idx]):
-        grp = src_idx[source_values[src_idx] == val]
-        d = geodesic_distance(N, edges, edge_len, grp)
-        # np.exp(-inf) == 0.0 handles disconnected components
-        out = np.maximum(out, val * np.exp(-d / length_mm))
-    return out.astype(np.float32)
-
-
-def compute_bc_reach_channels(N: int, edges: np.ndarray, edge_len: np.ndarray,
-                              adiab_only: np.ndarray, bc_config: str,
-                              mold_eff: float, chill_eff: float,
-                              exo_q0: float) -> np.ndarray:
-    """The five v5 multi-scale part-normalized BC-influence channels:
-    `[exo_mid, exo_long, chill_short, chill_mid, chill_long]`, shape (N, 5)
-    float32.
+                    dtype=np.float64)
 
-    `D_geo` (the part-normalizing length) is the part's GEODESIC DIAMETER --
-    `geodesic_diameter_two_sweep(N, edges, edge_len, start=0)`, a two-sweep
-    Dijkstra estimate over the SAME mesh-edge graph the other geodesic
-    features use. This is patch-INDEPENDENT (unlike `geodesic_d_cool`,
-    whose max is wall DEPTH, not part size -- see the CORRECTION note above
-    `BC_REACH_MID_FRAC`) and identical across every BC row of the same
-    geometry. exo_mid/exo_long reuse the exo_sleeve's `adiab_only` source
-    patch and peak (`exo_q0/1e6`, matching `exothermic_input`'s "exo_short"
-    scale) at `BC_REACH_MID_FRAC`/`BC_REACH_LONG_FRAC` * D_geo.
-    chill_short/chill_mid/chill_long give the chills bc_config the same
-    three-scale treatment, sourced at the chill contact nodes (`adiab_only`)
-    with peak magnitude `(chill_eff - mold_eff) / mold_eff` -- chill_short
-    reuses the fixed `EXO_DIFFUSION_LAMBDA_MM` short scale, chill_mid/
-    chill_long use the same part-normalized fractions as the exo channels.
 
-    D_geo (and therefore the two extra Dijkstra sweeps) is computed ONLY
-    when `bc_config` is `"exo_sleeve"` or `"chills"` -- an all-sand
-    (`"none"`) or `"insulating_sleeve"` part returns all-zeros without any
-    solver call, same contract as `exothermic_input`/`geodesic_diffuse`. A
-    degenerate/disconnected mesh (D_geo <= 0) logs a warning and leaves the
-    mid/long channels at zero (short-scale channels are unaffected, since
-    they don't depend on D_geo).
-    """
-    out = np.zeros((N, 5), dtype=np.float32)
-    if bc_config not in ("exo_sleeve", "chills"):
-        return out
-
-    d_geo = geodesic_diameter_two_sweep(N, edges, edge_len, start=0)
-    if d_geo <= 0:
-        logger.warning(
-            "compute_bc_reach_channels: D_geo (geodesic diameter) <= 0 for "
-            "bc_config=%r -- degenerate/disconnected mesh, zeroing the "
-            "part-normalized mid/long BC-reach channels", bc_config)
-    lambda_mid = BC_REACH_MID_FRAC * d_geo
-    lambda_long = BC_REACH_LONG_FRAC * d_geo
-
-    if bc_config == "exo_sleeve":
-        src = np.zeros(N, dtype=np.float64)
-        src[adiab_only] = exo_q0 / 1e6
-        out[:, 0] = geodesic_diffuse(src, lambda_mid, N, edges, edge_len)
-        out[:, 1] = geodesic_diffuse(src, lambda_long, N, edges, edge_len)
-
-    if bc_config == "chills":
-        chill_val = (chill_eff - mold_eff) / mold_eff if mold_eff != 0 else 0.0
-        src = np.zeros(N, dtype=np.float64)
-        src[adiab_only] = chill_val
-        out[:, 2] = geodesic_diffuse(src, EXO_DIFFUSION_LAMBDA_MM, N, edges, edge_len)
-        out[:, 3] = geodesic_diffuse(src, lambda_mid, N, edges, edge_len)
-        out[:, 4] = geodesic_diffuse(src, lambda_long, N, edges, edge_len)
-
-    return out
-
-
-def _node_features(mesh, version=2, spec=None):
-    """Compute the (N,9/13/15/20) feature matrix (v2/v3/v4/v5) + pos (mm) +
-    mesh edges for a mesh. No target."""
+def node_features(mesh, mold_effusivity: float) -> dict:
+    """The (N, 15) float32 feature matrix `X`, node positions `pos` (mm) and
+    the unique tetrahedral mesh edges `mesh_edges` for one casting."""
     xyz = mesh.node_xyz
     N = xyz.shape[0]
     cool_tri = mesh.tris.get("Cooling", np.empty((0, 3), np.int64))
@@ -240,6 +78,18 @@ def _node_features(mesh, version=2, spec=None):
     cool_frac = area_cool / area_total if area_total > 0 else 0.0
     volume = mesh_volume(xyz, mesh.tets)
 
+    g_adiab = geodesic_distance(N, edges, edge_len, np.flatnonzero(adiab_mask))
+    g_adiab = np.where(np.isinf(g_adiab), d_adiab, g_adiab)
+    lm = local_modulus(xyz, mesh.tets, cool_tri, edges, edge_len, LOCAL_MODULUS_RADII)
+
+    # Boundary channels. Surface nodes carry the mold effusivity (how fast the
+    # mold draws heat); interior nodes carry 0. The heat-input channel is for
+    # heat sources on the boundary, and a plain sand-mold surface has none.
+    boundary_extraction = np.zeros(N, dtype=np.float32)
+    boundary_extraction[cool_mask] = mold_effusivity
+    boundary_extraction[adiab_only] = mold_effusivity
+    boundary_heat_input = np.zeros(N, dtype=np.float32)
+
     cols = {
         "geodesic_d_cool": g_cool, "d_surf": d_surf, "d_adiab": d_adiab,
         "is_cooling": cool_mask.astype(np.float32),
@@ -248,87 +98,12 @@ def _node_features(mesh, version=2, spec=None):
         "volume_mm3": np.full(N, volume),
         "surface_area_mm2": np.full(N, area_total),
         "cooling_area_frac": np.full(N, cool_frac),
+        "geodesic_d_adiab": g_adiab,
+        "local_modulus_r5": lm[:, 0],
+        "local_modulus_r15": lm[:, 1],
+        "local_modulus_r40": lm[:, 2],
+        "boundary_extraction": boundary_extraction,
+        "boundary_heat_input": boundary_heat_input,
     }
-    if version >= 3:
-        g_adiab = geodesic_distance(N, edges, edge_len, np.flatnonzero(adiab_mask))
-        g_adiab = np.where(np.isinf(g_adiab), d_adiab, g_adiab)
-        lm = local_modulus(xyz, mesh.tets, cool_tri, edges, edge_len, LOCAL_MODULUS_RADII)
-        cols.update({
-            "geodesic_d_adiab": g_adiab,
-            "local_modulus_r5": lm[:, 0],
-            "local_modulus_r15": lm[:, 1],
-            "local_modulus_r40": lm[:, 2],
-        })
-    if version >= 4:
-        if spec is None:
-            raise ValueError("_node_features(version=4) requires a `spec` dict "
-                              "(BC info: mold_effusivity, bc_config, "
-                              "chill_effusivity, sleeve_MEF, exo_q0)")
-        mold_eff = float(spec["mold_effusivity"])
-        bc_config = spec["bc_config"]
-        chill_eff = float(spec.get("chill_effusivity", 0.0))
-        sleeve_mef = float(spec.get("sleeve_MEF", 0.0))
-        exo_q0 = float(spec.get("exo_q0", 0.0))
-
-        # Mirrors tools/doe/fenicsx_solver.py's tag convention: tag "Cooling"
-        # is always the ordinary sand-mold Robin term; tag "Adiabatic" is the
-        # BC-patch footprint, reinterpreted per bc_config. Nodes shared
-        # between the two triangle groups follow the same cooling-takes-
-        # priority convention as `is_adiabatic` (adiab_only).
-        boundary_extraction = np.zeros(N, dtype=np.float32)
-        boundary_extraction[cool_mask] = mold_eff
-        if bc_config == "chills":
-            adiab_val = chill_eff
-        elif bc_config in ("insulating_sleeve", "exo_sleeve"):
-            adiab_val = mold_eff / max(sleeve_mef ** 2, 1.0)
-        else:  # "none" -- BC patch treated as ordinary mold (matches solver)
-            adiab_val = mold_eff
-        boundary_extraction[adiab_only] = adiab_val
-
-        exothermic_input = np.zeros(N, dtype=np.float32)
-        if bc_config == "exo_sleeve":
-            exo_src_vals = np.zeros(N, dtype=np.float64)
-            exo_src_vals[adiab_only] = exo_q0 / 1e6
-            exothermic_input = geodesic_diffuse(exo_src_vals, EXO_DIFFUSION_LAMBDA_MM,
-                                                N, edges, edge_len)
-
-        cols.update({
-            "boundary_extraction": boundary_extraction,
-            "exothermic_input": exothermic_input,
-        })
-
-        if version >= 5:
-            bc_reach = compute_bc_reach_channels(
-                N, edges, edge_len, adiab_only, bc_config,
-                mold_eff, chill_eff, exo_q0)
-            cols.update({
-                "exo_mid": bc_reach[:, 0], "exo_long": bc_reach[:, 1],
-                "chill_short": bc_reach[:, 2], "chill_mid": bc_reach[:, 3],
-                "chill_long": bc_reach[:, 4],
-            })
-
-    names = (FEATURE_NAMES_V5 if version >= 5 else
-             FEATURE_NAMES_V4 if version >= 4 else
-             FEATURE_NAMES_V3 if version >= 3 else FEATURE_NAMES)
-    X = np.stack([cols[n] for n in names], axis=1).astype(np.float32)
-    return X, xyz.astype(np.float32), edges.astype(np.int64)
-
-
-def compute_features_inference(mesh, version=2, spec=None) -> dict:
-    X, pos, edges = _node_features(mesh, version, spec=spec)
-    return {"X": X, "pos": pos, "mesh_edges": edges}
-
-
-def compute_features(mesh, temps, times, version=2, spec=None) -> dict:
-    X, pos, edges = _node_features(mesh, version, spec=spec)
-    y = solidification_time(temps, times)
-    out = {"X": X, "y": y.astype(np.float32), "y_global": global_time(y),
-           "pos": pos, "mesh_edges": edges}
-    if version >= 4:
-        g = alloy_g_vector(spec)
-        volume = float(X[0, FEATURE_NAMES.index("volume_mm3")])
-        area = float(X[0, FEATURE_NAMES.index("surface_area_mm2")])
-        t_ref = chvorinov_t_ref(volume, area, spec["C_material"])
-        y_nondim = y.astype(np.float64) / max(t_ref, 1e-9)
-        out.update({"g": g, "t_ref": t_ref, "y_nondim": y_nondim.astype(np.float32)})
-    return out
+    X = np.stack([cols[n] for n in FEATURE_NAMES], axis=1).astype(np.float32)
+    return {"X": X, "pos": xyz.astype(np.float32), "mesh_edges": edges.astype(np.int64)}

src/gbm_ft2b.py

@@ -1,598 +0,0 @@
-"""FT2b Task 6: multi-alloy GBM scale-head refit.
-
-The production predictor is a HYBRID: GNN field *shape* x a per-graph
-*scale* from a HistGradientBoostingRegressor (see `src/hybrid.py`,
-`runs/baseline_v2/hgb_A.joblib`). That GBM was trained on `src/features.py`
-v2 physics features from the SINGLE-alloy dataset and knows nothing about
-material properties, so it cannot generalize its scale prediction across
-alloys/mold systems.
-
-This module refits a NEW scale-head GBM directly on raw material properties
-+ geometry scale + mold-system effusivity -- exactly the quantities that
-determine Chvorinov-style total solidification time -- resolved per training
-run_id from the FT2a/FT2b DOE manifest via the SAME resolution path used to
-build the v4 feature cache (`src.build_dataset._resolve_material_props_v4`).
-
-Feature row layout (see FEATURE_NAMES): [k, rho, cp, L_kJ_kg, TL, TS, Tpour,
-V_over_A, mold_effusivity] (9 material/geometry columns) + [bc_none, bc_chills,
-bc_insulating_sleeve, bc_exo_sleeve] (4 BC one-hot from manifest.bc_config) +
-[n_chills, sleeve_MEF, exo_q0_MW] (3 BC scalar fields) = 16 columns. `Tpour` IS a
-per-material property in this project's DB (`tools.doe.load_properties.load_materials`
-midpoints Tpour_lo_C/Tpour_hi_C for the 12 named anchor grades; the LHS-sampled
-synthetic alloys carry their own `Tpour_C` in `materials_synthetic.csv`) -- so it is
-resolved per material_id, not hardcoded as a constant.
-
-Target: per-graph `log1p(y_raw.max())` -- the total solidification time
-(scale), read straight from the v4 graph cache (`y_raw`, physical seconds).
-"""
-from __future__ import annotations
-
-import csv
-import os
-import os.path as osp
-
-import joblib
-import numpy as np
-from sklearn.ensemble import HistGradientBoostingRegressor
-from sklearn.model_selection import train_test_split
-
-from src.build_dataset import (FEATURE_NAMES_V4, _resolve_exo_q0_v4,
-                               _resolve_material_props_v4)
-from src.features import FEATURE_NAMES as _NODE_FEATURE_NAMES
-
-# volume_mm3/surface_area_mm2 are constant-per-node columns of the v2/v3/v4
-# node-feature matrix X (see src/features.py FEATURE_NAMES and the v4
-# comment there: "v3's 13 columns UNCHANGED (X[:, :9] == v2, X[:, :13] ==
-# v3)" -- so these indices are valid for the v4 cache too). Verified live
-# against src/build_dataset.py:622-623, which reads the SAME two indices
-# out of the SAME FEATURE_NAMES list to get volume/area for the Chvorinov
-# t_ref fit:
-#     volume = float(X[0, FEATURE_NAMES.index("volume_mm3")])
-#     area = float(X[0, FEATURE_NAMES.index("surface_area_mm2")])
-_VOL_IDX = _NODE_FEATURE_NAMES.index("volume_mm3")
-_AREA_IDX = _NODE_FEATURE_NAMES.index("surface_area_mm2")
-
-# bc_area_frac (soft-anchor Task 2, std head v2 17th feature) column indices --
-# resolved off `src.build_dataset.FEATURE_NAMES_V4` (never hardcoded silently)
-# with a cross-check assertion against the brief's stated indices
-# (is_cooling=3, is_adiabatic=4): if FEATURE_NAMES_V4 ever changes upstream,
-# this fails loudly instead of silently reading the wrong columns.
-_IS_COOLING_IDX = FEATURE_NAMES_V4.index("is_cooling")
-_IS_ADIABATIC_IDX = FEATURE_NAMES_V4.index("is_adiabatic")
-assert (_IS_COOLING_IDX, _IS_ADIABATIC_IDX) == (3, 4), (
-    f"bc_area_frac column indices drifted from the Global Constraints formula's "
-    f"is_cooling=3/is_adiabatic=4 (FEATURE_NAMES_V4 now gives "
-    f"is_cooling={_IS_COOLING_IDX}, is_adiabatic={_IS_ADIABATIC_IDX}) -- "
-    f"FEATURE_NAMES_V4 changed upstream; update this module before trusting "
-    f"bc_area_frac_from_graph")
-
-# GBM row feature layout -- exposed for the Task 8 evaluator so it can
-# rebuild rows the exact same way at inference time.
-#
-# BC-aware extension (exo-fix Task 3): the original 9 material/geometry
-# columns are blind to boundary conditions (chills/sleeves change total
-# solidification time independently of material/geometry), so the scale
-# head could not tell a chilled casting from a plain sand one. The 7 new
-# columns make it BC-aware: a one-hot over the manifest's `bc_config`
-# values (`none`/`chills`/`insulating_sleeve`/`exo_sleeve` -- the only
-# values `tools/doe/doe.py` ever writes), plus the raw `n_chills` and
-# `sleeve_MEF` manifest fields, plus the calibrated exothermic-sleeve heat
-# release `exo_q0_MW` (0.0 for non-exo rows, by `_resolve_exo_q0_v4`'s own
-# contract).
-FEATURE_NAMES = [
-    "k", "rho", "cp", "L_kJ_kg", "TL", "TS", "Tpour",
-    "V_over_A", "mold_effusivity",
-    "bc_none", "bc_chills", "bc_insulating_sleeve", "bc_exo_sleeve",
-    "n_chills", "sleeve_MEF", "exo_q0_MW",
-]
-
-# The only bc_config values `tools/doe/doe.py` ever writes to the manifest
-# (verified live against doe_out/manifest.csv: Counter gives exactly these
-# 4 keys). An empty string means "no BC row was ever assigned" and is
-# treated as "none" (mirrors `_resolve_exo_q0_v4`'s own
-# `row.get("bc_config") or "none"` convention). Anything else is an
-# unrecognized/typo'd value -- raise loudly rather than silently ignoring it
-# in the one-hot.
-_BC_VALUES = ("none", "chills", "insulating_sleeve", "exo_sleeve")
-
-_HGB_DEFAULTS = dict(learning_rate=0.1, max_leaf_nodes=63,
-                      l2_regularization=1.0, random_state=0)
-
-# HGB's own early_stopping carves off validation_fraction (default 0.1) of
-# whatever it's given -- on tiny row counts (e.g. a ~20-stem unit test) that
-# leaves 1-2 validation rows, which stops training almost immediately and
-# badly underfits. Only turn early stopping on once there's enough data for
-# it to be a meaningful signal; production refits (thousands of stems) still
-# get it.
-_MIN_ROWS_FOR_EARLY_STOPPING = 50
-
-
-def _load_manifest_rows(manifest_path: str) -> dict:
-    rows = {}
-    with open(manifest_path, newline="") as f:
-        for row in csv.DictReader(f):
-            rows[row["run_id"]] = row
-    return rows
-
-
-def _num_or_zero(raw, field: str, stem: str) -> float:
-    """float(raw), except '' / None -> 0.0 (the brief's explicit rule for
-    n_chills/sleeve_MEF, mirroring how the manifest encodes N/A for rows
-    with no chill/sleeve BC). A present-but-non-numeric value still raises
-    loudly (e.g. a typo) -- only blank is treated as "not applicable"."""
-    if raw is None or str(raw).strip() == "":
-        return 0.0
-    try:
-        return float(raw)
-    except (TypeError, ValueError) as e:
-        raise ValueError(f"run_id={stem!r}: {field}={raw!r} is not a number") from e
-
-
-def _row_features_and_graph(stem: str, manifest_rows: dict, manifest_dir: str,
-                             features_dir: str) -> tuple[np.ndarray, "np.lib.npyio.NpzFile"]:
-    """Shared row-assembly: one BC-aware feature row (see `FEATURE_NAMES`) +
-    the loaded graph-cache npz. Factored out of `_row_features` so the
-    std-target head (`_row_features_std`) reuses the IDENTICAL X (same
-    material/BC resolution, same I/O) and swaps only the target it derives
-    from the already-open graph -- never re-deriving or duplicating the
-    resolution logic.
-
-    Loud-fails (raises) on anything unresolvable -- a missing manifest row,
-    a missing/unresolvable material_id, a missing mold_effusivity, an
-    unrecognized bc_config value, an uncalibrated exo_sleeve row, or a
-    missing graph-cache file -- NEVER silently drops or fabricates a value.
-    The one EXCEPTION (per the brief) is n_chills/sleeve_MEF: '' -> 0.0,
-    mirroring how the manifest encodes "no chill/sleeve BC" for this row.
-    """
-    row = manifest_rows.get(stem)
-    if row is None:
-        raise KeyError(f"stem {stem!r} not found in manifest ({manifest_dir!r})")
-
-    material_id = (row.get("material_id") or "").strip()
-    if not material_id:
-        raise KeyError(f"run_id={stem!r}: manifest row has no material_id")
-    # Raises KeyError if unresolvable via any of the 3 fallback sources
-    # (owner-verified CSV / materials_synthetic.csv / named anchor grades) --
-    # see src.build_dataset._resolve_material_props_v4's own docstring.
-    props = _resolve_material_props_v4(material_id, manifest_dir)
-
-    mold_eff_raw = row.get("mold_effusivity")
-    if mold_eff_raw is None or str(mold_eff_raw).strip() == "":
-        raise KeyError(f"run_id={stem!r}: manifest row has no mold_effusivity")
-    mold_effusivity = float(mold_eff_raw)
-
-    bc_config = (row.get("bc_config") or "").strip() or "none"
-    if bc_config not in _BC_VALUES:
-        raise ValueError(
-            f"run_id={stem!r}: unrecognized bc_config={bc_config!r} "
-            f"(expected one of {_BC_VALUES})")
-    one_hot = [1.0 if bc_config == v else 0.0 for v in _BC_VALUES]
-
-    n_chills = _num_or_zero(row.get("n_chills"), "n_chills", stem)
-    sleeve_MEF = _num_or_zero(row.get("sleeve_MEF"), "sleeve_MEF", stem)
-
-    # 0.0 for non-exo rows by its own contract; raises loudly for
-    # bc_config=="exo_sleeve" rows with no calibrated MEF match -- see
-    # src.build_dataset._resolve_exo_q0_v4's docstring.
-    exo_q0_MW = _resolve_exo_q0_v4(row) / 1e6
-
-    gpath = osp.join(features_dir, "graphs", f"{stem}.npz")
-    if not osp.exists(gpath):
-        raise FileNotFoundError(f"graph cache missing for stem={stem!r}: {gpath}")
-    d = np.load(gpath)
-    volume = float(d["X"][0, _VOL_IDX])
-    area = float(d["X"][0, _AREA_IDX])
-    if area <= 0:
-        raise ValueError(f"run_id={stem!r}: non-positive surface_area_mm2={area}")
-    v_over_a = volume / area
-
-    x = np.array([
-        props["k"], props["rho"], props["cp"], props["L_kJ_kg"],
-        props["TL"], props["TS"], props["Tpour"],
-        v_over_a, mold_effusivity,
-        *one_hot,
-        n_chills, sleeve_MEF, exo_q0_MW,
-    ], dtype=np.float64)
-    return x, d
-
-
-def _row_features(stem: str, manifest_rows: dict, manifest_dir: str,
-                   features_dir: str) -> tuple[np.ndarray, float]:
-    """One BC-aware feature row (see `FEATURE_NAMES`) + log1p(y_global).
-
-    See `_row_features_and_graph` for the loud-fail contract this shares.
-    """
-    x, d = _row_features_and_graph(stem, manifest_rows, manifest_dir, features_dir)
-    y_raw = d["y_raw"] if "y_raw" in d.files else d["y"]
-    y_log = float(np.log1p(float(np.max(y_raw))))
-    return x, y_log
-
-
-def _row_features_std(stem: str, manifest_rows: dict, manifest_dir: str,
-                       features_dir: str) -> tuple[np.ndarray, float]:
-    """The SAME BC-aware feature row as `_row_features` (see
-    `_row_features_and_graph`) + log1p(std(y_nodal)) -- the per-graph
-    field-spread target for the GBM std head (variance-recomposition Task 2).
-
-    `y_nodal` is the graph npz's `y` field (per-node solidification time,
-    seconds) -- raises loudly (KeyError) if a requested stem's graph cache
-    has no `y`, never silently skips it.
-    """
-    x, d = _row_features_and_graph(stem, manifest_rows, manifest_dir, features_dir)
-    if "y" not in d.files:
-        raise KeyError(f"run_id={stem!r}: graph cache in {features_dir!r} has no 'y' field")
-    y_std_log = float(np.log1p(float(np.std(d["y"]))))
-    return x, y_std_log
-
-
-# std head v2 (soft-anchor Task 2): the v1 row (`FEATURE_NAMES`, 16 columns)
-# + one graph-derived 17th column, `bc_area_frac`. Kept as its own module
-# constant (rather than mutating `FEATURE_NAMES`) so the 16-feature total
-# head stays UNTOUCHED -- see `std_feature_row`'s width-dispatch contract.
-FEATURE_NAMES_STD_V2 = FEATURE_NAMES + ["bc_area_frac"]
-
-
-def bc_area_frac_from_graph(graph) -> float:
-    """The std-head-v2 17th feature (Global Constraints formula, soft-anchor
-    spec Sec 2): the fraction of BC-tagged surface nodes that are adiabatic
-    (vs. cooling) -- `sum(is_adiabatic) / max(sum(is_adiabatic) +
-    sum(is_cooling), 1)`, read straight off the graph's node feature matrix
-    `graph["X"]` at `_IS_ADIABATIC_IDX`/`_IS_COOLING_IDX` (columns 4/3 of
-    `FEATURE_NAMES_V4`, cross-checked at import time above -- never
-    hardcoded silently a second time here).
-
-    `graph` is whatever `_row_features_and_graph` returns as its second
-    element (an open graph-cache npz) -- or, for tests, any object
-    supporting `graph["X"]` indexing (e.g. a plain dict).
-
-    A node-count proxy for the true tri-area fraction (no per-node area
-    weight is available in `X`), acceptable per the spec. All-sand rows
-    (no adiabatic surface at all, e.g. the fully-cooled regime) -> exactly
-    0.0, since the numerator is 0 regardless of the denominator; a graph
-    with neither cooling nor adiabatic nodes tagged also -> 0.0 (the `max(
-    ..., 1)` floor avoids a 0/0 division).
-    """
-    X = graph["X"]
-    n_adiabatic = float(np.sum(X[:, _IS_ADIABATIC_IDX]))
-    n_cooling = float(np.sum(X[:, _IS_COOLING_IDX]))
-    return n_adiabatic / max(n_adiabatic + n_cooling, 1.0)
-
-
-def std_feature_row(stem: str, manifest_rows: dict, manifest_dir: str,
-                    features_dir: str, n_features: int) -> np.ndarray:
-    """Width-dispatching feature-row builder for the std head (soft-anchor
-    Task 2) -- lets a caller (Task 3's `predict_casting`/`evaluate_ft2b`)
-    build the right-width row for whichever std-head model it loaded,
-    keyed off that model's own `n_features_in_`, without duplicating
-    `_row_features_and_graph`'s resolution logic.
-
-    - `n_features == 16` -> the v1 row (`_row_features_and_graph`'s X,
-      unchanged -- byte-identical to `_row_features`'s own row).
-    - `n_features == 17` -> the v1 row + `bc_area_frac_from_graph` of the
-      SAME already-open graph appended (never a second graph load).
-    - anything else -> a loud `ValueError` naming the two supported widths
-      (never a silent truncation/pad).
-    """
-    if n_features not in (16, 17):
-        raise ValueError(
-            f"std_feature_row: unsupported n_features={n_features!r} for "
-            f"stem={stem!r} -- the std head only supports 16 (v1) or 17 "
-            f"(v2, v1 + bc_area_frac) feature columns")
-
-    x, d = _row_features_and_graph(stem, manifest_rows, manifest_dir, features_dir)
-    if n_features == 16:
-        return x
-    area_frac = bc_area_frac_from_graph(d)
-    return np.concatenate([x, [area_frac]])
-
-
-def build_v2_feature_row(props: dict, v_over_a: float,
-                         bc_spec: dict | None = None) -> np.ndarray:
-    """One BC-aware v2-total-GBM feature row (`FEATURE_NAMES` order) -- the
-    bare-mesh-inference counterpart to `_row_features_and_graph`'s
-    manifest/graph-cache-driven row (that one resolves `props`/`v_over_a`
-    from a manifest row + graph cache; this one takes them straight from
-    the caller, for real-part inference where there is no manifest row).
-
-    Public (no leading underscore) and factored out of
-    `gbm_v2_total_seconds` -- rather than duplicated wherever a caller needs
-    the total head's row (`src.predict_casting.predict_field`'s optional std
-    head, variance-recomposition Task 3) -- so the total head
-    (`gbm_v2_total_seconds`) and any std head scored from the SAME row
-    (`predict_std_seconds`) can never silently desync on a `FEATURE_NAMES`
-    reorder. Mirrors `_row_features_and_graph`'s own "so the two heads can
-    never drift apart" rationale (Task 2).
-
-    Builds the row from `props` (Task 1's `src.inference_props.
-    resolve_props()` 8-key dict: k/rho/cp/L_kJ_kg/TL/TS/Tpour/
-    mold_effusivity), `v_over_a` (volume/surface-area, mm), and `bc_spec`:
-        - `None` (the default, "plain sand mold" case) -> `bc_config`
-          "none" one-hot, `n_chills`/`sleeve_MEF`/`exo_q0_MW` all 0.0.
-        - a dict with keys `bc_config`/`n_chills`/`sleeve_MEF`/`exo_q0_MW`
-          (Part B's BC-aware inference entry point) -> that `bc_config`'s
-          one-hot + its 3 raw scalar fields, verbatim.
-
-    An unrecognized `bc_config` raises `ValueError` (mirrors `_row_features`'
-    own convention for a manifest row with a typo'd `bc_config`).
-    """
-    if bc_spec is None:
-        bc_config = "none"
-        n_chills = 0.0
-        sleeve_MEF = 0.0
-        exo_q0_MW = 0.0
-    else:
-        bc_config = bc_spec["bc_config"]
-        n_chills = float(bc_spec["n_chills"])
-        sleeve_MEF = float(bc_spec["sleeve_MEF"])
-        exo_q0_MW = float(bc_spec["exo_q0_MW"])
-
-    if bc_config not in _BC_VALUES:
-        raise ValueError(
-            f"unrecognized bc_config={bc_config!r} (expected one of {_BC_VALUES})")
-    one_hot = [1.0 if bc_config == v else 0.0 for v in _BC_VALUES]
-
-    return np.array([
-        props["k"], props["rho"], props["cp"], props["L_kJ_kg"],
-        props["TL"], props["TS"], props["Tpour"],
-        float(v_over_a), props["mold_effusivity"],
-        *one_hot,
-        n_chills, sleeve_MEF, exo_q0_MW,
-    ], dtype=np.float64)
-
-
-def gbm_v2_total_seconds(hgb, props: dict, v_over_a: float,
-                         bc_spec: dict | None = None) -> float:
-    """Bare-mesh inference entry point: one feature row -> total seconds.
-
-    Builds the row via `build_v2_feature_row` (see that function for the
-    full `props`/`v_over_a`/`bc_spec` contract) and returns
-    `float(np.expm1(hgb.predict(row)[0]))` -- physical seconds, inverting
-    the same `log1p(y_raw.max())` target `_row_features`/`refit_gbm` train
-    on.
-    """
-    x = build_v2_feature_row(props, v_over_a, bc_spec)
-    log_pred = float(hgb.predict(x.reshape(1, -1))[0])
-    return float(np.expm1(log_pred))
-
-
-def _assemble_matrix(stems, manifest_path: str, features_dir: str):
-    """[len(stems), len(FEATURE_NAMES)] feature matrix + log1p(y_global) target vector.
-
-    Raises (does not skip) on any row that cannot be resolved -- see
-    `_row_features`.
-    """
-    manifest_dir = osp.dirname(osp.abspath(manifest_path)) or "."
-    manifest_rows = _load_manifest_rows(manifest_path)
-
-    X_rows, y_rows, used = [], [], []
-    for stem in stems:
-        x, y_log = _row_features(stem, manifest_rows, manifest_dir, features_dir)
-        X_rows.append(x)
-        y_rows.append(y_log)
-        used.append(stem)
-
-    X = np.stack(X_rows) if X_rows else np.empty((0, len(FEATURE_NAMES)))
-    y = np.array(y_rows, dtype=np.float64)
-    return X, y, used
-
-
-def _assemble_matrix_std(stems, manifest_path: str, features_dir: str):
-    """[len(stems), len(FEATURE_NAMES)] feature matrix + log1p(std(y_nodal))
-    target vector -- the std-head counterpart to `_assemble_matrix`. Same X
-    (via `_row_features_std` -> `_row_features_and_graph`), different target.
-
-    Raises (does not skip) on any row that cannot be resolved -- see
-    `_row_features_std`.
-    """
-    manifest_dir = osp.dirname(osp.abspath(manifest_path)) or "."
-    manifest_rows = _load_manifest_rows(manifest_path)
-
-    X_rows, y_rows, used = [], [], []
-    for stem in stems:
-        x, y_log = _row_features_std(stem, manifest_rows, manifest_dir, features_dir)
-        X_rows.append(x)
-        y_rows.append(y_log)
-        used.append(stem)
-
-    X = np.stack(X_rows) if X_rows else np.empty((0, len(FEATURE_NAMES)))
-    y = np.array(y_rows, dtype=np.float64)
-    return X, y, used
-
-
-def _fit_hgb_with_metrics(X: np.ndarray, y: np.ndarray, out_path: str,
-                           max_iter: int, holdout_frac: float,
-                           min_samples_leaf: int, random_state: int,
-                           caller: str, loss_kwargs: dict | None = None) -> dict:
-    """Shared HGB-fit/holdout/save logic behind `refit_gbm`, `refit_gbm_std`,
-    and `refit_gbm_std_v2` -- identical modeling code, only the assembled
-    (X, y) (and, for std v2, the loss) differ between heads. Factored out so
-    the heads can never drift apart in their fitting/metrics logic;
-    `refit_gbm`'s and `refit_gbm_std`'s own public behavior is unchanged by
-    this factoring (see the guard tests in `tests/test_gbm_ft2b.py`).
-
-    `min_samples_leaf` (sklearn's HGB default is 20) must be turned down for
-    very small row counts (e.g. a ~20-row unit test) -- HGB cannot split a
-    node at all once a candidate split would leave either side with fewer
-    than `min_samples_leaf` samples, so with the sklearn default and only
-    ~20 training rows every tree degenerates to a single leaf (the
-    unconditional mean). Production refits (thousands of stems) should keep
-    the default.
-
-    `loss_kwargs` (soft-anchor Task 2): optional extra
-    `HistGradientBoostingRegressor` constructor kwargs (e.g.
-    `{"loss": "quantile", "quantile": 0.5}` for the std-head-v2 median
-    objective), merged into BOTH the holdout probe and the final model so a
-    caller never has to duplicate this fit/holdout/save logic just to swap
-    the loss. `None` (the default) reproduces the exact pre-existing
-    behavior (sklearn's default squared-error loss) -- `refit_gbm`'s and
-    `refit_gbm_std`'s own output is therefore byte-identical to before this
-    parameter existed.
-
-    Returns a dict of finite metrics: `n_train`, `train_mae_log1p`,
-    `train_mae_seconds`, and (when there are enough rows to spare a
-    holdout split) `holdout_mae_log1p`.
-    """
-    if X.shape[0] == 0:
-        raise ValueError(f"{caller}: no training stems resolved -- nothing to fit")
-
-    metrics: dict = {"n_train": X.shape[0]}
-
-    # Bonus: a quick internal holdout split for a generalization signal.
-    # Skipped for tiny row counts (both HGB's own internal validation split
-    # and this outer split would starve each other of data).
-    if X.shape[0] >= 10 and holdout_frac > 0:
-        X_tr, X_ho, y_tr, y_ho = train_test_split(
-            X, y, test_size=holdout_frac, random_state=random_state)
-        probe_params = dict(_HGB_DEFAULTS)
-        probe_params["random_state"] = random_state
-        if loss_kwargs:
-            probe_params.update(loss_kwargs)
-        probe = HistGradientBoostingRegressor(
-            max_iter=max_iter,
-            early_stopping=(len(X_tr) >= _MIN_ROWS_FOR_EARLY_STOPPING),
-            validation_fraction=0.1, min_samples_leaf=min_samples_leaf,
-            **probe_params)
-        probe.fit(X_tr, y_tr)
-        ho_mae = float(np.mean(np.abs(probe.predict(X_ho) - y_ho)))
-        if np.isfinite(ho_mae):
-            metrics["holdout_mae_log1p"] = ho_mae
-
-    # Final production model: fit on ALL resolved training rows.
-    model_params = dict(_HGB_DEFAULTS)
-    model_params["random_state"] = random_state
-    if loss_kwargs:
-        model_params.update(loss_kwargs)
-    model = HistGradientBoostingRegressor(
-        max_iter=max_iter,
-        early_stopping=(X.shape[0] >= _MIN_ROWS_FOR_EARLY_STOPPING),
-        validation_fraction=0.1, min_samples_leaf=min_samples_leaf,
-        **model_params)
-    model.fit(X, y)
-
-    train_pred = model.predict(X)
-    metrics["train_mae_log1p"] = float(np.mean(np.abs(train_pred - y)))
-    metrics["train_mae_seconds"] = float(
-        np.mean(np.abs(np.expm1(train_pred) - np.expm1(y))))
-
-    out_dir = osp.dirname(out_path)
-    if out_dir:
-        os.makedirs(out_dir, exist_ok=True)
-    joblib.dump(model, out_path)
-    metrics["out_path"] = out_path
-
-    return metrics
-
-
-def refit_gbm(features_dir: str, train_stems, manifest_path: str, out_path: str,
-              max_iter: int = 300, holdout_frac: float = 0.2,
-              min_samples_leaf: int = 20, random_state: int = 0) -> dict:
-    """Refit the multi-alloy GBM scale-head (total solidification time target,
-    `log1p(y_global)`) and save it to `out_path`. See `_fit_hgb_with_metrics`
-    for the shared fitting/metrics logic.
-
-    Returns a dict of finite metrics: `n_train`, `train_mae_log1p`,
-    `train_mae_seconds`, and (when there are enough rows to spare a
-    holdout split) `holdout_mae_log1p`.
-    """
-    X, y, used = _assemble_matrix(train_stems, manifest_path, features_dir)
-    return _fit_hgb_with_metrics(
-        X, y, out_path, max_iter=max_iter, holdout_frac=holdout_frac,
-        min_samples_leaf=min_samples_leaf, random_state=random_state,
-        caller="refit_gbm")
-
-
-def refit_gbm_std(features_dir: str, train_stems, manifest_path: str, out_path: str,
-                   max_iter: int = 300, holdout_frac: float = 0.2,
-                   min_samples_leaf: int = 20, random_state: int = 0) -> dict:
-    """Refit the GBM std-head (per-row field-spread target,
-    `log1p(std(y_nodal))`) and save it to `out_path` -- the counterpart to
-    `refit_gbm`'s total-time head, for variance-preserving hybrid
-    recomposition (see `src.hybrid.hybrid_field_affine`, Task 1).
-
-    Consumes the SAME feature row layout (`FEATURE_NAMES`) as `refit_gbm`,
-    built via `_assemble_matrix_std` (which reuses `_row_features_and_graph`
-    -- identical material/BC resolution, only the target differs). See
-    `_fit_hgb_with_metrics` for the shared fitting/metrics logic.
-
-    Returns a dict of finite metrics: `n_train`, `train_mae_log1p`,
-    `train_mae_seconds`, and (when there are enough rows to spare a
-    holdout split) `holdout_mae_log1p`.
-    """
-    X, y, used = _assemble_matrix_std(train_stems, manifest_path, features_dir)
-    return _fit_hgb_with_metrics(
-        X, y, out_path, max_iter=max_iter, holdout_frac=holdout_frac,
-        min_samples_leaf=min_samples_leaf, random_state=random_state,
-        caller="refit_gbm_std")
-
-
-def _assemble_matrix_std_v2(stems, manifest_path: str, features_dir: str):
-    """[len(stems), 17] feature matrix (v1 row + `bc_area_frac`) +
-    `log1p(std(y_nodal))` target vector -- the std-head-v2 counterpart to
-    `_assemble_matrix_std`. Reuses `_row_features_and_graph` for the shared
-    v1 row AND derives `bc_area_frac_from_graph` off that SAME already-open
-    graph (never a second graph load per row, mirroring
-    `_row_features_and_graph`'s own "so the two heads can never drift
-    apart" rationale).
-
-    Raises (does not skip) on any row that cannot be resolved -- see
-    `_row_features_and_graph`; also raises `KeyError` if a stem's graph
-    cache has no `y` field (mirrors `_row_features_std`).
-    """
-    manifest_dir = osp.dirname(osp.abspath(manifest_path)) or "."
-    manifest_rows = _load_manifest_rows(manifest_path)
-
-    X_rows, y_rows, used = [], [], []
-    for stem in stems:
-        x, d = _row_features_and_graph(stem, manifest_rows, manifest_dir, features_dir)
-        if "y" not in d.files:
-            raise KeyError(f"run_id={stem!r}: graph cache in {features_dir!r} has no 'y' field")
-        area_frac = bc_area_frac_from_graph(d)
-        x17 = np.concatenate([x, [area_frac]])
-        y_log = float(np.log1p(float(np.std(d["y"]))))
-        X_rows.append(x17)
-        y_rows.append(y_log)
-        used.append(stem)
-
-    X = np.stack(X_rows) if X_rows else np.empty((0, len(FEATURE_NAMES_STD_V2)))
-    y = np.array(y_rows, dtype=np.float64)
-    return X, y, used
-
-
-def refit_gbm_std_v2(features_dir: str, train_stems, manifest_path: str, out_path: str,
-                     max_iter: int = 300, holdout_frac: float = 0.2,
-                     min_samples_leaf: int = 20, random_state: int = 0) -> dict:
-    """Refit the std-head v2 (soft-anchor Task 2): the SAME per-row
-    field-spread target as `refit_gbm_std` (`log1p(std(y_nodal))`), but on
-    17 features (`FEATURE_NAMES_STD_V2` = the v1 16 + graph-derived
-    `bc_area_frac`) and with a median (quantile=0.5) loss instead of squared
-    error -- tail-robust against the per-stem noise that squared-error
-    over-weights (see the soft-anchor design spec Sec 2).
-
-    Built via `_assemble_matrix_std_v2` (identical material/BC/graph
-    resolution to `_assemble_matrix_std`, plus the extra column). See
-    `_fit_hgb_with_metrics` for the shared fitting/metrics logic; this is
-    the one caller that passes `loss_kwargs`.
-
-    Returns a dict of finite metrics: `n_train`, `train_mae_log1p`,
-    `train_mae_seconds`, and (when there are enough rows to spare a
-    holdout split) `holdout_mae_log1p`.
-    """
-    X, y, used = _assemble_matrix_std_v2(train_stems, manifest_path, features_dir)
-    return _fit_hgb_with_metrics(
-        X, y, out_path, max_iter=max_iter, holdout_frac=holdout_frac,
-        min_samples_leaf=min_samples_leaf, random_state=random_state,
-        caller="refit_gbm_std_v2", loss_kwargs={"loss": "quantile", "quantile": 0.5})
-
-
-def predict_std_seconds(model, X_row) -> float:
-    """One feature row (`FEATURE_NAMES` order, shape `(len(FEATURE_NAMES),)`
-    or `(1, len(FEATURE_NAMES))`) -> predicted std(y_nodal) in seconds.
-
-    Inverts the `log1p(std(y_nodal))` target `refit_gbm_std` trains on via
-    `expm1`, floored at 0.0 (a predicted std can never be physically
-    negative; `expm1` of a very negative log-prediction could underflow to
-    a tiny negative float).
-    """
-    x = np.asarray(X_row, dtype=np.float64).reshape(1, -1)
-    log_pred = float(model.predict(x)[0])
-    return max(0.0, float(np.expm1(log_pred)))

src/gbm_heads.py

@@ -0,0 +1,69 @@
+"""The two gradient-boosted heads: total solidification time and field spread.
+
+Both are scikit-learn HistGradientBoostingRegressor models that see one row
+per part and predict log1p of their target in seconds.
+
+Row layout (16 columns for the total head, 17 for the spread head):
+  0-6   alloy: k, rho, cp, L_kJ_kg, TL, TS, Tpour
+  7     part volume / surface area (mm)
+  8     mold effusivity
+  9-15  boundary-condition descriptors from training. A casting that cools
+        into the sand mold over its whole surface is the regime encoded as
+        [1, 0, 0, 0, 0, 0, 0]; this release always uses it.
+  16    (spread head only) fraction of surface nodes tagged adiabatic
+"""
+import numpy as np
+
+from src.features import FEATURE_NAMES
+
+_IS_COOLING_IDX = FEATURE_NAMES.index("is_cooling")
+_IS_ADIABATIC_IDX = FEATURE_NAMES.index("is_adiabatic")
+_VOL_IDX = FEATURE_NAMES.index("volume_mm3")
+_AREA_IDX = FEATURE_NAMES.index("surface_area_mm2")
+
+_SAND_MOLD_BOUNDARY = (1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
+N_TOTAL_FEATURES = 16
+N_SPREAD_FEATURES = 17
+
+
+def total_row(props: dict, X: np.ndarray) -> np.ndarray:
+    """The 16-column row for the total-time head."""
+    v_over_a = float(X[0, _VOL_IDX]) / float(X[0, _AREA_IDX])
+    return np.array([
+        props["k"], props["rho"], props["cp"], props["L_kJ_kg"],
+        props["TL"], props["TS"], props["Tpour"],
+        float(v_over_a), props["mold_effusivity"],
+        *_SAND_MOLD_BOUNDARY,
+    ], dtype=np.float64)
+
+
+def adiabatic_node_fraction(X: np.ndarray) -> float:
+    """Adiabatic surface nodes / (adiabatic + cooling surface nodes)."""
+    n_adiabatic = float(np.sum(X[:, _IS_ADIABATIC_IDX]))
+    n_cooling = float(np.sum(X[:, _IS_COOLING_IDX]))
+    return n_adiabatic / max(n_adiabatic + n_cooling, 1.0)
+
+
+def spread_row(row_total: np.ndarray, X: np.ndarray) -> np.ndarray:
+    """The 17-column row for the spread head."""
+    return np.concatenate([row_total, [adiabatic_node_fraction(X)]])
+
+
+def predict_total_seconds(model, row: np.ndarray) -> float:
+    total = float(np.expm1(float(model.predict(row.reshape(1, -1))[0])))
+    if total <= 0:
+        raise ValueError(f"total-time head returned a non-positive total ({total!r})")
+    return total
+
+
+def predict_spread_seconds(model, row: np.ndarray) -> float:
+    x = np.asarray(row, dtype=np.float64).reshape(1, -1)
+    log_pred = float(model.predict(x)[0])
+    return max(0.0, float(np.expm1(log_pred)))
+
+
+def check_head(model, n_features: int, name: str) -> None:
+    n = getattr(model, "n_features_in_", None)
+    if n is not None and int(n) != n_features:
+        raise ValueError(f"{name} head expects {n} input columns, this code builds "
+                         f"{n_features}")

src/geometry.py

@@ -80,53 +80,6 @@ def geodesic_distance(n_nodes: int, edges: np.ndarray, edge_len: np.ndarray,
     return d
 
 
-def geodesic_diameter_two_sweep(n_nodes: int, edges: np.ndarray,
-                                edge_len: np.ndarray, start: int = 0) -> float:
-    """Two-sweep (double-Dijkstra) approximation of the mesh's geodesic
-    diameter, over the SAME mesh-edge graph `geodesic_distance` already uses
-    (not a per-BC-patch quantity -- patch-independent, through-metal, a
-    single per-part constant): sweep 1 finds the farthest node from `start`;
-    sweep 2 finds the farthest node from THAT node; the sweep-2 distance is
-    returned as D_geo. This is the standard double-sweep heuristic for graph
-    diameter (exact on trees, a lower bound in general, but a good practical
-    estimate for a roughly part-shaped mesh graph).
-
-    `start` is a fixed node INDEX (default 0), not searched or randomized --
-    deterministic and reproducible for the same mesh (same node ordering).
-
-    Returns 0.0 (never NaN/inf) for an edgeless mesh, or if either sweep's
-    farthest reachable node has zero finite distance (degenerate/single-node
-    or fully-disconnected mesh) -- callers must treat 0.0 as "no valid
-    part-normalized length available" rather than a real zero-size part.
-
-    CAVEAT (bc-reach re-review N2): the two-sweep heuristic is a LOWER
-    BOUND on the true graph diameter in general (exact only on trees), and
-    is start-sensitive -- on a BUKSA-like real casting mesh, re-running
-    with different random start nodes measured ~15% spread in the reported
-    D_geo across starts. Fixing `start=0` makes the result perfectly
-    DETERMINISTIC for a given mesh (bit-identical across repeated calls,
-    verified in tests/test_geometry.py), which is what the cache-build and
-    inference paths both rely on (the same mesh must always yield the same
-    D_geo) -- but it does NOT make the estimate stable under a RE-MESH of
-    the same part (a different tetrahedralization can shift which node
-    lands at index 0, and therefore the estimate, by roughly that same
-    ~15%). Closing that gap would need iterating sweeps to convergence
-    (repeat from the new farthest node until the distance stops growing)
-    or a multi-start-and-take-max scheme -- deliberately NOT done here:
-    the part-normalized reach channels only need a stable-per-mesh
-    O(part-size) scale, not a tight diameter bound, and re-mesh stability
-    was judged out of scope for this task."""
-    if edges.shape[0] == 0 or n_nodes == 0:
-        return 0.0
-    start = int(min(max(start, 0), n_nodes - 1))
-    d1 = geodesic_distance(n_nodes, edges, edge_len, np.array([start], dtype=np.int64))
-    finite1 = np.isfinite(d1)
-    if not finite1.any():
-        return 0.0
-    far1 = int(np.flatnonzero(finite1)[np.argmax(d1[finite1])])
-    d2 = geodesic_distance(n_nodes, edges, edge_len, np.array([far1], dtype=np.int64))
-    finite2 = d2[np.isfinite(d2)]
-    return float(finite2.max()) if finite2.size else 0.0
 
 
 def _local_modulus_for_query(tets: np.ndarray, cool_tri: np.ndarray, radii: np.ndarray,
@@ -149,79 +102,17 @@ def _local_modulus_for_query(tets: np.ndarray, cool_tri: np.ndarray, radii: np.n
     return vals
 
 
-_IDW_K = 4             # neighbors for inverse-distance-weighted subsample fill
-_IDW_EPS = 1e-9        # distance floor to avoid divide-by-zero
-
-
-def _idw_geodesic_fill(vals: np.ndarray, dist_from_computed: np.ndarray,
-                       missing: np.ndarray, bbox_diag: float,
-                       k: int = _IDW_K) -> np.ndarray:
-    """Inverse-distance-weighted fill for `missing` nodes from the values
-    (`vals`, one row per computed node) at the computed/query nodes, using
-    geodesic (mesh-edge) distances (`dist_from_computed`, shape
-    (n_computed, N)) to the k nearest computed nodes. Smoother than a
-    single-nearest-node copy for a continuously-varying field. Nodes
-    unreachable from every computed node (inf distance -- disconnected mesh
-    component) fall back to the bbox-diagonal sentinel."""
-    n_computed = vals.shape[0]
-    k_eff = max(1, min(k, n_computed))
-    sub = dist_from_computed[:, missing]                        # (n_computed, n_missing)
-    part = np.argpartition(sub, k_eff - 1, axis=0)[:k_eff]       # (k_eff, n_missing)
-    d = np.take_along_axis(sub, part, axis=0)                   # (k_eff, n_missing)
-    finite = np.isfinite(d)
-    w = np.where(finite, 1.0 / np.maximum(d, _IDW_EPS), 0.0)     # (k_eff, n_missing)
-    wsum = w.sum(axis=0)                                        # (n_missing,)
-    reachable = wsum > 0
-    w_norm = w / np.where(reachable, wsum, 1.0)                 # (k_eff, n_missing)
-    gathered = vals[part]                                       # (k_eff, n_missing, R)
-    result = np.einsum('kn,knr->nr', w_norm, gathered)          # (n_missing, R)
-    result[~reachable] = bbox_diag
-    return result
 
 
 def local_modulus(node_xyz: np.ndarray, tets: np.ndarray, cool_tri: np.ndarray,
                   edges: np.ndarray, edge_len: np.ndarray, radii,
-                  chunk: int = 128, exact: bool = True, subsample=None) -> np.ndarray:
+                  chunk: int = 128) -> np.ndarray:
     """Local cooling modulus per node: V/A over geodesic balls (mm).
 
     For each node and radius r: (sum of volumes of tets whose 4 nodes all lie
     within geodesic distance r of the node) / (area of cooling triangles whose
     3 nodes all lie within r). Balls with no cooling area get the bbox-diagonal
-    sentinel (consistent with the other distance sentinels).
-
-    `exact=True` (default): the reference per-node computation -- byte-identical
-    to the original (pre-subsample) implementation regardless of `subsample`.
-
-    `exact=False` with `subsample=k` (int): compute the exact ball modulus only
-    on every k-th node ("computed" nodes), then fill every other node by
-    inverse-distance-weighted (IDW) interpolation over its `_IDW_K` (=4)
-    nearest computed nodes, measured by mesh-edge geodesic distance
-    (multi-source Dijkstra over `edges`/`edge_len`, reusing `_edge_csr`; see
-    `_idw_geodesic_fill`). This trades a small amount of accuracy for roughly
-    a `k`-fold cut in the number of expensive geodesic-ball evaluations (the
-    dense-mesh bottleneck). If `subsample >= N` (or is None/<=1) this
-    degenerates to the exact path.
-
-    HONEST ACCURACY NOTE (see tests/test_local_modulus.py::
-    test_subsample_per_radius_column_error and task-2-report.md for the full
-    sweep): at the production radii (5/15/40 mm) and subsample=4, per-radius-
-    column median relative error on a real ~4k-node training mesh (casting-
-    scale bbox) is roughly 0% at r=5mm but ~25-30% at r=15mm and ~6% at
-    r=40mm -- i.e. the mid radius does NOT meet a 5% target. This is because
-    local_modulus is a ball-membership STEP FUNCTION of geodesic position
-    (a tet/triangle is either fully inside a ball or not), not a smoothly-
-    varying field -- averaging several dissimilar neighbors (larger IDW k)
-    measurably makes the mid-radius error WORSE, not better (single-nearest,
-    k=1, beats IDW k=4-8 there); k=4 is used as the best compromise within
-    the requested range, not because it minimizes error in isolation.
-    CONSEQUENCE: this subsample path is intended for DENSE-mesh inference
-    (FT3), where it is a useful order-of-magnitude speed/accuracy trade
-    against an otherwise-prohibitive `exact` cost -- it is NOT a validated
-    drop-in replacement for `exact=True` in the training-density feature
-    cache (FT2a), where `exact` is already fast enough (seconds per mesh)
-    and should stay the default. Do not switch `build_dataset`'s v4 feature
-    cache to `exact=False` without re-validating per-radius accuracy on the
-    actual dense meshes it will be used on."""
+    sentinel (consistent with the other distance sentinels)."""
     N = node_xyz.shape[0]
     radii = np.asarray(radii, dtype=float)
     bbox_diag = float(np.linalg.norm(node_xyz.max(0) - node_xyz.min(0)))
@@ -231,22 +122,7 @@ def local_modulus(node_xyz: np.ndarray, tets: np.ndarray, cool_tri: np.ndarray,
     tvol = tet_volumes(node_xyz, tets)
     tarea = triangle_areas(node_xyz, cool_tri)
     g = _edge_csr(N, edges, edge_len)
-
-    use_subsample = (not exact) and subsample is not None and 1 < subsample < N
-    query_idx = (np.arange(0, N, subsample, dtype=np.int64) if use_subsample
-                 else np.arange(N, dtype=np.int64))
-
-    vals = _local_modulus_for_query(tets, cool_tri, radii, chunk, query_idx, g,
-                                    tvol, tarea, bbox_diag)
-    out[query_idx] = vals
-
-    if use_subsample and query_idx.shape[0] < N:
-        missing = np.setdiff1d(np.arange(N), query_idx, assume_unique=True)
-        if missing.shape[0]:
-            # inverse-distance-weighted fill from the k nearest computed
-            # nodes, along mesh-edge geodesics (smoother than single-nearest)
-            dist_from_computed = dijkstra(g, directed=False, indices=query_idx,
-                                          min_only=False)              # (M, N)
-            out[missing] = _idw_geodesic_fill(vals, dist_from_computed, missing, bbox_diag)
-
+    query_idx = np.arange(N, dtype=np.int64)
+    out[query_idx] = _local_modulus_for_query(tets, cool_tri, radii, chunk, query_idx, g,
+                                              tvol, tarea, bbox_diag)
     return out.astype(np.float32)

src/gnn_dataset.py

@@ -1,51 +0,0 @@
-"""PyG dataset over the v2/v3 feature cache. Builds either a kNN graph or the
-SP1 union graph (mesh edges union kNN, `graph="union"`, read from the cache's
-`mesh_edges`) with mm edge attributes on load; optional random-rotation
-augmentation. Features and the log1p target are standardized with the cache's
-global stats."""
-import os.path as osp
-import numpy as np
-import torch
-from torch.utils.data import Dataset
-from torch_geometric.data import Data
-
-from src.graph_build import (knn_edge_index, edge_attr, random_rotation_matrix, rotate,
-                             union_edge_index, union_edge_attr)
-from src.dataload import load_norm
-
-
-class SolidGraphDataset(Dataset):
-    def __init__(self, features_dir, stems, k=16, augment=False, norm=None, graph="knn"):
-        self.dir = features_dir
-        self.stems = list(stems)
-        self.k = k
-        self.augment = augment
-        self.graph = graph
-        self.norm = norm if norm is not None else load_norm(features_dir)
-        self.fmean = torch.tensor(self.norm["feat_mean"], dtype=torch.float32)
-        self.fstd = torch.tensor(self.norm["feat_std"], dtype=torch.float32)
-        self.ym = float(self.norm["y_log_mean"]); self.ys = float(self.norm["y_log_std"])
-
-    def __len__(self):
-        return len(self.stems)
-
-    def __getitem__(self, i):
-        d = np.load(osp.join(self.dir, "graphs", f"{self.stems[i]}.npz"))
-        x = (torch.tensor(d["X"], dtype=torch.float32) - self.fmean) / self.fstd
-        pos = torch.tensor(d["pos"], dtype=torch.float32)
-        y_log = torch.log1p(torch.tensor(d["y"], dtype=torch.float32))
-        y = (y_log - self.ym) / self.ys
-        g = torch.tensor(d["g"], dtype=torch.float).reshape(1, -1) if "g" in d.files else None
-        if self.augment:
-            pos = rotate(pos, random_rotation_matrix())
-        if self.graph == "union":
-            if "mesh_edges" not in d.files:
-                raise KeyError(f"{self.stems[i]}: cache lacks mesh_edges — "
-                               f"rebuild with build_dataset --feature-version 3")
-            me = torch.tensor(d["mesh_edges"], dtype=torch.long)
-            ei, is_mesh = union_edge_index(pos, me, k=self.k)
-            ea = union_edge_attr(pos, ei, is_mesh)
-        else:
-            ei = knn_edge_index(pos, k=self.k)
-            ea = edge_attr(pos, ei)
-        return Data(x=x, edge_index=ei, edge_attr=ea, y=y, pos=pos, g=g)

src/graph_build.py

@@ -1,52 +0,0 @@
-"""Mesh-agnostic graph construction: kNN edges with millimetre edge attributes,
-and random 3D-rotation augmentation (physics is rotation-invariant)."""
-import torch
-from torch_geometric.nn import knn_graph
-from torch_geometric.utils import to_undirected
-
-
-def knn_edge_index(pos: torch.Tensor, k: int = 16) -> torch.Tensor:
-    ei = knn_graph(pos, k=k, loop=False)      # directed k-nearest
-    return to_undirected(ei)                   # symmetric
-
-
-def edge_attr(pos: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor:
-    src, dst = edge_index
-    d = pos[dst] - pos[src]                     # (E,3) mm
-    return torch.cat([d, d.norm(dim=1, keepdim=True)], dim=1)   # (E,4)
-
-
-def random_rotation_matrix(seed: int | None = None) -> torch.Tensor:
-    g = torch.Generator().manual_seed(seed) if seed is not None else None
-    # QR of a random Gaussian matrix gives a Haar-random orthogonal matrix
-    a = torch.randn(3, 3, generator=g)
-    q, r = torch.linalg.qr(a)
-    q = q * torch.sign(torch.diagonal(r))       # fix signs
-    if torch.det(q) < 0:                        # ensure proper rotation (det +1)
-        q[:, 0] = -q[:, 0]
-    return q
-
-
-def rotate(pos: torch.Tensor, R: torch.Tensor) -> torch.Tensor:
-    return pos @ R.to(pos.device, pos.dtype).t()
-
-
-def union_edge_index(pos: torch.Tensor, mesh_edges: torch.Tensor, k: int = 16):
-    """Undirected union of tet-mesh edges and kNN edges, deduplicated.
-
-    mesh_edges: (M,2) long tensor of unique node pairs (any orientation).
-    Returns (edge_index (2,E) long, is_mesh (E,) float32 flag)."""
-    n = pos.shape[0]
-    knn = knn_edge_index(pos, k=k).to(torch.long)
-    me = mesh_edges.t().to(torch.long)
-    me = torch.cat([me, me.flip(0)], dim=1)                    # both directions
-    mesh_key = me[0] * n + me[1]
-    keys = torch.unique(torch.cat([mesh_key, knn[0] * n + knn[1]]))
-    ei = torch.stack([keys // n, keys % n])
-    is_mesh = torch.isin(keys, mesh_key).to(torch.float32)
-    return ei, is_mesh
-
-
-def union_edge_attr(pos: torch.Tensor, edge_index: torch.Tensor,
-                    is_mesh: torch.Tensor) -> torch.Tensor:
-    return torch.cat([edge_attr(pos, edge_index), is_mesh.unsqueeze(1)], dim=1)  # (E,5)

src/hybrid.py

@@ -1,15 +1,17 @@
-"""Hybrid predictor: GNN field *shape* + baseline absolute *scale*.
+"""Recomposition: combine the GNN field shape with the two heads' targets.
 
-The GNN captures the within-graph solidification field shape almost perfectly
-(per-graph corr ~0.98, hot-spot ~0.88) but has a per-graph absolute-scale bias.
-The physics-feature baseline (HGB) predicts the per-graph magnitude well. The
-hybrid keeps the GNN's deviations from its own per-graph mean and re-anchors the
-per-graph level to the baseline's per-graph mean:
+The GNN predicts the spatial pattern of the field well but is less reliable
+about its absolute level and spread. The heads predict the part's total time
+(used as the 99.5% quantile of the field) and the standard deviation of the
+field. `hybrid_field_soft` maps the GNN field affinely,
 
-    pred_i = gnn_i - mean_g(gnn) + mean_g(baseline)
+    pred(a, m) = clip(a * (gnn - mean(gnn)) + m, 0, None)
 
-This preserves the GNN's hot-spot/variance advantage while fixing the offset,
-and beats both parents on MAE and hot-spot across all splits.
+and picks (a, m) to minimize
+
+    J = w_total * r_total**2 + w_std * r_std**2,
+    r_total = (quantile(pred, q) - total) / total
+    r_std   = (std(pred) - std_target) / std_target.
 """
 import logging
 
@@ -18,52 +20,9 @@ import numpy as np
 logger = logging.getLogger(__name__)
 
 
-def hybrid_field(gnn_seconds: np.ndarray, baseline_seconds: np.ndarray) -> np.ndarray:
-    """Combine one casting's per-node predictions (both in seconds, same node order).
-
-    Shape from `gnn_seconds`, absolute level from `baseline_seconds`. Clipped at 0.
-    """
-    g = np.asarray(gnn_seconds, dtype=np.float64)
-    b = np.asarray(baseline_seconds, dtype=np.float64)
-    out = g - g.mean() + b.mean()
-    return np.clip(out, 0.0, None)
-
-
 def hybrid_field_max_anchor(gnn_seconds: np.ndarray, total_seconds: float,
                             q: float = 0.995) -> np.ndarray:
-    """v4 hybrid recomposition: GNN field *shape* + a per-graph quantile-anchor scale.
-
-    `pred_i = gnn_i - anchor + total`, clipped at 0, where
-    `anchor = np.quantile(gnn_seconds, q)`. Unlike `hybrid_field`'s
-    MEAN-anchor (correct for a genuine per-node baseline), the FT2b GBM scale
-    head (`src.gbm_ft2b`) predicts a single per-graph SCALAR -- the total
-    solidification time (the field's true max), not a per-node mean.
-    Broadcasting that scalar and MEAN-anchoring would force
-    `mean(pred) == total`, inflating every node by `(true_max - true_mean)`.
-    Anchoring to a high quantile of the GNN field instead makes
-    `max(pred) ~= total` (the quantity actually predicted) while preserving
-    the GNN's field shape (its deviations from that anchor).
-
-    `q=1.0` reproduces the previous exact-max behavior bit-for-bit
-    (`np.quantile(x, 1.0) == x.max()`): `max(pred) == total` exactly. The
-    DEFAULT `q=0.995` exists because the exact-max anchor is fragile to a
-    handful of outlier nodes -- a promotion-gate smoke test found a few
-    inflated GNN log-space predictions on a real casting dragging ~90% of
-    the field below the zero clip, because the single max node set the
-    entire graph's offset. Anchoring to the 99.5th percentile instead makes
-    the offset robust to a small number of outliers, at the cost that nodes
-    ABOVE the anchor (up to ~0.5% of nodes, by construction) may now exceed
-    `total` slightly. This does not affect hot-spot ranking or field shape:
-    the anchor choice is still a per-graph ADDITIVE CONSTANT shift, so
-    within-graph node ordering (`argsort`) is unchanged, and only the
-    reported `float(pred.max())` total-time summary can differ slightly
-    from `total`. See `src/evaluate_ft2b.py`'s module docstring ("v4 hybrid
-    recomposition: QUANTILE-anchor (q=0.995), not MEAN-anchor") for the full anchor-choice
-    rationale -- this is the single source; `evaluate_ft2b` imports this
-    function under its former private name, and the same default `q` applies
-    everywhere (gate evaluation, checkpoint selection, and production
-    inference all share this one function).
-    """
+    """Shift the GNN field so its q-quantile equals `total_seconds`, clip at 0."""
     g = np.asarray(gnn_seconds, dtype=np.float64)
     anchor = np.quantile(g, q)
     out = g - anchor + float(total_seconds)
@@ -72,169 +31,31 @@ def hybrid_field_max_anchor(gnn_seconds: np.ndarray, total_seconds: float,
 
 def hybrid_field_affine(gnn_seconds: np.ndarray, total_seconds: float,
                         std_seconds: float, q: float = 0.995,
-                        a_clamp: tuple = (0.3, 5.0),
-                        lambda_shrink: float = 1.0) -> tuple:
-    """Affine hybrid recomposition: GNN field *shape* rescaled to match TWO scalar
-    targets from the GBM scale head -- a per-graph total AND a per-graph spread --
-    instead of `hybrid_field_max_anchor`'s single total-only anchor.
-
-    `hybrid_field_max_anchor` only re-anchors the GNN field's level (an additive
-    shift); it inherits the GNN's own spread (std) verbatim. FT3 gate 4
-    (`runs/ft3_stage1/FT3_GATE_REPORT.md`) found that spread can be badly wrong
-    even when the level and hot-spot ranking are fixed: the GNN's std can
-    collapse relative to ground truth on regimes it under-fits, and a pure
-    level-shift cannot repair that -- e.g. the report's BUKSA case, where a
-    refit GBM anchor overshot the true total by ~+190 s because the GNN field's
-    shape carried a whole-field bias that a level-only anchor could not absorb.
-    This function additionally SCALES the field's deviations from its own mean
-    by a factor `a`, then re-anchors so `np.quantile(pred, q) == total_seconds`
-    (same quantile-anchor rationale as `hybrid_field_max_anchor` -- robust to a
-    handful of outlier nodes, at the cost that nodes above the anchor may
-    slightly exceed `total_seconds`). Scaling deviations by a positive constant
-    and then shifting by an additive constant is a strictly increasing affine
-    map, so within-graph node ordering (`argsort`, hot-spot ranking) is
-    unchanged from the raw GNN field -- exactly as for `hybrid_field_max_anchor`.
-
-    AMENDMENT (owner-approved 2026-07-31, see the "AMENDMENT" block in
-    `docs/superpowers/specs/2026-07-31-variance-recomposition-design.md`): the
-    field is 0-bounded (solidification time can't be negative) but this affine
-    map is symmetric around the field's own mean, so on real (skewed, clip-prone)
-    parts a naive `a = std_seconds / std(gnn_seconds)` collapses under its own
-    clip. The e_val eval matrix found 21-34% of nodes clipping and 33-35% of
-    chills/exo stems clamp-saturated at the old `a<=3.0` ceiling, leaving group
-    std_ratio at ~0.83 against the 0.9 floor. `a` is therefore solved so the
-    POST-CLIP field matches `std_seconds`, not the naive pre-clip ratio:
-
-        c      = gnn_seconds - mean(gnn_seconds)
-        d      = c - quantile(c, q)         # constant vector; independent of a
-        raw(a) = a * d + total_seconds      # quantile(raw(a), q) == total_seconds for any a>0
-        f(a)   = std(clip(raw(a), 0, None))
-        a      = solve f(a) == std_seconds over a in a_clamp
-
-    (`raw(a) = a*d + total_seconds` follows from the quantile-anchor property:
-    anchoring `a*c` at `total_seconds` gives `m(a) = total_seconds - a*quantile(c, q)`,
-    so `a*c + m(a) = a*(c - quantile(c, q)) + total_seconds = a*d + total_seconds`.)
-
-    CORRECTION (2026-07-31 review): `f` is **NOT monotone** on realistic
-    fields. Below `a_clipstart = total_seconds / max(-d.min(), eps)` nothing
-    clips and `f(a) = a * std(gnn_seconds)` exactly (the pre-amendment closed
-    form). PAST `a_clipstart`, clipping starts eating the field's low tail;
-    on real (23/25 sampled `data/features_v4` graphs) `y` fields this makes
-    `f` RISE to an interior peak just past `a_clipstart` and then FALL as
-    growing clip mass collapses variance faster than the surviving nodes'
-    growing spread can compensate. A naive `f(hi) < std_target => clamp to
-    hi` pre-check is therefore wrong: it can fire even when a perfectly good
-    (much-less-clipped) root exists on the RISING part of the curve, close to
-    `a_clipstart` -- on real data this false-clamp inflated MAE from 85.1 (at
-    the true root) to 359.5 (at the falsely-clamped `a=5.0`), because clamping
-    to `hi` clips 61-90% of the field instead of the true root's ~5-30%.
-
-    Solve semantics, in order:
-      1. If `a_raw = std_seconds / std(gnn_seconds) <= a_clipstart`: nothing
-         clips at `a_raw`, so use the closed form directly (clamped to
-         `a_clamp` -- fast path, identical to the pre-amendment closed form
-         to floating-point precision).
-      2. Otherwise, coarse-scan `a` over `[max(a_clipstart, a_clamp[0]),
-         a_clamp[1]]` (64 points) for the FIRST sign change of
-         `f(a) - std_seconds`. Because `f` starts below `std_seconds` at
-         `a_clipstart` (by construction of this branch) and is scanned
-         low-to-high, the first crossing is the SMALLEST root -- i.e. the
-         LEAST-clipped `a` that hits the target, always preferred over any
-         larger root further out on the peak's falling side. Bisect inside
-         that bracket to refine it; `info["a_clamped"]` is `False`.
-      3. If no crossing exists anywhere on the scan (`std_seconds` exceeds
-         even the interior peak, or the scan range is empty), the target is
-         genuinely unreachable: pick the SCANNED POINT whose `f` is closest
-         to `std_seconds` (not blindly `a_clamp[1]` -- the peak is often a
-         much better fit than the ceiling, exactly the failure mode above)
-         and set `info["a_clamped"] = True`.
-
-    `info["clip_frac"]` reports how much of the FINAL field hit the 0 floor,
-    so callers can flag graphs where the GBM's total/std pair is inconsistent
-    with the GNN's shape; if `clip_frac > 0.5` (a majority of the field
-    clipped -- a strong scale/anchor-inconsistency signal) this also logs a
-    warning.
-
-    Degenerate fallback: if the GNN field is (numerically) flat, `a` is
-    undefined (division by ~0), so this falls back to
-    `hybrid_field_max_anchor(gnn_seconds, total_seconds, q=q)` and logs a
-    warning; `info["degenerate_fallback"]` is `True` in that case and `a` is
-    reported as `1.0` (a no-op scale, since there is no shape to rescale).
-    This guard is unaffected by the amendment above -- it is checked before
-    any `a`-solve is attempted.
-
-    AMENDMENT 2 (owner-approved 2026-07-31, see the "AMENDMENT 2" block in
-    `docs/superpowers/specs/2026-07-31-variance-recomposition-design.md`):
-    the corrected solver exposed the next layer of error -- the GBM std
-    head's own per-part noise (p10-p90 ~ 0.66-1.36x; BUKSA gold
-    over-predicted at ~1.3x) feeds straight into MAE through the spread it
-    now honestly hits. `lambda_shrink` (default `1.0`, fully trusting
-    `std_seconds`) lets a caller partially discount that noisy target via
-    GEOMETRIC interpolation between the GNN's OWN spread and the GBM's
-    predicted spread, applied BEFORE the (unchanged) clip-aware solve above:
-
-        std_eff = std(gnn_seconds)**(1 - lambda_shrink) * std_seconds**lambda_shrink
-
-    `lambda_shrink=1.0` reproduces the full-trust behavior exactly (`std_eff
-    == std_seconds`) -- this is the ONLY code path change, gated by `if
-    lambda_shrink != 1.0:` so the default is untouched byte-for-bit.
-    `lambda_shrink=0.0` degenerates to `std_eff == std(gnn_seconds)`, i.e.
-    the solve targets the GNN's own (unscaled) spread -- in a no-clip case
-    this makes `a == 1.0` and the result IS `hybrid_field_max_anchor`'s
-    output exactly (no spread rescaling at all, quantile-anchor-only level
-    shift). Intermediate `lambda_shrink` partially shrinks the GBM's spread
-    prediction toward the GNN's own, geometric-mean-style. Selection
-    protocol (owner): `lambda_shrink` is tuned on `e_val` only (per-BC
-    `std_ratio` vs. the 0.9 floor + MAE); the GOLD assemblies stay a pure
-    held-out test, evaluated once at the chosen value.
-
-    `std_seconds <= 0` is a degenerate SCALE target (the geometric
-    interpolation is undefined/nonsensical for a non-positive base) --
-    ONLY WHEN `lambda_shrink != 1.0` (the default path never inspects
-    `std_seconds`'s sign, matching pre-amendment-2 behavior exactly), this
-    falls back to `hybrid_field_max_anchor` with the same loud-warning
-    pattern as the degenerate-GNN-field guard above.
-
-    Returns `(pred, info)` where `info = {"a": float, "a_clamped": bool,
-    "degenerate_fallback": bool, "clip_frac": float, "lambda_shrink": float,
-    "std_target_effective": float}`. NOTE `a_clamped` semantics since the
-    smallest-root correction: True means "the requested std target was
-    UNREACHABLE" (closest-fit `a` returned, which may be an interior scan
-    point) -- it no longer implies `a` sits exactly on an `a_clamp`
-    boundary. `std_target_effective` is `std_eff` above (`== std_seconds`
-    when `lambda_shrink == 1.0`).
-    """
+                        a_clamp: tuple = (0.3, 5.0)) -> tuple:
+    """Hard two-target fit: the q-quantile is anchored exactly at
+    `total_seconds` and the scale `a` is solved so that the clipped field's
+    standard deviation matches `std_seconds`.
+
+    Anchoring the quantile gives raw(a) = a * d + total with d fixed by the
+    GNN field. Below a_clipstart nothing is clipped and a = std/std(gnn) in
+    closed form. Past it, the post-clip spread is not monotone in a (it rises
+    to an interior peak and then falls), so the smallest root is found by a
+    coarse scan followed by bisection; if no root exists the closest scanned
+    point is used and `a_clamped` is set. Returns (field, info)."""
     g = np.asarray(gnn_seconds, dtype=np.float64)
     s = float(np.std(g))
-    lambda_shrink = float(lambda_shrink)
     if s < 1e-9:
         logger.warning("hybrid_field_affine: degenerate flat GNN field (std=%.3g) -- "
                        "falling back to quantile-anchor recomposition", s)
         pred = hybrid_field_max_anchor(g, total_seconds, q=q)
         return pred, {"a": 1.0, "a_clamped": False, "degenerate_fallback": True,
                       "clip_frac": float((pred == 0.0).mean()),
-                      "lambda_shrink": lambda_shrink,
                       "std_target_effective": float(std_seconds)}
 
     std_target = float(std_seconds)
-    if lambda_shrink != 1.0:
-        if std_target <= 0.0:
-            logger.warning("hybrid_field_affine: degenerate std target "
-                           "(std_seconds=%.3g, lambda_shrink=%.3g) -- falling back to "
-                           "quantile-anchor recomposition", std_target, lambda_shrink)
-            pred = hybrid_field_max_anchor(g, total_seconds, q=q)
-            return pred, {"a": 1.0, "a_clamped": False, "degenerate_fallback": True,
-                          "clip_frac": float((pred == 0.0).mean()),
-                          "lambda_shrink": lambda_shrink,
-                          "std_target_effective": std_target}
-        std_target = (s ** (1.0 - lambda_shrink)) * (std_target ** lambda_shrink)
-    std_target_effective = std_target
-
     lo, hi = float(a_clamp[0]), float(a_clamp[1])
     total = float(total_seconds)
 
-    # `d` and `a_clipstart` depend only on the (fixed) GNN field, the quantile
-    # `q`, and `total` -- computed ONCE, not inside any search loop.
     c = g - g.mean()
     shift = float(np.quantile(c, q))
     d = c - shift
@@ -250,17 +71,12 @@ def hybrid_field_affine(gnn_seconds: np.ndarray, total_seconds: float,
     a_raw = std_target / s
 
     if a_raw <= a_clipstart:
-        # At a_raw <= a_clipstart the SOLUTION a_raw clips nothing, so
-        # f(a_raw) = a_raw*s exactly (closed form; identical to the
-        # pre-amendment behavior to floating-point precision). NOTE: if
-        # a_raw < lo the clamp to lo can still clip (total << spread).
+        # Nothing is clipped at a_raw, so the closed form is exact. The clamp
+        # to `lo` can still clip when the total is small next to the spread.
         a = float(np.clip(a_raw, lo, hi))
         a_clamped = bool(a != a_raw)
         raw = _raw_for_a(a)
     else:
-        # a_raw would clip. f is not monotone past a_clipstart -- it rises to
-        # an interior peak, then falls. Coarse-scan for the smallest (first,
-        # least-clipping) root and bisect to refine it.
         lo_search = max(a_clipstart, lo)
         if lo_search >= hi:
             a, a_clamped = hi, True
@@ -294,9 +110,8 @@ def hybrid_field_affine(gnn_seconds: np.ndarray, total_seconds: float,
                 a_clamped = False
                 raw = _raw_for_a(a)
             else:
-                # No crossing anywhere on the scan: std_target is unreachable
-                # (e.g. it exceeds the interior peak). Pick the scanned point
-                # CLOSEST to the target rather than blindly a_clamp[1].
+                # The target exceeds the reachable spread: take the closest
+                # scanned point instead of the upper clamp.
                 best_idx = int(np.argmin(np.abs(diffs)))
                 a = float(scan_a[best_idx])
                 a_clamped = True
@@ -309,166 +124,40 @@ def hybrid_field_affine(gnn_seconds: np.ndarray, total_seconds: float,
                        "with the GNN shape", clip_frac)
     return np.clip(raw, 0.0, None), {"a": a, "a_clamped": a_clamped,
                                      "degenerate_fallback": False, "clip_frac": clip_frac,
-                                     "lambda_shrink": lambda_shrink,
-                                     "std_target_effective": std_target_effective}
+                                     "std_target_effective": std_target}
 
 
 def hybrid_field_soft(gnn_seconds: np.ndarray, total_seconds: float, std_seconds: float,
                       q: float = 0.995, a_clamp: tuple = (0.3, 5.0),
-                      w_total: float = 1.0, w_std: float = 1.0,
-                      lambda_shrink: float = 1.0) -> tuple:
-    """Soft-anchor recomposition: `hybrid_field_affine` generalized to a WEIGHTED
-    two-target objective, so the total anchor may itself deviate a few percent
-    when that trade buys a much closer spread match.
-
-    WHY (see `docs/superpowers/specs/2026-07-31-soft-anchor-design.md` §1):
-    `hybrid_field_affine`'s quantile-anchor is a HARD constraint
-    (`quantile(pred, q) == total_seconds` exactly, by construction of its
-    `m(a)`), and the field is 0-clipped from below. Those two constraints
-    together form a ceiling: on some real parts (e_val exo stems, std_ratio
-    0.874 < the 0.9 floor) there is NO `a` in `a_clamp` for which the
-    post-clip field can reach `std_seconds` while ALSO satisfying the hard
-    top-anchor exactly -- the affine solve already returns the closest-fit
-    root, and that is provably the best available under the hard constraint.
-    Relaxing the top anchor into a soft (penalized) term recovers a strictly
-    larger solution space: allowing `quantile(pred, q)` to sit a bit off
-    `total_seconds` can unlock a much less-clipped field whose spread lands
-    far closer to `std_seconds`, at a small, bounded cost to the total.
-
-    Objective (both residuals normalized by their own target, so the weights
-    are scale-free across parts of wildly different absolute magnitude):
-
-        pred(a, m) = clip(a * (gnn - mean(gnn)) + m, 0, None)
-        r_total    = (quantile(pred, q) - total_seconds) / total_seconds
-        r_std      = (std(pred) - std_target_eff) / std_target_eff
-        J(a, m)    = w_total * r_total**2 + w_std * r_std**2
-
-    where `std_target_eff` is `std_seconds` shrunk toward the GNN's own
-    spread by `lambda_shrink` (identical convention to `hybrid_field_affine`;
-    delegated to it -- see its AMENDMENT 2 docstring). `a` ranges over
-    `a_clamp`; `m` is a free additive constant (no longer tied to the hard
-    anchor's `m(a) = total - a*quantile(centered_gnn, q)`, though that
-    formula is exactly where the SEED point comes from, below).
-
-    Solver: `hybrid_field_affine`'s own clip-aware smallest-root solve
-    supplies the SEED `(a_seed, m_seed)` -- `a_seed` is its returned `a`;
-    `m_seed` is recovered from the identity `m(a) = total_seconds -
-    a*quantile(gnn - mean(gnn), q)` that its solve always satisfies (true in
-    every branch, including both degenerate fallbacks: `hybrid_field_max_anchor`
-    is itself `a=1, m = total - quantile(centered_gnn, q)` after the same
-    quantile-shift-equivariance identity). A coarse grid -- 64 log-spaced `a`
-    over `a_clamp` x 33 linear `m` CENTERED ON THE SEED, `m_seed +
-    total_seconds*[-0.3, +0.3]` -- is evaluated (vectorized per row: one
-    clip+quantile+std pass per `a`, across all `m` at once), PLUS the single
-    exact seed point `(a_seed, m_seed)` folded into the running best (not
-    cross-producted against either axis -- that would silently widen the
-    search past the intended band). REVIEW FIX (2026-07-31, round 1,
-    CRITICAL-1): the window was originally the ABSOLUTE
-    `total_seconds*[0.7,1.3]` (centered on `total_seconds`, not on
-    `m_seed`) -- measured on 58 real `features_v4` fields, `m_seed/total in
-    [0.03, 0.61]`, i.e. `m_seed` was NEVER inside that window, so the grid
-    could only ever explore points strictly worse than the seed and the
-    solver never engaged on the exo regime it exists for. Centering the
-    window on `m_seed` is what actually lets the search reach nearby, less-
-    clipped `(a, m)` combinations. ONE 9x9 local refinement grid (log-spaced
-    `a` / linear `m`, windowed to one coarse grid cell around whichever point
-    -- grid or seed -- is best so far, bounded within the same centered
-    window) is then evaluated the same way.
-
-    DEGENERATION GUARANTEE: `J_seed` is computed EXACTLY at `(a_seed, m_seed)`
-    before any grid search runs (not interpolated from a grid cell -- it is a
-    direct, separate evaluation). A grid or refinement candidate is accepted
-    ONLY if it beats the seed by a strict margin (`J_candidate < J_seed -
-    1e-12`); otherwise this returns the SEED array unchanged (same object,
-    not a recomputation) with the seed's own residuals. This is the guard
-    against the v5 track's earlier critical finding that `f` (the affine
-    solve's post-clip std, as a function of `a`) is NOT monotone/convex on
-    real fields -- the same non-convexity applies to `J(a, m)` here, so the
-    grid+refinement scheme never ASSUMES a single interior optimum; it only
-    ever accepts a point that is verifiably better than the seed by direct
-    evaluation. On any field where the affine seed already satisfies both
-    targets (no-clip case), this makes `hybrid_field_soft` numerically
-    identical to `hybrid_field_affine` -- and, as a PERFORMANCE consequence
-    (REVIEW FIX, CRITICAL-2a), when `J_seed < 1e-12` this returns immediately
-    WITHOUT building or evaluating any grid at all (healthy/no-clip parts pay
-    only the affine solve's own cost). When the grid does run (REVIEW FIX,
-    CRITICAL-2b), its per-`a` quantile term is computed via the closed-form
-    `a*shift + m` (exact whenever `clip_frac < 1 - q`, which always holds in
-    practice -- see the in-code comment above `_grid_eval`) instead of a per-
-    point `np.quantile`/`np.partition` call, and the grid's winning candidate
-    is re-evaluated with the exact `_residuals()` helper before the
-    accept/reject decision, so no approximation reaches the returned
-    `info`/prediction -- this keeps a 10k-node solve inside the
-    ~200ms budget even when the search engages.
-
-    `lambda_shrink` and the degenerate-field/degenerate-std-target fallbacks
-    are delegated verbatim to `hybrid_field_affine` (same warnings, same
-    `hybrid_field_max_anchor` escape hatch) -- when `info["degenerate_fallback"]`
-    comes back `True` there is no field shape left to trade off, so this
-    returns that fallback's prediction unchanged (no grid search).
-
-    Returns `(pred, info)` where `info = {"a": float, "m": float,
-    "target_unreachable": bool, "degenerate_fallback": bool,
-    "clip_frac": float, "total_deviation_frac": float, "j_total_resid": float,
-    "j_std_resid": float, "lambda_shrink": float,
-    "std_target_effective": float}`. `target_unreachable` repurposes
-    `hybrid_field_affine`'s `a_clamped` (`True` means the HARD two-target
-    problem the affine seed solved was itself unreachable at some `a` -- the
-    signal that motivated relaxing to a soft objective in the first place).
-    `clip_frac` is measured on the FINAL (returned) field. `total_deviation_frac
-    = r_total` of the final field (positive = the top anchor rose above
-    `total_seconds`). `j_total_resid`/`j_std_resid` are `J`'s two weighted,
-    squared terms at the final `(a, m)` (`J == j_total_resid + j_std_resid`).
-
-    REVIEW FIX (round 1, LOW): `w_total`/`w_std` must both be `>= 0` (a
-    negative weight REWARDS moving away from that target instead of
-    penalizing deviation from it -- a reviewer-demonstrated `w_total=-1.0`
-    silently drove `a` to the `a_clamp` ceiling with a +273% total
-    deviation, since minimizing `J` then means MAXIMIZING `r_total**2`) and
-    may not both be exactly `0` (the objective `J(a, m)` would be
-    identically zero everywhere, leaving the grid/refinement search no
-    signal to optimize against at all) -- either violation raises a loud
-    `ValueError` before any solve is attempted.
-    """
+                      w_total: float = 1.0, w_std: float = 1.0) -> tuple:
+    """Soft two-target fit (see the module docstring for J).
+
+    Seeded by `hybrid_field_affine`. When the seed already meets both
+    targets it is returned unchanged; otherwise a 64 x 33 grid over (a, m),
+    a 9 x 9 refinement around the best point and an exact re-check decide
+    whether letting the total move buys a better spread match. Returns
+    (field, info) with info keys a, m, target_unreachable,
+    degenerate_fallback, clip_frac, total_deviation_frac, j_total_resid,
+    j_std_resid, std_target_effective."""
     g = np.asarray(gnn_seconds, dtype=np.float64)
     total = float(total_seconds)
     q = float(q)
     w_total = float(w_total)
     w_std = float(w_std)
 
-    # REVIEW FIX (round 1, LOW): a negative weight would reward the search
-    # for moving AWAY from that target rather than penalizing deviation from
-    # it (minimizing J then means MAXIMIZING that residual's square), and
-    # both-zero weights make J identically zero everywhere -- fail loudly
-    # rather than silently returning a meaningless/adversarial optimum.
     if w_total < 0.0 or w_std < 0.0:
         raise ValueError(
             f"hybrid_field_soft: w_total/w_std must be >= 0 (got "
-            f"w_total={w_total!r}, w_std={w_std!r}) -- a negative weight would "
-            f"reward moving AWAY from that target instead of penalizing "
-            f"deviation from it")
+            f"w_total={w_total!r}, w_std={w_std!r})")
     if w_total == 0.0 and w_std == 0.0:
-        raise ValueError(
-            "hybrid_field_soft: w_total and w_std cannot both be 0 -- the "
-            "objective J(a, m) would be identically zero everywhere, leaving "
-            "the search no signal to optimize against")
-
-    # REVIEW FIX (LOW-6): np.logspace(log10(a_clamp[0]), ...) below silently
-    # produces NaN/-inf for a non-positive lower clamp; fail loudly instead,
-    # regardless of whether this call ends up taking the no-grid fast path.
+        raise ValueError("hybrid_field_soft: w_total and w_std cannot both be 0")
     if float(a_clamp[0]) <= 0.0:
         raise ValueError(f"hybrid_field_soft: a_clamp[0] must be > 0 (got {a_clamp[0]!r}) "
-                         "-- the a-grid is log-spaced and cannot include a non-positive "
-                         "lower bound")
+                         "-- the a-grid is log-spaced")
 
-    pred_seed, info_affine = hybrid_field_affine(g, total, std_seconds, q=q,
-                                                 a_clamp=a_clamp,
-                                                 lambda_shrink=lambda_shrink)
+    pred_seed, info_affine = hybrid_field_affine(g, total, std_seconds, q=q, a_clamp=a_clamp)
 
-    # `c`/`shift` recover m(a) = total - a*shift for ANY branch of the affine
-    # solve, including both its degenerate fallbacks (each is exactly a=1,
-    # m = total - shift, by the quantile-shift-equivariance identity used in
-    # `hybrid_field_affine`'s own AMENDMENT derivation).
+    # For every branch of the affine solve, m(a) = total - a * quantile(c, q).
     c = g - g.mean()
     shift = float(np.quantile(c, q))
     a_seed = float(info_affine["a"])
@@ -494,7 +183,6 @@ def hybrid_field_soft(gnn_seconds: np.ndarray, total_seconds: float, std_seconds
             "total_deviation_frac": final_rt,
             "j_total_resid": float(w_total * final_rt ** 2),
             "j_std_resid": float(w_std * final_rs ** 2),
-            "lambda_shrink": info_affine["lambda_shrink"],
             "std_target_effective": std_target_eff,
         }
 
@@ -508,18 +196,13 @@ def hybrid_field_soft(gnn_seconds: np.ndarray, total_seconds: float, std_seconds
             "total_deviation_frac": r_total_seed,
             "j_total_resid": float(w_total * r_total_seed ** 2),
             "j_std_resid": float(w_std * r_std_seed ** 2),
-            "lambda_shrink": info_affine["lambda_shrink"],
             "std_target_effective": std_target_eff,
         }
 
     target_unreachable = bool(info_affine["a_clamped"])
     seed_clip_frac = float((c * a_seed + m_seed < 0.0).mean())
 
-    # REVIEW FIX (CRITICAL-2a): when the seed already satisfies both targets
-    # (J_seed ~ 0), skip the grid/refinement search entirely -- healthy,
-    # no-clip parts then pay only the affine solve's own cost (~ms), not the
-    # grid's. This is the common case in production (most parts do not need
-    # the soft trade-off at all).
+    # Both targets already met: no search needed (the common case).
     if j_seed < 1e-12:
         return pred_seed, _build_info(a_seed, m_seed, r_total_seed, r_std_seed,
                                       seed_clip_frac, target_unreachable)
@@ -527,39 +210,15 @@ def hybrid_field_soft(gnn_seconds: np.ndarray, total_seconds: float, std_seconds
     lo, hi = float(a_clamp[0]), float(a_clamp[1])
     n_a, n_m = 64, 33
     a_grid = np.logspace(np.log10(lo), np.log10(hi), n_a)
-    # REVIEW FIX (CRITICAL-1): the m-window is CENTERED ON THE SEED
-    # (`m_seed + total*[-0.3, +0.3]`), not on the absolute `total*[0.7, 1.3]`
-    # -- see the docstring's "REVIEW FIX (2026-07-31, round 1, CRITICAL-1)"
-    # note for why the absolute window never contained `m_seed` on real
-    # fields and so never let the solver engage.
+    # The m-window is centred on the seed.
     m_lo, m_hi = sorted((m_seed - 0.3 * total, m_seed + 0.3 * total))
     m_grid = np.linspace(m_lo, m_hi, n_m)
 
-    # REVIEW FIX (CRITICAL-2b): the grid's hot loop no longer calls
-    # `np.quantile`/`np.partition` PER GRID POINT for the top-quantile term.
-    # `raw(a, m) = a*c + m` is a POSITIVE-scale-then-shift map, which is
-    # order-preserving; clipping at 0 only ever affects the LOW tail. The
-    # q-quantile of `clip(raw, 0, None)` therefore equals the q-quantile of
-    # the UNCLIPPED `raw` -- which is just `a*shift + m` (`shift =
-    # quantile(c, q)`, computed once above) -- exactly, as long as the number
-    # of clipped nodes is <= floor((n-1)*q), i.e. clip_frac up to ~q = 0.995
-    # (review-verified exact at clip_frac 0.00 / 0.43 / 0.83 on real fields;
-    # rel err <= 1.7e-16). The identity only breaks at clip_frac == 1.0, and
-    # there it is ONE-SIDED SAFE: the closed form returns the negative raw
-    # quantile vs the true 0, so r_total < -1 and J is OVER-estimated --
-    # such a point can never be falsely selected for any total > 0. This
-    # turns an O(n_nodes log n_nodes)-per-point cost into an O(1)-per-point
-    # closed form; the winning candidate is re-scored through the exact
-    # `_residuals()` before acceptance regardless. `np.std` still needs the actual clipped per-node values (clip
-    # changes variance nontrivially), so it remains the only per-node
-    # computation in the search; the search's own (a, m) argmin is therefore
-    # APPROXIMATE only in the sense that its objective is now built from an
-    # exact-closed-form quantile + an exact std -- i.e. not approximate at
-    # all for the quantile term, and exact for std. Regardless, the FINAL
-    # accept/reject decision below re-evaluates the winning candidate's
-    # residuals via the exact `_residuals()` helper (a single `np.quantile`
-    # call, not a grid-wide one) before comparing to `j_seed`, so no
-    # approximation ever reaches the returned `info`/decision.
+    # raw(a, m) = a*c + m is order-preserving and clipping only touches the
+    # low tail, so quantile(clip(raw), q) == a*shift + m in closed form as
+    # long as fewer than ~q of the nodes are clipped. Where that fails (all
+    # nodes clipped) J is over-estimated, so such a point is never selected.
+    # The winning candidate is re-scored exactly below.
     def _grid_eval(a_vals, m_vals):
         n_av, n_mv = len(a_vals), len(m_vals)
         J = np.empty((n_av, n_mv))
@@ -577,12 +236,6 @@ def hybrid_field_soft(gnn_seconds: np.ndarray, total_seconds: float, std_seconds
             Rs[i] = rs
         return J, Rt, Rs
 
-    # Coarse 64x33 (a, m) grid, PLUS the single explicit seed point (NOT a
-    # cross-product extension of either axis -- crossing m_seed against every
-    # `a` would silently widen the m-search far outside the intended window).
-    # The seed's own (J_seed, r_total_seed, r_std_seed) were already computed
-    # exactly above, so it is folded into the "best so far" comparison
-    # directly rather than re-evaluated here.
     J_coarse, Rt_coarse, Rs_coarse = _grid_eval(a_grid, m_grid)
     idx = np.unravel_index(int(np.argmin(J_coarse)), J_coarse.shape)
     best_a, best_m = float(a_grid[idx[0]]), float(m_grid[idx[1]])
@@ -593,9 +246,8 @@ def hybrid_field_soft(gnn_seconds: np.ndarray, total_seconds: float, std_seconds
         best_a, best_m, best_J = a_seed, m_seed, j_seed
         best_rt, best_rs = r_total_seed, r_std_seed
 
-    # One local refinement pass: a finer 9x9 sub-grid windowed to roughly one
-    # coarse grid cell around whichever point (grid or seed) is best so far
-    # (log-spaced in a, linear in m), bounded within the same centered window.
+    # One refinement pass on a 9 x 9 sub-grid about one coarse cell wide
+    # around the best point so far (log-spaced in a, linear in m).
     log_lo, log_hi = np.log10(lo), np.log10(hi)
     a_step = (log_hi - log_lo) / max(n_a - 1, 1)
     log_best_a = np.log10(best_a)
@@ -613,11 +265,7 @@ def hybrid_field_soft(gnn_seconds: np.ndarray, total_seconds: float, std_seconds
         best_J = float(J_ref[idx_r])
         best_rt, best_rs = float(Rt_ref[idx_r]), float(Rs_ref[idx_r])
 
-    # Re-evaluate the winning candidate's residuals EXACTLY (a single
-    # `np.quantile` + `np.std` call, not a grid-wide pass) before the
-    # accept/reject decision -- the search above is already exact for
-    # everything but relies on the closed-form quantile identity holding, so
-    # this is a cheap final confirmation, not a correction of a known bias.
+    # Exact re-check of the winner before accepting it over the seed.
     if (best_a, best_m) != (a_seed, m_seed):
         best_pred = np.clip(best_a * c + best_m, 0.0, None)
         best_rt, best_rs = _residuals(best_pred)

src/inference_props.py

@@ -1,112 +1,125 @@
-"""Conditioning resolver for v4 multi-alloy inference.
+"""Alloy and mold properties.
 
-Turns a named alloy grade + mold system (or a raw-props JSON override) into
-the 8-key property dict consumed by `src.features.alloy_g_vector` and the
-FT2b-conditioned model/GBM:
+Turns a named alloy and mold (rows of `data/properties/materials.csv` and
+`molds.csv`), or a JSON file with raw values, into the 8-key dict the model
+reads:
 
     {k, rho, cp, L_kJ_kg, TL, TS, Tpour, mold_effusivity}
 
-Named alloys come from `tools.doe.load_properties.load_materials()` (reads
-`data/properties/materials.csv`); mold effusivities come from
-`data/properties/molds.csv`. A `props_json` path bypasses both tables
-entirely and must already contain all 8 keys.
+Units: W/(m K), kg/m^3, J/(kg K), kJ/kg, deg C, W s^0.5/(m^2 K). The pouring
+temperature of a named alloy is the midpoint of its tabulated range.
 """
 from __future__ import annotations
 
 import csv
 import json
+from dataclasses import dataclass
 from pathlib import Path
 
-from tools.doe.load_properties import load_materials
-
 DEFAULT_ALLOY = "EN-GJS-400-15"
+DEFAULT_MOLD = "green_sand"
 
 _REQUIRED_KEYS = ("k", "rho", "cp", "L_kJ_kg", "TL", "TS", "Tpour", "mold_effusivity")
 
 _MATERIALS_PATH = "data/properties/materials.csv"
 _MOLDS_PATH = "data/properties/molds.csv"
 
-_materials_cache: dict[str, object] | None = None
+# Family used for shrinkage severity when the alloy family is unknown (a raw
+# property file has no grade). Steel has the largest shrinkage, so this errs
+# on the side of flagging.
+_UNKNOWN_ALLOY_FAMILY_FALLBACK = "steel"
+
+
+@dataclass(frozen=True)
+class Material:
+    material_id: str
+    alloy_family: str
+    k: float
+    rho: float
+    cp: float
+    L: float
+    TL: float
+    TS: float
+    Tpour: float
+
+
+def load_materials(path: str = _MATERIALS_PATH) -> list[Material]:
+    """Rows of the materials table; rejects rows that are not TS < TL < Tpour."""
+    materials: list[Material] = []
+    with Path(path).open(newline="") as handle:
+        for row in csv.DictReader(handle):
+            material_id = row["material_id"]
+            TL = float(row["TL_C"])
+            TS = float(row["TS_C"])
+            tpour = 0.5 * (float(row["Tpour_lo_C"]) + float(row["Tpour_hi_C"]))
+            if not TS < TL < tpour:
+                raise ValueError(
+                    f"material {material_id!r}: expected TS < TL < Tpour, "
+                    f"got {TS}, {TL}, {tpour}")
+            materials.append(Material(
+                material_id=material_id,
+                alloy_family=row["family"],
+                k=float(row["k_W_mK"]),
+                rho=float(row["rho_kg_m3"]),
+                cp=float(row["cp_J_kgK"]),
+                L=float(row["L_kJ_kg"]),
+                TL=TL,
+                TS=TS,
+                Tpour=tpour,
+            ))
+    return materials
+
+
+def load_molds(path: str = _MOLDS_PATH) -> dict[str, float]:
+    """{mold_system: effusivity} from the molds table."""
+    with Path(path).open(newline="") as f:
+        return {row["mold_system"]: float(row["effusivity_W_s05_m2K"])
+                for row in csv.DictReader(f)}
+
+
+_materials_cache: dict[str, Material] | None = None
 _molds_cache: dict[str, float] | None = None
 
 
-def _materials() -> dict[str, object]:
-    """Module-level cached {material_id: Material}."""
+def _materials() -> dict[str, Material]:
     global _materials_cache
     if _materials_cache is None:
-        _materials_cache = {m.material_id: m for m in load_materials(_MATERIALS_PATH)}
+        _materials_cache = {m.material_id: m for m in load_materials()}
     return _materials_cache
 
 
 def _molds() -> dict[str, float]:
-    """Module-level cached {mold_system: effusivity_W_s05_m2K}."""
     global _molds_cache
     if _molds_cache is None:
-        molds: dict[str, float] = {}
-        p = Path(_MOLDS_PATH)
-        if p.exists():
-            with p.open(newline="") as f:
-                for row in csv.DictReader(f):
-                    system = row.get("mold_system")
-                    raw = row.get("effusivity_W_s05_m2K")
-                    if system and raw not in (None, ""):
-                        molds[system] = float(raw)
-        _molds_cache = molds
+        _molds_cache = load_molds()
     return _molds_cache
 
 
 def list_alloys() -> list[str]:
-    """Named material_ids available (from `data/properties/materials.csv`)."""
     return list(_materials().keys())
 
 
 def list_molds() -> list[str]:
-    """Named mold systems available (from `data/properties/molds.csv`)."""
     return list(_molds().keys())
 
 
 def _validate_tl_ts(props: dict, source: str) -> None:
-    """Review finding I2 (final whole-branch review, 2026-08-01): `TL <= TS`
-    is physically inverted (liquidus at or below solidus) and was previously
-    UNvalidated for a `--props` custom-JSON override -- `resolve_props` only
-    checked key *presence*, never ordering. `src.shrinkage.niyama_proxy`
-    clamps its `tdot_proxy` term to `>= 0.0`, so an inverted pair silently
-    yields `Ny == 0` for every node, and `porosity_risk(mode="percentile")`
-    then flags 100% of nodes as at-risk -- no warning, no exception, just a
-    self-contradictory "100% flagged (bottom decile by construction)"
-    summary line. Named grades are protected by `tools.doe.load_properties.
-    load_materials`'s own `TS < TL < Tpour` check at CSV-load time, but that
-    guard never runs for a `--props` override (it bypasses the materials
-    table entirely) -- and SP2 makes shrinkage screening default-ON, so this
-    path is live in production, not dead code. Raises loudly, naming both
-    values and which of `props_json`/`alloy` triggered it.
-    """
+    """An inverted or equal liquidus/solidus pair would make the Niyama proxy
+    zero everywhere and flag every node, so it is rejected up front."""
     TL, TS = props["TL"], props["TS"]
     if TL <= TS:
         raise ValueError(
             f"resolve_props ({source}): TL ({TL!r}) must be strictly greater than "
-            f"TS ({TS!r}) -- an inverted or degenerate liquidus/solidus pair makes "
-            f"the Niyama proxy (src.shrinkage.niyama_proxy) collapse to 0 "
-            f"everywhere, silently flagging 100% of nodes as at-risk")
-
-
-def resolve_props(alloy: str | None = None, mold: str = "green_sand",
-                   props_json: str | None = None) -> dict:
-    """Resolve the 8-key conditioning-property dict.
+            f"TS ({TS!r})")
 
-    `props_json`, if given, overrides `alloy`/`mold` entirely: the file must
-    be a JSON object containing all 8 required keys (missing keys raise
-    ValueError naming them).
 
-    Otherwise resolves `alloy` (default `DEFAULT_ALLOY` if None) against the
-    named-material table and `mold` against the mold-effusivity table;
-    unknown names raise ValueError listing the available options.
+def resolve_props(alloy: str | None = None, mold: str = DEFAULT_MOLD,
+                  props_json: str | None = None) -> dict:
+    """The 8-key property dict.
 
-    Both branches validate `TL > TS` (review finding I2) before returning --
-    a `--props` override with an inverted or degenerate pair raises a loud
-    `ValueError` naming both values, rather than silently degrading the
-    shrinkage screening's Niyama proxy to "0 everywhere / 100% flagged".
-    """
+    `props_json` replaces `alloy` and `mold`; the file must hold all 8 keys.
+    Otherwise `alloy` (default `DEFAULT_ALLOY`) and `mold` are looked up in
+    the tables, and unknown names raise ValueError listing the options."""
     if props_json is not None:
         with open(props_json) as f:
             data = json.load(f)
@@ -130,7 +143,6 @@ def resolve_props(alloy: str | None = None, mold: str = "green_sand",
     if mold not in molds:
         raise ValueError(
             f"unknown mold system '{mold}'; available mold systems: {sorted(molds.keys())}")
-    mold_effusivity = molds[mold]
 
     props = {
         "k": float(material.k),
@@ -140,63 +152,19 @@ def resolve_props(alloy: str | None = None, mold: str = "green_sand",
         "TL": float(material.TL),
         "TS": float(material.TS),
         "Tpour": float(material.Tpour),
-        "mold_effusivity": float(mold_effusivity),
+        "mold_effusivity": float(molds[mold]),
     }
     _validate_tl_ts(props, f"alloy={alloy_id!r}")
     return props
 
 
-# --------------------------------------------------------------------------- #
-# resolve_alloy_family (SP2 shrinkage Task 4, controller directive 2):
-# `resolve_props`'s returned 8-key dict has no alloy-family field -- it is
-# BUILT from `data/properties/materials.csv`'s own `family` column
-# (`tools.doe.load_properties.load_materials` stores it on `Material.
-# alloy_family`), but that value never makes it into `resolve_props`'s
-# return dict, and `src.shrinkage.screen_casting` needs a family key from
-# `src.shrinkage.SHRINKAGE_PCT` to look up the alloy's volumetric shrinkage
-# percentage. Rather than reaching back into the materials table (which
-# doesn't exist at all for a `--props` custom-JSON override), this maps the
-# alloy GRADE NAME directly by prefix -- every named grade in
-# `materials.csv` today falls into one of `SHRINKAGE_PCT`'s 4 keys via this
-# same prefix convention (EN-GJL-* -> grey_iron, EN-GJS-* -> ductile_iron,
-# GS-*/G17*/G20* -> low_alloy_steel), so this is not a guess for any grade
-# actually in the property DB -- only for a name this helper has never seen
-# (a future grade, a typo, or a `--props` override with no named grade at
-# all) does it fall back, loudly (the caller records `alloy_family_assumed`
-# rather than silently proceeding as if the mapping were exact).
-# --------------------------------------------------------------------------- #
-_ALLOY_FAMILY_PREFIXES = (
-    ("EN-GJS", "ductile_iron"),
-    ("EN-GJL", "grey_iron"),
-    ("GS-", "low_alloy_steel"),
-    ("G17", "low_alloy_steel"),
-    ("G20", "low_alloy_steel"),
-)
-
-# The conservative (highest-shrinkage-percentage) fallback family for a grade
-# name this helper doesn't recognize, or for a `--props` custom-JSON override
-# (no grade name to map from at all) -- see the module comment above.
-_UNKNOWN_ALLOY_FAMILY_FALLBACK = "steel"
-
-
 def resolve_alloy_family(alloy: str | None, props_json: str | None = None) -> tuple[str, bool]:
-    """Map an alloy grade name to one of `src.shrinkage.SHRINKAGE_PCT`'s
-    family keys, for the shrinkage-screening integration.
-
-    Returns `(family, assumed)`: `assumed` is `True` when the mapping is a
-    documented fallback rather than an exact name-prefix match -- either
-    because `props_json` was given (a raw-properties override has no named
-    grade to map from) or because `alloy` (or `DEFAULT_ALLOY`, if `alloy`
-    is `None`) doesn't start with any of the known prefixes. The caller is
-    expected to surface `assumed` (e.g. as `alloy_family_assumed` in the
-    shrinkage result) rather than silently treat a defaulted family as if it
-    were an exact match.
-    """
+    """(family, assumed) for shrinkage severity. The family comes from the
+    materials table; `assumed` is True when it had to fall back to steel
+    (a raw property file, or an alloy that is not in the table)."""
     if props_json is not None:
         return _UNKNOWN_ALLOY_FAMILY_FALLBACK, True
-
-    label = alloy if alloy is not None else DEFAULT_ALLOY
-    for prefix, family in _ALLOY_FAMILY_PREFIXES:
-        if label.startswith(prefix):
-            return family, False
-    return _UNKNOWN_ALLOY_FAMILY_FALLBACK, True
+    material = _materials().get(alloy if alloy is not None else DEFAULT_ALLOY)
+    if material is None:
+        return _UNKNOWN_ALLOY_FAMILY_FALLBACK, True
+    return material.alloy_family, False

src/mesh_cache.py

@@ -1,57 +1,11 @@
-"""Content-keyed mesh cache for the STEP -> tet-mesh pipeline.
-
-Inference meshing (`tools.doe.prep_parts.heal_and_mesh`) takes minutes and is
-routinely repeated for byte-identical STEP inputs (re-runs, batch backends,
-concurrent workers). This module memoizes that call on disk, keyed on the
-SHA-256 of the STEP file's bytes + `target_nodes` + `MESH_PIPELINE_VERSION`,
-so any change to the meshing pipeline invalidates the cache automatically
-(bump `MESH_PIPELINE_VERSION` in `tools/doe/prep_parts.py` whenever the mesh
-output could change).
-
-`heal_and_mesh` and `MESH_PIPELINE_VERSION` are imported at module level and
-referenced dynamically (via `globals()`/module attribute lookup) so tests can
-monkeypatch `mesh_cache.heal_and_mesh` / `mesh_cache.MESH_PIPELINE_VERSION`
-without needing a real gmsh install.
-
-Cache layout (under `cache_dir`, default `runs/mesh_cache/`):
-    <key>.msh   -- the cached gmsh 2.2 mesh (only written when ok=True)
-    <key>.json  -- meta dict (ok, n_nodes, ..., mesh_wall_s, pipeline_version,
-                   source_path, created_at)
-
-Writes are atomic and per-writer-unique (`<key>.msh.<pid>.<uuid>.tmp` /
-`<key>.json.<pid>.<uuid>.tmp`, then `os.replace` onto the shared final path)
-so concurrent workers racing on the same cache miss never share a tmp inode
-(which would let their writes interleave and corrupt it) and never observe a
-partially written file -- whichever replace lands last simply wins, and both
-writers produced an equally valid, uncorrupted mesh for the same key.
-
-A failed `heal_and_mesh` call (`ok=False`) is never cached: a fixed STEP file
-(or a differently-tuned pipeline, which would change the key anyway) must be
-retried, not permanently remembered as unmeshable.
-
-Defeaturing auto-trigger (mesh-accel Task 5, 2026-07-29): `get_or_mesh`
-decides whether to pass `defeature=True` to `heal_and_mesh` via a
-deterministic, input-only policy, symmetric in both directions -- a STEP
-file whose on-disk byte size exceeds `_BIG_FILE_THRESHOLD_BYTES` (8 MiB)
-goes straight to a `defeature=True` attempt first (large real parts are,
-empirically, the ones that need it -- see
-`tools/doe/prep_parts.py::_apply_defeature`), and
-only if that attempt fails outright (`ok=False`) retries ONCE more with
-`defeature=False`; a smaller file does the mirror image -- plain
-(`defeature=False`) first, retrying ONCE with `defeature=True` only on
-outright failure. Either way there is at most one retry, ever. Because the
-policy is a pure function of the input (file size, then a same-input
-retry-on-failure), a given STEP file always resolves to the same
-`defeature` choice across runs/workers -- so `defeature` does NOT need to
-be (and is not) part of `mesh_cache_key`; the cache key stays exactly as
-before this task. The returned meta dict always carries `defeature` (bool,
-whichever value actually produced the returned result) and `retry` (bool,
-True only when a second attempt actually ran, i.e. the first attempt --
-big-file `defeature=True` or small-file `defeature=False` -- failed
-outright; False whenever the first attempt already succeeded or the input
-size never triggers the alternate branch).
-"""
+"""Content-addressed on-disk cache for `src.meshing.heal_and_mesh`.
 
+The key is sha256(STEP bytes) + target node count + `MESH_PIPELINE_VERSION`,
+so a byte-identical STEP file is meshed once. Each entry is a pair
+`<key>.msh` / `<key>.json` (the mesh and its meta dict), written through a
+per-writer temporary file and `os.replace`, so concurrent writers never leave
+a half-written entry. Failed meshes are not cached.
+"""
 from __future__ import annotations
 
 import hashlib
@@ -61,15 +15,13 @@ import shutil
 import time
 import uuid
 
-from tools.doe.prep_parts import DEFAULT_TARGET_NODES, MESH_PIPELINE_VERSION, heal_and_mesh
+# MESH_PIPELINE_VERSION and heal_and_mesh are read through globals() so tests
+# can monkeypatch them.
+from src.meshing import DEFAULT_TARGET_NODES, MESH_PIPELINE_VERSION, heal_and_mesh  # noqa: F401
 
 _HASH_CHUNK_SIZE = 1 << 20  # 1 MiB
-# Participates in mesh-output selection (which `defeature` attempt is tried
-# first -- see the "Defeaturing auto-trigger" section above). It is NOT part
-# of `mesh_cache_key`, so changing this value changes what mesh a given STEP
-# file resolves to WITHOUT changing its cache key -- any change here requires
-# bumping `MESH_PIPELINE_VERSION` (tools/doe/prep_parts.py) to invalidate
-# stale cache entries. Cross-file contract; keep in sync with prep_parts.py.
+# STEP files above this size get the defeaturing attempt first. Changing it
+# changes which mesh a file resolves to, so bump MESH_PIPELINE_VERSION too.
 _BIG_FILE_THRESHOLD_BYTES = 8 * 1024 * 1024  # 8 MiB
 
 
@@ -84,35 +36,16 @@ def _sha256_file(path: str) -> str:
     return h.hexdigest()
 
 
-def mesh_cache_key(step_path: str, target_nodes: int = DEFAULT_TARGET_NODES,
-                   bodies_key: str = "none") -> str:
-    """Content-addressed cache key: sha256(file bytes) + target_nodes +
-    MESH_PIPELINE_VERSION + bodies_key, so the key changes if any of those
-    change. Reads `MESH_PIPELINE_VERSION` dynamically from module globals so
-    tests can monkeypatch it.
-
-    `bodies_key` (Part B Task 4): a caller-supplied string folding in the
-    Part-B multi-body role assignment (see
-    `src.predict_casting._bodies_cache_key`) -- default `"none"` for the
-    ordinary single-body inference path, so an ordinary call's key is
-    unaffected by this parameter's mere existence (a bodies-aware and a
-    bodies-less call for the SAME step_path/target_nodes only collide if
-    the bodies-aware caller explicitly passes `bodies_key="none"`, which it
-    never does -- `_bodies_cache_key` always produces a real hash).
-    """
+def mesh_cache_key(step_path: str, target_nodes: int = DEFAULT_TARGET_NODES) -> str:
+    """sha256 over the file digest, `target_nodes` and the pipeline version.
+    The version is read from module globals so tests can monkeypatch it."""
     digest = _sha256_file(step_path)
     pipeline_version = globals()["MESH_PIPELINE_VERSION"]
-    key_input = f"{digest}:{int(target_nodes)}:{int(pipeline_version)}:{bodies_key}".encode("utf-8")
+    key_input = f"{digest}:{int(target_nodes)}:{int(pipeline_version)}".encode("utf-8")
     return hashlib.sha256(key_input).hexdigest()
 
 
 def _unique_tmp_path(final_path: str) -> str:
-    """A per-writer-unique tmp sibling of `final_path`. Two workers racing on
-    the same cache miss must never write the same tmp inode (that would let
-    their writes interleave and corrupt it before either `os.replace` runs) --
-    each caller gets its own pid+uuid-suffixed tmp file, then atomically
-    replaces the SAME shared final path (last replace wins, and both writers
-    produced a fully-formed, uncorrupted file)."""
     return f"{final_path}.{os.getpid()}.{uuid.uuid4().hex}.tmp"
 
 
@@ -124,45 +57,17 @@ def _paths(cache_dir: str, key: str) -> tuple[str, str]:
 
 
 def get_or_mesh(step_path: str, out_msh: str, target_nodes: int = DEFAULT_TARGET_NODES,
-                cache_dir: str = "runs/mesh_cache", keep_solid_id: int | None = None,
-                bodies_key: str = "none", post_mesh=None) -> dict:
-    """Return a `heal_and_mesh`-shaped meta dict for `step_path`, using the
-    on-disk cache when possible.
-
-    - Hit: copy the cached `.msh` to `out_msh`, return the stored meta dict
-      plus `{"cache": "hit"}` (no `heal_and_mesh` call, no `post_mesh`
-      call -- the cached mesh is already whatever `post_mesh` produced on
-      the miss that created it, see below).
-    - Miss + `ok=True`: call `heal_and_mesh` (forwarding `keep_solid_id`),
-      then -- if `post_mesh` is given -- call `post_mesh(out_msh) -> dict`
-      and merge its return into the meta, THEN atomically store both the
-      (now possibly `post_mesh`-modified) mesh and its meta json under the
-      cache key, and return the meta dict plus `{"cache": "miss"}`.
-    - Miss + `ok=False`: store nothing, return `heal_and_mesh`'s dict
-      unchanged plus `{"cache": "miss"}` -- failures are never cached, and
-      `post_mesh` is never called on a failed mesh.
-
-    `keep_solid_id` (Part B Task 4): the user-designated casting solid, when
-    a multi-body STEP's role assignment has already picked which solid IS
-    the casting -- forwarded verbatim to `heal_and_mesh`, in place of the
-    largest-solid default (`None`, unchanged).
-
-    `bodies_key` (Part B Task 4): folded into `mesh_cache_key` (see there)
-    so a role-assignment change (a different `bodies.json`, a registry edit
-    that changes a resolved BC scalar, etc.) is a genuinely different cache
-    entry, never a stale hit.
-
-    `post_mesh` (Part B Task 4): an optional `(out_msh) -> dict` hook run
-    on a cache MISS, after `heal_and_mesh` succeeds and BEFORE the mesh is
-    stored -- e.g. `src.assembly.tag_contact_patches`, which rewrites
-    `out_msh` in place (Cooling -> Adiabatic contact patch) and returns
-    `{"per_body_patch_area_mm2": ..., "n_adiab_tris": ...}`. Tagging cost is
-    not trivial on big meshes, so the TAGGED mesh is what gets cached (the
-    bodies hash is already in the key) -- a later hit for the same
-    `bodies_key` returns the tagged mesh directly, with no re-tagging.
+                cache_dir: str = "runs/mesh_cache") -> dict:
+    """`heal_and_mesh`'s meta dict for `step_path`, plus a `"cache"` key.
+
+    On a hit the cached mesh is copied to `out_msh` (`"cache": "hit"`). On a
+    miss the part is meshed; a successful mesh is stored and returned with
+    `"cache": "miss"`, a failed one is returned unchanged and not stored.
+    Small files try a plain pass first and retry once with defeaturing; files
+    above 8 MiB try the defeatured pass first and retry once without it.
     """
     os.makedirs(cache_dir, exist_ok=True)
-    key = mesh_cache_key(step_path, target_nodes, bodies_key)
+    key = mesh_cache_key(step_path, target_nodes)
     msh_path, json_path = _paths(cache_dir, key)
 
     if os.path.exists(msh_path) and os.path.exists(json_path):
@@ -177,26 +82,16 @@ def get_or_mesh(step_path: str, out_msh: str, target_nodes: int = DEFAULT_TARGET
     t0 = time.time()
     big_file = os.path.getsize(step_path) > _BIG_FILE_THRESHOLD_BYTES
     if big_file:
-        # Large STEP files go straight to the defeatured attempt -- no
-        # point burning a near-certain-to-fail plain pass first. But if
-        # even that fails outright, retry ONCE with defeature=False --
-        # symmetric with the small-file policy below, so a big file is
-        # never left with only one attempt while a small file gets two.
-        meta = heal_and_mesh_fn(step_path, out_msh, target_nodes=target_nodes, defeature=True,
-                                keep_solid_id=keep_solid_id)
+        meta = heal_and_mesh_fn(step_path, out_msh, target_nodes=target_nodes, defeature=True)
         defeature_used, retry = True, False
         if not meta.get("ok"):
-            meta = heal_and_mesh_fn(step_path, out_msh, target_nodes=target_nodes, defeature=False,
-                                    keep_solid_id=keep_solid_id)
+            meta = heal_and_mesh_fn(step_path, out_msh, target_nodes=target_nodes, defeature=False)
             defeature_used, retry = False, True
     else:
-        meta = heal_and_mesh_fn(step_path, out_msh, target_nodes=target_nodes, defeature=False,
-                                keep_solid_id=keep_solid_id)
+        meta = heal_and_mesh_fn(step_path, out_msh, target_nodes=target_nodes, defeature=False)
         defeature_used, retry = False, False
         if not meta.get("ok"):
-            # Exactly one retry, ever -- a defeature=True failure is final.
-            meta = heal_and_mesh_fn(step_path, out_msh, target_nodes=target_nodes, defeature=True,
-                                    keep_solid_id=keep_solid_id)
+            meta = heal_and_mesh_fn(step_path, out_msh, target_nodes=target_nodes, defeature=True)
             defeature_used, retry = True, True
     wall_s = time.time() - t0
 
@@ -207,10 +102,7 @@ def get_or_mesh(step_path: str, out_msh: str, target_nodes: int = DEFAULT_TARGET
         result["retry"] = retry
         return result
 
-    post_meta = post_mesh(out_msh) if post_mesh is not None else {}
-
     stored_meta = dict(meta)
-    stored_meta.update(post_meta)
     stored_meta["mesh_wall_s"] = wall_s
     stored_meta["pipeline_version"] = globals()["MESH_PIPELINE_VERSION"]
     stored_meta["source_path"] = os.path.abspath(step_path)
@@ -220,7 +112,7 @@ def get_or_mesh(step_path: str, out_msh: str, target_nodes: int = DEFAULT_TARGET
 
     msh_tmp = _unique_tmp_path(msh_path)
     json_tmp = _unique_tmp_path(json_path)
-    shutil.copyfile(out_msh, msh_tmp)  # copies the (possibly post_mesh-tagged) mesh
+    shutil.copyfile(out_msh, msh_tmp)
     os.replace(msh_tmp, msh_path)
     with open(json_tmp, "w") as f:
         json.dump(stored_meta, f)

src/meshing.py

@@ -0,0 +1,368 @@
+"""STEP to tetrahedral mesh with Gmsh (OpenCASCADE kernel).
+
+`heal_and_mesh` is the default path: import the STEP file, keep the largest
+solid, tag every surface "Cooling" and the volume "Volume", and search the
+element size until the mesh lands near `target_nodes` (the density the model
+was trained on). If the unhealed geometry cannot be meshed, a second pass heals
+it and retries. `mesh_step` meshes at one fixed element size instead.
+
+Both write gmsh 2.2 ASCII meshes, the format `src.msh_parse` reads. All
+lengths are in millimetres.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+DEFAULT_TARGET_NODES = 8500
+# Part of the mesh-cache key (src/mesh_cache.py): bump it whenever a change
+# here would produce a different mesh for the same STEP file.
+MESH_PIPELINE_VERSION = 3
+
+# Element-size seed: sqrt(K_SURF * surface_area / target_nodes), calibrated so
+# the first guess usually lands within the density search's tolerance.
+K_SURF = 2.5
+
+# Physical-group tags are pinned (not auto-assigned) so that repeated meshing
+# attempts in one Gmsh session keep the "Volume"/"Cooling" names bound.
+_VOLUME_PHYS_TAG = 1
+_COOLING_PHYS_TAG = 2
+
+
+def _reduce_to_largest_solid():
+    """Keep only the largest solid of the imported OCC model.
+
+    Returns (n_solids_in, kept_tag, kept_volume_mm3)."""
+    import gmsh
+
+    vols = [t for (d, t) in gmsh.model.occ.getEntities(3)]
+    n_in = len(vols)
+    if n_in == 0:
+        raise ValueError("no 3D solids found in STEP import")
+    masses = [(t, gmsh.model.occ.getMass(3, t)) for t in vols]
+    kept_tag, kept_vol = max(masses, key=lambda tv: tv[1])
+    if n_in > 1:
+        drop = [(3, t) for (t, _v) in masses if t != kept_tag]
+        gmsh.model.occ.remove(drop, recursive=True)
+        gmsh.model.occ.synchronize()
+    return n_in, kept_tag, float(kept_vol)
+
+
+def _import_and_reduce(step_path: str) -> tuple[int, float, str | None]:
+    """Import `step_path` into the current Gmsh session and keep its largest
+    solid. Returns (n_solids_in, kept_volume_mm3, error_or_None)."""
+    import gmsh
+
+    try:
+        gmsh.model.occ.importShapes(step_path)
+        gmsh.model.occ.synchronize()
+    except Exception as e:
+        return 0, 0.0, f"import failed: {e}"
+    try:
+        n_solids, _kept_tag, volume_mm3 = _reduce_to_largest_solid()
+    except Exception as e:
+        return 0, 0.0, f"largest-solid reduction failed: {e}"
+    return n_solids, volume_mm3, None
+
+
+def _elem_size(bbox_diag_mm: float) -> float:
+    """Fallback element-size seed: bounding-box diagonal / 28."""
+    return bbox_diag_mm / 28.0
+
+
+def _elem_size_from_area(area_mm2: float, target_nodes: int) -> float:
+    """Element-size seed from the total surface area of the model."""
+    return float(np.sqrt(K_SURF * area_mm2 / target_nodes))
+
+
+def _total_surface_area_mm2() -> float:
+    """Total area of all surface entities of the current OCC model."""
+    import gmsh
+
+    surfs = [t for (d, t) in gmsh.model.occ.getEntities(2)]
+    return float(sum(gmsh.model.occ.getMass(2, t) for t in surfs))
+
+
+def _first_guess_elem_size(bbox_diag_mm: float, target_nodes: int) -> float:
+    """Surface-area seed, or the bounding-box seed if the area query fails."""
+    try:
+        area_mm2 = _total_surface_area_mm2()
+    except Exception as e:
+        print(f"[meshing] WARNING: OCC surface-area query failed ({e}); "
+              f"falling back to bbox_diag/28 seed", flush=True)
+        return _elem_size(bbox_diag_mm)
+    if area_mm2 <= 0:
+        print(f"[meshing] WARNING: OCC surface-area query returned "
+              f"{area_mm2!r} (<=0); falling back to bbox_diag/28 seed", flush=True)
+        return _elem_size(bbox_diag_mm)
+    return _elem_size_from_area(area_mm2, target_nodes)
+
+
+def _search_elem_size(build_fn, out_msh: str, target_nodes: int,
+                      max_iter: int = 3) -> tuple[int, float]:
+    """Refine the element size assuming the node count scales as h^-3."""
+    n_nodes, elem_size = build_fn(out_msh, None)
+    for _ in range(max_iter - 1):
+        if n_nodes <= 0:
+            break
+        elem_size = elem_size * (n_nodes / target_nodes) ** (1.0 / 3.0)
+        n_nodes, elem_size = build_fn(out_msh, elem_size)
+    return n_nodes, elem_size
+
+
+def _search_target_nodes(build_fn, out_msh: str, target_nodes: int,
+                         tol: float = 0.35, max_iter: int = 3) -> tuple[int, float, int]:
+    """Density search towards `target_nodes`.
+
+    `build_fn(out_msh, elem_size) -> (n_nodes, elem_size_used)`, with
+    `elem_size=None` meaning "use the seed" and `n_nodes == 0` meaning the
+    mesh failed outright. Stops after the first attempt if it is already
+    within `tol` of the target (or failed); otherwise refines for up to
+    `max_iter` attempts. Returns (n_nodes, elem_size_used, attempts)."""
+    calls = {"n": 0}
+
+    def counted_build_fn(msh_path, elem_size):
+        calls["n"] += 1
+        return build_fn(msh_path, elem_size)
+
+    n0, es0 = counted_build_fn(out_msh, None)
+    if n0 <= 0 or abs(n0 - target_nodes) <= tol * target_nodes:
+        return n0, es0, calls["n"]
+    n, es = _search_elem_size(counted_build_fn, out_msh, target_nodes, max_iter=max_iter)
+    return n, es, calls["n"]
+
+
+def _tag_and_mesh(size: float, out_msh: str) -> tuple[bool, int, str]:
+    """Tag the current model ("Volume" + every surface as "Cooling") and mesh
+    it at element size `size`. Returns (ok, n_nodes, note).
+
+    The tag numbers are pinned: Gmsh keeps a removed group's *name* bound to
+    its old tag, so auto-assigned tags on a second attempt in the same
+    session would leave the written mesh with unnamed groups."""
+    import gmsh
+
+    gmsh.model.mesh.clear()
+    gmsh.model.removePhysicalGroups()
+    vols = [t for (d, t) in gmsh.model.getEntities(3)]
+    surfs = [t for (d, t) in gmsh.model.getEntities(2)]
+    if not vols:
+        return False, 0, "no volume entities to mesh"
+    pv = gmsh.model.addPhysicalGroup(3, vols, tag=_VOLUME_PHYS_TAG)
+    gmsh.model.setPhysicalName(3, pv, "Volume")
+    pc = gmsh.model.addPhysicalGroup(2, surfs, tag=_COOLING_PHYS_TAG)
+    gmsh.model.setPhysicalName(2, pc, "Cooling")
+    # The mesh reader matches elements by group name, so an unnamed group
+    # would silently look like a mesh without boundary. Fail loudly instead.
+    if (gmsh.model.getPhysicalName(3, pv) != "Volume"
+            or gmsh.model.getPhysicalName(2, pc) != "Cooling"):
+        return False, 0, ("physical-group tagging did not register "
+                          "(Volume/Cooling name binding missing)")
+    gmsh.option.setNumber("Mesh.MeshSizeMin", size)
+    gmsh.option.setNumber("Mesh.MeshSizeMax", size)
+    try:
+        gmsh.model.mesh.generate(3)
+        node_tags, _, _ = gmsh.model.mesh.getNodes()
+        try:
+            tet_tags = gmsh.model.mesh.getElementsByType(4)[0]
+        except Exception:
+            tet_tags = []
+        if len(node_tags) > 0 and len(tet_tags) > 0:
+            gmsh.option.setNumber("Mesh.MshFileVersion", 2.2)
+            gmsh.write(out_msh)
+            return True, len(node_tags), ""
+        return False, 0, f"empty mesh at size={size:.2f}mm"
+    except Exception as e:
+        return False, 0, f"mesh generate failed at size={size:.2f}mm: {e}"
+
+
+def _apply_heal() -> tuple[bool, str]:
+    """Run OCC shape healing. `healShapes()` can finish without raising and
+    still leave no solid, so a heal only counts if a volume survives.
+    Returns (effective, note)."""
+    import gmsh
+
+    try:
+        gmsh.model.occ.healShapes()
+        gmsh.model.occ.synchronize()
+    except Exception as e:
+        return False, f"healShapes raised: {e}"
+    n_vols = len(gmsh.model.getEntities(3))
+    if n_vols == 0:
+        return False, ("healShapes ran but yielded 0 volume entities (solid could not be "
+                       "reconstructed -- likely a self-intersecting B-rep)")
+    return True, ""
+
+
+def _apply_defeature() -> tuple[bool, str]:
+    """Merge literally duplicate OCC entities (`removeAllDuplicates`), a
+    cheap clean-up for CAD exports. Returns (ok, note); not ok if the call
+    raises or the solid is lost."""
+    import gmsh
+
+    n_vols_before = len(gmsh.model.getEntities(3))
+    try:
+        gmsh.model.occ.removeAllDuplicates()
+        gmsh.model.occ.synchronize()
+    except Exception as e:
+        return False, f"defeature (removeAllDuplicates) raised: {e}"
+    n_vols_after = len(gmsh.model.getEntities(3))
+    if n_vols_before > 0 and n_vols_after == 0:
+        return False, "defeature (removeAllDuplicates) yielded 0 volume entities (solid lost)"
+    return True, ""
+
+
+def heal_and_mesh(step_path: str, out_msh: str, target_nodes: int = DEFAULT_TARGET_NODES,
+                  defeature: bool = False) -> dict:
+    """Mesh the largest solid of `step_path` near `target_nodes` nodes.
+
+    Pass 1 meshes the geometry as imported (healing can destroy a solid
+    that would have meshed fine, so it is not applied by default). Only if
+    pass 1 fails outright does pass 2, in a fresh Gmsh session, heal the
+    geometry and retry the same density search. `defeature=True` runs
+    `_apply_defeature` before meshing in both passes.
+
+    Returns {n_solids, volume_mm3, bbox_diag_mm, n_nodes, ok, healed, notes,
+    attempts}; `notes` collects human-readable diagnostics.
+    """
+    import gmsh
+
+    notes: list[str] = []
+
+    # ---- Pass 1: unhealed ----
+    gmsh.initialize()
+    try:
+        gmsh.option.setNumber("General.Terminal", 0)
+        gmsh.model.add("part")
+        n_solids, volume_mm3, err = _import_and_reduce(step_path)
+        if err is not None:
+            notes.append(err)
+            return dict(n_solids=0, volume_mm3=0.0, bbox_diag_mm=0.0, n_nodes=0,
+                        ok=False, healed=False, notes="; ".join(notes), attempts=0)
+        if n_solids > 1:
+            notes.append(f"multi-body: kept largest of {n_solids} solids (vol={volume_mm3:.1f} mm3)")
+
+        if defeature:
+            defeature_ok, defeature_note = _apply_defeature()
+            if defeature_note:
+                notes.append(defeature_note)
+            if not defeature_ok:
+                return dict(n_solids=n_solids, volume_mm3=volume_mm3, bbox_diag_mm=0.0,
+                            n_nodes=0, ok=False, healed=False, notes="; ".join(notes),
+                            attempts=0)
+
+        xmin, ymin, zmin, xmax, ymax, zmax = gmsh.model.getBoundingBox(-1, -1)
+        bbox_diag_mm = float(np.linalg.norm([xmax - xmin, ymax - ymin, zmax - zmin]))
+        base_size = _first_guess_elem_size(bbox_diag_mm, target_nodes)
+
+        fail_notes: list[str] = []
+
+        def build_fn(msh_path, elem_size):
+            size = elem_size if elem_size is not None else base_size
+            ok_, n_nodes_, note_ = _tag_and_mesh(size, msh_path)
+            if not ok_:
+                fail_notes.append(note_)
+            return (n_nodes_ if ok_ else 0), size
+
+        n_nodes, _used_size, attempts = _search_target_nodes(build_fn, out_msh, target_nodes)
+        if n_nodes > 0:
+            return dict(n_solids=n_solids, volume_mm3=volume_mm3, bbox_diag_mm=bbox_diag_mm,
+                        n_nodes=n_nodes, ok=True, healed=False, notes="; ".join(notes),
+                        attempts=attempts)
+        notes.append(f"unhealed mesh attempt failed: {fail_notes[-1] if fail_notes else 'unknown'}")
+    finally:
+        gmsh.finalize()
+
+    # ---- Pass 2: heal + density-searched retry (fresh session) ----
+    healed = False
+    gmsh.initialize()
+    try:
+        gmsh.option.setNumber("General.Terminal", 0)
+        gmsh.model.add("part_healed")
+        n_solids2, volume_mm32, err2 = _import_and_reduce(step_path)
+        if err2 is not None:
+            notes.append(f"pass-2 {err2}")
+            return dict(n_solids=n_solids, volume_mm3=volume_mm3, bbox_diag_mm=bbox_diag_mm,
+                        n_nodes=0, ok=False, healed=False, notes="; ".join(notes),
+                        attempts=attempts)
+
+        if defeature:
+            defeature_ok2, defeature_note2 = _apply_defeature()
+            if defeature_note2:
+                notes.append(defeature_note2)
+            if not defeature_ok2:
+                return dict(n_solids=n_solids, volume_mm3=volume_mm3, bbox_diag_mm=bbox_diag_mm,
+                            n_nodes=0, ok=False, healed=False, notes="; ".join(notes),
+                            attempts=attempts)
+
+        heal_ok, heal_note = _apply_heal()
+        if heal_note:
+            notes.append(heal_note)
+        if not heal_ok:
+            notes.append("meshing failed after heal attempt")
+            return dict(n_solids=n_solids, volume_mm3=volume_mm3, bbox_diag_mm=bbox_diag_mm,
+                        n_nodes=0, ok=False, healed=False, notes="; ".join(notes),
+                        attempts=attempts)
+        healed = True
+
+        fail_notes2: list[str] = []
+
+        def build_fn2(msh_path, elem_size):
+            size = elem_size if elem_size is not None else base_size
+            ok_, n_nodes_, note_ = _tag_and_mesh(size, msh_path)
+            if not ok_:
+                fail_notes2.append(note_)
+            return (n_nodes_ if ok_ else 0), size
+
+        n_nodes2, _used_size2, attempts2 = _search_target_nodes(build_fn2, out_msh, target_nodes)
+        if n_nodes2 > 0:
+            notes.append("meshed after heal")
+            return dict(n_solids=n_solids, volume_mm3=volume_mm3, bbox_diag_mm=bbox_diag_mm,
+                        n_nodes=n_nodes2, ok=True, healed=True, notes="; ".join(notes),
+                        attempts=attempts + attempts2)
+        notes.append(f"post-heal mesh attempt failed: {fail_notes2[-1] if fail_notes2 else 'unknown'}")
+        notes.append("meshing failed after heal attempt")
+        return dict(n_solids=n_solids, volume_mm3=volume_mm3, bbox_diag_mm=bbox_diag_mm,
+                    n_nodes=0, ok=False, healed=healed, notes="; ".join(notes),
+                    attempts=attempts + attempts2)
+    finally:
+        gmsh.finalize()
+
+
+def auto_elem_size(bbox_diag: float) -> float:
+    """Default fixed element size for `mesh_step`: bounding-box diagonal / 60."""
+    return bbox_diag / 60.0
+
+
+def mesh_step(step_path: str, out_msh: str, elem_size: float | None = None) -> dict:
+    """Mesh the whole STEP model at one fixed element size (no healing, no
+    density search). Every surface is tagged "Cooling"."""
+    import gmsh
+
+    gmsh.initialize()
+    try:
+        gmsh.option.setNumber("General.Terminal", 0)
+        gmsh.model.add("casting")
+        gmsh.model.occ.importShapes(step_path)
+        gmsh.model.occ.synchronize()
+
+        xmin, ymin, zmin, xmax, ymax, zmax = gmsh.model.getBoundingBox(-1, -1)
+        bbox = (xmax - xmin, ymax - ymin, zmax - zmin)
+        diag = float(np.linalg.norm(bbox))
+        size = float(elem_size) if elem_size is not None else auto_elem_size(diag)
+        gmsh.option.setNumber("Mesh.MeshSizeMin", size)
+        gmsh.option.setNumber("Mesh.MeshSizeMax", size)
+
+        vols = [t for (d, t) in gmsh.model.getEntities(3)]
+        surfs = [t for (d, t) in gmsh.model.getEntities(2)]
+        pv = gmsh.model.addPhysicalGroup(3, vols); gmsh.model.setPhysicalName(3, pv, "Volume")
+        pc = gmsh.model.addPhysicalGroup(2, surfs); gmsh.model.setPhysicalName(2, pc, "Cooling")
+
+        gmsh.model.mesh.generate(3)
+        gmsh.option.setNumber("Mesh.MshFileVersion", 2.2)
+        gmsh.write(out_msh)
+
+        node_tags, _, _ = gmsh.model.mesh.getNodes()
+        tet_tags = gmsh.model.mesh.getElementsByType(4)[0]
+        return {"n_nodes": len(node_tags), "n_tets": len(tet_tags),
+                "bbox_mm": bbox, "elem_size": size}
+    finally:
+        gmsh.finalize()

src/metrics.py

@@ -1,50 +0,0 @@
-"""Evaluation metrics for the solidification-time field (all in seconds)."""
-import numpy as np
-
-
-def mae(y_true, y_pred) -> float:
-    return float(np.mean(np.abs(y_true - y_pred)))
-
-
-def rmse(y_true, y_pred) -> float:
-    return float(np.sqrt(np.mean((y_true - y_pred) ** 2)))
-
-
-def median_rel_error(y_true, y_pred, eps: float = 1.0) -> float:
-    return float(np.median(np.abs(y_true - y_pred) / (y_true + eps)))
-
-
-def _per_graph(graph_id):
-    order = np.argsort(graph_id, kind="stable")
-    gid = graph_id[order]
-    bounds = np.flatnonzero(np.diff(gid)) + 1
-    return order, np.split(order, bounds)
-
-
-def pred_std_ratio(y_true, y_pred, graph_id) -> float:
-    _, groups = _per_graph(np.asarray(graph_id))
-    ratios = []
-    for idx in groups:
-        st = np.std(y_true[idx])
-        if st > 1e-6:
-            ratios.append(np.std(y_pred[idx]) / st)
-    return float(np.mean(ratios)) if ratios else 0.0
-
-
-def hotspot_accuracy(y_true, y_pred, xyz, graph_id,
-                     top_frac: float = 0.10, tol_mm: float = 5.0) -> float:
-    _, groups = _per_graph(np.asarray(graph_id))
-    accs = []
-    for idx in groups:
-        n = len(idx)
-        k = max(1, int(round(n * top_frac)))
-        true_top = idx[np.argsort(y_true[idx])[-k:]]
-        pred_top = idx[np.argsort(y_pred[idx])[-k:]]
-        true_xyz = xyz[true_top]
-        hits = 0
-        for p in pred_top:
-            d = np.linalg.norm(true_xyz - xyz[p], axis=1)
-            if d.min() <= tol_mm:
-                hits += 1
-        accs.append(hits / k)
-    return float(np.mean(accs)) if accs else 0.0

src/model_api.py

@@ -1,53 +1,80 @@
-"""Mesh-in, field-out entry point for the published v5.1 model.
+"""Mesh in, solidification-time field out.
 
     from src.model_api import CastSolidModel
     model = CastSolidModel()                       # loads weights/ once
     seconds = model.predict("part.msh", alloy="G20Mn5")
 
-`predict` takes a gmsh 2.2 ASCII tetrahedral mesh (path or a parsed
-`src.msh_parse.Mesh`) and returns the solidification time in seconds for every
-node, in the mesh's node order. Getting from a STEP file to such a mesh (CAD
-healing, density-targeted meshing) is the job of `src.predict_casting`.
+`predict` takes a gmsh 2.2 ASCII tetrahedral mesh (a path or a parsed
+`src.msh_parse.Mesh`) whose cooling surface is tagged "Cooling", and returns
+the solidification time in seconds for every node, in the mesh's node order.
 """
 from __future__ import annotations
 
 import os
 
 import numpy as np
+import torch
 
-from src.inference_props import resolve_props
+from src import gbm_heads
+from src.features import alloy_g_vector, node_features
+from src.hybrid import hybrid_field_soft
+from src.inference_props import DEFAULT_MOLD, resolve_props
+from src.model_gnn import load_gnn, predict_graph_seconds
 from src.msh_parse import Mesh, parse_msh
-from src.predict_casting import DEFAULT_W_STD, DEFAULT_W_TOTAL, load_hybrid, predict_field
+from src.release_io import load_gbm, load_gnn_state
+
+# Recomposition weights of the published model (see src/hybrid.py).
+W_TOTAL = 1.0
+W_STD = 4.0
+
+GNN_FILE = "gnn_v5.1.safetensors"
+TOTAL_HEAD_FILE = "gbm_total_v3.skops"
+SPREAD_HEAD_FILE = "gbm_std_v2.skops"
 
 
 class CastSolidModel:
     def __init__(self, weights_dir: str = "weights", device: str | None = None):
-        w = lambda name: os.path.join(weights_dir, name)
-        self.device = device
-        (self.gnn, self.gbm_total, self.norm, self.cfg, _,
-         self.gbm_std) = load_hybrid(w("gnn_v5.1.safetensors"), w("gbm_total_v3.skops"),
-                                     device, hgb_std_path=w("gbm_std_v2.skops"))
-        self.device = device or next(self.gnn.parameters()).device.type
-
-    def predict(self, mesh: str | Mesh, alloy: str | None = None, mold: str = "green_sand",
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        state = load_gnn_state(os.path.join(weights_dir, GNN_FILE), self.device)
+        self.gnn, self.cfg, self.norm = load_gnn(state, self.device)
+        self.total_head = load_gbm(os.path.join(weights_dir, TOTAL_HEAD_FILE))
+        self.spread_head = load_gbm(os.path.join(weights_dir, SPREAD_HEAD_FILE))
+        gbm_heads.check_head(self.total_head, gbm_heads.N_TOTAL_FEATURES, "total-time")
+        gbm_heads.check_head(self.spread_head, gbm_heads.N_SPREAD_FEATURES, "spread")
+
+    def predict(self, mesh: str | Mesh, alloy: str | None = None, mold: str = DEFAULT_MOLD,
                 props_json: str | None = None) -> np.ndarray:
-        """Per-node solidification time (s) for `mesh`.
+        """Per-node solidification time (s).
 
         `alloy`/`mold` name rows of `data/properties/` (default EN-GJS-400-15
-        in green sand); `props_json` replaces both with a custom property set.
-        """
+        in green sand); `props_json` replaces both with your own values."""
+        props = resolve_props(alloy=alloy, mold=mold, props_json=props_json)
+        return self.predict_with_props(mesh, props)[0]
+
+    def predict_with_props(self, mesh: str | Mesh, props: dict) -> tuple[np.ndarray, dict]:
+        """(seconds, details) for an already-resolved property dict.
+
+        `details` holds the node feature matrix `X`, the heads' `total_s` and
+        `spread_s`, and the recomposition diagnostics `recomposition`."""
         if not isinstance(mesh, Mesh):
             mesh = parse_msh(os.fspath(mesh))
-        props = resolve_props(alloy=alloy, mold=mold, props_json=props_json)
-        seconds, _ = predict_field(mesh, self.gnn, self.gbm_total, self.norm, cfg=self.cfg,
-                                   device=self.device, props=props, hgb_std=self.gbm_std,
-                                   w_total=DEFAULT_W_TOTAL, w_std=DEFAULT_W_STD)
-        return np.asarray(seconds)
+        feat = node_features(mesh, props["mold_effusivity"])
+        X = feat["X"]
+        gnn_seconds = predict_graph_seconds(self.gnn, self.cfg, self.norm, X, feat["pos"],
+                                            feat["mesh_edges"], alloy_g_vector(props),
+                                            device=self.device)
+        row = gbm_heads.total_row(props, X)
+        total = gbm_heads.predict_total_seconds(self.total_head, row)
+        spread = gbm_heads.predict_spread_seconds(self.spread_head,
+                                                  gbm_heads.spread_row(row, X))
+        seconds, info = hybrid_field_soft(gnn_seconds, total, spread,
+                                          w_total=W_TOTAL, w_std=W_STD)
+        return seconds, {"X": X, "total_s": total, "spread_s": spread,
+                         "recomposition": info}
 
     @staticmethod
     def hotspots(seconds: np.ndarray, top_frac: float = 0.1) -> np.ndarray:
-        """Boolean mask of the `top_frac` latest-solidifying nodes (the same
-        rule as the `hotspot` field `src.predict_casting.write_vtu` writes)."""
+        """Boolean mask of the `top_frac` latest-solidifying nodes."""
         seconds = np.asarray(seconds, dtype=np.float32)
         k = max(1, int(round(seconds.shape[0] * top_frac)))
         return seconds >= np.sort(seconds)[-k]

src/model_gnn.py

@@ -1,29 +1,25 @@
-"""MeshGraphNet-style encode-process-decode GNN predicting standardized
-log1p(solidification_time) per node; global scalar = max over nodes."""
+"""The graph network: an encode-process-decode GNN in the style of
+MeshGraphNets, conditioned on the alloy/mold vector through FiLM.
+
+It predicts the standardized log1p solidification time of every node. The
+graph is the union of the tetrahedral mesh edges and each node's k nearest
+neighbours; an edge carries its displacement vector (mm), its length and a
+flag telling whether it is a mesh edge.
+"""
+import numpy as np
 import torch
 import torch.nn as nn
-from torch_geometric.utils import scatter
+from torch_geometric.nn import knn_graph
+from torch_geometric.utils import scatter, to_undirected
 
 
 def _mlp(i, h, o):
     return nn.Sequential(nn.Linear(i, h), nn.ReLU(), nn.Linear(h, o))
 
 
-def zero_init_input_columns(model, n_new):
-    """Zero the trailing n_new weight columns of the first Linear layer in node_encoder.
-
-    This is used when extending the input dim (e.g., 13 → 15 for new BC channels)
-    to ensure new columns start at zero and contribute nothing at initialization.
-
-    Args:
-        model: SolidGNN instance
-        n_new: Number of trailing columns to zero in the first Linear weight matrix
-    """
-    with torch.no_grad():
-        model.node_encoder[0][0].weight[:, -n_new:] = 0
-
-
 class ProcessorBlock(nn.Module):
+    """One message-passing step with residual edge and node updates."""
+
     def __init__(self, h):
         super().__init__()
         self.edge_mlp = _mlp(3 * h, h, h)
@@ -40,12 +36,12 @@ class ProcessorBlock(nn.Module):
 
 
 class FiLM(nn.Module):
-    """g -> (gamma, beta) per hidden channel; zero-init => gamma=0,beta=0 => identity."""
+    """g -> per-channel (gamma, beta); applied as h * (1 + gamma) + beta."""
+
     def __init__(self, g_dim, hidden):
         super().__init__()
         self.net = nn.Sequential(nn.Linear(g_dim, hidden), nn.SiLU(),
-                                  nn.Linear(hidden, 2 * hidden))
-        nn.init.zeros_(self.net[-1].weight); nn.init.zeros_(self.net[-1].bias)
+                                 nn.Linear(hidden, 2 * hidden))
 
     def forward(self, h, g_nodes):
         gamma, beta = self.net(g_nodes).chunk(2, dim=-1)
@@ -53,84 +49,88 @@ class FiLM(nn.Module):
 
 
 class SolidGNN(nn.Module):
-    def __init__(self, node_in=9, edge_in=4, hidden=128, n_blocks=12, two_head=False,
-                 film=False, g_dim=5, deriv_heads=False):
+    """12 processor blocks (hidden width 128 in the published model), FiLM
+    after each block, and a two-head decoder: a per-node deviation head and
+    a graph-level scale head on the mean node embedding."""
+
+    def __init__(self, node_in=15, edge_in=5, hidden=128, n_blocks=12, g_dim=5):
         super().__init__()
         self.node_encoder = nn.Sequential(_mlp(node_in, hidden, hidden), nn.LayerNorm(hidden))
         self.edge_encoder = nn.Sequential(_mlp(edge_in, hidden, hidden), nn.LayerNorm(hidden))
         self.blocks = nn.ModuleList([ProcessorBlock(hidden) for _ in range(n_blocks)])
         self.decoder = _mlp(hidden, hidden // 2, 1)
-        self.two_head = two_head
-        if two_head:
-            self.scale_head = _mlp(hidden, hidden // 2, 1)
-        self.film = film
-        self.films = nn.ModuleList([FiLM(g_dim, hidden) for _ in range(n_blocks)]) if film else None
-        self.deriv_heads = deriv_heads
-        if deriv_heads:
-            # Optional gradT/dTdt output heads. They hang off the SAME final
-            # hidden state h that the decoder consumes (after the FiLM/blocks
-            # loop, before the two_head mean/scale recomposition) -- they read
-            # h but never modify it, so the frozen y_nodal/y_global path below
-            # is unaffected by their presence.
-            self.grad_head = _mlp(hidden, hidden // 2, 1)
-            self.dtdt_head = _mlp(hidden, hidden // 2, 1)
-
-    def forward(self, x, edge_index, edge_attr, batch, g=None):
+        self.scale_head = _mlp(hidden, hidden // 2, 1)
+        self.films = nn.ModuleList([FiLM(g_dim, hidden) for _ in range(n_blocks)])
+
+    def forward(self, x, edge_index, edge_attr, batch, g):
         h = self.node_encoder(x)
         e = self.edge_encoder(edge_attr)
-        g_nodes = g[batch] if (self.film and g is not None) else None
+        g_nodes = g[batch]
         for i, blk in enumerate(self.blocks):
             h, e = blk(h, edge_index, e)
-            if g_nodes is not None:
-                h = self.films[i](h, g_nodes)
-        if self.deriv_heads:
-            deriv = torch.cat([self.grad_head(h), self.dtdt_head(h)], dim=-1)
+            h = self.films[i](h, g_nodes)
         y_nodal = self.decoder(h).squeeze(-1)
-        if self.two_head:
-            dev = y_nodal - scatter(y_nodal, batch, dim=0, reduce="mean")[batch]
-            scale = self.scale_head(scatter(h, batch, dim=0, reduce="mean")).squeeze(-1)
-            y_nodal = dev + scale[batch]
-        y_global = scatter(y_nodal, batch, dim=0, reduce="max")
-        if self.deriv_heads:
-            return y_nodal, y_global, deriv
-        return y_nodal, y_global
-
-
-def attach_deriv_heads(model, state_dict):
-    """Load a (possibly headless) state_dict into a `deriv_heads=True` model
-    via `load_state_dict(strict=False)`, asserting the ONLY missing keys are
-    the two new heads' parameters (`grad_head.*` / `dtdt_head.*`).
-
-    This is how a production checkpoint trained before the deriv-heads
-    feature existed gets loaded into a `deriv_heads=True` model: everything
-    else must match exactly, or this raises loudly instead of silently
-    leaving some backbone tensor at its random init.
-
-    Args:
-        model: a `SolidGNN(deriv_heads=True)` instance.
-        state_dict: source state_dict (e.g. from an old checkpoint's `model`
-            entry). May or may not already contain grad_head/dtdt_head keys.
-
-    Returns:
-        list of missing key names (all `grad_head.*`/`dtdt_head.*`).
-
-    Raises:
-        ValueError: if any non-head (backbone) key is missing after the load,
-            or if the state_dict contains keys the model doesn't recognize.
-    """
-    result = model.load_state_dict(state_dict, strict=False)
-    if result.unexpected_keys:
-        raise ValueError(
-            f"attach_deriv_heads: unexpected extra keys in state_dict "
-            f"(not present in model): {list(result.unexpected_keys)}"
-        )
-    non_head_missing = [
-        k for k in result.missing_keys
-        if not k.startswith("grad_head.") and not k.startswith("dtdt_head.")
-    ]
-    if non_head_missing:
-        raise ValueError(
-            f"attach_deriv_heads: missing backbone keys after load (expected "
-            f"only grad_head.*/dtdt_head.* to be missing): {non_head_missing}"
-        )
-    return list(result.missing_keys)
+        dev = y_nodal - scatter(y_nodal, batch, dim=0, reduce="mean")[batch]
+        scale = self.scale_head(scatter(h, batch, dim=0, reduce="mean")).squeeze(-1)
+        return dev + scale[batch]
+
+
+def knn_edge_index(pos: torch.Tensor, k: int = 16) -> torch.Tensor:
+    ei = knn_graph(pos, k=k, loop=False)      # directed k-nearest
+    return to_undirected(ei)                   # symmetric
+
+
+def union_edge_index(pos: torch.Tensor, mesh_edges: torch.Tensor, k: int = 16):
+    """Undirected union of the mesh edges and kNN edges, deduplicated.
+
+    mesh_edges: (M, 2) long tensor of unique node pairs (any orientation).
+    Returns (edge_index (2, E) long, is_mesh (E,) float32 flag)."""
+    n = pos.shape[0]
+    knn = knn_edge_index(pos, k=k).to(torch.long)
+    me = mesh_edges.t().to(torch.long)
+    me = torch.cat([me, me.flip(0)], dim=1)                    # both directions
+    mesh_key = me[0] * n + me[1]
+    keys = torch.unique(torch.cat([mesh_key, knn[0] * n + knn[1]]))
+    ei = torch.stack([keys // n, keys % n])
+    is_mesh = torch.isin(keys, mesh_key).to(torch.float32)
+    return ei, is_mesh
+
+
+def union_edge_attr(pos: torch.Tensor, edge_index: torch.Tensor,
+                    is_mesh: torch.Tensor) -> torch.Tensor:
+    """(E, 5): displacement (3), length, mesh-edge flag."""
+    src, dst = edge_index
+    d = pos[dst] - pos[src]
+    return torch.cat([d, d.norm(dim=1, keepdim=True), is_mesh.unsqueeze(1)], dim=1)
+
+
+def load_gnn(state: dict, device):
+    """Build a `SolidGNN` from a `{model, norm, cfg}` dict (see
+    `src.release_io.load_gnn_state`). Returns (model, cfg, norm)."""
+    cfg, norm = state["cfg"], state["norm"]
+    model = SolidGNN(node_in=cfg["node_in"], edge_in=cfg["edge_in"], hidden=cfg["hidden"],
+                     n_blocks=cfg["n_blocks"], g_dim=cfg.get("g_dim", 5)).to(device)
+    model.load_state_dict(state["model"])
+    model.eval()
+    return model, cfg, norm
+
+
+def predict_graph_seconds(model, cfg, norm, X, pos_np, mesh_edges, g_vec,
+                          device="cpu") -> np.ndarray:
+    """Per-node solidification time (s) predicted by the GNN alone.
+
+    `X` is the raw (N, 15) feature matrix, `g_vec` the raw (5,) alloy/mold
+    vector; standardization uses the statistics stored with the weights."""
+    fmean = torch.tensor(norm["feat_mean"], dtype=torch.float32)
+    fstd = torch.tensor(norm["feat_std"], dtype=torch.float32)
+    x = (torch.tensor(X, dtype=torch.float32) - fmean) / fstd
+    pos = torch.tensor(pos_np, dtype=torch.float32)
+    me = torch.tensor(mesh_edges, dtype=torch.long)
+    ei, is_mesh = union_edge_index(pos, me, k=cfg["k"])
+    ea = union_edge_attr(pos, ei, is_mesh)
+    zeros = torch.zeros(x.shape[0], dtype=torch.long, device=device)
+    g = torch.tensor(np.asarray(g_vec), dtype=torch.float32).reshape(1, -1).to(device)
+    with torch.no_grad():
+        yn = model(x.to(device), ei.to(device), ea.to(device), zeros, g)
+    pred_log = yn.cpu().numpy() * norm["y_log_std"] + norm["y_log_mean"]
+    return np.clip(np.expm1(pred_log), 0.0, None)

src/msh_parse.py

@@ -1,8 +1,8 @@
-"""Parser for gmsh 2.2 ASCII meshes used by the Foundry dataset.
+"""Reader for gmsh 2.2 ASCII meshes.
 
-Extracts node coordinates (mm, in node-id order — matches the paired .pt),
-surface triangles grouped by physical name (Cooling/Adiabatic), and volume
-tetrahedra. All connectivity is remapped to 0-based node indices.
+Extracts node coordinates (mm, in the file's node order), surface triangles
+grouped by physical name ("Cooling", optionally "Adiabatic") and volume
+tetrahedra. Connectivity is remapped to 0-based node indices.
 """
 from dataclasses import dataclass
 import numpy as np

src/ood.py

@@ -1,33 +1,7 @@
-"""Out-of-distribution size check: compare a casting's scale quantities against
-the training-set maxima so predictions beyond the training envelope are flagged."""
+"""Out-of-distribution size check: compare a casting's volume, bounding-box
+diagonal and largest distance to cooling with the training maxima
+(`weights/ood_ranges.json`) and warn when any of them is exceeded."""
 import json
-import os.path as osp
-from glob import glob
-
-import numpy as np
-
-from src.features import FEATURE_NAMES
-
-_VOL = FEATURE_NAMES.index("volume_mm3")
-_GC = FEATURE_NAMES.index("geodesic_d_cool")
-
-
-def build_ood_ranges(features_dir: str) -> dict:
-    vol_max = bbox_max = gc_max = 0.0
-    for p in glob(osp.join(features_dir, "graphs", "*.npz")):
-        d = np.load(p)
-        X, pos = d["X"], d["pos"]
-        vol_max = max(vol_max, float(X[0, _VOL]))
-        gc_max = max(gc_max, float(X[:, _GC].max()))
-        diag = float(np.linalg.norm(pos.max(0) - pos.min(0)))
-        bbox_max = max(bbox_max, diag)
-    return {"volume_mm3_max": vol_max, "bbox_diag_max": bbox_max,
-            "geodesic_d_cool_max": gc_max}
-
-
-def save_ranges(ranges: dict, path: str):
-    with open(path, "w") as f:
-        json.dump(ranges, f, indent=2)
 
 
 def load_ranges(path: str) -> dict:

src/predict_casting.py

@@ -1,387 +1,40 @@
-"""Real-casting inference: hybrid field prediction + VTU output.
+"""STEP file in, solidification field out.
 
-predict_field combines the GNN field shape with the baseline per-graph scale
-(src/hybrid.py). write_vtu exports a ParaView file. load_hybrid loads the
-Phase-1 checkpoints. The predict_casting orchestrator (Task 5) ties in meshing.
+    python -m src.predict_casting part.step --out-dir out [--alloy G20Mn5]
 
-v4 (film) path (Task 3): `cfg.get("film")` truthy routes `predict_field`
-through the multi-alloy/multi-BC stack instead -- v4 node features (BC
-spec fixed to a plain sand mold, per this module's bare-mesh-inference
-contract: no chills/sleeves info comes from a STEP file), the alloy/mould
-`g` vector, the v4 GNN (`src.evaluate_ft2b`), and the v2-total GBM
-(`src.gbm_ft2b`) recomposed via the quantile-anchor (`hybrid_field_max_anchor`,
-default `q=0.995` -- outlier-robust; `q=1.0` is the original exact-MAX
-anchor) -- NOT the SP1 MEAN-anchor, which assumes a genuine per-node
-baseline (see `src/evaluate_ft2b.py`'s module docstring for the full
-rationale). The non-film path below is functionally unchanged (plus an explicit HGB-dimension guard).
+Steps: mesh the STEP file (healed, density-searched and cached, or at a fixed
+element size with --elem-size), run the model, check the training envelope,
+screen for shrinkage, and write `<name>.vtu`, `<name>_summary.txt` and the
+mesh `<name>.msh` to the output directory.
+"""
+from __future__ import annotations
 
-variance-recomposition Task 3: an optional GBM std-head (`src.gbm_ft2b.
-refit_gbm_std`'s output, loaded via `load_hybrid`'s `hgb_std_path` /
-`predict_casting`'s `--hgb-std`) can be threaded through the film path so
-`predict_field` recomposes via `src.hybrid.hybrid_field_affine` (rescales
-the GNN field's deviations to match a per-graph SCALE target too, not just
-a per-graph LEVEL) instead of `hybrid_field_max_anchor`. `None` (the
-default, everywhere) is byte-identical to the pre-Task-3 code.
-
-variance-recomposition amendment 2: `predict_casting`'s `--affine-lambda` /
-`predict_field`'s `lambda_shrink` (default `1.0`, forwarded straight to
-`hybrid_field_affine`) geometrically shrinks the affine spread target
-toward the GNN's own spread -- see `src.hybrid.hybrid_field_affine`'s
-docstring for the math/rationale. `1.0` everywhere is byte-identical to
-before this argument existed.
-
-soft-anchor Task 3: `hybrid_field_affine` at the std-head call site above
-is REPLACED by `src.hybrid.hybrid_field_soft` (a weighted two-target
-objective that lets the total anchor itself deviate a bit when that trade
-buys a much closer spread match -- see that function's docstring). The
-degeneration guarantee built into `hybrid_field_soft` means any case where
-the hard affine solve already satisfied both targets recomposes
-IDENTICALLY (same array object, same `a`) -- only genuinely stuck cases
-(the ones that motivated this track) actually change. `feat["recomposition"]`
-is now `"soft"` (was `"affine"`) whenever `hgb_std` is active; the info dict
-gains `"m"`/`"target_unreachable"`/`"total_deviation_frac"`/`"j_total_resid"`/
-`"j_std_resid"` (see `hybrid_field_soft`'s own docstring for the full key
-list) in place of affine's `"a_clamped"`. New `w_total`/`w_std` params
-(`predict_field`, `predict_casting`, CLI `--w-total`/`--w-std`; production
-defaults `DEFAULT_W_TOTAL=1.0` / `DEFAULT_W_STD=4.0` since the v5.1
-promotion) forward straight to `hybrid_field_soft`'s own weights, letting a
-caller trade total-anchor fidelity against spread fidelity. The GBM std
-head's own row width is now dispatched off the loaded model's
-`n_features_in_` (16 -> the v1 row unchanged, 17 -> the v1 row + a
-graph-derived `bc_area_frac` 17th column computed from the SAME in-memory
-per-node feature array `predict_field` already has -- see
-`_std_row_for_predict_field`); anything else raises a loud `ValueError`.
-The no-std-head (quantile-anchor) path is completely untouched by any of
-this.
-
-variance-recomposition Task 6 (the production flip, 2026-07-31):
-`DEFAULT_GNN_CKPT`/`DEFAULT_HGB` now point at the ft3 stack
-(`runs/ft3_stage1/ckpt_ep021.pt` x `runs/ft3_gbm/hgb_v3.joblib`) and a new
-`DEFAULT_HGB_STD` (now `runs/ft3_gbm/hgb_std_v2.joblib` since the v5.1
-promotion) exists. `hgb_std_path`
-no longer defaults to plain `None` on `predict_casting`/CLI `--hgb-std`:
-both default to the module-level `_USE_DEFAULT` sentinel, which
-`load_hybrid` resolves to `DEFAULT_HGB_STD` for a film-capable checkpoint
-or to `None` for a non-film one (see `load_hybrid`'s docstring) -- so the
-production call, with zero extra flags, now recomposes via
-`hybrid_field_affine` (`lambda_shrink=1.0`) instead of the quantile-anchor
-(superseded by soft-anchor Task 3 below -- the zero-flag production call
-recomposes via `hybrid_field_soft` today; this paragraph is kept as the
-Task 6 historical record of when the std head first became the default),
-while the SP1 rollback (`--gnn`/`--hgb` pointed at a non-film checkpoint)
-is untouched: the sentinel silently resolves to `None` there, exactly as
-if `hgb_std_path` had never been passed. The case-insensitive string
-`"none"` always disables the std head regardless of checkpoint kind (the
-explicit rollback switch to the quantile-anchor, usable on any film
-checkpoint, including the new default one) -- passing an explicit path or
-loaded model object still loads/uses it unconditionally, exactly as
-before Task 6.
-
-learned-heads Task 3 (`--gnn` pointed at a `deriv_heads`-capable checkpoint):
-`predict_field`'s film branch also runs `src.evaluate_ft2b.predict_graph_derivs`
-and stashes the result in `feat["gradT"]`/`feat["dTdt"]` (K/m / K/s); `predict_casting`
-then feeds those into `src.shrinkage.screen_casting`'s exact Niyama branch
-instead of the time-only proxy, and relabels the result `niyama_source=
-"learned"` (screen_casting's own "exact" label is reserved for a real cached
-solver field). Absent entirely for a headless checkpoint -- byte-identical
-to the pre-Task-3 proxy-only behavior.
-
-bc-reach Task 4 (cfg-driven v5 feature building at inference): the film
-branch's `compute_features_inference` call now reads
-`version=cfg.get("feature_version", 4)` instead of hardcoding `version=4`
--- the SAME `cfg.get("feature_version", ...)` pattern the non-film branch
-already used (this module, the SP1/non-film path). A v5 (bc-reach)
-checkpoint's saved cfg carries `feature_version=5` (`src.train_gnn`'s
-`_build_checkpoint_dict`, echoed verbatim from `train(feature_version=...)`,
-independent of but in practice equal to `node_in`), so loading one through
-`load_hybrid`/`load_ft2b_gnn` and calling `predict_field` now builds the
-20-column v5 `X` (`src.features.FEATURE_NAMES_V5`: v4's 15 columns +
-`exo_mid`/`exo_long`/`chill_short`/`chill_mid`/`chill_long`) with NO other
-change to this module -- `compute_bc_reach_channels` (features v5) already
-resolves its chill/exo source patch off the SAME `spec`/`bc_spec` dicts
-this module already builds from the all-sand literal, so the reach channels see EXACTLY the same source magnitudes the
-training builder (`src.build_dataset.build_v5`) does. A v4 (or absent)
-`cfg["feature_version"]` is byte-identical to the pre-Task-4 code (defaults
-to `4`, same as before this dispatch existed)."""
-import hashlib
-import json
+import argparse
 import logging
+import os
 import os.path as osp
-import numpy as np
-import torch
-import joblib
 
-logger = logging.getLogger(__name__)
+import numpy as np
 
-from src.features import (compute_features_inference, FEATURE_NAMES, FEATURE_NAMES_V4,
-                          alloy_g_vector)
-from src.dataload import standardize
-from src.baseline import predict_seconds
-from src.evaluate_gnn import predict_graph_seconds, load_gnn, N_BASELINE_FEATURES
-from src.evaluate_ft2b import predict_graph_seconds as predict_graph_seconds_v4, load_ft2b_gnn
-from src.evaluate_ft2b import predict_graph_derivs as predict_graph_derivs_v4
-from src.gbm_ft2b import (gbm_v2_total_seconds, FEATURE_NAMES as GBM_FEATURE_NAMES,
-                          predict_std_seconds, build_v2_feature_row,
-                          bc_area_frac_from_graph)
-from src.hybrid import hybrid_field, hybrid_field_max_anchor, hybrid_field_soft
-from src.inference_props import resolve_props, DEFAULT_ALLOY, resolve_alloy_family
+from src import ood as _ood
+from src.features import FEATURE_NAMES
+from src.inference_props import DEFAULT_ALLOY, DEFAULT_MOLD, resolve_alloy_family, resolve_props
+from src.mesh_cache import get_or_mesh
+from src.meshing import DEFAULT_TARGET_NODES, mesh_step
+from src.model_api import CastSolidModel
+from src.msh_parse import parse_msh
 from src.shrinkage import screen_casting
-from src.release_io import resolve_artifact, load_gbm, load_gnn_any
-
-N_GBM_V2_FEATURES = len(GBM_FEATURE_NAMES)
-_V4_VOL_IDX = FEATURE_NAMES_V4.index("volume_mm3")
-_V4_AREA_IDX = FEATURE_NAMES_V4.index("surface_area_mm2")
-
-# Production v5.1 hybrid stack (promoted 2026-07-31, soft-anchor Task 6 -- see
-# docs/superpowers/specs/2026-07-31-soft-anchor-design.md's ADDENDUM) -- module
-# constants so CLI/kwarg defaults live in one place instead of being duplicated
-# string literals. GNN field *shape* x GBM *scale* (total) x GBM *spread* (std,
-# std_v2 -- median-loss HGB, 17 features incl. bc_area_frac), recomposed via the
-# soft-anchor map (src.hybrid.hybrid_field_soft) at w_total=1.0, w_std=4.0
-# (the e_val sweep winner: v2+soft takes gold decisively -- BUKSA 90.2 s/0.700,
-# H17 120.1 s/0.895, best-ever MAE on both -- missing only the (owner-revised,
-# see tools/ft3_gates.py's DEFAULT_E_VAL_STD_RATIO_FLOOR) e_val exo std_ratio floor).
-#
-# Rollback #1 (v5, 2026-07-31 same day, pre-soft-anchor-sweep): ckpt_ep021 x
-# hgb_v3 x hgb_std_v1 (squared-error, 16 features), w_total=1.0, w_std=1.0 --
-# reachable via `--hgb-std runs/ft3_gbm/hgb_std_v1.joblib --w-std 1.0`
-# (hgb_std_v1.joblib is left untouched on disk by this promotion).
-# Rollback #2 (previous production, 2026-07-28..2026-07-31): ckpt_ep054 x hgb_v2,
-# quantile-anchor recomposition (no std head) -- reachable via
-# `--gnn runs/ft2b_stage2c/ckpt_ep054.pt --hgb runs/ft2b_gbm/hgb_v2.joblib --hgb-std none`.
-# Rollback #3 (SP1, still on disk): runs/sp1_run3_final/best_gnn.pt +
-# runs/baseline_v2/hgb_A.joblib, non-film -- reachable via --gnn/--hgb; --hgb-std needs
-# no extra flag here, `load_hybrid` auto-resolves the defaulted std head to None for a
-# non-film checkpoint (see `load_hybrid`'s docstring / the `_USE_DEFAULT` sentinel below).
-DEFAULT_GNN_CKPT = "runs/ft3_stage1/ckpt_ep021.pt"
-DEFAULT_HGB = "runs/ft3_gbm/hgb_v3.joblib"
-DEFAULT_HGB_STD = "runs/ft3_gbm/hgb_std_v2.joblib"
-DEFAULT_W_TOTAL = 1.0
-DEFAULT_W_STD = 4.0
-
-
-class _UseDefaultHgbStd:
-    """Sentinel for `predict_casting`'s `hgb_std_path` param / CLI `--hgb-std`:
-    means "the caller did not specify" -- `load_hybrid` resolves this to
-    `DEFAULT_HGB_STD` when the selected checkpoint is film-capable, or to
-    `None` when it isn't (so the SP1 rollback via `--gnn`/`--hgb` needs no
-    extra flag). Distinct from an explicit `None` or the string `"none"`
-    (case-insensitive), both of which ALWAYS disable the std head -- the
-    explicit rollback switch to the quantile-anchor recomposition, usable on
-    any film checkpoint including the new default one."""
-
-    def __repr__(self):
-        return "<DEFAULT_HGB_STD>"
-
-
-_USE_DEFAULT = _UseDefaultHgbStd()
-
-
-def _check_hgb_dim(hgb_model, expected: int, ckpt_kind: str) -> None:
-    """Loud, clear failure on a mismatched (film vs. SP1) GNN/HGB pairing.
-
-    `hgb.n_features_in_` is set by sklearn on `.fit()`; a stub without it
-    (shouldn't happen for a real fitted estimator) is treated as "unknown"
-    and NOT checked, rather than crashing the attribute lookup itself.
-    """
-    n_have = getattr(hgb_model, "n_features_in_", None)
-    if n_have is None or n_have == expected:
-        return
-    raise ValueError(
-        f"predict_field: {ckpt_kind} checkpoint expects an HGB baseline with "
-        f"{expected} input features, but hgb_model.n_features_in_={n_have!r} -- "
-        f"this looks like the wrong baseline for this checkpoint (a v2-total "
-        f"GBM has {N_GBM_V2_FEATURES} features, an SP1 per-node HGB has "
-        f"{N_BASELINE_FEATURES}); pass the matching --hgb for this --gnn checkpoint.")
 
+logger = logging.getLogger(__name__)
 
-def _std_row_for_predict_field(row_16: np.ndarray, X: np.ndarray, n_features: int
-                               ) -> np.ndarray:
-    """`predict_field`'s in-memory counterpart to `src.gbm_ft2b.std_feature_row`
-    (soft-anchor Task 3): width-dispatches the std-head row off the loaded
-    std model's own `n_features_in_`, without a manifest stem or graph-cache
-    file to read from -- `predict_field` already has everything both widths
-    need in memory.
-
-    - `n_features == 16` -> `row_16` unchanged (the SAME 16-column row
-      `predict_field`'s total head already scores -- byte-identical to the
-      pre-soft-anchor code).
-    - `n_features == 17` -> `row_16` + `src.gbm_ft2b.bc_area_frac_from_graph`'s
-      17th column, computed from the SAME in-memory per-node v4 feature
-      array `X` (`bc_area_frac_from_graph` only ever indexes `graph["X"]`, so
-      a plain `{"X": X}` dict satisfies its contract without a real graph
-      npz).
-    - anything else -> a loud `ValueError` naming the two supported widths
-      (mirrors `std_feature_row`'s own contract).
-    """
-    if n_features not in (16, 17):
-        raise ValueError(
-            f"predict_field: unsupported std-head n_features_in_={n_features!r} -- "
-            f"the std head only supports 16 (v1) or 17 (v2, v1 + bc_area_frac) "
-            f"feature columns")
-    if n_features == 16:
-        return row_16
-    area_frac = bc_area_frac_from_graph({"X": X})
-    return np.concatenate([row_16, [area_frac]])
-
-
-def predict_field(mesh, gnn_model, hgb_model, norm, cfg=None, device=None, tta=0,
-                  k=16, baseline_norm=None, props=None, hgb_std=None,
-                  lambda_shrink: float = 1.0, w_total: float = DEFAULT_W_TOTAL,
-                  w_std: float = DEFAULT_W_STD) -> tuple:
-    """The film branch always uses the all-sand boundary spec: the whole
-    outer surface cools into the mold, and the model's chill/sleeve input
-    channels stay at zero (this release supports the sand-mold regime only).
-
-    `hgb_std` (variance-recomposition Task 3, recomposition swapped to
-    `hybrid_field_soft` by soft-anchor Task 3): an optional loaded GBM
-    std-head model (`src.gbm_ft2b.refit_gbm_std`/`refit_gbm_std_v2`'s
-    output -- e.g. loaded via `load_hybrid`'s `hgb_std_path`). Only
-    consulted on the film (v4) path, and only once the total head has
-    produced a valid `total`: the 16-column total-head row is built ONCE
-    (`src.gbm_ft2b.build_v2_feature_row`, the SAME public row-builder
-    `gbm_v2_total_seconds` itself uses, so a `FEATURE_NAMES` reorder can
-    never silently desync the two heads) and fed to the total head
-    (`hgb_model.predict` directly, mathematically identical to
-    `gbm_v2_total_seconds`); the std head is scored on a SEPARATE row,
-    width-dispatched off `hgb_std.n_features_in_` (16 -> that same row
-    unchanged, 17 -> that row + a graph-derived `bc_area_frac` 17th column,
-    via `_std_row_for_predict_field` -- else a loud `ValueError`), then
-    recomposed via `src.hybrid.hybrid_field_soft` instead of
-    `hybrid_field_max_anchor`.
-    `None` (the default) is BYTE-IDENTICAL to the pre-Task-3 code (the
-    `total = gbm_v2_total_seconds(...)` / `hybrid_field_max_anchor` path
-    below is untouched). When active, the returned `feat` dict gains
-    `feat["recomposition"] = "soft"` and `feat["recomposition_info"]`
-    (`hybrid_field_soft`'s `info`: `a`/`m`/`target_unreachable`/
-    `degenerate_fallback`/`clip_frac`/`total_deviation_frac`/
-    `j_total_resid`/`j_std_resid`/`lambda_shrink`/`std_target_effective`)
-    -- absent entirely when `hgb_std` is `None`. Ignored entirely by the
-    non-film (SP1) path -- SP1 has no scale-head std target.
-
-    `lambda_shrink` (variance-recomposition amendment 2): forwarded straight
-    to `hybrid_field_soft` (ignored when `hgb_std` is `None`, same as every
-    other soft-anchor-only knob here). `1.0` (the default) is byte-identical
-    to before this argument existed.
-
-    `w_total`/`w_std` (soft-anchor Task 3; defaults changed to `DEFAULT_W_TOTAL`/
-    `DEFAULT_W_STD` by the soft-anchor Task 6 promotion, 2026-07-31): forwarded
-    straight to `src.hybrid.hybrid_field_soft`'s own weights (ignored when
-    `hgb_std` is `None`). The defaults here (`1.0`/`4.0`) are the e_val-sweep
-    winner recorded in the spec addendum -- NOT `hybrid_field_soft`'s own
-    function-level default of `1.0`/`1.0` (that default is unchanged; a
-    caller invoking `hybrid_field_soft` directly, bypassing `predict_field`,
-    still gets equal weighting unless it passes these explicitly).
-
-    Learned-heads Task 3: when `cfg["deriv_heads"]` is truthy (a heads
-    checkpoint, film path only -- SP1 has no deriv heads at all), `feat`
-    additionally carries `feat["gradT"]`/`feat["dTdt"]` -- the checkpoint's
-    OWN learned per-node gradT (K/m) / dTdt (K/s) predictions
-    (`src.evaluate_ft2b.predict_graph_derivs`), on the SAME features the
-    shape head just scored. This is a `feat`-dict addition only -- the
-    `(pred, feat)` 2-tuple return contract is unchanged, so every existing
-    `out, feat = predict_field(...)` call site keeps working unmodified.
-    Absent entirely for a headless checkpoint (the default): `"gradT"`/
-    `"dTdt"` simply aren't keys of `feat` then, exactly as before this
-    feature existed. `predict_casting` reads these (if present) and feeds
-    them to `screen_casting`'s exact Niyama branch, labeling the result
-    `niyama_source="learned"` (see that function's own docstring)."""
-    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-    cfg = cfg or {"k": k, "graph": "knn", "feature_version": 2}
-    gnn_model = gnn_model.to(device).eval()
-
-    if cfg.get("film", False):
-        if props is None:
-            raise ValueError(
-                "predict_field: cfg indicates a film (v4) checkpoint but no "
-                "`props` (alloy/mould conditioning dict) were given")
-        if hgb_model is None:
-            raise ValueError(
-                "predict_field: a film (v4) checkpoint requires an hgb_model "
-                "(the v2-total GBM scale head) -- got None")
-        _check_hgb_dim(hgb_model, N_GBM_V2_FEATURES, "film (v4)")
-
-        spec = {"mold_effusivity": props["mold_effusivity"], "bc_config": "none",
-                "chill_effusivity": 0.0, "sleeve_MEF": 0.0, "exo_q0": 0.0}
-        bc_spec = None
-        feat = compute_features_inference(mesh, version=cfg.get("feature_version", 4), spec=spec)
-        X, pos = feat["X"], feat["pos"]
-        g_vec = alloy_g_vector(props)
-        gnn_seconds = predict_graph_seconds_v4(gnn_model, cfg, norm, X, pos,
-                                               feat["mesh_edges"], g_vec=g_vec,
-                                               device=device)
-
-        # Learned-heads Task 3: a deriv_heads-capable checkpoint additionally
-        # predicts per-node gradT (K/m) / dTdt (K/s) off the SAME features --
-        # stashed into `feat` (not the return tuple's arity, keeping every
-        # existing `out, feat = predict_field(...)` call site unchanged) so
-        # `predict_casting` can feed them to `screen_casting`'s exact branch
-        # instead of the time-only proxy. Absent entirely (no "gradT"/"dTdt"
-        # keys in `feat`) for a headless checkpoint -- byte-identical to the
-        # pre-Task-3 code for every existing (non-heads) checkpoint.
-        if cfg.get("deriv_heads", False):
-            gradT_pred, dTdt_pred = predict_graph_derivs_v4(
-                gnn_model, cfg, norm, X, pos, feat["mesh_edges"], g_vec=g_vec,
-                device=device)
-            feat["gradT"] = gradT_pred
-            feat["dTdt"] = dTdt_pred
-
-        v_over_a = float(X[0, _V4_VOL_IDX]) / float(X[0, _V4_AREA_IDX])
-
-        if hgb_std is not None:
-            row = build_v2_feature_row(props, v_over_a, bc_spec)
-            total = float(np.expm1(float(hgb_model.predict(row.reshape(1, -1))[0])))
-            if total <= 0:
-                raise ValueError(
-                    f"predict_field: the v2-total GBM head returned a non-positive "
-                    f"total ({total!r}) for props={props!r} -- the v2-total GBM is "
-                    f"broken for this input")
-            n_std_features = int(getattr(hgb_std, "n_features_in_", 16))
-            std_row = _std_row_for_predict_field(row, X, n_std_features)
-            std_seconds = predict_std_seconds(hgb_std, std_row)
-            pred, info = hybrid_field_soft(gnn_seconds, total, std_seconds,
-                                           lambda_shrink=lambda_shrink,
-                                           w_total=w_total, w_std=w_std)
-            feat["recomposition"] = "soft"
-            feat["recomposition_info"] = info
-            return pred, feat
-
-        total = gbm_v2_total_seconds(hgb_model, props, v_over_a, bc_spec=bc_spec)
-        if total <= 0:
-            raise ValueError(
-                f"predict_field: gbm_v2_total_seconds returned a non-positive "
-                f"total ({total!r}) for props={props!r} -- the v2-total GBM is "
-                f"broken for this input")
-        return hybrid_field_max_anchor(gnn_seconds, total), feat
-
-    # --- non-film (SP1) path: functionally unchanged (plus explicit HGB-dimension guard) ---
-    if hgb_model is not None:
-        _check_hgb_dim(hgb_model, N_BASELINE_FEATURES, "non-film (SP1)")
-    feat = compute_features_inference(mesh, version=cfg.get("feature_version", 2))
-    X, pos = feat["X"], feat["pos"]
-    Xs = standardize(X, norm)
-    x = torch.tensor(Xs, dtype=torch.float32)
-    pos_t = torch.tensor(pos, dtype=torch.float32)
-    me = (torch.tensor(feat["mesh_edges"], dtype=torch.long)
-          if cfg.get("graph", "knn") == "union" else None)
-    g = predict_graph_seconds(gnn_model, cfg, norm, x, pos_t, me, device=device, tta=tta)
-    if hgb_model is None:
-        return g, feat
-    bn = baseline_norm if baseline_norm is not None else norm
-    b = predict_seconds(hgb_model, standardize(X[:, :N_BASELINE_FEATURES], bn), bn)
-    return hybrid_field(g, b), feat
+MODEL_VERSION = "v5.1"
+_VOL_IDX = FEATURE_NAMES.index("volume_mm3")
+_GCOOL_IDX = FEATURE_NAMES.index("geodesic_d_cool")
 
 
 def write_vtu(out_path, node_xyz, tets, seconds, top_frac=0.1, extra_point_data=None):
-    """`extra_point_data` (SP2 shrinkage Task 4): an optional
-    `{field_name: (N,) array}` dict merged into the written `point_data`
-    dict alongside the original `solidification_time`/`hotspot` fields --
-    e.g. `{"niyama": ..., "porosity_risk": ..., "macro_region_id": ...}`
-    from `src.shrinkage.screen_casting`. `None` (the default) is
-    byte-identical to the pre-Task-4 contract (only the original two
-    fields are written) -- this is a generic extension point, not a
-    hardcoded shrinkage-only special case, so any future per-node field
-    can reuse it the same way."""
+    """Write the mesh with `solidification_time` (s) and `hotspot` (the
+    `top_frac` latest-solidifying nodes) point fields, plus any extras."""
     import meshio
     seconds = np.asarray(seconds, dtype=np.float32)
     n = seconds.shape[0]
@@ -397,37 +50,20 @@ def write_vtu(out_path, node_xyz, tets, seconds, top_frac=0.1, extra_point_data=
     m.write(out_path)
 
 
-# --------------------------------------------------------------------------- #
-# SP2 shrinkage Task 4: `predict_casting` integration helpers -- deriving
-# screen_casting's mesh-topology inputs from the parsed `.msh`, trimming its
-# result to a return-dict-friendly shape, and formatting the summary's risk
-# block. Kept as free functions (rather than inlined in `predict_casting`)
-# so each piece is independently testable and the main function stays
-# readable.
-# --------------------------------------------------------------------------- #
-def _shrinkage_surface_and_feed_nodes(mesh):
-    """Controller directive 1 (owner-adjudicated 2026-08-01): `surface_nodes`
-    = unique nodes of ALL surface triangles (every physical-group tag the
-    mesh carries -- "Cooling" union "Adiabatic" -- matching `screen_casting`'s
-    own contract of "the whole free surface", used for sink/feeder
-    placement). `feed_nodes` = unique nodes of "Adiabatic"-tagged triangles
-    ONLY (riser/sleeve faces, per `src.shrinkage.is_fed`'s documented
-    convention) -- empty when the group is absent or empty (a bare sand
-    casting with no risers, conservatively "everything unfed")."""
+def _surface_and_feed_nodes(mesh):
+    """All tagged surface nodes, and the adiabatic-tagged ones that can feed
+    a shrinkage region (none on a STEP meshed by this pipeline)."""
     tri_arrays = [tri for tri in mesh.tris.values() if tri.size]
     surface_nodes = (np.unique(np.concatenate(tri_arrays, axis=0).reshape(-1))
                      if tri_arrays else np.zeros(0, dtype=np.int64))
     adiab_tri = mesh.tris.get("Adiabatic")
     feed_nodes = (np.unique(adiab_tri.reshape(-1))
-                 if adiab_tri is not None and adiab_tri.size else np.zeros(0, dtype=np.int64))
+                  if adiab_tri is not None and adiab_tri.size else np.zeros(0, dtype=np.int64))
     return surface_nodes.astype(np.int64), feed_nodes.astype(np.int64)
 
 
 def _shrinkage_result_for_return(screen: dict) -> dict:
-    """Directive 6: the `screen_casting` result MINUS its bulky per-node
-    arrays (`niyama`/`porosity_risk` dropped), plus summary counts +
-    `flagged_frac` -- everything else (`macro_shrinkage`/`top_sinks`/
-    `feeders`/`material`/`niyama_source`/`per_orientation`) is kept as-is."""
+    """The screening result without its per-node arrays, plus counts."""
     result = {k: v for k, v in screen.items() if k not in ("niyama", "porosity_risk")}
     risk = screen["porosity_risk"]
     result["flagged_frac"] = float(np.mean(risk)) if risk.size else 0.0
@@ -437,10 +73,7 @@ def _shrinkage_result_for_return(screen: dict) -> dict:
     return result
 
 
-def _shrinkage_extra_point_data(mesh, screen: dict) -> dict:
-    """The 3 VTU point fields directive 4 asks for: `niyama` (float),
-    `porosity_risk` (uint8 0/1), and `macro_region_id` (-1 = no macro
-    region, else that region's own `region_id`)."""
+def _shrinkage_point_data(mesh, screen: dict) -> dict:
     n_nodes = mesh.node_xyz.shape[0]
     macro_region_id = np.full(n_nodes, -1, dtype=np.int64)
     for region in screen["macro_shrinkage"]:
@@ -452,44 +85,19 @@ def _shrinkage_extra_point_data(mesh, screen: dict) -> dict:
     }
 
 
-def _write_shrinkage_summary_block(f, shrinkage) -> None:
-    """Directive 5's exact skeleton: a `shrinkage screening:` section with
-    niyama_source, flagged-node %, N macro regions (worst severity +
-    centroid), N sinks, and feeder recommendations (location + min
-    modulus) -- or a single `shrinkage screening: disabled` line.
-
-    Only called when screening was ATTEMPTED (`screen_shrinkage=True`) but
-    either produced a result (`shrinkage` is the trimmed dict) or was
-    internally skipped (`shrinkage` is `None`, e.g. no surface nodes to
-    place a feeder against) -- `screen_shrinkage=False` (the explicit
-    opt-out) never calls this at all, so that path stays byte-identical to
-    the pre-Task-4 summary (see `predict_casting`'s own docstring)."""
+def _write_shrinkage_summary(f, shrinkage) -> None:
     if shrinkage is None:
         f.write("shrinkage screening: disabled\n")
         return
     f.write("shrinkage screening:\n")
     f.write(f"  niyama_source: {shrinkage['niyama_source']}\n")
-    flagged_line = f"  flagged nodes: {shrinkage['flagged_frac'] * 100.0:.1f}%"
-    if shrinkage["niyama_source"] == "proxy":
-        # Review round 1, finding 4: the proxy path's risk mask is
-        # `porosity_risk(mode="percentile")` -- ALWAYS ~the bottom decile by
-        # construction (src.shrinkage's default `pct=10.0`), a within-part
-        # relative ranking, not an absolute risk level. Without this note a
-        # reader could mistake "10.0%" for a measured severity rather than
-        # the fixed cutoff it structurally always is on this path.
-        flagged_line += " (bottom decile by construction)"
-    f.write(flagged_line + "\n")
+    # The risk mask is the lowest Niyama decile of the part, so this
+    # percentage is fixed by construction, not a measured severity.
+    f.write(f"  flagged nodes: {shrinkage['flagged_frac'] * 100.0:.1f}%"
+            f" (bottom decile by construction)\n")
     macro = shrinkage["macro_shrinkage"]
     if macro:
-        # Review finding C1: regions with fewer than 4 nodes
-        # (`region["sub_resolution"]`, set by `src.shrinkage.hot_regions`)
-        # are still listed here (their volume/severity stay meaningful) but
-        # reported SEPARATELY from the well-resolved count -- they never
-        # get a feeder recommendation below (their V/A modulus diverges
-        # without bound at this resolution; see `feeder_recommendations`'s
-        # docstring), so folding them into one undifferentiated N would
-        # overstate how many of the N regions actually have a usable
-        # feeder placement.
+        # Regions under 4 nodes are listed but get no feeder suggestion.
         n_sub_resolution = sum(1 for r in macro if r.get("sub_resolution"))
         region_count = str(len(macro))
         if n_sub_resolution:
@@ -497,8 +105,8 @@ def _write_shrinkage_summary_block(f, shrinkage) -> None:
         worst = max(macro, key=lambda r: r["severity"])
         c = worst["centroid"]
         f.write(f"  macro shrinkage regions: {region_count} "
-               f"(worst severity={worst['severity']:.4g} mm^3 @ "
-               f"({c[0]:.1f},{c[1]:.1f},{c[2]:.1f}))\n")
+                f"(worst severity={worst['severity']:.4g} mm^3 @ "
+                f"({c[0]:.1f},{c[1]:.1f},{c[2]:.1f}))\n")
     else:
         f.write("  macro shrinkage regions: 0\n")
     f.write(f"  sinks: {len(shrinkage['top_sinks'])}\n")
@@ -508,424 +116,114 @@ def _write_shrinkage_summary_block(f, shrinkage) -> None:
         for feeder in feeders:
             loc = feeder["location"]
             f.write(f"    - location=({loc[0]:.1f},{loc[1]:.1f},{loc[2]:.1f}) mm  "
-                   f"min_modulus_mm={feeder['min_modulus_mm']:.4f}\n")
+                    f"min_modulus_mm={feeder['min_modulus_mm']:.4f}\n")
     else:
         f.write("  feeder recommendations: none\n")
     if shrinkage.get("alloy_family_assumed"):
         f.write(f"  alloy_family_assumed: {shrinkage['material']['alloy_family']} "
-               f"(alloy grade name not recognized -- conservative fallback)\n")
-
-
-def _parse_gravity(s: str) -> tuple:
-    """CLI `--gravity` parser: `"x,y,z"` -> `(float(x), float(y), float(z))`.
-    A small free function so it's unit-testable without going through
-    `argparse`/`sys.argv` -- and (review round 1, finding 5) passed directly
-    as `argparse`'s own `type=` callable, so a malformed `--gravity` value
-    surfaces as argparse's normal `ValueError`-driven usage error (exit
-    code 2, "argument --gravity: invalid _parse_gravity value: ...") rather
-    than an uncaught traceback deep inside `predict_casting`.
-
-    Rejects the zero vector loudly (`ValueError`): `screen_casting`'s
-    `top_sinks`/`feeder_recommendations` normalize `gravity_vec` by its own
-    norm and would divide by zero silently otherwise (`up = up / norm`
-    guarded only by `norm > 0`, which degrades to a NEVER-normalized,
-    physically meaningless direction rather than raising) -- a caller
-    typo'ing `0,0,0` deserves a clear error, not a silently-wrong feeder
-    placement."""
-    parts = [p.strip() for p in s.split(",")]
-    if len(parts) != 3:
-        raise ValueError(f"--gravity must be 'x,y,z' (3 comma-separated numbers), got {s!r}")
-    vec = tuple(float(p) for p in parts)
-    if all(v == 0.0 for v in vec):
-        raise ValueError(
-            f"--gravity must not be the zero vector (got {s!r}) -- gravity_vec is "
-            f"normalized by its own norm for top-sink/feeder placement, and the zero "
-            f"vector has no direction")
-    return vec
-
-
-def load_hybrid(gnn_ckpt, hgb_path, device=None,
-                baseline_features_dir="data/features_v2", hgb_std_path=None):
-    """Load the GNN + (optional) HGB baseline (+ optional GBM std-head) for
-    `predict_field`.
-
-    Film (v4) checkpoints load via `src.evaluate_ft2b.load_ft2b_gnn` and skip
-    `baseline_norm` entirely (v4's HGB is the v2-total scale GBM, standardized
-    by `src.gbm_ft2b` itself -- there is no separate per-node baseline-feature
-    norm to load). Non-film checkpoints are loaded exactly as before
-    (`src.evaluate_gnn.load_gnn`) -- byte-identical to the pre-v4 code.
-
-    `hgb_std_path` (variance-recomposition Task 3): an optional path to a
-    `joblib`-dumped GBM std-head model (`src.gbm_ft2b.refit_gbm_std`'s
-    output), loaded exactly like `hgb_path` (`joblib.load` if given, else
-    `None`). `None` (the default -- unchanged by Task 6, see below) means
-    the returned `hgb_std` is `None` and every other return value / code
-    path is byte-identical to before this argument existed.
-
-    `hgb_std_path` may also be (variance-recomposition Task 6):
-    - the module-level `_USE_DEFAULT` sentinel: resolved HERE, using the
-      GNN checkpoint's own `cfg["film"]` flag (already being determined
-      below to pick the loader) -- `DEFAULT_HGB_STD` for a film-capable
-      checkpoint, `None` for a non-film one. This is what lets
-      `predict_casting`'s own sentinel default reach the new production
-      std head with zero flags, while leaving a non-film (SP1) `--gnn`/
-      `--hgb` rollback untouched: no attempt is made to load
-      `DEFAULT_HGB_STD` at all in that case, so a missing/irrelevant v4-only
-      artifact never breaks a non-film call.
-    - the case-insensitive string `"none"`: always resolved to `None`,
-      regardless of film-ness -- the explicit rollback switch to the
-      quantile-anchor recomposition (usable on any film checkpoint,
-      including the new default one).
-    An explicit path (or an already-loaded model object, via `hgb_path`'s
-    own convention) is still loaded/used unconditionally either way, exactly
-    as before Task 6 -- `load_hybrid` never judges whether a checkpoint/
-    std-head pairing makes sense; that is `predict_casting`'s job (it raises
-    if the resulting `hgb_std` is non-`None` for a non-film checkpoint).
-
-    Returns a 6-tuple (was 5, before this argument existed): `(gnn, hgb,
-    norm, cfg, baseline_norm, hgb_std)` -- every existing call site was
-    updated to the new arity when this argument was added; any NEW call site
-    must unpack all 6.
-    """
-    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-    gnn_ckpt, hgb_path = resolve_artifact(gnn_ckpt), resolve_artifact(hgb_path)
-    hgb = load_gbm(hgb_path)
-    state = load_gnn_any(gnn_ckpt, device)
-    is_film = state.get("cfg", {}).get("film", False)
-
-    if hgb_std_path is _USE_DEFAULT:
-        hgb_std_path = DEFAULT_HGB_STD if is_film else None
-    elif isinstance(hgb_std_path, str) and hgb_std_path.strip().lower() == "none":
-        hgb_std_path = None
-    hgb_std_path = resolve_artifact(hgb_std_path)
-    try:
-        hgb_std = load_gbm(hgb_std_path)
-    except FileNotFoundError as e:
-        raise FileNotFoundError(
-            f"GBM std head not found: {hgb_std_path!r}. Since the 2026-07-31 v5 "
-            f"promotion the std artifact ({DEFAULT_HGB_STD}) is a PRODUCTION "
-            f"requirement of the default affine recomposition -- provision it "
-            f"alongside the GNN/GBM artifacts, or pass --hgb-std none (API: "
-            f"hgb_std_path=None) to fall back to the anchor recomposition.") from e
-
-    if is_film:
-        gnn, cfg, norm = load_ft2b_gnn(state, device)
-        return gnn, hgb, norm, cfg, None, hgb_std
-    gnn, cfg, norm = load_gnn(gnn_ckpt, device)
-    from src.dataload import load_norm
-    baseline_norm = load_norm(baseline_features_dir) if hgb is not None else None
-    return gnn, hgb, norm, cfg, baseline_norm, hgb_std
-
-
-import os
-import argparse
-
-from src.step_mesh import mesh_step
-from src.msh_parse import parse_msh
-from src import ood as _ood
-from src.mesh_cache import get_or_mesh
-from tools.doe.prep_parts import DEFAULT_TARGET_NODES
+                f"(alloy not in the materials table -- conservative fallback)\n")
 
 
-def predict_casting(step_path, out_dir="runs/castings", elem_size=None,
-                    mesh_mode="prep",
-                    gnn_ckpt=DEFAULT_GNN_CKPT,
-                    hgb_path=DEFAULT_HGB,
-                    ood_json="data/features_v2/ood_ranges.json", device=None,
-                    tta=0, alloy=None, mold="green_sand", props_json=None,
+def _mesh(step_path, out_msh, elem_size, target_nodes, cache_dir):
+    if elem_size is not None:
+        info = mesh_step(step_path, out_msh, elem_size)
+        mesh = parse_msh(out_msh)
+        return mesh, {"mode": "fixed", "n_nodes": info["n_nodes"], "n_tets": info["n_tets"],
+                      "elem_size": info["elem_size"], "cache": "n/a"}
+    info = get_or_mesh(step_path, out_msh, target_nodes=target_nodes, cache_dir=cache_dir)
+    if not info.get("ok", False):
+        raise RuntimeError(f"meshing failed for {step_path!r}: {info.get('notes', '(no notes)')}")
+    mesh = parse_msh(out_msh)
+    return mesh, {"mode": "density-searched", "n_nodes": info["n_nodes"],
+                  "n_tets": int(mesh.tets.shape[0]), "elem_size": None,
+                  "cache": info.get("cache", "unknown")}
+
+
+def predict_casting(step_path, out_dir="runs/castings", alloy=None, mold=DEFAULT_MOLD,
+                    props_json=None, elem_size=None,
                     target_nodes: int = DEFAULT_TARGET_NODES,
-                    hgb_std_path=_USE_DEFAULT,
-                    lambda_shrink: float = 1.0, w_total: float = DEFAULT_W_TOTAL,
-                    w_std: float = DEFAULT_W_STD, screen_shrinkage: bool = True,
-                    gravity: tuple = (0, 0, -1)) -> dict:
-    """`mesh_mode`: "prep" (default) routes meshing through
-    `tools.doe.prep_parts.heal_and_mesh` -- the SAME largest-solid + OCC
-    heal + density-searched pipeline used to prep the training/dataset-D
-    parts, so real-part inference and training use one honest meshing
-    stack. "raw" keeps the old `src.step_mesh.mesh_step` single-pass
-    uniform-density path (escape hatch). An explicit `elem_size` always
-    implies "raw" regardless of `mesh_mode` -- density-searching towards a
-    node-count target and pinning an exact element size are mutually
-    exclusive requests, and the explicit size wins. "prep" mode meshes via
-    `src.mesh_cache.get_or_mesh`, so a byte-identical STEP file (same
-    `target_nodes`) is served from `runs/mesh_cache/` instead of re-meshed.
-
-    `target_nodes` (app Phase A, 2026-08-01): the node-count target the
-    "prep" density search aims at, forwarded to `src.mesh_cache.get_or_mesh`
-    (and thus part of the mesh-cache key, so two densities of the same STEP
-    cache separately). Defaults to `DEFAULT_TARGET_NODES` (8500), which is
-    byte-identical to the previous hardcoded behaviour. The trained density
-    band is ~7.5-10k nodes (DECISIONS-2026-07-22 §9); values outside it put
-    the mesh outside what the model saw, and callers exposing this to a user
-    should bound it there. Passing a non-default value together with an
-    explicit `elem_size` raises `ValueError` rather than being ignored --
-    `elem_size` forces "raw", which never runs the density search at all.
-
-    `hgb_std_path` (variance-recomposition Task 3; default changed in Task 6;
-    recomposition swapped to `hybrid_field_soft` by soft-anchor Task 3): an
-    optional path to a `joblib`-dumped GBM std-head model
-    (`src.gbm_ft2b.refit_gbm_std`/`refit_gbm_std_v2`'s output, 16 or 17
-    features -- `predict_field` width-dispatches off the loaded model's own
-    `n_features_in_`), forwarded to `load_hybrid`/`predict_field`. Defaults
-    to the module-level `_USE_DEFAULT` sentinel (not `None`, since Task 6):
-    `load_hybrid` resolves this to `DEFAULT_HGB_STD` for a film-capable
-    checkpoint (the production call, zero extra flags, now recomposes via
-    `hybrid_field_soft`) or to `None` for a non-film one (the SP1 rollback
-    via `--gnn`/`--hgb` is unaffected -- no attempt is made to load
-    `DEFAULT_HGB_STD` for a non-film checkpoint at all). Pass `None`
-    explicitly, or the case-insensitive string `"none"` (the CLI-friendly
-    form), to force the quantile-anchor recomposition on ANY checkpoint,
-    including the new default film one -- this is byte-identical to the
-    pre-Task-3 code. Passing an explicit path (or an already-loaded model
-    object) still requires a film-capable checkpoint (SP1 has no
-    scale-head std target) -- raises `ValueError`
-    otherwise, BEFORE meshing. When the soft-anchor recomposition is active,
-    the summary file gains a `soft recomposition: a=..., m=...,
-    target_unreachable=..., clip_frac=...` line (`target_unreachable`
-    surfaces whether the underlying hard (affine) two-target problem the
-    soft solve seeds from was itself unreachable at some `a` -- see
-    `src.hybrid.hybrid_field_soft`), further appending a trailing
-    `total_dev=+x.x%` field (final review I-1, 2026-08-01) whenever
-    `abs(info["total_deviation_frac"]) > 0.001` -- signed, one decimal
-    (e.g. `total_dev=+4.7%`) -- surfacing the soft anchor's defining
-    trade-off (letting the total deviate from the GBM total to buy a closer
-    spread match) in the human-readable summary, not just the return dict's
-    `recomposition_info`; omitted entirely below that threshold so a healthy
-    part's summary stays clean. Otherwise (no active std head) it gains a
-    `recomposition: anchor ...` line (film checkpoints only).
-
-    `lambda_shrink` (variance-recomposition amendment 2, CLI `--affine-
-    lambda`): forwarded straight to `predict_field`/`hybrid_field_soft`
-    (ignored unless `hgb_std_path` is also given). `1.0` (the default) is
-    byte-identical to before this argument existed. When active AND
-    `lambda_shrink != 1.0`, the summary's recomposition line gains a
-    trailing `lambda_shrink=...` field.
-
-    `w_total`/`w_std` (soft-anchor Task 3, CLI `--w-total`/`--w-std`; defaults
-    changed to `DEFAULT_W_TOTAL`/`DEFAULT_W_STD` (`1.0`/`4.0`) by the
-    soft-anchor Task 6 promotion, 2026-07-31 -- the e_val-sweep winner
-    recorded in the spec addendum): forwarded straight to
-    `predict_field`/`hybrid_field_soft`'s own weights (ignored unless
-    `hgb_std_path` is also given). When either differs from `1.0`/`1.0` (NOT
-    from the new defaults -- this compares against `hybrid_field_soft`'s own
-    equal-weighting baseline, so the new production default line still
-    prints its weights), the summary's recomposition line gains a trailing
-    `w_total=..., w_std=...` field.
-
-    `screen_shrinkage` (SP2 shrinkage Task 4, CLI `--no-shrinkage`; learned
-    source added by learned-heads Task 3): runs `src.shrinkage.screen_casting`
-    on the recomposed hybrid field (`t_s`) AFTER recomposition, BEFORE the
-    VTU write -- Niyama (the time-only proxy, UNLESS `gnn_ckpt` is a
-    deriv_heads-capable checkpoint, in which case `predict_field` has already
-    populated `feat["gradT"]`/`feat["dTdt"]` with the checkpoint's own learned
-    per-node predictions, routing `screen_casting` through its exact branch
-    instead; the returned `shrinkage["niyama_source"]` is then relabeled
-    `"learned"` here -- NOT `screen_casting`'s own `"exact"`, which is
-    reserved for a real cached solver field, e.g. Gate C's per-row
-    comparison), macro-shrinkage region detection, top-surface sinks, and
-    feeder recommendations. `surface_nodes`/`feed_nodes` come straight off
-    the parsed mesh's physical-group triangles (`_shrinkage_surface_and_
-    feed_nodes` -- ALL surface tris union for `surface_nodes`, "Adiabatic"
-    tris only for `feed_nodes`, empty for a bare sand casting with no
-    risers); `TL`/`TS` come from the already-resolved `props`; `alloy_family`
-    comes from `src.inference_props.resolve_alloy_family` (a name-prefix map
-    onto `src.shrinkage.SHRINKAGE_PCT`'s family keys -- an unrecognized grade
-    name, or any `--props` custom-JSON override, falls back loudly to
-    `"steel"` and is recorded as `alloy_family_assumed` in the result rather
-    than crashing or silently mis-scaling severities). `True` (the default)
-    adds `niyama`/`porosity_risk`/`macro_region_id` VTU point fields, a
-    `shrinkage screening:` block in the summary, and a `"shrinkage"` key in
-    the return dict (the `screen_casting` result minus its bulky per-node
-    `niyama`/`porosity_risk` arrays, plus `flagged_frac`/`n_macro_regions`/
-    `n_sinks`/`n_feeders` counts). The shrinkage dict carries
-    `"niyama_confidence": "full"` (this release supports the all-sand regime
-    only, where the proxy passed its validation gate). `False` skips all of
-    that entirely -- `"shrinkage"` is `None` and the VTU/summary carry no
-    extra fields (the `shrinkage screening: disabled` line only appears when
-    screening was attempted but internally skipped, e.g. a mesh with no
-    surface nodes to place a feeder against, never for this explicit
-    opt-out).
-
-    `gravity` (CLI `--gravity x,y,z`, parsed by `_parse_gravity`): forwarded
-    straight to `screen_casting`'s `gravity_vec` (top-sink/feeder placement
-    direction; conduction/Niyama/macro-region *detection* don't depend on
-    it). `(0, 0, -1)` (the default) is the same "gravity along -z" convention
-    used elsewhere in this module. Ignored entirely when `screen_shrinkage`
-    is `False`.
-
-    Returns a dict that (soft-anchor Task 4, v5 review M-1) always carries
-    `"recomposition"` and `"recomposition_info"` alongside the existing keys:
-    `"recomposition"` is `"soft"` (the affine two-target solve via
-    `hybrid_field_soft` is active -- `hgb_std_path` resolved to a real std
-    head), `"anchor"` (a film/v4 checkpoint recomposed via the plain
-    quantile-anchor -- `hgb_std_path` resolved to `None`, e.g. via `"none"`),
-    or `"sp1"` (a non-film checkpoint -- no recomposition-mode concept
-    applies at all). `"recomposition_info"` is `hybrid_field_soft`'s own
-    info dict (`a`/`m`/`target_unreachable`/`clip_frac`/
-    `total_deviation_frac`/... -- see `predict_field`'s docstring) when
-    `"recomposition"` is `"soft"`, else `None` -- there is no equivalent info
-    dict on the anchor or SP1 paths.
-    """
+                    screen_shrinkage: bool = True, gravity: tuple = (0, 0, -1),
+                    model: CastSolidModel | None = None, weights_dir: str = "weights",
+                    ood_json: str | None = None, cache_dir: str = "runs/mesh_cache",
+                    device: str | None = None) -> dict:
+    """Mesh `step_path`, predict, and write the result files.
+
+    `alloy`/`mold` name rows of `data/properties/` (default EN-GJS-400-15 in
+    green sand); `props_json` replaces both. `elem_size` (mm) meshes at a
+    fixed size instead of searching for `target_nodes`. Pass a loaded
+    `model` to reuse it across calls.
+
+    Returns {name, vtu, summary, msh, total_time_s, hotspot_xyz (5 x 3 mm),
+    ood_warnings, n_nodes, mesh_mode, elem_size, cache, recomposition_info,
+    shrinkage}."""
+    if elem_size is not None and target_nodes != DEFAULT_TARGET_NODES:
+        raise ValueError("predict_casting: give either elem_size or target_nodes, not both")
     os.makedirs(out_dir, exist_ok=True)
     name = osp.splitext(osp.basename(step_path))[0]
     out_msh = osp.join(out_dir, f"{name}.msh")
 
-    # Resolve properties early (cheap, no mesh dependency) so typos/bad alloys
-    # fail fast, before potentially lengthy meshing.
+    # Resolve properties before meshing so a bad alloy name fails fast.
     props = resolve_props(alloy=alloy, mold=mold, props_json=props_json)
+    model = model or CastSolidModel(weights_dir=weights_dir, device=device)
 
-    effective_mode = "raw" if elem_size is not None else mesh_mode
-    if target_nodes != DEFAULT_TARGET_NODES and effective_mode != "prep":
-        raise ValueError(
-            f"predict_casting: `target_nodes` ({target_nodes}) applies to the "
-            f"prep/get_or_mesh density search only, but an explicit "
-            f"elem_size={elem_size!r} forces mesh_mode='raw' -- pinning an "
-            f"element size and searching towards a node count are mutually "
-            f"exclusive requests. Drop one.")
-    gnn_ckpt, hgb_path = resolve_artifact(gnn_ckpt), resolve_artifact(hgb_path)
-    gnn, hgb, norm, cfg, baseline_norm, hgb_std = load_hybrid(
-        gnn_ckpt, hgb_path, device, hgb_std_path=hgb_std_path)
-    # Provenance label for the summary txt: the RESOLVED std artifact (the
-    # v5.1 stack is defined by it + the weights), mirroring load_hybrid's own
-    # sentinel/"none" resolution without re-running it.
-    if hgb_std is None:
-        hgb_std_path_label = "none"
-    elif hgb_std_path is _USE_DEFAULT:
-        hgb_std_path_label = resolve_artifact(DEFAULT_HGB_STD)
-    else:
-        hgb_std_path_label = str(hgb_std_path)
-    if hgb_std is not None and not cfg.get("film", False):
-        raise ValueError(
-            "predict_casting: --hgb-std requires the film-conditioned v4 "
-            "stack; the selected checkpoint is not film-capable -- rerun "
-            "without --hgb-std or with the v4 checkpoints")
-
-    if effective_mode == "raw":
-        raw_info = mesh_step(step_path, out_msh, elem_size)
-        mesh = parse_msh(out_msh)
-        minfo = {"n_nodes": raw_info["n_nodes"], "n_tets": raw_info["n_tets"],
-                "elem_size": raw_info["elem_size"], "bbox_mm": raw_info["bbox_mm"]}
-    elif effective_mode == "prep":
-        prep_info = get_or_mesh(step_path, out_msh, target_nodes=target_nodes)
-        if not prep_info.get("ok", False):
-            raise RuntimeError(
-                f"predict_casting: prep meshing pipeline (get_or_mesh) failed for "
-                f"{step_path!r}: {prep_info.get('notes', '(no notes)')}")
-        mesh = parse_msh(out_msh)
-        bbox_mm = tuple((mesh.node_xyz.max(0) - mesh.node_xyz.min(0)).tolist())
-        cache = prep_info.get("cache", "unknown")
-        minfo = {"n_nodes": prep_info["n_nodes"], "n_tets": int(mesh.tets.shape[0]),
-                "elem_size": None, "bbox_mm": bbox_mm, "cache": cache}
-    else:
-        raise ValueError(f"predict_casting: mesh_mode must be 'prep' or 'raw', got {mesh_mode!r}")
-
-    seconds, feat = predict_field(mesh, gnn, hgb, norm, cfg=cfg, device=device, tta=tta,
-                                  baseline_norm=baseline_norm, props=props,
-                                  hgb_std=hgb_std, lambda_shrink=lambda_shrink,
-                                  w_total=w_total, w_std=w_std)
+    mesh, minfo = _mesh(step_path, out_msh, elem_size, target_nodes, cache_dir)
+    seconds, details = model.predict_with_props(mesh, props)
+    X = details["X"]
 
     warnings = []
-    ood_json = resolve_artifact(ood_json)
+    ood_json = ood_json or osp.join(weights_dir, "ood_ranges.json")
     if osp.exists(ood_json):
-        ranges = _ood.load_ranges(ood_json)
-        X = feat["X"][:, :N_BASELINE_FEATURES]
-        vol = float(X[0, FEATURE_NAMES.index("volume_mm3")])
-        gcmax = float(X[:, FEATURE_NAMES.index("geodesic_d_cool")].max())
         diag = float(np.linalg.norm(mesh.node_xyz.max(0) - mesh.node_xyz.min(0)))
-        warnings = _ood.ood_flags(vol, diag, gcmax, ranges)
+        warnings = _ood.ood_flags(float(X[0, _VOL_IDX]), diag, float(X[:, _GCOOL_IDX].max()),
+                                  _ood.load_ranges(ood_json))
 
-    # SP2 shrinkage Task 4: screen the recomposed hybrid field for shrinkage-
-    # porosity risk, AFTER recomposition, BEFORE the VTU write -- see
-    # predict_casting's own docstring for the full contract.
     shrinkage = None
-    shrinkage_extra_point_data = None
+    point_data = None
     if screen_shrinkage:
-        surface_nodes, feed_nodes = _shrinkage_surface_and_feed_nodes(mesh)
+        surface_nodes, feed_nodes = _surface_and_feed_nodes(mesh)
         if surface_nodes.size == 0:
-            logger.warning(
-                "predict_casting: no surface nodes found on the meshed part -- "
-                "shrinkage screening skipped (a feeder location requires a free "
-                "surface)")
+            logger.warning("predict_casting: no surface nodes on the mesh -- "
+                           "shrinkage screening skipped")
         else:
-            alloy_family, alloy_family_assumed = resolve_alloy_family(alloy, props_json)
-            # Learned-heads Task 3: `feat["gradT"]`/`feat["dTdt"]` are only
-            # present when `predict_field` ran a deriv_heads-capable (film)
-            # checkpoint (see that function's docstring) -- `None`/`None` for
-            # every other checkpoint, which routes `screen_casting` through
-            # its own time-only proxy branch exactly as before this feature
-            # existed.
-            learned_gradT = feat.get("gradT")
-            learned_dTdt = feat.get("dTdt")
-            shrinkage_full = screen_casting(
-                pos=mesh.node_xyz, tets=mesh.tets, surface_nodes=surface_nodes,
-                t_s=seconds, TL=props["TL"], TS=props["TS"], alloy_family=alloy_family,
-                gravity_vec=gravity, feed_nodes=feed_nodes,
-                gradT=learned_gradT, dTdt=learned_dTdt)
-            shrinkage = _shrinkage_result_for_return(shrinkage_full)
-            if learned_gradT is not None and learned_dTdt is not None:
-                # `screen_casting` itself only distinguishes "exact" (real
-                # cached solver fields, e.g. Gate C's per-row comparison)
-                # from "proxy" -- it has no notion of "predicted by a
-                # checkpoint at inference time". Relabel HERE, at the
-                # inference call site, rather than teaching `screen_casting`
-                # a third source string: that keeps `src.shrinkage`'s own
-                # tests (which pass real cached gradT/dTdt and expect
-                # "exact") untouched, while giving inference-time callers an
-                # honest label distinguishing a MODEL PREDICTION from a true
-                # solver field.
-                shrinkage["niyama_source"] = "learned"
-            if alloy_family_assumed:
+            alloy_family, assumed = resolve_alloy_family(alloy, props_json)
+            screen = screen_casting(pos=mesh.node_xyz, tets=mesh.tets,
+                                    surface_nodes=surface_nodes, t_s=seconds,
+                                    TL=props["TL"], TS=props["TS"],
+                                    alloy_family=alloy_family, gravity_vec=gravity,
+                                    feed_nodes=feed_nodes)
+            shrinkage = _shrinkage_result_for_return(screen)
+            if assumed:
                 shrinkage["alloy_family_assumed"] = True
-            # This release covers the all-sand regime only, where the
-            # proxy Niyama passed its validation gate.
-            shrinkage["niyama_confidence"] = "full"
-            shrinkage_extra_point_data = _shrinkage_extra_point_data(mesh, shrinkage_full)
+            point_data = _shrinkage_point_data(mesh, screen)
 
     vtu = osp.join(out_dir, f"{name}.vtu")
-    write_vtu(vtu, mesh.node_xyz, mesh.tets, seconds, extra_point_data=shrinkage_extra_point_data)
+    write_vtu(vtu, mesh.node_xyz, mesh.tets, seconds, extra_point_data=point_data)
 
-    # hot-spot summary: top-5 latest-solidifying nodes
     order = np.argsort(seconds)[::-1][:5]
     hot_xyz = mesh.node_xyz[order]
     total = float(seconds.max())
+    info = details["recomposition"]
     summary = osp.join(out_dir, f"{name}_summary.txt")
-    alloy_label = alloy if alloy is not None else DEFAULT_ALLOY
-    if props_json is not None:
-        alloy_label = f"custom ({props_json})"
+    alloy_label = (f"custom ({props_json})" if props_json is not None
+                   else alloy if alloy is not None else DEFAULT_ALLOY)
     mold_label = mold if props_json is None else "n/a (props_json)"
-    elem_size_label = (f"{minfo['elem_size']:.2f} mm" if minfo["elem_size"] is not None
-                       else "n/a (density-searched)")
-    cache_label = minfo.get("cache") if effective_mode == "prep" else "n/a (raw mode)"
+    bbox_mm = tuple(round(float(v), 1) for v in mesh.node_xyz.max(0) - mesh.node_xyz.min(0))
+    elem_label = (f"{minfo['elem_size']:.2f} mm" if minfo["elem_size"] is not None
+                  else "n/a (density-searched)")
     with open(summary, "w") as f:
         f.write(f"casting: {name}\n")
         f.write(f"total solidification time: {total:.1f} s\n")
-        f.write(f"alloy: {alloy_label}  mold: {mold_label}  gnn: {gnn_ckpt}  hgb: {hgb_path}  "
-                f"hgb_std: {hgb_std_path_label}\n")
-        f.write(f"mesh_mode: {effective_mode}  cache: {cache_label}  nodes: {minfo['n_nodes']}  "
-                f"tets: {minfo['n_tets']}  elem_size: {elem_size_label}  bbox_mm: "
-                f"{tuple(round(v,1) for v in minfo['bbox_mm'])}\n")
-        if feat.get("recomposition") == "soft":
-            info = feat["recomposition_info"]
-            line = (f"soft recomposition: a={info['a']:.4f}, m={info['m']:.4f}, "
-                   f"target_unreachable={info['target_unreachable']}, "
-                   f"clip_frac={info['clip_frac']:.4f}")
-            if info.get("lambda_shrink", 1.0) != 1.0:
-                line += f", lambda_shrink={info['lambda_shrink']:.4f}"
-            if w_total != 1.0 or w_std != 1.0:
-                line += f", w_total={w_total:.4f}, w_std={w_std:.4f}"
-            total_dev = info.get("total_deviation_frac")
-            if total_dev is not None and abs(total_dev) > 0.001:
-                line += f", total_dev={total_dev * 100.0:+.1f}%"
-            f.write(line + "\n")
-        elif cfg.get("film", False):
-            # variance-recomposition Task 6: the film path without an active
-            # std head (hgb_std_path="none", or a film checkpoint used
-            # without one) recomposes via the quantile-anchor -- state that
-            # explicitly, mirroring the affine line above, rather than
-            # leaving the recomposition mode unstated.
-            f.write("recomposition: anchor (quantile, q=0.995)\n")
+        f.write(f"alloy: {alloy_label}  mold: {mold_label}  model: {MODEL_VERSION}\n")
+        f.write(f"mesh: {minfo['mode']}  cache: {minfo['cache']}  nodes: {minfo['n_nodes']}  "
+                f"tets: {minfo['n_tets']}  elem_size: {elem_label}  bbox_mm: {bbox_mm}\n")
+        line = (f"recomposition: a={info['a']:.4f}, m={info['m']:.4f}, "
+                f"target_unreachable={info['target_unreachable']}, "
+                f"clip_frac={info['clip_frac']:.4f}")
+        if abs(info.get("total_deviation_frac", 0.0)) > 0.001:
+            line += f", total_dev={info['total_deviation_frac'] * 100.0:+.1f}%"
+        f.write(line + "\n")
         if screen_shrinkage:
-            _write_shrinkage_summary_block(f, shrinkage)
+            _write_shrinkage_summary(f, shrinkage)
         f.write("top hot spots (x,y,z mm : t_solidus s):\n")
         for xyz_i, i in zip(hot_xyz, order):
             f.write(f"  ({xyz_i[0]:.1f},{xyz_i[1]:.1f},{xyz_i[2]:.1f}) : {seconds[i]:.1f}\n")
@@ -934,108 +232,52 @@ def predict_casting(step_path, out_dir="runs/castings", elem_size=None,
             for w in warnings:
                 f.write(f"  - {w}\n")
 
-    # soft-anchor Task 4 (v5 review M-1): surface the recomposition mode +
-    # its info dict as structured data, mirroring the summary text above --
-    # "soft" (hybrid_field_soft active, feat carries its info dict), "anchor"
-    # (film checkpoint, quantile-anchor recomposition -- no info dict exists
-    # on this path), or "sp1" (non-film checkpoint -- no recomposition
-    # concept at all).
-    if feat.get("recomposition") == "soft":
-        recomposition = "soft"
-        recomposition_info = feat["recomposition_info"]
-    elif cfg.get("film", False):
-        recomposition = "anchor"
-        recomposition_info = None
-    else:
-        recomposition = "sp1"
-        recomposition_info = None
+    return {"name": name, "vtu": vtu, "summary": summary, "msh": out_msh,
+            "total_time_s": total, "hotspot_xyz": hot_xyz, "ood_warnings": warnings,
+            "n_nodes": minfo["n_nodes"], "mesh_mode": minfo["mode"],
+            "elem_size": minfo["elem_size"], "cache": minfo["cache"],
+            "recomposition_info": info, "shrinkage": shrinkage}
 
-    result = {"name": name, "vtu": vtu, "summary": summary, "total_time_s": total,
-              "mesh_mode": effective_mode, "n_nodes": minfo["n_nodes"],
-              "elem_size": minfo["elem_size"], "ood_warnings": warnings,
-              "hotspot_xyz": hot_xyz,
-              "recomposition": recomposition,
-              "recomposition_info": recomposition_info,
-              "shrinkage": shrinkage}
-    if "cache" in minfo:
-        result["cache"] = minfo["cache"]
-    return result
+
+def _parse_gravity(s: str) -> tuple:
+    parts = [p.strip() for p in s.split(",")]
+    if len(parts) != 3:
+        raise ValueError(f"--gravity must be 'x,y,z' (3 comma-separated numbers), got {s!r}")
+    vec = tuple(float(p) for p in parts)
+    if all(v == 0.0 for v in vec):
+        raise ValueError(f"--gravity must not be the zero vector (got {s!r})")
+    return vec
 
 
-if __name__ == "__main__":
-    ap = argparse.ArgumentParser()
+def main(argv=None):
+    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
     ap.add_argument("step_path")
     ap.add_argument("--out-dir", default="runs/castings")
-    ap.add_argument("--elem-size", type=float, default=None,
-                    help="explicit uniform element size (mm) -- implies --mesh-mode raw, "
-                         "regardless of --mesh-mode")
-    ap.add_argument("--mesh-mode", choices=["prep", "raw"], default="prep",
-                    help="'prep' (default): route meshing through the same largest-solid + "
-                         "OCC heal + density-searched pipeline used to prep training/dataset-D "
-                         "parts (tools.doe.prep_parts.heal_and_mesh), served through the "
-                         "on-disk mesh cache (runs/mesh_cache/) so a repeated byte-identical "
-                         "STEP file is not re-meshed. 'raw': the old single-pass "
-                         "uniform-density src.step_mesh.mesh_step path (escape hatch). Ignored "
-                         "(forced to 'raw') when --elem-size is given.")
-    ap.add_argument("--gnn", dest="gnn_ckpt", default=DEFAULT_GNN_CKPT,
-                    help="GNN checkpoint path (SP1 or a film/v4 checkpoint)")
-    ap.add_argument("--hgb", dest="hgb_path", default=DEFAULT_HGB,
-                    help="HGB baseline path (SP1 per-node baseline or the v4 v2-total GBM)")
-    ap.add_argument("--hgb-std", dest="hgb_std_path", default=_USE_DEFAULT,
-                    help="path to an optional GBM std-head model "
-                         "(src.gbm_ft2b.refit_gbm_std's output); when active, recomposition "
-                         "uses the affine (variance-preserving) map instead of the "
-                         "quantile-anchor -- only consulted for film/v4 checkpoints. Default: "
-                         f"{DEFAULT_HGB_STD} for a film/v4 checkpoint (auto-disabled, no extra "
-                         "flag needed, for a non-film/SP1 checkpoint e.g. via --gnn/--hgb "
-                         "rollback); pass 'none' (case-insensitive) to force the old "
-                         "quantile-anchor recomposition on a film/v4 stack too")
-    ap.add_argument("--affine-lambda", dest="lambda_shrink", type=float, default=1.0,
-                    help="geometric-interpolation shrinkage of the affine spread target "
-                         "(variance-recomposition amendment 2): 1.0 (default) fully trusts "
-                         "--hgb-std's predicted spread; 0.0 targets the GNN's own spread "
-                         "(no rescaling); only used together with --hgb-std on film/v4 "
-                         "checkpoints")
-    ap.add_argument("--w-total", dest="w_total", type=float, default=DEFAULT_W_TOTAL,
-                    help="soft-anchor weight on the total-time residual in "
-                         "src.hybrid.hybrid_field_soft's objective; default "
-                         f"{DEFAULT_W_TOTAL} (the e_val-sweep winner, soft-anchor Task 6 "
-                         "promotion, 2026-07-31) -- only used together with --hgb-std")
-    ap.add_argument("--w-std", dest="w_std", type=float, default=DEFAULT_W_STD,
-                    help="soft-anchor weight on the spread residual in "
-                         "src.hybrid.hybrid_field_soft's objective; default "
-                         f"{DEFAULT_W_STD} (the e_val-sweep winner, soft-anchor Task 6 "
-                         "promotion, 2026-07-31) -- only used together with --hgb-std")
     ap.add_argument("--alloy", default=None,
-                    help=f"named alloy grade (default: {DEFAULT_ALLOY}); only used by film/v4 checkpoints")
-    ap.add_argument("--mold", default="green_sand",
-                    help="named mold system; only used by film/v4 checkpoints")
+                    help=f"alloy from data/properties/materials.csv (default {DEFAULT_ALLOY})")
+    ap.add_argument("--mold", default=DEFAULT_MOLD,
+                    help=f"mold from data/properties/molds.csv (default {DEFAULT_MOLD})")
     ap.add_argument("--props", dest="props_json", default=None,
-                    help="path to a JSON file overriding alloy/mold with raw material properties")
+                    help="JSON file with k, rho, cp, L_kJ_kg, TL, TS, Tpour, mold_effusivity "
+                         "(replaces --alloy and --mold)")
+    ap.add_argument("--elem-size", type=float, default=None,
+                    help="mesh at this fixed element size in mm instead of searching for "
+                         f"about {DEFAULT_TARGET_NODES} nodes")
+    ap.add_argument("--gravity", type=_parse_gravity, default=(0, 0, -1),
+                    help="gravity direction 'x,y,z' for sinks and feeder placement "
+                         "(default 0,0,-1)")
     ap.add_argument("--no-shrinkage", dest="screen_shrinkage", action="store_false",
-                    help="disable shrinkage-porosity screening (SP2): by default, "
-                         "predict_casting runs src.shrinkage.screen_casting on the "
-                         "recomposed field and adds niyama/porosity_risk/macro_region_id "
-                         "VTU point fields, a 'shrinkage screening:' summary block, and "
-                         "a 'shrinkage' return-dict key. This flag skips all of that, "
-                         "byte-identical to the pre-shrinkage-screening contract.")
-    ap.add_argument("--gravity", dest="gravity", type=_parse_gravity, default=(0, 0, -1),
-                    help="gravity direction as 'x,y,z' (default '0,0,-1'), forwarded to "
-                         "screen_casting's top-sink/feeder placement; ignored with "
-                         "--no-shrinkage. Parsed by _parse_gravity as argparse's own "
-                         "type= callable (review round 1, finding 5), so a malformed "
-                         "value (wrong arity, non-numeric, or the zero vector) is a "
-                         "normal argparse usage error, not a traceback.")
-    a = ap.parse_args()
-    r = predict_casting(a.step_path, out_dir=a.out_dir, elem_size=a.elem_size,
-                        mesh_mode=a.mesh_mode,
-                        gnn_ckpt=a.gnn_ckpt, hgb_path=a.hgb_path,
-                        alloy=a.alloy, mold=a.mold, props_json=a.props_json,
-                        hgb_std_path=a.hgb_std_path,
-                        lambda_shrink=a.lambda_shrink, w_total=a.w_total, w_std=a.w_std,
-                        screen_shrinkage=a.screen_shrinkage,
-                        gravity=a.gravity)
-    print(f"{r['name']}: total {r['total_time_s']:.1f}s  mesh_mode {r['mesh_mode']}  "
-          f"nodes {r['n_nodes']}  -> {r['vtu']}")
+                    help="skip the shrinkage screening")
+    ap.add_argument("--weights-dir", default="weights")
+    a = ap.parse_args(argv)
+    r = predict_casting(a.step_path, out_dir=a.out_dir, alloy=a.alloy, mold=a.mold,
+                        props_json=a.props_json, elem_size=a.elem_size,
+                        screen_shrinkage=a.screen_shrinkage, gravity=a.gravity,
+                        weights_dir=a.weights_dir)
+    print(f"{r['name']}: total {r['total_time_s']:.1f}s  nodes {r['n_nodes']}  -> {r['vtu']}")
     for w in r["ood_warnings"]:
         print("WARNING:", w)
+
+
+if __name__ == "__main__":
+    main()

src/release_io.py

@@ -1,89 +1,25 @@
-"""Public-release artifact formats (Hugging Face model repo).
+"""Loading the published weight files.
 
-The internal artifacts are pickles: the GNN checkpoint is a `torch.save`d
-`{model, norm, cfg}` dict (loaded with `weights_only=False`) and the GBM heads
-are `joblib` dumps. The public release ships them in non-executable formats:
+- GNN: one safetensors file with tensors `model.<state_dict key>` and
+  `norm.<key>`; the architecture config (`cfg`) and the container type of
+  each normalization entry are JSON strings in the safetensors metadata.
+- Gradient-boosted heads: skops files. Before loading, every type a file
+  names is checked against an allowlist (scikit-learn and NumPy only).
 
-- GNN  -> one `.safetensors` file: tensors `model.<state_dict key>` +
-  `norm.<key>`, and a JSON `cfg` plus the norm entries' original container
-  type in the safetensors metadata header. `load_gnn_state` rebuilds the
-  exact `{model, norm, cfg}` dict `torch.load` would have returned, so every
-  downstream loader (`src.evaluate_ft2b.load_ft2b_gnn`, `load_hybrid`) takes
-  it unchanged.
-- GBM  -> `.skops` (sklearn's safe serialization). Loading checks the file's
-  untrusted types against a fixed sklearn/numpy allowlist instead of
-  trusting whatever the file names.
-
-`resolve_artifact` maps the internal default paths (`runs/...`) onto the
-release layout (`weights/...`) when only the latter exists, so the public
-snapshot runs `python -m src.predict_casting part.step` with zero flags
-while the internal checkout keeps using its original files.
+Neither format runs code when loaded.
 """
 from __future__ import annotations
 
 import json
 import os
 
-import numpy as np
-import torch
-
 GNN_FORMAT = "foundry-gnn-release/1"
 
-# Internal default path -> public release path (repo-root relative).
-RELEASE_PATHS = {
-    "runs/ft3_stage1/ckpt_ep021.pt": "weights/gnn_v5.1.safetensors",
-    "runs/ft3_gbm/hgb_v3.joblib": "weights/gbm_total_v3.skops",
-    "runs/ft3_gbm/hgb_std_v2.joblib": "weights/gbm_std_v2.skops",
-    "data/features_v2/ood_ranges.json": "weights/ood_ranges.json",
-}
-
-# Types a release `.skops` GBM file may name. Anything else fails the load.
 _SKOPS_ALLOWED_PREFIXES = ("sklearn.", "numpy.")
 
 
-def resolve_artifact(path):
-    """`path` unchanged unless it is a known internal default that is absent
-    on disk while its release counterpart exists."""
-    if isinstance(path, (str, os.PathLike)):
-        p = os.fspath(path)
-        alt = RELEASE_PATHS.get(p)
-        if alt is not None and not os.path.exists(p) and os.path.exists(alt):
-            return alt
-    return path
-
-
-def _norm_to_tensor(value):
-    if isinstance(value, torch.Tensor):
-        return value.detach().cpu().contiguous(), "torch"
-    if isinstance(value, np.ndarray):
-        return torch.from_numpy(np.ascontiguousarray(value)), "numpy"
-    return None, None
-
-
-def save_gnn_release(state: dict, out_path: str, extra_metadata: dict | None = None) -> None:
-    """Write a `{model, norm, cfg}` checkpoint dict as one safetensors file."""
-    from safetensors.torch import save_file
-    tensors, norm_kinds, norm_scalars = {}, {}, {}
-    for k, v in state["model"].items():
-        tensors[f"model.{k}"] = v.detach().cpu().contiguous()
-    for k, v in state["norm"].items():
-        t, kind = _norm_to_tensor(v)
-        if t is None:
-            norm_scalars[k] = v
-        else:
-            tensors[f"norm.{k}"] = t
-            norm_kinds[k] = kind
-    meta = {"format": GNN_FORMAT,
-            "cfg": json.dumps(state["cfg"], sort_keys=True),
-            "norm_kinds": json.dumps(norm_kinds, sort_keys=True),
-            "norm_scalars": json.dumps(norm_scalars, sort_keys=True)}
-    for k, v in (extra_metadata or {}).items():
-        meta[k] = v if isinstance(v, str) else json.dumps(v)
-    save_file(tensors, out_path, metadata=meta)
-
-
 def load_gnn_state(path: str, device="cpu") -> dict:
-    """Inverse of `save_gnn_release`: the `{model, norm, cfg}` dict."""
+    """The `{model, norm, cfg}` dict stored in a GNN safetensors file."""
     from safetensors import safe_open
     with safe_open(path, framework="pt", device="cpu") as f:
         meta = f.metadata() or {}
@@ -101,31 +37,11 @@ def load_gnn_state(path: str, device="cpu") -> dict:
     return {"model": model, "norm": norm, "cfg": json.loads(meta["cfg"])}
 
 
-def save_gbm_release(model, out_path: str) -> None:
-    import skops.io as sio
-    sio.dump(model, out_path)
-
-
 def load_gbm(path):
-    """Load a GBM head: `.skops` through the type allowlist, else joblib."""
-    if path is None or not isinstance(path, (str, os.PathLike)):
-        return path
-    if os.fspath(path).endswith(".skops"):
-        import skops.io as sio
-        untrusted = sio.get_untrusted_types(file=path)
-        bad = [t for t in untrusted if not t.startswith(_SKOPS_ALLOWED_PREFIXES)]
-        if bad:
-            raise ValueError(f"{path}: refusing to load unexpected types {bad}")
-        return sio.load(path, trusted=untrusted)
-    import joblib
-    return joblib.load(path)
-
-
-def load_gnn_any(path, device):
-    """A GNN checkpoint path (`.safetensors` release or internal `.pt`) as
-    the `{model, norm, cfg}` dict; an already-loaded dict passes through."""
-    if not isinstance(path, (str, os.PathLike)):
-        return path
-    if os.fspath(path).endswith(".safetensors"):
-        return load_gnn_state(path, device)
-    return torch.load(path, map_location=device, weights_only=False)
+    """A gradient-boosted head from a skops file, after the type check."""
+    import skops.io as sio
+    untrusted = sio.get_untrusted_types(file=os.fspath(path))
+    bad = [t for t in untrusted if not t.startswith(_SKOPS_ALLOWED_PREFIXES)]
+    if bad:
+        raise ValueError(f"{path}: refusing to load unexpected types {bad}")
+    return sio.load(os.fspath(path), trusted=untrusted)

src/run_baseline.py

@@ -1,150 +0,0 @@
-"""End-to-end Phase-0 baseline: build splits, train HGB (+ ablation) and MLP,
-evaluate across splits, write a Markdown report and per-prefix MAE plot."""
-import os
-import os.path as osp
-
-import numpy as np
-import torch
-import matplotlib
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-
-from src.dataload import load_matrix, load_norm, standardize
-from src.baseline import train_hgb, train_mlp, predict_seconds
-from src.splits import build_splits, save_splits, load_splits
-from src import metrics
-
-
-def _xyz_mm(features_dir, data_root, stem, n):
-    """Per-node positions in mm. Prefer the feature cache (v2 stores `pos` in mm);
-    fall back to the raw .pt (metres) when the cache lacks positions."""
-    gpath = osp.join(features_dir, "graphs", f"{stem}.npz")
-    d = np.load(gpath)
-    if "pos" in d.files:
-        return d["pos"][:n]                       # already mm
-    raw = torch.load(osp.join(data_root, "pt", f"{stem}.pt"), weights_only=False)
-    return raw["pos"].numpy()[:n] * 1000.0
-
-
-def evaluate_model(model, features_dir, test_stems, norm, data_root="data"):
-    X, y, gid = load_matrix(features_dir, test_stems)
-    # Both model types are trained on standardized features (see run()), so
-    # standardize uniformly here too.
-    Xs = standardize(X, norm)
-    pred = predict_seconds(model, Xs, norm)
-    xyz = np.concatenate([_xyz_mm(features_dir, data_root, s, (gid == i).sum())
-                          for i, s in enumerate(test_stems)])
-    return {
-        "mae": metrics.mae(y, pred),
-        "rmse": metrics.rmse(y, pred),
-        "median_rel_error": metrics.median_rel_error(y, pred),
-        "pred_std_ratio": metrics.pred_std_ratio(y, pred, gid),
-        "hotspot_accuracy": metrics.hotspot_accuracy(y, pred, xyz, gid),
-    }
-
-
-def per_prefix_mae(model, features_dir, test_stems, norm):
-    by = {}
-    for s in test_stems:
-        X, y, _ = load_matrix(features_dir, [s])
-        pred = predict_seconds(model, standardize(X, norm), norm)
-        by.setdefault(s.split("_")[0], []).append(metrics.mae(y, pred))
-    return {k: float(np.mean(v)) for k, v in sorted(by.items())}
-
-
-def _balanced(train_stems, meta_lookup):
-    """Config C: downsample so each n_adiab regime has ~equal stem count."""
-    from collections import defaultdict
-    import random
-    groups = defaultdict(list)
-    for s in train_stems:
-        groups[meta_lookup[s]].append(s)
-    if not groups:
-        return train_stems
-    k = min(len(v) for v in groups.values())
-    rng = random.Random(0)
-    out = []
-    for v in groups.values():
-        out.extend(rng.sample(v, k))
-    return out
-
-
-def run(features_dir="data/features", splits_path=None, out_dir="runs/baseline",
-        ablation=True, max_nodes_per_graph=2000, data_root="data"):
-    os.makedirs(out_dir, exist_ok=True)
-    splits_path = splits_path or osp.join(features_dir, "splits.json")
-    meta = np.load(osp.join(features_dir, "meta.npz"), allow_pickle=True)
-    if not osp.exists(splits_path):
-        save_splits(build_splits(meta), splits_path)
-    sp = load_splits(splits_path)
-    norm = load_norm(features_dir)
-    meta_nad = {str(s): int(n) for s, n in zip(meta["stem"], meta["n_adiab"])}
-    fully = {str(s) for s, f in zip(meta["stem"], meta["is_fully_cooled"]) if f}
-
-    configs = {"A_all": sp["interpolation"]["train"]}
-    if ablation:
-        configs["B_cooled_only"] = [s for s in sp["interpolation"]["train"] if s in fully]
-        configs["C_balanced"] = _balanced(sp["interpolation"]["train"], meta_nad)
-
-    eval_sets = {"interpolation": sp["interpolation"]["test"],
-                 "family": sp["family"]["test"],
-                 "fully_cooled": sp["fully_cooled"]["test"]}
-
-    lines = ["# Phase 0 Baseline Report\n"]
-    results = {}
-    for cfg, train_stems in configs.items():
-        if not train_stems:
-            print(f"  skipping config {cfg}: no training stems")
-            continue
-        Xtr, ytr, _ = load_matrix(features_dir, train_stems,
-                                  max_nodes_per_graph=max_nodes_per_graph)
-        model = train_hgb(standardize(Xtr, norm), np.log1p(ytr))
-        lines.append(f"\n## HGB — config {cfg}  (train graphs: {len(train_stems)})\n")
-        lines.append("| split | MAE(s) | RMSE(s) | med.rel | pred/GT std | hotspot |")
-        lines.append("|---|---|---|---|---|---|")
-        for name, stems in eval_sets.items():
-            if not stems:
-                continue
-            r = evaluate_model(model, features_dir, stems, norm, data_root)
-            results[(cfg, name)] = r
-            lines.append(f"| {name} | {r['mae']:.2f} | {r['rmse']:.2f} | "
-                         f"{r['median_rel_error']:.3f} | {r['pred_std_ratio']:.3f} | "
-                         f"{r['hotspot_accuracy']:.3f} |")
-        if cfg == "A_all":
-            import joblib
-            joblib.dump(model, osp.join(out_dir, "hgb_A.joblib"))
-            # per-prefix plot on interpolation test
-            pmae = per_prefix_mae(model, features_dir, eval_sets["interpolation"], norm)
-            plt.figure(figsize=(10, 4))
-            plt.bar(list(pmae.keys()), list(pmae.values()))
-            plt.ylabel("MAE (s)"); plt.xticks(rotation=90); plt.tight_layout()
-            plt.savefig(osp.join(out_dir, "per_prefix_mae.png"), dpi=120)
-            plt.close()
-
-    # MLP on config A
-    if not configs["A_all"]:
-        print(f"  skipping MLP training: config A_all has no training stems")
-    else:
-        Xtr, ytr, _ = load_matrix(features_dir, configs["A_all"],
-                                  max_nodes_per_graph=max_nodes_per_graph)
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        mlp = train_mlp(standardize(Xtr, norm), np.log1p(ytr), device=device)
-        lines.append("\n## MLP — config A_all\n")
-        lines.append("| split | MAE(s) | RMSE(s) | med.rel | pred/GT std | hotspot |")
-        lines.append("|---|---|---|---|---|---|")
-        for name, stems in eval_sets.items():
-            if not stems:
-                continue
-            r = evaluate_model(mlp, features_dir, stems, norm, data_root)
-            lines.append(f"| {name} | {r['mae']:.2f} | {r['rmse']:.2f} | "
-                         f"{r['median_rel_error']:.3f} | {r['pred_std_ratio']:.3f} | "
-                         f"{r['hotspot_accuracy']:.3f} |")
-
-    with open(osp.join(out_dir, "report.md"), "w") as f:
-        f.write("\n".join(lines) + "\n")
-    print(f"Report → {osp.join(out_dir, 'report.md')}")
-    return results
-
-
-if __name__ == "__main__":
-    run()

src/select_ckpt_ft2b.py

@@ -1,331 +0,0 @@
-"""FT2b Task 7: gate-metric checkpoint selector.
-
-`src.train_gnn.train(..., save_every_epoch=True)` (Task 6) writes one
-checkpoint per epoch (`<out_dir>/ckpt_ep{ep:03d}.pt`, payload
-`{model, norm, cfg}`). Picking among them by validation loss alone can pick a
-variance-collapsed checkpoint (near-constant predictions can still have low
-smooth-L1 loss on a small/skewed val set) -- exactly the failure mode the
-FT2b gate metrics (`src.evaluate_ft2b`'s `std_ratio`/`hotspot_overlap`) exist
-to catch. This module re-ranks the per-epoch checkpoints by those SAME gate
-metrics instead.
-
-=== Why no GBM is needed here (key insight from the Task 7 brief) ===
-`src.evaluate_ft2b._max_anchor_hybrid` recomposes the hybrid prediction as
-`pred_i = gnn_i - anchor + gbm_total`, i.e. a per-graph ADDITIVE SHIFT of
-the GNN field by a constant (`gbm_total - anchor`, the same shift for every
-node in that graph), where `anchor = np.quantile(gnn, q)` (default
-`q=0.995`, a high quantile rather than the exact max -- see
-`src.hybrid.hybrid_field_max_anchor`'s docstring for why: the exact-max
-anchor is fragile to a handful of outlier nodes). Both metrics this
-selector ranks by are shift-invariant per graph regardless of which
-quantile the anchor uses:
-  - `std_ratio = std(pred)/std(true)` (`src.metrics.pred_std_ratio`): std of a
-    per-graph-shifted array is unchanged by the shift.
-  - `hotspot_overlap` (`src.metrics.hotspot_accuracy`): picks the top-`k`
-    LATEST-solidifying nodes by rank; a per-graph constant shift does not
-    reorder nodes within a graph, so the top-k set (and thus the overlap) is
-    unchanged.
-So GNN-only seconds and the eventual hybrid produce IDENTICAL std_ratio/
-hotspot_overlap numbers -- this selector never needs the GBM (which also
-means it never needs the DOE manifest's alloy/BC feature-row resolution that
-`src.gbm_ft2b._row_features` performs; only `alloy_family`/`bc_config` for
-grouping are read from the manifest).
-
-=== Reuse, not reimplementation ===
-Everything gate-metric-shaped is imported from `src.evaluate_ft2b`, which is
-already exercised by `tests/test_evaluate_ft2b.py`:
-  - `load_ft2b_gnn` -- checkpoint -> `(model, cfg, norm)`, film-aware.
-  - `predict_graph_seconds` -- one graph's cached `X`/`pos`/`mesh_edges`/`g`
-    -> physical-seconds GNN prediction (the ckpt's own norm-inversion chain).
-  - `_group_metrics` / `_overall_metrics` -- the EXACT per-group (and
-    overall) `{mae, hotspot_overlap, std_ratio}` computation `evaluate_ft2b`
-    uses for its own `eval.json` (pool all nodes of every stem in the group,
-    build a per-stem `graph_id` array, hand the pooled arrays to
-    `src.metrics.pred_std_ratio`/`hotspot_overlap`). This selector builds the
-    same `{"family", "bc", "y_true", "y_hybrid", "xyz"}`-shaped records (with
-    `y_hybrid` set to the GNN-only seconds, per the shift-invariance above)
-    and calls those two functions UNCHANGED, so its numbers are directly
-    comparable to `evaluate_ft2b`'s `eval.json` for the same stems/ckpt.
-  - `src.gbm_ft2b._load_manifest_rows` -- manifest CSV -> `{run_id: row}`.
-
-=== Ranking (fix pass -- two-sided calibration, not one-sided "biggest std") ===
-For each checkpoint, pool `per_alloy_family` and `per_bc_config` group
-metrics (`bc_config UNION alloy_family`, i.e. every group from BOTH
-partitions, not their cross product) into one flat collection, then rank
-checkpoints by, in order:
-  1. FEWEST of those groups with `std_ratio < std_ratio_threshold` (default
-     0.9, matching `evaluate_ft2b`'s own `variance_collapse_threshold`
-     default) -- fewer variance-collapsed groups is better.
-  2. SMALLEST `calib = max over groups of abs(log(std_ratio))` -- a
-     TWO-SIDED mis-calibration score. `std_ratio` measures the same
-     miscalibration whether it collapses (std_ratio << 1) or over-disperses
-     (std_ratio >> 1); `log` makes the two directions symmetric
-     (`log(0.5) == -log(2)`) so neither a collapsed nor an inflated group can
-     win by having a numerically larger raw ratio.
-     ORIGINAL (WRONG) CRITERION, replaced here: "largest minimum group
-     std_ratio" -- that criterion is UNBOUNDED ABOVE and actively rewards
-     over-dispersion (a checkpoint where every group is 2.3x over-dispersed
-     beats one where every group is a well-calibrated ~0.95, because 2.3 >
-     0.95 as a raw ratio). Confirmed on a real 64-checkpoint run: the old
-     criterion picked ep019 (min_group_std_ratio=2.34, EVERY group
-     over-dispersed) over ep063 (all groups >= 0.9, calibrated). `calib`
-     fixes this: ep019's `calib = abs(log(2.34)) ~= 0.85`; a calibrated
-     ep063 with all ratios near 1 has `calib` near 0 and wins.
-  3. HIGHEST overall `hotspot_overlap` (pooled over every given stem).
-  4. Deterministic tie-break: LOWEST epoch number (prefers the earlier,
-     cheaper-to-have-trained checkpoint when the gate metrics are tied).
-
-`min_group_std_ratio`/`max_group_std_ratio` are still computed and stored on
-every ranking row (for human inspection / plotting the calibration spread)
--- they are just no longer part of the sort key beyond criterion (1)'s
-threshold count.
-
-=== `select_ckpt` vs `rerank_selection` ===
-`select_ckpt` runs the (slow -- GPU forward pass per checkpoint) evaluation
-and produces `selection.json`. `rerank_selection` re-applies the ranking
-rule above to an ALREADY-WRITTEN `selection.json` (the per-group numbers for
-every checkpoint are already in that file) with NO forward passes --
-needed because a real run over many checkpoints takes hours, and a ranking
-bugfix must not force re-running it. Both route through the same
-`_rank_entries` so the two can never silently diverge.
-"""
-from __future__ import annotations
-
-import glob
-import json
-import math
-import os
-import os.path as osp
-import re
-
-import numpy as np
-import torch
-
-from src.evaluate_ft2b import (_group_metrics, _overall_metrics,
-                               load_ft2b_gnn, predict_graph_seconds)
-from src.gbm_ft2b import _load_manifest_rows
-
-DEFAULT_STD_RATIO_THRESHOLD = 0.9
-
-# Floor applied to std_ratio before taking log(), so a fully-collapsed group
-# (std_ratio == 0.0 exactly) yields a large-but-finite `calib` instead of
-# `inf` (which `json.dump` would emit as the non-standard token `Infinity`).
-_CALIB_STD_RATIO_FLOOR = 1e-6
-
-_CKPT_EPOCH_RE = re.compile(r"ckpt_ep(\d+)\.pt$")
-
-
-def _epoch_of(ckpt_path: str) -> int:
-    m = _CKPT_EPOCH_RE.search(osp.basename(ckpt_path))
-    if not m:
-        raise ValueError(f"not a ckpt_ep###.pt path: {ckpt_path!r}")
-    return int(m.group(1))
-
-
-def _list_ckpts(ckpt_dir: str) -> list[str]:
-    paths = glob.glob(osp.join(ckpt_dir, "ckpt_ep*.pt"))
-    return sorted(paths, key=_epoch_of)
-
-
-def _default_stems(features_dir: str) -> list[str]:
-    with open(osp.join(features_dir, "splits.json")) as f:
-        splits_d = json.load(f)
-    return list(splits_d["extrap"]["test"])
-
-
-def _gnn_only_records(model, cfg, norm, stems, manifest_rows, features_dir,
-                      device) -> list[dict]:
-    """One record per stem: `{family, bc, y_true, y_hybrid, xyz}`, shaped
-    exactly for `src.evaluate_ft2b._group_metrics`/`_overall_metrics` --
-    `y_hybrid` here holds the GNN-ONLY seconds prediction (see this module's
-    docstring for why that's metric-equivalent to the eventual hybrid)."""
-    records = []
-    for stem in stems:
-        d = np.load(osp.join(features_dir, "graphs", f"{stem}.npz"))
-        X, pos_np = d["X"], d["pos"]
-        y_true = d["y_raw"] if "y_raw" in d.files else d["y"]
-        mesh_edges = d["mesh_edges"] if "mesh_edges" in d.files else None
-        g_vec = d["g"] if "g" in d.files else None
-        gnn_seconds = predict_graph_seconds(model, cfg, norm, X, pos_np, mesh_edges,
-                                            g_vec, device=device)
-        row = manifest_rows.get(stem, {})
-        records.append({
-            "stem": stem,
-            "family": row.get("alloy_family") or "unknown",
-            "bc": row.get("bc_config") or "none",
-            "y_true": y_true, "y_hybrid": gnn_seconds, "xyz": pos_np,
-        })
-    return records
-
-
-def _combined_groups(per_family: dict, per_bc: dict) -> dict:
-    """Union (not cross product) of the two partitions' groups, disambiguated
-    by a `"family:"`/`"bc:"` prefix so same-named groups in each partition
-    (e.g. a `bc_config` and an `alloy_family` both happening to be `"none"`)
-    don't collide."""
-    out = {}
-    for k, v in per_family.items():
-        out[f"family:{k}"] = v
-    for k, v in per_bc.items():
-        out[f"bc:{k}"] = v
-    return out
-
-
-def _calib_score(std_ratio: float) -> float:
-    """Two-sided mis-calibration score for one group's `std_ratio`.
-
-    `abs(log(std_ratio))`, floored at `_CALIB_STD_RATIO_FLOOR` before the log
-    so an exact-zero (fully collapsed) ratio gives a large finite number
-    instead of `inf`. 0 at perfect calibration (`std_ratio == 1`), symmetric
-    between collapse and over-dispersion (`calib(0.5) == calib(2.0)`)."""
-    r = max(float(std_ratio), _CALIB_STD_RATIO_FLOOR)
-    return abs(math.log(r))
-
-
-def _rank_entries(entries: list[dict], std_ratio_threshold: float) -> list[dict]:
-    """Shared scoring + sort for both `select_ckpt` and `rerank_selection`.
-
-    `entries`: one dict per checkpoint, each with (at least)
-    `{"ckpt", "epoch", "overall", "per_alloy_family", "per_bc_config"}` --
-    exactly what both callers already have (freshly computed, or loaded back
-    out of a previously-written `selection.json`). Recomputes the derived
-    scoring fields (`n_fail_groups`/`min_group_std_ratio`/
-    `max_group_std_ratio`/`calib`/`overall_hotspot`) from the stored
-    per-group numbers -- no forward pass, no re-evaluation -- and returns a
-    NEW list of entries (original dicts are not mutated), sorted best-first
-    per this module's ranking rule (see module docstring)."""
-    scored = []
-    for e in entries:
-        combined = _combined_groups(e["per_alloy_family"], e["per_bc_config"])
-        std_ratios = [m["std_ratio"] for m in combined.values()]
-        n_fail_groups = sum(1 for r in std_ratios if r < std_ratio_threshold)
-        min_group_std_ratio = float(min(std_ratios)) if std_ratios else 0.0
-        max_group_std_ratio = float(max(std_ratios)) if std_ratios else 0.0
-        calib = max((_calib_score(r) for r in std_ratios), default=0.0)
-        overall_hotspot = float(e["overall"]["hotspot_overlap"])
-
-        new_entry = dict(e)
-        new_entry.update({
-            "n_fail_groups": int(n_fail_groups),
-            "min_group_std_ratio": min_group_std_ratio,
-            "max_group_std_ratio": max_group_std_ratio,
-            "calib": float(calib),
-            "overall_hotspot": overall_hotspot,
-        })
-        scored.append(new_entry)
-
-    scored.sort(key=lambda r: (r["n_fail_groups"], r["calib"],
-                               -r["overall_hotspot"], r["epoch"]))
-    return scored
-
-
-def rerank_selection(json_path, out_json=None,
-                     std_ratio_threshold: float = DEFAULT_STD_RATIO_THRESHOLD) -> dict:
-    """Re-rank an ALREADY-WRITTEN `select_ckpt` output (`json_path`) with the
-    current ranking rule, WITHOUT re-running any GNN forward pass -- every
-    checkpoint's per-group numbers are already stored in `json_path`'s
-    `"ranking"` entries (`per_alloy_family`/`per_bc_config`/`overall`), so
-    this only recomputes the derived scoring fields and re-sorts. Exists so a
-    ranking-rule bugfix never forces re-running a real (hours-long,
-    many-checkpoint) `select_ckpt` call.
-
-    Args:
-        json_path: path to a JSON file shaped like `select_ckpt`'s return
-            value (`{"best", "ranking": [...]}`) -- typically its own
-            previous `out_json` output.
-        out_json: path to write the re-ranked result (falsy skips writing).
-        std_ratio_threshold: same meaning as `select_ckpt`'s.
-
-    Returns: `{"best": <ckpt path>, "ranking": [...]}`, re-sorted.
-    """
-    with open(json_path) as f:
-        data = json.load(f)
-    ranking = _rank_entries(data["ranking"], std_ratio_threshold)
-    result = {"best": ranking[0]["ckpt"], "ranking": ranking}
-
-    if out_json:
-        out_dir = osp.dirname(out_json)
-        if out_dir:
-            os.makedirs(out_dir, exist_ok=True)
-        with open(out_json, "w") as f:
-            json.dump(result, f, indent=2)
-
-    return result
-
-
-def select_ckpt(ckpt_dir, features_dir, manifest, stems=None, out_json=None,
-                device=None, std_ratio_threshold: float = DEFAULT_STD_RATIO_THRESHOLD) -> dict:
-    """Rank every `<ckpt_dir>/ckpt_ep*.pt` checkpoint by the gate metrics
-    (per-group `std_ratio` calibration + `hotspot_overlap`), NOT val loss.
-
-    Args:
-        ckpt_dir: directory containing `ckpt_ep{ep:03d}.pt` files (Task 6's
-            `train(..., save_every_epoch=True)` output).
-        features_dir: v4 feature cache directory (`graphs/<stem>.npz` +
-            `splits.json`, used only when `stems` is None).
-        manifest: path to the DOE manifest CSV (`run_id`==stem; needs
-            `alloy_family`/`bc_config` columns).
-        stems: stems to evaluate on. Default (None): `splits.json`'s
-            `extrap.test` list (the OOD-generalization split -- the one the
-            gate metrics matter most on).
-        out_json: path to write the ranked result as JSON (falsy skips
-            writing).
-        device: "cuda"/"cpu"; default picks cuda if available.
-        std_ratio_threshold: a group's `std_ratio` below this counts as
-            "failing" for ranking criterion (1). Default 0.9, matching
-            `evaluate_ft2b`'s own `variance_collapse_threshold` default.
-
-    Returns: `{"best": <ckpt path>, "ranking": [<per-ckpt dict>, ...]}`,
-        `ranking` sorted best-first per this module's docstring's 4-step
-        ranking rule. Each ranking entry also carries `per_alloy_family`/
-        `per_bc_config`/`overall` (same shape as `evaluate_ft2b`'s `eval.json`
-        groups) so the numbers are directly inspectable/comparable.
-    """
-    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-    ckpts = _list_ckpts(ckpt_dir)
-    if not ckpts:
-        raise FileNotFoundError(f"no ckpt_ep*.pt files found in {ckpt_dir!r}")
-
-    if stems is None:
-        stems = _default_stems(features_dir)
-    manifest_rows = _load_manifest_rows(manifest)
-
-    raw_entries = []
-    n = len(ckpts)
-    for i, ckpt_path in enumerate(ckpts):
-        epoch = _epoch_of(ckpt_path)
-        print(f"[select_ckpt] ({i + 1}/{n}) evaluating {osp.basename(ckpt_path)} "
-              f"on {len(stems)} stems ...", flush=True)
-        model, cfg, norm = load_ft2b_gnn(ckpt_path, device)
-        records = _gnn_only_records(model, cfg, norm, stems, manifest_rows,
-                                    features_dir, device)
-        entry = {
-            "ckpt": ckpt_path,
-            "epoch": epoch,
-            "overall": _overall_metrics(records),
-            "per_alloy_family": _group_metrics(records, "family"),
-            "per_bc_config": _group_metrics(records, "bc"),
-        }
-        # Per-entry pass through the shared scorer just for the progress
-        # print -- `select_ckpt`'s own final ranking below re-derives this
-        # (identically) via `_rank_entries` over the whole batch, so the two
-        # can never disagree.
-        scored = _rank_entries([entry], std_ratio_threshold)[0]
-        print(f"[select_ckpt]   epoch {epoch:3d}: n_fail_groups={scored['n_fail_groups']} "
-              f"calib={scored['calib']:.4f} "
-              f"min_group_std_ratio={scored['min_group_std_ratio']:.4f} "
-              f"max_group_std_ratio={scored['max_group_std_ratio']:.4f} "
-              f"overall_hotspot={scored['overall_hotspot']:.4f}", flush=True)
-        raw_entries.append(entry)
-
-    ranking = _rank_entries(raw_entries, std_ratio_threshold)
-    result = {"best": ranking[0]["ckpt"], "ranking": ranking}
-
-    if out_json:
-        out_dir = osp.dirname(out_json)
-        if out_dir:
-            os.makedirs(out_dir, exist_ok=True)
-        with open(out_json, "w") as f:
-            json.dump(result, f, indent=2)
-
-    return result

src/shrinkage.py

@@ -1,37 +1,20 @@
-"""Field derivatives and the Niyama shrinkage-porosity criterion on tet meshes.
-
-Pure numpy/scipy: no project imports. Mesh arrays (`pos` in mm, `tets` as
-int node-index quadruples) are passed in by the caller -- this module knows
-nothing about `.msh` parsing, features, or the GNN.
-
-Unit conventions (mirroring the cached raw-solver fields, see
-`docs/superpowers/specs/2026-07-24-shrinkage-porosity-screening-design.md`,
-ADDENDUM section "GNN-path derived fields"):
-  - `pos` -- node coordinates in **mm**.
-  - `gradT` -- thermal-gradient magnitude in **K/m** (as cached by the FEniCSx
-    solver, see `tools/doe/adapters.py`).
-  - `dTdt` -- cooling-rate in **K/s** (as cached).
-  - `t_s` / solidification time -- **s**.
-  - Niyama number `Ny` -- reported in the literature unit
-    **(degC*s)^0.5 / mm** (`G[degC/mm] / sqrt(|Tdot|[degC/s])`); note 1 K
-    of *difference* equals 1 degC of difference, so no additive offset is
-    needed when converting a gradient/rate, only the K/m -> degC/mm scale
-    (divide by 1000).
-
-Two ways to obtain Ny, per the ADDENDUM:
-  - **Exact path** (training/validation rows with cached `gradT`/`dTdt`):
-    `niyama_exact` computes Ny directly from the solver's own gradient and
-    cooling-rate fields, and is compared against literature absolute
-    thresholds (`NY_CRIT`, `porosity_risk(mode="absolute")`).
-  - **Proxy path** (GNN-inference rows, where only the predicted
-    solidification-time field `t_s` exists): the exact identity
-    `grad(T) = -Tdot * grad(t_s)` (differentiate `T(x, t_s(x)) = T_S`) gives
-    `G = |Tdot| * |grad(t_s)|` and `Ny = G / sqrt(|Tdot|) = sqrt(|Tdot|) *
-    |grad(t_s)|`. `Tdot` itself is not available on this path and is proxied
-    by the characteristic freezing-interval rate `(TL - TS) / max(t_s, eps)`.
-    `niyama_proxy` implements this. Because it is a proxy, it is only a
-    *relative* risk ranking within one part -- use
-    `porosity_risk(mode="percentile")`, never the absolute thresholds.
+"""Shrinkage screening on a predicted solidification-time field.
+
+Pure numpy/scipy, working on node positions `pos` (mm), tetrahedra `tets` and
+the per-node solidification time `t_s` (s). Four outputs:
+
+- A Niyama-style criterion. With only `t_s` available, the identity
+  grad(T) = -Tdot * grad(t_s) gives Ny = sqrt(|Tdot|) * |grad(t_s)|, and the
+  cooling rate is approximated by the freezing-interval rate
+  (TL - TS) / t_s. The result is a relative ranking within one part, so the
+  risk mask flags its lowest decile rather than applying absolute
+  literature thresholds.
+- Macro-shrinkage regions: connected groups of the latest-solidifying nodes
+  that no feed source can reach through metal that is still liquid.
+- Sinks: regions whose highest point (against gravity) is on the free
+  surface, which would draw the surface in rather than void internally.
+- Feeder suggestions: a surface point above each region and the minimum
+  feeder modulus (1.2 x the region's V/A, Chvorinov's rule).
 """
 import logging
 
@@ -41,12 +24,9 @@ from scipy.sparse.csgraph import connected_components
 
 logger = logging.getLogger(__name__)
 
-# Shrinkage percentage per alloy family (volumetric solidification shrinkage,
-# used as severity = shrinkage_pct/100 * region volume). Grey iron's graphite
-# expansion during eutectic solidification largely offsets liquid/solid
-# shrinkage, hence 0.0 -- macro-shrinkage regions are still detected/reported
-# for grey iron (the geometry-only isolation criterion doesn't depend on
-# alloy), just with zero severity.
+# Volumetric solidification shrinkage per alloy family, in percent. Severity
+# of a region = shrinkage_pct / 100 * region volume. Grey iron's graphite
+# expansion offsets its shrinkage, hence 0 (regions are still reported).
 SHRINKAGE_PCT = {
     "steel": 2.5,
     "low_alloy_steel": 2.5,
@@ -54,45 +34,16 @@ SHRINKAGE_PCT = {
     "grey_iron": 0.0,
 }
 
-# Absolute Niyama-criterion thresholds (exact path only), (degC*s)^0.5/mm.
-# Below this value -> flagged as at-risk of micro-shrinkage porosity.
-NY_CRIT = {
-    "steel": 1.0,
-    "low_alloy_steel": 1.0,
-    "grey_iron": 0.75,
-    "ductile_iron": 0.75,
-}
-
-# Small epsilon guards against division-by-zero / sqrt(0) without perturbing
-# any physically meaningful value (cooling rates and solidification times of
-# interest are always many orders of magnitude larger than this).
+# Guards divisions and square roots without moving any physical value.
 _EPS = 1e-9
 
 
 def node_gradient(pos: np.ndarray, tets: np.ndarray, field: np.ndarray) -> np.ndarray:
-    """(N,3) volume-weighted average of per-tet linear-shape-function gradients.
-
-    For a linear (P1) tetrahedral element the interpolated field is affine
-    over the tet, so its gradient is constant per tet. For tet vertices
-    x0..x3 with field values f0..f3, the gradient g solves the 3x3 system
-        J @ g = rhs,  J[i] = x_{i+1} - x0,  rhs[i] = f_{i+1} - f0   (i=0,1,2)
-    i.e. J's rows are the tet's three edge vectors from vertex 0, and rhs is
-    the corresponding field differences -- the standard closed form for a
-    linear tet element's gradient (no shape-function basis needed beyond
-    this, since a P1 gradient is uniquely determined by 3 independent edge
-    derivatives).
-
-    The constant per-tet gradient is then averaged onto each of its 4 nodes,
-    weighted by tet volume (a volume-weighted average over all tets incident
-    to a node), which is the standard nodal-recovery scheme for
-    element-constant gradients on unstructured tet meshes.
-
-    Units: `pos` in mm; the returned gradient is in [field units] / mm.
-
-    Degenerate (zero-volume) tets contribute zero weight and zero gradient
-    (guarded, never divide by zero); a node touched by no tets (or an empty
-    `tets` array) gets a zero gradient.
-    """
+    """(N, 3) nodal gradient of a P1 field on a tet mesh, in [field units]/mm.
+
+    The gradient is constant per tet (solve J g = df with J the tet's edge
+    vectors from vertex 0), then averaged onto the nodes weighted by tet
+    volume. Degenerate tets contribute nothing; isolated nodes get zero."""
     pos = np.asarray(pos, dtype=np.float64)
     tets = np.asarray(tets, dtype=np.int64)
     field = np.asarray(field, dtype=np.float64)
@@ -135,37 +86,12 @@ def node_gradient(pos: np.ndarray, tets: np.ndarray, field: np.ndarray) -> np.nd
     return node_grad_sum / np.maximum(node_vol_sum, _EPS)[:, None]
 
 
-def niyama_exact(gradT_K_per_m: np.ndarray, dTdt_K_per_s: np.ndarray) -> np.ndarray:
-    """Ny in (degC*s)^0.5/mm: G[degC/mm] / sqrt(|Tdot|[degC/s]), eps-guarded.
-
-    Exact path: uses the solver's own cached `gradT` (K/m) and `dTdt` (K/s)
-    fields directly (see module docstring). 1 K/m = 1e-3 degC/mm (a
-    difference/gradient, not an absolute temperature, so no offset).
-    """
-    gradT_K_per_m = np.asarray(gradT_K_per_m, dtype=np.float64)
-    dTdt_K_per_s = np.asarray(dTdt_K_per_s, dtype=np.float64)
-
-    G_degC_per_mm = gradT_K_per_m / 1000.0
-    tdot_abs = np.abs(dTdt_K_per_s)
-    return G_degC_per_mm / np.sqrt(np.maximum(tdot_abs, _EPS))
-
-
 def niyama_proxy(pos: np.ndarray, tets: np.ndarray, t_s: np.ndarray,
                  TL: float, TS: float) -> np.ndarray:
-    """sqrt((TL-TS)/max(t_s,eps)) * |node_gradient(t_s)| -- the ADDENDUM contract.
-
-    Proxy path: only the (predicted) solidification-time field `t_s` (s) is
-    available. Using the identity grad(T) = -Tdot * grad(t_s), a proxy
-    cooling rate `Tdot_proxy = (TL - TS) / max(t_s, eps)` (the characteristic
-    freezing-interval rate) stands in for the true `Tdot`, giving
-    `Ny_proxy = sqrt(Tdot_proxy) * |grad(t_s)|`. `TL`, `TS` in degC; `t_s` in
-    s; `pos` in mm (so `grad(t_s)` is in s/mm, matching the exact path's
-    (degC*s)^0.5/mm unit once multiplied by sqrt(degC/s)).
-
-    This is a *relative* ranking only, not comparable to the literature
-    `NY_CRIT` absolute thresholds -- flag risk within one part via
-    `porosity_risk(mode="percentile")`.
-    """
+    """sqrt((TL - TS) / max(t_s, eps)) * |grad(t_s)|, per node.
+
+    `TL`, `TS` in deg C, `t_s` in s, `pos` in mm. A relative ranking within
+    one part; see the module docstring."""
     t_s = np.asarray(t_s, dtype=np.float64)
     grad = node_gradient(pos, tets, t_s)
     grad_mag = np.linalg.norm(grad, axis=1)
@@ -175,76 +101,21 @@ def niyama_proxy(pos: np.ndarray, tets: np.ndarray, t_s: np.ndarray,
     return np.sqrt(np.maximum(tdot_proxy, 0.0)) * grad_mag
 
 
-def porosity_risk(ny: np.ndarray, mode: str, alloy_family: str = None,
-                  pct: float = 10.0) -> np.ndarray:
-    """Boolean at-risk mask from a Niyama field.
-
-    mode="absolute" (exact path): `ny < NY_CRIT[alloy_family]`. `alloy_family`
-    must be a key of `NY_CRIT`; an unknown or missing (`None`) family raises
-    `KeyError` -- there is no silent default, because guessing an alloy
-    family would silently mis-threshold the risk call.
-
-    mode="percentile" (proxy path): flags the within-part bottom-`pct`%
-    lowest Niyama values, i.e. `ny <= numpy.percentile(ny, pct)` (numpy's
-    default linear-interpolation percentile). Tie handling is deterministic
-    and order-independent: every value tied at (or below) the threshold is
-    flagged, so if duplicate values straddle the threshold the flagged
-    fraction can be slightly *more* than `pct`% (never less, and never
-    dependent on array order).
-    """
+def porosity_risk(ny: np.ndarray, pct: float = 10.0) -> np.ndarray:
+    """Boolean mask of the within-part lowest `pct`% Niyama values
+    (`ny <= percentile(ny, pct)`; ties at the threshold are all flagged)."""
     ny = np.asarray(ny, dtype=np.float64)
+    threshold = np.percentile(ny, pct)
+    return ny <= threshold
 
-    if mode == "absolute":
-        crit = NY_CRIT[alloy_family]
-        return ny < crit
-    elif mode == "percentile":
-        threshold = np.percentile(ny, pct)
-        return ny <= threshold
-    else:
-        raise ValueError(f"unknown porosity_risk mode: {mode!r}")
-
-
-# ---------------------------------------------------------------------------
-# Macro-shrinkage regions, top-surface sinks, feeder recommendations.
-#
-# See docs/superpowers/specs/2026-07-24-shrinkage-porosity-screening-design.md
-# section 2-4 and the OWNER-ADJUDICATED correction in
-# docs/superpowers/plans/2026-08-01-sp2-shrinkage.md (2026-08-01, review round
-# 1, finding 1): a hot spot is a local maximum of solidification time `t_s`
-# (last metal to freeze). The ORIGINAL "monotonically non-decreasing t_s path
-# to any surface node" feeding-path definition was proved degenerate (a
-# region can never escape its own hot-node superlevel-set component under
-# that rule, making it tautologically equivalent to "region touches the
-# domain surface" -- replacing the whole search with `return False` left
-# every test green). The adjudicated definition instead: a region is fed iff
-# a node path over tet adjacency connects it to a caller-supplied FEED node
-# (e.g. nodes on Adiabatic-tagged riser/sleeve faces -- NOT the generic mold
-# surface; a bare sand casting with no risers passes an empty `feed_nodes`,
-# so every region is conservatively unfed) with EVERY node on that path
-# (endpoints included) at or above `corridor_frac * region["t_max"]` -- a
-# warm corridor that stays liquid until the region itself finishes freezing,
-# with a 5% relative tolerance. This breaks the old degeneracy because the
-# corridor threshold is relative to the REGION's own peak, not the mesh-wide
-# hot cut, so a feed node need not itself be part of the region's hot-node
-# superlevel set.
-#
-# Severity scales with the alloy's volumetric shrinkage percentage
-# (`SHRINKAGE_PCT`) times the region's volume. A region whose highest point
-# against gravity is itself a free surface sinks that surface instead of
-# voiding internally. A feeder placed above the region must have modulus
-# (V/A, Chvorinov's rule) >= 1.2x the region's own modulus to freeze after it.
-# ---------------------------------------------------------------------------
-
-# Local vertex-index pairs/triples of a tetrahedron (0,1,2,3), used to walk a
-# tet's 6 edges / 4 faces without re-deriving the combinatorics inline.
+
+# Local vertex-index pairs/triples of a tetrahedron.
 _TET_EDGE_PAIRS = [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]
 _TET_FACES = [(0, 1, 2), (0, 1, 3), (0, 2, 3), (1, 2, 3)]
 
 
 def _tet_volumes(pos: np.ndarray, tets: np.ndarray) -> np.ndarray:
-    """(M,) unsigned tet volumes in mm^3 (independent copy of the closed form
-    used inside `node_gradient`, kept private and separate so Task 1's
-    function is never touched)."""
+    """(M,) unsigned tet volumes in mm^3."""
     if tets.shape[0] == 0:
         return np.zeros(0)
     p0, p1, p2, p3 = pos[tets[:, 0]], pos[tets[:, 1]], pos[tets[:, 2]], pos[tets[:, 3]]
@@ -253,14 +124,10 @@ def _tet_volumes(pos: np.ndarray, tets: np.ndarray) -> np.ndarray:
 
 
 def _induced_components(n_nodes: int, tets: np.ndarray, mask: np.ndarray):
-    """Connected components of the tet-adjacency graph induced on `mask`
-    (i.e. vertices = nodes with `mask[node]` True, edges = tet-adjacent pairs
-    both satisfying `mask`). Returns `(n_comp, labels)` where `labels` has
-    shape `(n_nodes,)`: the component id for masked nodes, `-1` elsewhere.
-    Shared by `hot_regions` (mask = the global top-`top_frac` hot cut) and
-    `is_fed` (mask = the per-region warm corridor, `t_s >= corridor_frac *
-    region_t_max`) -- same graph-labeling primitive, different masks.
-    """
+    """Connected components of the tet-adjacency graph restricted to `mask`.
+
+    Returns (n_comp, labels) with labels of shape (n_nodes,): the component
+    id for masked nodes, -1 elsewhere."""
     idx = np.flatnonzero(mask)
     labels_full = np.full(n_nodes, -1, dtype=np.int64)
     if idx.size == 0:
@@ -290,28 +157,8 @@ def _induced_components(n_nodes: int, tets: np.ndarray, mask: np.ndarray):
 
 
 def _boundary_faces(tets: np.ndarray) -> np.ndarray:
-    """(F, 3) int64 array of the mesh's own boundary-face vertex triples
-    (ORIGINAL per-tet vertex order, not a sorted dedup key) -- faces touched
-    by exactly one tet, i.e. with nothing "on the other side".
-
-    Review finding I4: `_region_surface_area`'s domain-boundary term used to
-    rebuild this WHOLE-MESH structure via a pure-Python double loop over
-    every tet's 4 faces, on EVERY call -- but it depends only on `tets`, not
-    on any particular region's `nodes`/`threshold`, so `feeder_recommendations`
-    (this module's only caller of `_region_surface_area`) was redoing the
-    exact same O(n_tets) work once per region. Measured on BUKSA (32138
-    tets, 16 regions): 4.74 s total screening time, ~8x the documented
-    ~0.5-0.6 s (a figure that had only ever been measured on a single-region
-    part). Hoisted here so it runs ONCE per `feeder_recommendations` call,
-    vectorized via `np.unique(..., axis=0, return_counts=True)` (both the
-    dict-free dedup AND -- inside `_region_surface_area` -- the boundary-term
-    summation itself are now array ops, not a Python loop).
-
-    Every tet contributes 4 candidate faces; a face shared by exactly 2 tets
-    (an interior face) is discarded, one shared by exactly 1 (nothing on the
-    other side) is kept -- a face can never be shared by >2 tets in a
-    conforming tet mesh, so `counts == 1` is the correct (and only
-    necessary) boundary test."""
+    """(F, 3) vertex triples of the mesh's outer faces: faces that belong to
+    exactly one tet, in their original per-tet vertex order."""
     tets = np.asarray(tets, dtype=np.int64)
     if tets.shape[0] == 0:
         return np.zeros((0, 3), dtype=np.int64)
@@ -327,66 +174,17 @@ def _boundary_faces(tets: np.ndarray) -> np.ndarray:
 def _region_surface_area(pos: np.ndarray, tets: np.ndarray, nodes: np.ndarray,
                          t_s: np.ndarray, threshold: float,
                          boundary_faces: np.ndarray = None) -> float:
-    """Boundary-surface area (mm^2) of the SAME body `hot_regions`' fraction-
-    weighted `volume_mm3` integrates over (review round 2, finding NEW-1: the
-    original all-4-vertices-hot rule measured a strictly SMALLER body than
-    the fraction-weighted volume, giving an inflated, inconsistent modulus).
-
-    Round-2 history / why this changed again (review round 3): the first fix
-    ran the co-area formula on a BINARIZED 0/1 indicator (`nodes` membership
-    only). That is an exact calculus identity for a piecewise-linear field,
-    but the field it measures is the indicator's own artificial "staircase"
-    step (always crossing at each tet edge's midpoint, regardless of where
-    the REAL field actually crosses its selection threshold) -- an
-    ANISOTROPIC discrete-perimeter measure that plateaus ~24-29% high on
-    curved/diagonal boundaries (axis-aligned cuts happened to look fine only
-    because a midpoint crossing is coincidentally correct when the true
-    level surface IS flat and axis-aligned).
-
-    This version runs marching tetrahedra directly on the CONTINUOUS `t_s`
-    field, cut at `threshold` -- literally "the same cut that defines the
-    region" (the percentile cutoff `hot_regions` used, stored on each region
-    dict as `hot_threshold`), so the interpolated crossing point on each tet
-    edge sits at its true fractional position (`(threshold-f_a)/(f_b-f_a)`),
-    not an assumed 0.5. Per tet, classified by how many of its 4 vertices
-    have `t_s >= threshold` (`k`):
-      - k in {0, 4}: no crossing, no contribution.
-      - k in {1, 3}: the crossing is a single triangle through the 3 edges
-        connecting the "odd one out" vertex to the other 3 (interpolated
-        edge points) -- the standard 1-3 marching-tetrahedra case.
-      - k == 2: the crossing is a planar quadrilateral through the 4 edges
-        connecting each of the 2 "above" vertices to each of the 2 "below"
-        vertices (`t_s` is affine per tet, so ANY level set within a tet is
-        a plane intersected with the tet -- always exactly planar,
-        regardless of the specific values -- computed via the standard
-        diagonal cross-product quadrilateral-area formula).
-    A tet is attributed to THIS region iff any of its vertices is in
-    `nodes` (well-defined: two vertices with `t_s >= threshold` sharing a
-    tet are adjacent by `hot_regions`' own definition, so a mixed tet's
-    "above" vertices, if any, all belong to exactly one region already).
-
-    Domain-boundary term (KEPT, unchanged, review-verified exact for a
-    region whose boundary coincides with the mesh's own outer surface): a
-    tet face touching no other tet (nothing "on the other side" for the
-    interior marching-tets term to see) is credited separately, weighted by
-    the (binary, `nodes`-membership) hot fraction of its own 3 vertices.
-
-    See tests/test_shrinkage.py for the acceptance suite this targets:
-    sphere-field area converging toward the analytic 4*pi*R^2 with
-    refinement (not plateauing), a diagonal-plane cut and an orientation
-    sweep both landing far below the ~29-36% anisotropy the binarized
-    co-area formula showed, and the axis-aligned hand-check (now computed
-    against a genuine continuous field, not an arbitrarily-thresholded
-    binary block) landing near the Chvorinov M=side/6 identity for a cube.
-
-    `boundary_faces` (review finding I4): the mesh's own domain-boundary
-    faces (`_boundary_faces(tets)`'s output), region-INDEPENDENT -- a single
-    call's worth of work shared by every region in one
-    `feeder_recommendations` invocation instead of being rebuilt (a
-    pure-Python double loop over ALL tets) on every one of this function's
-    calls. `None` (the default) computes it internally -- exactly the old
-    behavior, for any direct caller that doesn't have one precomputed.
-    """
+    """Surface area (mm^2) of a hot region, for its V/A modulus.
+
+    The interior part is the iso-surface `t_s == threshold` (the cut that
+    selected the region), built by marching tetrahedra on the continuous
+    field: a tet with 1 or 3 vertices above the threshold contributes a
+    triangle, one with 2 contributes a planar quadrilateral, and crossing
+    points sit at their interpolated position on each edge. A tet counts for
+    this region if any of its vertices is in `nodes`. Outer faces of the
+    mesh are added separately, weighted by the fraction of their 3 vertices
+    in the region. `boundary_faces` can be passed in to avoid recomputing
+    `_boundary_faces(tets)` for every region."""
     n_nodes = pos.shape[0]
     is_member = np.zeros(n_nodes, dtype=bool)
     is_member[np.asarray(nodes, dtype=np.int64)] = True
@@ -433,10 +231,8 @@ def _region_surface_area(pos: np.ndarray, tets: np.ndarray, nodes: np.ndarray,
     if boundary_faces is None:
         boundary_faces = _boundary_faces(tets)
     if boundary_faces.shape[0] > 0:
-        # Vectorized version of the old per-face Python loop: each boundary
-        # face's hot fraction is (# of its 3 vertices in `nodes`) / 3, and
-        # only faces with a nonzero fraction contribute a (fraction-
-        # weighted) triangle area.
+        # Each outer face counts with the fraction of its 3 vertices in
+        # the region.
         member_counts = is_member[boundary_faces].sum(axis=1).astype(np.float64)  # (F,)
         frac = member_counts / 3.0
         nonzero = frac > 0
@@ -450,48 +246,15 @@ def _region_surface_area(pos: np.ndarray, tets: np.ndarray, nodes: np.ndarray,
 
 def hot_regions(pos: np.ndarray, tets: np.ndarray, t_s: np.ndarray,
                 top_frac: float = 0.10) -> list:
-    """Connected components (via tet adjacency) of the top-`top_frac`
-    latest-solidifying nodes.
-
-    A node is "hot" if `t_s >= percentile(t_s, 100*(1-top_frac))` (inclusive,
-    same tie convention as `porosity_risk`'s percentile mode -- may include
-    slightly more than `top_frac` when values tie at the threshold).
-    Components are the connected components of the graph whose vertices are
-    the hot nodes and whose edges are tet-adjacency pairs restricted to hot
-    endpoints (two hot nodes sharing a tet).
-
-    Returns a list of dicts sorted by `t_max` descending, each carrying a
-    stable `region_id` = its index in THIS sorted list (0-based) -- the id
-    downstream functions (`macro_shrinkage`, `top_sinks`,
-    `feeder_recommendations`) read back off the dict rather than
-    re-deriving positionally, so it stays coherent across filtering:
-    `{"nodes": sorted int index array, "centroid": (3,) mm, "volume_mm3":
-    float, "t_max": float, "region_id": int, "hot_threshold": float,
-    "sub_resolution": bool}`. `sub_resolution` (review finding C1) is `True`
-    iff the region has fewer than 4 nodes -- below this floor
-    `_region_surface_area`'s iso-surface area and the lumped-nodal
-    `volume_mm3` computed below diverge without bound as `t_max` approaches
-    `hot_threshold`; `feeder_recommendations` reads this flag back to skip
-    such regions rather than report a physically meaningless modulus (see
-    its own docstring), while `macro_shrinkage`/`top_sinks` still list them
-    unfiltered (their volume/severity/centroid stay meaningful).
-    `hot_threshold` is the same percentile cutoff (`t_s >= hot_threshold`)
-    used to select every region in THIS call -- `feeder_recommendations`
-    (review round 3) reads it back to run marching tetrahedra on the
-    continuous `t_s` field at exactly "the cut that defines the region",
-    rather than re-deriving or guessing it downstream.
-
-    `volume_mm3` is a fraction-weighted sum over ALL tets touching the
-    region: each tet contributes `tet_volume * n_hot_vertices/4` to the
-    component of its hot vertices (well-defined -- any two hot vertices
-    sharing a tet are adjacent by this function's own definition, so a
-    tet's hot vertices, if any, all belong to the same component). This is
-    mesh-convergent (equivalent to summing each hot node's lumped nodal
-    volume, `sum(tet_volume/4)` over its incident tets) and replaces an
-    earlier all-4-vertices-hot rule that measurably undercounted volume
-    (35-65% on smooth fields, review round 1 finding 2) by discarding every
-    tet straddling the region boundary instead of weighting it.
-    """
+    """Connected groups (via tet adjacency) of the top-`top_frac` latest-
+    solidifying nodes, sorted by `t_max` descending.
+
+    Each region is a dict: `nodes` (sorted indices), `centroid` (mm),
+    `volume_mm3`, `t_max`, `region_id` (its index in this list),
+    `hot_threshold` (the percentile cut used) and `sub_resolution` (fewer
+    than 4 nodes, too small for a meaningful V/A modulus). The volume sums,
+    over all tets touching the region, tet volume times the fraction of its
+    vertices that are hot, which converges with mesh refinement."""
     top_frac = float(top_frac)
     if not (0.0 < top_frac <= 1.0):
         raise ValueError(f"hot_regions: top_frac must be in (0, 1], got {top_frac!r}")
@@ -536,48 +299,22 @@ def hot_regions(pos: np.ndarray, tets: np.ndarray, t_s: np.ndarray,
     for i, r in enumerate(regions):
         r["region_id"] = i
         r["hot_threshold"] = float(threshold)
-        # Review finding C1: a region with fewer than 4 nodes is below the
-        # resolution floor at which `_region_surface_area`'s iso-surface
-        # area and the lumped-nodal `volume_mm3` above agree even
-        # approximately -- as `t_max` approaches `hot_threshold` the true
-        # iso-surface area shrinks toward zero while the lumped volume stays
-        # O(h^3), so V/A (the feeder modulus) diverges without bound.
-        # `feeder_recommendations` reads this flag back to skip such regions
-        # rather than emit a physically-impossible modulus; `volume_mm3`/
-        # `severity` remain meaningful (only the surface-area-derived
-        # modulus is the part that diverges), so `macro_shrinkage` still
-        # lists these regions unfiltered.
+        # Below 4 nodes the iso-surface area and the lumped volume no longer
+        # agree, so the V/A modulus is not meaningful.
         r["sub_resolution"] = len(r["nodes"]) < 4
     return regions
 
 
 def is_fed(region: dict, t_s: np.ndarray, tets: np.ndarray,
           feed_nodes: np.ndarray, corridor_frac: float = 0.95) -> bool:
-    """OWNER-ADJUDICATED SEMANTICS (2026-08-01, review round 1 finding 1 --
-    see the module-level comment above for why the original monotone-escape
-    definition was degenerate). A region is fed iff some node of `region`
-    and some node of `feed_nodes` (e.g. Adiabatic-tagged riser/sleeve nodes;
-    an empty `feed_nodes` -- a bare sand casting with no risers -- always
-    gives False, the conservative screen) lie in the same connected
-    component of the tet-adjacency graph induced on the "warm corridor":
-    nodes with `t_s >= corridor_frac * region["t_max"]`.
-
-    The corridor threshold is relative to the REGION's own peak, not the
-    mesh-wide hot cut -- a feed node several hops away through nodes just
-    below the region's own hot-selection threshold (but still within
-    `corridor_frac` of its peak) now correctly counts as fed, which the old
-    "escape the hot superlevel set" definition could never recognize.
-
-    Guard (review round 2, finding NEW-2): if `region["t_max"] <= 0`, the
-    corridor threshold `corridor_frac * t_max` would be `<= 0`, which is
-    vacuously satisfied by every non-negative `t_s` value (solidification
-    times are never negative) -- i.e. the "warm corridor" would silently
-    become the WHOLE mesh, making every region trivially "fed" through any
-    non-empty `feed_nodes` regardless of actual connectivity. This isn't a
-    physically meaningful region in the first place (its peak already
-    solidified at t=0), so it's reported unfed (`False`) rather than let the
-    threshold degenerate.
-    """
+    """Whether a feed source can still reach the region through liquid metal.
+
+    True iff some region node and some `feed_nodes` node are connected
+    through nodes with `t_s >= corridor_frac * region["t_max"]` (a corridor
+    that stays liquid until the region itself is nearly solid). Feed nodes
+    are adiabatic-tagged surface nodes, for example a riser face. With no
+    feed nodes every region is unfed, which is the conservative answer. A
+    region with t_max <= 0 is reported unfed."""
     t_s = np.asarray(t_s, dtype=np.float64)
     tets = np.asarray(tets, dtype=np.int64)
     feed_nodes = np.asarray(feed_nodes, dtype=np.int64)
@@ -610,16 +347,8 @@ def is_fed(region: dict, t_s: np.ndarray, tets: np.ndarray,
 
 def macro_shrinkage(pos: np.ndarray, tets: np.ndarray, t_s: np.ndarray,
                     feed_nodes: np.ndarray, shrinkage_pct: float) -> list:
-    """Unfed hot regions (`hot_regions` with its default `top_frac`, filtered
-    by `is_fed` against `feed_nodes` -- nodes on Adiabatic-tagged riser/sleeve
-    faces, NOT the generic mold/cooling surface), each with `severity =
-    shrinkage_pct/100 * volume_mm3` added. `region_id` is propagated
-    unchanged from `hot_regions` (its index in the FULL sorted list, not
-    renumbered by this function's filtering) so ids stay coherent if this
-    output is chained into `top_sinks`/`feeder_recommendations`. A bare sand
-    casting with no risers passes an empty `feed_nodes` -- every hot region
-    is then unfed, the conservative screen (see `is_fed`).
-    """
+    """The hot regions that are not fed, each with `severity = shrinkage_pct /
+    100 * volume_mm3`. `region_id` keeps its value from `hot_regions`."""
     regions = hot_regions(pos, tets, t_s)
     out = []
     for r in regions:
@@ -632,13 +361,8 @@ def macro_shrinkage(pos: np.ndarray, tets: np.ndarray, t_s: np.ndarray,
 
 def top_sinks(regions: list, pos: np.ndarray, surface_nodes: np.ndarray,
              gravity_vec) -> list:
-    """Regions whose highest point along `-gravity_vec` is itself a surface
-    node sink that free surface (rather than voiding internally) -- returns
-    `{"region_id": the region dict's own "region_id" (propagated, NOT a
-    positional re-index -- see `hot_regions`/`macro_shrinkage`), "location":
-    (3,) mm of that surface node, "volume_mm3": the region's own volume}` for
-    each qualifying region.
-    """
+    """Regions whose highest node along -gravity is a surface node. Returns
+    `{"region_id", "location" (mm), "volume_mm3"}` for each."""
     pos = np.asarray(pos, dtype=np.float64)
     gravity_vec = np.asarray(gravity_vec, dtype=np.float64)
     up = -gravity_vec
@@ -663,51 +387,13 @@ def top_sinks(regions: list, pos: np.ndarray, surface_nodes: np.ndarray,
 
 def feeder_recommendations(regions: list, pos: np.ndarray, tets: np.ndarray,
                            surface_nodes: np.ndarray, gravity_vec, t_s: np.ndarray) -> list:
-    """Per region: `{"location": surface point above the region's centroid
-    along -gravity_vec, "min_modulus_mm": 1.2 * region_modulus, "region_id":
-    the region dict's own "region_id" (propagated, NOT a positional
-    re-index)}`, where `region_modulus = V/A` (Chvorinov's rule), `A` =
-    `_region_surface_area` of the region's own tets, cut at the region's own
-    `hot_threshold` (review round 3 -- see `_region_surface_area`'s
-    docstring; `t_s` is a NEW required parameter added in that round since
-    the marching-tetrahedra area computation needs the continuous field, not
-    just the region's own node membership).
-
-    `location` is the surface node minimizing perpendicular distance to the
-    vertical line (direction `-gravity_vec`) through the centroid, restricted
-    to surface nodes at or above the centroid along that direction (falling
-    back to all surface nodes if none qualify, e.g. a degenerate/zero-height
-    domain) -- a nearest-existing-mesh-node stand-in for "directly above".
-
-    Raises `ValueError` if `surface_nodes` is empty (no candidate feeder
-    location exists -- a mesh with no free surface is not a valid casting).
-    A region with zero volume or zero boundary area (a degenerate/thin
-    region under the fraction-weighted volume rule) gets a logged warning
-    and `min_modulus_mm = 0.0` rather than a silently meaningless value.
-
-    Review finding C1 (sub-resolution regions): a region with fewer than 4
-    nodes (`region.get("sub_resolution", False)`, set by `hot_regions`) is
-    SKIPPED entirely -- no feeder recommendation is produced for it. For such
-    a region the iso-surface area shrinks toward zero as `t_max` approaches
-    the region's own selection threshold (an unbounded-in-the-limit
-    `M = V/A`, since the lumped-nodal `volume_mm3` stays O(h^3) while the
-    true iso-surface area vanishes), so `min_modulus_mm` is not a meaningful
-    number at this resolution -- reproduced on the branch's own BUKSA smoke
-    artifact (`runs/gate_c/smoke/BUKSA_cleaned_summary.txt`): a 1-node
-    region's modulus of 185 mm printed a `min_modulus_mm=222` feeder
-    recommendation on a 382x317x254 mm part. `macro_shrinkage` still lists
-    these regions (their `volume_mm3`/`severity` are fine -- only the
-    surface-area-derived modulus is the part that diverges), and the region
-    still contributes to `top_sinks`. The number of regions skipped this way
-    is logged once (`logger.info`) rather than per-region, so a part with
-    many marginal regions doesn't flood the log.
-
-    Review finding I4 (performance): the mesh's boundary-face structure
-    (`_boundary_faces(tets)`) is region-INDEPENDENT and is now computed ONCE
-    here and passed to every `_region_surface_area` call, instead of being
-    rebuilt (a pure-Python double loop over every tet) once per region --
-    see `_region_surface_area`'s own docstring for the before/after cost.
-    """
+    """Per region: a feeder `location` and `min_modulus_mm = 1.2 * V/A`.
+
+    The location is the surface node closest to the vertical line (along
+    -gravity) through the region's centroid, among surface nodes at or above
+    the centroid (all surface nodes if none are). Sub-resolution regions are
+    skipped. A region with zero volume or area gets a warning and
+    `min_modulus_mm = 0.0`. Raises ValueError if `surface_nodes` is empty."""
     pos = np.asarray(pos, dtype=np.float64)
     tets = np.asarray(tets, dtype=np.int64)
     t_s = np.asarray(t_s, dtype=np.float64)
@@ -723,7 +409,7 @@ def feeder_recommendations(regions: list, pos: np.ndarray, tets: np.ndarray,
             "feeder location exists (a mesh with no free surface is not a "
             "valid casting)")
     surface_pos = pos[surface_nodes]
-    boundary_faces = _boundary_faces(tets)  # I4: computed once, shared by every region
+    boundary_faces = _boundary_faces(tets)  # shared by every region
 
     recs = []
     n_skipped_sub_resolution = 0
@@ -769,83 +455,17 @@ def feeder_recommendations(regions: list, pos: np.ndarray, tets: np.ndarray,
     return recs
 
 
-# ---------------------------------------------------------------------------
-# screen_casting -- one-call orchestrator (Task 3). See
-# docs/superpowers/specs/2026-07-24-shrinkage-porosity-screening-design.md
-# section "Multi-orientation augmentation": the conduction thermal field (and
-# hence Niyama, output 1, and macro-shrinkage region *detection*) does not
-# depend on gravity -- only the gravity-conditioned outputs (top sinks and
-# feeder placement, outputs 2-4 in the spec's numbering) do. `orientations=
-# "axes"` sweeps those over the 6 signed axis directions at zero extra
-# thermal-solve cost, per the spec.
-# ---------------------------------------------------------------------------
-
-# Signed unit axis directions for the `orientations="axes"` sweep, keyed by
-# the label used in `per_orientation`.
-_AXIS_SWEEP = {
-    "+x": (1.0, 0.0, 0.0), "-x": (-1.0, 0.0, 0.0),
-    "+y": (0.0, 1.0, 0.0), "-y": (0.0, -1.0, 0.0),
-    "+z": (0.0, 0.0, 1.0), "-z": (0.0, 0.0, -1.0),
-}
-
-
 def screen_casting(pos: np.ndarray, tets: np.ndarray, surface_nodes: np.ndarray,
                    t_s: np.ndarray, TL: float, TS: float, alloy_family: str,
-                   gravity_vec=(0, 0, -1), gradT: np.ndarray = None, dTdt: np.ndarray = None,
-                   orientations: str = None, feed_nodes: np.ndarray = None) -> dict:
-    """One-call shrinkage-porosity screening: Niyama (exact when `gradT`+`dTdt`
-    are both given, else the time-only proxy), macro-shrinkage region
-    detection, top-surface sinks, and feeder recommendations -- the single
-    entry point Task 4's `predict_casting` integration is meant to call.
-
-    `feed_nodes` (BRIEF DEVIATION, owner-adjudicated 2026-08-01): the plan's
-    signature (`docs/superpowers/plans/2026-08-01-sp2-shrinkage.md` Task 3)
-    does not list this parameter, but `macro_shrinkage`/`is_fed` (Task 2,
-    review round 1) require an explicit array of ADIABATIC-tagged riser/
-    sleeve nodes to know what "feeds" a region -- this is NOT derivable from
-    `surface_nodes` (the whole free surface, used only for sink/feeder
-    placement -- a different, larger set that includes ordinary mold-cooling
-    surface). Omitting it here would either silently misuse `surface_nodes`
-    as feed nodes (every region trivially "self-fed" wherever it touches any
-    surface, defeating the whole unfed/fed distinction) or hardcode an empty
-    array with no way for a caller to ever supply real risers. `feed_nodes=
-    None` defaults to an empty array -- the same "bare sand casting, no
-    risers, everything conservatively unfed" convention `macro_shrinkage`/
-    `is_fed` already document.
-
-    Niyama: uses `niyama_exact` + `porosity_risk(mode="absolute")` whenever
-    both `gradT` and `dTdt` are given (`niyama_source="exact"`); otherwise
-    `niyama_proxy` + `porosity_risk(mode="percentile")`
-    (`niyama_source="proxy"`) -- see the module docstring for why the two
-    paths use different risk-thresholding modes (the proxy is only a
-    within-part relative ranking, never comparable to the literature
-    absolute thresholds).
-
-    Material: `alloy_family` must be a key of `SHRINKAGE_PCT`. An unknown (or
-    missing) family raises a loud `KeyError` naming the known families --
-    mirroring `porosity_risk`'s own loud-failure style for the identical
-    reason: silently defaulting would silently mis-scale every downstream
-    severity/threshold number.
-
-    Orientations: `orientations=None` (default) computes `macro_shrinkage`,
-    `top_sinks`, and `feeders` once, using `gravity_vec` as given, and
-    `per_orientation` is `None`. `orientations="axes"` additionally sweeps
-    `top_sinks`/`feeders` (against the SAME `macro_shrinkage` regions --
-    region *detection* is gravity-independent, see the module-level comment
-    above) over the 6 signed axis directions, populating `per_orientation`
-    with one `{"macro_shrinkage", "top_sinks", "feeders"}` entry per axis
-    label (`"+x"`, `"-x"`, `"+y"`, `"-y"`, `"+z"`, `"-z"`); the top-level
-    `macro_shrinkage`/`top_sinks`/`feeders` keys are still populated from the
-    single `gravity_vec` argument regardless of `orientations`, so the
-    `"-z"` entry is numerically identical to calling `screen_casting` again
-    with `orientations=None, gravity_vec=(0, 0, -1)` (the default) --
-    verified in `tests/test_shrinkage.py`.
-
-    Returns `{"niyama": (N,), "porosity_risk": (N,) bool, "niyama_source":
-    str, "macro_shrinkage": [...], "top_sinks": [...], "feeders": [...],
-    "per_orientation": {axis: {...}} | None, "material": {"alloy_family",
-    "shrinkage_pct"}}`.
-    """
+                   gravity_vec=(0, 0, -1), feed_nodes: np.ndarray = None) -> dict:
+    """All four screening outputs for one part.
+
+    `alloy_family` must be a key of `SHRINKAGE_PCT`. `surface_nodes` are all
+    surface nodes (for sinks and feeder placement); `feed_nodes` are the
+    nodes that can feed a region (None means none).
+
+    Returns {"niyama", "porosity_risk", "niyama_source": "proxy",
+    "macro_shrinkage", "top_sinks", "feeders", "material"}."""
     if alloy_family not in SHRINKAGE_PCT:
         raise KeyError(
             f"screen_casting: unknown alloy_family {alloy_family!r} -- known "
@@ -859,43 +479,20 @@ def screen_casting(pos: np.ndarray, tets: np.ndarray, surface_nodes: np.ndarray,
     feed_nodes_arr = (np.zeros(0, dtype=np.int64) if feed_nodes is None
                       else np.asarray(feed_nodes, dtype=np.int64))
 
-    if gradT is not None and dTdt is not None:
-        ny = niyama_exact(gradT, dTdt)
-        risk = porosity_risk(ny, mode="absolute", alloy_family=alloy_family)
-        niyama_source = "exact"
-    else:
-        ny = niyama_proxy(pos, tets, t_s, TL, TS)
-        risk = porosity_risk(ny, mode="percentile")
-        niyama_source = "proxy"
+    ny = niyama_proxy(pos, tets, t_s, TL, TS)
+    risk = porosity_risk(ny)
 
     shrinkage_pct = SHRINKAGE_PCT[alloy_family]
     macro = macro_shrinkage(pos, tets, t_s, feed_nodes_arr, shrinkage_pct)
     sinks = top_sinks(macro, pos, surface_nodes, gravity_vec)
     feeders = feeder_recommendations(macro, pos, tets, surface_nodes, gravity_vec, t_s)
 
-    per_orientation = None
-    if orientations == "axes":
-        per_orientation = {}
-        for label, axis_vec in _AXIS_SWEEP.items():
-            per_orientation[label] = {
-                "macro_shrinkage": macro,
-                "top_sinks": top_sinks(macro, pos, surface_nodes, axis_vec),
-                "feeders": feeder_recommendations(macro, pos, tets, surface_nodes,
-                                                  axis_vec, t_s),
-            }
-    elif orientations is not None:
-        raise ValueError(
-            f"screen_casting: unknown orientations mode {orientations!r} "
-            f"(expected None or 'axes')"
-        )
-
     return {
         "niyama": ny,
         "porosity_risk": risk,
-        "niyama_source": niyama_source,
+        "niyama_source": "proxy",
         "macro_shrinkage": macro,
         "top_sinks": sinks,
         "feeders": feeders,
-        "per_orientation": per_orientation,
         "material": {"alloy_family": alloy_family, "shrinkage_pct": shrinkage_pct},
     }

src/splits.py

@@ -1,61 +0,0 @@
-"""Reproducible train/test splits over the feature cache.
-
-Three splits, each with an explicit train/test stem list:
-  interpolation  — hold out (prefix,dims) groups (interpolate within families)
-  family         — hold out whole geometry prefixes (generalize to new shapes)
-  fully_cooled   — evaluate only on 0-adiabatic castings (real-inference regime)
-"""
-import json
-import re
-
-import numpy as np
-
-HELDOUT_FAMILIES = ["BO", "T45", "WH", "RR"]
-_NUM_RE = re.compile(r"^\d+$")
-
-
-def dims_of(stem: str) -> tuple:
-    return tuple(t for t in stem.split("_")[1:] if _NUM_RE.match(t))
-
-
-def build_splits(meta, seed: int = 0, test_frac: float = 0.1) -> dict:
-    stems = list(map(str, meta["stem"]))
-    prefix = {s: p for s, p in zip(stems, map(str, meta["prefix"]))}
-    fully = {s: bool(v) for s, v in zip(stems, meta["is_fully_cooled"])}
-
-    non_family = [s for s in stems if prefix[s] not in HELDOUT_FAMILIES]
-    family_test = [s for s in stems if prefix[s] in HELDOUT_FAMILIES]
-
-    # interpolation: hold out groups keyed by (prefix, dims)
-    groups = {}
-    for s in non_family:
-        groups.setdefault((prefix[s], dims_of(s)), []).append(s)
-    keys = sorted(groups.keys())
-    rng = np.random.default_rng(seed)
-    rng.shuffle(keys)
-    n_test = max(1, int(round(len(keys) * test_frac)))
-    test_keys = set(keys[:n_test])
-
-    interp_train, interp_test = [], []
-    for k, members in groups.items():
-        (interp_test if k in test_keys else interp_train).extend(members)
-
-    # fully_cooled: restrict to 0-adiabatic; test = fully-cooled in the held-out groups
-    fc_train = [s for s in interp_train if fully[s]]
-    fc_test = [s for s in interp_test if fully[s]]
-
-    return {
-        "interpolation": {"train": sorted(interp_train), "test": sorted(interp_test)},
-        "family": {"train": sorted(non_family), "test": sorted(family_test)},
-        "fully_cooled": {"train": sorted(fc_train), "test": sorted(fc_test)},
-    }
-
-
-def save_splits(splits: dict, path: str):
-    with open(path, "w") as f:
-        json.dump(splits, f, indent=2)
-
-
-def load_splits(path: str) -> dict:
-    with open(path) as f:
-        return json.load(f)

src/step_mesh.py

@@ -1,45 +0,0 @@
-"""STEP -> tetrahedral gmsh 2.2 mesh with all external surface tagged 'Cooling'.
-
-Uses the gmsh Python API (bundled OpenCASCADE) to import a STEP solid, tag the
-volume and its full boundary, mesh uniformly, and export in the gmsh 2.2 format
-consumed by src/msh_parse.py. All lengths in mm."""
-import numpy as np
-import gmsh
-
-
-def auto_elem_size(bbox_diag: float) -> float:
-    """Target element size (mm): bbox diagonal / 60, the validated Phase 2A
-    density chosen to give training-comparable node density on real castings."""
-    return bbox_diag / 60.0
-
-
-def mesh_step(step_path: str, out_msh: str, elem_size: float | None = None) -> dict:
-    gmsh.initialize()
-    try:
-        gmsh.option.setNumber("General.Terminal", 0)
-        gmsh.model.add("casting")
-        gmsh.model.occ.importShapes(step_path)
-        gmsh.model.occ.synchronize()
-
-        xmin, ymin, zmin, xmax, ymax, zmax = gmsh.model.getBoundingBox(-1, -1)
-        bbox = (xmax - xmin, ymax - ymin, zmax - zmin)
-        diag = float(np.linalg.norm(bbox))
-        size = float(elem_size) if elem_size is not None else auto_elem_size(diag)
-        gmsh.option.setNumber("Mesh.MeshSizeMin", size)
-        gmsh.option.setNumber("Mesh.MeshSizeMax", size)
-
-        vols = [t for (d, t) in gmsh.model.getEntities(3)]
-        surfs = [t for (d, t) in gmsh.model.getEntities(2)]
-        pv = gmsh.model.addPhysicalGroup(3, vols); gmsh.model.setPhysicalName(3, pv, "Volume")
-        pc = gmsh.model.addPhysicalGroup(2, surfs); gmsh.model.setPhysicalName(2, pc, "Cooling")
-
-        gmsh.model.mesh.generate(3)
-        gmsh.option.setNumber("Mesh.MshFileVersion", 2.2)
-        gmsh.write(out_msh)
-
-        node_tags, _, _ = gmsh.model.mesh.getNodes()
-        tet_tags = gmsh.model.mesh.getElementsByType(4)[0]
-        return {"n_nodes": len(node_tags), "n_tets": len(tet_tags),
-                "bbox_mm": bbox, "elem_size": size}
-    finally:
-        gmsh.finalize()

src/targets.py

@@ -1,39 +0,0 @@
-"""Per-node solidification time from temperature snapshots.
-
-For each node, find the first time its temperature crosses T_SOLIDUS
-(cooling, high->low) and linearly interpolate the crossing time. Nodes that
-never cross use the last simulation time. Ported from the verified logic in
-the original dataset.py::_compute_y_nodal, vectorised in NumPy."""
-import numpy as np
-
-T_SOLIDUS = 1120.0
-
-
-def solidification_time(temps: np.ndarray, times: np.ndarray,
-                        t_solidus: float = T_SOLIDUS) -> np.ndarray:
-    temps = np.asarray(temps, dtype=np.float64)   # (T, N)
-    times = np.asarray(times, dtype=np.float64)   # (T,)
-    T, N = temps.shape
-
-    below = temps <= t_solidus                     # (T, N)
-    any_below = below.any(axis=0)                  # (N,)
-    first_s = below.argmax(axis=0)                 # (N,) first True row, 0 if none
-    s = np.clip(first_s, 1, None)                  # avoid index -1
-    ni = np.arange(N)
-
-    t_hi = temps[s, ni]        # temp at crossing snapshot
-    t_lo = temps[s - 1, ni]    # temp at previous snapshot
-    dt = t_hi - t_lo
-    dt = np.where(np.abs(dt) < 1e-10, -1.0, dt)
-    alpha = (t_solidus - t_lo) / dt
-    y = times[s - 1] + alpha * (times[s] - times[s - 1])
-
-    # already below at first snapshot -> time[0]
-    y = np.where((first_s == 0) & below[0], times[0], y)
-    # never crosses -> last time
-    y = np.where(~any_below, times[-1], y)
-    return y.astype(np.float32)
-
-
-def global_time(y_nodal: np.ndarray) -> float:
-    return float(np.asarray(y_nodal).max())

src/train_ft2b.py

@@ -1,627 +0,0 @@
-"""FT2b Task 7: 3-stage staged-training driver (scale -> BC -> calibrate).
-
-Orchestrates the already-built FT2b pieces on top of the SP1 production GNN:
-  - Task 4 (`src.warm_start`): zero-init FiLM/BC-column warm-start (+ the
-    stage-chaining `build_or_load_v4` entry point, used implicitly here since
-    `src.train_gnn.train`'s `init_ckpt` calls it).
-  - Task 5 (`src.train_gnn`): `train(...)`'s g/FiLM/freeze hooks.
-  - Task 6 (`src.gbm_ft2b`): `refit_gbm(...)`, the multi-alloy scale-head GBM.
-
-Three stages, each warm-starting from the previous stage's own checkpoint
-(NOT always from the original SP1 checkpoint -- see `src.warm_start.
-build_or_load_v4` for how `train()`'s `init_ckpt` handles both a genuine SP1
-source and an already-v4 FT2b source):
-
-  1. **scale/alloy adaptation** -- train stems = the v4 train pool restricted
-     to `dataset == "A"` (all-sand, no BC variation). FiLM-only trainable
-     (`freeze_backbone=True`), warm-started from `sp1_ckpt`, `lr=3e-4`. Also
-     refits the scale-head GBM on the same stems -> `runs/ft2b_gbm/hgb.joblib`
-     (the artifact Task 8/9 evaluate against by default).
-  2. **BC generalization** -- train stems = train pool with
-     `dataset in {A, B, C}` (chills / insulating & exothermic sleeves added).
-     Full backbone unfrozen, chained from stage 1's checkpoint, `lr=1e-4`.
-  3. **calibration** -- train stems = train pool rows that are either
-     named-grade (`alloy_family != "synthetic"`) or `dataset == "D"` (real
-     parts). Only the last 2 processor blocks + `scale_head` trainable
-     (`freeze_mode="all_but_last2"`), chained from stage 2's checkpoint,
-     `lr=1e-4`. Refits the GBM again on the calibration stems into a
-     SEPARATE artifact, `runs/ft2b_gbm/hgb_stage3.joblib` -- `hgb.joblib`
-     from stage 1 is intentionally left untouched; Task 8/9's evaluation
-     decides which of the two GBMs the final hybrid predictor should use.
-
-Reality check this driver must survive: in the actual FT2a/FT2b DOE
-manifest + `data/features_v4/splits.json`, named-grade and dataset-D rows
-are ENTIRELY reserved for validation (`val_interp`/`val_extrap`) -- none are
-in the train pool. Stage 3 can therefore legitimately have zero train stems.
-When that happens this driver logs loudly and SKIPS stage 3 (no crash, no
-silent no-op train() call, no stage-3 GBM refit): `paths["final_ckpt"]` falls
-back to stage 2's checkpoint.
-
-All early-stopping/checkpoint-selection across all 3 stages uses the same
-`extrap.test` split (`early_stop_val_stems`), per the FINETUNING plan.
-
-Artifact isolation: this module writes ONLY under
-`runs/ft2b_stage{1,2,3}/` and `runs/ft2b_gbm/` (all nested under `out_root`,
-default `"runs"`).
-"""
-from __future__ import annotations
-
-import csv
-import json
-import logging
-import os
-import os.path as osp
-
-import torch
-
-from src.gbm_ft2b import refit_gbm
-from src.train_gnn import train
-
-logger = logging.getLogger(__name__)
-
-DEFAULT_OUT_ROOT = "runs"
-
-# Loss config the FT2b plan mandates for every stage (train()'s own defaults
-# are graph="knn", global_weight=0.0 -- both wrong for FT2b, so pinned here
-# rather than relying on train()'s defaults ever staying in sync).
-_LOSS_DEFAULTS = {"graph": "union", "global_weight": 0.3, "feature_version": 4}
-
-_STAGE1_DATASETS = {"A"}
-_STAGE2_DATASETS = {"A", "B", "C"}
-
-
-def _load_manifest_rows(manifest_path: str) -> dict:
-    with open(manifest_path, newline="") as f:
-        return {row["run_id"]: row for row in csv.DictReader(f)}
-
-
-def _load_splits(features_dir: str) -> dict:
-    with open(osp.join(features_dir, "splits.json")) as f:
-        return json.load(f)
-
-
-def _stage1_stems(train_pool, manifest_rows, manifest_path: str = None) -> list:
-    """dataset == A (all-sand, no BC/sleeve/chill variation).
-
-    Raises KeyError if any stem in train_pool is missing from manifest_rows.
-    """
-    missing = [s for s in train_pool if s not in manifest_rows]
-    if missing:
-        msg = f"Missing from manifest"
-        if manifest_path:
-            msg += f" {manifest_path}"
-        msg += f": {missing}"
-        raise KeyError(msg)
-    return [s for s in train_pool
-            if manifest_rows[s].get("dataset") in _STAGE1_DATASETS]
-
-
-def _stage2_stems(train_pool, manifest_rows, manifest_path: str = None) -> list:
-    """dataset in {A, B, C} (adds chills + sleeve BC variation).
-
-    Raises KeyError if any stem in train_pool is missing from manifest_rows.
-    """
-    missing = [s for s in train_pool if s not in manifest_rows]
-    if missing:
-        msg = f"Missing from manifest"
-        if manifest_path:
-            msg += f" {manifest_path}"
-        msg += f": {missing}"
-        raise KeyError(msg)
-    return [s for s in train_pool
-            if manifest_rows[s].get("dataset") in _STAGE2_DATASETS]
-
-
-def _stage3_stems(train_pool, manifest_rows, manifest_path: str = None) -> list:
-    """Named-grade rows (any dataset) or dataset == D (real parts).
-
-    Raises KeyError if any stem in train_pool is missing from manifest_rows.
-    """
-    missing = [s for s in train_pool if s not in manifest_rows]
-    if missing:
-        msg = f"Missing from manifest"
-        if manifest_path:
-            msg += f" {manifest_path}"
-        msg += f": {missing}"
-        raise KeyError(msg)
-    stems = []
-    for s in train_pool:
-        row = manifest_rows[s]
-        if row.get("alloy_family", "synthetic") != "synthetic" or row.get("dataset") == "D":
-            stems.append(s)
-    return stems
-
-
-def run_ft2b(features_dir, manifest, sp1_ckpt, out_root: str = DEFAULT_OUT_ROOT,
-             epochs: int = 200, lr_stage1: float = 3e-4, lr_stage2: float = 1e-4,
-             lr_stage3: float = 1e-4, train_kwargs: dict | None = None,
-             gbm_kwargs: dict | None = None) -> dict:
-    """Run the 3-stage FT2b staged training, returning a dict of output paths.
-
-    Args:
-        features_dir: v4 feature cache directory (must contain `splits.json`
-            with `{"interp": {"train": [...], "test": [...]},
-            "extrap": {"train": [...], "test": [...]}}`; `interp.train` ==
-            `extrap.train`, the single training pool).
-        manifest: path to the DOE manifest CSV (`run_id` == stem; needs
-            `dataset` and `alloy_family` columns to partition stems by stage).
-        sp1_ckpt: path to the production SP1 checkpoint (node_in=13, no
-            FiLM) that stage 1 warm-starts from.
-        out_root: base directory; stage dirs and the GBM dir are created
-            under it (`<out_root>/ft2b_stage{1,2,3}/`, `<out_root>/ft2b_gbm/`).
-        epochs: shared epoch budget passed to every `train()` call (early
-            stopping/patience are `train()`'s own defaults unless overridden
-            via `train_kwargs`).
-        lr_stage{1,2,3}: per-stage learning rate.
-        train_kwargs: extra kwargs forwarded to every `train()` call (e.g.
-            `n_blocks`, `hidden`, `batch_size`, `patience`, `num_workers`,
-            `device`) -- must NOT include `lr` (set via `lr_stage*` instead)
-            or the stage-specific wiring args (`init_ckpt`, `freeze_backbone`,
-            `freeze_mode`, `out_dir`, `early_stop_val_stems`, `film`), which
-            this function always sets itself. `graph`/`global_weight`/
-            `feature_version` default to the FT2b plan's values but may be
-            overridden here.
-        gbm_kwargs: extra kwargs forwarded to every `refit_gbm()` call (e.g.
-            `max_iter`, `min_samples_leaf`, `holdout_frac`).
-
-    Returns:
-        dict with `stage{1,2,3}_dir`, `stage{1,2,3}_ckpt` (stage 3's may be
-        `None` if it was skipped -- see module docstring), `final_ckpt` (the
-        last stage actually run), `gbm_stage1_path`/`gbm_stage1_metrics`,
-        and `gbm_stage3_path`/`gbm_stage3_metrics` (also `None` if stage 3
-        was skipped).
-    """
-    train_kwargs = dict(train_kwargs or {})
-    for k, v in _LOSS_DEFAULTS.items():
-        train_kwargs.setdefault(k, v)
-    train_kwargs.setdefault("epochs", epochs)
-    gbm_kwargs = dict(gbm_kwargs or {})
-
-    manifest_rows = _load_manifest_rows(manifest)
-    splits = _load_splits(features_dir)
-    train_pool_interp = splits["interp"]["train"]
-    train_pool_extrap = splits["extrap"]["train"]
-
-    # Guard: interp and extrap train pools must be identical
-    if train_pool_interp != train_pool_extrap:
-        raise ValueError(
-            f"interp['train'] and extrap['train'] must be identical "
-            f"(single training pool assumption). Found {len(train_pool_interp)} "
-            f"interp and {len(train_pool_extrap)} extrap stems with "
-            f"{len(set(train_pool_interp) - set(train_pool_extrap))} only-in-interp "
-            f"and {len(set(train_pool_extrap) - set(train_pool_interp))} only-in-extrap.")
-
-    train_pool = train_pool_interp
-    val_stems = splits["interp"]["test"]
-    early_stop_val_stems = splits["extrap"]["test"]
-
-    gbm_dir = osp.join(out_root, "ft2b_gbm")
-    os.makedirs(gbm_dir, exist_ok=True)
-
-    paths: dict = {}
-
-    # ---- Stage 1: scale/alloy adaptation (FiLM only, warm-start from SP1) ----
-    stage1_dir = osp.join(out_root, "ft2b_stage1")
-    stage1_stems = _stage1_stems(train_pool, manifest_rows, manifest_path=manifest)
-    logger.info("FT2b stage 1 (scale/alloy, dataset=A): %d train stems", len(stage1_stems))
-    if not stage1_stems:
-        raise RuntimeError(
-            "FT2b stage 1: 0 train stems resolved for dataset=A in the train "
-            "pool -- cannot proceed (check the manifest/splits inputs)")
-    stage1_result = train(
-        features_dir, stage1_stems, val_stems, out_dir=stage1_dir, lr=lr_stage1,
-        film=True, init_ckpt=sp1_ckpt, freeze_backbone=True,
-        early_stop_val_stems=early_stop_val_stems, **train_kwargs,
-    )
-    paths["stage1_dir"] = stage1_dir
-    paths["stage1_ckpt"] = stage1_result["ckpt"]
-    paths["final_ckpt"] = paths["stage1_ckpt"]
-
-    hgb_path = osp.join(gbm_dir, "hgb.joblib")
-    paths["gbm_stage1_metrics"] = refit_gbm(
-        features_dir=features_dir, train_stems=stage1_stems,
-        manifest_path=manifest, out_path=hgb_path, **gbm_kwargs)
-    paths["gbm_stage1_path"] = hgb_path
-
-    # ---- Stage 2: BC generalization (full backbone, chained from stage 1) ----
-    stage2_dir = osp.join(out_root, "ft2b_stage2")
-    stage2_stems = _stage2_stems(train_pool, manifest_rows, manifest_path=manifest)
-    logger.info("FT2b stage 2 (BC, dataset in A/B/C): %d train stems", len(stage2_stems))
-    if not stage2_stems:
-        raise RuntimeError(
-            "FT2b stage 2: 0 train stems resolved for dataset in {A,B,C} in "
-            "the train pool -- cannot proceed (check the manifest/splits inputs)")
-    stage2_result = train(
-        features_dir, stage2_stems, val_stems, out_dir=stage2_dir, lr=lr_stage2,
-        film=True, init_ckpt=paths["stage1_ckpt"], freeze_backbone=False,
-        early_stop_val_stems=early_stop_val_stems, **train_kwargs,
-    )
-    paths["stage2_dir"] = stage2_dir
-    paths["stage2_ckpt"] = stage2_result["ckpt"]
-    paths["final_ckpt"] = paths["stage2_ckpt"]
-
-    # ---- Stage 3: calibration (named-grade + real parts) -- may be empty ----
-    stage3_stems = _stage3_stems(train_pool, manifest_rows, manifest_path=manifest)
-    logger.info("FT2b stage 3 (calibrate, named-grade + dataset=D): %d train stems",
-                len(stage3_stems))
-    if not stage3_stems:
-        logger.warning(
-            "FT2b stage 3 SKIPPED: 0 train stems resolved for named-grade/"
-            "dataset-D rows in the train pool (they may be entirely held out "
-            "for validation in this manifest/splits, as in the production "
-            "FT2a v4 cache). Final checkpoint stays stage 2's: %s",
-            paths["stage2_ckpt"])
-        paths["stage3_dir"] = None
-        paths["stage3_ckpt"] = None
-        paths["gbm_stage3_path"] = None
-        paths["gbm_stage3_metrics"] = None
-        return paths
-
-    stage3_dir = osp.join(out_root, "ft2b_stage3")
-    stage3_result = train(
-        features_dir, stage3_stems, val_stems, out_dir=stage3_dir, lr=lr_stage3,
-        film=True, init_ckpt=paths["stage2_ckpt"], freeze_mode="all_but_last2",
-        early_stop_val_stems=early_stop_val_stems, **train_kwargs,
-    )
-    paths["stage3_dir"] = stage3_dir
-    paths["stage3_ckpt"] = stage3_result["ckpt"]
-    paths["final_ckpt"] = paths["stage3_ckpt"]
-
-    hgb_stage3_path = osp.join(gbm_dir, "hgb_stage3.joblib")
-    paths["gbm_stage3_metrics"] = refit_gbm(
-        features_dir=features_dir, train_stems=stage3_stems,
-        manifest_path=manifest, out_path=hgb_stage3_path, **gbm_kwargs)
-    paths["gbm_stage3_path"] = hgb_stage3_path
-
-    return paths
-
-
-def run_stage2b(features_dir, manifest, stage1_ckpt, out_dir: str = "runs/ft2b_stage2b",
-                 oversample_exo: int = 3, lr: float = 1e-4, epochs: int = 200,
-                 train_kwargs: dict | None = None) -> dict:
-    """Exo-channel remediation: repeat FT2b stage 2 (BC generalization) on the
-    fixed feature cache, oversampling exo-sleeve stems so the (previously
-    corrupted) exothermic-sleeve BC channel gets proportionally more gradient
-    signal than its raw row count in the train pool would give it.
-
-    Same stem selection as `run_ft2b`'s stage 2 (`_stage2_stems`: dataset in
-    {A, B, C}), then the train list actually passed to `train()` is that
-    stem list PLUS `(oversample_exo - 1)` extra copies of every stem whose
-    manifest row has `bc_config == "exo_sleeve"` (duplicates appended after
-    the base list, so ordering is deterministic).
-
-    Chains from `stage1_ckpt` (a stage-1 FT2b checkpoint, NOT the original
-    SP1 checkpoint), full backbone unfrozen (`freeze_backbone=False`), same
-    as `run_ft2b`'s stage 2. Early-stop/checkpoint-selection uses
-    `extrap.test`, same as every other FT2b stage.
-
-    Args:
-        features_dir: v4 feature cache directory (must contain `splits.json`
-            with `interp.train == extrap.train`, the single training pool).
-        manifest: path to the DOE manifest CSV (`run_id` == stem; needs
-            `dataset` and `bc_config` columns).
-        stage1_ckpt: path to the checkpoint this run warm-starts/chains from.
-        out_dir: directory `train()` writes its checkpoint into.
-        oversample_exo: total multiplicity for each exo-sleeve stem in the
-            train list (1 = no oversampling; the base stem is always
-            present once, plus `oversample_exo - 1` duplicates).
-        lr: learning rate passed to `train()`.
-        epochs: epoch budget passed to `train()` (overridden by
-            `train_kwargs["epochs"]` if given there).
-        train_kwargs: extra kwargs forwarded to `train()` -- must NOT
-            include `lr` or the stage-wiring args (`init_ckpt`,
-            `freeze_backbone`, `out_dir`, `early_stop_val_stems`, `film`,
-            `val_stems`), which this function always sets itself. `epochs` IS
-            overridable via train_kwargs["epochs"] (setdefault).
-            `graph`/`global_weight`/`feature_version` default to the FT2b
-            plan's values (`_LOSS_DEFAULTS`) but may be overridden here.
-
-    Returns:
-        dict with `ckpt` (the trained checkpoint path), `best_val` (from
-        `train()`'s return), `n_train` (length of the train list, including
-        oversampled duplicates), and `n_exo` (count of distinct exo-sleeve
-        stems that were oversampled).
-    """
-    train_kwargs = dict(train_kwargs or {})
-    for k, v in _LOSS_DEFAULTS.items():
-        train_kwargs.setdefault(k, v)
-    train_kwargs.setdefault("epochs", epochs)
-
-    manifest_rows = _load_manifest_rows(manifest)
-    splits = _load_splits(features_dir)
-    train_pool_interp = splits["interp"]["train"]
-    train_pool_extrap = splits["extrap"]["train"]
-
-    # Guard: interp and extrap train pools must be identical
-    if train_pool_interp != train_pool_extrap:
-        raise ValueError(
-            f"interp['train'] and extrap['train'] must be identical "
-            f"(single training pool assumption). Found {len(train_pool_interp)} "
-            f"interp and {len(train_pool_extrap)} extrap stems with "
-            f"{len(set(train_pool_interp) - set(train_pool_extrap))} only-in-interp "
-            f"and {len(set(train_pool_extrap) - set(train_pool_interp))} only-in-extrap.")
-
-    train_pool = train_pool_interp
-    val_stems = splits["interp"]["test"]
-    early_stop_val_stems = splits["extrap"]["test"]
-
-    stage2_stems = _stage2_stems(train_pool, manifest_rows, manifest_path=manifest)
-    exo_stems = [s for s in stage2_stems if manifest_rows[s].get("bc_config") == "exo_sleeve"]
-    train_list = list(stage2_stems) + exo_stems * (oversample_exo - 1)
-    logger.info(
-        "FT2b stage 2b (exo-oversample, dataset in A/B/C): %d train stems "
-        "(%d distinct exo-sleeve, oversampled x%d)",
-        len(train_list), len(exo_stems), oversample_exo)
-
-    result = train(
-        features_dir, train_list, val_stems=val_stems, out_dir=out_dir, lr=lr,
-        film=True, init_ckpt=stage1_ckpt, freeze_backbone=False,
-        early_stop_val_stems=early_stop_val_stems, **train_kwargs,
-    )
-
-    return {
-        "ckpt": result["ckpt"],
-        "best_val": result["best_val"],
-        "n_train": len(train_list),
-        "n_exo": len(exo_stems),
-    }
-
-
-def run_ft3(features_dir, manifest, base_ckpt: str = "runs/ft2b_stage2c/ckpt_ep054.pt",
-            out_dir: str = "runs/ft3_stage1", oversample_e: int = 3, lr: float = 5e-5,
-            epochs: int = 200, train_kwargs: dict | None = None) -> dict:
-    """FT3: frozen-norm fine-tune of the production v4 checkpoint onto the new
-    dataset-E (multi-body assembly) rows, mixed into the existing FT2b stage-2
-    train pool (dataset in {A, B, C}) with dataset-E stems oversampled.
-
-    Loads `base_ckpt`'s own `norm` (the production checkpoint's frozen
-    feature/target normalization -- recomputing norm_stats from the mixed
-    A+B+C+E pool would shift every standardization the checkpoint was
-    trained with) and forwards it to `train()` as `norm_override`, so every
-    SolidGraphDataset construction inside `train()` (train/val/early-stop-val)
-    standardizes against the SAME stats the checkpoint already expects.
-
-    Train stems = `_stage2_stems` (A+B+C ∩ train pool, same selection as
-    FT2b stage 2/2b) PLUS dataset-E stems in the train pool, each repeated
-    `oversample_e` times. Unlike `run_stage2b`'s exo-sleeve subset (already
-    present once inside its A/B/C stage-2 selection, so it tops up with
-    `oversample_exo - 1` extra copies), dataset "E" is NOT one of
-    `_stage2_stems`'s selected datasets ({A, B, C}) -- so E stems are
-    entirely new content here, added as a straight `x oversample_e` (not
-    `x (oversample_e - 1)`).
-
-    Chains from `base_ckpt` (`init_ckpt=base_ckpt`), full backbone unfrozen
-    (`freeze_backbone=False` -- FT3 fine-tunes the whole model, not just
-    FiLM), `film=True`, `save_every_epoch=True` (per-epoch checkpoints for
-    gate-metric checkpoint selection, mirroring the exo-fix stage-2c
-    precedent). Early-stop/checkpoint-selection uses `extrap.test`, same as
-    every other FT2b/FT3 stage; the `run_ft2b`/`run_stage2b` interp==extrap
-    train-pool guard applies here too (Task 4's `extend_splits_with_e` keeps
-    it true by extending both pools equally with E-train stems).
-
-    Args:
-        features_dir: v4 feature cache directory (`splits.json` with
-            `interp.train == extrap.train`, the single training pool,
-            extended with dataset-E train stems by Task 4's
-            `extend_splits_with_e`).
-        manifest: path to the DOE manifest CSV (`run_id` == stem; needs a
-            `dataset` column, with `"E"` rows for the new assembly castings).
-        base_ckpt: the production checkpoint this run warm-starts/chains
-            from AND whose `norm` is frozen for training (default: the FT2b
-            exo-fix production checkpoint, `runs/ft2b_stage2c/ckpt_ep054.pt`).
-        out_dir: directory `train()` writes its checkpoints into.
-        oversample_e: total multiplicity for each dataset-E train stem in the
-            train list (1 = no oversampling).
-        lr: learning rate passed to `train()`.
-        epochs: epoch budget passed to `train()` (overridden by
-            `train_kwargs["epochs"]` if given there).
-        train_kwargs: extra kwargs forwarded to `train()` -- must NOT include
-            `lr`, `norm_override`, or the stage-wiring args (`init_ckpt`,
-            `freeze_backbone`, `out_dir`, `early_stop_val_stems`, `film`,
-            `val_stems`, `save_every_epoch`), which this function always sets
-            itself. `epochs` IS overridable via `train_kwargs["epochs"]`
-            (setdefault); same for `batch_size` (default 2 to prevent OOM: train()'s
-            own default batch_size=8 OOMs the RTX 3060 12GB, incident 2026-07-31 FT3
-            execution) and `num_workers` (default 4). Explicit caller values in
-            `train_kwargs` override all defaults. `graph`/`global_weight`/
-            `feature_version` default to the FT2b plan's values (`_LOSS_DEFAULTS`)
-            but may be overridden here.
-
-    Returns:
-        dict with `ckpt` (the trained checkpoint path), `best_val` (from
-        `train()`'s return), `n_train` (length of the train list, including
-        oversampled E duplicates), and `n_e` (count of distinct dataset-E
-        stems that were oversampled).
-    """
-    train_kwargs = dict(train_kwargs or {})
-    train_kwargs.setdefault("batch_size", 2)
-    train_kwargs.setdefault("num_workers", 4)
-    for k, v in _LOSS_DEFAULTS.items():
-        train_kwargs.setdefault(k, v)
-    train_kwargs.setdefault("epochs", epochs)
-
-    base_norm = torch.load(base_ckpt, map_location="cpu")["norm"]
-
-    manifest_rows = _load_manifest_rows(manifest)
-    splits = _load_splits(features_dir)
-    train_pool_interp = splits["interp"]["train"]
-    train_pool_extrap = splits["extrap"]["train"]
-
-    # Guard: interp and extrap train pools must be identical
-    if train_pool_interp != train_pool_extrap:
-        raise ValueError(
-            f"interp['train'] and extrap['train'] must be identical "
-            f"(single training pool assumption). Found {len(train_pool_interp)} "
-            f"interp and {len(train_pool_extrap)} extrap stems with "
-            f"{len(set(train_pool_interp) - set(train_pool_extrap))} only-in-interp "
-            f"and {len(set(train_pool_extrap) - set(train_pool_interp))} only-in-extrap.")
-
-    train_pool = train_pool_interp
-    val_stems = splits["interp"]["test"]
-    early_stop_val_stems = splits["extrap"]["test"]
-
-    stage2_stems = _stage2_stems(train_pool, manifest_rows, manifest_path=manifest)
-    e_train_stems = [s for s in train_pool if manifest_rows[s].get("dataset") == "E"]
-    train_list = list(stage2_stems) + e_train_stems * oversample_e
-    logger.info(
-        "FT3 (assembly fine-tune, dataset A/B/C + E oversampled x%d): %d train "
-        "stems (%d distinct dataset-E)", oversample_e, len(train_list), len(e_train_stems))
-
-    result = train(
-        features_dir, train_list, val_stems=val_stems, out_dir=out_dir, lr=lr,
-        film=True, init_ckpt=base_ckpt, freeze_backbone=False,
-        early_stop_val_stems=early_stop_val_stems, save_every_epoch=True,
-        norm_override=base_norm, **train_kwargs,
-    )
-
-    return {
-        "ckpt": result["ckpt"],
-        "best_val": result["best_val"],
-        "n_train": len(train_list),
-        "n_e": len(e_train_stems),
-    }
-
-
-# Loss config for bc-reach (Task 3): same graph/global_weight convention as
-# FT2b/FT3, but `feature_version=5` (the 20-column BC-reach cache, Task 1/2)
-# instead of FT2b/FT3's 4.
-_BCR_LOSS_DEFAULTS = {"graph": "union", "global_weight": 0.3, "feature_version": 5}
-
-
-def run_bcr(features_dir, manifest, base_ckpt: str = "runs/ft3_stage1/ckpt_ep021.pt",
-            out_dir: str = "runs/bcr_stage1", oversample_e: int = 3, lr: float = 5e-5,
-            epochs: int = 200, train_kwargs: dict | None = None) -> dict:
-    """bc-reach: warm-started fine-tune of the production v4 checkpoint onto
-    the v5 (20-column, BC-reach) feature cache -- node_in 15 -> 20.
-
-    Same stem assembly as `run_ft3` (train stems = `_stage2_stems`, A+B+C ∩
-    train pool, PLUS dataset-E train-pool stems each repeated `oversample_e`
-    times), the same `film=True`/`freeze_backbone=False`/`save_every_epoch=True`
-    wiring, and the same `extrap.test` early-stop/checkpoint-selection
-    convention as every other FT2b/FT3/bc-reach stage.
-
-    UNLIKE `run_ft3`, this does NOT pass `norm_override`: `run_ft3` freezes
-    the fine-tune onto its `base_ckpt`'s own (15-wide) norm because it keeps
-    `node_in` unchanged (15 -> 15). Here `node_in` genuinely changes (15 ->
-    20, Task 1/2's 5 new BC-reach columns), so `base_ckpt`'s 15-wide norm
-    stats cannot standardize the v5 cache's 20-wide `X` at all -- `train()`
-    is left to load `features_dir`'s OWN fresh `norm_stats.npz` (Task 2's
-    `build_v5` output), the default `norm_override=None` behavior.
-
-    `train()`'s own `init_ckpt` dispatch (see its docstring) routes this
-    warm-start through `src.warm_start.build_or_load_v5` instead of
-    `build_or_load_v4`, because the actual v5 cache data has `node_in=20`:
-    `base_ckpt` (`node_in=15`, film=True -- e.g. the production
-    `ckpt_ep021.pt`) has its node encoder's first Linear rebuilt for
-    `in_features=20` (existing 15 columns' weights copied, the 5 new
-    trailing columns zero-init, bias copied verbatim); every other weight
-    (edge_encoder, processor blocks, FiLM, decoder, scale_head) is
-    strict-copied. The resulting checkpoint's `cfg` therefore carries
-    `node_in=20` (computed by `train()` from the actual data, same as every
-    other stage) and `feature_version=5` (via `_BCR_LOSS_DEFAULTS` below).
-
-    Args:
-        features_dir: v5 feature cache directory (`splits.json` with
-            `interp.train == extrap.train`, the single training pool --
-            Task 2's `build_v5` copies this verbatim from the v4 cache, so
-            it already carries the same E-train stems FT3's
-            `extend_splits_with_e` added).
-        manifest: path to the DOE manifest CSV (`run_id` == stem; needs a
-            `dataset` column, with `"E"` rows for the assembly castings).
-        base_ckpt: the v4 (node_in=15, film=True) checkpoint this run warm-
-            starts from (default: the production FT3 checkpoint,
-            `runs/ft3_stage1/ckpt_ep021.pt`).
-        out_dir: directory `train()` writes its checkpoints into.
-        oversample_e: total multiplicity for each dataset-E train stem in the
-            train list (1 = no oversampling), same convention as `run_ft3`.
-        lr: learning rate passed to `train()` (default `5e-5`, per the
-            bc-reach plan).
-        epochs: epoch budget passed to `train()` (overridden by
-            `train_kwargs["epochs"]` if given there); `train()`'s own
-            `patience=25` default (checked against `extrap.test`, via
-            `early_stop_val_stems`) is left untouched, per the bc-reach plan
-            ("epochs<=200 patience 25 on extrap.test").
-        train_kwargs: extra kwargs forwarded to `train()` -- must NOT
-            include `lr`, `norm_override`, or the stage-wiring args
-            (`init_ckpt`, `freeze_backbone`, `out_dir`,
-            `early_stop_val_stems`, `film`, `val_stems`, `save_every_epoch`),
-            which this function always sets itself. `epochs` IS overridable
-            via `train_kwargs["epochs"]` (setdefault); same for `batch_size`
-            (default 2, mirroring FT3's OOM-safe pin) and `num_workers`
-            (default 4). Explicit caller values in `train_kwargs` override
-            all defaults. `graph`/`global_weight`/`feature_version` default
-            to `_BCR_LOSS_DEFAULTS` but may be overridden here. Every other
-            disallowed key above fails loudly on its own (a plain `TypeError`
-            from `train()`'s call below, "got multiple values for keyword
-            argument") since this function also passes it explicitly;
-            `norm_override` has no such corresponding explicit call kwarg to
-            collide with (this function never passes it at all), so it is
-            guarded by an explicit `ValueError` instead (review round 1, F1).
-
-    Returns:
-        dict with `ckpt` (the trained checkpoint path), `best_val` (from
-        `train()`'s return), `n_train` (length of the train list, including
-        oversampled E duplicates), and `n_e` (count of distinct dataset-E
-        stems that were oversampled).
-    """
-    train_kwargs = dict(train_kwargs or {})
-    train_kwargs.setdefault("batch_size", 2)
-    train_kwargs.setdefault("num_workers", 4)
-    for k, v in _BCR_LOSS_DEFAULTS.items():
-        train_kwargs.setdefault(k, v)
-    train_kwargs.setdefault("epochs", epochs)
-    # F1 (review round 1): every OTHER pinned kwarg (init_ckpt, freeze_backbone,
-    # out_dir, early_stop_val_stems, film, val_stems, save_every_epoch, lr) already
-    # fails loudly if train_kwargs duplicates it -- train()'s call below raises a
-    # plain TypeError ("got multiple values for keyword argument"). norm_override
-    # is the ONE pinned divergence from run_ft3 (this driver deliberately never
-    # passes it -- v5's fresh norm_stats.npz, not the 15-wide base_ckpt's, must be
-    # what train() standardizes against) that has NO corresponding call kwarg to
-    # collide with, so it would otherwise slip through train_kwargs silently and
-    # freeze the wrong (15-wide, base_ckpt) normalization onto the v5 run. Guard
-    # it explicitly, matching the loud-fail convention every other pinned kwarg
-    # already gets for free.
-    if "norm_override" in train_kwargs:
-        raise ValueError(
-            "run_bcr: train_kwargs must not include 'norm_override' -- v5 uses "
-            "its own fresh norm_stats.npz (Task 2's build_v5), never a frozen "
-            "norm from base_ckpt (unlike run_ft3's frozen-norm fine-tune).")
-
-    manifest_rows = _load_manifest_rows(manifest)
-    splits = _load_splits(features_dir)
-    train_pool_interp = splits["interp"]["train"]
-    train_pool_extrap = splits["extrap"]["train"]
-
-    # Guard: interp and extrap train pools must be identical
-    if train_pool_interp != train_pool_extrap:
-        raise ValueError(
-            f"interp['train'] and extrap['train'] must be identical "
-            f"(single training pool assumption). Found {len(train_pool_interp)} "
-            f"interp and {len(train_pool_extrap)} extrap stems with "
-            f"{len(set(train_pool_interp) - set(train_pool_extrap))} only-in-interp "
-            f"and {len(set(train_pool_extrap) - set(train_pool_interp))} only-in-extrap.")
-
-    train_pool = train_pool_interp
-    val_stems = splits["interp"]["test"]
-    early_stop_val_stems = splits["extrap"]["test"]
-
-    stage2_stems = _stage2_stems(train_pool, manifest_rows, manifest_path=manifest)
-    e_train_stems = [s for s in train_pool if manifest_rows[s].get("dataset") == "E"]
-    train_list = list(stage2_stems) + e_train_stems * oversample_e
-    logger.info(
-        "bc-reach (node_in 15->20, dataset A/B/C + E oversampled x%d): %d train "
-        "stems (%d distinct dataset-E)", oversample_e, len(train_list), len(e_train_stems))
-
-    result = train(
-        features_dir, train_list, val_stems=val_stems, out_dir=out_dir, lr=lr,
-        film=True, init_ckpt=base_ckpt, freeze_backbone=False,
-        early_stop_val_stems=early_stop_val_stems, save_every_epoch=True,
-        **train_kwargs,
-    )
-
-    return {
-        "ckpt": result["ckpt"],
-        "best_val": result["best_val"],
-        "n_train": len(train_list),
-        "n_e": len(e_train_stems),
-    }

src/train_gnn.py

@@ -1,300 +0,0 @@
-"""GNN training: loss, training loop, checkpointing. Trains on standardized
-log1p targets with rotation augmentation; validates on a held-out fraction."""
-import torch
-import torch.nn.functional as F
-from torch_geometric.utils import scatter
-
-
-def gnn_loss(pred, target, batch, spatial_weight: float = 0.3, global_weight: float = 0.0):
-    base = F.smooth_l1_loss(pred, target)
-    pm = scatter(pred, batch, dim=0, reduce="mean")
-    tm = scatter(target, batch, dim=0, reduce="mean")
-    spatial = F.mse_loss(pred - pm[batch], target - tm[batch])
-    glob = F.smooth_l1_loss(pm, tm)
-    return base + spatial_weight * spatial + global_weight * glob
-
-
-import os
-import os.path as osp
-from torch.optim import AdamW
-from torch.optim.lr_scheduler import ReduceLROnPlateau
-from torch_geometric.loader import DataLoader
-
-from src.gnn_dataset import SolidGraphDataset
-from src.model_gnn import SolidGNN
-from src.warm_start import build_or_load_v4, build_or_load_v5
-
-# bc-reach (Task 1/2): the v5 BC-reach feature cache's per-node width --
-# FEATURE_NAMES_V4 (15) + the 5 new multi-scale BC-reach channels
-# (exo_mid/exo_long/chill_short/chill_mid/chill_long, `src.features.
-# FEATURE_NAMES_V5`). `train()`'s `init_ckpt` dispatch below compares the
-# ACTUAL training data's own feature width (`node_in`, computed from the
-# built `tr` dataset -- never the checkpoint's) against this constant to
-# decide whether to warm-start through `build_or_load_v5` (rebuild the node
-# encoder for 20 columns) instead of `build_or_load_v4` (the existing
-# 15-column FT2b/FT3 path, unaffected). Named here (review round 1, F3)
-# instead of an inline literal `20`, so the pin's provenance is legible at
-# the call site.
-_V5_NODE_IN = 20
-
-
-def _freeze_backbone_params(model):
-    """Freeze everything except FiLM (used for FT2b stage-1 training: only the
-    per-block gamma/beta conditioning is trainable, so warm-started SP1 shape
-    knowledge is preserved while the model learns to respond to the material
-    g-vector). The new BC input columns live inside node_encoder's first
-    Linear, so they are frozen too in stage 1 -- they get unfrozen (along with
-    the rest of the backbone) in stage 2 per the FINETUNING plan's stage table.
-    """
-    for name, p in model.named_parameters():
-        p.requires_grad = name.startswith("films.")
-
-
-# Public name for tests/other callers. `train()` below takes a same-named
-# `freeze_backbone` bool parameter, which would shadow this at module scope
-# inside that function's body -- so `train()` calls `_freeze_backbone_params`
-# directly instead of this alias.
-freeze_backbone = _freeze_backbone_params
-
-
-def _freeze_all_but_last2_and_scale_head(model):
-    """FT2b stage-3 (calibrate) freeze scope: everything is frozen except the
-    LAST 2 processor blocks and the scale head -- the small, late-stage
-    capacity that plausibly still needs to adapt to a handful of named-grade
-    (steel/ductile-iron/grey-iron) + real-part calibration rows, without
-    letting so few rows overwrite the shape knowledge learned in stages 1-2.
-    FiLM is deliberately NOT included here (unlike stage 1's `freeze_backbone`,
-    which trains FiLM only) -- stage 3 is calibrating the backbone's late
-    representation + scale, not re-learning material conditioning.
-    """
-    n_blocks = len(model.blocks)
-    keep_prefixes = tuple(f"blocks.{i}." for i in range(max(n_blocks - 2, 0), n_blocks))
-    keep_prefixes = keep_prefixes + ("scale_head.",)
-    for name, p in model.named_parameters():
-        p.requires_grad = name.startswith(keep_prefixes)
-
-
-# Public name, mirroring `freeze_backbone` above.
-freeze_all_but_last2_and_scale_head = _freeze_all_but_last2_and_scale_head
-
-
-def _run_epoch(model, loader, device, opt=None, scaler=None, spatial_weight=0.3, global_weight=0.0):
-    train_mode = opt is not None
-    model.train(train_mode)
-    total, n = 0.0, 0
-    for batch in loader:
-        batch = batch.to(device)
-        g = getattr(batch, "g", None)
-        if train_mode:
-            opt.zero_grad()
-            with torch.autocast(device_type=device.split(":")[0], enabled=scaler is not None):
-                yn, _ = model(batch.x, batch.edge_index, batch.edge_attr, batch.batch, g=g)
-                loss = gnn_loss(yn, batch.y, batch.batch, spatial_weight, global_weight)
-            if scaler is not None:
-                scaler.scale(loss).backward(); scaler.step(opt); scaler.update()
-            else:
-                loss.backward(); opt.step()
-        else:
-            with torch.no_grad():
-                yn, _ = model(batch.x, batch.edge_index, batch.edge_attr, batch.batch, g=g)
-                loss = gnn_loss(yn, batch.y, batch.batch, spatial_weight, global_weight)
-        total += float(loss) * batch.num_graphs; n += batch.num_graphs
-    return total / max(n, 1)
-
-
-def _build_checkpoint_dict(model, norm, hidden, n_blocks, k, node_in, edge_in,
-                            spatial_weight, global_weight, graph, feature_version,
-                            two_head, film, g_dim):
-    """Build the checkpoint dict payload (same format as best_gnn.pt).
-
-    Args:
-        model: the SolidGNN model.
-        norm: normalization stats (SolidGraphDataset.norm).
-        All other args: hyperparameters/config to store.
-
-    Returns:
-        dict with keys {model, norm, cfg}.
-    """
-    return {
-        "model": model.state_dict(),
-        "norm": norm,
-        "cfg": {
-            "hidden": hidden,
-            "n_blocks": n_blocks,
-            "k": k,
-            "node_in": node_in,
-            "edge_in": edge_in,
-            "spatial_weight": spatial_weight,
-            "global_weight": global_weight,
-            "graph": graph,
-            "feature_version": feature_version,
-            "two_head": two_head,
-            "film": film,
-            "g_dim": g_dim,
-        },
-    }
-
-
-def train(features_dir, train_stems, val_stems, epochs=200, k=16, hidden=128,
-          n_blocks=12, batch_size=8, lr=1e-3, device=None, out_dir="runs/gnn",
-          augment=True, patience=25, spatial_weight=0.3, global_weight=0.0,
-          graph="knn", feature_version=2, num_workers=0, two_head=False,
-          film=False, g_dim=5, freeze_backbone: bool = False, init_ckpt=None,
-          early_stop_val_stems=None, freeze_mode=None, save_every_epoch=False,
-          norm_override=None):
-    """
-    film: build the model with FiLM (per-block gamma/beta from data.g); forced
-        True when init_ckpt is given (warm-start is always onto a FiLM model).
-    g_dim: FiLM conditioning-vector width (ignored unless film/init_ckpt).
-    freeze_backbone: freeze node_encoder/edge_encoder/blocks/decoder/scale_head,
-        leaving only FiLM params trainable (FT2b stage-1: shape/scale knowledge
-        from warm-start is preserved while FiLM learns to respond to `g`).
-        Back-compat boolean from Task 5; equivalent to `freeze_mode="backbone"`.
-    freeze_mode: FT2b's more general freeze-scope switch, superset of
-        `freeze_backbone`. One of:
-          - None (default): use `freeze_backbone`'s boolean behavior (freeze
-            everything but FiLM if True, nothing if False) -- unchanged from
-            the pre-FT2b/Task-5 code path.
-          - "backbone": same as `freeze_backbone=True` (only FiLM trainable).
-          - "all_but_last2": only the last 2 processor blocks + scale_head are
-            trainable (FT2b stage 3 / calibrate -- see
-            `_freeze_all_but_last2_and_scale_head`'s docstring).
-        `freeze_mode` takes precedence over `freeze_backbone` when both are given.
-    init_ckpt: path to an SP1 checkpoint (node_in=13, no FiLM) OR an earlier
-        FT2b-stage checkpoint (node_in=15, FiLM already present -- stage 2
-        chains from stage 1's `best_gnn.pt`, stage 3 from stage 2's) to warm-
-        start from via `src.warm_start.build_or_load_v4`, instead of a
-        randomly initialized model. `build_or_load_v4` dispatches on the
-        checkpoint's own `cfg` (splice-from-SP1 vs. direct-load-an-already-v4-
-        checkpoint) -- see its docstring. Implies film=True. bc-reach: when
-        the ACTUAL training data (`node_in`, computed below from the built
-        `tr` dataset's own feature width) equals `_V5_NODE_IN` (20) -- i.e.
-        `features_dir` is a v5 (BC-reach) cache -- this dispatches to
-        `src.warm_start.build_or_load_v5` instead, which rebuilds the node
-        encoder's first Linear for the 5 new trailing columns (node_in 15 ->
-        20) from a v4 `init_ckpt` (e.g. the production `ckpt_ep021.pt`), or
-        direct-loads an already-v5 checkpoint. Every other `node_in` (in
-        particular the existing 15-wide FT2b/FT3 path) is unaffected -- still
-        routed through `build_or_load_v4`, byte-identical to before this
-        dispatch existed.
-        NOTE on normalization: the checkpoint's own `norm` (the SOURCE
-        checkpoint's feature/target stats -- SP1's v3 stats for the original
-        13->15 warm-start, or a v4 checkpoint's 15-wide stats for the
-        15->20 bc-reach one) is NOT reused here -- the training data in
-        `features_dir` (whichever cache version it actually is: v3/v4/v5)
-        is normalized (by `SolidGraphDataset`) with *its own*
-        `norm_stats.npz` (`tr.norm`), since the model always consumes
-        THAT cache's own normalized features, not the source checkpoint's.
-        The saved checkpoint's `norm` field must match what the *data* was
-        actually standardized with, so we always save `tr.norm`, discarding
-        the `_sp1_norm` that `build_or_load_v4`/`build_or_load_v5` return
-        (that value exists for callers who need to reconcile the two spaces
-        explicitly, e.g. a sanity-gate check at warm-start time -- not for
-        the trained ckpt).
-    early_stop_val_stems: if given, checkpoint-selection/early-stopping use a
-        loss computed on this stem list instead of `val_stems` (e.g. a
-        val_extrap split held out for OOD generalization signal). When None,
-        behavior is identical to the pre-FT2b code path (single val set/loader
-        serving both roles).
-    save_every_epoch: if True, save a checkpoint to `<out_dir>/ckpt_ep{ep:03d}.pt`
-        after EVERY epoch (in the same format as best_gnn.pt: {model, norm, cfg}),
-        in addition to the unchanged best_gnn.pt logic. Default False preserves
-        byte-identical behavior (no ckpt_ep* files).
-    norm_override: if given (a norm dict, e.g. a production checkpoint's own
-        `ckpt["norm"]"), freeze normalization to it instead of computing/loading
-        the feature cache's own `norm_stats.npz` (FT3: fine-tuning must not
-        shift any standardization the production checkpoint was trained
-        with). Passed as `norm=norm_override` into `tr`'s SolidGraphDataset
-        construction; `va`/`va_es` already reuse `tr.norm` downstream, so all
-        three dataset constructions end up standardized with `norm_override`,
-        and the saved checkpoint's `norm` field (`tr.norm`) is `norm_override`
-        too. Default None preserves byte-identical behavior (tr loads the
-        cache's own norm_stats.npz, as before).
-    """
-    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-    os.makedirs(out_dir, exist_ok=True)
-    tr = SolidGraphDataset(features_dir, train_stems, k=k, augment=augment, graph=graph,
-                            norm=norm_override)
-    va = SolidGraphDataset(features_dir, val_stems, k=k, augment=False, norm=tr.norm, graph=graph)
-    trl = DataLoader(tr, batch_size=batch_size, shuffle=True, num_workers=num_workers, persistent_workers=num_workers > 0)
-    val = DataLoader(va, batch_size=batch_size, num_workers=num_workers, persistent_workers=num_workers > 0)
-
-    if early_stop_val_stems is not None:
-        va_es = SolidGraphDataset(features_dir, early_stop_val_stems, k=k, augment=False,
-                                   norm=tr.norm, graph=graph)
-        val_es = DataLoader(va_es, batch_size=batch_size, num_workers=num_workers,
-                             persistent_workers=num_workers > 0)
-    else:
-        val_es = None
-
-    sample = tr[0]
-    node_in, edge_in = sample.x.shape[1], sample.edge_attr.shape[1]
-
-    if init_ckpt is not None:
-        film = True
-        if node_in == _V5_NODE_IN:
-            model, _sp1_norm = build_or_load_v5(init_ckpt, hidden=hidden, n_blocks=n_blocks, g_dim=g_dim)
-        else:
-            model, _sp1_norm = build_or_load_v4(init_ckpt, hidden=hidden, n_blocks=n_blocks, g_dim=g_dim)
-        model = model.to(device)
-        # build_or_load_v4 (like build_warm_started_v4 before it) prefers the
-        # *source* checkpoint's own cfg (hidden/n_blocks/two_head) over the
-        # args passed here (two_head isn't even forwarded to it), so
-        # re-derive them from the actual constructed model -- the saved cfg
-        # below must describe the model we trained, not just echo back the
-        # caller's arguments.
-        hidden = model.decoder[0].in_features
-        n_blocks = len(model.blocks)
-        two_head = model.two_head
-    else:
-        model = SolidGNN(node_in=node_in, edge_in=edge_in, hidden=hidden, n_blocks=n_blocks,
-                          two_head=two_head, film=film, g_dim=g_dim).to(device)
-
-    if freeze_mode == "all_but_last2":
-        _freeze_all_but_last2_and_scale_head(model)
-    elif freeze_mode == "backbone" or (freeze_mode is None and freeze_backbone):
-        _freeze_backbone_params(model)
-    elif freeze_mode not in (None,):
-        raise ValueError(f"train(): unknown freeze_mode={freeze_mode!r}")
-
-    trainable_params = [p for p in model.parameters() if p.requires_grad]
-    opt = AdamW(trainable_params, lr=lr, weight_decay=1e-5)
-    sched = ReduceLROnPlateau(opt, mode="min", factor=0.5, patience=10)
-    scaler = torch.cuda.amp.GradScaler() if device.startswith("cuda") else None
-
-    best, best_ep, since = float("inf"), 0, 0
-    ckpt = osp.join(out_dir, "best_gnn.pt")
-    ep = -1
-    for ep in range(epochs):
-        tr_loss = _run_epoch(model, trl, device, opt=opt, scaler=scaler, spatial_weight=spatial_weight, global_weight=global_weight)
-        # Checkpoint selection/early-stop use shape-only val loss: the global term is
-        # noisy at batch=2 and masks shape improvements (Run-1 lesson).
-        if val_es is not None:
-            vloss = _run_epoch(model, val_es, device, spatial_weight=spatial_weight, global_weight=0.0)
-        else:
-            vloss = _run_epoch(model, val, device, spatial_weight=spatial_weight, global_weight=0.0)
-        sched.step(vloss)
-        lr = opt.param_groups[0]["lr"]
-        print(f"  epoch {ep:3d}  train {tr_loss:.4f}  val {vloss:.4f}  lr {lr:.2e}", flush=True)
-
-        # Build checkpoint dict (used for both best_gnn.pt and per-epoch saves)
-        ckpt_dict = _build_checkpoint_dict(
-            model, tr.norm, hidden, n_blocks, k, node_in, edge_in,
-            spatial_weight, global_weight, graph, feature_version,
-            two_head, film, g_dim
-        )
-
-        # Save per-epoch checkpoint if requested
-        if save_every_epoch:
-            ep_ckpt = osp.join(out_dir, f"ckpt_ep{ep:03d}.pt")
-            torch.save(ckpt_dict, ep_ckpt)
-
-        # Save best checkpoint if this is a new best
-        if vloss < best - 1e-5:
-            best, best_ep, since = vloss, ep, 0
-            torch.save(ckpt_dict, ckpt)
-        else:
-            since += 1
-            if since >= patience:
-                break
-    return {"best_val": best, "epochs_run": ep + 1, "ckpt": ckpt}

src/train_heads.py

@@ -1,278 +0,0 @@
-"""Learned-heads Task 2: frozen-backbone trainer for the gradT/dTdt deriv
-heads (Task 1's `SolidGNN(deriv_heads=True)` / `attach_deriv_heads`).
-
-Loads a headless production checkpoint (`{model, norm, cfg}`), attaches the
-two new heads (`grad_head`/`dtdt_head`), freezes EVERY other parameter, and
-trains only the ~8 head tensors to regress standardized log1p(gradT) /
-log1p(|dTdt|) targets read straight from the v4 feature cache (each
-`data/features_v4/graphs/*.npz` already carries raw per-node `gradT` (K/m)
-and `dTdt` (K/s) arrays, see `src/build_dataset.py`). Base weights are never
-touched (verified by an assertion here, and by the caller's own
-byte-identity check).
-"""
-import json
-import logging
-import os
-import os.path as osp
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torch.optim import AdamW
-from torch_geometric.loader import DataLoader
-
-from src.gnn_dataset import SolidGraphDataset
-from src.model_gnn import SolidGNN, attach_deriv_heads
-
-logger = logging.getLogger(__name__)
-
-
-class DerivGraphDataset(SolidGraphDataset):
-    """`SolidGraphDataset` plus a standardized `y_deriv` (N, 2) node target:
-    `[(log1p(gradT) - mean) / std, (log1p(|dTdt|) - mean) / std]`, read from
-    the same graph cache. Callers must pre-filter `stems` to ones known to
-    carry valid `gradT`/`dTdt` (see `_filter_stems_with_deriv`) -- this class
-    does not re-validate.
-    """
-
-    def __init__(self, features_dir, stems, deriv_mean, deriv_std, k=16,
-                 graph="union", norm=None):
-        super().__init__(features_dir, stems, k=k, augment=False, norm=norm, graph=graph)
-        self.deriv_mean = torch.tensor(np.asarray(deriv_mean), dtype=torch.float32)
-        self.deriv_std = torch.tensor(np.asarray(deriv_std), dtype=torch.float32)
-
-    def __getitem__(self, i):
-        data = super().__getitem__(i)
-        raw = np.load(osp.join(self.dir, "graphs", f"{self.stems[i]}.npz"))
-        gradT = torch.tensor(raw["gradT"], dtype=torch.float32).clamp(min=0.0)
-        dTdt = torch.tensor(raw["dTdt"], dtype=torch.float32).abs()
-        deriv_log = torch.stack([torch.log1p(gradT), torch.log1p(dTdt)], dim=-1)
-        data.y_deriv = (deriv_log - self.deriv_mean) / self.deriv_std
-        return data
-
-
-def _stem_has_valid_deriv(features_dir, stem):
-    """A stem is usable if its cached graph has finite gradT/dTdt arrays,
-    each with exactly one value per node (matching X's row count -- a
-    length-mismatched stem would otherwise broadcast silently into the
-    loss). A missing `.npz` file is also treated as an exclusion (counted
-    and logged by the caller), not a raw crash."""
-    path = osp.join(features_dir, "graphs", f"{stem}.npz")
-    try:
-        d = np.load(path)
-    except FileNotFoundError:
-        return False
-    if "gradT" not in d.files or "dTdt" not in d.files or "X" not in d.files:
-        return False
-    gradT, dTdt, X = d["gradT"], d["dTdt"], d["X"]
-    if gradT.shape != dTdt.shape or gradT.shape[0] != X.shape[0]:
-        return False
-    return bool(np.all(np.isfinite(gradT)) and np.all(np.isfinite(dTdt)))
-
-
-def _filter_stems_with_deriv(features_dir, stems):
-    """Split `stems` into (ok, n_excluded), logging loudly (never silently
-    dropping) every stem missing/invalid gradT/dTdt."""
-    ok, excluded = [], []
-    for stem in stems:
-        if _stem_has_valid_deriv(features_dir, stem):
-            ok.append(stem)
-        else:
-            excluded.append(stem)
-    if excluded:
-        logger.warning(
-            "train_heads: excluding %d/%d stem(s) missing/invalid gradT/dTdt: %s",
-            len(excluded), len(stems), excluded)
-    return ok, len(excluded)
-
-
-def _compute_deriv_stats(features_dir, stems):
-    """Streamed (single accumulation pass, one stem's arrays in memory at a
-    time -- never the whole train-pool's per-node values at once) mean/std
-    of `[log1p(gradT), log1p(|dTdt|)]` over every node in `stems`."""
-    n = 0
-    s1 = np.zeros(2, dtype=np.float64)
-    s2 = np.zeros(2, dtype=np.float64)
-    for stem in stems:
-        d = np.load(osp.join(features_dir, "graphs", f"{stem}.npz"))
-        gradT = np.clip(d["gradT"].astype(np.float64), 0.0, None)
-        dTdt = np.abs(d["dTdt"].astype(np.float64))
-        vals = np.stack([np.log1p(gradT), np.log1p(dTdt)], axis=-1)
-        n += vals.shape[0]
-        s1 += vals.sum(axis=0)
-        s2 += (vals ** 2).sum(axis=0)
-    if n == 0:
-        raise ValueError("train_heads: no valid stems to compute deriv_log_mean/std from")
-    mean = s1 / n
-    var = np.maximum(s2 / n - mean ** 2, 1e-12)
-    std = np.sqrt(var)
-    return mean.astype(np.float32), std.astype(np.float32)
-
-
-def _freeze_all_but_deriv_heads(model):
-    """Freeze every parameter except `grad_head.*`/`dtdt_head.*`. Exposed
-    standalone (mirrors `train_gnn._freeze_all_but_last2_and_scale_head`) so
-    the freeze scope can be unit-tested directly."""
-    for name, p in model.named_parameters():
-        p.requires_grad = name.startswith("grad_head.") or name.startswith("dtdt_head.")
-
-
-def _run_epoch_heads(model, loader, device, opt=None):
-    train_mode = opt is not None
-    model.train(train_mode)
-    total, n = 0.0, 0
-    for batch in loader:
-        batch = batch.to(device)
-        g = getattr(batch, "g", None)
-        if train_mode:
-            opt.zero_grad()
-            _, _, deriv = model(batch.x, batch.edge_index, batch.edge_attr, batch.batch, g=g)
-            loss = F.mse_loss(deriv, batch.y_deriv)
-            loss.backward()
-            opt.step()
-        else:
-            with torch.no_grad():
-                _, _, deriv = model(batch.x, batch.edge_index, batch.edge_attr, batch.batch, g=g)
-                loss = F.mse_loss(deriv, batch.y_deriv)
-        total += float(loss) * batch.num_graphs
-        n += batch.num_graphs
-    return total / max(n, 1)
-
-
-def train_deriv_heads(features_dir: str, base_ckpt: str, out_dir: str = "runs/lh_stage1",
-                       train_stems=None, val_stems=None, epochs: int = 60, lr: float = 1e-3,
-                       batch_size: int = 2, num_workers: int = 4, patience: int = 10,
-                       device=None) -> dict:
-    """Freeze the entire base model and train only the deriv (gradT/dTdt) heads.
-
-    Args:
-        features_dir: v4 feature cache dir (`graphs/*.npz` carry X/pos/y/
-            mesh_edges/g plus raw `gradT`/`dTdt`).
-        base_ckpt: path to a headless `{model, norm, cfg}` production
-            checkpoint (Task 1's `attach_deriv_heads` contract: only
-            grad_head/dtdt_head keys may be missing from it).
-        out_dir: where `heads_gnn.pt` (best-val checkpoint) and per-epoch
-            `ckpt_ep{ep:03d}.pt` snapshots are written.
-        train_stems, val_stems: stem lists; default to `splits.json`'s
-            `interp.train` / **`interp.test`** (NOT `extrap.test`: per the
-            spec's 2026-08-01 Task-2-review correction, `extrap.test` is
-            72% of the Gate C evaluation set, so early-stopping/checkpoint
-            selection on it would leak into that pre-authorized gate;
-            `interp.test` is disjoint from both Gate C evaluation lists
-            (`extrap.test` + `e_val`)).
-        epochs, lr, batch_size, num_workers, patience: as in `src.train_gnn.train`.
-        device: torch device string; default cuda if available else cpu.
-
-    Freezes ALL base parameters (`requires_grad=False`), asserting the
-    optimizer's trainable-parameter set is EXACTLY the head tensors
-    (`grad_head.*`/`dtdt_head.*`). Computes `deriv_log_mean`/`deriv_log_std`
-    over the (filtered) train pool ONCE, streamed, and stores them in the
-    saved checkpoint's `norm` dict (base norm + these two (2,) arrays). Loss
-    is MSE over the two standardized heads. Stems whose cached graph is
-    missing/length-mismatched gradT/dTdt (or has non-finite values), or
-    whose `.npz` file is simply absent, are excluded loudly and counted in
-    `n_excluded`, rather than silently dropped or crashing the run. The
-    graph mode/kNN-k are read verbatim from the base checkpoint's own cfg
-    (`cfg["graph"]`/`cfg["k"]`, no divergent defaults -- a cfg missing
-    either raises `KeyError` loudly rather than silently training/serving
-    with a mismatched graph construction). Saves a checkpoint every epoch
-    (`<out_dir>/ckpt_ep{ep:03d}.pt`, gate-metric checkpoint-selection
-    precedent from the exo-fix retrain) IN ADDITION to early-stopping on val
-    loss and keeping the MSE-best as `<out_dir>/heads_gnn.pt`, with `cfg` =
-    base cfg + `{"deriv_heads": True}`.
-
-    Returns:
-        {"ckpt": str, "best_val": float, "n_train": int, "n_excluded": int}
-    """
-    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-    os.makedirs(out_dir, exist_ok=True)
-
-    if train_stems is None or val_stems is None:
-        with open(osp.join(features_dir, "splits.json")) as f:
-            splits = json.load(f)
-        if train_stems is None:
-            train_stems = splits["interp"]["train"]
-        if val_stems is None:
-            val_stems = splits["interp"]["test"]
-
-    train_ok, n_exc_train = _filter_stems_with_deriv(features_dir, train_stems)
-    val_ok, n_exc_val = _filter_stems_with_deriv(features_dir, val_stems)
-    n_excluded = n_exc_train + n_exc_val
-    if not train_ok:
-        raise ValueError("train_deriv_heads: no train stems with valid gradT/dTdt")
-    if not val_ok:
-        raise ValueError("train_deriv_heads: no val stems with valid gradT/dTdt")
-
-    deriv_mean, deriv_std = _compute_deriv_stats(features_dir, train_ok)
-
-    base = torch.load(base_ckpt, map_location="cpu", weights_only=False)
-    cfg = dict(base["cfg"])
-    base_norm = base["norm"]
-
-    model = SolidGNN(node_in=cfg.get("node_in", 15), edge_in=cfg.get("edge_in", 4),
-                      hidden=cfg["hidden"], n_blocks=cfg["n_blocks"],
-                      two_head=cfg.get("two_head", False), film=cfg.get("film", False),
-                      g_dim=cfg.get("g_dim", 5), deriv_heads=True)
-    attach_deriv_heads(model, base["model"])
-    model = model.to(device)
-
-    _freeze_all_but_deriv_heads(model)
-
-    trainable_names = {n for n, p in model.named_parameters() if p.requires_grad}
-    expected_names = {n for n, _ in model.named_parameters()
-                       if n.startswith("grad_head.") or n.startswith("dtdt_head.")}
-    assert trainable_names == expected_names, (
-        f"train_deriv_heads: optimizer param set must be exactly the deriv-head "
-        f"tensors -- unexpected trainable: {sorted(trainable_names - expected_names)}, "
-        f"missing: {sorted(expected_names - trainable_names)}")
-    trainable_params = [p for n, p in model.named_parameters() if n in trainable_names]
-
-    opt = AdamW(trainable_params, lr=lr, weight_decay=1e-5)
-
-    # No divergent .get(default) here: the graph mode/k must match whatever
-    # the base checkpoint was actually trained/normalized with, or the deriv
-    # heads would be trained (and later served) on a mismatched graph
-    # construction. A base cfg missing either key is a loud KeyError.
-    graph = cfg["graph"]
-    k = cfg["k"]
-
-    tr_ds = DerivGraphDataset(features_dir, train_ok, deriv_mean, deriv_std,
-                              k=k, graph=graph, norm=base_norm)
-    va_ds = DerivGraphDataset(features_dir, val_ok, deriv_mean, deriv_std,
-                              k=k, graph=graph, norm=base_norm)
-    trl = DataLoader(tr_ds, batch_size=batch_size, shuffle=True,
-                      num_workers=num_workers, persistent_workers=num_workers > 0)
-    val = DataLoader(va_ds, batch_size=batch_size, num_workers=num_workers,
-                      persistent_workers=num_workers > 0)
-
-    saved_norm = dict(base_norm)
-    saved_norm["deriv_log_mean"] = deriv_mean
-    saved_norm["deriv_log_std"] = deriv_std
-    saved_cfg = dict(cfg)
-    saved_cfg["deriv_heads"] = True
-
-    ckpt_path = osp.join(out_dir, "heads_gnn.pt")
-    best, since = float("inf"), 0
-    ep = -1
-    for ep in range(epochs):
-        tr_loss = _run_epoch_heads(model, trl, device, opt=opt)
-        vloss = _run_epoch_heads(model, val, device, opt=None)
-        print(f"  epoch {ep:3d}  train {tr_loss:.4f}  val {vloss:.4f}", flush=True)
-
-        ckpt_dict = {"model": model.state_dict(), "norm": saved_norm, "cfg": saved_cfg}
-
-        # Per-epoch checkpoint (gate-metric checkpoint-selection precedent:
-        # the exo-fix retrain needed to compare 8 candidates), independent
-        # of the best-val bookkeeping below.
-        torch.save(ckpt_dict, osp.join(out_dir, f"ckpt_ep{ep:03d}.pt"))
-
-        if vloss < best - 1e-6:
-            best, since = vloss, 0
-            torch.save(ckpt_dict, ckpt_path)
-        else:
-            since += 1
-            if since >= patience:
-                break
-
-    return {"ckpt": ckpt_path, "best_val": best, "n_train": len(train_ok),
-            "n_excluded": n_excluded}

src/utils.py

@@ -1,115 +0,0 @@
-import re
-from pathlib import Path
-
-ALL_FACES = ['+X', '-X', '+Y', '-Y', '+Z', '-Z']
-_FACE_RE  = re.compile(r'^[+\-][XYZ]$')
-
-# Integer encoding for geometry type (used as global feature)
-GEOMETRY_ENCODING = {
-    'BC':  0, 'RIB': 1, 'L':   2, 'T':   3, 'X':   4,
-    'WL':  5, 'UL':  6, 'UM':  7, 'US':  8, 'BH':  9,
-    'BO': 10, 'PB': 11, 'PF': 12, 'PH': 13,
-    'RPB':14, 'RPF':15, 'RPH':16,
-    # everything else → 17 (other)
-}
-GEOMETRY_OTHER = 17
-
-
-def parse_filename(path: str) -> dict:
-    """Parse any <PREFIX>_<tokens...>.pt filename used in the Foundry dataset.
-
-    Token classification (applied to tokens after the prefix):
-      - numeric (all digits)  → appended to dims list
-      - matches [+-][XYZ]     → adiabatic face
-      - anything else         → ignored (angle suffixes like 30, 45, 135
-                                 are already absorbed into the prefix token)
-
-    Returns:
-        prefix          — geometry type string (e.g. 'BC', 'RIB', 'L')
-        dims            — list of numeric dim values in mm  (may be 0–3 ints)
-        adiabatic_faces — list of face labels that are adiabatic
-        cooling_faces   — complement of adiabatic_faces over ALL_FACES
-        volume_mm3      — product of dims (mm³); 0.0 if no dims
-        geometry_type   — integer encoding (GEOMETRY_ENCODING, 17 = other)
-    """
-    stem   = Path(path).stem
-    tokens = stem.split('_')
-    prefix = tokens[0]
-
-    dims       = []
-    adiabatic  = []
-    for tok in tokens[1:]:
-        if tok.isdigit():
-            dims.append(int(tok))
-        elif _FACE_RE.match(tok):
-            adiabatic.append(tok)
-        # else: skip (e.g. stray letters, already part of prefix)
-
-    cooling  = [f for f in ALL_FACES if f not in set(adiabatic)]
-    volume   = float(eval('*'.join(str(d) for d in dims))) if dims else 0.0
-    geo_type = GEOMETRY_ENCODING.get(prefix, GEOMETRY_OTHER)
-
-    return {
-        'prefix':          prefix,
-        'dims':            dims,
-        'adiabatic_faces': adiabatic,
-        'cooling_faces':   cooling,
-        'volume_mm3':      volume,
-        'geometry_type':   geo_type,
-    }
-
-
-def global_feature_vec(info: dict) -> list:
-    """6-element global feature vector from a parse_filename result.
-
-    Layout: [dim1, dim2, dim3, volume_mm3, n_cooling_faces, geometry_type]
-    dims are zero-padded to length 3.
-    """
-    d = (info['dims'] + [0, 0, 0])[:3]
-    return [
-        float(d[0]),
-        float(d[1]),
-        float(d[2]),
-        info['volume_mm3'],
-        float(len(info['cooling_faces'])),
-        float(info['geometry_type']),
-    ]
-
-
-if __name__ == '__main__':
-    test_cases = [
-        # BC — 3 numeric dims
-        'BC_06_08_14_+X_+Z.pt',
-        'BC_06_08_14.pt',
-        # RIB — 3 dims
-        'RIB_06_06_06.pt',
-        # L — 2 dims
-        'L_06_06.pt',
-        # T — 2 dims with faces
-        'T_06_08_+X_+Y.pt',
-        # X — 2 dims
-        'X_06_06_+Y.pt',
-        # WL — 2 dims
-        'WL_06_06.pt',
-        # BH — 2 dims
-        'BH_06_06_+X.pt',
-        # CYL — 1 dim
-        'CYL_06.pt',
-        # KON — 1 dim
-        'KON_25.pt',
-        # angle-suffix prefix
-        'L135_06_06.pt',
-        'TR30_06_06.pt',
-    ]
-
-    header = f"{'File':<35s} {'prefix':<8} {'dims':<14} {'adiab':<18} {'vol':>8} {'geo':>4}"
-    print(header)
-    print('-' * len(header))
-    for p in test_cases:
-        r = parse_filename(p)
-        gf = global_feature_vec(r)
-        print(
-            f"{Path(p).name:<35s} {r['prefix']:<8} {str(r['dims']):<14} "
-            f"{str(r['adiabatic_faces']):<18} {r['volume_mm3']:>8.0f} {r['geometry_type']:>4d}"
-        )
-        print(f"  global_features: {gf}")

src/warm_start.py

@@ -1,239 +0,0 @@
-"""Warm-start loader: build a v4 (node_in=15, FiLM) SolidGNN from an SP1
-(node_in=13) checkpoint with zero learned knowledge lost.
-
-The v4 model adds two trailing per-node input columns (boundary_extraction,
-exothermic_input) and FiLM conditioning (per-block gamma/beta from a material
-g-vector). Both extensions are zero-init: the new input columns are zeroed in
-the node-encoder's first Linear, and FiLM's last layer is zero-init (see
-`src.model_gnn.FiLM`) so gamma=0, beta=0 => identity. Together this means the
-freshly warm-started v4 model reproduces the SP1 model's output exactly (see
-`sanity_gate` / tests/test_warm_start.py), before any FT2b training happens.
-
-torch 2.3 note: `load_state_dict(..., strict=False)` does NOT silently skip
-shape-mismatched tensors -- it raises `RuntimeError: size mismatch` even with
-strict=False (verified empirically). So we pre-filter the SP1 state_dict to
-keys whose shape matches the v4 model before calling load_state_dict, and
-splice the one shape-mismatched tensor (the node_encoder's first Linear
-weight) by hand.
-"""
-import logging
-
-import torch
-
-from src.model_gnn import SolidGNN
-
-logger = logging.getLogger(__name__)
-
-_ENC_W = "node_encoder.0.0.weight"
-_ENC_B = "node_encoder.0.0.bias"
-
-
-def splice_node_encoder(v4, sp1_sd, n_old=13):
-    """Copy SP1's node-encoder first-Linear weight into v4's leading n_old
-    input columns; zero the trailing (new) columns. Bias copied verbatim.
-
-    Args:
-        v4: SolidGNN instance (node_in > n_old).
-        sp1_sd: SP1 state_dict (or any mapping with the same keys).
-        n_old: number of leading input columns coming from SP1 (13).
-    """
-    with torch.no_grad():
-        w_new = v4.node_encoder[0][0].weight
-        w_new[:, :n_old] = sp1_sd[_ENC_W]
-        w_new[:, n_old:] = 0.0
-        v4.node_encoder[0][0].bias.copy_(sp1_sd[_ENC_B])
-
-
-def build_warm_started_v4(sp1_ckpt_path, hidden=128, n_blocks=12, g_dim=5,
-                           node_in_new=15, n_old=13):
-    """Build a film=True, two_head=True SolidGNN(node_in=node_in_new) warm-
-    started from an SP1 checkpoint (node_in=n_old).
-
-    All shape-matching parameters (edge_encoder, processor blocks, decoder,
-    scale_head) are loaded verbatim from SP1. The node_encoder's first Linear
-    is spliced (SP1 columns 0:n_old, zeros n_old:). FiLM stays at its
-    zero-init (identity) construction.
-
-    Args:
-        sp1_ckpt_path: path to the SP1 `{model, norm, cfg}` checkpoint.
-        hidden, n_blocks, g_dim: v4 model hyperparameters. If the checkpoint's
-            `cfg` carries `hidden`/`n_blocks`/`edge_in`/`two_head`, those take
-            precedence (the warm-started model must match SP1's architecture
-            to load its weights) -- the arguments here are fallback defaults
-            for checkpoints without a `cfg`.
-        node_in_new: v4 input feature count (13 v3 features + 2 new BC cols).
-        n_old: SP1 input feature count.
-
-    Returns:
-        (model, sp1_norm): the warm-started v4 SolidGNN (train mode, as
-        constructed) and the SP1 checkpoint's `norm` dict (feature/target
-        normalization stats), so callers can decide how to reconcile it with
-        v4's own (different) normalization -- not this function's concern.
-    """
-    ckpt = torch.load(sp1_ckpt_path, map_location="cpu", weights_only=False)
-    sp1_sd = ckpt["model"]
-    cfg = ckpt.get("cfg", {})
-    edge_in = cfg.get("edge_in", 4)
-    two_head = cfg.get("two_head", True)
-    hidden = cfg.get("hidden", hidden)
-    n_blocks = cfg.get("n_blocks", n_blocks)
-
-    v4 = SolidGNN(node_in=node_in_new, edge_in=edge_in, hidden=hidden,
-                  n_blocks=n_blocks, two_head=two_head, film=True, g_dim=g_dim)
-
-    v4_sd = v4.state_dict()
-    matching = {k: v for k, v in sp1_sd.items()
-                if k in v4_sd and v4_sd[k].shape == v.shape}
-    load_result = v4.load_state_dict(matching, strict=False)
-
-    # Everything skipped by the shape filter above must be exactly the
-    # spliced input-encoder tensors (handled below by splice_node_encoder).
-    skipped_sp1_keys = [k for k in sp1_sd if k not in matching]
-    unexpected_skips = set(skipped_sp1_keys) - {_ENC_W, _ENC_B}
-    if unexpected_skips:
-        raise RuntimeError(
-            f"warm_start: unexpected SP1 keys skipped by shape filter "
-            f"(expected only the node_encoder input layer): {sorted(unexpected_skips)}"
-        )
-
-    # Everything load_state_dict reports "missing" in v4 must be exactly the
-    # spliced input-encoder tensors plus the new FiLM params (no SP1 counterpart).
-    unexpected_missing = [
-        k for k in load_result.missing_keys
-        if k not in (_ENC_W, _ENC_B) and not k.startswith("films.")
-    ]
-    if unexpected_missing:
-        raise RuntimeError(
-            f"warm_start: unexpected missing keys after load "
-            f"(expected only node_encoder input layer + films.*): {unexpected_missing}"
-        )
-    if load_result.unexpected_keys:
-        raise RuntimeError(
-            f"warm_start: unexpected extra keys in SP1 checkpoint: "
-            f"{load_result.unexpected_keys}"
-        )
-
-    logger.info("warm_start: SP1 keys spliced by hand (input encoder): %s", skipped_sp1_keys)
-    logger.info("warm_start: new v4-only keys left at zero-init (FiLM): %s",
-                [k for k in load_result.missing_keys if k.startswith("films.")])
-
-    splice_node_encoder(v4, sp1_sd, n_old=n_old)
-
-    return v4, ckpt["norm"]
-
-
-def build_or_load_v4(ckpt_path, hidden=128, n_blocks=12, g_dim=5,
-                      node_in_new=15, n_old=13):
-    """Warm-start entry point used by FT2b's staged-training driver, which
-    chains checkpoints across stages (stage 2 warm-starts from stage 1's own
-    output, stage 3 from stage 2's) -- NOT always from the original SP1
-    checkpoint.
-
-    Dispatches on the checkpoint's own `cfg`:
-      - `cfg["node_in"] == node_in_new and cfg["film"]` => this is ALREADY a
-        v4 FT2b checkpoint (e.g. a previous stage's `best_gnn.pt`): same
-        architecture end to end, so it is loaded directly with a plain
-        `strict=True` `load_state_dict` -- no splicing, since there is no
-        column-count mismatch to bridge.
-      - otherwise => this is an SP1-shaped checkpoint (`node_in=n_old`, no
-        FiLM), so it is warm-started via `build_warm_started_v4` exactly as
-        before (the node-encoder splice + zero-init FiLM path).
-
-    `build_warm_started_v4` itself is intentionally left untouched -- it always
-    assumes an SP1-shaped source and unconditionally splices, which would
-    corrupt (or crash on, due to the column-count mismatch) an
-    already-15-column source encoder. This function is the one place that
-    decides which path applies.
-
-    Returns:
-        (model, norm): same contract as `build_warm_started_v4` -- the
-        `norm` is the SOURCE checkpoint's own `norm` dict, for callers to
-        reconcile against the current training data's normalization (as
-        `train()` already does for the build_warm_started_v4 path).
-    """
-    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
-    cfg = ckpt.get("cfg", {})
-    if cfg.get("node_in") == node_in_new and cfg.get("film"):
-        edge_in = cfg.get("edge_in", 4)
-        two_head = cfg.get("two_head", True)
-        hidden = cfg.get("hidden", hidden)
-        n_blocks = cfg.get("n_blocks", n_blocks)
-        model = SolidGNN(node_in=node_in_new, edge_in=edge_in, hidden=hidden,
-                          n_blocks=n_blocks, two_head=two_head, film=True, g_dim=g_dim)
-        model.load_state_dict(ckpt["model"], strict=True)
-        logger.info("build_or_load_v4: %s is already a v4 FT2b checkpoint "
-                    "(node_in=%d, film=True) -- loaded directly, no splice.",
-                    ckpt_path, node_in_new)
-        return model, ckpt["norm"]
-    return build_warm_started_v4(ckpt_path, hidden=hidden, n_blocks=n_blocks,
-                                  g_dim=g_dim, node_in_new=node_in_new, n_old=n_old)
-
-
-def build_or_load_v5(ckpt_path, hidden=128, n_blocks=12, g_dim=5,
-                      node_in_new=20, n_old=15):
-    """bc-reach warm-start entry point: rebuild the node encoder for the 5 new
-    BC-reach columns (node_in 15 -> 20), or load an already-v5 checkpoint
-    directly.
-
-    `build_or_load_v4`'s dispatch logic (splice-from-a-smaller-node_in-source
-    vs. direct-load-an-already-this-shape checkpoint, `torch 2.3` shape-filter
-    pre-check, `load_state_dict(..., strict=False)` caveat) is already fully
-    generic in its `node_in_new`/`n_old` arguments -- it makes no assumption
-    that the source is specifically SP1-shaped, only that its node_encoder's
-    first-Linear weight is the sole shape mismatch. This function is a thin,
-    v5-flavored entry point over it: given the production `ckpt_ep021.pt`
-    (node_in=15, film=True -- itself already a fine-tuned FT2b/FT3 checkpoint,
-    NOT a fresh warm-start, so its FiLM/blocks/decoder/scale_head weights are
-    real learned values, not zero-init), it:
-      - copies `weight[:, :15]` (the node encoder's first Linear) verbatim,
-      - zero-inits the 5 new trailing columns (`weight[:, 15:20]`),
-      - copies the bias verbatim (unaffected by the input-width change --
-        Linear bias only depends on `hidden`, never on `node_in`),
-      - strict-copies everything else (edge_encoder, processor blocks, FiLM,
-        decoder, scale_head) verbatim, since `hidden`/`n_blocks`/`g_dim` are
-        unchanged from v4 to v5.
-    Dispatches to a direct `strict=True` load instead (no splice) when
-    `ckpt_path`'s own `cfg` already says `node_in==node_in_new and film` --
-    e.g. a later bc-reach stage chaining off an earlier bc-reach checkpoint.
-
-    Returns:
-        (model, norm): same contract as `build_or_load_v4` -- `norm` is the
-        SOURCE checkpoint's own `norm` dict (15-wide for a genuine v4 source),
-        for callers to reconcile against the CURRENT training data's own
-        (fresh, 20-wide) normalization -- `src.train_ft2b.run_bcr` does NOT
-        forward this as a `norm_override` (unlike FT3's frozen-norm
-        fine-tune): bc-reach's v5 feature cache gets its own fresh norm_stats
-        (Task 2, `build_v5`), so `train()` standardizes against that cache's
-        own stats, not `ckpt_ep021`'s 15-wide ones.
-    """
-    return build_or_load_v4(ckpt_path, hidden=hidden, n_blocks=n_blocks, g_dim=g_dim,
-                             node_in_new=node_in_new, n_old=n_old)
-
-
-def sanity_gate(model_v4, sp1_model, sample):
-    """Run both models (eval, no_grad) on the same reference sample; return
-    the max abs diff between their outputs (y_nodal and y_global combined).
-
-    `sample` is a dict with keys:
-        x_v4:       (N, node_in_v4) input features for the v4 model.
-        x_sp1:      (N, node_in_sp1) input features for the SP1 model.
-        edge_index: (2, E) shared graph connectivity.
-        edge_attr:  (E, edge_in) shared edge features.
-        batch:      (N,) graph-membership vector, shared.
-        g:          (B, g_dim) material g-vector for v4's FiLM, optional.
-
-    This function makes no assumption about how x_v4/x_sp1 relate (e.g.
-    whether x_v4 is x_sp1 zero-padded, or real v4-normalized features) --
-    that reconciliation is the caller's job.
-    """
-    model_v4.eval()
-    sp1_model.eval()
-    with torch.no_grad():
-        yn_sp1, yg_sp1 = sp1_model(sample["x_sp1"], sample["edge_index"],
-                                    sample["edge_attr"], sample["batch"])
-        yn_v4, yg_v4 = model_v4(sample["x_v4"], sample["edge_index"],
-                                 sample["edge_attr"], sample["batch"],
-                                 g=sample.get("g"))
-    node_diff = (yn_v4 - yn_sp1).abs().max().item()
-    global_diff = (yg_v4 - yg_sp1).abs().max().item()
-    return max(node_diff, global_diff)

src/webapp.py

@@ -1,5 +1,5 @@
-"""Local trame + PyVista web app: upload a STEP casting, run the Phase 2A hybrid
-prediction, and view the solidification field with a time slider (solidified =
+"""Local trame + PyVista web app: upload a STEP casting, run the prediction,
+and view the solidification field with a time slider (solidified =
 ghost, liquid = opaque) and a hot-spot toggle. Run: python -m src.webapp
 
 A multi-solid STEP is meshed through the same prep pipeline as the CLI, which

tests/test_features.py

@@ -0,0 +1,55 @@
+import numpy as np
+import pytest
+
+from src.features import FEATURE_NAMES, alloy_g_vector, node_features
+from src.msh_parse import parse_msh
+
+MOLD_EFF = 1018.0
+
+
+def test_shapes_and_order(mini_msh_file):
+    mesh = parse_msh(mini_msh_file)
+    f = node_features(mesh, MOLD_EFF)
+    assert f["X"].shape == (5, len(FEATURE_NAMES)) == (5, 15)
+    assert f["X"].dtype == np.float32
+    assert f["pos"].shape == (5, 3)
+    assert f["mesh_edges"].shape == (6, 2)          # one tet -> 6 edges
+
+
+def test_flags_and_boundary_channels(mini_msh_file):
+    X = node_features(parse_msh(mini_msh_file), MOLD_EFF)["X"]
+    col = {n: X[:, i] for i, n in enumerate(FEATURE_NAMES)}
+    # nodes 0..3 lie on the tagged surface, node 4 (index 4) is not in any tet
+    assert col["is_cooling"][:4].tolist() == [1, 1, 1, 1]
+    assert col["is_adiabatic"].sum() == 0           # adiabatic nodes are also cooling here
+    assert col["is_interior"][4] == 1
+    np.testing.assert_allclose(col["boundary_extraction"][:4], MOLD_EFF)
+    assert col["boundary_extraction"][4] == 0
+    assert np.all(col["boundary_heat_input"] == 0)
+
+
+def test_global_columns_are_constant(mini_msh_file):
+    X = node_features(parse_msh(mini_msh_file), MOLD_EFF)["X"]
+    for name in ("volume_mm3", "surface_area_mm2", "cooling_area_frac"):
+        c = X[:, FEATURE_NAMES.index(name)]
+        assert np.all(c == c[0])
+    assert X[0, FEATURE_NAMES.index("volume_mm3")] == pytest.approx(1000.0 / 6.0)
+
+
+def test_translation_invariance(mini_msh_file):
+    mesh = parse_msh(mini_msh_file)
+    a = node_features(mesh, MOLD_EFF)["X"]
+    mesh.node_xyz = mesh.node_xyz + np.array([100.0, -50.0, 7.0])
+    b = node_features(mesh, MOLD_EFF)["X"]
+    np.testing.assert_allclose(a, b, rtol=1e-5, atol=1e-4)
+
+
+def test_g_vector():
+    props = {"k": 32.9, "rho": 6800.0, "cp": 515.0, "L_kJ_kg": 220.0,
+             "TL": 1160.0, "TS": 1120.0, "Tpour": 1260.0, "mold_effusivity": 1250.0}
+    g = alloy_g_vector(props)
+    assert g.shape == (5,)
+    assert g[0] == pytest.approx(515.0 * 140.0 / 220000.0)        # Stefan number
+    assert g[1] == pytest.approx(0.04) and g[3] == pytest.approx(0.1)
+    assert g[2] == pytest.approx(32.9 / (6800.0 * 515.0) / 1e-5)
+    assert g[4] == pytest.approx(1.0)

tests/test_features_inference.py

@@ -1,16 +0,0 @@
-import numpy as np
-from src.msh_parse import parse_msh
-from src.features import compute_features, compute_features_inference, FEATURE_NAMES
-
-
-def test_inference_matches_training_features(mini_msh_file):
-    mesh = parse_msh(mini_msh_file)
-    N = mesh.node_xyz.shape[0]
-    times = np.array([0.0, 10.0]); temps = np.full((2, N), 1200.0); temps[1] = 1100.0
-    full = compute_features(mesh, temps, times)
-    inf = compute_features_inference(mesh)
-    assert inf["X"].shape == (N, len(FEATURE_NAMES))
-    assert "y" not in inf
-    np.testing.assert_array_equal(inf["X"], full["X"])       # identical feature columns
-    np.testing.assert_array_equal(inf["pos"], full["pos"])
-    assert np.isfinite(inf["X"]).all()

tests/test_features_v3.py

@@ -1,36 +0,0 @@
-import numpy as np
-
-from src.msh_parse import parse_msh
-from src.features import (_node_features, compute_features_inference,
-                          FEATURE_NAMES, FEATURE_NAMES_V3, LOCAL_MODULUS_RADII)
-
-
-def test_v3_names_extend_v2():
-    assert FEATURE_NAMES_V3[:9] == FEATURE_NAMES and len(FEATURE_NAMES_V3) == 13
-    assert LOCAL_MODULUS_RADII == (5.0, 15.0, 40.0)
-
-
-def test_v3_first_columns_match_v2(mini_msh_file):
-    mesh = parse_msh(mini_msh_file)
-    X2, pos2, e2 = _node_features(mesh, version=2)
-    X3, pos3, e3 = _node_features(mesh, version=3)
-    assert X2.shape[1] == 9 and X3.shape[1] == 13
-    assert np.allclose(X3[:, :9], X2) and np.array_equal(e2, e3)
-
-
-def test_geodesic_d_adiab_values(mini_msh_file):
-    mesh = parse_msh(mini_msh_file)
-    X3, _, _ = _node_features(mesh, version=3)
-    col = FEATURE_NAMES_V3.index("geodesic_d_adiab")
-    # 0-based nodes 1,2,3 lie on the Adiabatic triangle -> 0; node 0 reaches one
-    # of them along a 10 mm tet edge
-    assert np.allclose(X3[[1, 2, 3], col], 0.0)
-    assert np.isclose(X3[0, col], 10.0)
-
-
-def test_inference_dict_has_mesh_edges(mini_msh_file):
-    mesh = parse_msh(mini_msh_file)
-    out = compute_features_inference(mesh, version=3)
-    assert out["X"].shape[1] == 13
-    me = out["mesh_edges"]
-    assert me.ndim == 2 and me.shape[1] == 2 and me.dtype == np.int64

tests/test_features_v4.py

@@ -1,275 +0,0 @@
-import numpy as np
-import pytest
-from scipy.spatial import Delaunay
-
-from src.msh_parse import parse_msh, Mesh
-from src.geometry import edges_from_tets
-from src.features import (
-    _node_features, compute_features,
-    FEATURE_NAMES, FEATURE_NAMES_V3, FEATURE_NAMES_V4,
-    G_VECTOR_NAMES, T_REF, ALPHA_REF, MOLD_EFF_REF,
-    EXO_DIFFUSION_LAMBDA_MM,
-    alloy_g_vector, chvorinov_t_ref,
-)
-
-MOLD_EFFUSIVITY = 1250.0
-CHILL_EFFUSIVITY = 11000.0
-SLEEVE_MEF = 2.0
-EXO_Q0 = 5000.0
-
-
-def _base_spec(bc_config="none", exo_q0=0.0):
-    return {
-        "mold_effusivity": MOLD_EFFUSIVITY,
-        "bc_config": bc_config,
-        "chill_effusivity": CHILL_EFFUSIVITY,
-        "sleeve_MEF": SLEEVE_MEF,
-        "exo_q0": exo_q0,
-        # alloy props (for g / t_ref -- only needed by compute_features tests)
-        "k": 35.0, "rho": 7100.0, "cp": 500.0, "L_kJ_kg": 220.0,
-        "TL": 1180.0, "TS": 1140.0, "Tpour": 1380.0,
-        "C_material": 6.5e5,
-    }
-
-
-def _cube_with_bc_patch(n=5, a=20.0):
-    """A Delaunay-tetrahedralized cube where the z=0 face is tagged
-    'Adiabatic' (the BC-patch footprint for chills/sleeves) and every other
-    hull face is tagged 'Cooling'. Interior points of the z=0 face (not on
-    a cube edge) are touched ONLY by Adiabatic triangles -- a clean
-    adiab_only region, unlike the tiny conftest mini_msh where every
-    Adiabatic node is also a Cooling node."""
-    lin = np.linspace(0.0, a, n)
-    pts = np.array([[x, y, z] for x in lin for y in lin for z in lin], float)
-    dt = Delaunay(pts)
-    tets = dt.simplices.astype(np.int64)
-    hull = dt.convex_hull.astype(np.int64)
-    is_bottom = np.isclose(pts[hull, 2], 0.0).all(axis=1)
-    adiab_tri = hull[is_bottom]
-    cool_tri = hull[~is_bottom]
-    mesh = Mesh(node_xyz=pts, tris={"Cooling": cool_tri, "Adiabatic": adiab_tri}, tets=tets)
-    interior_bottom = (np.isclose(pts[:, 2], 0.0)
-                       & (pts[:, 0] > 1e-9) & (pts[:, 0] < a - 1e-9)
-                       & (pts[:, 1] > 1e-9) & (pts[:, 1] < a - 1e-9))
-    on_cool = ~np.isclose(pts[:, 2], 0.0) | ~interior_bottom
-    return mesh, interior_bottom
-
-
-# --------------------------------------------------------------------- #
-# Names / shape contract
-# --------------------------------------------------------------------- #
-def test_v4_names_extend_v3():
-    assert FEATURE_NAMES_V4[:9] == FEATURE_NAMES
-    assert FEATURE_NAMES_V4[:13] == FEATURE_NAMES_V3
-    assert len(FEATURE_NAMES_V4) == 15
-    assert FEATURE_NAMES_V4[13:] == ["boundary_extraction", "exothermic_input"]
-
-
-def test_g_vector_names_and_constants():
-    assert G_VECTOR_NAMES == ["Ste", "dTf_norm", "alpha_ratio", "superheat_norm", "mold_eff_norm"]
-    assert T_REF == 1000.0
-    assert ALPHA_REF == 1e-5
-    assert MOLD_EFF_REF == 1250.0
-
-
-# --------------------------------------------------------------------- #
-# version==4 requires spec
-# --------------------------------------------------------------------- #
-def test_version4_without_spec_raises(mini_msh_file):
-    mesh = parse_msh(mini_msh_file)
-    with pytest.raises(ValueError):
-        _node_features(mesh, version=4)
-
-
-def test_compute_features_version4_without_spec_raises(mini_msh_file):
-    mesh = parse_msh(mini_msh_file)
-    N = mesh.node_xyz.shape[0]
-    times = np.array([0.0, 10.0]); temps = np.full((2, N), 1200.0); temps[1] = 1100.0
-    with pytest.raises(ValueError):
-        compute_features(mesh, temps, times, version=4)
-
-
-# --------------------------------------------------------------------- #
-# Backward-compat: all-sand (bc_config="none", exo_q0=0)
-# --------------------------------------------------------------------- #
-def test_v4_backward_compat_first13_columns_match_v3(mini_msh_file):
-    mesh = parse_msh(mini_msh_file)
-    spec = _base_spec(bc_config="none", exo_q0=0.0)
-    X4, pos4, e4 = _node_features(mesh, version=4, spec=spec)
-    X3, pos3, e3 = _node_features(mesh, version=3)
-    assert X4.shape[1] == 15
-    assert np.array_equal(X4[:, :13], X3)
-    assert np.array_equal(pos4, pos3) and np.array_equal(e4, e3)
-
-
-def test_v4_backward_compat_boundary_extraction_and_exo(mini_msh_file):
-    mesh = parse_msh(mini_msh_file)
-    spec = _base_spec(bc_config="none", exo_q0=0.0)
-    X4, _, _ = _node_features(mesh, version=4, spec=spec)
-    idx = {n: i for i, n in enumerate(FEATURE_NAMES_V4)}
-    is_interior = X4[:, idx["is_interior"]].astype(bool)
-    extraction = X4[:, idx["boundary_extraction"]]
-    exo = X4[:, idx["exothermic_input"]]
-    # all boundary nodes get the baseline mold effusivity, interior nodes 0
-    assert np.allclose(extraction[~is_interior], MOLD_EFFUSIVITY)
-    assert np.allclose(extraction[is_interior], 0.0)
-    assert np.allclose(exo, 0.0)
-
-
-def test_v4_backward_compat_on_cube_with_distinct_adiab_patch():
-    # Same all-sand check, but on a mesh with a genuinely distinct Adiabatic
-    # footprint (not overlapping Cooling), to make sure "none" really treats
-    # the BC patch as ordinary mold (matches fenicsx_solver's bc_config="none"
-    # branch: patch shares the SAME h_mold term).
-    mesh, interior_bottom = _cube_with_bc_patch()
-    spec = _base_spec(bc_config="none", exo_q0=0.0)
-    X4, _, _ = _node_features(mesh, version=4, spec=spec)
-    X3, _, _ = _node_features(mesh, version=3)
-    assert np.array_equal(X4[:, :13], X3)
-    idx = {n: i for i, n in enumerate(FEATURE_NAMES_V4)}
-    extraction = X4[:, idx["boundary_extraction"]]
-    is_interior = X4[:, idx["is_interior"]].astype(bool)
-    assert np.allclose(extraction[interior_bottom], MOLD_EFFUSIVITY)
-    assert np.allclose(extraction[is_interior], 0.0)
-    assert np.allclose(X4[:, idx["exothermic_input"]], 0.0)
-
-
-# --------------------------------------------------------------------- #
-# BC channel mapping: chills / insulating_sleeve / exo_sleeve
-# --------------------------------------------------------------------- #
-def test_v4_chills_mapping():
-    mesh, interior_bottom = _cube_with_bc_patch()
-    spec = _base_spec(bc_config="chills", exo_q0=0.0)
-    X4, _, _ = _node_features(mesh, version=4, spec=spec)
-    idx = {n: i for i, n in enumerate(FEATURE_NAMES_V4)}
-    extraction = X4[:, idx["boundary_extraction"]]
-    is_interior = X4[:, idx["is_interior"]].astype(bool)
-    assert np.allclose(extraction[interior_bottom], CHILL_EFFUSIVITY)
-    assert np.allclose(extraction[is_interior], 0.0)
-    assert np.allclose(X4[:, idx["exothermic_input"]], 0.0)
-    # cooling-only nodes still get the mold baseline
-    cool_only = ~is_interior & ~interior_bottom
-    assert np.allclose(extraction[cool_only], MOLD_EFFUSIVITY)
-
-
-def test_v4_insulating_sleeve_mapping():
-    mesh, interior_bottom = _cube_with_bc_patch()
-    spec = _base_spec(bc_config="insulating_sleeve", exo_q0=0.0)
-    X4, _, _ = _node_features(mesh, version=4, spec=spec)
-    idx = {n: i for i, n in enumerate(FEATURE_NAMES_V4)}
-    extraction = X4[:, idx["boundary_extraction"]]
-    expected = MOLD_EFFUSIVITY / max(SLEEVE_MEF ** 2, 1.0)
-    assert np.allclose(extraction[interior_bottom], expected)
-    assert np.allclose(X4[:, idx["exothermic_input"]], 0.0)
-
-
-def test_v4_exo_sleeve_mapping():
-    # NEW contract (geodesic-diffused field, not a sparse indicator): peak
-    # exo_q0/1e6 AT the patch nodes, decaying smoothly with geodesic distance
-    # -- see test_exo_channel_is_diffused_field/test_exo_channel_zero_for_non_exo
-    # below for the full shape/decay assertions. This test still checks that
-    # boundary_extraction (unrelated column) keeps its sleeve-scaled mapping
-    # and that the peak value at the patch is exactly exo_q0/1e6.
-    mesh, interior_bottom = _cube_with_bc_patch()
-    spec = _base_spec(bc_config="exo_sleeve", exo_q0=EXO_Q0)
-    X4, _, _ = _node_features(mesh, version=4, spec=spec)
-    idx = {n: i for i, n in enumerate(FEATURE_NAMES_V4)}
-    extraction = X4[:, idx["boundary_extraction"]]
-    exo = X4[:, idx["exothermic_input"]]
-    expected = MOLD_EFFUSIVITY / max(SLEEVE_MEF ** 2, 1.0)
-    assert np.allclose(extraction[interior_bottom], expected)
-    assert np.isclose(exo[interior_bottom].max(), EXO_Q0 / 1e6, rtol=1e-5)
-    # diffused, not sparse: some nodes OUTSIDE the patch are nonzero too.
-    outside_patch = ~interior_bottom
-    assert (exo[outside_patch] > 0).any()
-
-
-def test_exo_channel_is_diffused_field():
-    mesh, interior_bottom = _cube_with_bc_patch()
-    X, _, _ = _node_features(mesh, version=4, spec=_base_spec("exo_sleeve", exo_q0=EXO_Q0))
-    ei = X[:, FEATURE_NAMES_V4.index("exothermic_input")]
-    peak = EXO_Q0 / 1e6
-    assert np.isclose(ei[interior_bottom].max(), peak, rtol=1e-5)   # peak at patch
-    far = np.isclose(mesh.node_xyz[:, 2], 20.0)                     # opposite face
-    assert (ei[far] < peak * np.exp(-20.0 / EXO_DIFFUSION_LAMBDA_MM) * 1.5).all()  # decayed far away
-    assert (ei[far] > 0).any()                                      # but still nonzero (diffused)
-    mid = np.isclose(mesh.node_xyz[:, 2], 10.0)
-    assert ei[mid].max() < ei[interior_bottom].max()                # monotone trend
-    assert ei[mid].max() > ei[far].max()
-
-
-def test_exo_channel_zero_for_non_exo():
-    mesh, _ = _cube_with_bc_patch()
-    for bc in ("none", "chills", "insulating_sleeve"):
-        X, _, _ = _node_features(mesh, version=4, spec=_base_spec(bc))
-        assert (X[:, FEATURE_NAMES_V4.index("exothermic_input")] == 0).all()
-
-
-def test_v4_sleeve_mef_floor_at_one():
-    # max(MEF**2, 1) guards MEF < 1 from amplifying extraction above baseline
-    mesh, interior_bottom = _cube_with_bc_patch()
-    spec = _base_spec(bc_config="insulating_sleeve", exo_q0=0.0)
-    spec["sleeve_MEF"] = 0.5
-    X4, _, _ = _node_features(mesh, version=4, spec=spec)
-    idx = {n: i for i, n in enumerate(FEATURE_NAMES_V4)}
-    extraction = X4[:, idx["boundary_extraction"]]
-    assert np.allclose(extraction[interior_bottom], MOLD_EFFUSIVITY)  # divided by 1, not 0.25
-
-
-# --------------------------------------------------------------------- #
-# alloy_g_vector
-# --------------------------------------------------------------------- #
-def test_alloy_g_vector_hand_computed():
-    props = {
-        "k": 35.0, "rho": 7100.0, "cp": 500.0, "L_kJ_kg": 220.0,
-        "TL": 1180.0, "TS": 1140.0, "Tpour": 1380.0,
-        "mold_effusivity": MOLD_EFFUSIVITY,
-    }
-    g = alloy_g_vector(props)
-    assert g.shape == (5,) and g.dtype == np.float64
-    expected_Ste = 500.0 * (1380.0 - 1140.0) / (220.0 * 1000.0)
-    expected_dTf_norm = (1180.0 - 1140.0) / T_REF
-    expected_alpha_ratio = (35.0 / (7100.0 * 500.0)) / ALPHA_REF
-    expected_superheat_norm = (1380.0 - 1180.0) / T_REF
-    expected_mold_eff_norm = MOLD_EFFUSIVITY / MOLD_EFF_REF
-    np.testing.assert_allclose(
-        g, [expected_Ste, expected_dTf_norm, expected_alpha_ratio,
-            expected_superheat_norm, expected_mold_eff_norm])
-
-
-# --------------------------------------------------------------------- #
-# chvorinov_t_ref
-# --------------------------------------------------------------------- #
-def test_chvorinov_t_ref_monotonic_in_modulus_squared():
-    C = 6.5e5
-    area = 100.0
-    t_small = chvorinov_t_ref(volume_mm3=500.0, area_mm2=area, C_material=C)
-    t_large = chvorinov_t_ref(volume_mm3=2000.0, area_mm2=area, C_material=C)
-    assert t_large > t_small
-    # exact formula
-    assert np.isclose(t_small, C * (500.0 / area) ** 2)
-    assert np.isclose(t_large, C * (2000.0 / area) ** 2)
-
-
-# --------------------------------------------------------------------- #
-# compute_features(version=4): g / t_ref / y_nondim
-# --------------------------------------------------------------------- #
-def test_compute_features_v4_adds_g_tref_ynondim(mini_msh_file):
-    mesh = parse_msh(mini_msh_file)
-    N = mesh.node_xyz.shape[0]
-    times = np.array([0.0, 10.0]); temps = np.full((2, N), 1200.0); temps[1] = 1100.0
-    spec = _base_spec(bc_config="none", exo_q0=0.0)
-    out = compute_features(mesh, temps, times, version=4, spec=spec)
-    assert out["X"].shape == (N, 15)
-    assert "y" in out and "g" in out and "t_ref" in out and "y_nondim" in out
-    assert out["g"].shape == (5,) and out["g"].dtype == np.float64
-    expected_g = alloy_g_vector(spec)
-    np.testing.assert_allclose(out["g"], expected_g)
-
-    volume = out["X"][0, FEATURE_NAMES.index("volume_mm3")]
-    area = out["X"][0, FEATURE_NAMES.index("surface_area_mm2")]
-    expected_t_ref = chvorinov_t_ref(volume, area, spec["C_material"])
-    assert np.isclose(out["t_ref"], expected_t_ref)
-
-    expected_y_nondim = out["y"] / max(expected_t_ref, 1e-9)
-    np.testing.assert_allclose(out["y_nondim"], expected_y_nondim, rtol=1e-5)

tests/test_film.py

@@ -1,46 +0,0 @@
-import torch
-from src.model_gnn import FiLM, SolidGNN
-
-
-def test_film_zero_init_is_identity():
-    f = FiLM(g_dim=5, hidden=16)
-    h = torch.randn(7, 16); g = torch.randn(7, 5)
-    assert torch.allclose(f(h, g), h, atol=1e-6)
-
-
-def test_solidgnn_film_zero_init_matches_no_film():
-    """SolidGNN(film=True) with zero-init FiLM params must reproduce
-    SolidGNN(film=False) output exactly, given identical base weights,
-    identical inputs, a real batch vector, and a real (non-None) g tensor."""
-    N, E, B = 20, 120, 3
-    x = torch.randn(N, 9)
-    edge_index = torch.randint(0, N, (2, E))
-    edge_attr = torch.randn(E, 4)
-    batch = torch.tensor([0] * 7 + [1] * 7 + [2] * 6, dtype=torch.long)
-    g = torch.randn(B, 5)
-
-    torch.manual_seed(42)
-    base = SolidGNN(n_blocks=3, hidden=16, film=False)
-
-    torch.manual_seed(42)
-    filmed = SolidGNN(n_blocks=3, hidden=16, film=True, g_dim=5)
-    # base weights (encoder/blocks/decoder) are seeded identically; FiLM layers
-    # are zero-init by construction (asserted separately in the test above), so
-    # filmed's non-FiLM parameters must match base's exactly.
-    for (n1, p1), (n2, p2) in zip(base.named_parameters(), filmed.named_parameters()):
-        assert n1 == n2
-        torch.testing.assert_close(p1, p2)
-
-    base.eval(); filmed.eval()
-    with torch.no_grad():
-        yn_base, yg_base = base(x, edge_index, edge_attr, batch)
-        yn_film, yg_film = filmed(x, edge_index, edge_attr, batch, g=g)
-
-    torch.testing.assert_close(yn_film, yn_base, atol=1e-6, rtol=1e-6)
-    torch.testing.assert_close(yg_film, yg_base, atol=1e-6, rtol=1e-6)
-
-    # also check g=None on the filmed model is identical (old positional call still works)
-    with torch.no_grad():
-        yn_film_noG, yg_film_noG = filmed(x, edge_index, edge_attr, batch)
-    torch.testing.assert_close(yn_film_noG, yn_base, atol=1e-6, rtol=1e-6)
-    torch.testing.assert_close(yg_film_noG, yg_base, atol=1e-6, rtol=1e-6)

tests/test_gbm_heads.py

@@ -0,0 +1,55 @@
+import numpy as np
+import pytest
+
+from src import gbm_heads
+from src.features import FEATURE_NAMES
+
+PROPS = {"k": 32.9, "rho": 6800.0, "cp": 515.0, "L_kJ_kg": 220.0,
+         "TL": 1160.0, "TS": 1120.0, "Tpour": 1260.0, "mold_effusivity": 1018.0}
+
+
+def _X(n=10, vol=8000.0, area=2400.0, n_adiab=0):
+    X = np.zeros((n, len(FEATURE_NAMES)), np.float32)
+    X[:, FEATURE_NAMES.index("volume_mm3")] = vol
+    X[:, FEATURE_NAMES.index("surface_area_mm2")] = area
+    X[:6, FEATURE_NAMES.index("is_cooling")] = 1
+    X[6:6 + n_adiab, FEATURE_NAMES.index("is_adiabatic")] = 1
+    return X
+
+
+def test_total_row_layout():
+    row = gbm_heads.total_row(PROPS, _X())
+    assert row.shape == (gbm_heads.N_TOTAL_FEATURES,)
+    assert row[:7].tolist() == [32.9, 6800.0, 515.0, 220.0, 1160.0, 1120.0, 1260.0]
+    assert row[7] == pytest.approx(8000.0 / 2400.0)
+    assert row[8] == 1018.0
+    assert row[9:].tolist() == [1, 0, 0, 0, 0, 0, 0]
+
+
+def test_spread_row_appends_adiabatic_fraction():
+    X = _X(n_adiab=2)
+    row = gbm_heads.spread_row(gbm_heads.total_row(PROPS, X), X)
+    assert row.shape == (gbm_heads.N_SPREAD_FEATURES,)
+    assert row[-1] == pytest.approx(2 / 8)
+
+
+class _Stub:
+    def __init__(self, value, n):
+        self.value, self.n_features_in_ = value, n
+
+    def predict(self, x):
+        return np.full(len(x), self.value)
+
+
+def test_predictions_are_expm1_of_log_output():
+    row = np.zeros(16)
+    assert gbm_heads.predict_total_seconds(_Stub(np.log1p(500.0), 16), row) == pytest.approx(500.0)
+    assert gbm_heads.predict_spread_seconds(_Stub(-5.0, 17), np.zeros(17)) == 0.0
+    with pytest.raises(ValueError):
+        gbm_heads.predict_total_seconds(_Stub(-1.0, 16), row)
+
+
+def test_check_head_rejects_wrong_width():
+    gbm_heads.check_head(_Stub(0.0, 16), 16, "total-time")
+    with pytest.raises(ValueError):
+        gbm_heads.check_head(_Stub(0.0, 9), 16, "total-time")

tests/test_geometry.py

@@ -1,9 +1,6 @@
 import numpy as np
-import pytest
-from unittest.mock import patch
 from src.geometry import (triangle_areas, total_area, mesh_volume,
-                          nearest_distance, boundary_node_mask,
-                          geodesic_diameter_two_sweep)
+                          nearest_distance, boundary_node_mask)
 
 
 def test_triangle_area_right_triangle():
@@ -37,68 +34,3 @@ def test_boundary_node_mask():
     tri = np.array([[0, 1, 2], [1, 2, 3]])
     m = boundary_node_mask(5, tri)
     np.testing.assert_array_equal(m, [True, True, True, True, False])
-
-
-# --------------------------------------------------------------------- #
-# geodesic_diameter_two_sweep (bc-reach re-review N1): direct unit tests.
-# The bc-reach features-v5 tests only ever use this helper as an
-# expected-value oracle on real/synthetic meshes -- a "d2 = d1" single-
-# sweep mutant (skip the second Dijkstra entirely, just report sweep-1's
-# own max) passes all of them, because none of those fixtures put the
-# fixed `start=0` at a graph MIDPOINT where a single sweep provably can't
-# see the true diameter. These tests target that property directly.
-# --------------------------------------------------------------------- #
-def _mid_start_path_graph(edge_len_mm=10.0):
-    """A 5-node path graph -- the chain order (by position) is
-    node1-node2-node0-node3-node4 -- where node INDEX 0 (this function's
-    default `start`) sits at the graph's MIDDLE, not at either end. True
-    diameter (node1 <-> node4, both ends) = 4 hops = 4*edge_len_mm. A
-    single Dijkstra sweep from node0 only reaches 2 hops in either
-    direction (both ends are tied at that distance) -- it structurally
-    cannot see the two ends' mutual distance -- so a "d2 = d1" mutant
-    (reporting sweep-1's own max instead of re-sweeping from sweep-1's
-    farthest node) returns 2*edge_len_mm here: HALF the true diameter."""
-    edges = np.array([[1, 2], [2, 0], [0, 3], [3, 4]], dtype=np.int64)
-    edge_len = np.full(4, edge_len_mm, dtype=float)
-    return 5, edges, edge_len
-
-
-def test_geodesic_diameter_two_sweep_mid_start_finds_true_diameter():
-    n_nodes, edges, edge_len = _mid_start_path_graph(edge_len_mm=10.0)
-    d_geo = geodesic_diameter_two_sweep(n_nodes, edges, edge_len)  # default start=0
-    # 4 hops * 10mm = 40.0 -- a single-sweep-from-the-middle mutant would
-    # instead report 20.0 (2 hops, the farthest EITHER end appears from
-    # the midpoint), so this assertion fails under that mutant.
-    assert d_geo == pytest.approx(40.0)
-
-
-def test_geodesic_diameter_two_sweep_edgeless_returns_zero():
-    edges = np.empty((0, 2), dtype=np.int64)
-    edge_len = np.empty(0, dtype=float)
-    assert geodesic_diameter_two_sweep(5, edges, edge_len) == 0.0
-
-
-def test_geodesic_diameter_two_sweep_zero_nodes_returns_zero():
-    edges = np.empty((0, 2), dtype=np.int64)
-    edge_len = np.empty(0, dtype=float)
-    assert geodesic_diameter_two_sweep(0, edges, edge_len) == 0.0
-
-
-def test_geodesic_diameter_two_sweep_no_finite_distances_returns_zero():
-    # The documented "no finite distances" guard is unreachable through any
-    # REAL Dijkstra call (a node's distance to itself is always 0, so
-    # sweep-1 always has at least one finite entry) -- exercised directly
-    # here by forcing the underlying primitive to report all-inf, to prove
-    # the guard degrades to 0.0 instead of crashing (e.g. np.argmax over an
-    # all-False finite mask) if that invariant is ever violated.
-    n_nodes, edges, edge_len = _mid_start_path_graph()
-    with patch("src.geometry.geodesic_distance",
-              return_value=np.full(n_nodes, np.inf)):
-        assert geodesic_diameter_two_sweep(n_nodes, edges, edge_len) == 0.0
-
-
-def test_geodesic_diameter_two_sweep_deterministic():
-    n_nodes, edges, edge_len = _mid_start_path_graph()
-    d_a = geodesic_diameter_two_sweep(n_nodes, edges, edge_len, start=0)
-    d_b = geodesic_diameter_two_sweep(n_nodes, edges, edge_len, start=0)
-    assert d_a == d_b   # bit-equal: deterministic, fixed start, no RNG

tests/test_graph_build.py

@@ -1,52 +0,0 @@
-import torch
-from src.graph_build import knn_edge_index, edge_attr, random_rotation_matrix, rotate
-
-
-def test_knn_symmetric_no_selfloops():
-    pos = torch.tensor([[0., 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]])
-    ei = knn_edge_index(pos, k=1)
-    assert ei.shape[0] == 2
-    assert (ei[0] != ei[1]).all()                       # no self-loops
-    pairs = {(int(a), int(b)) for a, b in ei.t()}
-    assert all((b, a) in pairs for a, b in pairs)        # symmetric
-
-
-def test_edge_attr_mm():
-    pos = torch.tensor([[0., 0, 0], [3, 4, 0]])
-    ei = torch.tensor([[0], [1]])
-    ea = edge_attr(pos, ei)
-    torch.testing.assert_close(ea[0], torch.tensor([3., 4., 0., 5.]))  # dst-src + norm
-
-
-def test_rotation_is_orthonormal_and_preserves_distances():
-    R = random_rotation_matrix(seed=0)
-    torch.testing.assert_close(R @ R.t(), torch.eye(3), atol=1e-5, rtol=1e-5)
-    pos = torch.randn(10, 3)
-    rp = rotate(pos, R)
-    d0 = torch.cdist(pos, pos); d1 = torch.cdist(rp, rp)
-    torch.testing.assert_close(d0, d1, atol=1e-4, rtol=1e-4)  # rotation preserves distances
-
-
-def test_union_contains_mesh_edges_flagged():
-    from src.graph_build import union_edge_index, union_edge_attr
-    pos = torch.tensor([[0., 0, 0], [1, 0, 0], [2, 0, 0], [10, 10, 10]])
-    me = torch.tensor([[0, 3]])                 # pair kNN(k=1) would never pick
-    ei, ism = union_edge_index(pos, me, k=1)
-    pairs = {(int(a), int(b)): float(f) for (a, b), f in zip(ei.t(), ism)}
-    assert pairs[(0, 3)] == 1.0 and pairs[(3, 0)] == 1.0       # both directions, flagged
-    assert all(f == 0.0 for p, f in pairs.items() if p not in {(0, 3), (3, 0)})
-    key = ei[0] * pos.shape[0] + ei[1]
-    assert len(torch.unique(key)) == ei.shape[1]               # no duplicates
-    assert all((b, a) in pairs for a, b in pairs)              # symmetric
-    ea = union_edge_attr(pos, ei, ism)
-    assert ea.shape == (ei.shape[1], 5)
-    torch.testing.assert_close(ea[:, 4], ism)
-
-
-def test_union_dedups_overlap():
-    from src.graph_build import union_edge_index
-    pos = torch.tensor([[0., 0, 0], [1, 0, 0]])
-    me = torch.tensor([[0, 1]])                 # same pair kNN also finds
-    ei, ism = union_edge_index(pos, me, k=1)
-    assert ei.shape[1] == 2                     # (0,1) and (1,0), once each
-    assert (ism == 1.0).all()                   # mesh flag wins on overlap

tests/test_hybrid.py

@@ -1,34 +1,14 @@
 import logging
-import os
 
 import numpy as np
 import pytest
-from src.hybrid import (hybrid_field, hybrid_field_affine, hybrid_field_max_anchor,
-                        hybrid_field_soft)
+from src.hybrid import hybrid_field_affine, hybrid_field_max_anchor, hybrid_field_soft
 
 
-def test_hybrid_takes_shape_from_gnn_level_from_baseline():
-    # GNN shape [-1,0,+1] around mean 5; baseline mean 20 (shape irrelevant)
-    gnn = np.array([4.0, 5.0, 6.0])          # deviations -1,0,+1
-    base = np.array([18.0, 22.0, 20.0])      # mean 20
-    out = hybrid_field(gnn, base)
-    # deviations preserved from GNN, level set to baseline mean (20)
-    np.testing.assert_allclose(out, [19.0, 20.0, 21.0])
 
 
-def test_hybrid_clips_negative():
-    gnn = np.array([0.0, 10.0])              # mean 5, devs -5,+5
-    base = np.array([1.0, 1.0])              # mean 1 -> level 1: -4, +6 -> clip -4 to 0
-    out = hybrid_field(gnn, base)
-    np.testing.assert_allclose(out, [0.0, 6.0])
 
 
-def test_hybrid_preserves_within_graph_ranking():
-    # hot-spot ranking (argsort) must come from the GNN shape, unchanged by re-leveling
-    gnn = np.array([3.0, 1.0, 9.0, 5.0])
-    base = np.array([100.0, 100.0, 100.0, 100.0])
-    out = hybrid_field(gnn, base)
-    assert list(np.argsort(out)) == list(np.argsort(gnn))
 
 
 def test_hybrid_affine_spread_match():
@@ -59,9 +39,8 @@ def test_hybrid_affine_preserves_ranking():
 
 def test_hybrid_affine_clamp_high():
     # gnn std is 10; std_seconds=100 -> a_raw=10.0, clamped down to the a_clamp
-    # ceiling (widened 3.0 -> 5.0 by the 2026-07-31 spec amendment). No clipping
-    # occurs at a=5.0 here (total=1000 is large), so this exercises the fast
-    # (pre-amendment closed-form) path's own clamp handling.
+    # ceiling 5.0. No clipping occurs at a=5.0 here (total=1000 is large), so
+    # this exercises the closed-form path's own clamp handling.
     gnn = np.array([0.0, 20.0])
     total = 1000.0
     pred, info = hybrid_field_affine(gnn, total, std_seconds=100.0)
@@ -96,8 +75,8 @@ def test_hybrid_affine_clip_negative():
 
 
 def test_hybrid_affine_heavy_clip_achievable():
-    # Real fields under FT3 gate 4 sometimes need heavy scaling to hit the
-    # GBM's std target; the pre-amendment closed form ignores clipping
+    # Real fields sometimes need heavy scaling to hit the spread head's
+    # target; a closed form that ignores clipping
     # entirely, so its implied `a` can leave the delivered std far short of
     # the target. Construct such a case (right-skewed field; the naive
     # closed-form `a` clips ~30% of nodes) and confirm the clip-aware solve
@@ -123,8 +102,7 @@ def test_hybrid_affine_heavy_clip_achievable():
 def test_hybrid_affine_unachievable_due_to_clipping():
     # std_seconds=8.0 exceeds this field's interior PEAK (f rises past
     # a_clipstart to a peak around a~1.78, std~3.82, then falls -- it never
-    # reaches 8.0 anywhere in the clamp range). Per the 2026-07-31 review
-    # correction, the unreachable case must pick the scanned point CLOSEST to
+    # reaches 8.0 anywhere in the clamp range). The unreachable case must pick the scanned point CLOSEST to
     # the target (here, at/near the peak) rather than blindly clamping to
     # a_clamp[1]=5.0 -- clamping to 5.0 would give a worse fit (std~3.45,
     # clip_frac 0.7) than the peak (std~3.82, clip_frac 0.4).
@@ -142,30 +120,6 @@ def test_hybrid_affine_unachievable_due_to_clipping():
     assert abs(achieved - 8.0) < abs(np.std(at_hi) - 8.0)
 
 
-_A_00968 = "data/features_v4/graphs/A_00968.npz"
-
-
-@pytest.mark.skipif(not os.path.exists(_A_00968),
-                    reason="requires the FT2a v4 feature cache "
-                           "(data/features_v4/graphs/A_00968.npz) on disk")
-def test_hybrid_affine_real_field_returns_smallest_root_not_ceiling():
-    # Regression for the 2026-07-31 review finding: on a REAL per-node `y`
-    # field, f(a) has an interior peak (~a=2.67, std~420.5) then falls
-    # (f(5.0)~390.8). A naive `f(hi) < std_target => clamp` pre-check would
-    # wrongly land on a=5.0 (clip_frac~0.48) even though a much-less-clipped
-    # true root exists near the peak's rising edge (a~2.6, clip_frac~0.07).
-    # The clip-aware solve must find that smaller root, not the ceiling.
-    data = np.load(_A_00968)
-    y = data["y"].astype(np.float64)
-    total = 1000.0
-    target = 419.0  # inside the narrow (f(a_clipstart), peak_f) crossing window
-
-    pred, info = hybrid_field_affine(y, total, target)
-
-    assert info["a_clamped"] is False
-    assert info["a"] != pytest.approx(5.0)
-    assert info["clip_frac"] < 0.3
-    assert np.std(pred) == pytest.approx(target, rel=1e-3)
 
 
 def test_hybrid_affine_no_clip_matches_pre_amendment_closed_form():
@@ -193,87 +147,18 @@ def test_hybrid_affine_no_clip_matches_pre_amendment_closed_form():
     np.testing.assert_allclose(pred_new, pred_old, rtol=1e-6)
 
 
-# --------------------------------------------------------------------------- #
-# AMENDMENT 2: lambda-shrinkage of the affine spread target (2026-07-31)
-# --------------------------------------------------------------------------- #
-def test_hybrid_affine_lambda_default_bit_identical_on_clip_heavy_case():
-    # lambda_shrink=1.0 (the default) must be byte-for-bit identical to
-    # calling hybrid_field_affine with no lambda_shrink argument at all --
-    # on a CLIP-HEAVY case (the heavy-clip-achievable field from round 1),
-    # to prove the guard skips the new code path even when the solve
-    # actually clips (not just in a trivial no-clip case).
-    g = np.random.default_rng(1).exponential(5.0, size=200)
-    total = 8.0
-    target = 0.3 * np.std(g)
 
-    pred_default, info_default = hybrid_field_affine(g, total, target)
-    pred_lambda1, info_lambda1 = hybrid_field_affine(g, total, target, lambda_shrink=1.0)
 
-    assert np.array_equal(pred_default, pred_lambda1)
-    assert info_default["a"] == info_lambda1["a"]
-    assert info_default["clip_frac"] == info_lambda1["clip_frac"]
-    assert info_lambda1["lambda_shrink"] == 1.0
-    assert info_lambda1["std_target_effective"] == pytest.approx(target, rel=1e-12)
 
 
-def test_hybrid_affine_lambda_zero_targets_gnn_own_spread():
-    # lambda_shrink=0.0 -> std_eff = std(gnn)**1 * std_target**0 = std(gnn)
-    # exactly, regardless of std_target's value. In a no-clip case this
-    # makes a==1.0 (no spread rescaling at all), so the result must equal
-    # hybrid_field_max_anchor's pure quantile-anchor level shift exactly.
-    gnn = np.array([3.0, 1.0, 9.0, 5.0])
-    total = 1000.0
-    pred, info = hybrid_field_affine(gnn, total, std_seconds=999.0, lambda_shrink=0.0)
-    expected = hybrid_field_max_anchor(gnn, total)
-    assert info["a"] == pytest.approx(1.0, rel=1e-9)
-    assert info["a_clamped"] is False
-    assert info["std_target_effective"] == pytest.approx(np.std(gnn), rel=1e-9)
-    np.testing.assert_allclose(pred, expected, rtol=1e-9)
 
 
-def test_hybrid_affine_lambda_half_is_geometric_mean_of_stds():
-    # lambda_shrink=0.5 -> std_eff = sqrt(std(gnn) * std_target). In a
-    # no-clip case the achieved std equals std_eff exactly (fast path).
-    gnn = np.array([0.0, 20.0])  # std(gnn) == 10
-    total = 1000.0  # large enough that nothing clips
-    std_target = 20.0
-    s = np.std(gnn)
 
-    pred, info = hybrid_field_affine(gnn, total, std_target, lambda_shrink=0.5)
 
-    expected_std_eff = np.sqrt(s * std_target)
-    assert info["std_target_effective"] == pytest.approx(expected_std_eff, rel=1e-9)
-    assert info["a_clamped"] is False
-    assert np.std(pred) == pytest.approx(expected_std_eff, rel=1e-9)
-
-
-def test_hybrid_affine_lambda_info_carries_new_keys():
-    gnn = np.array([3.0, 1.0, 9.0, 5.0])
-    pred, info = hybrid_field_affine(gnn, 1000.0, 5.0, lambda_shrink=0.7)
-    assert "lambda_shrink" in info
-    assert "std_target_effective" in info
-    assert info["lambda_shrink"] == pytest.approx(0.7)
-    assert isinstance(info["std_target_effective"], float)
-
-
-def test_hybrid_affine_lambda_nondefault_degenerate_std_target_falls_back(caplog):
-    # std_seconds <= 0 is only inspected when lambda_shrink != 1.0 (the
-    # default path never checks its sign) -- with a non-default lambda it
-    # must fall back to the quantile-anchor with the same loud-warning
-    # pattern as the degenerate-GNN-field guard.
-    gnn = np.array([3.0, 1.0, 9.0, 5.0])
-    total = 1000.0
-    with caplog.at_level(logging.WARNING):
-        pred, info = hybrid_field_affine(gnn, total, std_seconds=0.0, lambda_shrink=0.5)
-    expected = hybrid_field_max_anchor(gnn, total)
-    np.testing.assert_allclose(pred, expected)
-    assert info["degenerate_fallback"] is True
-    assert info["lambda_shrink"] == 0.5
-    assert any("degenerate" in rec.message.lower() for rec in caplog.records)
 
 
 # --------------------------------------------------------------------------- #
-# hybrid_field_soft: weighted two-target soft recomposition (2026-07-31)
+# hybrid_field_soft: weighted two-target soft recomposition
 # --------------------------------------------------------------------------- #
 def _squeezed_case():
     # A field with a small (3%) group of low-value outlier nodes and a tight
@@ -349,77 +234,20 @@ def test_hybrid_soft_degenerate_flat_field_falls_back(caplog):
     assert any("degenerate" in rec.message.lower() for rec in caplog.records)
 
 
-def test_hybrid_soft_info_dict_carries_all_keys():
-    g, total, target = _squeezed_case()
-    pred, info = hybrid_field_soft(g, total, target)
-    expected_keys = {"a", "m", "target_unreachable", "degenerate_fallback", "clip_frac",
-                     "total_deviation_frac", "j_total_resid", "j_std_resid",
-                     "lambda_shrink", "std_target_effective"}
-    assert expected_keys <= set(info.keys())
-    # clip_frac must be measured on the FINAL (returned) field.
-    a, m = info["a"], info["m"]
-    c = g - g.mean()
-    final_raw = a * c + m
-    assert info["clip_frac"] == pytest.approx(float((final_raw < 0.0).mean()))
 
 
-def test_hybrid_soft_lambda_shrink_matches_v5_convention():
-    gnn = np.array([0.0, 20.0])
-    total = 1000.0
-    std_seconds = 20.0
-    s = np.std(gnn)
-    pred, info = hybrid_field_soft(gnn, total, std_seconds, lambda_shrink=0.5)
-    expected = (s ** 0.5) * (std_seconds ** 0.5)
-    assert info["std_target_effective"] == pytest.approx(expected, rel=1e-9)
 
 
 # --------------------------------------------------------------------------- #
-# Review round 1 (2026-07-31) fix coverage: CRITICAL-1 (m-window centered on
-# the seed, not the absolute total*[0.7,1.3]), CRITICAL-2 (perf budget), LOW-6
-# (a_clamp[0]<=0 guard), LOW-7 (realistic-scale degeneration case).
+# Edge cases: the a_clamp[0] <= 0 guard and degeneration to the affine solve
+# at a realistic scale.
 # --------------------------------------------------------------------------- #
-_REAL_FIELD_STEMS = ["A_00968", "A_01541", "C_00681"]
-_REAL_FIELD_PATHS = [f"data/features_v4/graphs/{stem}.npz" for stem in _REAL_FIELD_STEMS]
-
-
-@pytest.mark.parametrize("path", _REAL_FIELD_PATHS)
-def test_hybrid_soft_engages_with_interior_optimum_on_real_fields(path):
-    # Reviewer-mandated re-verification for CRITICAL-1: on real per-node `y`
-    # fields with std_t just past the affine solve's own reach (1.1-1.5x
-    # std(g)), the soft solve at DEFAULT weights must engage (not just return
-    # the seed) with an INTERIOR `a` (strictly inside a_clamp, not pinned to
-    # either boundary) and a single-digit-percent total deviation -- verified
-    # interactively before encoding: e.g. A_00968 at 1.3x gives a~1.28 (vs
-    # a_clamp (0.3, 5.0)) and total_deviation_frac ~3.8%.
-    if not os.path.exists(path):
-        pytest.skip(f"requires the FT2a v4 feature cache ({path}) on disk")
-    data = np.load(path)
-    y = data["y"].astype(np.float64)
-    s = float(np.std(y))
-    total = float(np.quantile(y, 0.995))
-
-    engaged = False
-    for mult in (1.1, 1.2, 1.3, 1.4, 1.5):
-        target = mult * s
-        pred_affine, info_affine = hybrid_field_affine(y, total, target)
-        if not info_affine["a_clamped"]:
-            continue  # affine already reaches this target -- nothing to trade
-        engaged = True
-        pred_soft, info_soft = hybrid_field_soft(y, total, target)
-        assert info_soft["target_unreachable"] is True
-        a_clamp = (0.3, 5.0)
-        assert a_clamp[0] * 1.001 < info_soft["a"] < a_clamp[1] * 0.999, (
-            f"expected an interior optimum, got a={info_soft['a']!r}")
-        assert abs(info_soft["total_deviation_frac"]) < 0.10
-        # the soft solve must never do WORSE than the affine seed on std:
-        d_affine = abs(np.std(pred_affine) - target)
-        d_soft = abs(np.std(pred_soft) - target)
-        assert d_soft <= d_affine + 1e-9
-    assert engaged, f"{path}: no mult in (1.1..1.5) made the affine solve a_clamped=True"
+
+
 
 
 def test_hybrid_soft_realistic_scale_exact_degeneration_matches_affine():
-    # LOW-7: the original exact-degeneration test used a 2-node toy field;
+    # The exact-degeneration test above uses a 2-node toy field;
     # this adds a realistic-scale (n=5000), right-skewed synthetic field with
     # an EASILY achievable std target (no clipping needed at all), so the
     # affine seed hits both targets and the soft solve must reduce to it
@@ -447,7 +275,7 @@ def test_hybrid_soft_a_clamp_lo_nonpositive_raises():
 
 
 def test_hybrid_soft_negative_weight_raises():
-    # Review round 1, LOW: a negative weight rewards moving AWAY from that
+    # LOW: a negative weight rewards moving AWAY from that
     # target (minimizing J then maximizes the residual's square) -- must
     # raise loudly rather than silently drive a to a nonsensical optimum.
     gnn = np.array([1.0, 2.0, 3.0])
@@ -458,7 +286,7 @@ def test_hybrid_soft_negative_weight_raises():
 
 
 def test_hybrid_soft_both_weights_zero_raises():
-    # Review round 1, LOW: w_total == w_std == 0 makes J identically zero
+    # LOW: w_total == w_std == 0 makes J identically zero
     # everywhere -- no signal for the search to optimize against.
     gnn = np.array([1.0, 2.0, 3.0])
     with pytest.raises(ValueError, match="w_total|w_std"):
@@ -476,7 +304,7 @@ def test_hybrid_soft_one_weight_zero_is_allowed():
 
 
 def test_hybrid_soft_accepted_optimum_majority_clipped_warns(caplog):
-    # LOW-5: the affine path already warns when ITS OWN clip_frac > 0.5; the
+    # The affine path already warns when ITS OWN clip_frac > 0.5; the
     # soft path must emit the analogous warning when the ACCEPTED soft
     # optimum (which can differ from the affine seed) itself clips a
     # majority of the field.
@@ -492,25 +320,17 @@ def test_hybrid_soft_accepted_optimum_majority_clipped_warns(caplog):
               for rec in caplog.records)
 
 
-@pytest.mark.parametrize("path", _REAL_FIELD_PATHS)
-def test_hybrid_soft_perf_budget_10k_nodes(path):
-    # CRITICAL-2: < 200ms/10k-node budget (generous 2x margin here to avoid
-    # CI flakiness while still catching a gross regression). Measured on
-    # dev hardware: ~40-60ms per real field at n=5.4k-8k nodes engaged.
-    import time
-    if not os.path.exists(path):
-        pytest.skip(f"requires the FT2a v4 feature cache ({path}) on disk")
-    data = np.load(path)
-    y = data["y"].astype(np.float64)
-    s = float(np.std(y))
-    total = float(np.quantile(y, 0.995))
-    target = 1.4 * s  # engages the grid search on all three reference fields
-
-    hybrid_field_soft(y, total, target)  # warm-up (import/JIT-adjacent costs)
-    times = []
-    for _ in range(3):
-        t0 = time.perf_counter()
-        hybrid_field_soft(y, total, target)
-        times.append(time.perf_counter() - t0)
-    budget_s = 0.200 * (len(y) / 10_000) * 2  # scale budget to n, then 2x slack
-    assert min(times) < budget_s, f"{path}: {min(times)*1000:.1f}ms exceeds {budget_s*1000:.1f}ms budget"
+
+
+def test_hybrid_soft_info_dict_carries_all_keys():
+    g, total, target = _squeezed_case()
+    pred, info = hybrid_field_soft(g, total, target)
+    expected_keys = {"a", "m", "target_unreachable", "degenerate_fallback", "clip_frac",
+                     "total_deviation_frac", "j_total_resid", "j_std_resid",
+                     "std_target_effective"}
+    assert expected_keys <= set(info.keys())
+    # clip_frac must be measured on the FINAL (returned) field.
+    a, m = info["a"], info["m"]
+    c = g - g.mean()
+    final_raw = a * c + m
+    assert info["clip_frac"] == pytest.approx(float((final_raw < 0.0).mean()))

tests/test_hybrid_max_anchor.py

@@ -1,11 +1,6 @@
-"""Task 2/7: max/quantile-anchor hybrid recomposition, in src/hybrid.py.
+"""max/quantile-anchor hybrid recomposition, in src/hybrid.py.
 
-See src/evaluate_ft2b.py's module docstring ("v4 hybrid recomposition:
-MAX-anchor, not MEAN-anchor") for why this differs from hybrid_field's
-MEAN-anchor. This module is the single source; evaluate_ft2b re-exports it
-under its old private name via an alias import.
-
-Task 7: the exact-max anchor is fragile to a handful of outlier nodes (one
+The exact-max anchor is fragile to a handful of outlier nodes (one
 inflated GNN log-space prediction can drag the entire rest of the field
 below the zero clip). The anchor moved to a high quantile
 (`np.quantile(gnn, q)`, default `q=0.995`); `q=1.0` reproduces the old
@@ -17,7 +12,7 @@ from src.hybrid import hybrid_field_max_anchor
 
 
 def test_max_anchor_q1_sets_exact_max_and_preserves_shape():
-    """q=1.0 regression pin: reproduces the pre-Task-7 exact-max behavior."""
+    """q=1.0 regression pin: reproduces the earlier exact-max behavior."""
     g = np.array([1.0, 4.0, 2.5])
     out = hybrid_field_max_anchor(g, 10.0, q=1.0)
     assert out.max() == 10.0
@@ -37,13 +32,10 @@ def test_max_anchor_clips_at_zero():
     assert hybrid_field_max_anchor(np.array([0.0, 100.0]), 10.0)[0] == 0.0
 
 
-def test_evaluate_ft2b_uses_the_same_function():
-    from src import evaluate_ft2b, hybrid
-    assert evaluate_ft2b._max_anchor_hybrid is hybrid.hybrid_field_max_anchor
 
 
 def test_max_anchor_q1_reproduces_old_exact_max_on_random_vector():
-    """Regression pin: q=1.0 bit-for-bit equals the pre-Task-7 formula
+    """Regression pin: q=1.0 bit-for-bit equals the earlier formula
     `clip(g - g.max() + total, 0, None)` on an arbitrary random vector."""
     rng = np.random.default_rng(42)
     g = rng.normal(loc=50.0, scale=20.0, size=200)

tests/test_inference_props.py

@@ -36,8 +36,7 @@ def test_props_json_override_and_missing_keys(tmp_path):
 
 
 # --------------------------------------------------------------------------- #
-# Review finding I2 (final whole-branch review, 2026-08-01): a `--props`
-# custom-JSON override with TL <= TS was previously unvalidated -- silently
+# A `--props` custom-JSON override with TL <= TS was previously unvalidated -- silently
 # collapsing src.shrinkage.niyama_proxy to 0 everywhere (100% flagged).
 # --------------------------------------------------------------------------- #
 def test_props_json_tl_less_equal_ts_raises(tmp_path):
@@ -68,21 +67,6 @@ def test_named_alloy_tl_ts_still_valid_after_i2_guard():
         assert p["TL"] > p["TS"]
 
 
-# --------------------------------------------------------------------------- #
-# resolve_alloy_family (SP2 shrinkage Task 4, controller directive 2):
-# resolve_props's returned 8-key dict has no family field, so shrinkage
-# screening maps the ALLOY GRADE NAME to one of src.shrinkage.SHRINKAGE_PCT's
-# family keys via a small name-prefix helper, rather than crashing or
-# silently guessing on an unrecognized grade.
-# --------------------------------------------------------------------------- #
-def test_resolve_alloy_family_known_prefixes():
-    assert resolve_alloy_family("EN-GJL-250") == ("grey_iron", False)
-    assert resolve_alloy_family("EN-GJL-200") == ("grey_iron", False)
-    assert resolve_alloy_family("EN-GJS-400-15") == ("ductile_iron", False)
-    assert resolve_alloy_family("GS-38 (GE200)") == ("low_alloy_steel", False)
-    assert resolve_alloy_family("GS-45 (GP240GH)") == ("low_alloy_steel", False)
-    assert resolve_alloy_family("G20Mn5") == ("low_alloy_steel", False)
-    assert resolve_alloy_family("G17CrMo5-5") == ("low_alloy_steel", False)
 
 
 def test_resolve_alloy_family_default_alloy_when_none():
@@ -95,22 +79,12 @@ def test_resolve_alloy_family_unknown_prefix_falls_back_loud():
     assert family == "steel" and assumed is True
 
 
-def test_resolve_alloy_family_props_json_always_assumed_steel():
-    # A custom --props JSON bypasses the named-material table entirely -- no
-    # grade name to map from -- so the documented fallback is "steel",
-    # flagged as assumed regardless of what `alloy` happens to be.
-    family, assumed = resolve_alloy_family("EN-GJL-250", props_json="/some/path.json")
-    assert family == "steel" and assumed is True
 
 
-def test_resolve_alloy_family_every_mapped_family_is_a_known_shrinkage_pct_key():
-    for grade in ("EN-GJL-250", "EN-GJS-400-15", "GS-38 (GE200)", "unknown"):
-        family, _ = resolve_alloy_family(grade)
-        assert family in SHRINKAGE_PCT
 
 
 def test_resolve_alloy_family_covers_every_catalog_grade_exactly():
-    """Review round 1, finding 6: a data-driven guard over the REAL property
+    """a data-driven guard over the REAL property
     DB (`data/properties/materials.csv` via `list_alloys()`) -- every grade
     actually on file today must be an EXACT prefix match (`assumed is
     False`), not a fallback. This is the loud-failure trip-wire for a

tests/test_input_encoder_extension.py

@@ -1,9 +0,0 @@
-import torch
-from src.model_gnn import SolidGNN, zero_init_input_columns
-
-
-def test_zero_init_last_columns():
-    m = SolidGNN(node_in=15, hidden=16, n_blocks=1)
-    zero_init_input_columns(m, n_new=2)
-    w = m.node_encoder[0][0].weight  # first Linear of the node_encoder _mlp
-    assert torch.count_nonzero(w[:, -2:]) == 0 and torch.count_nonzero(w[:, :13]) > 0