diff --git a/.gitignore b/.gitignore index 1dc6ac797a4d4781433194a3df32bae293486bb0..82c37d8bbf6ff3b320cf692f6bb6fcc253d84114 100644 --- a/.gitignore +++ b/.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/ diff --git a/README.md b/README.md index 0ee993f274fbd5fd5d1500fe7ecc52d38ae74386..5c65356337ecd25c8ffa7ddfa1395f2f65aeda10 100644 --- a/README.md +++ b/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. diff --git a/requirements.txt b/requirements.txt index 354c9c233f65a905658447fe21542ea921577d8c..a8604f0e11594ab2cc019fa7f76c37387f952d28 100644 --- a/requirements.txt +++ b/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 diff --git a/src/baseline.py b/src/baseline.py deleted file mode 100644 index a15c1ab35101da3d9c0f8d28d510aad4445713bf..0000000000000000000000000000000000000000 --- a/src/baseline.py +++ /dev/null @@ -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) diff --git a/src/build_dataset.py b/src/build_dataset.py deleted file mode 100644 index 895381f207747ae31580539f6f4390075843bf68..0000000000000000000000000000000000000000 --- a/src/build_dataset.py +++ /dev/null @@ -1,1199 +0,0 @@ -"""Reprocess raw .pt + paired .msh into cached per-node feature files. - -Writes one .npz per casting under /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/.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/.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 `/mesh_work_v4/.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 (`/mesh_work/` vs this module's -# `/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 `/mesh_work/[_scaled] - .msh` (`adapters._work_dir()`'s naming convention). `manifest_dir` - (where manifest.csv lives) IS `` in the standard layout (both - default to "doe_out"), so no separate config knob is needed.""" - run_id = row.get("run_id", "") - 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", "") - 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 `/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 `/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 `/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) diff --git a/src/compare_d_parts.py b/src/compare_d_parts.py deleted file mode 100644 index 000907b5e2ebddc7f97e436ed543abbaa8dfe70e..0000000000000000000000000000000000000000 --- a/src/compare_d_parts.py +++ /dev/null @@ -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/.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 diff --git a/src/dataload.py b/src/dataload.py deleted file mode 100644 index 20ca33ff98620ee36f39fdc3dce8582de10be7d2..0000000000000000000000000000000000000000 --- a/src/dataload.py +++ /dev/null @@ -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"] diff --git a/src/evaluate_ft2b.py b/src/evaluate_ft2b.py deleted file mode 100644 index 2f100d078d324580b8f0d3bb9493f1ef406dbf9f..0000000000000000000000000000000000000000 --- a/src/evaluate_ft2b.py +++ /dev/null @@ -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/.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 diff --git a/src/evaluate_gnn.py b/src/evaluate_gnn.py deleted file mode 100644 index 892d722986e062eb2bda5b543161861da5c96233..0000000000000000000000000000000000000000 --- a/src/evaluate_gnn.py +++ /dev/null @@ -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: /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) diff --git a/src/features.py b/src/features.py index e4d5a6359bc0930f0a3578e43bacbacc14cfb6c9..8be7f286b455035d043d97138793eba873ca772f 100644 --- a/src/features.py +++ b/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 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)} diff --git a/src/gbm_ft2b.py b/src/gbm_ft2b.py deleted file mode 100644 index 09a14e141cf33c122dbbf8bb2c1ed039d6e27af9..0000000000000000000000000000000000000000 --- a/src/gbm_ft2b.py +++ /dev/null @@ -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))) diff --git a/src/gbm_heads.py b/src/gbm_heads.py new file mode 100644 index 0000000000000000000000000000000000000000..441d174a401b127ed1700712741f986c50f847da --- /dev/null +++ b/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}") diff --git a/src/geometry.py b/src/geometry.py index 665cf90c4337d98bffee180d64ff693770d2beeb..dcac65ca66c0aa46b671ac8defbb149eefd0840f 100644 --- a/src/geometry.py +++ b/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) diff --git a/src/gnn_dataset.py b/src/gnn_dataset.py deleted file mode 100644 index f7203ac41f443ac4880295a6fdd258b2d3c79564..0000000000000000000000000000000000000000 --- a/src/gnn_dataset.py +++ /dev/null @@ -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) diff --git a/src/graph_build.py b/src/graph_build.py deleted file mode 100644 index c8a6ec6bae6bdde14d2560bbb8089140f020acbb..0000000000000000000000000000000000000000 --- a/src/graph_build.py +++ /dev/null @@ -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) diff --git a/src/hybrid.py b/src/hybrid.py index d725672143ceb396717d10f85cdb72248a191497..94f6cd6dbde256eb19fd62ad5933505722709645 100644 --- a/src/hybrid.py +++ b/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) diff --git a/src/inference_props.py b/src/inference_props.py index 16294ac904477f49ad1106b22f8eef739589e4e5..0543f22feb415805d822d33ab4b3eb37eb688fc0 100644 --- a/src/inference_props.py +++ b/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 diff --git a/src/mesh_cache.py b/src/mesh_cache.py index d79811c62d185ff86cbfebfab41a9ba3e5bcce13..b42c009eb7a844653932fe19c72663ecfa2ebfba 100644 --- a/src/mesh_cache.py +++ b/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/`): - .msh -- the cached gmsh 2.2 mesh (only written when ok=True) - .json -- meta dict (ok, n_nodes, ..., mesh_wall_s, pipeline_version, - source_path, created_at) - -Writes are atomic and per-writer-unique (`.msh...tmp` / -`.json...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 +`.msh` / `.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) diff --git a/src/meshing.py b/src/meshing.py new file mode 100644 index 0000000000000000000000000000000000000000..68c83eb837b74be9b00d4e2808f070f08dd90f69 --- /dev/null +++ b/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() diff --git a/src/metrics.py b/src/metrics.py deleted file mode 100644 index 02aa001312c5a4948ef45025233fd3edd7bddf69..0000000000000000000000000000000000000000 --- a/src/metrics.py +++ /dev/null @@ -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 diff --git a/src/model_api.py b/src/model_api.py index 7fbb0ed29996509e32081c2b17f3562ee5eec7fb..728b938ddb2d6cbacc00c2f40769665446f5c4a3 100644 --- a/src/model_api.py +++ b/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] diff --git a/src/model_gnn.py b/src/model_gnn.py index f25fb276243a670186508b3d8be4cc47fa598733..16003bdf5addbc519197428c50e250e2c6185dbc 100644 --- a/src/model_gnn.py +++ b/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) diff --git a/src/msh_parse.py b/src/msh_parse.py index 644418d59f3dc7b13a0ed86fd7484c1b382dfd7e..ebc05ae00da8d242658b354b53cf6d14ca33fa93 100644 --- a/src/msh_parse.py +++ b/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 diff --git a/src/ood.py b/src/ood.py index b150dc1bb25bdb68e37d61411bb3b90a961e828e..04e39293ab62906fbbd9531a12ed448f9bb6ad98 100644 --- a/src/ood.py +++ b/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: diff --git a/src/predict_casting.py b/src/predict_casting.py index d4a0e425f2eb19691d62714923cbeed0cbf7b20d..b64af070246d52f9700a164cbcd92925f9f694d3 100644 --- a/src/predict_casting.py +++ b/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 `.vtu`, `_summary.txt` and the +mesh `.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 "" - - -_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() diff --git a/src/release_io.py b/src/release_io.py index 0f14bd86956c0077ec76dd11b8c1bb34a9a5a75e..ee243fd098fcbd26f031b4b86f18ab2a73d68a6e 100644 --- a/src/release_io.py +++ b/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.` and + `norm.`; 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.` + - `norm.`, 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) diff --git a/src/run_baseline.py b/src/run_baseline.py deleted file mode 100644 index 98fea41fbe61b8813e6c942e7d1eb1e05295fe23..0000000000000000000000000000000000000000 --- a/src/run_baseline.py +++ /dev/null @@ -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() diff --git a/src/select_ckpt_ft2b.py b/src/select_ckpt_ft2b.py deleted file mode 100644 index 78c5959482ca0291023e654bacdde5c34cea17df..0000000000000000000000000000000000000000 --- a/src/select_ckpt_ft2b.py +++ /dev/null @@ -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 (`/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": , "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_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/.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": , "ranking": [, ...]}`, - `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 diff --git a/src/shrinkage.py b/src/shrinkage.py index c82b3eecbdf718d207b4b8d522e2a4ec0ef0c27c..cc12225d2b7f8c90e809675f13294961d231ab10 100644 --- a/src/shrinkage.py +++ b/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}, } diff --git a/src/splits.py b/src/splits.py deleted file mode 100644 index 8eb3731af70b798a24fb6ffac9d85cd8f4dee98d..0000000000000000000000000000000000000000 --- a/src/splits.py +++ /dev/null @@ -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) diff --git a/src/step_mesh.py b/src/step_mesh.py deleted file mode 100644 index 39e07a37c3420617f0a7e99eba438b5857d9a066..0000000000000000000000000000000000000000 --- a/src/step_mesh.py +++ /dev/null @@ -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() diff --git a/src/targets.py b/src/targets.py deleted file mode 100644 index 505709de4b1acde633ac7e293e0cf5fe5ef7705d..0000000000000000000000000000000000000000 --- a/src/targets.py +++ /dev/null @@ -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()) diff --git a/src/train_ft2b.py b/src/train_ft2b.py deleted file mode 100644 index 25c7ce14d0a763e41a45392d095cf1ccf61db9bf..0000000000000000000000000000000000000000 --- a/src/train_ft2b.py +++ /dev/null @@ -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 (`/ft2b_stage{1,2,3}/`, `/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), - } diff --git a/src/train_gnn.py b/src/train_gnn.py deleted file mode 100644 index fa052a7a5ebd6b373819ed372497f11520ca64ba..0000000000000000000000000000000000000000 --- a/src/train_gnn.py +++ /dev/null @@ -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 `/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} diff --git a/src/train_heads.py b/src/train_heads.py deleted file mode 100644 index f4b1bcd7d1a102089c115af6b3efced4c99b20a5..0000000000000000000000000000000000000000 --- a/src/train_heads.py +++ /dev/null @@ -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 - (`/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 `/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} diff --git a/src/utils.py b/src/utils.py deleted file mode 100644 index 4ad0fd885455d6c07201812481db8c0b312d3c37..0000000000000000000000000000000000000000 --- a/src/utils.py +++ /dev/null @@ -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 _.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}") diff --git a/src/warm_start.py b/src/warm_start.py deleted file mode 100644 index 05ceeb0d68d7648280132200b934ba151c17478d..0000000000000000000000000000000000000000 --- a/src/warm_start.py +++ /dev/null @@ -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) diff --git a/src/webapp.py b/src/webapp.py index fd482f37457ba9180222f85998f64242d7410dcb..01547d8b1f5840625b793c7207c2754e10bc1b42 100644 --- a/src/webapp.py +++ b/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 diff --git a/tests/test_features.py b/tests/test_features.py new file mode 100644 index 0000000000000000000000000000000000000000..df90610170fa5f216cdeadc1aeaeb21898497a64 --- /dev/null +++ b/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) diff --git a/tests/test_features_inference.py b/tests/test_features_inference.py deleted file mode 100644 index 055e7a912414b25e2c8935c7cf9626afc3711cf7..0000000000000000000000000000000000000000 --- a/tests/test_features_inference.py +++ /dev/null @@ -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() diff --git a/tests/test_features_v3.py b/tests/test_features_v3.py deleted file mode 100644 index dfedd289085af93032f7174e5cecfe6802288c66..0000000000000000000000000000000000000000 --- a/tests/test_features_v3.py +++ /dev/null @@ -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 diff --git a/tests/test_features_v4.py b/tests/test_features_v4.py deleted file mode 100644 index a508d4eeb9cd80e47a6e48d50ff6d3cdb109660c..0000000000000000000000000000000000000000 --- a/tests/test_features_v4.py +++ /dev/null @@ -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) diff --git a/tests/test_film.py b/tests/test_film.py deleted file mode 100644 index 53c6155c7f9277d2502691779ba649254401d2c3..0000000000000000000000000000000000000000 --- a/tests/test_film.py +++ /dev/null @@ -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) diff --git a/tests/test_gbm_heads.py b/tests/test_gbm_heads.py new file mode 100644 index 0000000000000000000000000000000000000000..238b2331243d73accd43dabf050f068056539ce1 --- /dev/null +++ b/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") diff --git a/tests/test_geometry.py b/tests/test_geometry.py index 1ad7ffc04393d9e700c0d4cbc5abd399d9f57b0a..1064572249945f6b0a35071cd2b080373e803255 100644 --- a/tests/test_geometry.py +++ b/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 diff --git a/tests/test_graph_build.py b/tests/test_graph_build.py deleted file mode 100644 index a2c2a93cbd91f2dfc6012e057c124ae796cf5b86..0000000000000000000000000000000000000000 --- a/tests/test_graph_build.py +++ /dev/null @@ -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 diff --git a/tests/test_hybrid.py b/tests/test_hybrid.py index 0745f9512394191bf1d676d46e13eb4db960828e..bf9a90edc0c4925807a53d30f967a6473ba7f8f5 100644 --- a/tests/test_hybrid.py +++ b/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())) diff --git a/tests/test_hybrid_max_anchor.py b/tests/test_hybrid_max_anchor.py index 094d77d5bce6614a06f579d8170c083ad640c461..db05208cf16954a80ea1f79ed754d5e71310e30f 100644 --- a/tests/test_hybrid_max_anchor.py +++ b/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) diff --git a/tests/test_inference_props.py b/tests/test_inference_props.py index b91a146ea6df1132cde4c37737cc19661c9aa659..7c517d2e08fec8eb4d118f72e56ca2f3ed2f8182 100644 --- a/tests/test_inference_props.py +++ b/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 diff --git a/tests/test_input_encoder_extension.py b/tests/test_input_encoder_extension.py deleted file mode 100644 index 1b19ee257c6e49c0926af0ff1916614405cb873b..0000000000000000000000000000000000000000 --- a/tests/test_input_encoder_extension.py +++ /dev/null @@ -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 diff --git a/tests/test_local_modulus.py b/tests/test_local_modulus.py index 7f6ae6a7e24374b76009f450a64b760e947aacb0..6627e5dbb1b9d0c2c3c6182c7ab7c5bebdc6d8b7 100644 --- a/tests/test_local_modulus.py +++ b/tests/test_local_modulus.py @@ -1,13 +1,8 @@ -import glob -import os - import numpy as np from scipy.spatial import Delaunay from src.geometry import (local_modulus, tet_volumes, edges_from_tets, mesh_volume, total_area) -from src.features import LOCAL_MODULUS_RADII -from src.msh_parse import parse_msh def _cube(a=10.0): @@ -34,79 +29,6 @@ def _box_grid(n=4, a=10.0): return pts, tets, cool, edges, elen -def _jittered_box(n=9, a=40.0, jitter_frac=0.2, seed=0): - """An UNSTRUCTURED stand-in when the real dataset isn't present: a box - grid with interior points randomly jittered (fixed seed -> deterministic) - so Delaunay tetrahedralization doesn't produce the exact-tie coincidences - of a perfectly regular grid -- closer in character to a real gmsh mesh.""" - rng = np.random.default_rng(seed) - 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) - step = a / (n - 1) - interior = np.all((pts > 1e-9) & (pts < a - 1e-9), axis=1) - pts[interior] += rng.uniform(-jitter_frac * step, jitter_frac * step, - size=(interior.sum(), 3)) - dt = Delaunay(pts) - tets = dt.simplices.astype(np.int64) - cool = dt.convex_hull.astype(np.int64) - edges = edges_from_tets(tets) - elen = np.linalg.norm(pts[edges[:, 0]] - pts[edges[:, 1]], axis=1) - return pts, tets, cool, edges, elen - - -# Known-good real training meshes (verified casting-scale bbox, in order of -# increasing exact-computation cost), from the git-ignored ~1.9 GB -# data/msh_none/ dataset -- present on the dev machine but NOT guaranteed to -# exist on every machine/CI (see .gitignore: `data/*`). -_REAL_MESH_CANDIDATES = [ - "data/msh_none/geom/RO_80.msh", # bbox diag ~493 mm, N~4087 - "data/msh_none/L_shape/L135_08_19.msh", # bbox diag ~284 mm, N~5698 - "data/msh_none/T_shape/T_11_19.msh", # bbox diag ~167 mm, N~4610 - "data/msh_none/plate/PF_60_60.msh", # bbox diag ~434 mm, N~4490 -] -# Some geometry-family directories (e.g. `geom/`) mix tiny calibration -# primitives (e.g. a 12 mm cylinder) with real casting-scale meshes; a bare -# "first file alphabetically" glob can land on one of the tiny ones, which -# is too small to exercise the r=15/40 mm radii meaningfully (their balls -# would cover the whole tiny part, matching exact/approx trivially and -# understating real error). The size filter below guards against that. -_FALLBACK_SCAN_DIRS = ["data/msh_none/L_shape", "data/msh_none/T_shape", - "data/msh_none/plate", "data/msh_none/U_prof"] -_MIN_CASTING_BBOX_DIAG_MM = 100.0 - - -def _find_real_training_mesh_path(): - for p in _REAL_MESH_CANDIDATES: - if os.path.exists(p): - return p - for d in _FALLBACK_SCAN_DIRS: - for f in sorted(glob.glob(os.path.join(d, "*.msh")))[:5]: - try: - mesh = parse_msh(f) - except Exception: - continue - bbox_diag = np.linalg.norm(mesh.node_xyz.max(0) - mesh.node_xyz.min(0)) - if bbox_diag >= _MIN_CASTING_BBOX_DIAG_MM: - return f - return None - - -def _representative_unstructured_mesh(): - """A real gmsh training mesh if the dataset is present on this machine - (preferred -- genuinely unstructured, real element-size distribution), - else a jittered synthetic box (unstructured-like fallback). Returns - (xyz, tets, cool_tri, edges, edge_len, is_real).""" - path = _find_real_training_mesh_path() - if path is not None: - mesh = parse_msh(path) - cool = mesh.tris.get("Cooling", np.empty((0, 3), np.int64)) - edges = edges_from_tets(mesh.tets) - elen = np.linalg.norm(mesh.node_xyz[edges[:, 0]] - mesh.node_xyz[edges[:, 1]], axis=1) - return mesh.node_xyz, mesh.tets, cool, edges, elen, True - xyz, tets, cool, edges, elen = _jittered_box() - return xyz, tets, cool, edges, elen, False - - def test_tet_volumes_sum_to_mesh_volume(): xyz, tets, _, _, _ = _cube() assert np.isclose(tet_volumes(xyz, tets).sum(), mesh_volume(xyz, tets)) @@ -139,109 +61,3 @@ def test_multiple_radii_columns_ordered(): assert m.shape == (8, 2) -def _median_rel_err(approx, exact): - denom = np.maximum(np.abs(exact), 1e-9) - return float(np.median(np.abs(approx - exact) / denom)) - - -def test_exact_kwarg_default_unchanged(): - # exact=True (the default) must reproduce the pre-existing behavior - # byte-for-byte -- calling with exact=True explicitly must equal the - # bare (kwarg-free) call. - xyz, tets, cool, edges, elen = _cube() - m_default = local_modulus(xyz, tets, cool, edges, elen, radii=[0.1, 5.0, 1e6]) - m_exact = local_modulus(xyz, tets, cool, edges, elen, radii=[0.1, 5.0, 1e6], exact=True) - assert np.array_equal(m_default, m_exact) - - -def test_subsample_within_tolerance_of_exact_multiple_meshes(): - # Sanity/regression check on small synthetic fixtures: shape/dtype/sign - # are preserved and the combined (flattened) error stays low. NOTE: a - # combined-across-radii median is a weak correctness signal (see - # test_subsample_per_radius_column_error below for the real gate) -- - # kept here only as a cheap smoke test across several node densities. - radii = list(LOCAL_MODULUS_RADII) - geometries = [ - _cube(), # coarse (spacing 10 mm) - _box_grid(n=4, a=12.0), # spacing 4 mm - _box_grid(n=5, a=20.0), # spacing 5 mm - _box_grid(n=6, a=50.0), # spacing 10 mm, more nodes (N=216) - ] - for xyz, tets, cool, edges, elen in geometries: - exact = local_modulus(xyz, tets, cool, edges, elen, radii=radii, exact=True) - approx = local_modulus(xyz, tets, cool, edges, elen, radii=radii, - exact=False, subsample=4) - assert approx.shape == exact.shape - assert approx.dtype == np.float32 and exact.dtype == np.float32 - # no sign errors -- local modulus (V/A, or sentinel) is always >= 0 - assert np.all(approx >= 0) and np.all(exact >= 0) - - -def test_subsample_per_radius_column_error(capsys): - # The real correctness gate (per FT2a coordinator review): check each - # production radius column (5/15/40 mm) SEPARATELY, not a combined/ - # flattened median -- a combined median is dominated by whichever - # column happens to saturate (all-sentinel or all-global-ratio) on a - # given mesh, hiding much larger error at the diagnostic mid radius. - # - # Uses a real gmsh training mesh (genuinely unstructured Delaunay - # tetrahedralization) if the git-ignored ~1.9 GB data/msh_none/ dataset - # is present on this machine; falls back to a jittered synthetic box - # otherwise (see `_representative_unstructured_mesh`). - # - # FINDING (see task-2-report.md for the full k/power sweep): IDW - # smoothing (this task's fix, replacing single-nearest) does NOT bring - # every column under 5% on a representative unstructured mesh, and - # counter to the smoothing hypothesis, INCREASES error at the r=15mm - # column as k grows (k=1 nearest-only beats k=4-8 IDW there) -- because - # local_modulus is a ball-membership step function of geodesic position, - # not a smoothly-varying field, so averaging dissimilar neighbors hurts - # more than it helps. Per the review's branch-3 instruction, this is - # documented honestly (not force-fit to an artificial pass): the - # ceilings below are regression guards on today's measured/characterized - # behavior, calibrated separately for the real-mesh and fallback paths - # (each has a different problem column, both geometry-dependent). See - # the module-level note on `local_modulus` for the resulting guidance: - # this subsample path targets dense-mesh inference, not the - # training-density feature cache (which should keep using exact=True). - xyz, tets, cool, edges, elen, is_real = _representative_unstructured_mesh() - radii = list(LOCAL_MODULUS_RADII) # (5.0, 15.0, 40.0) mm - exact = local_modulus(xyz, tets, cool, edges, elen, radii=radii, exact=True) - approx = local_modulus(xyz, tets, cool, edges, elen, radii=radii, - exact=False, subsample=4) - - # regression ceilings, NOT the aspirational 5% target -- see docstring. - # measured (this repo, this fixture): real mesh r5/r15/r40 approx - # 0.00 / 0.27 / 0.06; fallback jittered box approx 0.22 / 0.21 / 0.01. - ceilings = ({5.0: 0.05, 15.0: 0.35, 40.0: 0.10} if is_real else - {5.0: 0.30, 15.0: 0.30, 40.0: 0.05}) - meets_5pct_target = {} - for j, r in enumerate(radii): - e = exact[:, j]; a = approx[:, j] - rel = np.abs(a - e) / np.maximum(np.abs(e), 1e-9) - med = float(np.median(rel)) - meets_5pct_target[r] = med <= 0.05 - assert med <= ceilings[r], ( - f"r={r}mm median rel-err {med:.4f} exceeds regression ceiling " - f"{ceilings[r]} (fixture={'real mesh' if is_real else 'fallback'})") - with capsys.disabled(): - print(f"\n[test_subsample_per_radius_column_error] fixture=" - f"{'real gmsh mesh' if is_real else 'jittered box fallback'} " - f"N={xyz.shape[0]} meets-5%-target={meets_5pct_target}") - - -def test_subsample_ge_n_falls_back_to_exact(): - xyz, tets, cool, edges, elen = _cube() - radii = [3.0, 1e6] - exact = local_modulus(xyz, tets, cool, edges, elen, radii=radii, exact=True) - approx = local_modulus(xyz, tets, cool, edges, elen, radii=radii, - exact=False, subsample=xyz.shape[0] + 10) - assert np.array_equal(approx, exact) - - -def test_subsample_shape_and_dtype_preserved(): - xyz, tets, cool, edges, elen = _box_grid(n=4) - m = local_modulus(xyz, tets, cool, edges, elen, radii=[1.0, 5.0, 1e6], - exact=False, subsample=3) - assert m.shape == (xyz.shape[0], 3) - assert m.dtype == np.float32 diff --git a/tests/test_mesh_cache.py b/tests/test_mesh_cache.py index 470e4140e9f9b726fc8e31b8850043a91dc6ffa7..6d50e73c82ee77fc4399da70c6ac2b31f3ca736b 100644 --- a/tests/test_mesh_cache.py +++ b/tests/test_mesh_cache.py @@ -1,5 +1,3 @@ -import json, os -import pytest import src.mesh_cache as mc def _fake_heal(calls, ok=True, n_nodes=5000): @@ -55,17 +53,13 @@ def test_failure_not_cached(tmp_path, monkeypatch): monkeypatch.setattr(mc, "heal_and_mesh", _fake_heal(calls, ok=False)) r1 = mc.get_or_mesh(str(step), str(tmp_path / "a.msh"), 8500, cache_dir=str(tmp_path / "c")) r2 = mc.get_or_mesh(str(step), str(tmp_path / "b.msh"), 8500, cache_dir=str(tmp_path / "c")) - # `_fake_heal` (ok=False) always fails regardless of `defeature`, so each - # of the 2 `get_or_mesh` invocations now makes 2 calls under the Task-5 - # auto-trigger policy (plain, then a defeature=True retry) -- 4 total, - # not 2 (pre-Task-5 baseline, before the retry-on-failure policy existed). + # `_fake_heal` (ok=False) always fails, so each of the 2 `get_or_mesh` + # calls tries twice (plain, then defeature=True): 4 calls in total. assert r1["ok"] is False and r2["cache"] == "miss" and len(calls) == 4 # -------------------------------------------------------------------------- -# Task 5: size/failure auto-trigger policy for `defeature=`. Big STEP files -# (>8 MiB) go straight to a single defeature=True call (no wasted plain -# attempt on a file already known, empirically, to need it -- see -# task-5-report.md); smaller files try plain first and retry ONCE with +# size/failure auto-trigger policy for `defeature=`. Big STEP files +# (>8 MiB) try defeature=True first; smaller files try plain first and retry ONCE with # defeature=True only if the plain attempt fails outright. Either way the # returned meta records which `defeature` value actually produced the # result and whether a retry happened, and the cache key is untouched by @@ -168,145 +162,15 @@ def test_tmp_names_are_unique_per_writer(tmp_path): assert tmp_a != final + ".tmp" and tmp_b != final + ".tmp" # no fixed shared name -# -------------------------------------------------------------------------- -# Part B Task 4: `keep_solid_id` / `bodies_key` / `post_mesh` (tagged-mesh -# caching) -- MESH_PIPELINE_VERSION pinned to 3. -# -------------------------------------------------------------------------- def test_pipeline_version_pinned_to_3(): assert mc.MESH_PIPELINE_VERSION == 3 -def test_mesh_cache_key_changes_with_bodies_key(tmp_path): - step = tmp_path / "p.step"; step.write_bytes(b"AAA") - k_none = mc.mesh_cache_key(str(step), 8500) - k_default = mc.mesh_cache_key(str(step), 8500, bodies_key="none") - k_other = mc.mesh_cache_key(str(step), 8500, bodies_key="somehash") - assert k_none == k_default # "none" is the implicit default - assert k_other != k_none - - -def test_get_or_mesh_forwards_keep_solid_id_to_heal_and_mesh(tmp_path, monkeypatch): - step = tmp_path / "p.step"; step.write_bytes(b"body-bytes") - calls = [] - - def fake(step_path, out_msh, target_nodes=8500, defeature=False, keep_solid_id=None, **kw): - calls.append({"keep_solid_id": keep_solid_id, "defeature": defeature}) - with open(out_msh, "w") as f: - f.write("MESH 42") - return {"ok": True, "n_nodes": 42, "notes": ""} - - monkeypatch.setattr(mc, "heal_and_mesh", fake) - r = mc.get_or_mesh(str(step), str(tmp_path / "a.msh"), 8500, cache_dir=str(tmp_path / "c"), - keep_solid_id=7, bodies_key="rolehash1") - assert len(calls) == 1 and calls[0]["keep_solid_id"] == 7 - assert r["cache"] == "miss" and r["n_nodes"] == 42 - - # A different keep_solid_id / bodies_key must miss the same cache_dir - # (different key), calling heal_and_mesh again. - r2 = mc.get_or_mesh(str(step), str(tmp_path / "b.msh"), 8500, cache_dir=str(tmp_path / "c"), - keep_solid_id=9, bodies_key="rolehash2") - assert len(calls) == 2 and calls[1]["keep_solid_id"] == 9 - - -def test_get_or_mesh_post_mesh_hook_runs_on_miss_before_store_and_merges_into_meta( - tmp_path, monkeypatch): - step = tmp_path / "p.step"; step.write_bytes(b"tag-bytes") - heal_calls = [] - post_calls = [] - - def fake_heal(step_path, out_msh, target_nodes=8500, defeature=False, keep_solid_id=None, **kw): - heal_calls.append(1) - with open(out_msh, "w") as f: - f.write("UNTAGGED") - return {"ok": True, "n_nodes": 10, "notes": ""} - - def fake_post_mesh(out_msh): - post_calls.append(out_msh) - # simulate tag_contact_patches rewriting the mesh in place - with open(out_msh, "a") as f: - f.write(" TAGGED") - return {"per_body_patch_area_mm2": {2: 12.5}, "n_adiab_tris": 3} - - monkeypatch.setattr(mc, "heal_and_mesh", fake_heal) - out1 = str(tmp_path / "a.msh") - r1 = mc.get_or_mesh(str(step), out1, 8500, cache_dir=str(tmp_path / "c"), - keep_solid_id=1, bodies_key="bh1", post_mesh=fake_post_mesh) - - assert len(post_calls) == 1 and post_calls[0] == out1 - assert r1["cache"] == "miss" - assert r1["per_body_patch_area_mm2"] == {2: 12.5} - assert r1["n_adiab_tris"] == 3 - assert open(out1).read() == "UNTAGGED TAGGED" # tagged version was stored - - # A cache HIT returns the already-tagged mesh + merged meta, WITHOUT - # calling post_mesh (or heal_and_mesh) again. - out2 = str(tmp_path / "b.msh") - r2 = mc.get_or_mesh(str(step), out2, 8500, cache_dir=str(tmp_path / "c"), - keep_solid_id=1, bodies_key="bh1", post_mesh=fake_post_mesh) - assert r2["cache"] == "hit" - assert len(post_calls) == 1 # not called again - assert len(heal_calls) == 1 # not re-meshed - # NOTE: a JSON round-trip through the on-disk cache turns dict int keys - # into strings (json.dump has no int-key mode) -- a cache HIT's - # per_body_patch_area_mm2 keys are therefore str, not int, unlike a - # fresh MISS's. Callers (predict_casting) normalize back to int. - assert r2["per_body_patch_area_mm2"] == {"2": 12.5} - assert open(out2).read() == "UNTAGGED TAGGED" # tagged mesh served from cache - - -def test_get_or_mesh_no_bodies_key_or_post_mesh_is_byte_identical_to_before(tmp_path, monkeypatch): - """Back-compat: omitting keep_solid_id/bodies_key/post_mesh entirely - must behave exactly as before this task (aside from the version-bump - key change already covered above).""" - step = tmp_path / "p.step"; step.write_bytes(b"plain-bytes") - calls = [] - monkeypatch.setattr(mc, "heal_and_mesh", _fake_heal(calls)) - r1 = mc.get_or_mesh(str(step), str(tmp_path / "a.msh"), 8500, cache_dir=str(tmp_path / "c")) - r2 = mc.get_or_mesh(str(step), str(tmp_path / "b.msh"), 8500, cache_dir=str(tmp_path / "c")) - assert r1["cache"] == "miss" and r2["cache"] == "hit" - assert "per_body_patch_area_mm2" not in r2 - - -def test_get_or_mesh_post_mesh_exception_propagates_and_stores_nothing(tmp_path, monkeypatch): - """Review fix (LOW): if `post_mesh` raises (e.g. `tag_contact_patches` - finds an empty/all-covering patch), `get_or_mesh` must propagate the - exception AND must not have stored anything under that cache key -- - a subsequent call re-meshes (and re-runs post_mesh) from scratch, - rather than silently caching an untagged mesh under a bodies-aware key.""" - step = tmp_path / "p.step"; step.write_bytes(b"boom-bytes") - cache_dir = str(tmp_path / "c") - heal_calls = [] - def fake_heal(step_path, out_msh, target_nodes=8500, defeature=False, keep_solid_id=None, **kw): - heal_calls.append(1) - with open(out_msh, "w") as f: - f.write("UNTAGGED") - return {"ok": True, "n_nodes": 10, "notes": ""} - def failing_post_mesh(out_msh): - raise ValueError("body 2 does not touch the casting") - monkeypatch.setattr(mc, "heal_and_mesh", fake_heal) - with pytest.raises(ValueError, match="does not touch"): - mc.get_or_mesh(str(step), str(tmp_path / "a.msh"), 8500, cache_dir=cache_dir, - keep_solid_id=1, bodies_key="bh-boom", post_mesh=failing_post_mesh) - assert len(heal_calls) == 1 - # Nothing was stored under this key -- no .msh/.json for it in cache_dir. - key = mc.mesh_cache_key(str(step), 8500, bodies_key="bh-boom") - assert not os.path.exists(os.path.join(cache_dir, f"{key}.msh")) - assert not os.path.exists(os.path.join(cache_dir, f"{key}.json")) - # A subsequent call re-meshes (and re-runs post_mesh) from scratch -- - # this time succeeding -- rather than serving a stale/partial entry. - ok_calls = [] - def ok_post_mesh(out_msh): - ok_calls.append(out_msh) - return {"per_body_patch_area_mm2": {2: 5.0}, "n_adiab_tris": 1} - r2 = mc.get_or_mesh(str(step), str(tmp_path / "b.msh"), 8500, cache_dir=cache_dir, - keep_solid_id=1, bodies_key="bh-boom", post_mesh=ok_post_mesh) - assert len(heal_calls) == 2 and len(ok_calls) == 1 - assert r2["cache"] == "miss" and r2["per_body_patch_area_mm2"] == {2: 5.0} diff --git a/tests/test_prep_parts.py b/tests/test_meshing.py similarity index 59% rename from tests/test_prep_parts.py rename to tests/test_meshing.py index 123a4b50659d118c22afe6acb89d3544b8bcbc71..dfc7db034e63ef95df715d1d7fbb7ba4856d1a28 100644 --- a/tests/test_prep_parts.py +++ b/tests/test_meshing.py @@ -3,10 +3,10 @@ from pathlib import Path import numpy as np import pytest -import tools.doe.prep_parts as pp -from tools.doe.prep_parts import ( - largest_solid, is_ood, heal_and_mesh, _elem_size_from_area, _elem_size, - _first_guess_elem_size, K_SURF, _import_and_reduce, +import src.meshing as pp +from src.meshing import ( + heal_and_mesh, _elem_size_from_area, _elem_size, + _first_guess_elem_size, K_SURF, ) from src.msh_parse import parse_msh @@ -40,98 +40,20 @@ def _make_two_box_step(path, big=40.0, small=10.0): gmsh.finalize() -def test_largest_solid_keeps_max_volume_box(tmp_path): - step = tmp_path / "two_box.step" - _make_two_box_step(step, big=40.0, small=10.0) - out_step = tmp_path / "kept.step" - n_in, kept_vol = largest_solid(str(step), str(out_step)) - assert n_in == 2 - assert kept_vol == pytest.approx(40.0 ** 3, rel=1e-3) - # Re-import the exported STEP and confirm exactly 1 solid remains. - import gmsh - gmsh.initialize() - try: - gmsh.option.setNumber("General.Terminal", 0) - gmsh.model.add("check") - gmsh.model.occ.importShapes(str(out_step)) - gmsh.model.occ.synchronize() - vols = gmsh.model.occ.getEntities(3) - assert len(vols) == 1 - vol = gmsh.model.occ.getMass(3, vols[0][1]) - assert vol == pytest.approx(40.0 ** 3, rel=1e-3) - finally: - gmsh.finalize() - - -def test_largest_solid_single_solid_passthrough(tmp_path): - step = tmp_path / "one_box.step" - import gmsh - gmsh.initialize() - try: - gmsh.option.setNumber("General.Terminal", 0) - gmsh.model.add("one_box") - gmsh.model.occ.addBox(0, 0, 0, 20, 20, 20) - gmsh.model.occ.synchronize() - gmsh.write(str(step)) - finally: - gmsh.finalize() - - out_step = tmp_path / "kept.step" - n_in, kept_vol = largest_solid(str(step), str(out_step)) - assert n_in == 1 - assert kept_vol == pytest.approx(20.0 ** 3, rel=1e-3) -def test_is_ood_flags_large_diag(): - ranges = {"volume_mm3_max": 1.0e12, "bbox_diag_max": 2177.7} - assert is_ood(volume_mm3=1.0, bbox_diag_mm=2663.5, ranges=ranges) is True - assert is_ood(volume_mm3=1.0, bbox_diag_mm=100.0, ranges=ranges) is False -def test_is_ood_flags_large_volume(): - ranges = {"volume_mm3_max": 390663008.0, "bbox_diag_max": 2177.7} - assert is_ood(volume_mm3=5.0e9, bbox_diag_mm=100.0, ranges=ranges) is True - assert is_ood(volume_mm3=1.0e6, bbox_diag_mm=100.0, ranges=ranges) is False -def test_heal_and_mesh_synthetic_box(tmp_path): - """Small synthetic box -- fast, not a real dataset-D part -- exercises - the heal+mesh+tag pipeline end to end.""" - step = tmp_path / "box.step" - import gmsh - gmsh.initialize() - try: - gmsh.option.setNumber("General.Terminal", 0) - gmsh.model.add("box") - gmsh.model.occ.addBox(0, 0, 0, 30, 20, 10) - gmsh.model.occ.synchronize() - gmsh.write(str(step)) - finally: - gmsh.finalize() - - out_msh = tmp_path / "box.msh" - out_step = tmp_path / "box_cleaned.step" - result = heal_and_mesh(str(step), str(out_msh), target_nodes=500, out_step=str(out_step)) - - assert result["ok"] is True - assert result["n_solids"] == 1 - assert result["n_nodes"] > 0 - assert result["volume_mm3"] == pytest.approx(30 * 20 * 10, rel=1e-2) - assert np.isfinite(result["bbox_diag_mm"]) and result["bbox_diag_mm"] > 0 - assert out_msh.exists() - assert out_step.exists() - # -------------------------------------------------------------------------- -# Surface-area density-search seed (mesh-accel Task 4): the first-attempt -# elem size is now seeded from the OCC surface area of the imported model -# (`_elem_size_from_area`, via `_first_guess_elem_size`) instead of a fixed -# bbox_diag/28 guess, since the latter misses badly on non-typical shapes -# (see module docstring). K_SURF calibrated against BUKSA/H17 + 2 synthetic -# boxes (see task-4-report.md). +# Surface-area density-search seed: the first element size comes from the +# OCC surface area of the imported model (`_elem_size_from_area`, via +# `_first_guess_elem_size`), with bbox_diag/28 as the fallback. # -------------------------------------------------------------------------- def test_elem_size_from_area_matches_formula(): area_mm2, target_nodes = 2200.0, 500 @@ -263,8 +185,7 @@ def test_heal_and_mesh_effective_heal_then_succeeds(tmp_path, monkeypatch): # -------------------------------------------------------------------------- -# Untagged-mesh regression (H17.STEP bug, found 2026-07-28): `_tag_and_mesh` -# can run more than once within a single gmsh session (density-search +# Untagged-mesh regression: `_tag_and_mesh` can run more than once within a single gmsh session (density-search # retries) -- `removePhysicalGroups()` frees the *tag number* but not the # *name registration*, so on the 2nd+ call an auto-assigned new tag number # silently fails to bind "Volume"/"Cooling" again. The fix pins explicit @@ -275,7 +196,7 @@ def test_heal_and_mesh_effective_heal_then_succeeds(tmp_path, monkeypatch): def test_heal_and_mesh_multi_round_density_search_keeps_tags(tmp_path, monkeypatch): """A part whose first-attempt guess misses `target_nodes` by more than the 35% early-exit tolerance forces `_search_target_nodes` into - `search_elem_size`'s multi-round refinement -- i.e. `_tag_and_mesh` runs + `_search_elem_size`'s multi-round refinement -- i.e. `_tag_and_mesh` runs MORE THAN ONCE in the same gmsh session (verified during triage: without the tag-pinning fix, this reliably ends with an untagged mesh -- `parse_msh(...).tris == {}` -- because only the FIRST round's @@ -284,7 +205,7 @@ def test_heal_and_mesh_multi_round_density_search_keeps_tags(tmp_path, monkeypat this exact 30x20x10 @ target_nodes=500 fixture converge in ONE attempt (see `test_heal_and_mesh_area_seed_converges_on_first_attempt` above), K_SURF is monkeypatched to a deliberately bad value (50.0, which misses - this fixture by >90% -- see task-4-report.md) purely to force the + this fixture by >90%) purely to force the multi-round path back open. Regardless of how many rounds ran, the final written mesh must have a nonempty "Cooling" surface group and a nonempty tet volume.""" @@ -314,46 +235,8 @@ def test_heal_and_mesh_multi_round_density_search_keeps_tags(tmp_path, monkeypat assert mesh.tets.shape[0] > 0 -# -------------------------------------------------------------------------- -# Fast meshing options (mesh-accel Task 3): HXT (Mesh.Algorithm3D=10) + -# multithreading + curvature-refinement off, flag-gated via `fast=` and -# OFF by default -- these tests must pass with the default (fast=False) -# behavior completely unchanged. -# -------------------------------------------------------------------------- -def test_heal_and_mesh_fast_true_synthetic_box(tmp_path): - """fast=True must still produce a usable mesh: ok, a nonempty "Cooling" - surface group, and node count within the same +-35% density-search - band used elsewhere in this file for target_nodes=500 on this box.""" - step = tmp_path / "box.step" - _make_one_box_step(step, 30, 20, 10) - out_msh = tmp_path / "box.msh" - result = heal_and_mesh(str(step), str(out_msh), target_nodes=500, fast=True) - assert result["ok"] is True - assert 500 * 0.65 <= result["n_nodes"] <= 500 * 1.35 - - mesh = parse_msh(str(out_msh)) - assert "Cooling" in mesh.tris - assert mesh.tris["Cooling"].shape[0] > 0 - assert mesh.tets.shape[0] > 0 - - -def test_apply_fast_options_sets_hxt_threads_and_curvature_off(): - """Direct test of `_apply_fast_options()` inside an initialized gmsh - session: HXT tet algorithm, >=1 threads, curvature-based refinement - off.""" - import gmsh - - gmsh.initialize() - try: - gmsh.option.setNumber("General.Terminal", 0) - pp._apply_fast_options() - assert gmsh.option.getNumber("Mesh.Algorithm3D") == 10 - assert gmsh.option.getNumber("General.NumThreads") >= 1 - assert gmsh.option.getNumber("Mesh.MeshSizeFromCurvature") == 0 - finally: - gmsh.finalize() def test_fast_options_not_applied_in_plain_session(): @@ -373,28 +256,11 @@ def test_fast_options_not_applied_in_plain_session(): gmsh.finalize() -def test_heal_and_mesh_fast_default_never_calls_apply_fast_options(tmp_path, monkeypatch): - """`fast` defaults to False -- `heal_and_mesh` must not invoke - `_apply_fast_options` at all in that case (zero behavior change from - before this flag existed).""" - step = tmp_path / "box.step" - _make_one_box_step(step, 30, 20, 10) - - def _boom(): - raise AssertionError("_apply_fast_options must not be called when fast=False (default)") - - monkeypatch.setattr(pp, "_apply_fast_options", _boom) - - out_msh = tmp_path / "box.msh" - result = heal_and_mesh(str(step), str(out_msh), target_nodes=500) - assert result["ok"] is True # -------------------------------------------------------------------------- -# Defeaturing pass (mesh-accel Task 5): best-effort `removeAllDuplicates()` -# simplification, flag-gated via `defeature=` and OFF by default -- see -# `_apply_defeature`'s docstring / task-5-report.md for the fuller virtual- -# topology investigation (why nothing beyond `removeAllDuplicates` shipped). +# Defeaturing pass: `removeAllDuplicates()`, enabled via `defeature=` and +# off by default. # -------------------------------------------------------------------------- def test_heal_and_mesh_defeature_true_synthetic_box(tmp_path): """defeature=True on a clean box must still succeed with a nonempty @@ -437,7 +303,7 @@ def test_heal_and_mesh_defeature_true_calls_apply_defeature(tmp_path, monkeypatc def test_heal_and_mesh_destructive_defeature_returns_ok_false_not_a_crash(tmp_path, monkeypatch): - """Regression for the review finding (2026-07-29): if `_apply_defeature` + """If `_apply_defeature` reports the solid was lost (ok=False), `heal_and_mesh` must branch on that and return a clean `ok=False` dict with the solid-loss note -- NOT fall through to `gmsh.model.getBoundingBox(-1, -1)` on an empty @@ -492,109 +358,20 @@ def test_apply_defeature_removes_duplicates_without_losing_solid(): gmsh.finalize() -# -------------------------------------------------------------------------- -# `keep_solid_id` (Part B Task 3): `_import_and_reduce`/`heal_and_mesh` can -# keep a CALLER-designated OCC solid instead of the largest-volume one -- -# needed when the user has assigned solid roles and the "casting" is not -# necessarily the biggest body in the assembly (e.g. a small casting next to -# a bulky chill). `keep_solid_id=None` (default) must remain byte-identical -# to the pre-existing largest-solid behavior. -# -------------------------------------------------------------------------- -def test_import_and_reduce_keep_solid_id_keeps_designated_solid(tmp_path): - step = tmp_path / "two_box.step" - _make_two_box_step(step, big=40.0, small=10.0) - - import gmsh - gmsh.initialize() - try: - gmsh.option.setNumber("General.Terminal", 0) - gmsh.model.add("part") - # tag 2 is the SMALL box (see _make_two_box_step: big added first). - n_solids, kept_vol, err = _import_and_reduce(str(step), keep_solid_id=2) - assert err is None - assert n_solids == 2 - assert kept_vol == pytest.approx(10.0 ** 3, rel=1e-3) - vols = gmsh.model.occ.getEntities(3) - assert len(vols) == 1 - assert vols[0][1] == 2 - finally: - gmsh.finalize() - - -def test_import_and_reduce_keep_solid_id_none_keeps_largest_unchanged(tmp_path): - """`keep_solid_id=None` (default) must reproduce the pre-existing - largest-solid behavior exactly.""" - step = tmp_path / "two_box.step" - _make_two_box_step(step, big=40.0, small=10.0) - - import gmsh - gmsh.initialize() - try: - gmsh.option.setNumber("General.Terminal", 0) - gmsh.model.add("part") - n_solids, kept_vol, err = _import_and_reduce(str(step)) - assert err is None - assert n_solids == 2 - assert kept_vol == pytest.approx(40.0 ** 3, rel=1e-3) - finally: - gmsh.finalize() -def test_import_and_reduce_keep_solid_id_missing_tag_returns_error_string(tmp_path): - step = tmp_path / "two_box.step" - _make_two_box_step(step, big=40.0, small=10.0) - import gmsh - gmsh.initialize() - try: - gmsh.option.setNumber("General.Terminal", 0) - gmsh.model.add("part") - n_solids, kept_vol, err = _import_and_reduce(str(step), keep_solid_id=999) - assert err is not None - assert "999" in err - assert n_solids == 0 - assert kept_vol == 0.0 - finally: - gmsh.finalize() -def test_heal_and_mesh_keep_solid_id_meshes_designated_small_solid(tmp_path): - """`heal_and_mesh(..., keep_solid_id=)` must mesh the - SMALL solid (volume/node-count consistent with its own size), while the - largest-solid default on the SAME input is untouched (still meshes the - big one) -- regression coverage for both branches on one fixture.""" - step = tmp_path / "two_box.step" - _make_two_box_step(step, big=40.0, small=10.0) - out_msh_small = tmp_path / "small.msh" - result_small = heal_and_mesh(str(step), str(out_msh_small), target_nodes=300, - keep_solid_id=2) - assert result_small["ok"] is True - assert result_small["volume_mm3"] == pytest.approx(10.0 ** 3, rel=1e-2) - assert "solid_id=2" in result_small["notes"] - out_msh_default = tmp_path / "default.msh" - result_default = heal_and_mesh(str(step), str(out_msh_default), target_nodes=300) - assert result_default["ok"] is True - assert result_default["volume_mm3"] == pytest.approx(40.0 ** 3, rel=1e-2) - assert "kept largest" in result_default["notes"] -def test_heal_and_mesh_keep_solid_id_missing_tag_ok_false(tmp_path): - step = tmp_path / "two_box.step" - _make_two_box_step(step, big=40.0, small=10.0) - - out_msh = tmp_path / "missing.msh" - result = heal_and_mesh(str(step), str(out_msh), target_nodes=300, keep_solid_id=999) - assert result["ok"] is False - assert "999" in result["notes"] - assert not out_msh.exists() - def test_tag_and_mesh_fails_loudly_when_name_binding_missing(tmp_path, monkeypatch): """Direct regression test for the loud-fail guard itself: if gmsh ever fails to bind the "Volume"/"Cooling" physical-group names (the exact - silent-failure mode that produced H17's untagged mesh), `_tag_and_mesh` + silent failure that once produced an untagged mesh), `_tag_and_mesh` must return ok=False and must NOT write `out_msh` -- never silently hand back an untagged mesh.""" import gmsh @@ -620,3 +397,47 @@ def test_tag_and_mesh_fails_loudly_when_name_binding_missing(tmp_path, monkeypat assert not out_msh.exists() finally: gmsh.finalize() + + + +def test_heal_and_mesh_synthetic_box(tmp_path): + """A small synthetic box runs the import + mesh + tag pipeline end to end.""" + step = tmp_path / "box.step" + _make_one_box_step(step) + out_msh = tmp_path / "box.msh" + result = heal_and_mesh(str(step), str(out_msh), target_nodes=500) + + assert result["ok"] is True + assert result["n_solids"] == 1 + assert result["n_nodes"] > 0 + assert result["volume_mm3"] == pytest.approx(30 * 20 * 10, rel=1e-2) + assert np.isfinite(result["bbox_diag_mm"]) and result["bbox_diag_mm"] > 0 + mesh = parse_msh(str(out_msh)) + assert set(mesh.tris) == {"Cooling"} and mesh.tets.shape[0] > 0 + + +def test_heal_and_mesh_keeps_largest_of_two_solids(tmp_path): + step = tmp_path / "two.step" + import gmsh + gmsh.initialize() + try: + gmsh.option.setNumber("General.Terminal", 0) + gmsh.model.add("two_box") + gmsh.model.occ.addBox(0, 0, 0, 40, 40, 40) + gmsh.model.occ.addBox(140, 0, 0, 10, 10, 10) + gmsh.model.occ.synchronize() + gmsh.write(str(step)) + finally: + gmsh.finalize() + result = heal_and_mesh(str(step), str(tmp_path / "two.msh"), target_nodes=500) + assert result["ok"] is True and result["n_solids"] == 2 + assert result["volume_mm3"] == pytest.approx(40 ** 3, rel=1e-2) + assert "kept largest of 2 solids" in result["notes"] + + +def test_mesh_step_fixed_size(tmp_path): + step = tmp_path / "box.step" + _make_one_box_step(step) + info = pp.mesh_step(str(step), str(tmp_path / "fixed.msh"), elem_size=5.0) + assert info["elem_size"] == 5.0 and info["n_nodes"] > 0 and info["n_tets"] > 0 + assert set(parse_msh(str(tmp_path / "fixed.msh")).tris) == {"Cooling"} diff --git a/tests/test_metrics.py b/tests/test_metrics.py deleted file mode 100644 index 73c63df6f99155929a8a861a4eb855901bc720c7..0000000000000000000000000000000000000000 --- a/tests/test_metrics.py +++ /dev/null @@ -1,34 +0,0 @@ -import numpy as np -from src.metrics import (mae, rmse, median_rel_error, pred_std_ratio, - hotspot_accuracy) - - -def test_basic_errors(): - yt = np.array([0.0, 10.0, 20.0]) - yp = np.array([1.0, 10.0, 18.0]) - np.testing.assert_allclose(mae(yt, yp), (1 + 0 + 2) / 3) - np.testing.assert_allclose(rmse(yt, yp), np.sqrt((1 + 0 + 4) / 3)) - - -def test_pred_std_ratio_collapse(): - gid = np.array([0, 0, 0, 0]) - yt = np.array([0.0, 10.0, 20.0, 30.0]) - yp_flat = np.full(4, 15.0) # total collapse - assert pred_std_ratio(yt, yp_flat, gid) == 0.0 - assert abs(pred_std_ratio(yt, yt, gid) - 1.0) < 1e-9 - - -def test_hotspot_accuracy_perfect(): - gid = np.zeros(10, int) - yt = np.arange(10.0) # node 9 is the hot spot (top 10%) - yp = np.arange(10.0) - xyz = np.zeros((10, 3)) - assert hotspot_accuracy(yt, yp, xyz, gid, top_frac=0.1, tol_mm=5.0) == 1.0 - - -def test_hotspot_accuracy_wrong_far(): - gid = np.zeros(10, int) - yt = np.arange(10.0) # true hot spot = node 9 - yp = 9.0 - np.arange(10.0) # predicted hot spot = node 0 - xyz = np.zeros((10, 3)); xyz[0] = [100, 0, 0] # node0 far from node9 - assert hotspot_accuracy(yt, yp, xyz, gid, top_frac=0.1, tol_mm=5.0) == 0.0 diff --git a/tests/test_model_gnn.py b/tests/test_model_gnn.py index 6b6d507331fb14fcecfe281367952494f44dcfd7..85dc349924b63bf4cc9357dcff641101c2ca7636 100644 --- a/tests/test_model_gnn.py +++ b/tests/test_model_gnn.py @@ -1,187 +1,66 @@ -import pytest +import numpy as np import torch -from src.model_gnn import SolidGNN, attach_deriv_heads -# Frozen regression guard: the exact state_dict key set for a -# SolidGNN(node_in=9, edge_in=4, hidden=16, n_blocks=2, two_head=True) model -# BEFORE the deriv_heads feature was added. deriv_heads=False must reproduce -# this set exactly (no grad_head/dtdt_head keys, no renamed backbone keys). -_PRE_DERIV_HEADS_KEYS = { - "blocks.0.edge_mlp.0.bias", "blocks.0.edge_mlp.0.weight", - "blocks.0.edge_mlp.2.bias", "blocks.0.edge_mlp.2.weight", - "blocks.0.edge_norm.bias", "blocks.0.edge_norm.weight", - "blocks.0.node_mlp.0.bias", "blocks.0.node_mlp.0.weight", - "blocks.0.node_mlp.2.bias", "blocks.0.node_mlp.2.weight", - "blocks.0.node_norm.bias", "blocks.0.node_norm.weight", - "blocks.1.edge_mlp.0.bias", "blocks.1.edge_mlp.0.weight", - "blocks.1.edge_mlp.2.bias", "blocks.1.edge_mlp.2.weight", - "blocks.1.edge_norm.bias", "blocks.1.edge_norm.weight", - "blocks.1.node_mlp.0.bias", "blocks.1.node_mlp.0.weight", - "blocks.1.node_mlp.2.bias", "blocks.1.node_mlp.2.weight", - "blocks.1.node_norm.bias", "blocks.1.node_norm.weight", - "decoder.0.bias", "decoder.0.weight", - "decoder.2.bias", "decoder.2.weight", - "edge_encoder.0.0.bias", "edge_encoder.0.0.weight", - "edge_encoder.0.2.bias", "edge_encoder.0.2.weight", - "edge_encoder.1.bias", "edge_encoder.1.weight", - "node_encoder.0.0.bias", "node_encoder.0.0.weight", - "node_encoder.0.2.bias", "node_encoder.0.2.weight", - "node_encoder.1.bias", "node_encoder.1.weight", - "scale_head.0.bias", "scale_head.0.weight", - "scale_head.2.bias", "scale_head.2.weight", -} +from src.model_gnn import (SolidGNN, knn_edge_index, load_gnn, predict_graph_seconds, + union_edge_attr, union_edge_index) +CFG = {"node_in": 15, "edge_in": 5, "hidden": 16, "n_blocks": 2, "g_dim": 5, "k": 4} -def test_forward_shapes(): - torch.manual_seed(0) - N, E, B = 50, 400, 2 - x = torch.randn(N, 9) - edge_index = torch.randint(0, N, (2, E)) - edge_attr = torch.randn(E, 4) - batch = torch.zeros(N, dtype=torch.long); batch[25:] = 1 - model = SolidGNN(n_blocks=3, hidden=32) - yn, yg = model(x, edge_index, edge_attr, batch) - assert yn.shape == (N,) - assert yg.shape == (B,) +def _norm(): + return {"feat_mean": np.zeros(15, np.float32), "feat_std": np.ones(15, np.float32), + "y_log_mean": 3.0, "y_log_std": 0.5} -def test_global_is_max_of_nodal(): - torch.manual_seed(1) - N, E = 30, 200 - x = torch.randn(N, 9); edge_index = torch.randint(0, N, (2, E)); edge_attr = torch.randn(E, 4) - batch = torch.zeros(N, dtype=torch.long) - model = SolidGNN(n_blocks=2, hidden=16).eval() - with torch.no_grad(): - yn, yg = model(x, edge_index, edge_attr, batch) - torch.testing.assert_close(yg[0], yn.max()) - - -def test_two_head_forward_backward_and_shape_decomposition(): - import torch - from torch_geometric.utils import scatter - from src.model_gnn import SolidGNN - m = SolidGNN(node_in=9, edge_in=4, hidden=16, n_blocks=2, two_head=True) - x = torch.randn(12, 9) - ei = torch.stack([torch.arange(11), torch.arange(1, 12)]) - ei = torch.cat([ei, ei.flip(0)], dim=1) - ea = torch.randn(ei.shape[1], 4) - batch = torch.tensor([0] * 6 + [1] * 6) - yn, yg = m(x, ei, ea, batch) - assert yn.shape == (12,) and yg.shape == (2,) - yn.sum().backward() # grads flow to both heads - assert m.scale_head[0].weight.grad is not None - # per-graph mean of the output equals the scale head's prediction: - # dev part is exactly zero-mean per graph by construction - pm = scatter(yn, batch, dim=0, reduce="mean") - with torch.no_grad(): - h = m.node_encoder(x); e = m.edge_encoder(ea) - for blk in m.blocks: - h, e = blk(h, ei, e) - scale = m.scale_head(scatter(h, batch, dim=0, reduce="mean")).squeeze(-1) - torch.testing.assert_close(pm, scale, atol=1e-5, rtol=1e-5) - - -def test_deriv_heads_false_is_2tuple_and_state_dict_keys_unchanged(): - """deriv_heads=False (the default) must be BIT-identical to the - pre-deriv_heads module: same 2-tuple return, same state_dict key set - (no grad_head/dtdt_head keys, no accidental backbone renames).""" - torch.manual_seed(0) - m = SolidGNN(node_in=9, edge_in=4, hidden=16, n_blocks=2, two_head=True) - N, E, B = 12, 60, 2 - x = torch.randn(N, 9) - edge_index = torch.randint(0, N, (2, E)) - edge_attr = torch.randn(E, 4) - batch = torch.tensor([0] * 6 + [1] * 6) - - out = m(x, edge_index, edge_attr, batch) - assert len(out) == 2 - yn, yg = out - assert yn.shape == (N,) - assert yg.shape == (B,) - - assert set(m.state_dict().keys()) == _PRE_DERIV_HEADS_KEYS - - -def test_deriv_heads_true_returns_3tuple_with_correct_deriv_shape(): - torch.manual_seed(0) - m = SolidGNN(node_in=9, edge_in=4, hidden=16, n_blocks=2, two_head=True, - deriv_heads=True) - N, E, B = 12, 60, 2 - x = torch.randn(N, 9) - edge_index = torch.randint(0, N, (2, E)) - edge_attr = torch.randn(E, 4) - batch = torch.tensor([0] * 6 + [1] * 6) - out = m(x, edge_index, edge_attr, batch) - assert len(out) == 3 - yn, yg, deriv = out - assert yn.shape == (N,) - assert yg.shape == (B,) - assert deriv.shape == (N, 2) +def _state(seed=0): + torch.manual_seed(seed) + model = SolidGNN(node_in=15, edge_in=5, hidden=16, n_blocks=2) + return {"model": model.state_dict(), "norm": _norm(), "cfg": CFG} - # new heads exist as real modules with grad-tracked params - deriv.sum().backward() - assert m.grad_head[0].weight.grad is not None - assert m.dtdt_head[0].weight.grad is not None +def test_knn_edges_are_symmetric(): + pos = torch.rand(30, 3) + ei = knn_edge_index(pos, k=4) + pairs = set(map(tuple, ei.t().tolist())) + assert all((b, a) in pairs for a, b in pairs) -def test_deriv_heads_true_frozen_path_is_bit_identical_to_false(): - """The deriv heads read h but must never alter y_nodal/y_global: a - deriv_heads=True model loaded from a deriv_heads=False model's state_dict - (via attach_deriv_heads) must produce torch.equal y_nodal/y_global on the - same inputs -- true bit-identity of the frozen path, not just close.""" - torch.manual_seed(7) - base = SolidGNN(node_in=9, edge_in=4, hidden=16, n_blocks=2, two_head=True) - base.eval() - m = SolidGNN(node_in=9, edge_in=4, hidden=16, n_blocks=2, two_head=True, - deriv_heads=True) - attach_deriv_heads(m, base.state_dict()) - m.eval() +def test_union_graph_contains_mesh_edges_and_flags_them(): + pos = torch.rand(20, 3) + mesh_edges = torch.tensor([[0, 19], [3, 7]]) + ei, is_mesh = union_edge_index(pos, mesh_edges, k=3) + pairs = {tuple(p): f for p, f in zip(ei.t().tolist(), is_mesh.tolist())} + for a, b in [(0, 19), (19, 0), (3, 7), (7, 3)]: + assert pairs[(a, b)] == 1.0 + assert len(pairs) == ei.shape[1] # no duplicates + ea = union_edge_attr(pos, ei, is_mesh) + assert ea.shape == (ei.shape[1], 5) + src, dst = ei + np.testing.assert_allclose(ea[:, 3], (pos[dst] - pos[src]).norm(dim=1), rtol=1e-6) - N, E, B = 14, 70, 2 - 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) +def test_film_conditioning_changes_output(): + model, cfg, norm = load_gnn(_state(), "cpu") with torch.no_grad(): - yn_base, yg_base = base(x, edge_index, edge_attr, batch) - yn_new, yg_new, deriv = m(x, edge_index, edge_attr, batch) - - assert torch.equal(yn_base, yn_new) - assert torch.equal(yg_base, yg_new) - assert deriv.shape == (N, 2) - - -def test_attach_deriv_heads_missing_keys_are_exactly_the_new_heads(): - torch.manual_seed(3) - base = SolidGNN(node_in=9, edge_in=4, hidden=16, n_blocks=2, two_head=True) - m = SolidGNN(node_in=9, edge_in=4, hidden=16, n_blocks=2, two_head=True, - deriv_heads=True) - - missing = attach_deriv_heads(m, base.state_dict()) - - # Derive the expected set from _mlp's own structure (grad_head/dtdt_head - # are each an nn.Sequential(Linear, ReLU, Linear) -- 4 param tensors - # apiece: {0,2}.{weight,bias}), rather than hardcoding a magic count. - expected = {k for k in m.state_dict() if k.startswith(("grad_head.", "dtdt_head."))} - assert expected, "expected grad_head/dtdt_head keys to exist on a deriv_heads=True model" - assert set(missing) == expected - - # every remaining (backbone) key must have been loaded, i.e. not missing - non_head_missing = [k for k in missing if k not in expected] - assert non_head_missing == [] - - -def test_attach_deriv_heads_missing_backbone_key_raises_loudly(): - torch.manual_seed(3) - base = SolidGNN(node_in=9, edge_in=4, hidden=16, n_blocks=2, two_head=True) - m = SolidGNN(node_in=9, edge_in=4, hidden=16, n_blocks=2, two_head=True, - deriv_heads=True) - - sd = dict(base.state_dict()) - del sd["decoder.0.weight"] # simulate a corrupted/mismatched checkpoint - - with pytest.raises(ValueError): - attach_deriv_heads(m, sd) + for film in model.films: # make FiLM non-trivial + film.net[-1].weight.normal_(0, 0.5) + rng = np.random.default_rng(0) + X, pos = rng.normal(size=(25, 15)), rng.uniform(0, 50, size=(25, 3)) + edges = np.array([[0, 1], [1, 2], [2, 3]]) + a = predict_graph_seconds(model, cfg, norm, X, pos, edges, np.zeros(5)) + b = predict_graph_seconds(model, cfg, norm, X, pos, edges, np.ones(5)) + assert a.shape == (25,) and np.all(a >= 0) + assert not np.allclose(a, b) + + +def test_prediction_ignores_node_order(): + model, cfg, norm = load_gnn(_state(1), "cpu") + rng = np.random.default_rng(1) + X, pos = rng.normal(size=(25, 15)), rng.uniform(0, 50, size=(25, 3)) + edges = np.array([[0, 1], [1, 2], [5, 9]]) + g = rng.normal(size=5) + a = predict_graph_seconds(model, cfg, norm, X, pos, edges, g) + perm = rng.permutation(25) + inv = np.argsort(perm) + b = predict_graph_seconds(model, cfg, norm, X[perm], pos[perm], inv[edges], g) + np.testing.assert_allclose(a[perm], b, rtol=1e-4) diff --git a/tests/test_ood.py b/tests/test_ood.py index 9a9c78b1336b06a67c508547a8ab9480a321f053..972e8b315359871c423e179a0b5f7b01e39f3ae4 100644 --- a/tests/test_ood.py +++ b/tests/test_ood.py @@ -1,33 +1,19 @@ -import numpy as np -from src.ood import build_ood_ranges, ood_flags, save_ranges, load_ranges +import json +from src.ood import load_ranges, ood_flags -def _cache(tmp_path): - g = tmp_path / "graphs"; g.mkdir(parents=True) - # FEATURE_NAMES index: volume_mm3=6, geodesic_d_cool=0 - for name, vol, gc, extent in [("A", 100.0, 5.0, 10.0), ("B", 400.0, 9.0, 40.0)]: - N = 8 - X = np.zeros((N, 9), np.float32); X[:, 6] = vol; X[:, 0] = np.linspace(0, gc, N) - pos = np.zeros((N, 3), np.float32); pos[-1] = extent # bbox diag ~ extent*sqrt3 - np.savez(g / f"{name}.npz", X=X, y=np.ones(N, np.float32), pos=pos) - return tmp_path +RANGES = {"volume_mm3_max": 400.0, "bbox_diag_max": 69.3, "geodesic_d_cool_max": 9.0} -def test_build_and_flags(tmp_path): - root = _cache(tmp_path) - r = build_ood_ranges(str(root)) - assert abs(r["volume_mm3_max"] - 400.0) < 1e-6 - assert abs(r["geodesic_d_cool_max"] - 9.0) < 1e-6 - assert r["bbox_diag_max"] > 40.0 # sqrt(3)*40 +def test_load_ranges(tmp_path): + p = tmp_path / "ood.json" + p.write_text(json.dumps(RANGES)) + assert load_ranges(str(p)) == RANGES - # inside range -> no warnings - assert ood_flags(300.0, 50.0, 8.0, r) == [] - # volume far above max -> a warning mentioning volume - w = ood_flags(10000.0, 50.0, 8.0, r) - assert len(w) == 1 and "volume" in w[0].lower() - -def test_roundtrip(tmp_path): - r = {"volume_mm3_max": 1.0, "bbox_diag_max": 2.0, "geodesic_d_cool_max": 3.0} - p = tmp_path / "r.json"; save_ranges(r, str(p)) - assert load_ranges(str(p)) == r +def test_flags_only_beyond_tolerance(): + assert ood_flags(400.0, 69.3, 9.0, RANGES) == [] + assert ood_flags(410.0, 70.0, 9.3, RANGES) == [] # within the 5% tolerance + w = ood_flags(1000.0, 200.0, 30.0, RANGES) + assert len(w) == 3 + assert "volume" in w[0] and "bounding-box" in w[1] and "distance-to-cooling" in w[2] diff --git a/tests/test_predict_casting.py b/tests/test_predict_casting.py new file mode 100644 index 0000000000000000000000000000000000000000..e1b24b8e3a7494183364388d3ad1271b788367ff --- /dev/null +++ b/tests/test_predict_casting.py @@ -0,0 +1,106 @@ +"""predict_casting orchestration with meshing and the model replaced by fakes, +plus an end-to-end run on a synthetic STEP when the weights are present.""" +import os +import shutil + +import meshio +import numpy as np +import pytest + +import src.predict_casting as pc +from src.features import FEATURE_NAMES +from src.msh_parse import parse_msh + + +class _FakeModel: + def predict_with_props(self, mesh, props): + n = mesh.node_xyz.shape[0] + X = np.zeros((n, len(FEATURE_NAMES)), np.float32) + X[:, FEATURE_NAMES.index("volume_mm3")] = 1000.0 / 6.0 + seconds = np.arange(n, dtype=np.float64) + 1.0 + info = {"a": 1.0, "m": 0.0, "target_unreachable": False, "clip_frac": 0.0, + "total_deviation_frac": 0.0} + return seconds, {"X": X, "total_s": float(n), "spread_s": 1.0, "recomposition": info} + + +def _fake_get_or_mesh(mini_msh_file, calls): + def _fn(step_path, out_msh, target_nodes=None, cache_dir=None): + calls.append(target_nodes) + shutil.copy(mini_msh_file, out_msh) + return {"ok": True, "n_nodes": 5, "cache": "miss"} + return _fn + + +def test_writes_vtu_and_summary(monkeypatch, mini_msh_file, tmp_path): + calls = [] + monkeypatch.setattr(pc, "get_or_mesh", _fake_get_or_mesh(mini_msh_file, calls)) + r = pc.predict_casting("part.step", out_dir=str(tmp_path), model=_FakeModel(), + ood_json="missing.json") + assert calls == [pc.DEFAULT_TARGET_NODES] + assert r["total_time_s"] == 5.0 and r["n_nodes"] == 5 + assert r["hotspot_xyz"].shape == (5, 3) + vtu = meshio.read(r["vtu"]) + assert {"solidification_time", "hotspot", "niyama", "porosity_risk", + "macro_region_id"} <= set(vtu.point_data) + text = open(r["summary"]).read() + assert "total solidification time: 5.0 s" in text + assert "alloy: EN-GJS-400-15 mold: green_sand model: v5.1" in text + assert "shrinkage screening:" in text + assert r["shrinkage"]["niyama_source"] == "proxy" + + +def test_bad_alloy_fails_before_meshing(monkeypatch, tmp_path): + def boom(*a, **k): + raise AssertionError("meshing must not run") + monkeypatch.setattr(pc, "get_or_mesh", boom) + with pytest.raises(ValueError, match="unknown alloy"): + pc.predict_casting("part.step", out_dir=str(tmp_path), alloy="nope", + model=_FakeModel()) + + +def test_no_shrinkage_and_ood_warning(monkeypatch, mini_msh_file, tmp_path): + monkeypatch.setattr(pc, "get_or_mesh", _fake_get_or_mesh(mini_msh_file, [])) + ood = tmp_path / "ood.json" + ood.write_text('{"volume_mm3_max": 1.0, "bbox_diag_max": 1000.0, ' + '"geodesic_d_cool_max": 1000.0}') + r = pc.predict_casting("part.step", out_dir=str(tmp_path), model=_FakeModel(), + screen_shrinkage=False, ood_json=str(ood)) + assert r["shrinkage"] is None + assert len(r["ood_warnings"]) == 1 and "volume" in r["ood_warnings"][0] + assert "porosity_risk" not in meshio.read(r["vtu"]).point_data + + +def test_elem_size_and_target_nodes_are_exclusive(tmp_path): + with pytest.raises(ValueError): + pc.predict_casting("part.step", out_dir=str(tmp_path), elem_size=4.0, + target_nodes=9000, model=_FakeModel()) + + +def test_parse_gravity(): + assert pc._parse_gravity("0, 0, -1") == (0.0, 0.0, -1.0) + for bad in ("1,2", "0,0,0", "a,b,c"): + with pytest.raises(ValueError): + pc._parse_gravity(bad) + + +@pytest.mark.skipif(not os.path.exists("weights/gnn_v5.1.safetensors"), + reason="published weights not present") +def test_end_to_end_on_a_box(tmp_path): + import gmsh + step = tmp_path / "box.step" + gmsh.initialize() + try: + gmsh.option.setNumber("General.Terminal", 0) + gmsh.model.add("box") + gmsh.model.occ.addBox(0, 0, 0, 120, 80, 40) + gmsh.model.occ.synchronize() + gmsh.write(str(step)) + finally: + gmsh.finalize() + r = pc.predict_casting(str(step), out_dir=str(tmp_path / "out"), + cache_dir=str(tmp_path / "cache"), device="cpu") + t = meshio.read(r["vtu"]).point_data["solidification_time"] + assert np.all(np.isfinite(t)) and t.min() >= 0 and r["total_time_s"] > 0 + assert set(parse_msh(r["msh"]).tris) == {"Cooling"} + # the hottest point of a box sits near its centre + assert np.linalg.norm(r["hotspot_xyz"][0] - np.array([60, 40, 20])) < 25 diff --git a/tests/test_predict_casting_v4.py b/tests/test_predict_casting_v4.py deleted file mode 100644 index 738985da24ae66542e9523189a7d2c09c9d48971..0000000000000000000000000000000000000000 --- a/tests/test_predict_casting_v4.py +++ /dev/null @@ -1,529 +0,0 @@ -"""Task 3: film-aware v4 inference path in predict_field/predict_casting. - -Real tiny in-process FiLM SolidGNN + a stub v2-total GBM (mirrors -tests/test_evaluate_ft2b.py / tests/test_select_ckpt_ft2b.py's fixture -style) exercise the new v4 (film) branch of `predict_field`, while a -plain SP1-shaped (non-film) call proves the OLD code path is untouched. - -Task 8: prep-aligned inference meshing -- `predict_casting`'s meshing -stage routes through `tools.doe.prep_parts.heal_and_mesh` (largest-solid -+ OCC heal + density search, same pipeline as the training/prep dataset) -by default (`mesh_mode="prep"`), with `mesh_mode="raw"` (or an explicit -`elem_size`) as an escape hatch to the old `src.step_mesh.mesh_step` -path. The routing tests below monkeypatch both meshing functions as -recorders (plus `load_hybrid`/`predict_field`) -- no real gmsh/model -checkpoints are touched. -""" -import json -import logging -import os.path as osp -import shutil - -import numpy as np -import pytest -import torch -from scipy.spatial import Delaunay -from sklearn.ensemble import HistGradientBoostingRegressor - -import src.predict_casting as pc_mod -from src.features import FEATURE_NAMES_V5 -from src.inference_props import resolve_props -from src.model_gnn import SolidGNN -from src.msh_parse import Mesh, parse_msh -from src.predict_casting import predict_field - -N_COLS_V4 = 15 - - -def _all_cooling_cube(n=4, a=20.0): - """A Delaunay-tetrahedralized cube with every hull face tagged - 'Cooling' (no Adiabatic patch) -- mirrors test_features_v4.py's - `_cube_with_bc_patch` mesh-construction style.""" - 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) - return Mesh(node_xyz=pts, tris={"Cooling": hull, "Adiabatic": np.empty((0, 3), np.int64)}, - tets=tets) - - -def _make_film_gnn(seed=0): - """A real tiny film SolidGNN with a NON-zero-init FiLM (the last FiLM - layer is zero-init by construction -- see model_gnn.FiLM -- so g would - be a no-op unless we perturb it here, exactly as the brief requires).""" - torch.manual_seed(seed) - model = SolidGNN(node_in=N_COLS_V4, edge_in=5, hidden=16, n_blocks=2, - two_head=True, film=True, g_dim=5) - rng = np.random.default_rng(seed) - with torch.no_grad(): - for film in model.films: - w = rng.normal(scale=2.0, size=tuple(film.net[-1].weight.shape)).astype(np.float32) - b = rng.normal(scale=2.0, size=tuple(film.net[-1].bias.shape)).astype(np.float32) - film.net[-1].weight.copy_(torch.tensor(w)) - film.net[-1].bias.copy_(torch.tensor(b)) - model.eval() - return model - - -def _norm15(): - # y_log_mean=3.0 (not 0.0): an untrained model's raw per-node output has - # only a tiny spread (~1e-2, this architecture's "dev" term); centering - # log1p(seconds) near 0 would expm1-clip that whole spread to exactly - # 0 for every node (seen in practice), masking any FiLM/g effect. A - # moderate y_log_mean keeps the (tiny but real) per-node spread inside - # expm1's non-degenerate region so alloy-driven differences survive. - return {"feat_mean": np.zeros(N_COLS_V4, dtype=np.float32), - "feat_std": np.ones(N_COLS_V4, dtype=np.float32), - "g_mean": np.zeros(5), "g_std": np.ones(5), - "y_log_mean": np.float64(3.0), "y_log_std": np.float64(1.0)} - - -def _film_cfg(): - return {"hidden": 16, "n_blocks": 2, "k": 4, "graph": "union", - "feature_version": 4, "film": True, "g_dim": 5, "two_head": True, - "node_in": N_COLS_V4, "edge_in": 5} - - -class _StubHGB: - """A v2-total GBM stand-in: `.predict` ignores its input and always - returns the SAME log1p(total_seconds) row -- enough to pin the - max-anchor recomposition exactly, without a real gbm_ft2b fit.""" - - def __init__(self, n_features_in_, total_seconds=50.0): - self.n_features_in_ = n_features_in_ - self._log_pred = float(np.log1p(total_seconds)) - - def predict(self, X): - return np.full(X.shape[0], self._log_pred) - - -# --------------------------------------------------------------------------- # -# film (v4) path -# --------------------------------------------------------------------------- # -def test_film_path_max_anchors_to_stub_gbm_total(): - mesh = _all_cooling_cube() - gnn = _make_film_gnn(seed=0) - hgb = _StubHGB(n_features_in_=16, total_seconds=50.0) - props = resolve_props() - - out, feat = predict_field(mesh, gnn, hgb, _norm15(), cfg=_film_cfg(), - device="cpu", props=props) - - assert out.shape == (mesh.node_xyz.shape[0],) - assert np.isfinite(out).all() and (out >= 0).all() - # Task 7: hybrid_field_max_anchor's default anchor is a high quantile - # (q=0.995), not the exact max, so `out.max()` need only be AT LEAST the - # stub total (the anchor node itself, and any node above it, gets - # shifted to >= total) -- it is no longer pinned to exactly 50.0. The - # excess is bounded by the top-quantile spread of the field, which is - # small here. - assert out.max() >= 50.0 - 1e-6 - assert out.max() <= 50.0 + 1.0 - assert feat["X"].shape == (mesh.node_xyz.shape[0], N_COLS_V4) - - -def test_film_path_changing_alloy_changes_field(): - mesh = _all_cooling_cube() - gnn = _make_film_gnn(seed=0) - hgb = _StubHGB(n_features_in_=16, total_seconds=50.0) - - props_a = resolve_props(alloy="EN-GJS-400-15") - props_b = resolve_props(alloy="G20Mn5") - out_a, _ = predict_field(mesh, gnn, hgb, _norm15(), cfg=_film_cfg(), - device="cpu", props=props_a) - out_b, _ = predict_field(mesh, gnn, hgb, _norm15(), cfg=_film_cfg(), - device="cpu", props=props_b) - - # Both anchor to the SAME stub total (50.0) -- only the field SHAPE - # (the FiLM-conditioned deviations from that quantile anchor) can - # differ. Per Task 7, the default quantile anchor means `max() >= total` - # rather than `== total` exactly (see the assertion above this fixture). - assert out_a.max() >= 50.0 - 1e-6 - assert out_b.max() >= 50.0 - 1e-6 - assert not np.allclose(out_a, out_b) - - -def test_film_path_without_bc_plan_spec_stays_all_sand_literal(monkeypatch): - """`bc_plan=None` (the default) -- byte-identical to the pre-Part-B - behavior: the all-sand literal spec, and `bc_spec=None` into the GBM - call (its own documented default, "plain sand mold").""" - mesh = _all_cooling_cube() - gnn = _make_film_gnn(seed=0) - hgb = _StubHGB(n_features_in_=16, total_seconds=50.0) - props = resolve_props() - - captured_spec = {} - real_compute = pc_mod.compute_features_inference - - def spy_compute(mesh_, version=2, spec=None): - captured_spec["spec"] = spec - return real_compute(mesh_, version=version, spec=spec) - - captured_bc_spec = {} - real_gbm = pc_mod.gbm_v2_total_seconds - - def spy_gbm(hgb_model, props_, v_over_a, bc_spec=None): - captured_bc_spec["bc_spec"] = bc_spec - return real_gbm(hgb_model, props_, v_over_a, bc_spec=bc_spec) - - monkeypatch.setattr(pc_mod, "compute_features_inference", spy_compute) - monkeypatch.setattr(pc_mod, "gbm_v2_total_seconds", spy_gbm) - - predict_field(mesh, gnn, hgb, _norm15(), cfg=_film_cfg(), device="cpu", props=props) - - assert captured_spec["spec"] == {"mold_effusivity": props["mold_effusivity"], - "bc_config": "none", "chill_effusivity": 0.0, - "sleeve_MEF": 0.0, "exo_q0": 0.0} - assert captured_bc_spec["bc_spec"] is None - - -def test_film_ckpt_with_wrong_dim_hgb_raises(): - mesh = _all_cooling_cube() - gnn = _make_film_gnn(seed=0) - hgb = _StubHGB(n_features_in_=9, total_seconds=50.0) # SP1-shaped baseline, wrong for v4 - props = resolve_props() - - with pytest.raises(ValueError): - predict_field(mesh, gnn, hgb, _norm15(), cfg=_film_cfg(), device="cpu", props=props) - - -def test_film_cfg_without_props_raises(): - mesh = _all_cooling_cube() - gnn = _make_film_gnn(seed=0) - hgb = _StubHGB(n_features_in_=16, total_seconds=50.0) - - with pytest.raises(ValueError): - predict_field(mesh, gnn, hgb, _norm15(), cfg=_film_cfg(), device="cpu", props=None) - - -# --------------------------------------------------------------------------- # -# Task 4 (bc-reach): cfg-driven v5 (20-column) feature building at inference. -# Mirrors the v4 fixtures above (_make_film_gnn/_norm15/_film_cfg) at -# node_in=20/feature_version=5 -- `predict_field`'s film branch must read -# `cfg.get("feature_version", 4)` (the SAME pattern the non-film branch -# already uses, predict_casting.py:354) instead of hardcoding `version=4`, -# so a v5 (bc-reach) checkpoint's cfg drives `compute_features_inference` -# into the 20-column builder. `cfg.feature_version` absent or ==4 must stay -# byte-identical to every test above (unmodified). -# --------------------------------------------------------------------------- # -N_COLS_V5 = 20 - - -def _make_film_gnn_v5(seed=0): - """`_make_film_gnn` at node_in=N_COLS_V5 (20) -- otherwise identical - (non-zero-init FiLM last layer, so `g` is not a no-op).""" - torch.manual_seed(seed) - model = SolidGNN(node_in=N_COLS_V5, edge_in=5, hidden=16, n_blocks=2, - two_head=True, film=True, g_dim=5) - rng = np.random.default_rng(seed) - with torch.no_grad(): - for film in model.films: - w = rng.normal(scale=2.0, size=tuple(film.net[-1].weight.shape)).astype(np.float32) - b = rng.normal(scale=2.0, size=tuple(film.net[-1].bias.shape)).astype(np.float32) - film.net[-1].weight.copy_(torch.tensor(w)) - film.net[-1].bias.copy_(torch.tensor(b)) - model.eval() - return model - - -def _norm20(): - return {"feat_mean": np.zeros(N_COLS_V5, dtype=np.float32), - "feat_std": np.ones(N_COLS_V5, dtype=np.float32), - "g_mean": np.zeros(5), "g_std": np.ones(5), - "y_log_mean": np.float64(3.0), "y_log_std": np.float64(1.0)} - - -def _film_cfg_v5(): - return {"hidden": 16, "n_blocks": 2, "k": 4, "graph": "union", - "feature_version": 5, "film": True, "g_dim": 5, "two_head": True, - "node_in": N_COLS_V5, "edge_in": 5} - - -def test_v4_cfg_still_builds_15_columns_unchanged(): - """(a) The existing v4 cfg (feature_version=4, node_in=15) must still - produce exactly the 15-column X every test above already asserts -- - the binding "cfg.feature_version==4 (or absent) -> byte-identical - current behavior" contract, re-stated as its own direct test rather - than relying only on the untouched tests above.""" - mesh = _all_cooling_cube() - gnn = _make_film_gnn(seed=0) - hgb = _StubHGB(n_features_in_=16, total_seconds=50.0) - props = resolve_props() - - out, feat = predict_field(mesh, gnn, hgb, _norm15(), cfg=_film_cfg(), - device="cpu", props=props) - assert feat["X"].shape == (mesh.node_xyz.shape[0], N_COLS_V4) - - -def test_film_cfg_missing_feature_version_key_defaults_to_15_columns(): - """(a) `cfg.get("feature_version", 4)` -- a film cfg with NO - "feature_version" key at all (not just one explicitly set to 4) must - still default to the v4 (15-column) builder.""" - mesh = _all_cooling_cube() - gnn = _make_film_gnn(seed=0) - hgb = _StubHGB(n_features_in_=16, total_seconds=50.0) - props = resolve_props() - cfg = {k: v for k, v in _film_cfg().items() if k != "feature_version"} - - out, feat = predict_field(mesh, gnn, hgb, _norm15(), cfg=cfg, - device="cpu", props=props) - assert feat["X"].shape == (mesh.node_xyz.shape[0], N_COLS_V4) - - -def test_v5_cfg_film_path_builds_20_column_x_all_sand_last_five_zero(): - """(b) A v5-cfg stub ckpt (cfg.feature_version=5, node_in=20) drives - `compute_features_inference` into the 20-column builder; a bare - all-sand mesh (no bc_plan, no Adiabatic patch) leaves the 5 new - BC-reach channels (exo_mid/exo_long/chill_short/chill_mid/chill_long) - exactly zero -- matching src.features's own all-sand contract.""" - mesh = _all_cooling_cube() - gnn = _make_film_gnn_v5(seed=0) - hgb = _StubHGB(n_features_in_=16, total_seconds=50.0) - props = resolve_props() - - out, feat = predict_field(mesh, gnn, hgb, _norm20(), cfg=_film_cfg_v5(), - device="cpu", props=props) - X = feat["X"] - assert X.shape == (mesh.node_xyz.shape[0], N_COLS_V5) - assert out.shape == (mesh.node_xyz.shape[0],) - assert np.isfinite(out).all() - assert (X[:, 15:] == 0.0).all() - - -# --------------------------------------------------------------------------- # -# SP1 (non-film) path: byte-identical regression -# --------------------------------------------------------------------------- # -def test_sp1_path_regression_still_routes_through_hybrid_field(monkeypatch, mini_msh_file): - mesh = parse_msh(mini_msh_file) - gnn = SolidGNN(node_in=9, edge_in=4, hidden=16, n_blocks=2) - norm = {"feat_mean": np.zeros(9), "feat_std": np.ones(9), - "y_log_mean": 0.0, "y_log_std": 1.0} - rng = np.random.default_rng(0) - hgb = HistGradientBoostingRegressor(max_iter=5).fit(rng.normal(size=(50, 9)), - rng.normal(size=50)) - - sentinel = np.array([111.0]) - called = {} - - def _fake_hybrid_field(g, b): - called["yes"] = True - return sentinel - - monkeypatch.setattr(pc_mod, "hybrid_field", _fake_hybrid_field) - - cfg = {"k": 3, "graph": "knn", "feature_version": 2} # no "film" key -> non-film path - out, _ = predict_field(mesh, gnn, hgb, norm, cfg=cfg, device="cpu") - - assert called.get("yes") is True - assert np.array_equal(out, sentinel) - - -# --------------------------------------------------------------------------- # -# load_hybrid: film-ckpt detection -# --------------------------------------------------------------------------- # -def test_load_hybrid_detects_film_ckpt_and_skips_baseline_norm(tmp_path): - gnn = _make_film_gnn(seed=0) - ck_path = tmp_path / "film.pt" - torch.save({"model": gnn.state_dict(), "norm": _norm15(), "cfg": _film_cfg()}, ck_path) - - loaded_gnn, hgb, norm, cfg, baseline_norm, hgb_std = pc_mod.load_hybrid( - str(ck_path), hgb_path=None, device="cpu") - - assert cfg.get("film") is True - assert baseline_norm is None - assert hgb is None - assert hgb_std is None - assert norm["y_log_mean"] == pytest.approx(3.0) - - -def test_load_hybrid_non_film_ckpt_unchanged(tmp_path): - gnn = SolidGNN(node_in=9, edge_in=4, hidden=16, n_blocks=2) - ck_path = tmp_path / "sp1.pt" - norm = {"feat_mean": np.zeros(9), "feat_std": np.ones(9), - "y_log_mean": 0.0, "y_log_std": 1.0} - torch.save({"model": gnn.state_dict(), "norm": norm, - "cfg": {"hidden": 16, "n_blocks": 2, "k": 6}}, ck_path) - - loaded_gnn, hgb, loaded_norm, cfg, baseline_norm, hgb_std = pc_mod.load_hybrid( - str(ck_path), hgb_path=None, device="cpu") - - assert not cfg.get("film", False) - assert baseline_norm is None # hgb_path is None here too, so unchanged either way - assert hgb_std is None - - -def test_sp1_path_with_wrong_dim_hgb_raises(mini_msh_file): - """Vice-versa guard: a v4 (16-feature) GBM fed to the SP1/non-film path - must also fail loudly, not silently misbehave or crash inside sklearn.""" - mesh = parse_msh(mini_msh_file) - gnn = SolidGNN(node_in=9, edge_in=4, hidden=16, n_blocks=2) - norm = {"feat_mean": np.zeros(9), "feat_std": np.ones(9), - "y_log_mean": 0.0, "y_log_std": 1.0} - hgb = _StubHGB(n_features_in_=16, total_seconds=50.0) - - cfg = {"k": 3, "graph": "knn", "feature_version": 2} - with pytest.raises(ValueError): - predict_field(mesh, gnn, hgb, norm, cfg=cfg, device="cpu") - - -# --------------------------------------------------------------------------- # -# Task 8: mesh_mode routing (prep vs. raw) -- recorder-only, no real gmsh -# --------------------------------------------------------------------------- # -def _stub_prediction_stage(monkeypatch, n_nodes): - """Bypass the model-loading + field-prediction stages entirely (no real - checkpoints touched) so these tests exercise ONLY predict_casting's - meshing-stage routing.""" - monkeypatch.setattr(pc_mod, "load_hybrid", - lambda *a, **kw: (None, None, {}, {}, None, None)) - monkeypatch.setattr(pc_mod, "predict_field", - lambda mesh, *a, **kw: (np.arange(n_nodes, dtype=np.float64) + 1.0, - {"X": np.zeros((n_nodes, 9))})) - - -def _fake_get_or_mesh(calls, mini_msh_file, ok=True, cache="miss"): - def _fn(step_path, out_msh, target_nodes=None, cache_dir=None, - keep_solid_id=None, bodies_key="none", post_mesh=None): - calls["get_or_mesh"] = dict(step_path=step_path, out_msh=out_msh, - target_nodes=target_nodes, cache_dir=cache_dir, - keep_solid_id=keep_solid_id, bodies_key=bodies_key, - post_mesh=post_mesh) - shutil.copy(mini_msh_file, out_msh) - result = dict(n_solids=1, volume_mm3=1.0, bbox_diag_mm=17.3, n_nodes=5, - ok=ok, healed=False, notes="", cache=cache) - # Part B Task 4: exercise the tagging-before-store contract exactly - # as the real get_or_mesh does -- call the post_mesh hook (if any) - # on this "miss" and merge its return into the meta. - if post_mesh is not None: - result.update(post_mesh(out_msh)) - return result - return _fn - - -def _fake_mesh_step(calls, mini_msh_file): - def _fn(step_path, out_msh, elem_size=None): - calls["mesh_step"] = dict(step_path=step_path, out_msh=out_msh, elem_size=elem_size) - shutil.copy(mini_msh_file, out_msh) - return dict(n_nodes=5, n_tets=1, bbox_mm=(10.0, 10.0, 10.0), elem_size=elem_size or 1.0) - return _fn - - -def test_predict_casting_default_routes_through_get_or_mesh( - monkeypatch, mini_msh_file, tmp_path): - calls = {} - monkeypatch.setattr(pc_mod, "get_or_mesh", _fake_get_or_mesh(calls, mini_msh_file)) - monkeypatch.setattr(pc_mod, "mesh_step", _fake_mesh_step(calls, mini_msh_file)) - _stub_prediction_stage(monkeypatch, n_nodes=5) - - out = pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path), - ood_json="does-not-exist.json") - - assert "get_or_mesh" in calls - assert "mesh_step" not in calls - assert calls["get_or_mesh"]["target_nodes"] == pc_mod.DEFAULT_TARGET_NODES - assert out["mesh_mode"] == "prep" - assert out["n_nodes"] == 5 - assert "cache" in out - assert out["cache"] == "miss" - - -def test_mesh_mode_raw_routes_through_mesh_step(monkeypatch, mini_msh_file, tmp_path): - calls = {} - monkeypatch.setattr(pc_mod, "get_or_mesh", _fake_get_or_mesh(calls, mini_msh_file)) - monkeypatch.setattr(pc_mod, "mesh_step", _fake_mesh_step(calls, mini_msh_file)) - _stub_prediction_stage(monkeypatch, n_nodes=5) - - out = pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path), - mesh_mode="raw", ood_json="does-not-exist.json") - - assert "mesh_step" in calls - assert "get_or_mesh" not in calls - assert out["mesh_mode"] == "raw" - - -def test_explicit_elem_size_implies_raw_even_with_default_mesh_mode( - monkeypatch, mini_msh_file, tmp_path): - calls = {} - monkeypatch.setattr(pc_mod, "get_or_mesh", _fake_get_or_mesh(calls, mini_msh_file)) - monkeypatch.setattr(pc_mod, "mesh_step", _fake_mesh_step(calls, mini_msh_file)) - _stub_prediction_stage(monkeypatch, n_nodes=5) - - out = pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path), - elem_size=5.0, ood_json="does-not-exist.json") - - assert "mesh_step" in calls - assert "get_or_mesh" not in calls - assert calls["mesh_step"]["elem_size"] == 5.0 - assert out["mesh_mode"] == "raw" - - -def test_prep_meshing_failure_raises_loudly(monkeypatch, mini_msh_file, tmp_path): - calls = {} - monkeypatch.setattr(pc_mod, "get_or_mesh", - _fake_get_or_mesh(calls, mini_msh_file, ok=False)) - monkeypatch.setattr(pc_mod, "mesh_step", _fake_mesh_step(calls, mini_msh_file)) - _stub_prediction_stage(monkeypatch, n_nodes=5) - - with pytest.raises(RuntimeError): - pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path), - ood_json="does-not-exist.json") - - -def test_predict_casting_includes_cache_hit_in_result( - monkeypatch, mini_msh_file, tmp_path): - """Verify that a cache hit result dict includes the 'cache' field.""" - calls = {} - monkeypatch.setattr(pc_mod, "get_or_mesh", _fake_get_or_mesh(calls, mini_msh_file, cache="hit")) - monkeypatch.setattr(pc_mod, "mesh_step", _fake_mesh_step(calls, mini_msh_file)) - _stub_prediction_stage(monkeypatch, n_nodes=5) - - out = pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path), - ood_json="does-not-exist.json") - - assert "cache" in out - assert out["cache"] == "hit" - - - - -# -------------------------------------------------------------------------- -# target_nodes pass-through (app Phase A, 2026-08-01) -# -# `get_or_mesh`/`heal_and_mesh` have honoured `target_nodes` all along and it is -# part of the mesh-cache key, but `predict_casting` pinned it to the default, so -# no caller could choose a mesh density without dropping to the legacy `raw` -# path via `elem_size`. -# -------------------------------------------------------------------------- - - -def test_target_nodes_is_forwarded_to_get_or_mesh(monkeypatch, mini_msh_file, tmp_path): - calls = {} - monkeypatch.setattr(pc_mod, "get_or_mesh", _fake_get_or_mesh(calls, mini_msh_file)) - monkeypatch.setattr(pc_mod, "mesh_step", _fake_mesh_step(calls, mini_msh_file)) - _stub_prediction_stage(monkeypatch, n_nodes=5) - - pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path), - ood_json="does-not-exist.json", target_nodes=9600) - - assert calls["get_or_mesh"]["target_nodes"] == 9600 - - -def test_target_nodes_with_an_explicit_elem_size_raises(monkeypatch, mini_msh_file, - tmp_path): - """Density search and a pinned element size are mutually exclusive requests. - - `elem_size` forces the `raw` path, which never calls `get_or_mesh`, so a - silently ignored `target_nodes` would mean the caller asked for a mesh - density and got a different one with no indication, so it fails loudly. - """ - monkeypatch.setattr(pc_mod, "get_or_mesh", _fake_get_or_mesh({}, mini_msh_file)) - monkeypatch.setattr(pc_mod, "mesh_step", _fake_mesh_step({}, mini_msh_file)) - _stub_prediction_stage(monkeypatch, n_nodes=5) - - with pytest.raises(ValueError) as excinfo: - pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path), - ood_json="does-not-exist.json", - target_nodes=9600, elem_size=4.0) - - message = str(excinfo.value) - assert "target_nodes" in message and "elem_size" in message diff --git a/tests/test_release_io.py b/tests/test_release_io.py index 8f3e5c820dc8e516561ae31cd3541c86e8cf5057..24bb39ed04d0f355a5f62f9434af1b7121625107 100644 --- a/tests/test_release_io.py +++ b/tests/test_release_io.py @@ -1,82 +1,55 @@ -"""Public-release artifact formats round-trip exactly (src/release_io.py).""" -import os +import collections +import json import numpy as np import pytest import torch +from safetensors.torch import save_file -from src.model_gnn import SolidGNN from src import release_io as rio -def _ckpt(): - torch.manual_seed(0) - cfg = {"hidden": 16, "n_blocks": 2, "node_in": 15, "edge_in": 5, "two_head": True, - "film": True, "g_dim": 5, "graph": "union", "k": 16, "feature_version": 4} - model = SolidGNN(node_in=15, edge_in=5, hidden=16, n_blocks=2, two_head=True, - film=True, g_dim=5) - norm = {"feat_mean": np.arange(15, dtype=np.float32), - "feat_std": np.ones(15, dtype=np.float64), - "g_mean": torch.zeros(5), "y_log_mean": 3.5} - return {"model": model.state_dict(), "norm": norm, "cfg": cfg} +def _write_gnn(path, fmt=rio.GNN_FORMAT): + tensors = {"model.w": torch.arange(6, dtype=torch.float32).reshape(2, 3), + "norm.feat_mean": torch.from_numpy(np.arange(3, dtype=np.float32))} + meta = {"format": fmt, "cfg": json.dumps({"hidden": 16, "k": 4}), + "norm_kinds": json.dumps({"feat_mean": "numpy"}), + "norm_scalars": json.dumps({"y_log_mean": 3.5})} + save_file(tensors, str(path), metadata=meta) -def test_gnn_safetensors_roundtrip(tmp_path): - pytest.importorskip("safetensors") - state = _ckpt() - out = str(tmp_path / "gnn.safetensors") - rio.save_gnn_release(state, out) - back = rio.load_gnn_state(out) - assert back["cfg"] == state["cfg"] - assert set(back["model"]) == set(state["model"]) - for k, v in state["model"].items(): - assert torch.equal(back["model"][k], v) - n = back["norm"] - assert isinstance(n["feat_mean"], np.ndarray) and n["feat_mean"].dtype == np.float32 - assert n["feat_std"].dtype == np.float64 - assert isinstance(n["g_mean"], torch.Tensor) - assert n["y_log_mean"] == 3.5 - np.testing.assert_array_equal(n["feat_mean"], state["norm"]["feat_mean"]) - # load_gnn_any dispatches on the extension - assert rio.load_gnn_any(out, "cpu")["cfg"] == state["cfg"] +def test_load_gnn_state(tmp_path): + p = tmp_path / "gnn.safetensors" + _write_gnn(p) + state = rio.load_gnn_state(str(p)) + assert state["cfg"] == {"hidden": 16, "k": 4} + assert torch.equal(state["model"]["w"], torch.arange(6, dtype=torch.float32).reshape(2, 3)) + assert isinstance(state["norm"]["feat_mean"], np.ndarray) + assert state["norm"]["feat_mean"].dtype == np.float32 + assert state["norm"]["y_log_mean"] == 3.5 -def test_gnn_rejects_foreign_safetensors(tmp_path): - st = pytest.importorskip("safetensors.torch") - out = str(tmp_path / "other.safetensors") - st.save_file({"w": torch.zeros(2)}, out) +def test_load_gnn_state_rejects_other_files(tmp_path): + p = tmp_path / "other.safetensors" + _write_gnn(p, fmt="something-else") with pytest.raises(ValueError, match="not a"): - rio.load_gnn_state(out) + rio.load_gnn_state(str(p)) -def test_gbm_skops_roundtrip(tmp_path): - pytest.importorskip("skops") +def test_load_gbm_roundtrip(tmp_path): + skops = pytest.importorskip("skops.io") from sklearn.ensemble import HistGradientBoostingRegressor rng = np.random.default_rng(0) X, y = rng.normal(size=(200, 4)), rng.normal(size=200) m = HistGradientBoostingRegressor(max_iter=20).fit(X, y) - out = str(tmp_path / "gbm.skops") - rio.save_gbm_release(m, out) - back = rio.load_gbm(out) - np.testing.assert_array_equal(back.predict(X), m.predict(X)) + p = tmp_path / "gbm.skops" + skops.dump(m, str(p)) + np.testing.assert_array_equal(rio.load_gbm(p).predict(X), m.predict(X)) -def test_load_gbm_passthrough(): - assert rio.load_gbm(None) is None - obj = object() - assert rio.load_gbm(obj) is obj - - -def test_resolve_artifact(tmp_path, monkeypatch): - monkeypatch.chdir(tmp_path) - internal = "runs/ft3_gbm/hgb_v3.joblib" - release = rio.RELEASE_PATHS[internal] - assert rio.resolve_artifact(internal) == internal # neither exists - os.makedirs(os.path.dirname(release)) - open(release, "w").close() - assert rio.resolve_artifact(internal) == release # only release exists - os.makedirs(os.path.dirname(internal)) - open(internal, "w").close() - assert rio.resolve_artifact(internal) == internal # internal wins - assert rio.resolve_artifact("some/other.joblib") == "some/other.joblib" - assert rio.resolve_artifact(None) is None +def test_load_gbm_refuses_foreign_types(tmp_path): + skops = pytest.importorskip("skops.io") + p = tmp_path / "bad.skops" + skops.dump({"c": collections.Counter(a=1)}, str(p)) + with pytest.raises(ValueError, match="unexpected types"): + rio.load_gbm(p) diff --git a/tests/test_shrinkage.py b/tests/test_shrinkage.py index 6adb90440a0c19d47da2f1d1d3caf4fbd87941c9..6126e7f094e862d8a64faf5bd07bb715460a2c76 100644 --- a/tests/test_shrinkage.py +++ b/tests/test_shrinkage.py @@ -2,14 +2,14 @@ import numpy as np import pytest from src.shrinkage import ( - node_gradient, niyama_exact, niyama_proxy, porosity_risk, NY_CRIT, + node_gradient, niyama_proxy, porosity_risk, hot_regions, is_fed, macro_shrinkage, top_sinks, feeder_recommendations, SHRINKAGE_PCT, screen_casting, _boundary_faces, ) # --------------------------------------------------------------------------- -# Structured tet-grid helpers (Task 2+) -- a Kuhn/Freudenthal 6-tet-per-cube +# Structured tet-grid helpers -- a Kuhn/Freudenthal 6-tet-per-cube # decomposition of a regular nx*ny*nz node lattice (unit spacing unless # overridden). This is a conforming triangulation (verified numerically: tet # volumes sum exactly to the lattice bounding volume, and internal faces @@ -103,28 +103,10 @@ def test_node_gradient_no_tets_returns_zeros(): # (b) niyama_exact # --------------------------------------------------------------------------- -def test_niyama_exact_unit_conversion(): - # gradT = 1000 K/m = 1 degC/mm ; dTdt = -4 K/s -> |Tdot| = 4 - # Ny = G / sqrt(|Tdot|) = 1 / sqrt(4) = 0.5 - gradT = np.array([1000.0]) - dTdt = np.array([-4.0]) - ny = niyama_exact(gradT, dTdt) - np.testing.assert_allclose(ny, [0.5]) -def test_niyama_exact_eps_guarded_at_zero_dTdt(): - gradT = np.array([1000.0]) - dTdt = np.array([0.0]) - ny = niyama_exact(gradT, dTdt) - assert np.isfinite(ny).all() -def test_niyama_exact_vectorized(): - gradT = np.array([1000.0, 2000.0]) - dTdt = np.array([-4.0, -16.0]) - ny = niyama_exact(gradT, dTdt) - # G = [1, 2] degC/mm ; sqrt(|Tdot|) = [2, 4] -> Ny = [0.5, 0.5] - np.testing.assert_allclose(ny, [0.5, 0.5]) # --------------------------------------------------------------------------- @@ -160,38 +142,14 @@ def test_niyama_proxy_eps_guarded_at_zero_t_s(): # (d) porosity_risk -- absolute mode # --------------------------------------------------------------------------- -def test_porosity_risk_absolute_straddling_crit(): - # grey_iron / ductile_iron NY_CRIT = 0.75 - ny = np.array([0.5, 0.75, 1.0]) - mask = porosity_risk(ny, mode="absolute", alloy_family="grey_iron") - np.testing.assert_array_equal(mask, [True, False, False]) -def test_porosity_risk_absolute_steel_crit(): - ny = np.array([0.9, 1.0, 1.1]) - mask = porosity_risk(ny, mode="absolute", alloy_family="steel") - np.testing.assert_array_equal(mask, [True, False, False]) -def test_porosity_risk_absolute_unknown_family_raises_keyerror(): - ny = np.array([0.5, 1.0]) - with pytest.raises(KeyError): - porosity_risk(ny, mode="absolute", alloy_family="unobtainium") -def test_porosity_risk_absolute_missing_family_raises_keyerror(): - ny = np.array([0.5, 1.0]) - with pytest.raises(KeyError): - porosity_risk(ny, mode="absolute") # alloy_family=None -> loud KeyError, no silent default -def test_ny_crit_table(): - assert NY_CRIT == { - "steel": 1.0, - "low_alloy_steel": 1.0, - "grey_iron": 0.75, - "ductile_iron": 0.75, - } # --------------------------------------------------------------------------- @@ -200,14 +158,14 @@ def test_ny_crit_table(): def test_porosity_risk_percentile_exact_bottom_10pct(): ny = np.arange(20, dtype=float) # 0..19, no ties - mask = porosity_risk(ny, mode="percentile", pct=10.0) + mask = porosity_risk(ny, pct=10.0) assert mask.sum() == 2 # exactly bottom 10% of 20 = 2 np.testing.assert_array_equal(np.where(mask)[0], [0, 1]) def test_porosity_risk_percentile_exact_bottom_10pct_n10(): ny = np.arange(10, dtype=float) # 0..9 - mask = porosity_risk(ny, mode="percentile", pct=10.0) + mask = porosity_risk(ny, pct=10.0) assert mask.sum() == 1 np.testing.assert_array_equal(np.where(mask)[0], [0]) @@ -219,15 +177,11 @@ def test_porosity_risk_percentile_tie_handling_documented(): # nominal pct% when duplicates sit at the boundary. Documented behavior, # not a bug: deterministic and order-independent. ny = np.array([0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]) - mask = porosity_risk(ny, mode="percentile", pct=10.0) + mask = porosity_risk(ny, pct=10.0) assert mask.sum() == 2 np.testing.assert_array_equal(np.where(mask)[0], [0, 1]) -def test_porosity_risk_unknown_mode_raises(): - ny = np.array([0.5, 1.0]) - with pytest.raises(ValueError): - porosity_risk(ny, mode="bogus") # --------------------------------------------------------------------------- @@ -311,13 +265,9 @@ def test_is_fed_empty_feed_nodes_is_always_false(): # --------------------------------------------------------------------------- # -# Review finding C1 (final whole-branch review, 2026-08-01): marginal -# ("sub-resolution", <4 nodes) hot regions mix a lumped-nodal V (stays -# O(h^3)) with a true iso-surface A (shrinks toward zero as t_max -> -# hot_threshold) -> M = V/A diverges without bound. Reproduced on the -# branch's own committed BUKSA smoke artifact (a 1-node region's M=185mm -# printed a min_modulus_mm=222 feeder recommendation on a 382mm part; 6/16 -# BUKSA regions were 1-3 nodes). Fix: `hot_regions` marks such regions +# Marginal ("sub-resolution", <4 nodes) hot regions mix a lumped-nodal V +# (stays O(h^3)) with a true iso-surface A (shrinks toward zero as t_max -> +# hot_threshold), so M = V/A diverges without bound. `hot_regions` marks such regions # "sub_resolution": True; `feeder_recommendations` skips them (logging the # count once); `macro_shrinkage` still lists them (volume/severity remain # meaningful, only the modulus diverges). @@ -396,7 +346,7 @@ def test_feeder_recommendations_mixed_regions_skips_only_sub_resolution_ones(cap # --------------------------------------------------------------------------- # -# Review finding I4 (performance): the mesh's boundary-face structure is +# the mesh's boundary-face structure is # region-independent and hoisted out of `_region_surface_area`'s per-call # work into a standalone `_boundary_faces(tets)` helper, computed once by # `feeder_recommendations` and shared across every region. A direct check @@ -428,10 +378,9 @@ def test_hot_regions_top_frac_out_of_range_raises(): # --------------------------------------------------------------------------- -# (a2) is_fed -- adjudicated corridor semantics (2026-08-01 review round 1, -# finding 1): a warm corridor (t within `corridor_frac` of the region's own +# (a2) is_fed corridor semantics: a warm corridor (t within `corridor_frac` of the region's own # t_max) bridges the region to a feed node several hops away, NOT required -# to be part of the region's own hot-node set. This is the reviewer's +# to be part of the region's own hot-node set. This is the # REQUIRED test: it must fail both under a `return False` stub (variant 1 # expects True) and under a naive "touch test" that only checks direct # region-to-feed_node adjacency or membership (variant 1's feed node is 6 @@ -482,7 +431,7 @@ def test_is_fed_false_when_feed_node_is_on_far_cold_side(): def test_is_fed_region_t_max_zero_is_unfed_not_vacuously_true(): - """review round 2, finding NEW-2: a region t_max <= 0 would make the + """a region t_max <= 0 would make the corridor threshold (corridor_frac * t_max) <= 0, vacuously satisfied by every non-negative t_s -- i.e. the whole (connected) mesh would count as "warm", trivially fed via any non-empty feed_nodes regardless of real @@ -502,18 +451,16 @@ def test_is_fed_region_t_max_zero_is_unfed_not_vacuously_true(): def _interior_cube_block_mesh(): """9x6x5 lattice with a fully-interior 3x3x3 node block (a 2x2x2 mm solid - cube of 8 sub-cells). Under the fraction-weighted volume rule (review - round 1 finding 2), volume_mm3 = 27.0 exactly: each of the 27 hot nodes + cube of 8 sub-cells). Under the fraction-weighted volume rule, volume_mm3 = 27.0 exactly: each of the 27 hot nodes is fully interior with a complete 8-cube neighborhood, so its lumped nodal volume (sum of tet_volume/4 over incident tets) is exactly 1 mm^3 on this unit-spacing lattice -- hand-verified numerically (independently confirmed: fraction-weighted volume == hot-node COUNT here, since summing volume*n_hot/4 over all tets is algebraically identical to - summing each hot node's own lumped volume). Boundary area (review round 2 - finding NEW-1) is now `_region_surface_area`'s co-area/domain-boundary + summing each hot node's own lumped volume). Boundary area is now `_region_surface_area`'s co-area/domain-boundary formula -- see that function's docstring -- which for this fixture comes to ~49.31 mm^2 (M = V/A ~= 0.5475, within 15% of the idealized "3x3x3 mm - dilated cube" reference of A=54/M=0.5 the reviewer hand-derived; an + dilated cube" reference of A=54/M=0.5 derived by hand; an exact closed-form match isn't expected for a graded/partial-tet field -- see `test_feeder_recommendations_modulus_exact` and the mesh-refinement drift test below for the actual acceptance checks). The block is exactly @@ -570,8 +517,8 @@ def _unfed_boss_and_fed_band_10pct_mesh(): boundary (unfed) -- hand-verified via the boss/band-touches-surface assertions below. Band's t_max is deliberately HIGHER than boss's so hot_regions sorts band first (region_id=0) and boss second (region_id=1) - -- this exposes a positional-reindex bug (region_id coherence, review - round 1 finding 3) that a same-valued t_max tie would hide.""" + -- this exposes a positional-reindex bug (region_id coherence) that a + same-valued t_max tie would hide.""" nx, ny, nz = 10, 5, 4 pos, tets = _grid_mesh(nx, ny, nz) idx = _node_idx(nx, ny) @@ -601,7 +548,7 @@ def test_macro_shrinkage_excludes_fed_regions(): def test_region_id_coherence_across_pipeline(): """region_id must be the STABLE hot_regions index, propagated (not positionally re-derived) through macro_shrinkage's filtering and further - through top_sinks/feeder_recommendations (review round 1 finding 3).""" + through top_sinks/feeder_recommendations.""" pos, tets, t_s, surface_nodes, boss_nodes, band_nodes = _unfed_boss_and_fed_band_10pct_mesh() raw_regions = hot_regions(pos, tets, t_s) @@ -625,8 +572,7 @@ def test_region_id_coherence_across_pipeline(): def test_macro_shrinkage_verdict_flips_via_genuine_external_corridor(): - """review round 2, finding NEW-3 (round 3 cheap follow-up: made the - corridor genuinely multi-hop): an end-to-end (not just is_fed-level) + """an end-to-end (not just is_fed-level) demonstration that a region's fed/unfed verdict can flip because of a node that is NOT part of the region (hot_regions' own top-10% cut excludes it) and NOT the feed node itself, but sits between @@ -686,7 +632,7 @@ def test_shrinkage_pct_table(): # one hot region spanning a top-face node, an interior bridge node, an # x-face node, and a y-face node (all connected via tet adjacency through # the bridge) -- also reused by feeder_recommendations' rotation tests below -# (review round 1 finding 5). +#. # --------------------------------------------------------------------------- def _rotatable_sink_mesh(): @@ -729,7 +675,7 @@ def test_top_sinks_rotated_gravity_relocates_to_x_face(): # top surface above the region's centroid; ValueError on empty surface_nodes; # a degenerate (zero-volume/area) region logs a warning instead of a silent 0. # -# review round 3: `_region_surface_area` now runs marching tetrahedra on the +# `_region_surface_area` now runs marching tetrahedra on the # CONTINUOUS `t_s` field cut at the region's own `hot_threshold`, replacing # round 2's co-area formula on a BINARIZED indicator -- the binarized version # always crosses each tet edge at its midpoint regardless of where the real @@ -832,14 +778,14 @@ def _interior_chebyshev_cube_mesh(): def test_feeder_recommendations_modulus_level_set_consistent(): - """review round 3: V and A must measure the same body via a genuine + """V and A must measure the same body via a genuine continuous-field level set (not an artificially-thresholded binary block). On this fixture the true level set is an exact 3x3x3mm cube (M = side/6 = 0.5 by the Chvorinov identity for any cube); marching tetrahedra on a real linear field still chamfers cube corners/edges (hand-verified: the same chamfering appears in an independent marching- tets reimplementation, not just the shipped code), so M lands near but - not exactly at 0.5 -- within the reviewer's stated 15% band.""" + not exactly at 0.5 -- within the stated 15% band.""" pos, tets, t_s, surface_nodes, region = _interior_chebyshev_cube_mesh() recs = feeder_recommendations([region], pos, tets, surface_nodes, gravity_vec=(0, 0, -1), t_s=t_s) @@ -917,7 +863,7 @@ def test_feeder_recommendations_degenerate_region_warns_instead_of_silent_zero(c def test_feeder_recommendations_region_touching_mesh_skin_exercises_boundary_term(caplog): - """review round 3 cheap follow-up: a dedicated, in-suite exercise of the + """a dedicated, in-suite exercise of the domain-boundary term (a region whose own boundary coincides with the mesh's outer surface, so `_region_surface_area`'s "no other side to integrate across" correction path actually fires) -- the gravity- @@ -950,12 +896,7 @@ def _smooth_blob_mesh(spacing, extent=10.0, top_frac=0.05, center_offset=(0.13, density (unlike the axis-aligned block fixtures elsewhere in this file, which are deliberately clean/hand-checkable but not representative of a "smooth field" mesh-convergence scenario). `center_offset` is a - deliberately non-grid-aligned (sub-cell) shift (review round 3: a - center placed exactly on a grid node/plane risks the sphere boundary - landing at unusually "favorable" axis-aligned tet crossings at every - tested density, which would NOT stress-test the anisotropy this - fixture exists to catch -- round 2's version used an exact-grid-point - center).""" + deliberately non-grid-aligned (sub-cell) shift.""" n = int(round(extent / spacing)) + 1 pos, tets = _grid_mesh(n, n, n, spacing=spacing) center = np.array([extent / 2.0] * 3) + np.array(center_offset) @@ -965,8 +906,7 @@ def _smooth_blob_mesh(spacing, extent=10.0, top_frac=0.05, center_offset=(0.13, def test_feeder_recommendations_modulus_refinement_drift_improves_nonaxis(caplog): - """review round 3: moved from an exact-grid-centered blob (round 2's - version) to an off-grid-centered spherical blob -- a genuinely NON-axis- + """moved from an exact-grid-centered blob to an off-grid-centered spherical blob -- a genuinely NON-axis- aligned fixture, since the round-2 co-area formula's axis-aligned cuts were "the one exact case" and could mask residual anisotropy even on a nominally spherical shape if its center happened to sit on a favorable @@ -974,7 +914,7 @@ def test_feeder_recommendations_modulus_refinement_drift_improves_nonaxis(caplog all-4-vertices-hot area rule, M = V/A drifted ~-29% between a coarse and 2x-refined mesh of a smooth field (worse than the pre-fix-round-1 baseline of -17%). The marching-tetrahedra area (this round) must keep - drift well under the reviewer's 15% bound.""" + drift well under the 15% bound.""" import logging as _logging def modulus_at(spacing): @@ -993,7 +933,7 @@ def test_feeder_recommendations_modulus_refinement_drift_improves_nonaxis(caplog # --------------------------------------------------------------------------- -# (g) review round 3 acceptance suite: marching-tets area vs. the ANALYTIC +# (g) Acceptance suite: marching-tets area vs. the ANALYTIC # true value on shapes the old (binarized co-area) formula got badly wrong -- # a sphere (curved, converging with refinement, not plateauing), a diagonal # plane cut through a box (exact via a convex-hull oracle), and an @@ -1019,9 +959,9 @@ def _sphere_region(spacing, extent=12.0, top_frac=0.03, center_offset=(0.13, 0.2 def test_region_area_sphere_converges_not_plateaus(): - """review round 3 acceptance check #1: sphere-field area must converge + """sphere-field area must converge toward the analytic 4*pi*R^2 as the mesh refines (error DECREASING), - not plateau at the ~24-29% the reviewer measured for the binarized + not plateau at the ~24-29% measured for the binarized co-area formula.""" errors = [] for spacing in (1.0, 0.5, 0.25): @@ -1089,9 +1029,9 @@ def _diagonal_cut_region(spacing, normal, extent=10.0, top_frac=0.3): def test_region_area_diagonal_plane_cut_error_far_below_old_bias(): - """review round 3 acceptance check #2: a half-space region cut by a + """a half-space region cut by a plane along the space diagonal (1,1,1) -- normal not aligned with any - mesh axis. Error must be far below the reviewer's reported +36% for the + mesh axis. Error must be far below the reported +36% for the prior (binarized co-area) formula.""" for spacing in (1.0, 0.5): pos, tets, nodes, t_s, threshold, true_area = _diagonal_cut_region(spacing, (1, 1, 1)) @@ -1100,31 +1040,10 @@ def test_region_area_diagonal_plane_cut_error_far_below_old_bias(): assert err < 0.10, f"spacing={spacing}: diagonal-cut error {err:.1%} too high" -def test_region_area_orientation_sweep_spread_far_below_old_bias(): - """review round 3 acceptance check #3: the SAME half-space-cut - construction at several normals (axis-aligned, face-diagonal, space- - diagonal, and a generic non-symmetric direction) -- the spread of - relative errors across orientations must be far below the reviewer's - reported ~29% orientation sensitivity for the prior formula.""" - orientations = { - "axis-x": (1, 0, 0), - "face-diag": (1, 1, 0), - "space-diag": (1, 1, 1), - "generic": (2, 1, 0.3), - } - errors = {} - for name, normal in orientations.items(): - pos, tets, nodes, t_s, threshold, true_area = _diagonal_cut_region(0.5, normal) - area = _reference_region_area_marching_tets(pos, tets, nodes, t_s, threshold) - errors[name] = abs(area - true_area) / true_area - - spread = max(errors.values()) - min(errors.values()) - assert spread < 0.10, f"orientation-sweep spread {spread:.1%} too high: {errors}" - assert all(e < 0.10 for e in errors.values()), f"an orientation exceeded 10% error: {errors}" # --------------------------------------------------------------------------- -# (f) hot_regions volume -- fraction-weighted rule (review round 1 finding 2): +# (f) hot_regions volume -- fraction-weighted rule: # each tet contributes tet_volume * n_hot_vertices/4, replacing an earlier # all-4-vertices-hot rule that discarded (undercounted) every boundary tet. # --------------------------------------------------------------------------- @@ -1193,46 +1112,12 @@ def test_determinism_hot_regions_and_downstream(): # --------------------------------------------------------------------------- -# Task 3: screen_casting orchestrator -- one call combining Niyama (exact when -# gradT+dTdt given, proxy otherwise), macro-shrinkage, top sinks, feeder -# recommendations, and an optional fixed-axis gravity sweep. -# -# NOTE (brief deviation, owner-adjudicated -- see src/shrinkage.py's -# `screen_casting` docstring): the brief's signature in -# docs/superpowers/plans/2026-08-01-sp2-shrinkage.md does not carry a -# `feed_nodes` parameter, but `macro_shrinkage`/`is_fed` (task 2, review round -# 1) require an explicit ADIABATIC-face feed-node array -- distinct from and -# not derivable from `surface_nodes` (the whole free surface, used only for -# sink/feeder placement). `feed_nodes=None` defaults to an empty array, the -# same "bare sand casting, no risers, everything unfed" convention already -# used by `macro_shrinkage`/`is_fed` directly. +# screen_casting: one call combining the Niyama proxy, macro-shrinkage, top +# sinks and feeder recommendations. `feed_nodes` (adiabatic-face nodes) is +# separate from `surface_nodes` (the whole free surface, used for sink and +# feeder placement); None means no feed source, so every region is unfed. # --------------------------------------------------------------------------- -def test_screen_casting_exact_niyama_when_fields_given(): - nx = ny = nz = 4 - pos, tets = _grid_mesh(nx, ny, nz) - n_nodes = pos.shape[0] - t_s = pos[:, 0].copy() + 1.0 # linear ramp along x, strictly positive - surface_nodes = _surface_nodes(nx, ny, nz) - gradT = np.full(n_nodes, 1000.0) # 1000 K/m = 1 degC/mm - dTdt = np.full(n_nodes, -4.0) # |Tdot|=4 -> Ny = 1/sqrt(4) = 0.5 everywhere - TL, TS = 1400.0, 1100.0 - - result = screen_casting(pos, tets, surface_nodes, t_s, TL, TS, "steel", - gradT=gradT, dTdt=dTdt) - - assert result["niyama_source"] == "exact" - expected_ny = niyama_exact(gradT, dTdt) - np.testing.assert_allclose(result["niyama"], expected_ny) - expected_risk = porosity_risk(expected_ny, mode="absolute", alloy_family="steel") - np.testing.assert_array_equal(result["porosity_risk"], expected_risk) - # Ny=0.5 < NY_CRIT["steel"]=1.0 everywhere -> every node at risk. - assert result["porosity_risk"].all() - assert result["material"] == {"alloy_family": "steel", "shrinkage_pct": 2.5} - assert result["per_orientation"] is None - assert isinstance(result["macro_shrinkage"], list) - assert isinstance(result["top_sinks"], list) - assert isinstance(result["feeders"], list) def test_screen_casting_proxy_niyama_when_fields_omitted(): @@ -1247,38 +1132,11 @@ def test_screen_casting_proxy_niyama_when_fields_omitted(): assert result["niyama_source"] == "proxy" expected_ny = niyama_proxy(pos, tets, t_s, TL, TS) np.testing.assert_allclose(result["niyama"], expected_ny) - expected_risk = porosity_risk(expected_ny, mode="percentile") + expected_risk = porosity_risk(expected_ny) np.testing.assert_array_equal(result["porosity_risk"], expected_risk) assert result["material"] == {"alloy_family": "grey_iron", "shrinkage_pct": 0.0} -def test_screen_casting_orientations_axes_sweep_has_six_entries_matching_single_call(): - pos, tets, t_s, surface_nodes, block_nodes = _interior_cube_block_mesh() - TL, TS = 1400.0, 1100.0 - - swept = screen_casting(pos, tets, surface_nodes, t_s, TL, TS, "steel", - orientations="axes") - assert swept["per_orientation"] is not None - assert set(swept["per_orientation"].keys()) == {"+x", "-x", "+y", "-y", "+z", "-z"} - for axis, entry in swept["per_orientation"].items(): - assert set(entry.keys()) == {"macro_shrinkage", "top_sinks", "feeders"} - assert len(entry["macro_shrinkage"]) == 1 # gravity-independent region detection - - single = screen_casting(pos, tets, surface_nodes, t_s, TL, TS, "steel", - gravity_vec=(0, 0, -1)) - assert single["per_orientation"] is None - - neg_z = swept["per_orientation"]["-z"] - assert len(neg_z["feeders"]) == len(single["feeders"]) == 1 - assert neg_z["feeders"][0]["region_id"] == single["feeders"][0]["region_id"] - np.testing.assert_allclose(neg_z["feeders"][0]["location"], single["feeders"][0]["location"]) - assert neg_z["feeders"][0]["min_modulus_mm"] == pytest.approx(single["feeders"][0]["min_modulus_mm"]) - assert len(neg_z["top_sinks"]) == len(single["top_sinks"]) - - # sanity: the sweep isn't a no-op -- flipping gravity along z relocates the - # feeder to the opposite face. - plus_z = swept["per_orientation"]["+z"] - assert plus_z["feeders"][0]["location"][2] != neg_z["feeders"][0]["location"][2] def test_screen_casting_unknown_alloy_family_raises_loud_error_naming_known_families(): @@ -1303,3 +1161,26 @@ def test_screen_casting_feed_nodes_default_is_empty_bare_casting_convention(): fed_call = screen_casting(pos, tets, surface_nodes, t_s, TL, TS, "steel", feed_nodes=surface_nodes) assert len(fed_call["macro_shrinkage"]) == 1 # band fed away via feed_nodes=surface_nodes + + +def test_region_area_orientation_sweep_spread_far_below_old_bias(): + """The SAME half-space-cut + construction at several normals (axis-aligned, face-diagonal, space- + diagonal, and a generic non-symmetric direction) -- the spread of + relative errors across orientations must be far below the ~29% orientation + sensitivity of a binarized co-area estimate.""" + orientations = { + "axis-x": (1, 0, 0), + "face-diag": (1, 1, 0), + "space-diag": (1, 1, 1), + "generic": (2, 1, 0.3), + } + errors = {} + for name, normal in orientations.items(): + pos, tets, nodes, t_s, threshold, true_area = _diagonal_cut_region(0.5, normal) + area = _reference_region_area_marching_tets(pos, tets, nodes, t_s, threshold) + errors[name] = abs(area - true_area) / true_area + + spread = max(errors.values()) - min(errors.values()) + assert spread < 0.10, f"orientation-sweep spread {spread:.1%} too high: {errors}" + assert all(e < 0.10 for e in errors.values()), f"an orientation exceeded 10% error: {errors}" diff --git a/tools/__init__.py b/tools/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/tools/doe/__init__.py b/tools/doe/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/tools/doe/load_properties.py b/tools/doe/load_properties.py deleted file mode 100644 index f4e4038767ba57cf6e8dd52316646b419cdd4209..0000000000000000000000000000000000000000 --- a/tools/doe/load_properties.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Load the public material, mold, chill and sleeve-calibration tables.""" - -from __future__ import annotations - -import csv -from dataclasses import dataclass -from pathlib import Path - - -@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 - - -_MATERIAL_REQUIRED_FIELDS = [ - "k_W_mK", "rho_kg_m3", "cp_J_kgK", "L_kJ_kg", - "TL_C", "TS_C", "Tpour_lo_C", "Tpour_hi_C", -] - - -def _is_missing(value) -> bool: - value = (value or "").strip() - return value == "" or value.upper() == "OWNER" - - -def load_materials(path: str = "data/properties/materials.csv") -> list[Material]: - """Return complete material rows and reject physically inverted ranges.""" - table = Path(path) - if not table.exists(): - print(f"[load_properties] WARNING: materials file not found: {path}") - return [] - - materials: list[Material] = [] - with table.open(newline="") as handle: - for row in csv.DictReader(handle): - material_id = row.get("material_id") or "" - missing = [key for key in _MATERIAL_REQUIRED_FIELDS - if _is_missing(row.get(key))] - if missing: - print(f"[load_properties] WARNING: skipping {material_id!r}; " - f"missing fields: {missing}") - continue - 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.get("family", "") or "", - 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_effusivity_dict(path: str, key_column: str) -> dict[str, float]: - table = Path(path) - if not table.exists(): - print(f"[load_properties] WARNING: file not found: {path}") - return {} - values: dict[str, float] = {} - with table.open(newline="") as handle: - for row in csv.DictReader(handle): - key = row.get(key_column) - raw = row.get("effusivity_W_s05_m2K") - if key and not _is_missing(raw): - values[key] = float(raw) - return values - - -def load_mold_effusivities(path: str = "data/properties/molds.csv") -> dict[str, float]: - return _load_effusivity_dict(path, "mold_system") - - -def load_chill_effusivities(path: str = "data/properties/chills.csv") -> dict[str, float]: - return _load_effusivity_dict(path, "material") - - -def resolve_exo_q0_by_mef( - mef: float, - csv_path: str = "data/properties/exo_calibration.csv", - tol: float = 0.05, -) -> float: - """Resolve an exothermic heat-flux input from the nearest calibrated MEF.""" - table = Path(csv_path) - if not table.exists(): - raise ValueError(f"exothermic calibration table does not exist: {csv_path}") - with table.open(newline="") as handle: - rows = list(csv.DictReader(handle)) - if not rows: - raise ValueError(f"exothermic calibration table has no rows: {csv_path}") - best = min(rows, key=lambda row: abs(float(row["MEF"]) - mef)) - difference = abs(float(best["MEF"]) - mef) - if difference > tol: - raise ValueError( - f"no calibrated MEF within {tol} of {mef}; " - f"nearest={best['MEF']}, difference={difference:.4f}") - return float(best["q0_W_m2"]) diff --git a/tools/doe/mesh_study.py b/tools/doe/mesh_study.py deleted file mode 100644 index 2dbaf86fc7846d6948f42531dc162439e6d55bc2..0000000000000000000000000000000000000000 --- a/tools/doe/mesh_study.py +++ /dev/null @@ -1,18 +0,0 @@ -"""Small public helper shared by the density-targeted meshing pipeline.""" - -from __future__ import annotations - - -def search_elem_size(build_fn, out_msh: str, target_nodes: int, - max_iter: int = 3) -> tuple[int, float]: - """Refine element size assuming node count scales approximately 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 - - -_search_elem_size = search_elem_size diff --git a/tools/doe/prep_parts.py b/tools/doe/prep_parts.py deleted file mode 100644 index f2b51789d48f9b2517a33c824251bad4be8a3c85..0000000000000000000000000000000000000000 --- a/tools/doe/prep_parts.py +++ /dev/null @@ -1,890 +0,0 @@ -"""Dataset-D part preparation and density-targeted STEP meshing. - -Two failure modes, both owner-decided (DECISIONS-2026-07-22 §11): - - Multi-body (Koleso_otl 4 solids, Track 3, TP 2): these are castings - accidentally split into multiple solids by the source CAD. Fix: keep - ONLY the largest solid by volume (`largest_solid`). - - Mesh-fail "Invalid boundary mesh (overlapping facets)" (20508M, SSG0750, - TP, DO678078): attempt an OCC heal (`gmsh.model.occ.healShapes()`) - before meshing. -DO678078 is additionally OOD (bbox diag 2663.5mm > envelope 2177.7mm, -`data/features_v2/ood_ranges.json`). In practice its heal attempt does NOT -recover a meshable solid (same genuine self-intersecting-B-rep class as -20508M/SSG0750/TP, confirmed 2026-07-23) -- it is reported FAILED for -meshing, same as the other 3 mesh-fail parts. It still gets a cleaned -(largest-solid) `.step` export and is tagged OOD in the report (useful if a -future CAD repair fixes the mesh, and as a record that this geometry sits -outside the training envelope), but it is NOT currently a usable training -geometry -- the "healed + KEPT (OOD)" framing in the original task brief -was aspirational and did not hold up against the real STEP file. 5D49.78.01 -is not broken, just slow/heavy -- it gets a longer per-part time-box (see -`SLOW_PART_TIMEOUTS_S` below) rather than any geometry fix. - -Mesh density: FT1 production density (DECISIONS-2026-07-22 §9, ~7.5-10k -nodes, target_nodes=8500 default). A single fixed `bbox_diag/28` element -size (the original approach) misses this badly across parts of very -different shapes -- observed 1855..18343 nodes in a real batch run at that -fixed size. Instead, `_search_target_nodes` reuses the validated iterative -density search from `tools.doe.mesh_study.search_elem_size` (same -node_count ~ elem_size^-3 rescaling, imported not reimplemented), seeded -with `bbox_diag/28` as the first guess, refining up to 2 more times to -converge near `target_nodes`; a node-count early exit skips the refinement -entirely when the first guess already lands within 35% of the target (the -common case). If the (density-searched) unhealed attempt still fails -outright (not a density problem -- a topology error), heal + one more -density-searched attempt is tried before giving up. - -Every part is time-boxed via a `--single`-mode SUBPROCESS (not an in-process -signal alarm: `gmsh.model.mesh.generate(3)` is a blocking C call that does -not yield to Python signal handlers, so `signal.alarm` cannot reliably -interrupt a hung mesh) -- `nice -n 10 -u -m tools.doe.prep_parts ---single ...` under `subprocess.run(..., timeout=...)`. Default time-box is -300s; 5D49.78.01 (known slow/heavy) gets a longer one. - -CLI: - python -m tools.doe.prep_parts --parts-dir pc_parts \\ - --out-dir doe_out/parts_prepped -writes, per part, a cleaned single-solid `_cleaned.step` (best-effort: -largest-solid + healed, written regardless of whether meshing itself -succeeds) and, when meshing succeeds, `.msh` (gmsh 2.2, all-surface -"Cooling" + "Volume"); plus a summary `prep_report.md`. -""" -from __future__ import annotations - -import argparse -import json -import os -import subprocess -import sys -import time -from pathlib import Path - -import numpy as np - -from tools.doe.mesh_study import search_elem_size - -REPO_ROOT = Path(__file__).resolve().parents[2] - -DEFAULT_TARGET_NODES = 8500 - -# Bump on ANY change that affects mesh output (heal/mesh logic, gmsh options, -# density search, etc.) -- the mesh cache (src/mesh_cache.py) keys on this. -# `fast` (HXT/multithreading/curvature-off, see `_apply_fast_options`) is -# OFF by default and thus does not change today's output, but it DOES -# participate in mesh output whenever a caller passes fast=True -- if -# `fast`'s DEFAULT is ever flipped to True, this version must bump (a -# cache keyed only on this version, not on the `fast` argument itself, -# would otherwise silently reuse a pre-flip mesh for a post-flip call). -# v1 -> v2 (Task 4, 2026-07-29): surface-area density-search seed -# (`_first_guess_elem_size`/`_elem_size_from_area`/`K_SURF`) replaced the -# bbox_diag/28 first-attempt guess UNCONDITIONALLY (default-on, not -# flag-gated like `fast`) -- this changes the first elem-size tried (and -# thus can change the converged mesh) for every part, so a cache entry -# written under v1 must not be served after this change. -# `defeature` (Task 5, 2026-07-29): a best-effort geometry-simplification -# pass (`_apply_defeature` -- see its docstring for the full virtual- -# topology investigation), gated OFF by default and flag-gated like `fast` -# -- does NOT bump this version, for the same reason `fast` doesn't: it -# only participates in mesh output when a caller explicitly passes -# `defeature=True`. `src/mesh_cache.py::get_or_mesh`'s auto-trigger policy -# decides `defeature` deterministically from the input alone (file size, -# then a same-input retry-on-failure), so a given STEP file always -# resolves to the same `defeature` choice across runs -- the cache key -# (which does not include `defeature` itself) stays consistent per input. -# v2 -> v3 (Part B Task 4, 2026-07-29): v3 = multi-body bodies path (Part B) -# -- `heal_and_mesh`'s `keep_solid_id` (user-designated casting solid, -# replacing the largest-solid default) and `src/mesh_cache.py::get_or_mesh`'s -# `bodies_key`/`post_mesh` (tagged-mesh caching) both participate in mesh -# output whenever `predict_casting(bodies=...)` is used -- bump so a -# pre-Part-B cache entry (meshed without any of this) is never served to a -# bodies-aware call. -MESH_PIPELINE_VERSION = 3 -DEFAULT_TIMEOUT_S = 300.0 -# 5D49.78.01 was killed at 180s in the inventory scan's coarser (diag/25) -# attempt; at FT1's finer diag/28 density it needs more room, not a geometry -# fix -- give it a longer per-part time-box instead of the 300s default. -SLOW_PART_TIMEOUTS_S = {"5D49.78.01": 1800.0} - -DEFAULT_OOD_RANGES_PATH = "data/features_v2/ood_ranges.json" - -# Fixed physical-group tags for the volume/cooling groups (see `_tag_and_mesh` -# docstring for why these must be pinned rather than auto-assigned). -_VOLUME_PHYS_TAG = 1 -_COOLING_PHYS_TAG = 2 - - -# -------------------------------------------------------------------------- -# Largest-solid extraction -# -------------------------------------------------------------------------- -def _reduce_to_largest_solid(): - """Assumes an OCC model is already imported + synchronized in the - current gmsh session. If >1 solid (dim=3 OCC entity) exists, keeps only - the max-volume one (removes the rest, recursively, from the OCC CAD - representation) and re-synchronizes. If exactly 1 (or 0) solids exist, - it's a no-op. 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 largest_solid(step_path: str, out_step: str) -> tuple[int, float]: - """Import `step_path` (STEP); if it contains >1 solid, keep only the - max-volume one; export a single-solid STEP to `out_step` either way - (a 1-solid input is copied through via a plain gmsh re-export). - - Returns (n_solids_in, kept_volume_mm3).""" - import gmsh - - gmsh.initialize() - try: - gmsh.option.setNumber("General.Terminal", 0) - gmsh.model.add("largest_solid") - gmsh.model.occ.importShapes(step_path) - gmsh.model.occ.synchronize() - n_in, _kept_tag, kept_vol = _reduce_to_largest_solid() - gmsh.write(out_step) - return n_in, kept_vol - finally: - gmsh.finalize() - - -# -------------------------------------------------------------------------- -# Heal + mesh -# -------------------------------------------------------------------------- -def _elem_size(bbox_diag_mm: float) -> float: - """Starting guess for the density search: bbox diagonal / 28 (roughly - the FT1 ~7.5-10k node density on a "typical" casting, DECISIONS-2026-07-22 - §9) -- `_search_target_nodes` then refines from here towards the actual - `target_nodes`, since this fixed heuristic alone was found to miss badly - across parts of very different shapes (see module docstring). Used as - the FALLBACK first guess (see `_first_guess_elem_size`) when the OCC - surface-area query is unavailable.""" - return bbox_diag_mm / 28.0 - - -# Calibrated 2026-07-29 (mesh-accel Task 4) against BUKSA.STEP + H17.STEP -# (real dataset-D parts, target_nodes=8500) + 2 synthetic box fixtures -# (30x20x10mm and 40x40x40mm, target_nodes=500) over K_SURF in -# {1.0, 1.5, 2.0, 2.5, 3.0} -- see -# The calibration table is retained in the private experiment record. -# K_SURF=2.5 was the outright winner: 4/4 fixtures land within -# +-35% of target_nodes on the FIRST attempt (K=2.0 and K=3.0 each get 3/4; -# K<=1.5 get <=1/4). -K_SURF = 2.5 - - -def _elem_size_from_area(area_mm2: float, target_nodes: int) -> float: - """Starting guess for the density search, seeded from the model's total - OCC surface area (mm^2) rather than the fixed bbox_diag/28 heuristic: - `sqrt(K_SURF * area_mm2 / target_nodes)`. Surface area tracks the - number of surface elements (and, via the tet-to-surface-element ratio - for a roughly-uniform mesh, the total node count) far more directly - than a bounding-box diagonal, which is blind to shape (a thin plate and - a compact block with the same diagonal need very different element - sizes for the same target_nodes). `K_SURF` is calibrated in this task - (see module-level comment above `K_SURF`).""" - return float(np.sqrt(K_SURF * area_mm2 / target_nodes)) - - -def _total_surface_area_mm2() -> float: - """Assumes an already-imported + synchronized OCC model in the current - gmsh session. Returns the total area (mm^2) of all dim=2 (surface) - entities via `gmsh.model.occ.getMass(2, tag)` -- confirmed to work - directly on an imported STEP model (no re-synchronize needed beyond the - caller's own `occ.synchronize()`). Raises whatever gmsh raises on - failure; callers must catch (see `_first_guess_elem_size`).""" - 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: - """First-attempt element-size guess for the density search: prefers the - OCC-surface-area-seeded guess (`_elem_size_from_area`, using the - CURRENTLY-imported/synchronized gmsh model), since the fixed - bbox_diag/28 heuristic alone was found to miss badly on non-typical - (thin-walled, elongated, etc.) shapes (see module docstring). Falls - back to `_elem_size(bbox_diag_mm)` -- loudly logged -- if the area query - raises or returns a non-positive area.""" - try: - area_mm2 = _total_surface_area_mm2() - except Exception as e: - print(f"[prep_parts] 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"[prep_parts] 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_target_nodes(build_fn, out_msh: str, target_nodes: int, - tol: float = 0.35, max_iter: int = 3) -> tuple[int, float, int]: - """Node-count-targeted density search: a thin wrapper around the - validated `tools.doe.mesh_study.search_elem_size` iterative refinement - (same node_count ~ elem_size^-3 rescaling law, imported and reused - as-is, not reimplemented) that adds a node-count early exit. - - `build_fn(out_msh, elem_size) -> (n_nodes, elem_size_used)` -- same - contract as `search_elem_size` expects (`elem_size=None` means "pick a - default"); `n_nodes` is 0 on outright mesh failure (not a density - problem). - - If the FIRST attempt (the surface-area-seeded guess, see - `_first_guess_elem_size`) is already within `tol` (relative) of - `target_nodes`, or fails outright, return immediately -- no point in - further (wasted) refinement iterations in either case. Otherwise, - delegate to `search_elem_size` for up to `max_iter` total attempts - (this redoes the first attempt once more as its own internal first - step -- one redundant mesh generation in this less-common branch, - accepted for simplicity over threading extra state through the reused - function). - - Returns `(n_nodes, elem_size_used, attempts)` -- `attempts` is the - total number of `build_fn` invocations across the whole search (1 when - the first guess already lands in-band; useful both for this module's - own tests and for a future bench harness).""" - 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]: - """Assumes an already-imported (and possibly reduced/healed) OCC model - is synchronized in the current gmsh session. (Re)tags all volume/surface - entities "Volume"/"Cooling" (fresh physical groups -- entity tags can - change after healing) and attempts `gmsh.model.mesh.generate(3)` at - element size `size`. Returns (ok, n_nodes, note). - - Pins the physical-group tag numbers to `_VOLUME_PHYS_TAG`/ - `_COOLING_PHYS_TAG` instead of letting gmsh auto-assign them. Root cause - (found 2026-07-28 debugging H17.STEP producing an untagged mesh): this - function can run MULTIPLE times within one gmsh session -- the - density-search retries in `_search_target_nodes`/`search_elem_size` call - it again with a different `size` when the first attempt misses the node - target. `removePhysicalGroups()` frees the *tag number* for reuse but - does NOT clear the *name registration* for that string -- so on the 2nd+ - call in a session, `gmsh.model.addPhysicalGroup(dim, ents)` auto-assigns - a NEW tag, but `setPhysicalName(dim, new_tag, "Volume"/"Cooling")` - silently no-ops (gmsh will not bind a name string that's already - registered to a different, even now-removed, tag). The result: the - LAST-written `.msh` has its actual `$Elements` referencing an unnamed - physical group, while `$PhysicalNames` still lists the stale first-round - (tag -> name) binding -- `parse_msh` matches elements by name and finds - none, i.e. `tris == {}`. Confirmed with a minimal gmsh repro (no meshing - involved) that reusing the same tag number every round keeps the name - bound correctly. BUKSA never hit this because its first density guess - already converged (one call, no retries); H17's did not (multiple - calls), which is why the two parts diverged. - """ - 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") - # Fail loudly rather than silently writing an untagged mesh: this is a - # defense-in-depth check, not a substitute for the tag-pinning above -- - # it catches ANY other way the name binding could fail to register, - # since a downstream consumer (`parse_msh`) trusts these two names - # unconditionally and would otherwise silently see empty groups. - 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]: - """Assumes a fresh gmsh session with an imported + largest-solid-reduced - OCC model synchronized. Runs `gmsh.model.occ.healShapes()` and checks - whether it actually yielded >=1 volume entity -- a real bug found during - development (see module docstring / task report) is that `healShapes()` - can run WITHOUT raising and still silently reduce the model to 0 volumes - (all surfaces, no solid), which must be treated as an ineffective heal, - not a successful one. Returns (effective, note); `note` is "" when - effective.""" - 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 genuine self-intersecting B-rep, not a " - "density artifact)") - return True, "" - - -def _apply_defeature() -> tuple[bool, str]: - """Assumes a fresh gmsh session with an imported (and possibly largest- - solid-reduced) OCC model synchronized. Best-effort geometry- - simplification pass ahead of meshing, gated behind `defeature=True` - (see `heal_and_mesh`). Investigated for this task (mesh-accel Task 5, - 2026-07-29; full probe log in - the meshing-acceleration experiment record) - against real dataset-D parts (4474_443.STEP: 3009 faces / 1 solid / - diag 1283mm; TP.STEP: 981 faces after largest-solid reduction): - - - `gmsh.model.occ.removeAllDuplicates()` -- APPLIED. Fast (<2s even - on the 3009-face part) and, in every part probed, never reduced - the volume count to 0. Its actual effect is narrow (it only merges - LITERALLY duplicate/coincident OCC entities -- e.g. a CAD-export - artifact) -- it was a no-op on face count on every real part - probed here (no exact duplicates present), but it is gmsh's own - recommended call whenever "the OCC model may contain redundant - geometry" and costs nothing to run unconditionally. - - `gmsh.model.occ.healShapes(tolerance=, ...)` as a - "merge faces smaller than the mesh scale" lever -- TRIED AND - REJECTED. Even the library's DEFAULT tolerance (1e-8, i.e. no - intentional defeaturing at all) silently reduced 4474_443.STEP - (a clean, single-solid part) to 0 volume entities after a 59s - call; scaling the tolerance up to merge small faces (tried 5mm on - TP.STEP, whose bbox-diag/28 scale is ~30mm) instead raised - "Could not fix wire in surface 54" outright. This is the exact - class of silent-destruction failure `_apply_heal` already guards - against for the pass-2 heal -- deliberately enlarging the - tolerance to defeature makes it MORE likely to trigger, not less. - Not a safe unconditional lever; not applied. - - `gmsh.option.setNumber("Geometry.Tolerance", ...)` ahead of - `removeAllDuplicates()` -- TRIED AND REJECTED. Swept up to 10mm - (already far above a typical small-part mesh element size) on - 4474_443.STEP with zero change in face count -- in this gmsh - version this option does not make `removeAllDuplicates` merge - near-but-not-exactly-coincident faces. - - `gmsh.model.mesh.setCompound(dim, tags)` -- NOT APPLIED. Per the - gmsh API this is a manual meshing CONSTRAINT (mesh a caller- - supplied group of face tags as one reparametrized patch), not an - automatic small-face detector. Using it for defeaturing would - need a from-scratch face-area + adjacency analysis (per-face area - via `occ.getMass(2, tag)`, threshold against the target element - scale, then group adjacent small faces via shared-edge topology - queries) -- real engineering beyond a best-effort pass; flagged as - a follow-up, not implemented. - - `gmsh.model.mesh.classifySurfaces(...)` -- NOT APPLICABLE. Per its - own docs it operates on an ALREADY-DISCRETIZED surface mesh and - (via `exportDiscrete=True`, the only way to get usable output) - REPLACES the OCC CAD entities with discrete ones -- this is gmsh's - STL/reverse-engineering workflow, not a pre-mesh OCC B-rep - operation, and would break `_tag_and_mesh`'s OCC-entity-based - physical-group tagging downstream. Confirmed inapplicable from the - API contract alone (angle-based reclassification assumes a - discrete/triangulated model); not probed further, per the task - brief's own suspicion. - - Net effect: the applied pass is `removeAllDuplicates()` only, with a - defensive "did we just lose the solid" check (never observed to - trigger during probing, but cheap insurance -- mirrors `_apply_heal`'s - `(ok, note)` return shape and MUST be branched on by the caller the - same way: a `removeAllDuplicates()` call that raises, or that zeroes - out the volume count, leaves the gmsh session with 0 solids -- letting - the caller fall through to `gmsh.model.getBoundingBox(-1, -1)` on an - empty model raises an uncaught `Exception: Empty bounding box` instead - of a clean `ok=False` (found in review, 2026-07-29; see task-5-report.md - "Review fix"). Returns `(ok, note)`; `note` is "" when `ok` (dedup ran - and the solid survived), and a human-readable diagnostic when not.""" - 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) -- not observed during development, but " - "reported honestly rather than silently proceeding") - return True, "" - - -def _reduce_to_solid_id(keep_solid_id: int) -> tuple[int, float]: - """Assumes an OCC model is already imported + synchronized in the current - gmsh session (mirrors `_reduce_to_largest_solid`, but keeps a CALLER- - designated OCC tag instead of the max-volume one -- used by Part B's - multi-body pipeline, where the user has already assigned a specific - solid the "casting" role and it need not be the largest body in the - assembly). Deletes every other dim=3 entity (recursively, from the OCC - CAD representation) and re-synchronizes; a no-op if `keep_solid_id` is - already the only solid. Raises `ValueError` if `keep_solid_id` is not - among the imported model's 3-D entities -- callers (here, - `_import_and_reduce`) turn that into a loud error string, never a - silent fallback to some other solid. Returns (n_solids_in, 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") - if keep_solid_id not in vols: - raise ValueError( - f"keep_solid_id={keep_solid_id} not found among 3D solids {sorted(vols)}") - kept_vol = float(gmsh.model.occ.getMass(3, keep_solid_id)) - if n_in > 1: - drop = [(3, t) for t in vols if t != keep_solid_id] - gmsh.model.occ.remove(drop, recursive=True) - gmsh.model.occ.synchronize() - return n_in, kept_vol - - -def _import_and_reduce(step_path: str, keep_solid_id: int | None = None) -> tuple[int, float, str | None]: - """Assumes a fresh gmsh session (already initialized). Imports - `step_path` and reduces multi-body parts to a single solid: the largest - by volume (`keep_solid_id=None`, existing/default behavior, byte- - identical) or the caller-designated OCC tag (`keep_solid_id=` -- - Part B's multi-body pipeline, where the "casting" solid is whichever one - the user assigned that role to, not necessarily the largest). A missing - `keep_solid_id` is reported the same way as any other reduction failure - -- a loud error string, not an exception -- so `heal_and_mesh` turns it - into `ok=False` with the reason in `notes`, never a silent fallback. - 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}" - - if keep_solid_id is None: - 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 - - try: - n_solids, volume_mm3 = _reduce_to_solid_id(keep_solid_id) - except Exception as e: - return 0, 0.0, f"keep_solid_id reduction failed: {e}" - return n_solids, volume_mm3, None - - -def _apply_fast_options() -> None: - """Assumes an already-initialized gmsh session. Sets the flag-gated - "fast" meshing option set (mesh-accel Task 3, default OFF -- see - `heal_and_mesh`'s `fast` kwarg): `Mesh.Algorithm3D=10` (HXT, gmsh's - multithreaded Delaunay-refinement tet algorithm, vs. the single- - threaded classical Delaunay default), `General.NumThreads` capped at - min(8, os.cpu_count()) (diminishing returns/oversubscription risk past - 8 on this box -- see the CPU-oversubscription rule), and - `Mesh.MeshSizeFromCurvature=0` (curvature-adaptive refinement is - unneeded for this project's uniform analytic-density workflow and only - costs solve time). Does NOT touch `Mesh.MeshSizeMin`/`Mesh.MeshSizeMax` - -- those are already set per-call by `_tag_and_mesh`, right before - `gmsh.model.mesh.generate(3)`.""" - import gmsh - - gmsh.option.setNumber("Mesh.Algorithm3D", 10) - gmsh.option.setNumber("General.NumThreads", min(8, os.cpu_count() or 1)) - gmsh.option.setNumber("Mesh.MeshSizeFromCurvature", 0) - - -def heal_and_mesh(step_path: str, out_msh: str, target_nodes: int = DEFAULT_TARGET_NODES, - out_step: str | None = None, fast: bool = False, - defeature: bool = False, keep_solid_id: int | None = None) -> dict: - """Import `step_path`, reduce multi-body parts to their largest solid, - tag all surfaces "Cooling" / the volume "Volume", and mesh at the FT1 - density via `_search_target_nodes` (node-count-targeted, seeded from - `_first_guess_elem_size`'s OCC-surface-area guess, falling back to - `_elem_size`'s bbox-diag/28 guess if the area query is unavailable -- - see module docstring). Writes gmsh 2.2 `out_msh` on success. - - Two passes, each a FRESH gmsh session (an OCC heal is not reliably - undoable within one session): - 1. Unhealed, density-searched towards `target_nodes`. This is tried - FIRST and is the only pass run for parts that already mesh fine -- - a real finding during development was that - `gmsh.model.occ.healShapes()` can silently reduce a perfectly - meshable solid to 0 volumes (all surfaces, no solid) even though it - raises no exception, so it is NOT applied unconditionally to every - part (that would regress the 9 already-clean parts). `out_step`, if - given, is exported here -- the single-solid (but unhealed) - geometry -- so a usable cleaned STEP exists regardless of what - pass 2 does. - 2. Only if pass 1's mesh attempt fails outright (a topology error, not - a density problem -- `_search_target_nodes` already handles density): - heal (`_apply_heal`, default tolerance) and, only if the heal - actually yielded a volume, retry the SAME density search. If - healing runs but yields 0 volume entities (observed on - 20508M/SSG0750/TP/DO678078 in this repo's pc_parts -- a genuine - self-intersecting B-rep, not a mesh-density artifact: confirmed by - testing heal tolerances from 1e-6 to 0.7, none recover a solid, and - tolerance=1.0 raises "Could not fix wire in surface ..." -- gmsh's - healer simply cannot repair it), that is reported honestly as a - failed heal, not silently treated as success -- meshing is not - even attempted in that case. - - Returns a dict: {n_solids, volume_mm3, bbox_diag_mm, n_nodes, ok, - healed, notes, attempts}. `ok` is True iff meshing produced a non-empty - tetrahedral mesh (in either pass). `healed` is True iff pass 2 ran AND - its heal yielded at least one volume entity (i.e. the heal was - actually effective, not just exception-free). `notes` is a single - human-readable string (diagnostics joined by "; ") for the prep report. - `attempts` is the total number of density-search `build_fn` calls - across whichever pass(es) actually ran a search (1 = the first, - surface-area-seeded guess already landed within tolerance; 0 if no - density search was ever attempted, e.g. import failed). - - `fast`, if True, applies `_apply_fast_options()` (HXT + multithreading - + curvature-refinement off) in BOTH passes' gmsh sessions, right after - each `gmsh.initialize()`. Default False -> zero behavior change (no - option calls at all, gmsh's own defaults apply, current single- - threaded classical-Delaunay behavior is unchanged) -- gated OFF until a - later benchmark task flips the default (see `MESH_PIPELINE_VERSION`). - - `defeature`, if True, applies `_apply_defeature()` (best-effort - geometry simplification -- see its docstring for what was investigated - and what actually ships) in BOTH passes, right after that pass's - import/largest-solid-reduction step. Default False -> zero behavior - change. This is opt-in and does not bump `MESH_PIPELINE_VERSION` (see - the comment above it); `src/mesh_cache.py::get_or_mesh` decides when to - pass `defeature=True` via its own size/failure auto-trigger policy, not - this function. The Volume/Cooling pinned-tag + loud-guard machinery in - `_tag_and_mesh` applies completely unchanged after defeaturing -- it is - the safety net that a defeatured model still gets tagged correctly (or - fails loudly if it doesn't), not something `defeature` needs to - duplicate. - - `keep_solid_id`, if given (Part B multi-body pipeline: the user has - assigned solid roles and the "casting" is whichever solid got that - role, not necessarily the largest), is threaded into `_import_and_reduce` - on BOTH passes in place of the largest-solid default -- see its - docstring. Default `None` -> existing largest-solid behavior, byte- - identical to before this parameter existed. An invalid/missing tag comes - back as an `_import_and_reduce` error string, which this function - already turns into a loud `ok=False` (same code path as an import - failure) -- no new branch needed. - """ - import gmsh - - notes: list[str] = [] - - # ---- Pass 1: unhealed ---- - gmsh.initialize() - try: - gmsh.option.setNumber("General.Terminal", 0) - if fast: - _apply_fast_options() - gmsh.model.add("part") - n_solids, volume_mm3, err = _import_and_reduce(step_path, keep_solid_id=keep_solid_id) - 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: - if keep_solid_id is None: - notes.append(f"multi-body: kept largest of {n_solids} solids (vol={volume_mm3:.1f} mm3)") - else: - notes.append(f"multi-body: kept solid_id={keep_solid_id} of {n_solids} solids " - f"(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) - - if out_step is not None: - try: - gmsh.write(out_step) - except Exception as e: - notes.append(f"cleaned-STEP export failed: {e}") - - 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) - if fast: - _apply_fast_options() - gmsh.model.add("part_healed") - n_solids2, volume_mm32, err2 = _import_and_reduce(step_path, keep_solid_id=keep_solid_id) - 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() - - -# -------------------------------------------------------------------------- -# OOD check -# -------------------------------------------------------------------------- -def load_ood_ranges(path: str = DEFAULT_OOD_RANGES_PATH) -> dict: - with open(path) as f: - return json.load(f) - - -def is_ood(volume_mm3: float, bbox_diag_mm: float, ranges: dict, tol: float = 1.05) -> bool: - """True if either the kept volume or the bbox diagonal exceeds the - training envelope (`data/features_v2/ood_ranges.json`) by more than - `tol` (default 1.05 = 5% margin) -- same convention as `src.ood.ood_flags`, - which applies the identical 5% tolerance before flagging, so this - parts-prep OOD check and the real-casting-inference OOD check agree.""" - vol_max = ranges.get("volume_mm3_max", float("inf")) - diag_max = ranges.get("bbox_diag_max", float("inf")) - return bool(volume_mm3 > vol_max * tol or bbox_diag_mm > diag_max * tol) - - -# -------------------------------------------------------------------------- -# Per-part subprocess isolation (real 300s-class time-box) -# -------------------------------------------------------------------------- -def _timeout_for(stem: str, default_timeout_s: float) -> float: - return SLOW_PART_TIMEOUTS_S.get(stem, default_timeout_s) - - -def _run_one_bounded(step_path: str, out_msh: str, out_step: str, target_nodes: int, - timeout_s: float) -> dict: - """Run `heal_and_mesh` for ONE part in a fresh subprocess (`nice -n 10`), - bounded by `timeout_s`. A hang or crash in this one part never blocks - the rest of the batch -- it's reported as a failed/timed-out row.""" - cmd = [ - "nice", "-n", "10", sys.executable, "-u", "-m", "tools.doe.prep_parts", - "--single", step_path, "--out-msh", out_msh, "--out-step", out_step, - "--target-nodes", str(target_nodes), - ] - try: - proc = subprocess.run(cmd, capture_output=True, text=True, - timeout=timeout_s, cwd=str(REPO_ROOT)) - except subprocess.TimeoutExpired: - return dict(n_solids=0, volume_mm3=0.0, bbox_diag_mm=0.0, n_nodes=0, - ok=False, healed=False, notes=f"TIMEOUT after {timeout_s:.0f}s") - - if proc.returncode != 0: - tail = (proc.stderr or "").strip().splitlines()[-5:] - return dict(n_solids=0, volume_mm3=0.0, bbox_diag_mm=0.0, n_nodes=0, - ok=False, healed=False, - notes=f"worker exit {proc.returncode}: {' | '.join(tail)}") - - lines = [ln for ln in proc.stdout.splitlines() if ln.strip()] - if not lines: - return dict(n_solids=0, volume_mm3=0.0, bbox_diag_mm=0.0, n_nodes=0, - ok=False, healed=False, notes="worker produced no output") - try: - return json.loads(lines[-1]) - except json.JSONDecodeError: - return dict(n_solids=0, volume_mm3=0.0, bbox_diag_mm=0.0, n_nodes=0, - ok=False, healed=False, notes=f"worker output not JSON: {lines[-1][:200]}") - - -# -------------------------------------------------------------------------- -# Batch CLI -# -------------------------------------------------------------------------- -_STEP_SUFFIXES = {".step", ".stp"} - - -def run_all(parts_dir: str, out_dir: str, target_nodes: int = DEFAULT_TARGET_NODES, - timeout_s: float = DEFAULT_TIMEOUT_S, - ood_ranges_path: str = DEFAULT_OOD_RANGES_PATH) -> list[dict]: - out_path = Path(out_dir) - out_path.mkdir(parents=True, exist_ok=True) - - ranges = load_ood_ranges(ood_ranges_path) if Path(ood_ranges_path).exists() else {} - if not ranges: - print(f"[prep_parts] WARNING: OOD ranges not found at {ood_ranges_path}; " - f"OOD column will read 'unknown' for every part", flush=True) - - step_files = sorted(p for p in Path(parts_dir).iterdir() - if p.suffix.lower() in _STEP_SUFFIXES) - - rows: list[dict] = [] - for step_path in step_files: - stem = step_path.stem - out_msh = out_path / f"{stem}.msh" - out_step = out_path / f"{stem}_cleaned.step" - part_timeout = _timeout_for(stem, timeout_s) - - print(f"[prep_parts] {step_path.name} (timeout={part_timeout:.0f}s) ...", flush=True) - t0 = time.time() - try: - result = _run_one_bounded(str(step_path), str(out_msh), str(out_step), - target_nodes, part_timeout) - except Exception as e: - result = dict(n_solids=0, volume_mm3=0.0, bbox_diag_mm=0.0, n_nodes=0, - ok=False, healed=False, notes=f"controller error: {e}") - dt = time.time() - t0 - - if ranges: - ood = is_ood(result["volume_mm3"], result["bbox_diag_mm"], ranges) - ood_str = "yes" if ood else "no" - else: - ood = False - ood_str = "unknown" - - status = "OK" if result["ok"] else "FAILED" - if result["ok"] and ood: - status = "OK (OOD)" - - row = dict(file=step_path.name, n_solids_in=result["n_solids"], - kept_vol_mm3=result["volume_mm3"], diag_mm=result["bbox_diag_mm"], - healed=result["healed"], mesh_nodes=result["n_nodes"], - ood=ood_str, status=status, notes=result.get("notes", ""), - seconds=dt) - rows.append(row) - print(f" -> {status}, nodes={result['n_nodes']}, {dt:.1f}s -- {result.get('notes', '')}", - flush=True) - - write_report(rows, str(out_path / "prep_report.md")) - return rows - - -def write_report(rows: list[dict], out_path: str) -> None: - n_ok = sum(1 for r in rows if r["status"].startswith("OK")) - lines = [ - "# Dataset-D Parts Prep Report\n\n", - f"Processed {len(rows)} parts -- {n_ok} usable (mesh OK), {len(rows) - n_ok} failed.\n\n", - "| file | n_solids_in | kept_vol_mm3 | diag_mm | healed? | mesh_nodes | OOD? | status | notes |\n", - "|---|---|---|---|---|---|---|---|---|\n", - ] - for r in rows: - lines.append( - f"| {r['file']} | {r['n_solids_in']} | {r['kept_vol_mm3']:.1f} | " - f"{r['diag_mm']:.1f} | {r['healed']} | {r['mesh_nodes']} | {r['ood']} | " - f"{r['status']} | {r['notes']} |\n" - ) - Path(out_path).write_text("".join(lines)) - - -# -------------------------------------------------------------------------- -# argparse entry point -# -------------------------------------------------------------------------- -def main(argv=None) -> int: - ap = argparse.ArgumentParser(description=__doc__.splitlines()[0] if __doc__ else "") - ap.add_argument("--parts-dir", default="pc_parts") - ap.add_argument("--out-dir", default="doe_out/parts_prepped") - ap.add_argument("--target-nodes", type=int, default=DEFAULT_TARGET_NODES) - ap.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT_S) - ap.add_argument("--ood-ranges", default=DEFAULT_OOD_RANGES_PATH) - # Hidden single-part worker mode, used internally via subprocess for a - # real per-part time-box (see `_run_one_bounded`). Not part of the - # documented CLI surface. - ap.add_argument("--single", default=None, help=argparse.SUPPRESS) - ap.add_argument("--out-msh", default=None, help=argparse.SUPPRESS) - ap.add_argument("--out-step", default=None, help=argparse.SUPPRESS) - args = ap.parse_args(argv) - - if args.single: - result = heal_and_mesh(args.single, args.out_msh, target_nodes=args.target_nodes, - out_step=args.out_step) - print(json.dumps(result)) - return 0 - - run_all(args.parts_dir, args.out_dir, target_nodes=args.target_nodes, - timeout_s=args.timeout, ood_ranges_path=args.ood_ranges) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tools/export_release.py b/tools/export_release.py deleted file mode 100644 index 462a540c855a912a93b2327710cc6ab8091bf817..0000000000000000000000000000000000000000 --- a/tools/export_release.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Convert the internal v5.1 artifacts into the public release layout. - - python -m tools.export_release --out weights - -Writes `gnn_v5.1.safetensors`, `gbm_total_v3.skops`, `gbm_std_v2.skops`, -`ood_ranges.json` and `SHA256SUMS` into `--out`, then reloads every file -through `src.release_io` and checks it against the original: -GNN tensors / norm / cfg bit-identical, GBM predictions identical on -randomized feature rows spanning each feature's training split range. -""" -import argparse -import hashlib -import json -import os -import shutil - -import joblib -import numpy as np -import torch - -from src import release_io as rio - -SOURCES = { - "gnn": "runs/ft3_stage1/ckpt_ep021.pt", - "gbm_total": "runs/ft3_gbm/hgb_v3.joblib", - "gbm_std": "runs/ft3_gbm/hgb_std_v2.joblib", - "ood": "data/features_v2/ood_ranges.json", -} - - -def _check_gnn(orig: dict, back: dict) -> None: - assert back["cfg"] == orig["cfg"], "cfg mismatch" - assert set(back["model"]) == set(orig["model"]), "state_dict keys differ" - for k, v in orig["model"].items(): - assert torch.equal(back["model"][k].cpu(), v.cpu()), f"tensor {k} differs" - assert set(back["norm"]) == set(orig["norm"]), "norm keys differ" - for k, v in orig["norm"].items(): - b = back["norm"][k] - if isinstance(v, torch.Tensor): - assert torch.equal(b, v.cpu()), f"norm {k} differs" - elif isinstance(v, np.ndarray): - assert b.dtype == v.dtype and np.array_equal(b, v), f"norm {k} differs" - else: - assert b == v, f"norm {k} differs" - - -def _gbm_probe_rows(model, n=10_000, seed=0) -> np.ndarray: - """Random rows inside each feature's histogram bin range (the binner's - thresholds are exactly the training split points), plus the edges.""" - rng = np.random.default_rng(seed) - thresholds = model._bin_mapper.bin_thresholds_ - cols = [] - for t in thresholds: - lo, hi = (float(t[0]), float(t[-1])) if len(t) else (0.0, 1.0) - pad = 0.1 * (hi - lo + 1e-9) - cols.append(rng.uniform(lo - pad, hi + pad, size=n)) - X = np.stack(cols, axis=1) - edges = np.array([[float(t[0]) if len(t) else 0.0 for t in thresholds], - [float(t[-1]) if len(t) else 1.0 for t in thresholds]]) - return np.vstack([X, edges]) - - -def _check_gbm(orig, back) -> int: - X = _gbm_probe_rows(orig) - assert np.array_equal(orig.predict(X), back.predict(X)), "GBM predictions differ" - return len(X) - - -def _sha256(path: str) -> str: - h = hashlib.sha256() - with open(path, "rb") as f: - for chunk in iter(lambda: f.read(1 << 20), b""): - h.update(chunk) - return h.hexdigest() - - -def export(out_dir: str = "weights") -> dict: - os.makedirs(out_dir, exist_ok=True) - out = { - "gnn": os.path.join(out_dir, "gnn_v5.1.safetensors"), - "gbm_total": os.path.join(out_dir, "gbm_total_v3.skops"), - "gbm_std": os.path.join(out_dir, "gbm_std_v2.skops"), - "ood": os.path.join(out_dir, "ood_ranges.json"), - } - state = torch.load(SOURCES["gnn"], map_location="cpu", weights_only=False) - rio.save_gnn_release(state, out["gnn"], extra_metadata={ - "release": "v5.1", "source_checkpoint": "ft3_stage1/ckpt_ep021", - "license": "apache-2.0"}) - _check_gnn(state, rio.load_gnn_state(out["gnn"])) - report = {"gnn_params": int(sum(v.numel() for v in state["model"].values()))} - - for key in ("gbm_total", "gbm_std"): - model = joblib.load(SOURCES[key]) - rio.save_gbm_release(model, out[key]) - report[f"{key}_rows_checked"] = _check_gbm(model, rio.load_gbm(out[key])) - report[f"{key}_n_features"] = int(model.n_features_in_) - - with open(SOURCES["ood"]) as f: - json.load(f) - shutil.copyfile(SOURCES["ood"], out["ood"]) - - with open(os.path.join(out_dir, "SHA256SUMS"), "w") as f: - for path in out.values(): - f.write(f"{_sha256(path)} {os.path.basename(path)}\n") - return report - - -if __name__ == "__main__": - ap = argparse.ArgumentParser() - ap.add_argument("--out", default="weights") - print(json.dumps(export(ap.parse_args().out), indent=2)) diff --git a/tools/recovery/README.md b/tools/recovery/README.md deleted file mode 100644 index e5f86bf48acacd984eb9e36b1c1b2186fb035cc1..0000000000000000000000000000000000000000 --- a/tools/recovery/README.md +++ /dev/null @@ -1,15 +0,0 @@ -# Node-ordering recovery (Phase 0, Task 0) - -Raw `snapshots_nodal`/`temps` are in dolfinx DOF order, scrambled vs gmsh node -order. These scripts reconstruct the per-mesh DOF->gmsh -permutation via dolfinx and validate physically. Run in the `fenicsx` env. - -Paths default to the repository's `data/` and `runs/` directories and can be -overridden with the `RECOVERY_*` environment variables defined at the top of -each script. - -1. `compute_perms.py` — all stems, coord-hash match -> data/features/perms/.npy - + runs/recovery_report.csv -2. `reprocess_failures.py` — re-do `no_bijection` stems with KDTree nearest match. - -Result: 9710/9710 perms, all physically valid. build_dataset applies the perm. diff --git a/tools/recovery/compute_perms.py b/tools/recovery/compute_perms.py deleted file mode 100644 index 2eedcede5f0efe9120bd82677676808feb1ba1ba..0000000000000000000000000000000000000000 --- a/tools/recovery/compute_perms.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Compute DOF->gmsh node permutation per casting via dolfinx and validate -physics (interior hotter than cooling after reorder). Saves compact int32 -perm per stem + a CSV report. Run in the fenicsx env. - -perm g: temps_gmsh[:, g] = temps_dof (g[dof] = gmsh_node_index) -""" -import os, glob, warnings, csv, sys, tempfile -import numpy as np -warnings.filterwarnings("ignore") -from mpi4py import MPI -from dolfinx import io, fem -import meshio - -RAW=os.environ.get("RECOVERY_RAW", "data/raw") -MSHROOT=os.environ.get("RECOVERY_MESH_ROOT", "data") -PTDIR=os.environ.get("RECOVERY_PT_DIR", "data/pt") -OUT=os.environ.get("RECOVERY_OUT", "data/features_v2/perms") -REPORT=os.environ.get("RECOVERY_REPORT", "runs/recovery_report.csv") -TMP=os.path.join(tempfile.gettempdir(), f"recovery_perm_{os.getpid()}.xdmf") -T_SOL=1120.0 -os.makedirs(OUT, exist_ok=True) - -def find_msh(stem): - for s in ("msh_none","msh_half","msh_all"): - h=glob.glob(os.path.join(MSHROOT,s,"*",stem+".msh")) - if h: return h[0] -def find_npz(stem): - h=glob.glob(os.path.join(RAW,"*",stem+".npz")); return h[0] if h else None - -def parse_msh(path): - lines=open(path).read().splitlines(); i=0; phys={}; ids=[]; xyz=[] - while i=1: - nm=phys.get(tags[0]) - if nm=="Cooling": cool.update(id2[v] for v in conn) - elif nm=="Adiabatic": adi.update(id2[v] for v in conn) - N=len(ids); cm=np.zeros(N,bool); am=np.zeros(N,bool); cm[list(cool)]=True; am[list(adi)]=True - return np.array(xyz),cm,am - i+=1 - -def compute(stem): - mp=find_msh(stem) - if mp is None: return dict(stem=stem, status="no_msh") - xyz,cool,adi=parse_msh(mp); inter=~(cool|adi); N=xyz.shape[0] - m=meshio.read(mp); tet=[cb.data for cb in m.cells if cb.type=="tetra"] - if not tet: return dict(stem=stem, status="no_tet") - meshio.write(TMP, meshio.Mesh(points=m.points, cells=[("tetra",tet[0])])) - with io.XDMFFile(MPI.COMM_SELF, TMP, "r") as xf: dom=xf.read_mesh(name="Grid") - V=fem.functionspace(dom,("Lagrange",1)); dof=V.tabulate_dof_coordinates()[:,:3] - if dof.shape[0]!=N: return dict(stem=stem, status=f"N_mismatch_{dof.shape[0]}_{N}") - gk={tuple(np.round(p,4)):j for j,p in enumerate(xyz)} - g=np.array([gk.get(tuple(np.round(p,4)),-1) for p in dof]) - if (g<0).any() or len(np.unique(g))!=N: - return dict(stem=stem, status="no_bijection") - # validate physics on raw temps - npzp=find_npz(stem) - sn=np.load(npzp)["snapshots_nodal"] - mid=sn.shape[0]//2 - tg=np.empty_like(sn); tg[:,g]=sn - db=float(sn[mid][inter].mean()-sn[mid][cool].mean()) if cool.any() and inter.any() else float('nan') - da=float(tg[mid][inter].mean()-tg[mid][cool].mean()) if cool.any() and inter.any() else float('nan') - status="valid" if (da>0 or np.isnan(da)) else "phys_fail" - np.save(os.path.join(OUT, f"{stem}.npy"), g.astype(np.int32)) - return dict(stem=stem, status=status, delta_before=round(db,2), delta_after=round(da,2)) - -def main(): - stems=sorted(os.path.splitext(os.path.basename(p))[0] for p in glob.glob(os.path.join(PTDIR,"*.pt"))) - print(f"stems: {len(stems)}", flush=True) - rows=[]; counts={} - with open(REPORT,"w",newline="") as f: - w=csv.DictWriter(f, fieldnames=["stem","status","delta_before","delta_after"]); w.writeheader() - for i,s in enumerate(stems): - try: r=compute(s) - except Exception as e: r=dict(stem=s, status=f"err_{type(e).__name__}") - w.writerow({k:r.get(k,"") for k in ["stem","status","delta_before","delta_after"]}) - counts[r["status"]]=counts.get(r["status"],0)+1 - if (i+1)%500==0: - f.flush(); print(f" {i+1}/{len(stems)} {counts}", flush=True) - print("DONE", counts, flush=True) - if os.path.exists(TMP): os.remove(TMP) - -if __name__=="__main__": main() diff --git a/tools/recovery/reprocess_failures.py b/tools/recovery/reprocess_failures.py deleted file mode 100644 index 31b48e840fc8c89fc0fdeb17672f2c0a62553480..0000000000000000000000000000000000000000 --- a/tools/recovery/reprocess_failures.py +++ /dev/null @@ -1,99 +0,0 @@ -"""Re-process no_bijection stems with KDTree nearest matching (robust to the -float noise from the XDMF roundtrip). Greedy top-k assignment guarantees a -bijection for near-identical point sets. Updates perms + a second report.""" -import os, glob, csv, warnings, tempfile -import numpy as np -warnings.filterwarnings("ignore") -from mpi4py import MPI -from dolfinx import io, fem -import meshio -from scipy.spatial import cKDTree - -MSHROOT=os.environ.get("RECOVERY_MESH_ROOT", "data") -RAW=os.environ.get("RECOVERY_RAW", "data/raw") -OUT=os.environ.get("RECOVERY_OUT", "data/features_v2/perms") -REPORT=os.environ.get("RECOVERY_REPORT", "runs/recovery_report.csv") -OUT2=os.environ.get("RECOVERY_FIX_REPORT", "runs/recovery_report_fix.csv") -TMP=os.path.join(tempfile.gettempdir(), f"recovery_fix_{os.getpid()}.xdmf") -T_SOL=1120.0 - -def find_msh(s): - for d in ("msh_none","msh_half","msh_all"): - h=glob.glob(f"{MSHROOT}/{d}/*/{s}.msh") - if h: return h[0] -def find_npz(s): - h=glob.glob(f"{RAW}/*/{s}.npz"); return h[0] if h else None - -def parse_msh(path): - lines=open(path).read().splitlines(); i=0; phys={}; ids=[]; xyz=[] - while i=1: - nm=phys.get(tags[0]) - if nm=="Cooling": cool.update(id2[v] for v in conn) - elif nm=="Adiabatic": adi.update(id2[v] for v in conn) - N=len(ids); cm=np.zeros(N,bool); am=np.zeros(N,bool); cm[list(cool)]=True; am[list(adi)]=True - return np.array(xyz),cm,am - i+=1 - -def greedy_match(dof, gxyz): - N=dof.shape[0]; tree=cKDTree(gxyz) - k=min(6,N); d,cand=tree.query(dof,k=k) - if k==1: cand=cand[:,None]; d=d[:,None] - order=np.argsort(d[:,0]) # assign closest-first - g=-np.ones(N,dtype=np.int64); used=np.zeros(N,bool) - ok=True - for i in order: - placed=False - for c in cand[i]: - if not used[c]: - g[i]=c; used[c]=True; placed=True; break - if not placed: ok=False; break - return (g if ok and (g>=0).all() and len(np.unique(g))==N else None), float(d[:,0].max()) - -def compute(stem): - mp=find_msh(stem) - if mp is None: return dict(stem=stem,status="no_msh") - xyz,cool,adi=parse_msh(mp); inter=~(cool|adi); N=xyz.shape[0] - m=meshio.read(mp); tet=[cb.data for cb in m.cells if cb.type=="tetra"] - if not tet: return dict(stem=stem,status="no_tet") - meshio.write(TMP, meshio.Mesh(points=m.points, cells=[("tetra",tet[0])])) - with io.XDMFFile(MPI.COMM_SELF,TMP,"r") as xf: dom=xf.read_mesh(name="Grid") - V=fem.functionspace(dom,("Lagrange",1)); dof=V.tabulate_dof_coordinates()[:,:3] - if dof.shape[0]!=N: return dict(stem=stem,status=f"N_mismatch") - g,maxd=greedy_match(dof,xyz) - if g is None: return dict(stem=stem,status="no_bijection2",delta_after=round(maxd,4)) - sn=np.load(find_npz(stem))["snapshots_nodal"]; mid=sn.shape[0]//2 - tg=np.empty_like(sn); tg[:,g]=sn - da=float(tg[mid][inter].mean()-tg[mid][cool].mean()) if cool.any() and inter.any() else float('nan') - status="valid" if (da>0 or np.isnan(da)) else "phys_fail" - if status=="valid": np.save(os.path.join(OUT,f"{stem}.npy"), g.astype(np.int32)) - return dict(stem=stem,status=status,delta_after=round(da,2),maxd=round(maxd,5)) - -fails=[r[0] for r in csv.reader(open(REPORT)) if len(r)>1 and r[1]=="no_bijection"] -print(f"reprocessing {len(fails)} failures", flush=True) -counts={} -with open(OUT2,"w",newline="") as f: - w=csv.DictWriter(f,fieldnames=["stem","status","delta_after","maxd"]); w.writeheader() - for i,s in enumerate(fails): - try: r=compute(s) - except Exception as e: r=dict(stem=s,status=f"err_{type(e).__name__}") - w.writerow({k:r.get(k,"") for k in ["stem","status","delta_after","maxd"]}) - counts[r["status"]]=counts.get(r["status"],0)+1 - if (i+1)%200==0: f.flush(); print(f" {i+1}/{len(fails)} {counts}",flush=True) -print("DONE",counts,flush=True) -if os.path.exists(TMP): os.remove(TMP)