diff --git a/.gitattributes b/.gitattributes
index a6344aac8c09253b3b630fb776ae94478aa0275b..8b26e119274859dc77dd34a481726fccfd61868d 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
*.zst filter=lfs diff=lfs merge=lfs -text
*tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/3dforms.png filter=lfs diff=lfs merge=lfs -text
+assets/mesh.png filter=lfs diff=lfs merge=lfs -text
+assets/solidification.gif filter=lfs diff=lfs merge=lfs -text
+weights/gbm_std_v2.skops filter=lfs diff=lfs merge=lfs -text
+weights/gbm_total_v3.skops filter=lfs diff=lfs merge=lfs -text
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..1dc6ac797a4d4781433194a3df32bae293486bb0
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,15 @@
+# Raw and derived data are large. Keep only the small public property tables.
+data/*
+!data/properties/
+!data/properties/*.csv
+# Training artifacts, checkpoints, logs
+runs/
+# Model checkpoints anywhere
+*.pt
+# SDD scratch (progress ledger, briefs, review packages)
+.superpowers/
+# Python
+__pycache__/
+*.pyc
+.ipynb_checkpoints/
+.pytest_cache/
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..e06d2081865a766a8668acc12878f98b27fc9ea0
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,202 @@
+Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "{}"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. We also recommend that a
+ file or class name and description of purpose be included on the
+ same "printed page" as the copyright notice for easier
+ identification within third-party archives.
+
+ Copyright {yyyy} {name of copyright owner}
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+
diff --git a/NOTICE b/NOTICE
new file mode 100644
index 0000000000000000000000000000000000000000..80050bd01ab06ba72f5facfc3164e37d48ceb4ba
--- /dev/null
+++ b/NOTICE
@@ -0,0 +1,9 @@
+CastSolid GNN (casting solidification surrogate)
+Copyright 2026 Eugen Miknevic
+
+This product includes software and trained model weights developed by
+Eugen Miknevic, released under the Apache License, Version 2.0.
+
+The graph network follows the encode-process-decode design of MeshGraphNets
+(Pfaff et al., "Learning Mesh-Based Simulation with Graph Networks", ICLR 2021).
+Reference solidification fields were generated with FEniCSx (dolfinx).
diff --git a/README.md b/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..27ae3df72751b68eb32a77b7915186a500edc7d8
--- /dev/null
+++ b/README.md
@@ -0,0 +1,464 @@
+---
+license: apache-2.0
+library_name: pytorch
+pipeline_tag: graph-ml
+language:
+ - en
+tags:
+ - graph-neural-network
+ - gnn
+ - meshgraphnet
+ - film
+ - surrogate-model
+ - scientific-machine-learning
+ - physics-ml
+ - heat-transfer
+ - casting
+ - metal-casting
+ - foundry
+ - solidification
+ - shrinkage
+ - finite-element
+ - fenicsx
+ - pytorch-geometric
+ - scikit-learn
+---
+
+# CastSolid GNN v5.1: casting solidification surrogate
+
+CastSolid GNN is a graph neural network surrogate for screening how metal castings solidify. You give it
+a STEP model, an alloy and a mold condition. The pipeline meshes the part, predicts the
+solidification time at every mesh node and marks the regions that freeze last. Those hot spots
+are where feeding problems and shrinkage defects usually start.
+
+A transient finite-element solve of the same field takes minutes to hours. Once the part is
+meshed, this model returns the field in seconds, which makes it useful early in design: you can
+compare concepts, see where the hot regions are, and decide which variants are worth a full
+simulation.
+
+It does not replace a calibrated process simulation or physical validation. A hot spot marks
+feeding difficulty and shrinkage risk. It is not a porosity label.
+
+
+ 
+ Predicted solidification playback. Solidified material turns translucent while the
+ remaining liquid contracts toward the hot spot.
+
+
+| | |
+|---|---|
+| Version | v5.1, frozen on 31 July 2026 |
+| Task | Per-node regression on a tetrahedral mesh (graph ML) |
+| Inputs | STEP geometry in mm, alloy grade, mold system, optional chill and sleeve body roles |
+| Outputs | Solidification time per node, total time, top-10% hot-spot mask, out-of-distribution warnings, shrinkage screening fields |
+| Architecture | 12-block MeshGraphNet-style GNN with FiLM conditioning, plus two gradient-boosted heads |
+| Training data | FEniCSx enthalpy-method conduction simulations: about 9.7k primitive cases, 4.7k multi-alloy and multi-BC cases, 500 real-part cases |
+| Framework | PyTorch 2.3, PyTorch Geometric 2.7, scikit-learn 1.5.2, NumPy 2 |
+| License | Apache-2.0 for code and weights |
+| Code | Included in this repository. Showcase on [GitHub](https://github.com/eugenmik/casting-gnn-solidification) |
+| Author | Eugen Miknevic, independent foundry engineer |
+
+## Contents
+
+- [What is in this repository](#what-is-in-this-repository)
+- [Quick start](#quick-start)
+- [How the predictor works](#how-the-predictor-works)
+- [Training data](#training-data)
+- [Training procedure](#training-procedure)
+- [Evaluation](#evaluation)
+- [Intended use and limitations](#intended-use-and-limitations)
+- [The node-order bug](#the-node-order-bug)
+- [Reproducibility](#reproducibility)
+- [License and citation](#license-and-citation)
+
+## What is in this repository
+
+The predictor combines three trained models. Use them together; none of them is meant to run
+alone.
+
+| File | Size | Role |
+|---|---:|---|
+| `weights/gnn_v5.1.safetensors` | 7.4 MB | GNN weights, feature and target normalization, and the architecture config (stored in the safetensors metadata). Predicts the shape of the field |
+| `weights/gbm_total_v3.skops` | 0.6 MB | `HistGradientBoostingRegressor` on 16 features. Predicts the total solidification time of the part, which sets the absolute scale |
+| `weights/gbm_std_v2.skops` | 0.8 MB | `HistGradientBoostingRegressor` with absolute-error loss on 17 features. Predicts the spread (standard deviation) of the field |
+| `weights/ood_ranges.json` | <1 KB | Training envelope used by the geometry out-of-distribution check |
+| `weights/SHA256SUMS` | <1 KB | Checksums of the files above |
+| `src/` | | The full pipeline: CAD healing and meshing, features, GNN, recomposition, shrinkage screening, VTU export, training scripts and a local web viewer |
+| `tools/` | | Density-targeted STEP meshing, the property loader, the sleeve registry and the scripts that recover solver-to-mesh node order |
+| `data/properties/` | | Material, mold, chill and exothermic-sleeve tables read at inference |
+| `tests/` | | Deterministic unit tests that need neither weights nor simulation data |
+
+None of the weight files can execute code when loaded. The GNN is stored as
+[safetensors](https://github.com/huggingface/safetensors) and the gradient-boosted 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. All three files were
+converted from the internal PyTorch and joblib pickles and checked against them (see
+[Reproducibility](#reproducibility)).
+
+## Quick start
+
+### 1. Install
+
+Python 3.11 is recommended. Install PyTorch 2.3 for your hardware first (CPU works for
+inference), then the matching `torch-cluster` wheel, then the rest:
+
+```bash
+hf download eugenmik/castsolid-gnn --local-dir castsolid-gnn
+cd castsolid-gnn
+
+python -m venv .venv && source .venv/bin/activate
+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
+pip install -r requirements.txt
+pytest -q # optional, no GPU needed
+```
+
+For a CUDA build, replace `cpu` with your CUDA tag (for example `cu121`) in both URLs.
+The `gmsh` Python wheel handles CAD healing and tetrahedral meshing. `pyvista` and `trame` are
+only needed for the browser viewer.
+
+### 2. Predict from the command line
+
+Run this from the repository root. The weight paths resolve on their own.
+
+```bash
+python -m src.predict_casting path/to/part.step --out-dir out \
+ --alloy EN-GJS-500-7 --mold green_sand
+```
+
+It writes three files:
+
+- `out/.vtu` is the tetrahedral mesh with point fields `solidification_time` (s),
+ `hotspot` (the 10% of nodes that solidify last), `niyama`, `porosity_risk` and
+ `macro_region_id`. Open it in ParaView.
+- `out/_summary.txt` holds the total time, the coordinates of the five latest nodes, any
+ OOD warnings and the shrinkage screening summary.
+- `out/.msh` is the mesh the prediction actually used.
+
+Meshing runs through the same pipeline that prepared the training parts: keep the largest
+solid, heal it with OpenCASCADE, then search the element size until the mesh has about 8,500
+nodes. A second run on the same STEP file reuses the cached mesh.
+
+### 3. Predict from Python
+
+```python
+from src.predict_casting import predict_casting
+from src.inference_props import list_alloys, list_molds
+
+print(list_alloys()) # 12 named grades: EN-GJL-200/250/300, EN-GJS-400-15/500-7/600-3/700-2,
+ # GS-38, GS-45, GS-52, G20Mn5, G17CrMo5-5
+print(list_molds()) # ['green_sand', 'no_bake_furan_HTS']
+
+r = predict_casting("part.step", out_dir="out", alloy="EN-GJS-500-7", mold="green_sand")
+r["total_time_s"] # predicted total solidification time, s
+r["hotspot_xyz"] # (5, 3) mm coordinates of the latest-solidifying nodes
+r["ood_warnings"] # an empty list means the part is inside the training envelope
+r["shrinkage"] # Niyama and macro-shrinkage screening, with a confidence label
+```
+
+To use an alloy that is not in the table, pass a JSON file with exactly these keys (SI units,
+temperatures in °C):
+
+```json
+{"k": 41.5, "rho": 7150.0, "cp": 535.0, "L_kJ_kg": 240.0,
+ "TL": 1220.0, "TS": 1150.0, "Tpour": 1320.0, "mold_effusivity": 1018.0}
+```
+
+```bash
+python -m src.predict_casting part.step --props my_alloy.json
+```
+
+### 4. Assemblies with chills or a feeder sleeve
+
+By default the whole outer surface cools into a sand mold. For a multi-body STEP you can say
+which solid is the casting and which ones are chills or a sleeve:
+
+```bash
+python -m src.assembly part.step # lists solids with id, volume, bounding box
+python -m src.predict_casting part.step --bodies bodies.json
+```
+
+```json
+[{"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "chill", "params": {"material": "cast_iron"}},
+ {"solid_id": 3, "role": "chill", "params": {"material": "cast_iron"}}]
+```
+
+The roles are `casting`, `chill`, `insulating_sleeve`, `exo_sleeve` and `ignore`. Every
+training case had chills only, one sleeve only, or neither, so the code rejects any other
+combination instead of guessing. A sleeve is either named from
+`tools/sleeves/sleeve_registry.csv` (`{"sleeve": "FOSECO_P12_15K"}`) or given a manual modulus
+extension factor (`{"MEF": 1.55}`). Only `cast_iron` chills appear in the training data.
+
+### 5. Browser viewer (optional)
+
+```bash
+python -m src.webapp # http://localhost:8080
+```
+
+This is a local trame and PyVista app. You upload a STEP file and run the prediction, then move
+a time slider (solidified material turns translucent), toggle the hot spots, or export a
+solidification GIF from the current camera angle.
+
+### Runtime
+
+The GNN itself runs in seconds on a CPU or GPU. CAD healing and meshing take most of the time,
+from seconds to a few minutes for typical parts. Very large or very detailed STEP files, with
+thousands of small faces, can run past the meshing time limit.
+
+## How the predictor works
+
+```mermaid
+flowchart LR
+ A[STEP + alloy + mold
+ optional body roles] --> B[CAD healing +
tetrahedral mesh]
+ B --> C[15 node features
+ 5-value alloy/BC vector]
+ C --> D1[12-block GNN
field shape]
+ C --> D2[GBM total head
absolute scale]
+ C --> D3[GBM spread head
field contrast]
+ D1 --> E[Soft-anchor
recomposition]
+ D2 --> E
+ D3 --> E
+ E --> F[Time field + hot spots
+ OOD + shrinkage screening]
+ F --> G[VTU / browser viewer]
+```
+
+### Why three models
+
+The GNN gets the spatial pattern of solidification right more reliably than its absolute size.
+The size depends heavily on alloy, mold and part dimensions, and two small gradient-boosted
+models estimate it better from a handful of global descriptors. On the original benchmark with
+held-out geometry families, the hybrid brought nodal MAE down from 5.77 s (GNN alone) to
+4.26 s.
+
+### Graph network
+
+Nodes are mesh vertices. Edges are the union of the tetrahedral mesh edges and the 16 nearest
+neighbours of each node. Each edge carries its displacement vector, its length and a flag that
+says whether it is a mesh edge, so the model sees geometry in millimetres but never absolute
+coordinates.
+
+Each node has 15 features. All of them come from the mesh and its boundary labels; file names
+and part identity are never used.
+
+- Geodesic distance through the mesh to the nearest cooling surface and to the nearest
+ adiabatic surface, plus straight-line distance to the surface and to adiabatic faces.
+- Cooling, adiabatic and interior flags.
+- 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. This is
+ a local version of Chvorinov's modulus.
+- Two boundary-condition channels. `boundary_extraction` encodes heat drawn out by chills.
+ `exothermic_input` encodes heat from an exothermic sleeve, spread along the mesh with a
+ 15 mm decay length.
+
+A vector `g` of five dimensionless numbers describes the alloy and the mold: Stefan number,
+normalized freezing range, thermal diffusivity ratio, normalized superheat and normalized mold
+effusivity. It conditions every processor block through FiLM (`h ← h·(1+γ(g)) + β(g)`).
+
+The network follows the encode-process-decode design of
+[MeshGraphNets](https://arxiv.org/abs/2010.03409). MLP encoders embed nodes and edges, 12
+message-passing blocks (hidden width 128, residual connections, LayerNorm, mean aggregation)
+process them, and a two-head decoder predicts the per-node deviation and a graph-level scale
+separately. The model has about 1.85 M parameters. Its target is the standardized `log1p` of the
+time (s) at which each node cools through the solidus.
+
+### Gradient-boosted heads
+
+Both heads are scikit-learn `HistGradientBoostingRegressor` models. They see one row per part:
+alloy properties (k, ρ, cp, latent heat, liquidus, solidus, pouring temperature), the part's
+volume-to-area ratio, mold effusivity, a one-hot boundary regime (none, chills, insulating
+sleeve or exothermic sleeve), the number of chills, the sleeve MEF and the exothermic heat flux.
+The spread head also gets the fraction of boundary area covered by chills or a sleeve. The total
+head predicts `log1p` of the total solidification time; the spread head predicts `log1p` of the
+standard deviation of the nodal field and is trained with absolute-error loss.
+
+### Soft-anchor recomposition
+
+The GNN field is scaled, shifted and clipped at zero:
+
+```
+pred(a, m) = max(0, a · (gnn − mean(gnn)) + m)
+```
+
+The pair `(a, m)` minimizes
+
+```
+J = w_total · ((q99.5(pred) − T_total) / T_total)² + w_std · ((std(pred) − S) / S)²
+```
+
+where `T_total` and `S` come from the two heads, `w_total = 1` and `w_std = 4`. When both
+targets can be met exactly, the result equals the closed-form affine solve. When clipping makes
+that impossible, the total anchor is allowed to move by a few percent so the spread stays close.
+The post-clip spread is not monotone in `a`, so the solver brackets the smallest root
+explicitly. The code is in `src/hybrid.py::hybrid_field_soft`.
+
+### Shrinkage screening
+
+`src/shrinkage.py` post-processes the predicted field. It computes a time-based Niyama proxy,
+finds macro-shrinkage regions (hot regions that liquid can no longer reach to feed them), places
+gravity-aware sinks and suggests feeder moduli. It is deterministic and adds about half a second
+per part. On parts with chills the Niyama result carries `niyama_confidence: "low"`, because a
+proxy built only from time cannot see the extra heat a chill pulls out locally. Pass
+`--no-shrinkage` to skip it.
+
+## Training data
+
+Every reference field was computed with [FEniCSx](https://fenicsproject.org/) (dolfinx) using an
+enthalpy-method model of transient heat conduction with latent heat. The model has no mold
+filling, flow, gas or stress.
+
+
+ 
+ Some of the primitive geometry families from the first training regime.
+
+
+| Set | Cases | Content |
+|---|---:|---|
+| Primitives (original regime) | 9,710 | 42 parametric geometry families, a single alloy, sand mold, different assignments of cooling and adiabatic faces |
+| Multi-alloy and multi-BC (v4) | 4,727 | Realistic thick-walled castings in 12 named iron and steel grades plus sampled synthetic alloys. Boundary regimes: sand only, chills, insulating sleeves, exothermic sleeves |
+| Real-part assemblies (E) | 500 | Prepared production geometries with virtual chill and sleeve contact patches |
+
+
+ 
+ Training meshes carry three physical groups: Volume, Cooling (mold contact) and
+ Adiabatic.
+
+
+The meshes are gmsh tetrahedral meshes with roughly 7,500 to 10,000 nodes. The target at each
+node is the time its temperature drops through the solidus. Material properties come from
+handbook data, and the subset used at inference is in `data/properties/`. The exothermic sleeve
+heat input is calibrated so that the modulus extension falls in vendor-published ranges. The
+raw simulation corpus (about 100 GB) is not distributed.
+
+## Training procedure
+
+v5.1 came out of a chain of fine-tunes, not a single training run.
+
+1. SP1 trained a 12-block two-head GNN from scratch on the primitive set, with the union graph,
+ local-modulus features and a loss on both nodal and total time.
+2. FT2b fine-tuned it on the multi-alloy, multi-BC set. The FiLM layers started at zero, so
+ training began from exactly the SP1 function.
+3. Exo-fix re-encoded the exothermic channel. As a sparse indicator with +19σ spikes it had
+ been suppressed by the network; as a geodesic-diffused field it is used. Exothermic cases were
+ oversampled three times, and the checkpoint was picked on gate metrics instead of validation
+ loss alone.
+4. FT3 fine-tuned on the real-part assemblies (set E) with the normalization statistics frozen.
+ Its epoch-21 checkpoint is the v5.1 GNN.
+5. The BC-aware total head (v3), the absolute-error spread head (v2) and the soft anchor were
+ selected on held-out real-part validation rows.
+
+Training ran on one RTX 3060 (12 GB) with batch size 2, at 15 to 20 minutes per epoch. To be
+promoted, v5.1 had to pass four gates on held-out data: two gold-standard real assemblies, 24
+held-out real parts, an extrapolation split and a check against collapsed field variance.
+
+## Evaluation
+
+Do not compare numbers across regimes. Total times in the primitive regime are tens of seconds;
+on real geometry they run to hundreds or thousands.
+
+| Evaluation | n | Nodal MAE | Hot-spot overlap | What it shows |
+|---|---:|---:|---:|---|
+| Held-out primitive geometry families | family split | 4.26 s | 0.917 | Generalization to unseen synthetic shape families in the original regime |
+| Unseen real single-body castings | 6 | 80.3 s median | 0.738 median | Generalization to new real geometry, scored against a FEniCSx reference on the same mesh |
+| Unseen real sleeve assemblies | 2 | 443.7 s median | 0.932 median | The hot regions rank well, but the field is too smooth (predicted spread is about 0.23 of the reference) and the absolute times are off |
+
+Hot-spot overlap compares the 10% of nodes that solidify last in the prediction with the same
+set in the reference. 1.0 means the two sets are identical.
+
+The real-geometry study started from 10 production CAD models. Two had self-intersecting
+B-reps and could not be meshed, which left eight scored parts. One single-body part reached a
+hot-spot overlap of only 0.035, so the median says nothing certain about any one part.
+
+> **Important.** The training data and the real-geometry references come from the same FEniCSx
+> heat-transfer model, the same material database and the same boundary assumptions. These
+> numbers measure how well the surrogate carries over to new geometry. They are not evidence of
+> physical accuracy against instrumented castings.
+
+## Intended use and limitations
+
+Good uses:
+
+- comparing early casting concepts and gating or feeding variants;
+- finding candidate last-to-solidify regions;
+- deciding which designs need a full process simulation;
+- early quoting and design discussions, as long as the limitations stay visible.
+
+Out of scope:
+
+- final process sign-off or contractual acceptance;
+- guaranteed defect or porosity prediction;
+- mold filling, flow, gas, stress or distortion;
+- unattended decisions on geometry outside the training envelope;
+- safety-critical decisions without an independent engineering review.
+
+Known limitations:
+
+1. The physics is pure conduction with latent heat, and the reference solver shares the
+ surrogate's property and boundary assumptions.
+2. STEP files must be in millimetres; units are not detected. STL is not supported, because
+ healing and face tagging need a B-rep.
+3. Invalid or very detailed B-reps can fail to mesh or run past the meshing time limit.
+4. The default density of about 8,500 nodes can under-resolve large thin-walled parts.
+5. On assemblies the predicted field is over-smoothed, even when the hot regions are ranked
+ well.
+6. The OOD check only covers part size and distance to cooling. It is not a calibrated
+ uncertainty estimate.
+7. The two mold systems currently share the same interim thermal properties, so choosing one
+ over the other does not change the prediction.
+8. The Niyama proxy has low confidence on parts with chills and is flagged that way.
+
+Treat each result as a screening hypothesis: where should an engineer look first, and is the
+concept ready for a full simulation? Keep the mesh, the inputs, the model version and the
+warnings together with every result. OOD, low-confidence and high-consequence cases belong in a
+calibrated simulation and a physical review.
+
+## The node-order bug
+
+The worst defect in this project was outside the neural network. The solver stored
+temperatures in dolfinx degree-of-freedom order, while coordinates and geometric features were
+in gmsh node order. About 60% of the nodal targets ended up attached to the wrong positions,
+silently, and every early model failure traced back to that.
+
+`tools/recovery/` rebuilds the permutation between the two orderings, and
+`src/build_dataset.py` applies it when it builds the graphs. If you work on simulation-to-ML
+pipelines, check node identity across every solver and mesh interface before you start tuning a
+model.
+
+## Reproducibility
+
+`tools/export_release.py` converted the internal checkpoints with the functions in
+`src/release_io.py`:
+
+| Release file | Internal artifact | Check |
+|---|---|---|
+| `gnn_v5.1.safetensors` | `ckpt_ep021.pt` (FT3, epoch 21) | Every tensor bit-identical; same config and normalization |
+| `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 part, the release files and the internal pickles give the same predicted field.
+The training scripts are in `src/` (`train_gnn.py`, `train_ft2b.py`, `gbm_ft2b.py`), but the
+simulation corpus is not part of this release.
+
+Software used: Python 3.11, PyTorch 2.3.0, PyTorch Geometric 2.7.0, scikit-learn 1.5.2, NumPy 2.x,
+gmsh 4.13 or newer and SciPy. The gradient-boosted heads were trained with scikit-learn 1.5.2,
+and other versions may refuse to load them or load them with changed behaviour.
+
+## License and citation
+
+Code and weights are released under the [Apache License 2.0](LICENSE). The tables in
+`data/properties/` are modelling inputs, not certified material specifications. Check them
+against your actual alloy and process.
+
+```bibtex
+@software{miknevic_castsolid_gnn_2026,
+ author = {Miknevic, Eugen},
+ title = {CastSolid GNN: a graph neural network surrogate for
+ rapid solidification screening of metal castings},
+ version = {5.1},
+ year = {2026},
+ url = {https://huggingface.co/eugenmik/castsolid-gnn}
+}
+```
+
+Contact: Eugen Miknevic, eugenmiknevic@gmail.com. If you have failure cases, questions or
+measured data from real castings, please post them in the
+[Community tab](https://huggingface.co/eugenmik/castsolid-gnn/discussions).
diff --git a/assets/3dforms.png b/assets/3dforms.png
new file mode 100644
index 0000000000000000000000000000000000000000..b7b5580d85c19a15d05a1fc624ae423a7256e0c8
--- /dev/null
+++ b/assets/3dforms.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f560e08a9d1d0dd0ee37bebc2398e96078d3eb7d70498ee9de73630c4125f9c6
+size 204096
diff --git a/assets/mesh.png b/assets/mesh.png
new file mode 100644
index 0000000000000000000000000000000000000000..1714c2c6b076c33cba5f48c609e0c0ff82963011
--- /dev/null
+++ b/assets/mesh.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f47539fc931f4f38f417992b323d7271de0ad2ef36dc4e108a0453ffad449cb6
+size 1065363
diff --git a/assets/solidification.gif b/assets/solidification.gif
new file mode 100644
index 0000000000000000000000000000000000000000..0f9facea39ecc5ed95aedc5ba93e8a3c9c48f222
--- /dev/null
+++ b/assets/solidification.gif
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32aa99b45277161680d10e51f62b9a0665b36fc09ab1f20795415baa0a225d80
+size 4577374
diff --git a/data/properties/README.md b/data/properties/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..4b9bb44f2c5f23368d936cc926cfb95fd6c9bc9d
--- /dev/null
+++ b/data/properties/README.md
@@ -0,0 +1,13 @@
+# Public runtime property tables
+
+These compact CSV files contain the numeric rows required by the public inference and unit-test
+code. They are a curated runtime subset of the project's larger property record.
+
+- `materials.csv`: 12 named iron and steel grades.
+- `molds.csv`: two UI labels currently sharing the same interim thermal properties.
+- `chills.csv`: example chill materials; incomplete rows are skipped by the loader.
+- `exo_calibration.csv`: the MEF-to-heat-flux calibration used by the example sleeve path.
+
+The values remain modeling inputs, not certified material specifications. A production user
+should replace or verify them for the actual alloy, mold and process and should retain the exact
+property row with every result.
diff --git a/data/properties/chills.csv b/data/properties/chills.csv
new file mode 100644
index 0000000000000000000000000000000000000000..2aef40f1135240285d609a54fe6aadb9f325bb56
--- /dev/null
+++ b/data/properties/chills.csv
@@ -0,0 +1,5 @@
+material,k_W_mK,rho_kg_m3,cp_J_kgK,effusivity_W_s05_m2K
+graphite,68.97,1922,1692,14978
+cast_iron,38.9,7150,515,11968
+steel,51.0,7860,486,13958
+copper,OWNER,OWNER,OWNER,OWNER
diff --git a/data/properties/exo_calibration.csv b/data/properties/exo_calibration.csv
new file mode 100644
index 0000000000000000000000000000000000000000..3ccb72b9b415765eb76d07de75a9902153060ab7
--- /dev/null
+++ b/data/properties/exo_calibration.csv
@@ -0,0 +1,13 @@
+MEF,q0_W_m2,duration_s,t_sand_s,t_exo_s,ratio,ratio_target
+1.4,640000.0,105.68312000817032,352.2770666939011,639.1643789347195,1.814379757766347,1.9599999999999997
+1.45,960000.0,105.68312000817032,352.2770666939011,754.0987323561276,2.1406410000892153,2.1025
+1.49,960000.0,105.68312000817032,352.2770666939011,760.0587908504301,2.1575596674047954,2.2201
+1.54,960000.0,105.68312000817032,352.2770666939011,766.9076476958378,2.1770013441216016,2.3716
+1.58,1120000.0,105.68312000817032,352.2770666939011,827.4879309999664,2.348969062238114,2.4964000000000004
+1.63,1280000.0,105.68312000817032,352.2770666939011,890.1572505029782,2.5268668745798584,2.6569
+1.67,1280000.0,105.68312000817032,352.2770666939011,895.2697715970885,2.541379658912063,2.7889
+1.72,1920000.0,105.68312000817032,352.2770666939011,1131.8937331840514,3.2130781143569904,2.9583999999999997
+1.76,1920000.0,105.68312000817032,352.2770666939011,1138.063508699037,3.2305921000753584,3.0976
+1.81,1920000.0,105.68312000817032,352.2770666939011,1145.2364712168662,3.2509538073677096,3.2761
+1.85,1920000.0,105.68312000817032,352.2770666939011,1150.5882262573784,3.266145699053246,3.4225000000000003
+1.9,1920000.0,105.68312000817032,352.2770666939011,1156.8270113469869,3.2838555805058167,3.61
diff --git a/data/properties/materials.csv b/data/properties/materials.csv
new file mode 100644
index 0000000000000000000000000000000000000000..95781345a4d30ea2a0264147b79d75c6dc21fb33
--- /dev/null
+++ b/data/properties/materials.csv
@@ -0,0 +1,13 @@
+material_id,family,k_W_mK,rho_kg_m3,cp_J_kgK,L_kJ_kg,TL_C,TS_C,Tpour_lo_C,Tpour_hi_C
+EN-GJL-200,grey_iron,51.9,7000,535,240,1175,1150,1265,1285
+EN-GJL-250,grey_iron,41.5,7150,535,240,1220,1150,1310,1330
+EN-GJL-300,grey_iron,36.3,7300,535,240,1245,1150,1335,1355
+EN-GJS-400-15,ductile_iron,32.9,6800,515,220,1160,1120,1250,1270
+EN-GJS-500-7,ductile_iron,36.0,7100,515,220,1160,1120,1250,1270
+EN-GJS-600-3,ductile_iron,27.7,7250,515,220,1160,1120,1250,1270
+EN-GJS-700-2,ductile_iron,26.0,7400,515,220,1160,1120,1250,1270
+GS-38 (GE200),low_alloy_steel,51.0,7800,599,260,1530,1495,1620,1640
+GS-45 (GP240GH),low_alloy_steel,50.8,7800,586,260,1530,1495,1620,1640
+GS-52 (GE300),low_alloy_steel,50.7,7800,586,260,1530,1495,1620,1640
+G20Mn5,low_alloy_steel,48.4,7800,590,260,1530,1495,1620,1640
+G17CrMo5-5,low_alloy_steel,28.1,7800,595,251,1520,1420,1610,1630
diff --git a/data/properties/molds.csv b/data/properties/molds.csv
new file mode 100644
index 0000000000000000000000000000000000000000..1505f10718cd63aeea019a8168a280c65ccd7cb4
--- /dev/null
+++ b/data/properties/molds.csv
@@ -0,0 +1,3 @@
+mold_system,k_W_mK,rho_kg_m3,cp_J_kgK,effusivity_W_s05_m2K
+green_sand,0.607,1520,1123,1018
+no_bake_furan_HTS,0.607,1520,1123,1018
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..354c9c233f65a905658447fe21542ea921577d8c
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,21 @@
+# Install PyTorch 2.3 for your platform first (see README), plus the matching
+# torch-cluster wheel, e.g. for 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
+gmsh>=4.13
+meshio>=5.3
+pytest>=8.3
+# Optional: local browser viewer (python -m src.webapp)
+pyvista>=0.44
+trame>=3.6
+trame-vtk>=2.8
+trame-vuetify>=2.7
+trame-pyvista>=0.1
+imageio>=2.34
diff --git a/src/assembly.py b/src/assembly.py
new file mode 100644
index 0000000000000000000000000000000000000000..9284aa761185b6b9b774b731265ce3cc145126ba
--- /dev/null
+++ b/src/assembly.py
@@ -0,0 +1,904 @@
+"""Multi-body STEP assembly tooling -- Part B Task 1.
+
+Two pieces:
+ - `list_solids(step_path)` -- the assembly inspector: fresh gmsh/OCC
+ session, imports a (possibly multi-solid) STEP file with entity names
+ preserved (`Geometry.OCCImportLabels=1`, set BEFORE `importShapes` --
+ setting it after import is a no-op), and reports one row per 3-D solid
+ (id/name/volume/bbox/center-of-mass), sorted by volume descending. This
+ is the data source for the later "role menu" (casting / chill /
+ insulating sleeve / exothermic sleeve) that a human or heuristic
+ assigns per solid.
+ - `compose_assembly(base_step, attachments, out_step)` -- a test/fixture
+ helper (used by this task's own tests and by all later Part B fixtures,
+ e.g. proximity classification, BC-patch tagging) that builds a
+ synthetic multi-body STEP: import a base part, then place one OCC
+ primitive (box/cylinder) per attachment spec at one of the base's
+ bounding-box faces. NO boolean operations are performed -- the
+ attachment is written as an independent solid, deliberately overlapping
+ the base by `overlap_mm` (default 2.0mm) along the face normal. This is
+ intentional, not a bug: downstream, these bodies are never meshed
+ together (only the casting solid is meshed); the attachments exist
+ purely so their bounding boxes/faces can drive proximity/contact
+ classification against the casting, and a few mm of geometric overlap
+ is harmless for that purpose while guaranteeing the two solids actually
+ register as "touching" rather than floating a hair apart.
+
+All lengths in mm, all gmsh sessions are fresh (`gmsh.initialize()` /
+`gmsh.finalize()` in a try/finally per call) -- never share state across
+calls, matching the pattern in `tools/doe/prep_parts.py` /
+`src/step_mesh.py`.
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Sequence
+
+import numpy as np
+
+_AXIS_INDEX = {"x": 0, "y": 1, "z": 2}
+_VALID_SHAPES = {"box", "cylinder"}
+
+
+# --------------------------------------------------------------------------
+# Assembly inspection
+# --------------------------------------------------------------------------
+def list_solids(step_path: str) -> list[dict]:
+ """Import `step_path` in a fresh gmsh/OCC session and return one dict
+ per 3-D solid entity, sorted by `volume_mm3` descending (largest first
+ -- in practice the casting itself, since attached chills/sleeves are
+ normally much smaller):
+
+ {"solid_id": , "name": ,
+ "volume_mm3": , "bbox": [xmin, ymin, zmin, xmax, ymax, zmax],
+ "com": [cx, cy, cz]}
+
+ `Geometry.OCCImportLabels` MUST be set to 1 BEFORE `importShapes` --
+ that's the only way gmsh preserves STEP entity names (product names
+ from the CAD export) for `gmsh.model.getEntityName` to read; setting it
+ afterwards has no effect on an already-imported model. Solids with no
+ STEP name (or on any importer that doesn't preserve one) get `name=""`
+ rather than raising -- callers decide how to label an unnamed solid in
+ a role menu.
+ """
+ import gmsh
+
+ gmsh.initialize()
+ try:
+ gmsh.option.setNumber("General.Terminal", 0)
+ gmsh.option.setNumber("Geometry.OCCImportLabels", 1)
+ gmsh.model.add("assembly_inspect")
+ gmsh.model.occ.importShapes(step_path)
+ gmsh.model.occ.synchronize()
+
+ rows: list[dict] = []
+ for (_dim, tag) in gmsh.model.getEntities(3):
+ name = gmsh.model.getEntityName(3, tag) or ""
+ volume_mm3 = float(gmsh.model.occ.getMass(3, tag))
+ bbox = list(gmsh.model.getBoundingBox(3, tag))
+ com = list(gmsh.model.occ.getCenterOfMass(3, tag))
+ rows.append({
+ "solid_id": tag,
+ "name": name,
+ "volume_mm3": volume_mm3,
+ "bbox": bbox,
+ "com": com,
+ })
+ rows.sort(key=lambda r: r["volume_mm3"], reverse=True)
+ return rows
+ finally:
+ gmsh.finalize()
+
+
+# --------------------------------------------------------------------------
+# Assembly composer (test/fixture helper)
+# --------------------------------------------------------------------------
+def _parse_face(at_face: str) -> tuple[int, float]:
+ """`"+x"/"-x"/"+y"/"-y"/"+z"/"-z"` -> (axis_index, sign)."""
+ if len(at_face) != 2 or at_face[0] not in "+-" or at_face[1] not in _AXIS_INDEX:
+ raise ValueError(
+ f"invalid at_face {at_face!r}; expected one of +x -x +y -y +z -z")
+ return _AXIS_INDEX[at_face[1]], (1.0 if at_face[0] == "+" else -1.0)
+
+
+def _place_attachment(att: dict, bbox_min: Sequence[float], bbox_max: Sequence[float],
+ center: Sequence[float]) -> None:
+ """Add ONE primitive (box or cylinder) to the currently-open OCC model,
+ positioned at the base's bbox face named by `att["at_face"]`, sunk
+ `att.get("overlap_mm", 2.0)` INTO the base along the face normal (see
+ module docstring for why a positive overlap is intentional). Does not
+ synchronize/write -- the caller does that once after all attachments
+ are placed."""
+ import gmsh
+
+ shape = att.get("shape")
+ if shape not in _VALID_SHAPES:
+ raise ValueError(f"invalid shape {shape!r}; expected one of {sorted(_VALID_SHAPES)}")
+ if "size" not in att:
+ raise ValueError(f"attachment missing required 'size': {att!r}")
+ if "at_face" not in att:
+ raise ValueError(f"attachment missing required 'at_face': {att!r}")
+
+ axis, sign = _parse_face(att["at_face"])
+ overlap_mm = float(att.get("overlap_mm", 2.0))
+ du, dv = att.get("offset", (0.0, 0.0))
+ perp_axes = [i for i in range(3) if i != axis]
+ face_coord = bbox_max[axis] if sign > 0 else bbox_min[axis]
+
+ if shape == "box":
+ dx, dy, dz = att["size"]
+ size_by_axis = [dx, dy, dz]
+ perp_size = size_by_axis[axis]
+
+ if sign > 0:
+ lo_axis = face_coord - overlap_mm
+ else:
+ lo_axis = face_coord + overlap_mm - perp_size
+
+ coords = [0.0, 0.0, 0.0]
+ extents = [0.0, 0.0, 0.0]
+ coords[axis] = lo_axis
+ extents[axis] = perp_size
+ for i, u in zip(perp_axes, (du, dv)):
+ extents[i] = size_by_axis[i]
+ coords[i] = center[i] - size_by_axis[i] / 2.0 + u
+
+ gmsh.model.occ.addBox(coords[0], coords[1], coords[2],
+ extents[0], extents[1], extents[2])
+ else: # cylinder
+ radius, height = att["size"]
+ base_point = [0.0, 0.0, 0.0]
+ axis_vec = [0.0, 0.0, 0.0]
+ base_point[axis] = face_coord - sign * overlap_mm
+ axis_vec[axis] = sign * height
+ for i, u in zip(perp_axes, (du, dv)):
+ base_point[i] = center[i] + u
+
+ gmsh.model.occ.addCylinder(base_point[0], base_point[1], base_point[2],
+ axis_vec[0], axis_vec[1], axis_vec[2], radius)
+
+
+def compose_assembly(base_step: str, attachments: list[dict], out_step: str) -> int:
+ """Build a synthetic multi-body STEP: import `base_step`, then add one
+ OCC primitive per `attachments` entry (see module docstring for the
+ overlap/no-boolean-ops contract), and write the result (all solids,
+ unmerged) to `out_step`. Returns the number of 3-D solids written
+ (1 + len(attachments) in the normal case where `base_step` itself is
+ single-solid).
+
+ Each attachment dict:
+ {"shape": "box" | "cylinder",
+ "size": (dx, dy, dz) for a box, (radius, height) for a cylinder,
+ "at_face": "+x" | "-x" | "+y" | "-y" | "+z" | "-z",
+ "overlap_mm": float, default 2.0,
+ "offset": (du, dv) in the face plane, default (0.0, 0.0) (centered)}
+
+ For a box, `size[axis]` (the dimension along the face normal) is the
+ part that gets sunk `overlap_mm` into the base; the other two `size`
+ entries lie in the face plane. For a cylinder, the axis runs along the
+ face normal (`height`) and `radius` is in the face plane. `offset`
+ shifts the attachment's center within the face plane, in the same
+ (u, v) order as the two non-normal world axes (e.g. for "+z"/"-z",
+ (u, v) = (x, y); for "+x"/"-x", (u, v) = (y, z); for "+y"/"-y",
+ (u, v) = (x, z)).
+
+ NOTE for `tag_contact_patches` callers: a positive `overlap_mm` is a
+ fixed geometric gap that `tol_mm` must exceed just to detect the patch
+ at all, and any `tol_mm` beyond that dilates the classified patch past
+ the true footprint edge (see `tag_contact_patches`'s CAVEAT). A
+ synthetic fixture whose `overlap_mm` is large relative to its footprint
+ size will over-classify noticeably; real assemblies (near-zero contact
+ or a sub-mm clearance, not several mm sunk in) only need a small
+ `tol_mm`, where this dilation is negligible.
+ """
+ import gmsh
+
+ gmsh.initialize()
+ try:
+ gmsh.option.setNumber("General.Terminal", 0)
+ gmsh.model.add("assembly_compose")
+ gmsh.model.occ.importShapes(base_step)
+ gmsh.model.occ.synchronize()
+
+ xmin, ymin, zmin, xmax, ymax, zmax = gmsh.model.getBoundingBox(-1, -1)
+ bbox_min = [xmin, ymin, zmin]
+ bbox_max = [xmax, ymax, zmax]
+ center = [(bbox_min[i] + bbox_max[i]) / 2.0 for i in range(3)]
+
+ for att in attachments:
+ _place_attachment(att, bbox_min, bbox_max, center)
+
+ gmsh.model.occ.synchronize()
+ gmsh.write(out_step)
+ return len(gmsh.model.getEntities(3))
+ finally:
+ gmsh.finalize()
+
+
+# --------------------------------------------------------------------------
+# Body spec -- Part B Task 2: user-assigned solid roles -> the exact BC
+# parameters the v4 model was trained on (BcPlan), via the same registries
+# `tools/doe/doe.py` used to build the training manifest.
+# --------------------------------------------------------------------------
+_VALID_ROLES = {"casting", "chill", "insulating_sleeve", "exo_sleeve", "ignore"}
+_SLEEVE_ROLES = {"insulating_sleeve", "exo_sleeve"}
+_DEFAULT_CHILL_MATERIAL = "cast_iron"
+
+
+@dataclass
+class BcPlan:
+ """The resolved boundary-condition plan for one casting body, in the
+ exact shape `_node_features(version=4)` and `gbm_v2_total_seconds`
+ expect (see `feature_spec`/`gbm_bc_spec`). Produced by `validate_bodies`
+ -- never construct this by hand from unvalidated user input."""
+ bc_config: str # "none" | "chills" | "insulating_sleeve" | "exo_sleeve"
+ casting_solid_id: int
+ role_solid_ids: list[int] = field(default_factory=list) # bodies forming the patch
+ n_chills: int = 0
+ chill_effusivity: float = 0.0
+ sleeve_MEF: float = 0.0
+ exo_q0: float = 0.0 # W/m^2 (training units)
+ # Part-B roles task (2026-08-03): per-BODY role, `{solid_id: role}` for
+ # every id in `role_solid_ids` (e.g. all-"chill" for the chills branch,
+ # the single sleeve id -> its bc_config for the sleeve branch, empty for
+ # "none"). Lets a caller (e.g. `predict_casting`'s feed-node selection)
+ # decide per BODY ROLE rather than per global `bc_config` -- necessary
+ # groundwork for a future mixed-role regime even though `validate_bodies`
+ # itself still only ever produces a single-role mapping today.
+ role_by_solid: dict[int, str] = field(default_factory=dict)
+ # SP4 ADDENDUM (2026-08-08): several sleeve bodies are the norm on real
+ # castings. When they all resolve to one MEF it is used exactly and these
+ # stay at their defaults. When they differ, `sleeve_MEF` is the
+ # contact-area-weighted mean, `mef_weighted` is True, and `mef_spread`
+ # carries max(MEF) - min(MEF) so the report can state the spread the
+ # single scalar channel is standing in for rather than hiding it.
+ mef_weighted: bool = False
+ mef_spread: float = 0.0
+
+ def feature_spec(self, mold_effusivity: float) -> dict:
+ """The `spec` dict `src.features._node_features(version=4, spec=...)`
+ reads -- keys match `Global Constraints` in
+ `docs/superpowers/plans/2026-07-29-part-b-multibody.md` exactly."""
+ return {
+ "mold_effusivity": mold_effusivity,
+ "bc_config": self.bc_config,
+ "chill_effusivity": self.chill_effusivity,
+ "sleeve_MEF": self.sleeve_MEF,
+ "exo_q0": self.exo_q0,
+ }
+
+ def gbm_bc_spec(self) -> dict:
+ """The `bc_spec` dict `src.gbm_ft2b.gbm_v2_total_seconds` reads."""
+ return {
+ "bc_config": self.bc_config,
+ "n_chills": self.n_chills,
+ "sleeve_MEF": self.sleeve_MEF,
+ "exo_q0_MW": self.exo_q0 / 1e6,
+ }
+
+
+def _load_sleeve_registry(csv_path: str) -> dict[str, dict]:
+ """{sleeve name: {"MEF": float, "type": str}} from
+ `tools/sleeves/sleeve_registry.csv` -- `type` is "exo" or "insulating"
+ and is used to cross-check a registry-resolved sleeve's physics against
+ the declared role (see `validate_bodies`)."""
+ out: dict[str, dict] = {}
+ with open(csv_path, newline="") as f:
+ for row in csv.DictReader(f):
+ out[row["name"]] = {"MEF": float(row["MEF"]), "type": row["type"]}
+ return out
+
+
+# Declared role -> the sleeve_registry.csv `type` value it requires when the
+# sleeve is resolved BY REGISTRY NAME. Without this check, a registry
+# collision (two products landing on the same MEF, e.g. FOSECO_V88E
+# type=insulating MEF=1.45 vs the exo_calibration.csv row at MEF=1.45) lets
+# a plain insulating sleeve silently resolve a confident nonzero exo_q0 --
+# a silent-wrong-physics path, not merely a cosmetic mismatch.
+_ROLE_TO_SLEEVE_TYPE = {"exo_sleeve": "exo", "insulating_sleeve": "insulating"}
+
+# Two sleeves whose MEFs differ by less than this are treated as identical --
+# the registry stores MEF to 2 decimals, so anything below it is float noise,
+# not a real difference in sleeve product.
+_MEF_AGREEMENT_TOL = 1e-6
+
+
+def _resolve_chill_material(params: dict | None) -> str:
+ params = params or {}
+ return (params.get("material") or _DEFAULT_CHILL_MATERIAL).strip()
+
+
+def validate_bodies(bodies: list[dict],
+ chills_csv: str = "data/properties/chills.csv",
+ sleeve_registry_csv: str = "tools/sleeves/sleeve_registry.csv",
+ exo_calibration_csv: str = "data/properties/exo_calibration.csv",
+ tol: float = 0.05,
+ mef_weights: dict[int, float] | None = None) -> BcPlan:
+ """Validate a user-assigned role list and resolve it into a `BcPlan`.
+
+ `bodies` rows: `{"solid_id": int, "role": ,
+ "params": dict (optional, role-dependent)}`. `ignore` rows are dropped
+ before anything else. Enforces the training-regime restriction (owner,
+ 2026-07-29): EXACTLY one casting; then either >=1 `chill` bodies
+ (bc_config="chills") OR one or more sleeve bodies that share a single role
+ (`insulating_sleeve`/`exo_sleeve`, bc_config = that role; mixing the two
+ roles raises) OR neither (bc_config="none").
+ Chills+sleeve mixed, two sleeves, or a casting count != 1 all raise a
+ loud `ValueError` explaining the restriction -- the trained model has no
+ BC channel for anything else, so silently coercing one regime into
+ another would produce a confidently wrong prediction.
+
+ Chill params: `{"material": str}` (default "cast_iron"); an unknown
+ material raises `ValueError` listing `load_chill_effusivities()`'s known
+ materials. Multiple chills MUST share one material -- the trained model
+ has a single scalar `chill_effusivity` channel, not one per chill, so
+ differing materials raise `ValueError` (single chill_effusivity per
+ training convention).
+
+ Sleeve params: exactly one of `{"sleeve": }`
+ (MEF looked up in the registry) or `{"MEF": float}` (manual override) --
+ both or neither present raises `ValueError`. An unknown registry name
+ raises `ValueError` listing the known names. When resolved BY NAME, the
+ registry row's `type` column ("exo" or "insulating") MUST match the
+ declared role (`exo_sleeve` requires type=="exo", `insulating_sleeve`
+ requires type=="insulating") -- a mismatch raises `ValueError` naming
+ both, since two products can share the same MEF (a registry collision)
+ while being physically different (an insulating sleeve has no
+ exothermic heat release), and resolving one under the other's role
+ would silently produce a confidently wrong prediction. Manual
+ `{"MEF": float}` overrides BYPASS this check entirely (no registry row
+ to check against) -- this is expert mode: the caller is asserting the
+ physics directly and is trusted to pick the right role. For `bc_config
+ == "exo_sleeve"`, the resolved MEF is turned into a calibrated `exo_q0`
+ via `tools.doe.load_properties.resolve_exo_q0_by_mef` (same nearest-MEF
+ lookup `src.build_dataset._resolve_exo_q0_v4` uses for training data --
+ raises `ValueError` if no calibration point is within `tol`).
+
+ Several sleeves (SP4 ADDENDUM, 2026-08-08): real castings carry 3-8
+ sleeves, so N bodies of one role are accepted, mirroring the chills
+ branch. When every sleeve resolves to the same MEF it is used exactly
+ (`mef_weighted=False`, `mef_spread=0.0`). When they differ -- e.g. an
+ assembly with two sleeve sizes -- the model's single scalar channel
+ cannot hold both, so `mef_weights` (contact-patch area per solid_id,
+ exactly `tag_contact_patches`'s `per_body_patch_area_mm2`) MUST be
+ supplied and the effective MEF is their area-weighted mean, with
+ `mef_weighted=True` and `mef_spread` recording what was approximated.
+ Omitting the weights on a mixed set raises rather than silently picking
+ one sleeve's MEF or an unweighted mean.
+ """
+ from tools.doe.load_properties import (load_chill_effusivities,
+ resolve_exo_q0_by_mef)
+
+ active = [b for b in bodies if (b.get("role") or "") != "ignore"]
+
+ for b in active:
+ role = b.get("role")
+ if role not in _VALID_ROLES:
+ raise ValueError(
+ f"body solid_id={b.get('solid_id')!r}: unknown role {role!r} "
+ f"(expected one of {sorted(_VALID_ROLES)})")
+
+ castings = [b for b in active if b["role"] == "casting"]
+ chills = [b for b in active if b["role"] == "chill"]
+ sleeves = [b for b in active if b["role"] in _SLEEVE_ROLES]
+
+ if len(castings) != 1:
+ raise ValueError(
+ f"training-regime restriction: expected exactly one body with "
+ f"role='casting', found {len(castings)}")
+ casting_solid_id = castings[0]["solid_id"]
+
+ if chills and sleeves:
+ raise ValueError(
+ "training-regime restriction: chills and a sleeve cannot be "
+ "combined -- the trained model has no mixed chills+sleeve BC "
+ "channel in the current training regimes")
+ if sleeves:
+ sleeve_roles = {b["role"] for b in sleeves}
+ if len(sleeve_roles) > 1:
+ raise ValueError(
+ f"training-regime restriction: all sleeve bodies must share "
+ f"one role, got {sorted(sleeve_roles)} -- the trained model "
+ f"has a single sleeve BC channel and cannot represent an "
+ f"insulating and an exothermic sleeve at once")
+
+ if chills:
+ materials = {_resolve_chill_material(b.get("params")) for b in chills}
+ effs = load_chill_effusivities(chills_csv)
+ for material in materials:
+ if material not in effs:
+ raise ValueError(
+ f"unknown chill material {material!r} (known: "
+ f"{sorted(effs)})")
+ if len(materials) > 1:
+ raise ValueError(
+ f"training-regime restriction: all chills must share one "
+ f"material (single chill_effusivity per training "
+ f"convention), got {sorted(materials)}")
+ material = next(iter(materials))
+ chill_ids = [b["solid_id"] for b in chills]
+ return BcPlan(
+ bc_config="chills",
+ casting_solid_id=casting_solid_id,
+ role_solid_ids=chill_ids,
+ n_chills=len(chills),
+ chill_effusivity=float(effs[material]),
+ role_by_solid={sid: "chill" for sid in chill_ids},
+ )
+
+ if sleeves:
+ bc_config = sleeves[0]["role"]
+ registry = None
+ mef_by_solid: dict[int, float] = {}
+ for sleeve in sleeves:
+ params = sleeve.get("params") or {}
+ has_name = "sleeve" in params
+ has_mef = "MEF" in params
+ if has_name and has_mef:
+ raise ValueError(
+ f"sleeve solid_id={sleeve['solid_id']!r}: specify either "
+ f"'sleeve' (registry name) or 'MEF' (manual override), not "
+ f"both")
+ if not has_name and not has_mef:
+ raise ValueError(
+ f"sleeve solid_id={sleeve['solid_id']!r}: params must "
+ f"contain either 'sleeve' (registry name) or 'MEF' (manual "
+ f"override)")
+ if has_name:
+ if registry is None:
+ registry = _load_sleeve_registry(sleeve_registry_csv)
+ name = params["sleeve"]
+ if name not in registry:
+ raise ValueError(
+ f"unknown sleeve name {name!r} (known: "
+ f"{sorted(registry)})")
+ entry = registry[name]
+ expected_type = _ROLE_TO_SLEEVE_TYPE[bc_config]
+ if entry["type"] != expected_type:
+ raise ValueError(
+ f"sleeve {name!r} is type={entry['type']} but was "
+ f"assigned role {bc_config} (expected "
+ f"type={expected_type!r})")
+ mef_by_solid[sleeve["solid_id"]] = float(entry["MEF"])
+ else:
+ mef_by_solid[sleeve["solid_id"]] = float(params["MEF"])
+
+ sleeve_ids = [b["solid_id"] for b in sleeves]
+ mefs = [mef_by_solid[sid] for sid in sleeve_ids]
+ spread = float(max(mefs) - min(mefs))
+ if spread <= _MEF_AGREEMENT_TOL:
+ sleeve_mef = float(mefs[0])
+ weighted = False
+ spread = 0.0
+ else:
+ if mef_weights is None:
+ raise ValueError(
+ f"sleeve bodies {sleeve_ids} resolve to differing MEF "
+ f"{sorted(set(mefs))} (spread {spread:.4f}) -- the model has "
+ f"a single scalar sleeve_MEF channel, so a mixed set needs "
+ f"explicit `mef_weights` (contact-patch area per solid_id, "
+ f"as returned by `tag_contact_patches`'s "
+ f"'per_body_patch_area_mm2'). Refusing to guess a "
+ f"weighting. Contact-patch areas only exist after meshing "
+ f"and tagging, so a caller that only has `bodies` (e.g. "
+ f"`predict_casting`'s CLI/webapp entry point) cannot supply "
+ f"them here -- mesh and tag the part first (`get_or_mesh` "
+ f"+ `tag_contact_patches`), then call `validate_bodies` "
+ f"again with `mef_weights=per_body_patch_area_mm2`, as "
+ f"`tools/sp4/prep.py` and `tools/sp4/score.py` do for the "
+ f"SP4 gold set.")
+ missing = [sid for sid in sleeve_ids if sid not in mef_weights]
+ if missing:
+ raise ValueError(
+ f"mef_weights is missing solid_id(s) {missing} -- every "
+ f"sleeve body needs a weight, got keys "
+ f"{sorted(mef_weights)}")
+ weights = [float(mef_weights[sid]) for sid in sleeve_ids]
+ if any(w <= 0.0 for w in weights):
+ raise ValueError(
+ f"mef_weights must all be positive contact areas, got "
+ f"{dict(zip(sleeve_ids, weights))}")
+ total_w = sum(weights)
+ sleeve_mef = float(sum(m * w for m, w in zip(mefs, weights)) / total_w)
+ weighted = True
+
+ exo_q0 = 0.0
+ if bc_config == "exo_sleeve":
+ exo_q0 = resolve_exo_q0_by_mef(sleeve_mef, exo_calibration_csv, tol)
+
+ return BcPlan(
+ bc_config=bc_config,
+ casting_solid_id=casting_solid_id,
+ role_solid_ids=sleeve_ids,
+ sleeve_MEF=sleeve_mef,
+ exo_q0=exo_q0,
+ role_by_solid={sid: bc_config for sid in sleeve_ids},
+ mef_weighted=weighted,
+ mef_spread=spread,
+ )
+
+ return BcPlan(bc_config="none", casting_solid_id=casting_solid_id, role_solid_ids=[])
+
+
+# --------------------------------------------------------------------------
+# Contact-patch tagging -- Part B Task 3: mesh the user-designated casting
+# solid, classify its surface triangles by proximity to the role bodies'
+# BReps, and rewrite the mesh's facet groups so the contact patch becomes
+# "Adiabatic" (the exact BC-patch convention `_node_features(version=4)` /
+# `fenicsx_solver.solve_case` read -- see the module docstring's B3 design
+# reference).
+# --------------------------------------------------------------------------
+_ADIABATIC_PHYS_TAG = 3 # alongside prep_parts._VOLUME_PHYS_TAG=1 / _COOLING_PHYS_TAG=2
+
+
+def _triangle_areas(node_xyz: np.ndarray, tris: np.ndarray) -> np.ndarray:
+ """(M,) triangle areas (mm^2) via the standard half cross-product
+ magnitude -- `tris` is (M, 3) 0-based vertex indices into `node_xyz`."""
+ p0 = node_xyz[tris[:, 0]]
+ p1 = node_xyz[tris[:, 1]]
+ p2 = node_xyz[tris[:, 2]]
+ cross = np.cross(p1 - p0, p2 - p0)
+ return 0.5 * np.linalg.norm(cross, axis=1)
+
+
+def _sample_solid_surface_points(solid_id: int) -> np.ndarray:
+ """Assumes a fresh gmsh/OCC session already has the model imported,
+ synchronized, AND 2-D-meshed (`gmsh.model.mesh.generate(2)` already ran
+ -- done ONCE per `tag_contact_patches` call, not per body, since it
+ covers every solid's boundary surfaces in one pass). Returns (P, 3) the
+ coordinates of every mesh node lying on any boundary surface
+ (`gmsh.model.getBoundary`) of 3-D entity `solid_id`."""
+ import gmsh
+
+ boundary = gmsh.model.getBoundary([(3, solid_id)], oriented=False)
+ surf_tags = [t for (d, t) in boundary if d == 2]
+ if not surf_tags:
+ raise ValueError(f"role body {solid_id}: no boundary surfaces found (not a 3-D solid?)")
+
+ pts = []
+ for surf_tag in surf_tags:
+ _node_tags, coords, _param = gmsh.model.mesh.getNodes(2, surf_tag, includeBoundary=True)
+ pts.append(np.asarray(coords, dtype=np.float64).reshape(-1, 3))
+ points = np.concatenate(pts, axis=0)
+ if points.shape[0] == 0:
+ raise ValueError(f"role body {solid_id}: surface sampling yielded 0 points")
+ return points
+
+
+def _distances_to_solid(points: np.ndarray, centroids: np.ndarray) -> np.ndarray:
+ """(M,) nearest-neighbor distance from each row of `centroids` (M, 3, mm)
+ to the sampled boundary point cloud `points` (P, 3, mm) of one role
+ solid, via `scipy.spatial.cKDTree`.
+
+ WHY sampled points + KDTree, not `gmsh.model.getClosestPoint` directly
+ on the centroids: probed first (as the brief suggested trying), and
+ found to be not just slower but WRONG for axis-aligned primitive faces
+ (the chill/sleeve boxes and cylinders this module builds and that real
+ role bodies resemble) -- `getClosestPoint`'s orthogonal projection lands
+ on the face's UNDERLYING (untrimmed) surface, not the actual bounded
+ patch (`isInside()` on the returned point came back False in every case
+ probed), so a query point far outside a face's true extent can still
+ project to a deceptively small distance whenever it happens to share
+ that face's plane (e.g. a base box's OWN top face reported ~2mm to a
+ chill's bottom-face plane from a corner ~15mm from the chill's actual
+ footprint). Meshing the role solids' real (trimmed) boundary at a size
+ tied to `tol_mm` and taking nearest-neighbor distance to those sampled
+ points respects the true finite geometry and reproduces the expected
+ footprint exactly in this task's synthetic fixtures (measured
+ corner-case distance corrected from a wrong ~2mm to the correct
+ ~11.2mm). Timing (composer fixture, `Mesh.MeshSizeMax = tol_mm/2`,
+ 2-D-meshing the WHOLE imported model once): 0.16-0.9s total for
+ 408-11019-node castings (656-8328 Cooling triangles) -- see
+ task-3-report.md for the full measurement."""
+ from scipy.spatial import cKDTree
+
+ tree = cKDTree(points)
+ dist, _ = tree.query(centroids)
+ return np.asarray(dist, dtype=np.float64)
+
+
+def _parse_physical_names_block(lines: list[str], start_idx: int):
+ """`lines[start_idx]` must be the bare `"$PhysicalNames"` marker line.
+ Returns `(count_idx, entries, end_idx)`: `count_idx` is the index of the
+ entry-count line, `entries` is a list of `(dim, tag, name)` (name
+ unquoted), and `end_idx` is the index of the `"$EndPhysicalNames"`
+ line."""
+ count_idx = start_idx + 1
+ cnt = int(lines[count_idx].strip())
+ entries = []
+ for k in range(cnt):
+ dim_s, tag_s, name_s = lines[count_idx + 1 + k].split(maxsplit=2)
+ entries.append((int(dim_s), int(tag_s), name_s.strip().strip('"')))
+ end_idx = count_idx + 1 + cnt
+ return count_idx, entries, end_idx
+
+
+def rewrite_msh_adiabatic(msh_path: str, is_adiabatic: np.ndarray) -> None:
+ """Rewrite `msh_path` in place so that the Cooling-tagged surface
+ triangles flagged True in `is_adiabatic` (same length AND order as
+ `src.msh_parse.parse_msh(msh_path).tris["Cooling"]` -- both this
+ function's own `$Elements` scan and `parse_msh` walk the file's
+ triangle elements in on-disk order, so the two arrays line up
+ positionally) move to a new "Adiabatic" physical group (pinned tag
+ `_ADIABATIC_PHYS_TAG`, reusing an existing "Adiabatic" group's tag if
+ the mesh already has one).
+
+ Deliberately edits the raw text rather than reimporting through gmsh:
+ `$Nodes` is never touched and every `$Elements` line is copied through
+ unchanged except the physical-group tag (first tag) of the flagged
+ triangle lines -- this is what guarantees node ids/coords and tet
+ connectivity survive EXACTLY (features/geodesics depend on it; verified
+ by the caller re-parsing and comparing `node_xyz`/`tets`). Raises
+ loudly (`ValueError`/`AssertionError`) on any structural surprise --
+ missing sections, a malformed `$PhysicalNames` block, tag-3 already
+ bound to a different name, or an `is_adiabatic` length that doesn't
+ match the number of Cooling-tagged triangle elements actually found --
+ rather than silently writing a mismatched mesh."""
+ from src.msh_parse import _TRI_TYPE
+
+ lines = Path(msh_path).read_text().splitlines()
+
+ try:
+ phn_idx = lines.index("$PhysicalNames")
+ except ValueError:
+ raise ValueError(f"{msh_path}: missing $PhysicalNames section")
+ count_idx, entries, end_idx = _parse_physical_names_block(lines, phn_idx)
+ if lines[end_idx].strip() != "$EndPhysicalNames":
+ raise ValueError(f"{msh_path}: malformed $PhysicalNames block (no $EndPhysicalNames "
+ f"at the expected offset)")
+
+ cooling_entries = [e for e in entries if e[0] == 2 and e[2] == "Cooling"]
+ if not cooling_entries:
+ raise ValueError(f"{msh_path}: no dim=2 'Cooling' physical group found")
+ cooling_tag = cooling_entries[0][1]
+
+ existing_adiabatic = [e for e in entries if e[0] == 2 and e[2] == "Adiabatic"]
+ if existing_adiabatic:
+ adiabatic_tag = existing_adiabatic[0][1]
+ new_entries = entries
+ else:
+ collision = [e for e in entries if e[0] == 2 and e[1] == _ADIABATIC_PHYS_TAG]
+ if collision:
+ raise ValueError(f"{msh_path}: physical tag {_ADIABATIC_PHYS_TAG} is already "
+ f"bound to {collision[0][2]!r}; cannot pin 'Adiabatic' there")
+ adiabatic_tag = _ADIABATIC_PHYS_TAG
+ new_entries = entries + [(2, adiabatic_tag, "Adiabatic")]
+
+ new_phys_lines = [str(len(new_entries))] + [
+ f'{d} {t} "{nm}"' for (d, t, nm) in new_entries
+ ]
+
+ try:
+ elem_idx = lines.index("$Elements")
+ except ValueError:
+ raise ValueError(f"{msh_path}: missing $Elements section")
+ elem_count_idx = elem_idx + 1
+ elem_cnt = int(lines[elem_count_idx].strip())
+ elem_lines = lines[elem_count_idx + 1: elem_count_idx + 1 + elem_cnt]
+
+ flags = is_adiabatic.tolist()
+ n_flags = len(flags)
+ flag_i = 0
+ new_elem_lines = []
+ for ln in elem_lines:
+ parts = ln.split()
+ etype = int(parts[1])
+ if etype == _TRI_TYPE and int(parts[3]) == cooling_tag:
+ if flag_i >= n_flags:
+ raise AssertionError(
+ f"{msh_path}: is_adiabatic has only {n_flags} flag(s) -- fewer than the "
+ f"Cooling-tagged triangle elements found in $Elements -- refusing to write "
+ f"a mismatched mesh")
+ if flags[flag_i]:
+ parts[3] = str(adiabatic_tag)
+ ln = " ".join(parts)
+ flag_i += 1
+ new_elem_lines.append(ln)
+ if flag_i != n_flags:
+ raise AssertionError(
+ f"{msh_path}: is_adiabatic has {n_flags - flag_i} more flag(s) than Cooling-tagged "
+ f"triangle elements found in $Elements -- refusing to write a mismatched mesh")
+
+ new_lines = (
+ lines[:count_idx] + new_phys_lines + lines[end_idx:elem_count_idx]
+ + [str(elem_cnt)] + new_elem_lines + lines[elem_count_idx + 1 + elem_cnt:]
+ )
+ Path(msh_path).write_text("\n".join(new_lines) + "\n")
+
+
+# Back-compat alias -- this function was originally private (`_rewrite_msh_
+# adiabatic`); FT3's `tools/doe/virtual_patches.py` reuses it under the
+# promoted public name `rewrite_msh_adiabatic`, but existing call sites/tests
+# (this module's own `tag_contact_patches`, `tests/test_assembly.py`) still
+# reference the old private name, so both names resolve to the same function.
+_rewrite_msh_adiabatic = rewrite_msh_adiabatic
+
+
+def tag_contact_patches(msh_path: str, step_path: str, role_solid_ids: Sequence[int],
+ tol_mm: float | None = None) -> dict:
+ """Classify the casting mesh's (`msh_path`) "Cooling" surface triangles
+ by proximity to each of `role_solid_ids`'s BReps in `step_path` (a fresh
+ gmsh/OCC session, distinct from the casting mesh's own session --
+ `step_path` is the ORIGINAL multi-body assembly STEP the role bodies'
+ tags came from, e.g. `list_solids`' `solid_id`, not the reduced
+ single-solid STEP the casting was meshed from); triangles within
+ `tol_mm` of ANY role body move from "Cooling" to a new "Adiabatic"
+ facet group (see `_rewrite_msh_adiabatic`), rewriting `msh_path` IN
+ PLACE. Distance-to-body samples each role solid's real (trimmed)
+ boundary surfaces into a point cloud (one shared `gmsh.model.mesh.
+ generate(2)` pass over the whole imported model, sized off `tol_mm`)
+ and takes the nearest-neighbor distance via `scipy.spatial.cKDTree`
+ (`_sample_solid_surface_points` + `_distances_to_solid`) -- see
+ `_distances_to_solid`'s docstring for why a direct
+ `gmsh.model.getClosestPoint` probe was rejected (not just slower, but
+ wrong for bounded planar faces).
+
+ `tol_mm`, if `None` (default), is `0.5 * elem_est` where `elem_est =
+ bbox_diag / n_nodes**(1/3)` of the casting mesh (all nodes, matching
+ `heal_and_mesh`'s own whole-model bbox convention) -- a rough local
+ element-size estimate, so the tolerance scales with mesh density rather
+ than being a fixed physical distance that's too tight for a coarse mesh
+ or too loose for a fine one.
+
+ CAVEAT -- tol_mm dilates the detected patch past the true contact
+ footprint: since classification is "distance to body <= tol_mm", any
+ positive tol_mm grows the patch outward from the geometric contact edge
+ by roughly `perimeter * tol_mm + pi * tol_mm**2` (a rectangular-
+ footprint Minkowski-sum estimate). This matters most for the SYNTHETIC
+ fixtures `compose_assembly` builds with a large POSITIVE `overlap_mm`
+ (sunk INTO the base -- see that function's docstring): `tol_mm` must
+ exceed `overlap_mm` just to detect the patch at all, and when
+ `overlap_mm` is a large fraction of the footprint size, the resulting
+ dilation can be a large fraction of the footprint's own area (measured
+ up to +140% on a 10x10mm footprint with a 2mm overlap -- see
+ task-3-report.md). Real assemblies -- where role bodies sit at
+ near-zero contact or a sub-mm clearance rather than several mm sunk in
+ -- only need a small `tol_mm` to bridge that gap, so this dilation
+ stays a small fraction of a realistic contact-patch area; it is a
+ synthetic-fixture-scale concern, not a production one.
+
+ A role body whose patch comes out EMPTY (no Cooling triangle centroid
+ within `tol_mm`) raises `ValueError` naming that body and `tol_mm` --
+ never a silent all-sand prediction for a body the caller explicitly
+ assigned a role to. Symmetrically, if EVERY Cooling triangle ends up
+ classified Adiabatic (role bodies + tol_mm cover the entire casting
+ surface), this raises `ValueError` too, rather than silently handing
+ downstream feature code (`geodesic_d_cool` et al., see
+ `src.features`) a casting mesh with no cooling surface at all to
+ measure geodesic distance to -- which would silently degrade to
+ bbox_diag-derived constants (confidently wrong physics, not a loud
+ failure).
+
+ Returns `{"per_body_patch_area_mm2": {solid_id: float}, "n_adiab_tris":
+ int, "nodes_by_solid": {solid_id: sorted list[int]}}`.
+ `per_body_patch_area_mm2` sums each body's own classified
+ triangle areas INDEPENDENTLY -- a triangle within `tol_mm` of more than
+ one role body (their tolerance zones overlap) is counted once per such
+ body, i.e. DOUBLE-COUNTED across `per_body_patch_area_mm2`'s entries,
+ even though it is relabeled to "Adiabatic" exactly once; `n_adiab_tris`
+ counts each such triangle once (via the union `is_adiabatic` mask), so
+ `sum(per_body_patch_area_mm2.values())` can exceed the true total
+ Adiabatic area when role bodies' patches overlap.
+
+ `nodes_by_solid` (Part-B roles task, 2026-08-03): the sorted-unique node
+ ids (0-based, into `mesh.node_xyz` -- the SAME indexing as every other
+ node id this module and `src.msh_parse` use) of each role body's own
+ classified triangles -- computed from the SAME per-body `close` mask as
+ `per_body_patch_area_mm2`, so it shares its double-counting property (a
+ triangle within `tol_mm` of more than one body contributes its nodes to
+ EVERY such body's entry). This is an IN-MEMORY-only piece of metadata:
+ it is never written into `msh_path` itself (the `.msh` physical groups,
+ node order, and this function's cache-relevant hashes/tags are all
+ unchanged by its addition) -- callers needing a role-aware node
+ selection (e.g. `predict_casting`'s shrinkage-screening feed-node logic)
+ combine this per-SOLID mapping with their own solid_id -> role lookup
+ (e.g. `BcPlan.role_by_solid`) themselves; `tag_contact_patches` has no
+ opinion on role semantics, only geometry. After rewriting,
+ re-parses `msh_path` and asserts `node_xyz` is unchanged (`np.allclose`)
+ and `tets` is unchanged (`np.array_equal`) -- the node-order/tet-
+ connectivity invariant features/geodesics depend on (see the module
+ docstring's caution and `_rewrite_msh_adiabatic`)."""
+ from src.msh_parse import parse_msh
+ import gmsh
+
+ mesh = parse_msh(msh_path)
+ cooling_tris = mesh.tris.get("Cooling")
+ if cooling_tris is None or cooling_tris.shape[0] == 0:
+ raise ValueError(f"{msh_path}: no 'Cooling' surface triangles to classify")
+
+ centroids = mesh.node_xyz[cooling_tris].mean(axis=1)
+ tri_areas = _triangle_areas(mesh.node_xyz, cooling_tris)
+
+ if tol_mm is None:
+ bbox_min = mesh.node_xyz.min(axis=0)
+ bbox_max = mesh.node_xyz.max(axis=0)
+ bbox_diag_mm = float(np.linalg.norm(bbox_max - bbox_min))
+ n_nodes = mesh.node_xyz.shape[0]
+ elem_est = bbox_diag_mm / (n_nodes ** (1.0 / 3.0))
+ tol_mm = 0.5 * elem_est
+
+ gmsh.initialize()
+ try:
+ gmsh.option.setNumber("General.Terminal", 0)
+ gmsh.model.add("tag_contact_patches")
+ gmsh.model.occ.importShapes(step_path)
+ gmsh.model.occ.synchronize()
+
+ # One shared 2-D surface mesh (not per-body) -- covers every role
+ # solid's boundary in a single pass. Sized off tol_mm/2 so sampled
+ # points are dense enough relative to the tolerance to not miss or
+ # blur the true footprint.
+ sample_size_mm = max(tol_mm / 2.0, 1e-3)
+ gmsh.option.setNumber("Mesh.MeshSizeMin", sample_size_mm)
+ gmsh.option.setNumber("Mesh.MeshSizeMax", sample_size_mm)
+ gmsh.model.mesh.generate(2)
+
+ dist_by_body = {}
+ for solid_id in role_solid_ids:
+ points = _sample_solid_surface_points(solid_id)
+ dist_by_body[solid_id] = _distances_to_solid(points, centroids)
+ finally:
+ gmsh.finalize()
+
+ is_adiabatic = np.zeros(cooling_tris.shape[0], dtype=bool)
+ per_body_patch_area_mm2: dict[int, float] = {}
+ nodes_by_solid: dict[int, list[int]] = {}
+ for solid_id, dist in dist_by_body.items():
+ close = dist <= tol_mm
+ area = float(tri_areas[close].sum())
+ if area <= 0.0:
+ raise ValueError(f"body {solid_id} does not touch the casting within "
+ f"tol={tol_mm:.4f}mm")
+ per_body_patch_area_mm2[solid_id] = area
+ nodes_by_solid[solid_id] = sorted(
+ int(n) for n in np.unique(cooling_tris[close].reshape(-1)))
+ is_adiabatic |= close
+
+ if bool(is_adiabatic.all()):
+ raise ValueError(
+ "role bodies cover the entire casting surface; no cooling surface remains "
+ f"-- check roles/tol (tol={tol_mm:.4f}mm)")
+
+ rewrite_msh_adiabatic(msh_path, is_adiabatic)
+
+ mesh_after = parse_msh(msh_path)
+ if not np.allclose(mesh_after.node_xyz, mesh.node_xyz):
+ raise AssertionError(f"{msh_path}: node_xyz changed after Adiabatic rewrite "
+ f"(node order/coords must be byte-identical)")
+ if mesh_after.tets.shape != mesh.tets.shape or not np.array_equal(mesh_after.tets, mesh.tets):
+ raise AssertionError(f"{msh_path}: tets changed after Adiabatic rewrite "
+ f"(tet connectivity must be byte-identical)")
+
+ return {"per_body_patch_area_mm2": per_body_patch_area_mm2,
+ "n_adiab_tris": int(is_adiabatic.sum()),
+ "nodes_by_solid": nodes_by_solid}
+
+
+# --------------------------------------------------------------------------
+# CLI printer
+# --------------------------------------------------------------------------
+def main(argv: list[str] | None = None) -> int:
+ ap = argparse.ArgumentParser(
+ description="List the 3-D solids in a (possibly multi-body) STEP file.")
+ ap.add_argument("step_path", help="path to a .step/.stp file")
+ args = ap.parse_args(argv)
+
+ rows = list_solids(args.step_path)
+ header = f"{'solid_id':>8} {'name':<24} {'volume_mm3':>14} {'bbox (mm)':<48} {'com (mm)':<30}"
+ print(header)
+ print("-" * len(header))
+ for row in rows:
+ bbox_str = "[" + ", ".join(f"{v:.1f}" for v in row["bbox"]) + "]"
+ com_str = "[" + ", ".join(f"{v:.1f}" for v in row["com"]) + "]"
+ name = row["name"] or ""
+ print(f"{row['solid_id']:>8} {name:<24} {row['volume_mm3']:>14.2f} "
+ f"{bbox_str:<48} {com_str:<30}")
+ print(f"\n{len(rows)} solid(s) total.")
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/src/baseline.py b/src/baseline.py
new file mode 100644
index 0000000000000000000000000000000000000000..a15c1ab35101da3d9c0f8d28d510aad4445713bf
--- /dev/null
+++ b/src/baseline.py
@@ -0,0 +1,63 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..895381f207747ae31580539f6f4390075843bf68
--- /dev/null
+++ b/src/build_dataset.py
@@ -0,0 +1,1199 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..000907b5e2ebddc7f97e436ed543abbaa8dfe70e
--- /dev/null
+++ b/src/compare_d_parts.py
@@ -0,0 +1,214 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..20ca33ff98620ee36f39fdc3dce8582de10be7d2
--- /dev/null
+++ b/src/dataload.py
@@ -0,0 +1,28 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..2f100d078d324580b8f0d3bb9493f1ef406dbf9f
--- /dev/null
+++ b/src/evaluate_ft2b.py
@@ -0,0 +1,630 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..892d722986e062eb2bda5b543161861da5c96233
--- /dev/null
+++ b/src/evaluate_gnn.py
@@ -0,0 +1,150 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..e4d5a6359bc0930f0a3578e43bacbacc14cfb6c9
--- /dev/null
+++ b/src/features.py
@@ -0,0 +1,334 @@
+"""Assemble per-node physics features (v2) + target for one casting.
+
+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.
+"""
+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__)
+
+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",
+]
+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).
+
+ `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"])
+ L = float(props["L_kJ_kg"]); TL = float(props["TL"]); TS = float(props["TS"])
+ Tpour = float(props["Tpour"]); b_mold = float(props["mold_effusivity"])
+
+ Ste = cp * (Tpour - TS) / (L * 1000.0)
+ dTf_norm = (TL - TS) / T_REF
+ alpha_ratio = (k / (rho * cp)) / ALPHA_REF
+ 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.
+
+ `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."""
+ xyz = mesh.node_xyz
+ N = xyz.shape[0]
+ cool_tri = mesh.tris.get("Cooling", np.empty((0, 3), np.int64))
+ adiab_tri = mesh.tris.get("Adiabatic", np.empty((0, 3), np.int64))
+
+ cool_mask = boundary_node_mask(N, cool_tri)
+ adiab_mask = boundary_node_mask(N, adiab_tri)
+ adiab_only = adiab_mask & ~cool_mask
+ surf_mask = cool_mask | adiab_mask
+ interior_mask = ~surf_mask
+
+ bbox = xyz.max(0) - xyz.min(0)
+ bbox_diag = float(np.linalg.norm(bbox))
+
+ edges = edges_from_tets(mesh.tets)
+ edge_len = np.linalg.norm(xyz[edges[:, 0]] - xyz[edges[:, 1]], axis=1) if edges.shape[0] else np.empty(0)
+ src = np.flatnonzero(cool_mask)
+ g_cool = geodesic_distance(N, edges, edge_len, src)
+ eucl_cool = nearest_distance(xyz, xyz[cool_mask])
+ eucl_cool = np.where(np.isinf(eucl_cool), bbox_diag, eucl_cool)
+ g_cool = np.where(np.isinf(g_cool), eucl_cool, g_cool)
+
+ d_surf = nearest_distance(xyz, xyz[surf_mask])
+ d_adiab = nearest_distance(xyz, xyz[adiab_mask])
+ d_surf = np.where(np.isinf(d_surf), bbox_diag, d_surf)
+ d_adiab = np.where(np.isinf(d_adiab), bbox_diag, d_adiab)
+
+ area_cool = total_area(xyz, cool_tri)
+ area_total = area_cool + total_area(xyz, adiab_tri)
+ cool_frac = area_cool / area_total if area_total > 0 else 0.0
+ volume = mesh_volume(xyz, mesh.tets)
+
+ cols = {
+ "geodesic_d_cool": g_cool, "d_surf": d_surf, "d_adiab": d_adiab,
+ "is_cooling": cool_mask.astype(np.float32),
+ "is_adiabatic": adiab_only.astype(np.float32),
+ "is_interior": interior_mask.astype(np.float32),
+ "volume_mm3": np.full(N, volume),
+ "surface_area_mm2": np.full(N, area_total),
+ "cooling_area_frac": np.full(N, cool_frac),
+ }
+ 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
diff --git a/src/gbm_ft2b.py b/src/gbm_ft2b.py
new file mode 100644
index 0000000000000000000000000000000000000000..09a14e141cf33c122dbbf8bb2c1ed039d6e27af9
--- /dev/null
+++ b/src/gbm_ft2b.py
@@ -0,0 +1,598 @@
+"""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/geometry.py b/src/geometry.py
new file mode 100644
index 0000000000000000000000000000000000000000..665cf90c4337d98bffee180d64ff693770d2beeb
--- /dev/null
+++ b/src/geometry.py
@@ -0,0 +1,252 @@
+"""Pure-geometry helpers: triangle areas, tet volumes, nearest-surface
+distances (via KD-tree), and boundary-node masks. All lengths in mm."""
+import numpy as np
+from scipy.spatial import cKDTree
+from scipy.sparse import coo_matrix
+from scipy.sparse.csgraph import dijkstra
+
+
+def triangle_areas(node_xyz: np.ndarray, tri: np.ndarray) -> np.ndarray:
+ if tri.shape[0] == 0:
+ return np.zeros(0)
+ a = node_xyz[tri[:, 0]]
+ b = node_xyz[tri[:, 1]]
+ c = node_xyz[tri[:, 2]]
+ cross = np.cross(b - a, c - a)
+ return 0.5 * np.linalg.norm(cross, axis=1)
+
+
+def total_area(node_xyz: np.ndarray, tri: np.ndarray) -> float:
+ return float(triangle_areas(node_xyz, tri).sum())
+
+
+def tet_volumes(node_xyz: np.ndarray, tets: np.ndarray) -> np.ndarray:
+ if tets.shape[0] == 0:
+ return np.zeros(0)
+ a = node_xyz[tets[:, 0]]
+ b = node_xyz[tets[:, 1]]
+ c = node_xyz[tets[:, 2]]
+ d = node_xyz[tets[:, 3]]
+ return np.abs(np.einsum('ij,ij->i', b - a, np.cross(c - a, d - a))) / 6.0
+
+
+def mesh_volume(node_xyz: np.ndarray, tets: np.ndarray) -> float:
+ if tets.shape[0] == 0:
+ return 0.0
+ return float(tet_volumes(node_xyz, tets).sum())
+
+
+def nearest_distance(points: np.ndarray, targets: np.ndarray) -> np.ndarray:
+ """Euclidean distance from each point to the nearest target point.
+ Returns +inf for every point when there are no targets."""
+ if targets.shape[0] == 0:
+ return np.full(points.shape[0], np.inf)
+ tree = cKDTree(targets)
+ d, _ = tree.query(points, k=1)
+ return d
+
+
+def boundary_node_mask(n_nodes: int, tri: np.ndarray) -> np.ndarray:
+ mask = np.zeros(n_nodes, dtype=bool)
+ if tri.shape[0]:
+ mask[np.unique(tri.reshape(-1))] = True
+ return mask
+
+
+_TET_EDGE_PAIRS = ((0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3))
+
+
+def edges_from_tets(tets: np.ndarray) -> np.ndarray:
+ if tets.shape[0] == 0:
+ return np.empty((0, 2), dtype=np.int64)
+ e = np.concatenate([tets[:, [a, b]] for a, b in _TET_EDGE_PAIRS], axis=0)
+ e = np.sort(e, axis=1)
+ return np.unique(e, axis=0).astype(np.int64)
+
+
+def _edge_csr(n_nodes: int, edges: np.ndarray, edge_len: np.ndarray):
+ r = np.concatenate([edges[:, 0], edges[:, 1]])
+ c = np.concatenate([edges[:, 1], edges[:, 0]])
+ w = np.concatenate([edge_len, edge_len])
+ return coo_matrix((w, (r, c)), shape=(n_nodes, n_nodes)).tocsr()
+
+
+def geodesic_distance(n_nodes: int, edges: np.ndarray, edge_len: np.ndarray,
+ source_idx: np.ndarray) -> np.ndarray:
+ if source_idx.shape[0] == 0 or edges.shape[0] == 0:
+ return np.full(n_nodes, np.inf)
+ g = _edge_csr(n_nodes, edges, edge_len)
+ d = dijkstra(g, directed=False, indices=source_idx, min_only=True)
+ 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,
+ chunk: int, query_idx: np.ndarray, g, tvol: np.ndarray,
+ tarea: np.ndarray, bbox_diag: float) -> np.ndarray:
+ """Exact V/A geodesic-ball modulus, computed only for the nodes in
+ `query_idx` (as Dijkstra sources). Returns (len(query_idx), len(radii))."""
+ rmax = float(radii.max())
+ vals = np.full((query_idx.shape[0], len(radii)), bbox_diag, dtype=np.float64)
+ for start in range(0, query_idx.shape[0], chunk):
+ sel = query_idx[start:start + chunk]
+ dmat = dijkstra(g, directed=False, indices=sel, limit=rmax) # (c, N)
+ for j, r in enumerate(radii):
+ in_ball = dmat <= r # (c, N)
+ vol = in_ball[:, tets].all(axis=2) @ tvol # (c,)
+ area = (in_ball[:, cool_tri].all(axis=2) @ tarea
+ if cool_tri.shape[0] else np.zeros(len(sel)))
+ vals[start:start + len(sel), j] = np.where(
+ area > 0, vol / np.maximum(area, 1e-12), bbox_diag)
+ 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:
+ """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."""
+ 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)))
+ out = np.full((N, len(radii)), bbox_diag, dtype=np.float64)
+ if edges.shape[0] == 0 or tets.shape[0] == 0:
+ return out.astype(np.float32)
+ 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)
+
+ return out.astype(np.float32)
diff --git a/src/gnn_dataset.py b/src/gnn_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..f7203ac41f443ac4880295a6fdd258b2d3c79564
--- /dev/null
+++ b/src/gnn_dataset.py
@@ -0,0 +1,51 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..c8a6ec6bae6bdde14d2560bbb8089140f020acbb
--- /dev/null
+++ b/src/graph_build.py
@@ -0,0 +1,52 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..d725672143ceb396717d10f85cdb72248a191497
--- /dev/null
+++ b/src/hybrid.py
@@ -0,0 +1,639 @@
+"""Hybrid predictor: GNN field *shape* + baseline absolute *scale*.
+
+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:
+
+ pred_i = gnn_i - mean_g(gnn) + mean_g(baseline)
+
+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.
+"""
+import logging
+
+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).
+ """
+ g = np.asarray(gnn_seconds, dtype=np.float64)
+ anchor = np.quantile(g, q)
+ out = g - anchor + float(total_seconds)
+ return np.clip(out, 0.0, None)
+
+
+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`).
+ """
+ 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
+ eps = 1e-9
+ a_clipstart = total / max(-float(d.min()), eps)
+
+ def _raw_for_a(a):
+ return a * d + total
+
+ def _post_clip_std(a):
+ return float(np.std(np.clip(_raw_for_a(a), 0.0, None)))
+
+ 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).
+ 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
+ raw = _raw_for_a(a)
+ else:
+ n_scan = 64
+ scan_a = np.linspace(lo_search, hi, n_scan)
+ scan_f = np.array([_post_clip_std(x) for x in scan_a])
+ diffs = scan_f - std_target
+ crossing = None
+ for i in range(1, n_scan):
+ if diffs[i - 1] == 0.0 or diffs[i - 1] * diffs[i] <= 0.0:
+ crossing = i
+ break
+ if crossing is not None:
+ blo, bhi = float(scan_a[crossing - 1]), float(scan_a[crossing])
+ flo = float(diffs[crossing - 1])
+ if flo == 0.0:
+ a = blo
+ else:
+ for _ in range(50):
+ mid = 0.5 * (blo + bhi)
+ fm = _post_clip_std(mid) - std_target
+ if (fm > 0.0) == (flo > 0.0):
+ blo, flo = mid, fm
+ else:
+ bhi = mid
+ if (bhi - blo) <= 1e-9 * max(abs(bhi), 1e-12):
+ break
+ a = 0.5 * (blo + bhi)
+ 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].
+ best_idx = int(np.argmin(np.abs(diffs)))
+ a = float(scan_a[best_idx])
+ a_clamped = True
+ raw = _raw_for_a(a)
+
+ clip_frac = float((raw < 0.0).mean())
+ if clip_frac > 0.5:
+ logger.warning("hybrid_field_affine: majority of field clipped "
+ "(clip_frac=%.3f) -- GBM total/std may be inconsistent "
+ "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}
+
+
+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.
+ """
+ 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")
+ 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.
+ 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")
+
+ pred_seed, info_affine = hybrid_field_affine(g, total, std_seconds, q=q,
+ a_clamp=a_clamp,
+ lambda_shrink=lambda_shrink)
+
+ # `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).
+ c = g - g.mean()
+ shift = float(np.quantile(c, q))
+ a_seed = float(info_affine["a"])
+ m_seed = total - a_seed * shift
+ std_target_eff = float(info_affine["std_target_effective"])
+
+ denom_total = total if abs(total) > 1e-9 else (1e-9 if total >= 0.0 else -1e-9)
+ denom_std = std_target_eff if abs(std_target_eff) > 1e-9 else 1e-9
+
+ def _residuals(pred):
+ r_total = (float(np.quantile(pred, q)) - total) / denom_total
+ r_std = (float(np.std(pred)) - std_target_eff) / denom_std
+ return r_total, r_std
+
+ def _build_info(final_a, final_m, final_rt, final_rs, clip_frac, target_unreachable):
+ if clip_frac > 0.5:
+ logger.warning("hybrid_field_soft: majority of the accepted soft-anchor "
+ "field clipped (clip_frac=%.3f) -- GBM total/std may be "
+ "inconsistent with the GNN shape", clip_frac)
+ return {
+ "a": final_a, "m": final_m, "target_unreachable": target_unreachable,
+ "degenerate_fallback": False, "clip_frac": clip_frac,
+ "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,
+ }
+
+ r_total_seed, r_std_seed = _residuals(pred_seed)
+ j_seed = w_total * r_total_seed ** 2 + w_std * r_std_seed ** 2
+
+ if info_affine["degenerate_fallback"]:
+ return pred_seed, {
+ "a": a_seed, "m": m_seed, "target_unreachable": False,
+ "degenerate_fallback": True, "clip_frac": info_affine["clip_frac"],
+ "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).
+ 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)
+
+ 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.
+ 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.
+ def _grid_eval(a_vals, m_vals):
+ n_av, n_mv = len(a_vals), len(m_vals)
+ J = np.empty((n_av, n_mv))
+ Rt = np.empty((n_av, n_mv))
+ Rs = np.empty((n_av, n_mv))
+ for i, av in enumerate(a_vals):
+ raw = m_vals[:, None] + av * c[None, :] # (n_mv, n_nodes)
+ clipped = np.clip(raw, 0.0, None)
+ qv = av * shift + m_vals
+ sv = np.std(clipped, axis=-1)
+ rt = (qv - total) / denom_total
+ rs = (sv - std_target_eff) / denom_std
+ J[i] = w_total * rt ** 2 + w_std * rs ** 2
+ Rt[i] = rt
+ 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]])
+ best_J = float(J_coarse[idx])
+ best_rt, best_rs = float(Rt_coarse[idx]), float(Rs_coarse[idx])
+
+ if j_seed < best_J:
+ 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.
+ 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)
+ ref_log_a = np.clip(np.linspace(log_best_a - a_step, log_best_a + a_step, 9),
+ log_lo, log_hi)
+ ref_a = 10.0 ** ref_log_a
+
+ m_step = (m_hi - m_lo) / max(n_m - 1, 1)
+ ref_m = np.clip(np.linspace(best_m - m_step, best_m + m_step, 9), m_lo, m_hi)
+
+ J_ref, Rt_ref, Rs_ref = _grid_eval(ref_a, ref_m)
+ idx_r = np.unravel_index(int(np.argmin(J_ref)), J_ref.shape)
+ if float(J_ref[idx_r]) < best_J:
+ best_a, best_m = float(ref_a[idx_r[0]]), float(ref_m[idx_r[1]])
+ 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.
+ 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)
+ best_J = w_total * best_rt ** 2 + w_std * best_rs ** 2
+
+ if best_J < j_seed - 1e-12:
+ final_a, final_m = best_a, best_m
+ final_raw = final_a * c + final_m
+ final_pred = np.clip(final_raw, 0.0, None)
+ final_rt, final_rs = best_rt, best_rs
+ else:
+ final_a, final_m = a_seed, m_seed
+ final_raw = final_a * c + final_m
+ final_pred = pred_seed
+ final_rt, final_rs = r_total_seed, r_std_seed
+
+ clip_frac = float((final_raw < 0.0).mean())
+ return final_pred, _build_info(final_a, final_m, final_rt, final_rs, clip_frac,
+ target_unreachable)
diff --git a/src/inference_props.py b/src/inference_props.py
new file mode 100644
index 0000000000000000000000000000000000000000..653bb42f587a4eebcf870a15fd048362f469b1fd
--- /dev/null
+++ b/src/inference_props.py
@@ -0,0 +1,202 @@
+"""Conditioning resolver for v4 multi-alloy inference.
+
+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:
+
+ {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.
+"""
+from __future__ import annotations
+
+import csv
+import json
+from pathlib import Path
+
+from tools.doe.load_properties import load_materials
+
+DEFAULT_ALLOY = "EN-GJL-250"
+
+_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
+_molds_cache: dict[str, float] | None = None
+
+
+def _materials() -> dict[str, object]:
+ """Module-level cached {material_id: Material}."""
+ global _materials_cache
+ if _materials_cache is None:
+ _materials_cache = {m.material_id: m for m in load_materials(_MATERIALS_PATH)}
+ 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
+ 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.
+ """
+ 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.
+
+ `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.
+
+ 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".
+ """
+ if props_json is not None:
+ with open(props_json) as f:
+ data = json.load(f)
+ missing = [k for k in _REQUIRED_KEYS if k not in data]
+ if missing:
+ raise ValueError(
+ f"props_json '{props_json}' is missing required key(s): {missing} "
+ f"(required: {list(_REQUIRED_KEYS)})")
+ props = {k: float(data[k]) for k in _REQUIRED_KEYS}
+ _validate_tl_ts(props, f"props_json={props_json!r}")
+ return props
+
+ alloy_id = alloy if alloy is not None else DEFAULT_ALLOY
+ materials = _materials()
+ if alloy_id not in materials:
+ raise ValueError(
+ f"unknown alloy '{alloy_id}'; available alloys: {sorted(materials.keys())}")
+ material = materials[alloy_id]
+
+ molds = _molds()
+ 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),
+ "rho": float(material.rho),
+ "cp": float(material.cp),
+ "L_kJ_kg": float(material.L),
+ "TL": float(material.TL),
+ "TS": float(material.TS),
+ "Tpour": float(material.Tpour),
+ "mold_effusivity": float(mold_effusivity),
+ }
+ _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.
+ """
+ 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
diff --git a/src/mesh_cache.py b/src/mesh_cache.py
new file mode 100644
index 0000000000000000000000000000000000000000..d79811c62d185ff86cbfebfab41a9ba3e5bcce13
--- /dev/null
+++ b/src/mesh_cache.py
@@ -0,0 +1,236 @@
+"""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).
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import shutil
+import time
+import uuid
+
+from tools.doe.prep_parts import DEFAULT_TARGET_NODES, MESH_PIPELINE_VERSION, heal_and_mesh
+
+_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.
+_BIG_FILE_THRESHOLD_BYTES = 8 * 1024 * 1024 # 8 MiB
+
+
+def _sha256_file(path: str) -> str:
+ h = hashlib.sha256()
+ with open(path, "rb") as f:
+ while True:
+ chunk = f.read(_HASH_CHUNK_SIZE)
+ if not chunk:
+ break
+ h.update(chunk)
+ 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).
+ """
+ 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")
+ 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"
+
+
+def _paths(cache_dir: str, key: str) -> tuple[str, str]:
+ return (
+ os.path.join(cache_dir, f"{key}.msh"),
+ os.path.join(cache_dir, f"{key}.json"),
+ )
+
+
+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.
+ """
+ os.makedirs(cache_dir, exist_ok=True)
+ key = mesh_cache_key(step_path, target_nodes, bodies_key)
+ msh_path, json_path = _paths(cache_dir, key)
+
+ if os.path.exists(msh_path) and os.path.exists(json_path):
+ with open(json_path) as f:
+ meta = json.load(f)
+ shutil.copyfile(msh_path, out_msh)
+ result = dict(meta)
+ result["cache"] = "hit"
+ return result
+
+ heal_and_mesh_fn = globals()["heal_and_mesh"]
+ 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)
+ 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)
+ 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)
+ 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)
+ defeature_used, retry = True, True
+ wall_s = time.time() - t0
+
+ if not meta.get("ok"):
+ result = dict(meta)
+ result["cache"] = "miss"
+ result["defeature"] = defeature_used
+ 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)
+ stored_meta["created_at"] = time.time()
+ stored_meta["defeature"] = defeature_used
+ stored_meta["retry"] = retry
+
+ 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
+ os.replace(msh_tmp, msh_path)
+ with open(json_tmp, "w") as f:
+ json.dump(stored_meta, f)
+ os.replace(json_tmp, json_path)
+
+ result = dict(stored_meta)
+ result["cache"] = "miss"
+ return result
+
+
+def clear_mesh_cache(cache_dir: str = "runs/mesh_cache") -> None:
+ """Remove the entire cache directory, if present."""
+ shutil.rmtree(cache_dir, ignore_errors=True)
diff --git a/src/metrics.py b/src/metrics.py
new file mode 100644
index 0000000000000000000000000000000000000000..02aa001312c5a4948ef45025233fd3edd7bddf69
--- /dev/null
+++ b/src/metrics.py
@@ -0,0 +1,50 @@
+"""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_gnn.py b/src/model_gnn.py
new file mode 100644
index 0000000000000000000000000000000000000000..f25fb276243a670186508b3d8be4cc47fa598733
--- /dev/null
+++ b/src/model_gnn.py
@@ -0,0 +1,136 @@
+"""MeshGraphNet-style encode-process-decode GNN predicting standardized
+log1p(solidification_time) per node; global scalar = max over nodes."""
+import torch
+import torch.nn as nn
+from torch_geometric.utils import scatter
+
+
+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):
+ def __init__(self, h):
+ super().__init__()
+ self.edge_mlp = _mlp(3 * h, h, h)
+ self.edge_norm = nn.LayerNorm(h)
+ self.node_mlp = _mlp(2 * h, h, h)
+ self.node_norm = nn.LayerNorm(h)
+
+ def forward(self, hnode, edge_index, e):
+ src, dst = edge_index
+ e_new = self.edge_norm(self.edge_mlp(torch.cat([hnode[src], hnode[dst], e], dim=-1)) + e)
+ agg = scatter(e_new, dst, dim=0, dim_size=hnode.shape[0], reduce="mean")
+ h_new = self.node_norm(self.node_mlp(torch.cat([hnode, agg], dim=-1)) + hnode)
+ return h_new, e_new
+
+
+class FiLM(nn.Module):
+ """g -> (gamma, beta) per hidden channel; zero-init => gamma=0,beta=0 => identity."""
+ 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)
+
+ def forward(self, h, g_nodes):
+ gamma, beta = self.net(g_nodes).chunk(2, dim=-1)
+ return h * (1.0 + gamma) + beta
+
+
+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):
+ 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):
+ 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
+ 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)
+ 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)
diff --git a/src/msh_parse.py b/src/msh_parse.py
new file mode 100644
index 0000000000000000000000000000000000000000..644418d59f3dc7b13a0ed86fd7484c1b382dfd7e
--- /dev/null
+++ b/src/msh_parse.py
@@ -0,0 +1,67 @@
+"""Parser for gmsh 2.2 ASCII meshes used by the Foundry dataset.
+
+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.
+"""
+from dataclasses import dataclass
+import numpy as np
+
+_TRI_TYPE = 2 # 3-node triangle
+_TET_TYPE = 4 # 4-node tetrahedron
+
+
+@dataclass
+class Mesh:
+ node_xyz: np.ndarray # (N, 3) float64, mm
+ tris: dict # name -> (M, 3) int64, 0-based
+ tets: np.ndarray # (K, 4) int64, 0-based
+
+
+def parse_msh(path: str) -> Mesh:
+ lines = open(path).read().splitlines()
+ i, n = 0, len(lines)
+ phys = {} # physical tag -> name
+ node_ids = [] # id order
+ coords = [] # xyz in that order
+ tri_by_name = {} # name -> list of [i,j,k]
+ tets = []
+
+ while i < n:
+ s = lines[i].strip()
+ if s == "$PhysicalNames":
+ cnt = int(lines[i + 1])
+ for k in range(cnt):
+ dim, tag, name = lines[i + 2 + k].split(maxsplit=2)
+ phys[int(tag)] = name.strip().strip('"')
+ i += 2 + cnt
+ continue
+ if s == "$Nodes":
+ cnt = int(lines[i + 1])
+ for k in range(cnt):
+ p = lines[i + 2 + k].split()
+ node_ids.append(int(p[0]))
+ coords.append((float(p[1]), float(p[2]), float(p[3])))
+ i += 2 + cnt
+ continue
+ if s == "$Elements":
+ cnt = int(lines[i + 1])
+ id2idx = {nid: idx for idx, nid in enumerate(node_ids)}
+ for k in range(cnt):
+ p = [int(x) for x in lines[i + 2 + k].split()]
+ etype, ntags = p[1], p[2]
+ tags = p[3:3 + ntags]
+ conn = [id2idx[v] for v in p[3 + ntags:]]
+ name = phys.get(tags[0]) if ntags else None
+ if etype == _TRI_TYPE and name is not None:
+ tri_by_name.setdefault(name, []).append(conn)
+ elif etype == _TET_TYPE:
+ tets.append(conn)
+ i += 2 + cnt
+ continue
+ i += 1
+
+ node_xyz = np.asarray(coords, dtype=np.float64)
+ tris = {name: np.asarray(v, dtype=np.int64) for name, v in tri_by_name.items()}
+ tets_arr = np.asarray(tets, dtype=np.int64) if tets else np.empty((0, 4), np.int64)
+ return Mesh(node_xyz=node_xyz, tris=tris, tets=tets_arr)
diff --git a/src/ood.py b/src/ood.py
new file mode 100644
index 0000000000000000000000000000000000000000..b150dc1bb25bdb68e37d61411bb3b90a961e828e
--- /dev/null
+++ b/src/ood.py
@@ -0,0 +1,49 @@
+"""Out-of-distribution size check: compare a casting's scale quantities against
+the training-set maxima so predictions beyond the training envelope are flagged."""
+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:
+ with open(path) as f:
+ return json.load(f)
+
+
+def ood_flags(volume_mm3, bbox_diag, geodesic_d_cool_max, ranges, tol=1.05) -> list:
+ w = []
+ if volume_mm3 > ranges["volume_mm3_max"] * tol:
+ w.append(f"casting volume {volume_mm3:.0f} mm^3 exceeds training max "
+ f"{ranges['volume_mm3_max']:.0f} mm^3 — prediction unreliable")
+ if bbox_diag > ranges["bbox_diag_max"] * tol:
+ w.append(f"bounding-box diagonal {bbox_diag:.0f} mm exceeds training max "
+ f"{ranges['bbox_diag_max']:.0f} mm — prediction unreliable")
+ if geodesic_d_cool_max > ranges["geodesic_d_cool_max"] * tol:
+ w.append(f"max distance-to-cooling {geodesic_d_cool_max:.1f} mm exceeds training max "
+ f"{ranges['geodesic_d_cool_max']:.1f} mm — prediction unreliable")
+ return w
diff --git a/src/predict_casting.py b/src/predict_casting.py
new file mode 100644
index 0000000000000000000000000000000000000000..cb5b977873b7c458d2e3d7fc0b0178631b3d0e91
--- /dev/null
+++ b/src/predict_casting.py
@@ -0,0 +1,1409 @@
+"""Real-casting inference: hybrid field prediction + VTU output.
+
+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.
+
+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).
+
+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. The chill low-confidence
+`niyama_confidence` marking (spec ADDENDUM 2) is UNCHANGED by this -- still
+keyed only on `plan.bc_config == "chills"`, regardless of `niyama_source`;
+retiring that scoping is a later (promotion-step) task.
+
+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 `bc_plan.feature_spec(...)`/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 logging
+import os.path as osp
+import numpy as np
+import torch
+import joblib
+
+logger = logging.getLogger(__name__)
+
+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.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.")
+
+
+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, bc_plan=None, hgb_std=None,
+ lambda_shrink: float = 1.0, w_total: float = DEFAULT_W_TOTAL,
+ w_std: float = DEFAULT_W_STD) -> tuple:
+ """`bc_plan` (Part B Task 4): an optional `src.assembly.BcPlan`, resolved
+ by `validate_bodies` from a multi-body role assignment. `None` (the
+ default) is byte-identical to the pre-Part-B behavior -- the film
+ branch's `spec` stays the all-sand literal and `gbm_v2_total_seconds` is
+ called with `bc_spec=None` (its own "plain sand mold" default). When
+ given, the film branch's `spec` comes from `bc_plan.feature_spec(...)`
+ (carries `chill_effusivity`/`sleeve_MEF`/`exo_q0`/`bc_config`) and
+ `gbm_v2_total_seconds` gets `bc_spec=bc_plan.gbm_bc_spec()`. Ignored
+ entirely by the non-film (SP1) path -- SP1 has no BC channel at all.
+
+ `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)")
+
+ if bc_plan is not None:
+ spec = bc_plan.feature_spec(props["mold_effusivity"])
+ bc_spec = bc_plan.gbm_bc_spec()
+ else:
+ 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
+
+
+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."""
+ import meshio
+ seconds = np.asarray(seconds, dtype=np.float32)
+ n = seconds.shape[0]
+ k = max(1, int(round(n * top_frac)))
+ thresh = np.sort(seconds)[-k]
+ hotspot = (seconds >= thresh).astype(np.int32)
+ point_data = {"solidification_time": seconds, "hotspot": hotspot}
+ if extra_point_data:
+ point_data.update(extra_point_data)
+ m = meshio.Mesh(points=np.asarray(node_xyz, dtype=np.float64),
+ cells=[("tetra", np.asarray(tets))],
+ point_data=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")."""
+ 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))
+ return surface_nodes.astype(np.int64), feed_nodes.astype(np.int64)
+
+
+# --------------------------------------------------------------------------- #
+# Part-B roles task (2026-08-03): role-aware feed-node selection, replacing
+# review finding I1's coarse `plan.bc_config == "chills"` special case (which
+# forced feed_nodes empty for the WHOLE input whenever any chill role was
+# present, rather than excluding just the chill bodies' own patch nodes).
+# `src.assembly.tag_contact_patches` relabels EVERY role body's contact patch
+# "Adiabatic" regardless of role (chills included) -- these two pure helpers
+# recover the per-BODY-role distinction from `tag_contact_patches`'s
+# `nodes_by_solid` output (per-solid, geometry-only) + a solid_id -> role
+# lookup (`BcPlan.role_by_solid`), so the exclusion decision is made per
+# body, never per global `bc_config`.
+#
+# Investigation note (binding on the design below): in bodies mode the
+# casting mesh handed to `tag_contact_patches` never carries a PRE-EXISTING
+# "Adiabatic" group of its own. `bodies` requires `mesh_mode="prep"` (see the
+# guard above `predict_casting` raises), and `tools.doe.prep_parts.
+# heal_and_mesh` -- the only meshing path "prep" mode uses -- tags every
+# casting surface "Cooling" uniformly with no Adiabatic seed at all (verified
+# in `tools/doe/prep_parts.py` and `src/step_mesh.py`: neither ever calls
+# `setPhysicalName(2, ..., "Adiabatic")`; the ONLY producer of an "Adiabatic"
+# physical group anywhere in this codebase is `src.assembly.
+# rewrite_msh_adiabatic`, i.e. THIS tagging pass itself, or FT3's training-
+# data patch simulator `tools/doe/virtual_patches.py`, which never runs at
+# inference time). So every "Adiabatic" triangle on a bodies-mode inference
+# mesh is a role-body contact patch -- there is no separate "genuine riser"
+# contribution that could coexist with role patches and would need unioning
+# in on top of the per-role mapping below. The bare/all-sand path (no
+# `bodies`) is therefore UNCHANGED by this task: `_shrinkage_surface_and_
+# feed_nodes`'s "Adiabatic" union (if the mesh has one at all -- only
+# possible via `mesh_mode="raw"` or a training `.msh` with real user-tagged
+# riser faces) keeps meaning exactly what it always has.
+# --------------------------------------------------------------------------- #
+def _nodes_by_role_from_tagging(nodes_by_solid: dict, role_by_solid: dict) -> dict:
+ """Group `tag_contact_patches`'s per-solid contact-patch node ids
+ (`nodes_by_solid`, `{solid_id: node ids}` -- solid_id keys may arrive as
+ `str` after a mesh-cache JSON round-trip on a cache HIT, handled here via
+ `int()`) by each solid's declared role (`role_by_solid`, `{solid_id:
+ role}`, e.g. `BcPlan.role_by_solid`).
+
+ A solid_id present in `nodes_by_solid` with no entry in `role_by_solid`
+ is dropped rather than raising, but LOUDLY (`logger.warning`, naming the
+ solid_id) -- review round 1, M1: silently dropping it would silently
+ mis-attribute (or simply lose) that body's contact-patch nodes. In the
+ normal `predict_casting` call path this warning never actually fires:
+ `_role_aware_feed_nodes` (review round 1, H1) only calls this function
+ after confirming `nodes_by_solid` covers every id in `plan.
+ role_solid_ids`, and `role_by_solid` covers exactly those same ids by
+ construction in `validate_bodies` -- this is a defense-in-depth net for
+ any future/direct caller that skips that completeness check, not a code
+ path this module exercises today.
+
+ Returns `{role: sorted node ids}`.
+
+ NOTE on shared-node precedence (L3, review round 1): if a node's
+ contact-patch membership spans more than one role (two role bodies'
+ `tol_mm` zones overlap on the casting surface -- see `tag_contact_
+ patches`'s documented double-counting of `per_body_patch_area_mm2`),
+ this function keeps it under EVERY role it was classified into, and
+ `_feed_nodes_from_role_patches` below then unions in whichever of those
+ roles is a sleeve role -- i.e. a node shared between a chill patch and a
+ sleeve patch counts as fed ("sleeve wins"), matching the spec's literal
+ per-role union rule. `validate_bodies` forbids a chills+sleeve mix today
+ so this can never actually occur via the validated path (only via a
+ directly-constructed mapping that bypasses it, e.g. this module's own
+ mixed-mapping unit test). A FUTURE mixed-role regime might instead want
+ "chill wins" for a shared node (the more physically conservative choice
+ -- proximity to a chill accelerates freezing there regardless of a
+ nearby sleeve) -- flagged here for that future design, not implemented
+ now."""
+ grouped: dict[str, set] = {}
+ for sid_key, nodes in nodes_by_solid.items():
+ sid = int(sid_key)
+ role = role_by_solid.get(sid)
+ if role is None:
+ logger.warning(
+ f"_nodes_by_role_from_tagging: solid_id={sid} has classified "
+ f"contact-patch nodes but no entry in role_by_solid -- "
+ f"dropping it from the per-role feed-node mapping rather than "
+ f"silently mis-attributing it to some other role")
+ continue
+ grouped.setdefault(role, set()).update(int(n) for n in nodes)
+ return {role: sorted(ids) for role, ids in grouped.items()}
+
+
+def _feed_nodes_from_role_patches(nodes_by_role: dict) -> np.ndarray:
+ """The sorted-unique union of every SLEEVE role's nodes
+ (`src.assembly._SLEEVE_ROLES` = `insulating_sleeve`/`exo_sleeve`) in
+ `nodes_by_role` (`{role: iterable of node ids}`) -- a chill patch, or any
+ other non-sleeve role, NEVER contributes a feed node. Pure function of
+ its argument, with NO `BcPlan`/`bc_config` involved at all: it stays
+ correct even fed a directly-constructed mapping mixing roles that
+ `validate_bodies` itself would reject (chills + a sleeve combined) --
+ the decision is per body role, not per global config. See
+ `_nodes_by_role_from_tagging`'s docstring (L3, review round 1) for the
+ shared-node ("sleeve wins") precedence this union implies when a node is
+ classified into more than one role."""
+ feed: set = set()
+ for role, nodes in nodes_by_role.items():
+ if role in _SLEEVE_ROLES:
+ feed.update(int(n) for n in nodes)
+ return np.asarray(sorted(feed), dtype=np.int64)
+
+
+def _role_aware_feed_nodes(step_path: str, nodes_by_solid: dict, plan,
+ fallback_feed_nodes: np.ndarray) -> np.ndarray:
+ """Review round 1, H1: `nodes_by_solid` is only present in `prep_info`
+ when `tag_contact_patches` actually ran on THIS call. A mesh-cache HIT
+ served from a side-car JSON written BEFORE the Part-B roles task never
+ had `nodes_by_solid` written into it at all (the key didn't exist yet),
+ and even a post-task entry could in principle cover only some of
+ `plan.role_solid_ids`. `MESH_PIPELINE_VERSION` is deliberately NOT
+ bumped for this task (cache keys are binding, per the task brief), so
+ such LEGACY entries are genuinely reachable today -- `runs/mesh_cache/`
+ already holds pre-task assembly entries, including both gold parts.
+ Naively grouping incomplete/absent `nodes_by_solid` data would silently
+ degrade `feed_nodes` to empty for EVERY `bc_config` (not just chills),
+ a silent regression on exactly the sleeve inputs this task's own byte-
+ identity test binds. `plan.role_solid_ids` MUST all be present as keys
+ of `nodes_by_solid` (after normalizing str keys) for the new per-role
+ path to run at all; otherwise this loudly (`logger.warning`, naming
+ `step_path`) falls back to the PRE-CHANGE rule: empty for `bc_config ==
+ "chills"`, the passed-through `fallback_feed_nodes` (the plain
+ "Adiabatic"-tag node union `_shrinkage_surface_and_feed_nodes` already
+ computed) for every other config -- i.e. exactly review finding I1's
+ original fix, byte-identical to `predict_casting`'s behavior before this
+ task. Recommends refreshing/clearing the mesh-cache entry for
+ `step_path` (see `src.mesh_cache.clear_mesh_cache`) so a later run gets
+ full per-role accuracy."""
+ required_ids = set(plan.role_solid_ids)
+ covered_ids = {int(k) for k in nodes_by_solid}
+ if not required_ids.issubset(covered_ids):
+ logger.warning(
+ f"predict_casting: role-aware contact-patch node metadata is "
+ f"missing or incomplete for step_path={step_path!r} "
+ f"(nodes_by_solid covers solid_id(s) {sorted(covered_ids)}, but "
+ f"plan.role_solid_ids requires {sorted(required_ids)}) -- this is "
+ f"almost certainly a LEGACY runs/mesh_cache/ entry written before "
+ f"the Part-B roles task (nodes_by_solid didn't exist yet). "
+ f"Falling back to the pre-change feed-node rule (empty for "
+ f"bc_config=='chills', the Adiabatic-tag union otherwise) rather "
+ f"than silently zeroing every config's feed_nodes. Recommend "
+ f"clearing/refreshing the mesh-cache entry for this STEP file "
+ f"(see src.mesh_cache.clear_mesh_cache) to restore full per-role "
+ f"feed-node accuracy on the next run.")
+ return (np.zeros(0, dtype=np.int64) if plan.bc_config == "chills"
+ else fallback_feed_nodes)
+ nodes_by_role = _nodes_by_role_from_tagging(nodes_by_solid, plan.role_by_solid)
+ return _feed_nodes_from_role_patches(nodes_by_role)
+
+
+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."""
+ 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
+ result["n_macro_regions"] = len(screen["macro_shrinkage"])
+ result["n_sinks"] = len(screen["top_sinks"])
+ result["n_feeders"] = len(screen["feeders"])
+ 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`)."""
+ n_nodes = mesh.node_xyz.shape[0]
+ macro_region_id = np.full(n_nodes, -1, dtype=np.int64)
+ for region in screen["macro_shrinkage"]:
+ macro_region_id[np.asarray(region["nodes"], dtype=np.int64)] = region["region_id"]
+ return {
+ "niyama": np.asarray(screen["niyama"], dtype=np.float32),
+ "porosity_risk": np.asarray(screen["porosity_risk"], dtype=np.uint8),
+ "macro_region_id": macro_region_id,
+ }
+
+
+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)."""
+ 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")
+ if shrinkage.get("niyama_confidence") == "low":
+ # Spec ADDENDUM 2 (2026-08-01, owner-approved scoped ship): Gate C
+ # passed both floors on the all-sand slice (sleeves owner-accepted as
+ # intermediate: rho 0.768/0.691, below the 0.8 floor) but failed on the
+ # chill-bearing slice (Spearman rho 0.313) -- the proxy Niyama
+ # identity `Tdot_proxy = (TL-TS)/t_s` has no way to see a chill's
+ # locally-enhanced cooling rate, so its Niyama ranking on a
+ # chill-bearing input is not trustworthy the way it is on a bare or
+ # sleeve-only casting. Macro-shrinkage/sinks/feeders don't depend on
+ # Niyama at all and keep full confidence regardless.
+ #
+ # Learned-heads Task 3 review round 1, F3: the marking itself
+ # (`niyama_confidence == "low"` on any chill-bearing input) STAYS as
+ # is until the promote step (per the spec's own retirement plan) --
+ # only the CAVEAT TEXT is now conditional on `niyama_source`. The
+ # proxy's specific failure mode (and its Gate C rho number) is a
+ # claim about the PROXY identity, not about a learned checkpoint's
+ # heads -- citing "rho 0.313" next to a learned-source result would
+ # misattribute the proxy's own measured shortfall to a different
+ # (not yet slice-validated) source. A learned checkpoint gets a
+ # neutral caveat instead, pending its own Gate C learned-mode slice
+ # numbers (`tools/gate_c.py field_source="learned"`).
+ if shrinkage["niyama_source"] == "proxy":
+ f.write(" niyama confidence: LOW (chill-bearing input; proxy Niyama cannot see "
+ "chill-local cooling rates — Gate C chills slice rho 0.313)\n")
+ elif shrinkage["niyama_source"] == "learned":
+ f.write(" niyama confidence: LOW (chill-bearing input; learned-head Niyama "
+ "pending slice validation)\n")
+ else:
+ f.write(f" niyama confidence: LOW (chill-bearing input; "
+ f"niyama_source={shrinkage['niyama_source']!r})\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")
+ 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.
+ n_sub_resolution = sum(1 for r in macro if r.get("sub_resolution"))
+ region_count = str(len(macro))
+ if n_sub_resolution:
+ region_count += f" (+{n_sub_resolution} sub-resolution)"
+ 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")
+ else:
+ f.write(" macro shrinkage regions: 0\n")
+ f.write(f" sinks: {len(shrinkage['top_sinks'])}\n")
+ feeders = shrinkage["feeders"]
+ if feeders:
+ f.write(" feeder recommendations:\n")
+ 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")
+ 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 src.assembly import validate_bodies, tag_contact_patches, _SLEEVE_ROLES
+from tools.doe.prep_parts import DEFAULT_TARGET_NODES
+
+
+def _bodies_cache_key(bodies: list, plan) -> str:
+ """Part B Task 4: a normalized-content hash of `bodies`, folding in the
+ resolved `BcPlan` scalars (`chill_effusivity`/`sleeve_MEF`/`exo_q0`) for
+ whichever body ids are `plan.role_solid_ids` -- NOT the raw role/params
+ text the caller typed -- so a `sleeve_registry.csv`/`chills.csv` edit
+ changes the mesh-cache key even for a byte-identical `bodies.json` (the
+ key means "this exact resolved BC plan", not just "this exact JSON
+ text"). Order-invariant (sorted by `solid_id`) and role-sensitive (a
+ role change always changes that row, even for a body with no numeric
+ params of its own, e.g. `"casting"` or `"ignore"`)."""
+ role_ids = set(plan.role_solid_ids)
+ rows = []
+ for i, b in enumerate(bodies):
+ if "solid_id" not in b:
+ raise ValueError(
+ f"_bodies_cache_key: bodies[{i}] (role={b.get('role')!r}) is "
+ f"missing 'solid_id' -- every row, including 'ignore' rows, "
+ f"must carry a solid_id")
+ row = {"solid_id": b["solid_id"], "role": b.get("role")}
+ if b["solid_id"] in role_ids:
+ if plan.bc_config == "chills":
+ row["chill_effusivity"] = plan.chill_effusivity
+ elif plan.bc_config in ("insulating_sleeve", "exo_sleeve"):
+ row["sleeve_MEF"] = plan.sleeve_MEF
+ row["exo_q0"] = plan.exo_q0
+ rows.append(row)
+ rows.sort(key=lambda r: r["solid_id"])
+ return hashlib.sha256(json.dumps(rows, sort_keys=True).encode("utf-8")).hexdigest()
+
+
+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,
+ bodies=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.
+
+ `bodies` (Part B Task 4): a multi-body STEP's role assignment -- either
+ a list of `{"solid_id": int, "role": ..., "params": {...}}` dicts (see
+ `src.assembly.validate_bodies`) or a `str`/path-like path to a JSON file
+ holding that list (loaded via `json.load`). `None` (the default) is
+ byte-identical to the pre-Part-B behavior: no `validate_bodies` call,
+ `keep_solid_id=None`/`bodies_key="none"` into `get_or_mesh`, no
+ `tag_contact_patches` call, and the film spec/gbm bc_spec stay at their
+ all-sand-mold defaults (see `predict_field`). When given, the flow is
+ `validate_bodies` -> `get_or_mesh(..., keep_solid_id=plan.casting_solid_id,
+ bodies_key=, post_mesh=)`
+ -> `predict_field(..., bc_plan=plan)`; the resolved `bc_config` and any
+ contact-patch areas are added to the result dict and the summary file.
+ Requires `mesh_mode="prep"` (the effective mode after any `elem_size`
+ override) -- role-based casting-solid selection and contact tagging
+ only run through that pipeline; combining `bodies` with the "raw" mesh
+ path raises `ValueError` rather than silently ignoring the role
+ assignment.
+
+ `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 (same reasoning as
+ `bodies`: 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).
+
+ Shrinkage-screening spec ADDENDUM 2 (2026-08-01, owner-approved SCOPED
+ SHIP): the shrinkage dict always carries `"niyama_confidence"` (`"full"`
+ or `"low"`). Gate C (the proxy-Niyama validation) passed both floors on
+ all-sand only (sleeve slices are an owner-accepted intermediate: rho
+ 0.768/0.691, below the 0.8 floor) and FAILED on the chill-bearing slice (Spearman rho
+ 0.313 -- the proxy `Tdot ~ (TL-TS)/t_s` identity has no way to see a
+ chill's local cooling-rate enhancement). Whenever `bodies` resolves to a
+ `BcPlan` with `bc_config == "chills"`, `"niyama_confidence"` is `"low"`
+ and the summary gains a caveat line naming the Gate C number; bare/
+ all-sand and sleeve-only inputs (`bc_config` `"none"`/
+ `"insulating_sleeve"`/`"exo_sleeve"`) stay `"full"`. The `niyama`/
+ `porosity_risk` VTU fields and the return dict still populate either way
+ (screening stays informative, just honestly labeled) -- macro-shrinkage/
+ top-sink/feeder outputs never depended on Niyama and are unaffected by
+ this label.
+ `False` skips all of that entirely --
+ `"shrinkage"` is `None` and the VTU/summary are byte-identical to the
+ pre-Task-4 contract (no extra fields, no extra summary lines at all --
+ 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.
+ """
+ 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.
+ props = resolve_props(alloy=alloy, mold=mold, props_json=props_json)
+
+ plan = None
+ bodies_key = "none"
+ if bodies is not None:
+ if isinstance(bodies, (str, os.PathLike)):
+ with open(osp.expanduser(bodies)) as f:
+ bodies = json.load(f)
+ # No `mef_weights` here (SP4 review I3, 2026-08-08): a mixed-MEF
+ # sleeve set needs contact-patch areas per solid_id, but those only
+ # exist after meshing + `tag_contact_patches`, which have not run
+ # yet at this point -- `validate_bodies` is called here specifically
+ # to build the mesh-cache key (`_bodies_cache_key` below) before any
+ # mesh exists. This is not a regression (every mixed-MEF assembly
+ # was rejected outright before the multi-sleeve ADDENDUM), but it
+ # does mean this CLI/webapp entry point cannot resolve a mixed-MEF
+ # assembly -- it raises naming `mef_weights`, a parameter neither
+ # surface accepts. A caller that needs a mixed-MEF assembly must
+ # mesh and tag first, then call `validate_bodies` again with
+ # `mef_weights=per_body_patch_area_mm2`, exactly as `tools/sp4/
+ # prep.py` + `tools/sp4/score.py` do for the SP4 gold set. Known
+ # gap, not fixed here -- restructuring this function's meshing flow
+ # to do that two-pass resolution is separate work; see the spec's
+ # multi-sleeve ADDENDUM.
+ plan = validate_bodies(bodies)
+ bodies_key = _bodies_cache_key(bodies, plan)
+
+ 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.")
+ if bodies is not None and effective_mode != "prep":
+ raise ValueError(
+ f"predict_casting: `bodies` requires mesh_mode='prep' (got "
+ f"mesh_mode={mesh_mode!r}, effective_mode={effective_mode!r} -- "
+ f"an explicit elem_size forces 'raw') -- role-based casting-solid "
+ f"selection and contact-patch tagging only run through the "
+ f"prep/get_or_mesh pipeline")
+
+ # Loaded BEFORE meshing so a bodies/non-film-checkpoint mismatch (below)
+ # is caught before tag_contact_patches ever runs -- not merely before
+ # its result is used. A resolved BcPlan (chills/sleeve roles) means
+ # tag_contact_patches would relabel the contact patch "Adiabatic", which
+ # a non-film (SP1) checkpoint's geometry features read AS PHYSICAL
+ # INSULATION. For a chill, that's backwards physics (a chill COOLS,
+ # Adiabatic INSULATES); for a sleeve it's merely uncalibrated. Either
+ # way, SP1 has no BC channel to honor the role assignment at all, so
+ # silently proceeding would both report bc_config as if honored AND
+ # leave a wrong-physics tagged mesh sitting in the mesh cache.
+ 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 plan is not None and not cfg.get("film", False):
+ raise ValueError(
+ "predict_casting: bodies/BC roles require the film-conditioned "
+ "v4 stack; the selected checkpoint is not film-capable -- rerun "
+ "without --bodies or with the v4 checkpoints")
+ 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":
+ keep_solid_id = plan.casting_solid_id if plan is not None else None
+ post_mesh = None
+ if plan is not None and plan.role_solid_ids:
+ def post_mesh(msh_path, _step_path=step_path, _role_ids=plan.role_solid_ids):
+ return tag_contact_patches(msh_path, _step_path, _role_ids)
+
+ prep_info = get_or_mesh(step_path, out_msh, target_nodes=target_nodes,
+ keep_solid_id=keep_solid_id, bodies_key=bodies_key,
+ post_mesh=post_mesh)
+ 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}
+ if "per_body_patch_area_mm2" in prep_info:
+ # A cache HIT round-trips through JSON, turning int dict keys
+ # into str -- normalize back to int so callers (and the summary
+ # below) see one consistent shape regardless of hit vs. miss.
+ minfo["per_body_patch_area_mm2"] = {
+ int(k): v for k, v in prep_info["per_body_patch_area_mm2"].items()}
+ 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, bc_plan=plan,
+ hgb_std=hgb_std, lambda_shrink=lambda_shrink,
+ w_total=w_total, w_std=w_std)
+
+ warnings = []
+ ood_json = resolve_artifact(ood_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)
+
+ # 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
+ if screen_shrinkage:
+ surface_nodes, feed_nodes = _shrinkage_surface_and_feed_nodes(mesh)
+ if plan is not None and plan.role_solid_ids:
+ # Part-B roles task (2026-08-03, replaces review finding I1's
+ # coarse `bc_config == "chills"` special case): recompute
+ # feed_nodes PER BODY ROLE rather than zeroing it out for the
+ # whole input -- only the SLEEVE roles' nodes are ever a feed
+ # source, a chill patch never is, regardless of `bc_config`. See
+ # `_feed_nodes_from_role_patches`'s docstring and this module's
+ # own investigation note (above `_nodes_by_role_from_tagging`)
+ # for why no separate "pre-existing Adiabatic" contribution
+ # needs unioning in here. `_role_aware_feed_nodes` (review round
+ # 1, H1) additionally guards against a LEGACY mesh-cache HIT
+ # whose side-car predates this task (no `nodes_by_solid` at
+ # all) -- see its own docstring.
+ nodes_by_solid = prep_info.get("nodes_by_solid", {})
+ feed_nodes = _role_aware_feed_nodes(step_path, nodes_by_solid, plan, feed_nodes)
+ 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)")
+ 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:
+ shrinkage["alloy_family_assumed"] = True
+ # Spec ADDENDUM 2 (2026-08-01, owner-approved scoped ship): Gate C
+ # failed on the chill-bearing slice (rho 0.313) -- proxy Niyama
+ # cannot see a chill's local cooling-rate enhancement, so mark it
+ # low-confidence whenever the input's BC plan has a CHILL role
+ # (`plan.bc_config == "chills"`; bare/all-sand and sleeve-only
+ # inputs, whose bc_config is "none"/"insulating_sleeve"/
+ # "exo_sleeve", are unaffected and stay "full"). Macro-shrinkage/
+ # sinks/feeders don't depend on Niyama at all and keep full scope
+ # regardless -- only this confidence LABEL changes, no output is
+ # withheld ("screening stays informative").
+ shrinkage["niyama_confidence"] = (
+ "low" if (plan is not None and plan.bc_config == "chills") else "full")
+ shrinkage_extra_point_data = _shrinkage_extra_point_data(mesh, shrinkage_full)
+
+ vtu = osp.join(out_dir, f"{name}.vtu")
+ write_vtu(vtu, mesh.node_xyz, mesh.tets, seconds, extra_point_data=shrinkage_extra_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())
+ 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})"
+ 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)"
+ bc_config = plan.bc_config if plan is not None else "none"
+ patch_areas = minfo.get("per_body_patch_area_mm2", {})
+ 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")
+ f.write(f"bc_config: {bc_config}\n")
+ if plan is not None:
+ for b in bodies:
+ sid = b.get("solid_id")
+ line = f" body solid_id={sid} role={b.get('role')}"
+ if b.get("params"):
+ line += f" params={b['params']}"
+ if sid in patch_areas:
+ line += f" patch_area_mm2={patch_areas[sid]:.1f}"
+ f.write(line + "\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")
+ if screen_shrinkage:
+ _write_shrinkage_summary_block(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")
+ if warnings:
+ f.write("WARNINGS:\n")
+ 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
+
+ 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, "bc_config": bc_config,
+ "per_body_patch_area_mm2": patch_areas,
+ "recomposition": recomposition,
+ "recomposition_info": recomposition_info,
+ "shrinkage": shrinkage}
+ if "cache" in minfo:
+ result["cache"] = minfo["cache"]
+ return result
+
+
+if __name__ == "__main__":
+ ap = argparse.ArgumentParser()
+ 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")
+ ap.add_argument("--props", dest="props_json", default=None,
+ help="path to a JSON file overriding alloy/mold with raw material properties")
+ ap.add_argument("--bodies", dest="bodies_json", default=None,
+ help="path to a JSON file describing a multi-body STEP's role "
+ "assignment (Part B): a list of {\"solid_id\": int, \"role\": "
+ "\"casting\"|\"chill\"|\"insulating_sleeve\"|\"exo_sleeve\"|"
+ "\"ignore\", \"params\": {...} (optional, role-dependent)}. "
+ "Example: [{\"solid_id\": 1, \"role\": \"casting\"}, "
+ "{\"solid_id\": 2, \"role\": \"chill\", \"params\": "
+ "{\"material\": \"cast_iron\"}}]. Requires --mesh-mode prep "
+ "(the default) -- role-based casting-solid selection and "
+ "contact-patch tagging only run through that pipeline.")
+ 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,
+ bodies=a.bodies_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']} bc_config {r['bc_config']} -> {r['vtu']}")
+ for w in r["ood_warnings"]:
+ print("WARNING:", w)
diff --git a/src/release_io.py b/src/release_io.py
new file mode 100644
index 0000000000000000000000000000000000000000..0f14bd86956c0077ec76dd11b8c1bb34a9a5a75e
--- /dev/null
+++ b/src/release_io.py
@@ -0,0 +1,131 @@
+"""Public-release artifact formats (Hugging Face model repo).
+
+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: 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.
+"""
+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."""
+ from safetensors import safe_open
+ with safe_open(path, framework="pt", device="cpu") as f:
+ meta = f.metadata() or {}
+ if meta.get("format") != GNN_FORMAT:
+ raise ValueError(f"{path}: not a {GNN_FORMAT} file "
+ f"(format={meta.get('format')!r})")
+ tensors = {k: f.get_tensor(k) for k in f.keys()}
+ norm_kinds = json.loads(meta["norm_kinds"])
+ model = {k[len("model."):]: v.to(device) for k, v in tensors.items()
+ if k.startswith("model.")}
+ norm = dict(json.loads(meta["norm_scalars"]))
+ for k, kind in norm_kinds.items():
+ t = tensors[f"norm.{k}"]
+ norm[k] = t.numpy() if kind == "numpy" else t
+ 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)
diff --git a/src/run_baseline.py b/src/run_baseline.py
new file mode 100644
index 0000000000000000000000000000000000000000..98fea41fbe61b8813e6c942e7d1eb1e05295fe23
--- /dev/null
+++ b/src/run_baseline.py
@@ -0,0 +1,150 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..78c5959482ca0291023e654bacdde5c34cea17df
--- /dev/null
+++ b/src/select_ckpt_ft2b.py
@@ -0,0 +1,331 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..c82b3eecbdf718d207b4b8d522e2a4ec0ef0c27c
--- /dev/null
+++ b/src/shrinkage.py
@@ -0,0 +1,901 @@
+"""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.
+"""
+import logging
+
+import numpy as np
+from scipy.sparse import csr_matrix
+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.
+SHRINKAGE_PCT = {
+ "steel": 2.5,
+ "low_alloy_steel": 2.5,
+ "ductile_iron": 0.75,
+ "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).
+_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.
+ """
+ pos = np.asarray(pos, dtype=np.float64)
+ tets = np.asarray(tets, dtype=np.int64)
+ field = np.asarray(field, dtype=np.float64)
+ n_nodes = pos.shape[0]
+
+ if tets.shape[0] == 0:
+ return np.zeros((n_nodes, 3))
+
+ p0 = pos[tets[:, 0]]
+ p1 = pos[tets[:, 1]]
+ p2 = pos[tets[:, 2]]
+ p3 = pos[tets[:, 3]]
+ f0 = field[tets[:, 0]]
+ f1 = field[tets[:, 1]]
+ f2 = field[tets[:, 2]]
+ f3 = field[tets[:, 3]]
+
+ e1 = p1 - p0
+ e2 = p2 - p0
+ e3 = p3 - p0
+
+ jac = np.stack([e1, e2, e3], axis=1) # (M,3,3), rows = edge vectors
+ rhs = np.stack([f1 - f0, f2 - f0, f3 - f0], axis=1) # (M,3)
+
+ vol = np.abs(np.einsum('ij,ij->i', e1, np.cross(e2, e3))) / 6.0 # (M,)
+
+ grad_tet = np.zeros_like(rhs)
+ nondeg = vol > 0
+ if nondeg.any():
+ grad_tet[nondeg] = np.linalg.solve(jac[nondeg], rhs[nondeg, :, None])[..., 0]
+
+ node_grad_sum = np.zeros((n_nodes, 3))
+ node_vol_sum = np.zeros(n_nodes)
+ weighted = grad_tet * vol[:, None]
+ for k in range(4):
+ idx = tets[:, k]
+ np.add.at(node_grad_sum, idx, weighted)
+ np.add.at(node_vol_sum, idx, vol)
+
+ 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")`.
+ """
+ t_s = np.asarray(t_s, dtype=np.float64)
+ grad = node_gradient(pos, tets, t_s)
+ grad_mag = np.linalg.norm(grad, axis=1)
+
+ t_safe = np.maximum(t_s, _EPS)
+ tdot_proxy = (TL - TS) / t_safe
+ 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).
+ """
+ ny = np.asarray(ny, dtype=np.float64)
+
+ 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.
+_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)."""
+ 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]]
+ e1, e2, e3 = p1 - p0, p2 - p0, p3 - p0
+ return np.abs(np.einsum('ij,ij->i', e1, np.cross(e2, e3))) / 6.0
+
+
+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.
+ """
+ idx = np.flatnonzero(mask)
+ labels_full = np.full(n_nodes, -1, dtype=np.int64)
+ if idx.size == 0:
+ return 0, labels_full
+
+ local_of = -np.ones(n_nodes, dtype=np.int64)
+ local_of[idx] = np.arange(idx.size)
+
+ rows_list, cols_list = [], []
+ for a, b in _TET_EDGE_PAIRS:
+ va, vb = tets[:, a], tets[:, b]
+ both = mask[va] & mask[vb]
+ if both.any():
+ rows_list.append(local_of[va[both]])
+ cols_list.append(local_of[vb[both]])
+ n_local = idx.size
+ if rows_list:
+ rows = np.concatenate(rows_list)
+ cols = np.concatenate(cols_list)
+ data = np.ones(rows.shape[0])
+ graph = csr_matrix((data, (rows, cols)), shape=(n_local, n_local))
+ else:
+ graph = csr_matrix((n_local, n_local))
+ n_comp, labels_local = connected_components(graph, directed=False)
+ labels_full[idx] = labels_local
+ return n_comp, labels_full
+
+
+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."""
+ tets = np.asarray(tets, dtype=np.int64)
+ if tets.shape[0] == 0:
+ return np.zeros((0, 3), dtype=np.int64)
+ face_idx = np.asarray(_TET_FACES, dtype=np.int64) # (4, 3)
+ all_faces = tets[:, face_idx].reshape(-1, 3) # (4*M, 3), original order
+ sorted_faces = np.sort(all_faces, axis=1)
+ _, first_idx, counts = np.unique(sorted_faces, axis=0, return_index=True,
+ return_counts=True)
+ boundary = counts == 1
+ return all_faces[first_idx[boundary]]
+
+
+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.
+ """
+ n_nodes = pos.shape[0]
+ is_member = np.zeros(n_nodes, dtype=bool)
+ is_member[np.asarray(nodes, dtype=np.int64)] = True
+ t_s = np.asarray(t_s, dtype=np.float64)
+ threshold = float(threshold)
+
+ area = 0.0
+ if tets.shape[0] > 0:
+ hot_mask = t_s >= threshold
+ member_per_tet = hot_mask[tets]
+ n_above = member_per_tet.sum(axis=1)
+ mixed = np.flatnonzero((n_above >= 1) & (n_above <= 3))
+ for ti in mixed:
+ tet = tets[ti]
+ if not is_member[tet].any():
+ continue # this crossing belongs to a different region
+ verts = pos[tet]
+ vals = t_s[tet]
+ above_local = member_per_tet[ti]
+ k = int(n_above[ti])
+ if k in (1, 3):
+ odd = int(np.flatnonzero(above_local)[0]) if k == 1 else \
+ int(np.flatnonzero(~above_local)[0])
+ others = [i for i in range(4) if i != odd]
+ pts = []
+ for o in others:
+ frac = (threshold - vals[odd]) / (vals[o] - vals[odd])
+ pts.append(verts[odd] + frac * (verts[o] - verts[odd]))
+ area += 0.5 * float(np.linalg.norm(np.cross(pts[1] - pts[0], pts[2] - pts[0])))
+ else: # k == 2
+ hi = np.flatnonzero(above_local)
+ lo = np.flatnonzero(~above_local)
+
+ def _cross_point(h, c, verts=verts, vals=vals):
+ frac = (threshold - vals[h]) / (vals[c] - vals[h])
+ return verts[h] + frac * (verts[c] - verts[h])
+
+ a_pt = _cross_point(hi[0], lo[0])
+ b_pt = _cross_point(hi[0], lo[1])
+ c_pt = _cross_point(hi[1], lo[1])
+ d_pt = _cross_point(hi[1], lo[0])
+ area += 0.5 * float(np.linalg.norm(np.cross(c_pt - a_pt, d_pt - b_pt)))
+
+ 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.
+ member_counts = is_member[boundary_faces].sum(axis=1).astype(np.float64) # (F,)
+ frac = member_counts / 3.0
+ nonzero = frac > 0
+ if nonzero.any():
+ v = boundary_faces[nonzero]
+ p0f, p1f, p2f = pos[v[:, 0]], pos[v[:, 1]], pos[v[:, 2]]
+ tri_areas = 0.5 * np.linalg.norm(np.cross(p1f - p0f, p2f - p0f), axis=1)
+ area += float(np.sum(frac[nonzero] * tri_areas))
+ return area
+
+
+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.
+ """
+ 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}")
+
+ pos = np.asarray(pos, dtype=np.float64)
+ tets = np.asarray(tets, dtype=np.int64)
+ t_s = np.asarray(t_s, dtype=np.float64)
+ n_nodes = pos.shape[0]
+
+ if n_nodes == 0 or tets.shape[0] == 0:
+ return []
+
+ threshold = np.percentile(t_s, 100.0 * (1.0 - top_frac))
+ hot_mask = t_s >= threshold
+ if not hot_mask.any():
+ return []
+
+ n_comp, labels = _induced_components(n_nodes, tets, hot_mask)
+
+ vol = _tet_volumes(pos, tets)
+ n_hot_per_tet = hot_mask[tets].sum(axis=1)
+ has_hot = n_hot_per_tet > 0
+ region_vol = np.zeros(n_comp)
+ if has_hot.any():
+ # any one hot vertex identifies the tet's (unique) component.
+ first_hot_col = np.argmax(hot_mask[tets], axis=1)
+ first_hot_node = tets[np.arange(tets.shape[0]), first_hot_col]
+ comp_of_tet = labels[first_hot_node]
+ weight = vol * (n_hot_per_tet / 4.0)
+ np.add.at(region_vol, comp_of_tet[has_hot], weight[has_hot])
+
+ regions = []
+ for c in range(n_comp):
+ nodes = np.sort(np.flatnonzero(labels == c))
+ regions.append({
+ "nodes": nodes,
+ "centroid": pos[nodes].mean(axis=0),
+ "volume_mm3": float(region_vol[c]),
+ "t_max": float(t_s[nodes].max()),
+ })
+ regions.sort(key=lambda r: r["t_max"], reverse=True)
+ 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.
+ 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.
+ """
+ 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)
+ region_nodes = np.asarray(region["nodes"], dtype=np.int64)
+
+ if feed_nodes.size == 0:
+ return False
+ if tets.shape[0] == 0:
+ return False
+ region_t_max = float(region["t_max"])
+ if region_t_max <= 0:
+ logger.debug(
+ "is_fed: region t_max=%.6g <= 0 -- corridor threshold would be "
+ "vacuous (every non-negative t_s qualifies); reporting unfed.",
+ region_t_max)
+ return False
+
+ n_nodes = t_s.shape[0]
+ threshold = corridor_frac * region_t_max
+ warm_mask = t_s >= threshold
+
+ _, labels = _induced_components(n_nodes, tets, warm_mask)
+
+ region_labels = set(int(l) for l in labels[region_nodes] if l >= 0)
+ if not region_labels:
+ return False
+ feed_labels = set(int(l) for l in labels[feed_nodes] if l >= 0)
+ return len(region_labels & feed_labels) > 0
+
+
+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`).
+ """
+ regions = hot_regions(pos, tets, t_s)
+ out = []
+ for r in regions:
+ if is_fed(r, t_s, tets, feed_nodes):
+ continue
+ severity = shrinkage_pct / 100.0 * r["volume_mm3"]
+ out.append({**r, "severity": severity})
+ return out
+
+
+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.
+ """
+ pos = np.asarray(pos, dtype=np.float64)
+ gravity_vec = np.asarray(gravity_vec, dtype=np.float64)
+ up = -gravity_vec
+ norm = np.linalg.norm(up)
+ if norm > 0:
+ up = up / norm
+ surface_set = set(np.asarray(surface_nodes, dtype=np.int64).tolist())
+
+ sinks = []
+ for r in regions:
+ nodes = np.asarray(r["nodes"], dtype=np.int64)
+ proj = pos[nodes] @ up
+ top_node = int(nodes[np.argmax(proj)])
+ if top_node in surface_set:
+ sinks.append({
+ "region_id": r["region_id"],
+ "location": pos[top_node].copy(),
+ "volume_mm3": r["volume_mm3"],
+ })
+ return sinks
+
+
+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.
+ """
+ pos = np.asarray(pos, dtype=np.float64)
+ tets = np.asarray(tets, dtype=np.int64)
+ t_s = np.asarray(t_s, dtype=np.float64)
+ gravity_vec = np.asarray(gravity_vec, dtype=np.float64)
+ up = -gravity_vec
+ norm = np.linalg.norm(up)
+ if norm > 0:
+ up = up / norm
+ surface_nodes = np.asarray(surface_nodes, dtype=np.int64)
+ if surface_nodes.size == 0:
+ raise ValueError(
+ "feeder_recommendations: surface_nodes is empty -- no candidate "
+ "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
+
+ recs = []
+ n_skipped_sub_resolution = 0
+ for r in regions:
+ if r.get("sub_resolution", False):
+ n_skipped_sub_resolution += 1
+ continue
+ nodes = np.asarray(r["nodes"], dtype=np.int64)
+ volume = r["volume_mm3"]
+ threshold = r["hot_threshold"]
+ area = _region_surface_area(pos, tets, nodes, t_s, threshold,
+ boundary_faces=boundary_faces)
+ if area > 0 and volume > 0:
+ region_modulus = volume / area
+ else:
+ logger.warning(
+ "feeder_recommendations: region_id=%s has degenerate "
+ "volume/area (V=%.6g mm^3, A=%.6g mm^2) -- min_modulus_mm is "
+ "not meaningful and is reported as 0.0",
+ r.get("region_id"), volume, area)
+ region_modulus = 0.0
+
+ centroid = np.asarray(r["centroid"], dtype=np.float64)
+ rel = surface_pos - centroid
+ axis_comp = rel @ up
+ perp = rel - np.outer(axis_comp, up)
+ perp_dist = np.linalg.norm(perp, axis=1)
+
+ above = np.flatnonzero(axis_comp >= 0)
+ candidates = above if above.size > 0 else np.arange(surface_pos.shape[0])
+ best = candidates[np.argmin(perp_dist[candidates])]
+
+ recs.append({
+ "location": surface_pos[best].copy(),
+ "min_modulus_mm": 1.2 * region_modulus,
+ "region_id": r["region_id"],
+ })
+ if n_skipped_sub_resolution:
+ logger.info(
+ "feeder_recommendations: skipped %d sub-resolution region(s) (<4 "
+ "nodes) -- V/A modulus is not meaningful at this resolution",
+ n_skipped_sub_resolution)
+ 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"}}`.
+ """
+ if alloy_family not in SHRINKAGE_PCT:
+ raise KeyError(
+ f"screen_casting: unknown alloy_family {alloy_family!r} -- known "
+ f"families: {sorted(SHRINKAGE_PCT)}"
+ )
+
+ pos = np.asarray(pos, dtype=np.float64)
+ tets = np.asarray(tets, dtype=np.int64)
+ surface_nodes = np.asarray(surface_nodes, dtype=np.int64)
+ t_s = np.asarray(t_s, dtype=np.float64)
+ 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"
+
+ 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,
+ "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
new file mode 100644
index 0000000000000000000000000000000000000000..8eb3731af70b798a24fb6ffac9d85cd8f4dee98d
--- /dev/null
+++ b/src/splits.py
@@ -0,0 +1,61 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..39e07a37c3420617f0a7e99eba438b5857d9a066
--- /dev/null
+++ b/src/step_mesh.py
@@ -0,0 +1,45 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..505709de4b1acde633ac7e293e0cf5fe5ef7705d
--- /dev/null
+++ b/src/targets.py
@@ -0,0 +1,39 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..25c7ce14d0a763e41a45392d095cf1ccf61db9bf
--- /dev/null
+++ b/src/train_ft2b.py
@@ -0,0 +1,627 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..fa052a7a5ebd6b373819ed372497f11520ca64ba
--- /dev/null
+++ b/src/train_gnn.py
@@ -0,0 +1,300 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..f4b1bcd7d1a102089c115af6b3efced4c99b20a5
--- /dev/null
+++ b/src/train_heads.py
@@ -0,0 +1,278 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..4ad0fd885455d6c07201812481db8c0b312d3c37
--- /dev/null
+++ b/src/utils.py
@@ -0,0 +1,115 @@
+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/viz.py b/src/viz.py
new file mode 100644
index 0000000000000000000000000000000000000000..8a50ad9a9fed5a3a5aafb57065b1620f3ce842a1
--- /dev/null
+++ b/src/viz.py
@@ -0,0 +1,90 @@
+"""PyVista helpers for the solidification viewer. Pure/testable: load a
+prediction VTU and derive the ghost mesh, the still-liquid region above a time
+threshold, and the hot-spot points."""
+import os
+
+import numpy as np
+import pyvista as pv
+
+
+def load_prediction(vtu_path: str) -> pv.UnstructuredGrid:
+ grid = pv.read(vtu_path)
+ for name in ("solidification_time", "hotspot"):
+ if name not in grid.point_data:
+ raise ValueError(f"VTU missing point array '{name}'")
+ return grid
+
+
+def time_range(grid):
+ t = np.asarray(grid.point_data["solidification_time"])
+ return float(t.min()), float(t.max())
+
+
+def liquid_region(grid, tau: float):
+ # cells still liquid at time tau (solidification_time >= tau)
+ return grid.threshold(tau, scalars="solidification_time")
+
+
+def hotspot_points(grid):
+ mask = np.asarray(grid.point_data["hotspot"]) > 0.5
+ return pv.PolyData(np.asarray(grid.points)[mask])
+
+
+def points_above_time(grid, t: float) -> pv.PolyData:
+ """Points whose solidification_time is >= t. Used to thin out the hot-spot
+ cloud with a user-controlled cutoff/threshold slider."""
+ times = np.asarray(grid.point_data["solidification_time"])
+ mask = times >= t
+ pts = pv.PolyData(np.asarray(grid.points)[mask])
+ pts.point_data["solidification_time"] = times[mask]
+ return pts
+
+
+def solidification_animation(grid, out_path, *, n_frames=40, fps=10,
+ cmap="turbo_r", ghost_opacity=0.15,
+ scalar_bar_args=None, bg="#f7f7f5",
+ window_size=(1280, 720), camera_position=None):
+ """Render an animated GIF of progressive solidification: tau sweeps from
+ tmin to tmax; already-solidified material (t < tau) shows as a low-opacity
+ ghost of the full body, still-liquid material (t >= tau) is opaque and
+ colored by solidification time. Each frame shows the color legend and the
+ current time. Returns out_path. Off-screen; independent of any live view,
+ except that when `camera_position` is given (e.g. captured from a live
+ plotter) it is used verbatim so the GIF renders from that exact viewpoint;
+ otherwise the camera is auto-framed via reset_camera()."""
+ import pyvista as pv
+ tmin, tmax = time_range(grid)
+ taus = np.linspace(tmin, tmax, int(n_frames))
+ p = pv.Plotter(off_screen=True, window_size=list(window_size))
+ p.set_background(bg)
+ try:
+ # permanent ghost of the full body (owns the one legend)
+ p.add_mesh(grid, scalars="solidification_time", opacity=ghost_opacity,
+ cmap=cmap, clim=[tmin, tmax], smooth_shading=True,
+ show_scalar_bar=True, scalar_bar_args=dict(scalar_bar_args or {}))
+ p.add_mesh(liquid_region(grid, tmin), scalars="solidification_time",
+ cmap=cmap, clim=[tmin, tmax], smooth_shading=True,
+ name="liquid", show_scalar_bar=False)
+ if camera_position is not None:
+ p.camera_position = camera_position
+ else:
+ p.reset_camera()
+ p.open_gif(out_path, fps=int(fps))
+ for tau in taus:
+ liq = liquid_region(grid, float(tau))
+ if liq.n_cells > 0:
+ p.add_mesh(liq, scalars="solidification_time", cmap=cmap,
+ clim=[tmin, tmax], smooth_shading=True,
+ name="liquid", show_scalar_bar=False)
+ else:
+ p.remove_actor("liquid") # fully solidified -> only the ghost
+ p.add_text(f"t = {float(tau):6.1f} s", name="tlabel",
+ position="upper_left", font_size=14, color="black")
+ p.write_frame()
+ except Exception:
+ if os.path.exists(out_path):
+ os.unlink(out_path) # don't leave a truncated/corrupt gif behind
+ raise
+ finally:
+ p.close()
+ return out_path
diff --git a/src/warm_start.py b/src/warm_start.py
new file mode 100644
index 0000000000000000000000000000000000000000..05ceeb0d68d7648280132200b934ba151c17478d
--- /dev/null
+++ b/src/warm_start.py
@@ -0,0 +1,239 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..253c82f10d51ddf9235aebece1a7eff313e93c4c
--- /dev/null
+++ b/src/webapp.py
@@ -0,0 +1,927 @@
+"""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 =
+ghost, liquid = opaque) and a hot-spot toggle. Run: python -m src.webapp
+
+Part B Task 5 adds an assembly mode: a multi-solid STEP upload shows a body
+table (role + params per solid) and a role-colored 3-D preview instead of
+auto-running the prediction; a single-solid upload keeps the original flow
+untouched."""
+import csv
+import os
+import os.path as osp
+
+import numpy as np
+import pyvista as pv
+from trame.app import get_server
+from trame.ui.vuetify3 import SinglePageWithDrawerLayout
+from trame.widgets import html, vuetify3
+from pyvista.trame.ui import plotter_ui
+
+from src import viz
+from src.assembly import list_solids, validate_bodies
+from src.inference_props import DEFAULT_ALLOY, list_alloys, list_molds
+from src.predict_casting import predict_casting
+from tools.doe.load_properties import load_chill_effusivities
+
+DEFAULT_MOLD = "green_sand"
+SLEEVE_REGISTRY_CSV = "tools/sleeves/sleeve_registry.csv"
+
+pv.OFF_SCREEN = True
+
+GHOST_OPACITY = 0.15
+DEFAULT_HOTSPOT_SIZE = 12.0
+HOTSPOT_CUTOFF_PERCENTILE = 90.0
+
+# --- visual identity -------------------------------------------------------
+# "Neuralcast" = a neural-network surrogate for a foundry casting process, so
+# the palette stays in the same molten-metal family as before but is rebuilt
+# for a light, legible workbench: warm parchment chrome, white control
+# surfaces, and a near-white VTK stage so the inferno field is what draws the
+# eye. The one recurring accent is a deeper burnt-orange (readable at text
+# contrast on white, unlike the original pale version). Wired into vuetify's
+# theme system so every stock control (slider, checkbox, buttons) inherits it.
+# The app bar is a deliberate exception: a dark charcoal strip (independent of
+# the light theme) so the header reads as instrument-panel chrome holding just
+# the title and a status line, not another content surface.
+BG_APP = "#f4f2ee" # app chrome background (warm parchment)
+SURFACE = "#ffffff" # drawer / card surfaces
+BORDER = "#ddd8cf" # hairline dividers
+VTK_BG = "#f7f7f5" # 3D viewport background -- near-white so the
+ # colored field is the highest-contrast thing on screen
+ACCENT_MOLTEN = "#c1440e" # burnt-orange -- the one recurring accent (readable on white)
+TEXT_PRIMARY = "#23262b"
+TEXT_MUTED = "#6b7078"
+HEADER_BG = "#1c1a17" # dark charcoal app-bar -- deliberately dark in the light theme
+HEADER_TEXT = "#f4f2ee"
+HEADER_TEXT_MUTED = "#b6afa1"
+FIELD_CMAP = "turbo_r" # reversed turbo (red -> blue): wide, high-contrast gradient
+ # with far more distinguishable steps than a single-hue map,
+ # so small time differences across the field stay legible
+HOTSPOT_COLOR = "#0092b8" # deep cyan/teal marker: falls in the middle of turbo's
+ # range but is rendered as small spheres (not a field
+ # color), so it still reads as an overlay/alert
+
+# --- assembly mode (Part B Task 5) -----------------------------------------
+# One row per solid in a multi-body STEP upload: role dropdown (defaults
+# largest solid = casting, everything else = ignore) + role-dependent params
+# (chill material / sleeve name) feeding `src.assembly.validate_bodies` and
+# `predict_casting(..., bodies=...)`. The 3-D preview colors each solid's
+# coarse tessellation by its currently-selected role; "ignore" bodies stay
+# the casting's neutral grey but at the same faint GHOST_OPACITY used for the
+# solidified-region ghost actor elsewhere in this file.
+ROLE_OPTIONS = ["casting", "chill", "insulating_sleeve", "exo_sleeve", "ignore"]
+ROLE_COLORS = {
+ "casting": "#9e9e9e",
+ "chill": "#2196f3",
+ "insulating_sleeve": "#ffc107",
+ "exo_sleeve": "#f44336",
+ "ignore": "#9e9e9e",
+}
+
+_MONO_STACK = ('"JetBrains Mono", ui-monospace, "SFMono-Regular", '
+ "Menlo, Consolas, monospace")
+_COND_STACK = '"Roboto Condensed", "Arial Narrow", sans-serif'
+
+# Scalar bar (legend) for the solidification-time color field: a vertical bar
+# pinned to the right edge of the viewport, with generous title/label padding
+# so the title never collides with the tick numbers, and integer tick labels
+# (fractional seconds add no information at casting scale).
+SCALAR_BAR_ARGS = dict(
+ title="Solidification Time, s",
+ fmt="%.0f",
+ color="black",
+ bold=True,
+ font_family="arial",
+ vertical=True,
+ position_x=0.91,
+ position_y=0.06,
+ width=0.055,
+ height=0.80,
+ title_font_size=20,
+ label_font_size=16,
+ n_labels=6,
+)
+
+_CUSTOM_CSS = f"""
+.neuralcast-title {{
+ display: flex;
+ align-items: center;
+ gap: 10px;
+}}
+.neuralcast-title-text {{
+ font-family: {_COND_STACK};
+ font-weight: 700;
+ letter-spacing: .04em;
+ text-transform: uppercase;
+ font-size: 1.05rem;
+ color: {HEADER_TEXT};
+}}
+.neuralcast-status {{
+ font-family: {_COND_STACK};
+ letter-spacing: .08em;
+ text-transform: uppercase;
+ font-size: .7rem;
+ color: {HEADER_TEXT_MUTED};
+ white-space: nowrap;
+}}
+.neuralcast-drawer-content {{
+ overflow-y: auto;
+}}
+.neuralcast-section-label {{
+ font-family: {_COND_STACK};
+ letter-spacing: .12em;
+ text-transform: uppercase;
+ font-size: .7rem;
+ color: {TEXT_MUTED};
+ margin-bottom: 10px;
+ flex: none;
+}}
+.neuralcast-section {{
+ margin-bottom: 4px;
+}}
+.neuralcast-section + .neuralcast-section {{
+ border-top: 1px solid {BORDER};
+ padding-top: 20px;
+ margin-top: 20px;
+}}
+.neuralcast-tau-badge {{
+ display: inline-block;
+ font-family: {_MONO_STACK};
+ font-size: .75rem;
+ color: {ACCENT_MOLTEN};
+ background: {BG_APP};
+ border: 1px solid {BORDER};
+ border-radius: 4px;
+ padding: 3px 10px;
+ white-space: nowrap;
+}}
+.neuralcast-summary-box {{
+ min-height: 0;
+}}
+.neuralcast-summary-box .v-input {{
+ flex: 1 1 auto;
+ min-height: 0;
+}}
+.neuralcast-summary-box .v-field__input,
+.neuralcast-summary-box textarea {{
+ font-family: {_MONO_STACK} !important;
+ font-size: .78rem;
+ line-height: 1.5;
+ color: {TEXT_PRIMARY};
+}}
+.neuralcast-footer-text {{
+ width: 100%;
+ text-align: center;
+ font-family: {_COND_STACK};
+ letter-spacing: .06em;
+ font-size: .72rem;
+ color: {TEXT_MUTED};
+}}
+"""
+
+
+def summary_text(result: dict, alloy: str = None, mold: str = None) -> str:
+ lines = [f"casting: {result['name']}",
+ f"total solidification time: {result['total_time_s']:.1f} s",
+ f"nodes: {result['n_nodes']}"]
+ if alloy is not None or mold is not None:
+ cond_parts = []
+ if alloy is not None:
+ cond_parts.append(f"alloy: {alloy}")
+ if mold is not None:
+ cond_parts.append(f"mold: {mold}")
+ lines.append(" " + ", ".join(cond_parts))
+ lines.append("top hot spots (x,y,z mm):")
+ for xyz in result["hotspot_xyz"]:
+ lines.append(f" ({xyz[0]:.1f}, {xyz[1]:.1f}, {xyz[2]:.1f})")
+ if result.get("ood_warnings"):
+ lines.append("WARNINGS:")
+ lines += [f" - {w}" for w in result["ood_warnings"]]
+ return "\n".join(lines)
+
+
+INITIAL_SUMMARY = "(upload a STEP file)"
+
+
+def _run_prediction(step_path, out_dir, alloy, mold, bodies=None):
+ """Thin wrapper around `predict_casting` forwarding alloy/mold conditioning
+ and (Part B Task 5) an optional multi-body role assignment.
+
+ Extracted so the upload handler's predict call is unit-testable without
+ driving the full trame upload flow (module-level lookup of
+ `predict_casting` so tests can monkeypatch `src.webapp.predict_casting`).
+ `bodies` is only forwarded as a keyword when given (not merely `None`),
+ so the pre-Part-B single-solid call shape (no `bodies` kwarg at all)
+ stays byte-identical -- some callers (this module's own existing tests)
+ monkeypatch `predict_casting` with a fake that doesn't accept `bodies`.
+ """
+ if bodies is not None:
+ return predict_casting(step_path, out_dir=out_dir, alloy=alloy, mold=mold, bodies=bodies)
+ return predict_casting(step_path, out_dir=out_dir, alloy=alloy, mold=mold)
+
+
+def _load_sleeve_names_by_type(csv_path=SLEEVE_REGISTRY_CSV):
+ """{"exo": [name, ...], "insulating": [name, ...]} from
+ `tools/sleeves/sleeve_registry.csv`'s `name`/`type` columns, sorted --
+ the source of the sleeve-name dropdown options for the assembly panel,
+ pre-filtered by role so the UI never offers a sleeve `validate_bodies`
+ would reject for a type mismatch (see `src.assembly.validate_bodies`'s
+ `_ROLE_TO_SLEEVE_TYPE` cross-check)."""
+ out: dict[str, list[str]] = {"exo": [], "insulating": []}
+ try:
+ with open(csv_path, newline="") as f:
+ for row in csv.DictReader(f):
+ names = out.setdefault(row["type"], [])
+ names.append(row["name"])
+ except FileNotFoundError:
+ pass
+ for names in out.values():
+ names.sort()
+ return out
+
+
+def _assemble_bodies_from_state(state_like) -> list:
+ """Pure builder: the assembly panel's per-row state (`state_like.
+ assembly_bodies`, one dict per solid with at least `solid_id` and
+ `role`, plus an optional `params`) -> the plain `bodies` list
+ `src.assembly.validate_bodies`/`predict_casting(..., bodies=...)`
+ expect (`{"solid_id", "role", "params"}`, UI-only fields like `name`/
+ `volume_mm3` dropped). Takes anything with attribute access
+ (`state_like.assembly_bodies`) so it is testable with a plain
+ `SimpleNamespace`/fake as well as the real trame state -- no gmsh, no
+ validation (that's `validate_bodies`'s job, called separately as a
+ pre-flight check by the Run-assembly handler)."""
+ rows = state_like.assembly_bodies or []
+ bodies = []
+ for row in rows:
+ body = {"solid_id": row["solid_id"], "role": row.get("role", "ignore")}
+ params = row.get("params")
+ if params:
+ body["params"] = params
+ bodies.append(body)
+ return bodies
+
+
+def _solid_preview_meshes(step_path):
+ """Coarse per-solid surface tessellation for the assembly 3-D preview:
+ one fresh gmsh/OCC session imports `step_path` and generates a SINGLE
+ 2-D surface mesh over the whole model, but sized PER SOLID (each
+ solid's own boundary points get their own target size, its OWN bbox
+ diagonal / 15, via `gmsh.model.mesh.setSize` + gmsh's default
+ `Mesh.MeshSizeExtendFromBoundary` sizing field) -- a large casting
+ tessellates at its own coarse scale while a small chill/sleeve still
+ gets enough triangles to read as a shape, instead of one uniform size
+ forcing every solid into whichever end of the size range is wrong for
+ it (a small uniform size on a big casting is a triangle-count blowup;
+ a uniform size sized for the casting leaves a small role body a
+ handful of degenerate facets). Returns one `(solid_id,
+ pyvista.PolyData)` per 3-D solid (its own boundary surfaces' triangles
+ + nodes). `solid_id` matches what `list_solids` reports for the SAME
+ STEP file, since both walk a fresh `gmsh.model.occ.importShapes` +
+ `gmsh.model.getEntities(3)` and gmsh assigns solid tags
+ deterministically from import order.
+
+ Real-gmsh, no mocking -- callers should expect this to be slow-ish (a
+ full 2-D mesh pass) but bounded for the coarse sizes used here."""
+ import gmsh
+
+ gmsh.initialize()
+ try:
+ gmsh.option.setNumber("General.Terminal", 0)
+ gmsh.model.add("assembly_preview")
+ gmsh.model.occ.importShapes(step_path)
+ gmsh.model.occ.synchronize()
+
+ solids = gmsh.model.getEntities(3)
+ if not solids:
+ return []
+
+ size_by_tag = {}
+ for (_dim, tag) in solids:
+ xmin, ymin, zmin, xmax, ymax, zmax = gmsh.model.getBoundingBox(3, tag)
+ diag = float(np.linalg.norm([xmax - xmin, ymax - ymin, zmax - zmin]))
+ size_by_tag[tag] = max(diag / 15.0, 1e-3)
+
+ # Global bounds only fence in gmsh's implicit sizing AWAY from any
+ # explicitly-sized point (e.g. deep in a solid's interior); the
+ # per-solid setSize calls below are what actually drive each
+ # solid's own coarse scale.
+ gmsh.option.setNumber("Mesh.MeshSizeMin", min(size_by_tag.values()))
+ gmsh.option.setNumber("Mesh.MeshSizeMax", max(size_by_tag.values()))
+ gmsh.option.setNumber("Mesh.MeshSizeExtendFromBoundary", 1)
+ for (_dim, tag) in solids:
+ boundary_pts = gmsh.model.getBoundary([(3, tag)], recursive=True)
+ point_dimtags = [(d, t) for (d, t) in boundary_pts if d == 0]
+ if point_dimtags:
+ gmsh.model.mesh.setSize(point_dimtags, size_by_tag[tag])
+ gmsh.model.mesh.generate(2)
+
+ out = []
+ for (_dim, tag) in solids:
+ boundary = gmsh.model.getBoundary([(3, tag)], oriented=False)
+ surf_tags = [t for (d, t) in boundary if d == 2]
+ node_xyz_by_id = {}
+ tri_blocks = []
+ for surf_tag in surf_tags:
+ node_tags, coords, _param = gmsh.model.mesh.getNodes(
+ 2, surf_tag, includeBoundary=True)
+ coords = np.asarray(coords, dtype=np.float64).reshape(-1, 3)
+ for nid, xyz in zip(node_tags, coords):
+ node_xyz_by_id.setdefault(int(nid), xyz)
+ elem_types, _elem_tags, elem_node_tags = gmsh.model.mesh.getElements(2, surf_tag)
+ for etype, enodes in zip(elem_types, elem_node_tags):
+ if etype != 2: # gmsh type 2 = 3-node triangle
+ continue
+ tri_blocks.append(np.asarray(enodes, dtype=np.int64).reshape(-1, 3))
+ if not tri_blocks or not node_xyz_by_id:
+ continue
+ tris = np.concatenate(tri_blocks, axis=0)
+ node_ids = sorted(node_xyz_by_id)
+ local_index = {nid: i for i, nid in enumerate(node_ids)}
+ points = np.array([node_xyz_by_id[nid] for nid in node_ids], dtype=np.float64)
+ local_tris = np.vectorize(local_index.get)(tris)
+ faces = np.hstack(
+ [np.full((local_tris.shape[0], 1), 3, dtype=np.int64), local_tris])
+ out.append((tag, pv.PolyData(points, faces)))
+ return out
+ finally:
+ gmsh.finalize()
+
+
+def _rebuild_hotspots(app):
+ """(Re)build the hot-spot point actor from the current cutoff/size state.
+ Called on load and whenever the size or cutoff slider moves."""
+ plotter, server = app["plotter"], app["server"]
+ state, ctrl = server.state, server.controller
+ if app["grid"] is None:
+ return
+ pts = viz.points_above_time(app["grid"], float(state.hotspot_cutoff))
+ actor = plotter.add_mesh(pts, color=HOTSPOT_COLOR, point_size=float(state.hotspot_size),
+ name="hotspots", render_points_as_spheres=True)
+ actor.SetVisibility(bool(state.show_hotspots))
+ app["hotspot_actor"] = actor
+ ctrl.view_update()
+
+
+def _remove_scalar_bar(plotter):
+ """Drop any existing scalar bar(s), regardless of how many are present.
+ Guarded because pyvista raises if none exist (or if a title is required
+ with more than one) -- callers just want a clean slate either way."""
+ try:
+ plotter.remove_scalar_bar()
+ except Exception:
+ pass
+ # Belt-and-suspenders: clear() takes the whole registry in one shot, in
+ # case remove_scalar_bar() above only got one of several.
+ try:
+ plotter.scalar_bars.clear()
+ except Exception:
+ pass
+
+
+def _render_assembly_preview(app):
+ """(Re)draw every cached assembly-preview mesh (`app["assembly_meshes"]`,
+ built once by `_enter_assembly_mode`), colored/opacity'd by each body's
+ CURRENT role in `state.assembly_bodies` (see `ROLE_COLORS`/
+ `GHOST_OPACITY`). Called on entering assembly mode and again whenever
+ `assembly_bodies` changes (role edits) -- cheap, since it only touches
+ actors, never re-tessellates."""
+ plotter, server = app["plotter"], app["server"]
+ state, ctrl = server.state, server.controller
+ plotter.clear_actors()
+ _remove_scalar_bar(plotter)
+ roles = {b["solid_id"]: b.get("role", "ignore") for b in (state.assembly_bodies or [])}
+ meshes = app.get("assembly_meshes") or []
+ for solid_id, mesh in meshes:
+ role = roles.get(solid_id, "ignore")
+ color = ROLE_COLORS.get(role, ROLE_COLORS["ignore"])
+ opacity = GHOST_OPACITY if role == "ignore" else 1.0
+ plotter.add_mesh(mesh, color=color, opacity=opacity,
+ name=f"assembly_solid_{solid_id}", smooth_shading=True)
+ if meshes:
+ plotter.reset_camera()
+ ctrl.view_update()
+
+
+def _enter_assembly_mode(app, step_path, rows):
+ """Switch the viewer into assembly mode for a multi-solid STEP upload:
+ builds the role-colored 3-D preview meshes FIRST (`_solid_preview_meshes`
+ into `app["assembly_meshes"]`), THEN seeds `state.assembly_bodies` from
+ `list_solids`' rows (default roles -- largest solid = casting, everything
+ else = ignore, per the brief) and renders (`_render_assembly_preview`).
+ This order matters: `state.assembly_bodies` is watched
+ (`_on_assembly_bodies_change` -> `_render_assembly_preview`), so setting
+ it BEFORE `app["assembly_meshes"]` is refreshed would briefly recolor
+ whatever meshes happened to be cached from a PRIOR upload (e.g. a
+ multi -> single -> multi upload sequence) against the new roles -- a
+ stale-preview flash. A preview failure (e.g. a STEP gmsh can list solids
+ for but not re-tessellate) is surfaced in `state.assembly_error` rather
+ than raised -- the body table/role panel still works without a 3-D
+ preview."""
+ server = app["server"]
+ state = server.state
+ app["grid"] = None
+ app["hotspot_actor"] = None
+ state.loaded = False
+ state.summary = INITIAL_SUMMARY
+ state.assembly_step_path = step_path
+ state.assembly_error = ""
+
+ try:
+ app["assembly_meshes"] = _solid_preview_meshes(step_path)
+ except Exception as e:
+ app["assembly_meshes"] = []
+ state.assembly_error = f"3D preview failed: {e}"
+
+ bodies = []
+ for i, row in enumerate(rows):
+ bodies.append({
+ "solid_id": row["solid_id"],
+ "name": row["name"] or f"solid {row['solid_id']}",
+ "volume_mm3": row["volume_mm3"],
+ "role": "casting" if i == 0 else "ignore",
+ "params": {},
+ })
+ state.assembly_mode = True
+ state.assembly_bodies = bodies
+ _render_assembly_preview(app)
+
+
+def _process_upload(app, step_path, alloy, mold):
+ """Core upload logic, extracted from the `uploaded_file` state-change
+ handler so it is testable without a browser/file-upload round trip:
+ inspects `step_path` via `list_solids`; >1 solid enters assembly mode
+ (`_enter_assembly_mode` -- role panel + 3-D preview, prediction deferred
+ to the Run-assembly button); exactly 1 solid (or a STEP `list_solids`
+ itself can't parse -- fails soft here so it falls through to the SAME
+ meshing/healing pipeline `predict_casting` always used) runs the
+ original single-solid predict-and-render flow, byte-identical to
+ pre-Part-B behavior."""
+ server = app["server"]
+ state = server.state
+ try:
+ rows = list_solids(step_path)
+ except Exception:
+ rows = [] # fall back to the pre-Part-B single-body flow untouched
+
+ if len(rows) > 1:
+ _enter_assembly_mode(app, step_path, rows)
+ return
+
+ # Single-solid (or unlistable) upload: leave no assembly-mode residue
+ # behind -- mirrors clear_casting()'s reset of the same four fields, so
+ # a LATER multi-solid upload starts from a clean slate rather than
+ # _enter_assembly_mode briefly recoloring stale meshes/rows from a
+ # previous assembly session (see _enter_assembly_mode's docstring).
+ state.assembly_mode = False
+ state.assembly_bodies = []
+ state.assembly_error = ""
+ state.assembly_step_path = None
+ app["assembly_meshes"] = []
+ plotter = app["plotter"]
+ out_dir = "runs/castings"
+ try:
+ result = _run_prediction(step_path, out_dir, alloy, mold)
+ grid = viz.load_prediction(result["vtu"])
+ app["name"] = result["name"]
+ render_prediction(app, grid, summary_text(result, alloy=alloy, mold=mold))
+ except Exception as e: # keep the UI alive on bad/unmeshable STEP
+ plotter.clear_actors()
+ _remove_scalar_bar(plotter)
+ app["grid"] = None
+ app["hotspot_actor"] = None
+ state.loaded = False
+ state.summary = f"Prediction failed for {osp.basename(step_path)}:\n{e}"
+
+
+def render_prediction(app, grid, summary):
+ plotter, server = app["plotter"], app["server"]
+ state, ctrl = server.state, server.controller
+ plotter.clear_actors()
+ _remove_scalar_bar(plotter)
+ tmin, tmax = viz.time_range(grid)
+ app["grid"] = grid
+ # The "ghost" actor is the single, permanent owner of the scalar bar/legend:
+ # it is added once here and never touched again by the tau slider, so the
+ # legend persists across playback without being re-added or dropped. The
+ # "liquid" actor (re-added on every tau change) must NEVER own a scalar bar,
+ # or scrubbing the slider would repeatedly add/replace it.
+ plotter.add_mesh(grid, scalars="solidification_time", opacity=GHOST_OPACITY,
+ cmap=FIELD_CMAP, name="ghost", clim=[tmin, tmax],
+ smooth_shading=True, show_scalar_bar=True,
+ scalar_bar_args=SCALAR_BAR_ARGS)
+ plotter.add_mesh(viz.liquid_region(grid, tmin), scalars="solidification_time",
+ cmap=FIELD_CMAP, name="liquid", clim=[tmin, tmax],
+ smooth_shading=True, show_scalar_bar=False)
+ state.tmin, state.tmax, state.tau = tmin, tmax, tmin
+ times = np.asarray(grid.point_data["solidification_time"])
+ state.hotspot_cutoff = float(np.percentile(times, HOTSPOT_CUTOFF_PERCENTILE))
+ _rebuild_hotspots(app)
+ state.anim_ready = False
+ state.anim_uri = ""
+ state.summary, state.loaded = summary, True
+ plotter.reset_camera()
+ ctrl.view_update()
+
+
+def build_app(server=None):
+ if server is None:
+ server = get_server(client_type="vue3")
+ state, ctrl = server.state, server.controller
+ plotter = pv.Plotter()
+ plotter.set_background(VTK_BG)
+ app = {"server": server, "plotter": plotter, "grid": None, "hotspot_actor": None,
+ "name": "casting"}
+ server.app = app # exposes the app dict for tests / headless callers that
+ # need render_prediction()/clear_casting() without the UI
+
+ state.trame__title = "Neuralcast Viewer"
+
+ # Custom light vuetify theme (molten-metal accent) applied through
+ # Vuetify's own theme system, so every stock control (slider, checkbox,
+ # buttons) picks up the palette automatically instead of one-off color
+ # props.
+ vuetify_config = {
+ "theme": {
+ "defaultTheme": "neuralcastLight",
+ "themes": {
+ "neuralcastLight": {
+ "dark": False,
+ "colors": {
+ "background": BG_APP,
+ "surface": SURFACE,
+ "primary": ACCENT_MOLTEN,
+ "secondary": TEXT_MUTED,
+ "on-background": TEXT_PRIMARY,
+ "on-surface": TEXT_PRIMARY,
+ },
+ },
+ },
+ },
+ }
+ sleeve_names_by_type = _load_sleeve_names_by_type()
+ for k, v in dict(tau=0.0, tmin=0.0, tmax=1.0, show_hotspots=False,
+ summary=INITIAL_SUMMARY, loaded=False, uploaded_file=None,
+ loading=False, hotspot_size=DEFAULT_HOTSPOT_SIZE,
+ hotspot_cutoff=0.0, animating=False, anim_ready=False,
+ anim_uri="", anim_name="solidification.gif",
+ alloy=DEFAULT_ALLOY, mold=DEFAULT_MOLD,
+ alloy_options=list_alloys(), mold_options=list_molds(),
+ # Assembly mode (Part B Task 5): populated by
+ # _enter_assembly_mode() on a multi-solid STEP upload.
+ assembly_mode=False, assembly_bodies=[], assembly_error="",
+ assembly_step_path=None, role_options=ROLE_OPTIONS,
+ chill_material_options=sorted(load_chill_effusivities()),
+ exo_sleeve_options=sleeve_names_by_type["exo"],
+ insulating_sleeve_options=sleeve_names_by_type["insulating"],
+ ).items():
+ state.setdefault(k, v)
+
+ @state.change("tau")
+ def _on_tau(tau, **kwargs):
+ if app["grid"] is None:
+ return
+ # show_scalar_bar=False: the "ghost" actor owns the one legend (see
+ # render_prediction); re-adding it here would duplicate/replace it on
+ # every slider tick.
+ liq = viz.liquid_region(app["grid"], tau)
+ if liq.n_cells > 0:
+ plotter.add_mesh(liq, scalars="solidification_time",
+ cmap=FIELD_CMAP, name="liquid", smooth_shading=True,
+ clim=[state.tmin, state.tmax], show_scalar_bar=False)
+ else:
+ plotter.remove_actor("liquid")
+ ctrl.view_update()
+
+ @state.change("show_hotspots")
+ def _on_hot(show_hotspots, **kwargs):
+ if app["hotspot_actor"] is not None:
+ app["hotspot_actor"].SetVisibility(bool(show_hotspots))
+ ctrl.view_update()
+
+ @state.change("hotspot_cutoff", "hotspot_size")
+ def _on_hotspot_controls(**kwargs):
+ _rebuild_hotspots(app)
+
+ @state.change("uploaded_file")
+ def _on_upload(uploaded_file, **kwargs):
+ if not uploaded_file:
+ return
+ f = uploaded_file[0] if isinstance(uploaded_file, list) else uploaded_file
+ out_dir = "runs/castings"
+ os.makedirs(out_dir, exist_ok=True)
+ step_path = osp.join(out_dir, f.get("name", "part.step"))
+ with open(step_path, "wb") as fh:
+ fh.write(f["content"])
+ state.loading = True
+ state.flush()
+ try:
+ _process_upload(app, step_path, state.alloy, state.mold)
+ finally:
+ state.loading = False
+ ctrl.view_update()
+
+ @state.change("assembly_bodies")
+ def _on_assembly_bodies_change(**kwargs):
+ # Fires on every role/param edit in the body table (trame syncs the
+ # whole reactive list back on any nested-item mutation) -- just
+ # recolor the already-tessellated preview, never re-tessellate.
+ if state.assembly_mode:
+ _render_assembly_preview(app)
+
+ def run_assembly_prediction():
+ """The assembly panel's "Run" button: build the bodies list from the
+ current panel state, pre-flight it through `validate_bodies` so a
+ training-regime violation (see `src.assembly.validate_bodies`) shows
+ up in the panel (`state.assembly_error`) WITHOUT ever calling
+ `predict_casting` -- never a crash for a bad role combination.
+ `_assemble_bodies_from_state` itself is inside this same try/except:
+ a malformed row (e.g. a missing `solid_id`) is a caller/state bug,
+ not a training-regime violation, but it should surface in the panel
+ exactly the same way -- never a raw traceback, and never a
+ `predict_casting` call on a bodies list that failed to even build."""
+ try:
+ bodies = _assemble_bodies_from_state(state)
+ validate_bodies(bodies)
+ except Exception as e:
+ state.assembly_error = str(e)
+ ctrl.view_update()
+ return
+ state.assembly_error = ""
+ state.loading = True
+ state.flush()
+ try:
+ result = _run_prediction(state.assembly_step_path, "runs/castings",
+ state.alloy, state.mold, bodies=bodies)
+ grid = viz.load_prediction(result["vtu"])
+ app["name"] = result["name"]
+ state.assembly_mode = False
+ render_prediction(app, grid, summary_text(result, alloy=state.alloy, mold=state.mold))
+ except Exception as e: # keep the UI alive on a bad assembly/prediction
+ plotter.clear_actors()
+ _remove_scalar_bar(plotter)
+ app["grid"] = None
+ app["hotspot_actor"] = None
+ state.loaded = False
+ state.summary = f"Prediction failed:\n{e}"
+ finally:
+ state.loading = False
+ ctrl.view_update()
+
+ # Exposed on the controller so the Run button (and tests) can invoke it
+ # without a browser round trip.
+ ctrl.run_assembly_prediction = run_assembly_prediction
+
+ def clear_casting():
+ """Drop the loaded casting and return to a ready-for-upload state."""
+ plotter.clear_actors()
+ _remove_scalar_bar(plotter)
+ app["grid"] = None
+ app["hotspot_actor"] = None
+ app["assembly_meshes"] = []
+ state.loaded = False
+ state.summary = INITIAL_SUMMARY
+ state.uploaded_file = None
+ state.tau, state.tmin, state.tmax = 0.0, 0.0, 1.0
+ state.show_hotspots = False
+ state.anim_ready = False
+ state.anim_uri = ""
+ state.animating = False
+ state.assembly_mode = False
+ state.assembly_bodies = []
+ state.assembly_error = ""
+ state.assembly_step_path = None
+ ctrl.view_update()
+
+ # Exposed on the controller so tests (and any future non-UI caller) can
+ # invoke the same clear logic the "Clear" button uses.
+ ctrl.clear_casting = clear_casting
+
+ def generate_animation():
+ """Render the solidification GIF for the loaded casting and expose it
+ to the browser as a base64 data URI (preview + download). Blocking, like
+ the upload prediction -- fine for this local single-user app."""
+ if app["grid"] is None:
+ return
+ state.animating = True
+ state.flush()
+ try:
+ out_dir = "runs/castings"
+ os.makedirs(out_dir, exist_ok=True)
+ gif_path = osp.join(out_dir, f"{app['name']}_solidification.gif")
+ viz.solidification_animation(app["grid"], gif_path,
+ cmap=FIELD_CMAP, ghost_opacity=GHOST_OPACITY,
+ scalar_bar_args=SCALAR_BAR_ARGS, bg=VTK_BG,
+ camera_position=plotter.camera_position)
+ import base64
+ with open(gif_path, "rb") as fh:
+ b64 = base64.b64encode(fh.read()).decode("ascii")
+ state.anim_uri = f"data:image/gif;base64,{b64}"
+ state.anim_name = f"{app['name']}_solidification.gif"
+ state.anim_ready = True
+ except Exception as e:
+ state.anim_ready = False
+ state.summary = f"Animation failed:\n{e}"
+ finally:
+ state.animating = False
+ ctrl.view_update()
+ ctrl.generate_animation = generate_animation
+
+ with SinglePageWithDrawerLayout(server, vuetify_config=vuetify_config,
+ width=360) as layout:
+ html.Style(_CUSTOM_CSS)
+
+ # The nav-icon toggles the drawer, but the drawer is `permanent` (it
+ # stays open as the fixed left control panel), so the toggle has no
+ # visible effect -- hide it rather than show a button that does nothing.
+ layout.icon.hide()
+
+ # Dark, compact app-bar: it holds only the title and a status line, so
+ # it is styled as an instrument-panel strip distinct from the light
+ # control surfaces below it, and kept slim since it carries no controls.
+ layout.toolbar.color = HEADER_BG
+ layout.toolbar.density = "compact"
+ layout.toolbar.flat = True
+
+ layout.title.set_text(
+ ''
+ 'Neuralcast Viewer'
+ "
"
+ )
+ with layout.toolbar:
+ vuetify3.VSpacer()
+ html.Div(
+ "{{ loading ? 'Predicting…' : (loaded ? 'Prediction loaded' "
+ ": 'Awaiting upload') }}",
+ classes="neuralcast-status mr-2",
+ )
+ # Thin indeterminate progress bar pinned to the bottom edge of the
+ # header -- visible only while a prediction is running, so the user
+ # can see the model is working during the blocking STEP -> VTU call.
+ vuetify3.VProgressLinear(
+ indeterminate=True, absolute=True, location="bottom",
+ height=3, color=ACCENT_MOLTEN, v_show=("loading",),
+ )
+
+ with layout.drawer:
+ with html.Div(classes="d-flex flex-column fill-height pa-4 "
+ "neuralcast-drawer-content"):
+ with html.Div(classes="neuralcast-section"):
+ html.Div("Upload", classes="neuralcast-section-label")
+ vuetify3.VSelect(v_model=("alloy",), items=("alloy_options",),
+ label="Alloy", density="compact",
+ variant="outlined", hide_details=True,
+ classes="mb-2")
+ vuetify3.VSelect(v_model=("mold",), items=("mold_options",),
+ label="Mold system", density="compact",
+ variant="outlined", hide_details=True,
+ classes="mb-2")
+ vuetify3.VFileInput(v_model=("uploaded_file", None), label="STEP file",
+ prepend_inner_icon="mdi-file-upload-outline",
+ accept=".step,.stp,.STEP,.STP", density="compact",
+ variant="outlined", hide_details=True)
+ vuetify3.VBtn("Clear", click=clear_casting, variant="text",
+ size="small", color="secondary", disabled=("!loaded",),
+ classes="mt-2", block=True)
+
+ # Assembly mode (Part B Task 5): shown instead of an
+ # auto-run prediction whenever the STEP upload has >1
+ # solid. One row per body -- role dropdown (defaults
+ # largest=casting/rest=ignore, set by _enter_assembly_mode)
+ # + a role-dependent params select (chill material / sleeve
+ # name, the latter pre-filtered by sleeve type so the UI
+ # never offers a pairing validate_bodies would reject).
+ # `v_model="body."` (a literal string, not a
+ # (name, default) tuple) is trame's standard per-row
+ # editable-table idiom: it binds directly to the `body`
+ # loop variable from `v_for`, and edits to `assembly_bodies`
+ # sync back to the server as a whole (see
+ # `_on_assembly_bodies_change`).
+ with html.Div(classes="neuralcast-section", v_show=("assembly_mode",)):
+ html.Div("Assembly (multi-body STEP)",
+ classes="neuralcast-section-label")
+ with html.Template(v_for="body in assembly_bodies", key="body.solid_id"):
+ with html.Div(classes="mb-3"):
+ html.Div(
+ "{{ body.name || ('solid ' + body.solid_id) }} "
+ "— {{ Number(body.volume_mm3).toFixed(0) }} mm³",
+ style=f"font-size:.75rem; color:{TEXT_MUTED}; "
+ "margin-bottom:2px;",
+ )
+ vuetify3.VSelect(v_model="body.role", items=("role_options",),
+ label="role", density="compact",
+ variant="outlined", hide_details=True,
+ classes="mb-1")
+ vuetify3.VSelect(v_model="body.params.material",
+ items=("chill_material_options",),
+ label="chill material", density="compact",
+ variant="outlined", hide_details=True,
+ v_if="body.role === 'chill'")
+ vuetify3.VSelect(v_model="body.params.sleeve",
+ items=("exo_sleeve_options",),
+ label="sleeve (exo)", density="compact",
+ variant="outlined", hide_details=True,
+ v_if="body.role === 'exo_sleeve'")
+ vuetify3.VSelect(v_model="body.params.sleeve",
+ items=("insulating_sleeve_options",),
+ label="sleeve (insulating)", density="compact",
+ variant="outlined", hide_details=True,
+ v_if="body.role === 'insulating_sleeve'")
+ html.Div("{{ assembly_error }}", v_show=("assembly_error",),
+ style=f"color:{ACCENT_MOLTEN}; font-size:.75rem; "
+ "margin:4px 0 8px;")
+ vuetify3.VBtn("Run Assembly Prediction",
+ click=ctrl.run_assembly_prediction,
+ variant="tonal", size="small", color="primary",
+ disabled=("loading",), loading=("loading",), block=True)
+
+ with html.Div(classes="neuralcast-section"):
+ html.Div("Playback", classes="neuralcast-section-label")
+ vuetify3.VSlider(v_model=("tau",), min=("tmin",), max=("tmax",),
+ label="time (s)",
+ hide_details=True, disabled=("!loaded",),
+ thumb_label=True)
+ html.Div(
+ "{{ Number(tau).toFixed(1) }} / {{ Number(tmax).toFixed(1) }} s",
+ classes="neuralcast-tau-badge mt-1",
+ )
+
+ with html.Div(classes="neuralcast-section"):
+ html.Div("Hot Spots", classes="neuralcast-section-label")
+ vuetify3.VCheckbox(v_model=("show_hotspots",), label="show hot spots",
+ hide_details=True, disabled=("!loaded",))
+ vuetify3.VSlider(v_model=("hotspot_size",), min=2, max=40, step=1,
+ label="size", hide_details=True,
+ disabled=("!loaded",), thumb_label=True,
+ classes="mt-4")
+ vuetify3.VSlider(v_model=("hotspot_cutoff",), min=("tmin",),
+ max=("tmax",), label="cutoff (s)", hide_details=True,
+ disabled=("!loaded",), thumb_label=True,
+ classes="mt-2")
+
+ with html.Div(classes="neuralcast-section"):
+ html.Div("Animation", classes="neuralcast-section-label")
+ vuetify3.VBtn("Generate GIF", click=generate_animation,
+ variant="tonal", size="small", color="primary",
+ disabled=("!loaded || animating",), block=True,
+ loading=("animating",))
+ vuetify3.VProgressLinear(indeterminate=True, height=3,
+ color=ACCENT_MOLTEN, v_show=("animating",),
+ classes="mt-2")
+ html.Img(src=("anim_uri",), v_show=("anim_ready",),
+ style="max-width:100%; margin-top:10px; "
+ f"border:1px solid {BORDER}; border-radius:4px;")
+ html.A("Download GIF", href=("anim_uri",), download=("anim_name",),
+ v_show=("anim_ready",),
+ style=f"display:block; margin-top:8px; color:{ACCENT_MOLTEN}; "
+ "font-size:.8rem; text-decoration:underline;")
+
+ with html.Div(classes="neuralcast-section flex-grow-1 d-flex "
+ "flex-column neuralcast-summary-box"):
+ html.Div("Summary", classes="neuralcast-section-label")
+ vuetify3.VTextarea(v_model=("summary",), readonly=True,
+ variant="outlined", density="compact",
+ hide_details=True, no_resize=True,
+ classes="flex-grow-1")
+
+ with layout.content:
+ with vuetify3.VContainer(fluid=True, classes="fill-height pa-0"):
+ # add_menu=True: keep pyvista's default view-toolbar overlay --
+ # it gives reset-camera / view-axis / screenshot controls that
+ # users want when inspecting the casting.
+ # mode="server": render server-side (VTK on the Python process,
+ # frames streamed to the browser) rather than client-side
+ # (geometry shipped to the browser and rendered there). This
+ # keeps the server-side VTK camera in sync with the user's
+ # mouse rotations/pans/zooms, so generate_animation() below can
+ # read plotter.camera_position and reuse the exact view for
+ # the GIF -- with client-side rendering the server-side camera
+ # would stay frozen regardless of what the user sees.
+ view = plotter_ui(plotter, add_menu=True, mode="server")
+ ctrl.view_update = view.update
+
+ # Replace trame's default footer ("Powered by trame" / Kitware credit)
+ # with our own.
+ layout.footer.clear()
+ with layout.footer:
+ html.Div("© 2026 Eugen Miknevic", classes="neuralcast-footer-text")
+
+ # NOTE (deviation from brief): trame's `state.to_dict()` only reflects state
+ # that has been *flushed*, and flushing is a no-op until the state is marked
+ # `ready()` (normally done when a client connects, i.e. inside `.start()`).
+ # Since `build_app()` must be unit-constructible without starting the server,
+ # we mark the state ready here so the initial `setdefault(...)` values are
+ # visible via `to_dict()`/attribute access immediately after construction.
+ # This performs no networking (no client is connected yet) and is idempotent
+ # (a later real `.start()` calling `ready()` again is a no-op).
+ state.ready()
+
+ return server
+
+
+def main():
+ build_app().start(port=8080)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000000000000000000000000000000000000..1ac94b7023abcc35075c3bdddcb861566e20114c
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,49 @@
+"""Shared pytest fixtures: a minimal gmsh-2.2 mesh for unit tests.
+
+The mesh is a unit-ish tetrahedron-based shape with 5 nodes:
+ - one internal-ish apex, 4 base nodes forming two tets
+ - two boundary triangles tagged 'Cooling', one tagged 'Adiabatic'
+Coordinates are in mm.
+"""
+import textwrap
+import pytest
+
+# node ids 1..5 ; coords in mm
+MINI_MSH = textwrap.dedent("""\
+ $MeshFormat
+ 2.2 0 8
+ $EndMeshFormat
+ $PhysicalNames
+ 3
+ 2 1 "Adiabatic"
+ 2 2 "Cooling"
+ 3 3 "Volume"
+ $EndPhysicalNames
+ $Nodes
+ 5
+ 1 0 0 0
+ 2 10 0 0
+ 3 0 10 0
+ 4 0 0 10
+ 5 3 3 3
+ $EndNodes
+ $Elements
+ 4
+ 1 2 2 2 1 1 2 3
+ 2 2 2 2 1 1 2 4
+ 3 2 2 1 1 2 3 4
+ 4 4 2 3 1 1 2 3 4
+ $EndElements
+ """)
+# Element lines: "id type ntags phys geom n1 n2 n3[ n4]"
+# el1: triangle (type 2), phys 2 = Cooling, nodes 1,2,3
+# el2: triangle, phys 2 = Cooling, nodes 1,2,4
+# el3: triangle, phys 1 = Adiabatic, nodes 2,3,4
+# el4: tetrahedron (type 4), phys 3 = Volume, nodes 1,2,3,4
+
+
+@pytest.fixture
+def mini_msh_file(tmp_path):
+ p = tmp_path / "MINI_00.msh"
+ p.write_text(MINI_MSH)
+ return str(p)
diff --git a/tests/test_assembly.py b/tests/test_assembly.py
new file mode 100644
index 0000000000000000000000000000000000000000..ba4678eb04f4da835fa8e42e4dda0f14c2d88efc
--- /dev/null
+++ b/tests/test_assembly.py
@@ -0,0 +1,732 @@
+"""Tests for src/assembly.py -- Part B Task 1: assembly inspection
+(`list_solids`) + the `compose_assembly` test/fixture helper used by later
+Part B fixtures (proximity classification, BC-patch tagging, etc.), and
+Part B Task 2: `validate_bodies -> BcPlan` (no gmsh needed for Task 2)."""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from src.assembly import (BcPlan, list_solids, compose_assembly, main, validate_bodies,
+ tag_contact_patches, _triangle_areas, _rewrite_msh_adiabatic)
+from src.msh_parse import parse_msh
+from tests.test_prep_parts import _make_one_box_step
+from tools.doe.load_properties import resolve_exo_q0_by_mef
+from tools.doe.prep_parts import heal_and_mesh
+
+
+def test_compose_and_list_two_solid_assembly(tmp_path):
+ """Base box (30x20x10, the "casting") + a smaller box (10x10x6, a
+ "chill") attached at the +z face. `compose_assembly` places the chill
+ sunk `overlap_mm` into the base (default 2.0mm) -- NO boolean op, so the
+ chill's own reported volume is its full analytic 10*10*6, unaffected by
+ the intentional overlap with the base solid. `list_solids` must report
+ both solids, sorted by volume descending (casting first)."""
+ base_step = tmp_path / "base.step"
+ _make_one_box_step(base_step, 30.0, 20.0, 10.0)
+
+ out_step = tmp_path / "assembly.step"
+ attachments = [
+ {"shape": "box", "size": (10.0, 10.0, 6.0), "at_face": "+z"},
+ ]
+ n = compose_assembly(str(base_step), attachments, str(out_step))
+ assert n == 2
+
+ rows = list_solids(str(out_step))
+ assert len(rows) == 2
+ assert rows[0]["volume_mm3"] == pytest.approx(30.0 * 20.0 * 10.0, rel=0.05)
+ assert rows[1]["volume_mm3"] == pytest.approx(10.0 * 10.0 * 6.0, rel=0.05)
+ # sorted desc by volume
+ assert rows[0]["volume_mm3"] >= rows[1]["volume_mm3"]
+ for row in rows:
+ assert set(row.keys()) == {"solid_id", "name", "volume_mm3", "bbox", "com"}
+ assert len(row["bbox"]) == 6
+ assert len(row["com"]) == 3
+
+
+def test_list_solids_single_solid(tmp_path):
+ """A plain single-solid STEP (reusing the existing box fixture from
+ tests/test_prep_parts.py) must yield exactly 1 row."""
+ step = tmp_path / "one_box.step"
+ _make_one_box_step(step, 20.0, 20.0, 20.0)
+
+ rows = list_solids(str(step))
+ assert len(rows) == 1
+ assert rows[0]["volume_mm3"] == pytest.approx(20.0 ** 3, rel=0.05)
+
+
+def test_compose_assembly_cylinder_attachment(tmp_path):
+ """Cylinder-shaped attachment (radius, height) at a different face,
+ exercised end to end -- also confirms non-"+z" faces work."""
+ base_step = tmp_path / "base.step"
+ _make_one_box_step(base_step, 30.0, 20.0, 10.0)
+
+ out_step = tmp_path / "assembly.step"
+ attachments = [
+ {"shape": "cylinder", "size": (3.0, 15.0), "at_face": "+x", "overlap_mm": 1.5},
+ ]
+ n = compose_assembly(str(base_step), attachments, str(out_step))
+ assert n == 2
+
+ rows = list_solids(str(out_step))
+ assert len(rows) == 2
+ import math
+ cyl_vol = math.pi * 3.0 ** 2 * 15.0
+ volumes = sorted(row["volume_mm3"] for row in rows)
+ assert volumes[0] == pytest.approx(cyl_vol, rel=0.05)
+ assert volumes[1] == pytest.approx(30.0 * 20.0 * 10.0, rel=0.05)
+
+
+def test_main_prints_solid_ids(tmp_path, capsys):
+ """`python -m src.assembly part.step` (exercised here via a direct
+ `main()` call rather than a real subprocess, to keep the test fast)
+ prints a table that includes both solids' ids."""
+ base_step = tmp_path / "base.step"
+ _make_one_box_step(base_step, 30.0, 20.0, 10.0)
+ out_step = tmp_path / "assembly.step"
+ compose_assembly(str(base_step),
+ [{"shape": "box", "size": (10.0, 10.0, 6.0), "at_face": "+z"}],
+ str(out_step))
+
+ rows = list_solids(str(out_step))
+ rc = main([str(out_step)])
+ captured = capsys.readouterr()
+
+ assert rc == 0
+ for row in rows:
+ assert str(row["solid_id"]) in captured.out
+
+
+# ---------------------------------------------------------------------- #
+# Part B Task 2 -- validate_bodies -> BcPlan (no gmsh; pure dict/CSV logic).
+# All tests use the real, committed registries (data/properties/chills.csv,
+# tools/sleeves/sleeve_registry.csv, data/properties/exo_calibration.csv)
+# since those ARE the production registries and their values are pinned
+# (see module docstrings): cast_iron effusivity 11968, sleeve "SD_4_7K"
+# MEF 1.547 (2026-08-09 sleeve-mef-correction: catalog_klx/catalog_geom =
+# 1.16/0.75, was the synthetic-spread value 1.49 before the fix -- see
+# tools/sleeves/gen_sleeves.py's module docstring).
+# ---------------------------------------------------------------------- #
+def test_validate_bodies_casting_plus_two_chills():
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "chill"},
+ {"solid_id": 3, "role": "chill"},
+ ]
+ plan = validate_bodies(bodies)
+ assert isinstance(plan, BcPlan)
+ assert plan.bc_config == "chills"
+ assert plan.casting_solid_id == 1
+ assert sorted(plan.role_solid_ids) == [2, 3]
+ assert plan.n_chills == 2
+ assert plan.chill_effusivity == pytest.approx(11968.0)
+ assert plan.sleeve_MEF == 0.0
+ assert plan.exo_q0 == 0.0
+
+
+def test_validate_bodies_chills_explicit_material_matches_default():
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "chill", "params": {"material": "cast_iron"}},
+ ]
+ plan = validate_bodies(bodies)
+ assert plan.bc_config == "chills"
+ assert plan.n_chills == 1
+ assert plan.chill_effusivity == pytest.approx(11968.0)
+
+
+def test_validate_bodies_chills_differing_materials_raises():
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "chill", "params": {"material": "cast_iron"}},
+ {"solid_id": 3, "role": "chill", "params": {"material": "graphite"}},
+ ]
+ with pytest.raises(ValueError, match="all chills must share one material"):
+ validate_bodies(bodies)
+
+
+def test_validate_bodies_unknown_chill_material_raises_listing_options():
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "chill", "params": {"material": "unobtainium"}},
+ ]
+ with pytest.raises(ValueError) as exc_info:
+ validate_bodies(bodies)
+ msg = str(exc_info.value)
+ assert "unobtainium" in msg
+ assert "cast_iron" in msg # a known material is listed
+
+
+def test_validate_bodies_exo_sleeve_by_registry_name():
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "exo_sleeve", "params": {"sleeve": "SD_4_7K"}},
+ ]
+ plan = validate_bodies(bodies)
+ assert plan.bc_config == "exo_sleeve"
+ assert plan.role_solid_ids == [2]
+ assert plan.sleeve_MEF == pytest.approx(1.547)
+ assert plan.exo_q0 > 0.0
+ assert plan.exo_q0 == pytest.approx(resolve_exo_q0_by_mef(1.547))
+
+
+def test_validate_bodies_exo_sleeve_manual_mef_override():
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "exo_sleeve", "params": {"MEF": 1.49}},
+ ]
+ plan = validate_bodies(bodies)
+ assert plan.bc_config == "exo_sleeve"
+ assert plan.sleeve_MEF == pytest.approx(1.49)
+ assert plan.exo_q0 == pytest.approx(resolve_exo_q0_by_mef(1.49))
+
+
+def test_validate_bodies_insulating_sleeve_no_exo_q0():
+ # FOSECO_V90L is registry type=insulating -- matches role insulating_sleeve.
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "insulating_sleeve", "params": {"sleeve": "FOSECO_V90L"}},
+ ]
+ plan = validate_bodies(bodies)
+ assert plan.bc_config == "insulating_sleeve"
+ assert plan.sleeve_MEF == pytest.approx(1.15)
+ assert plan.exo_q0 == 0.0
+
+
+def test_validate_bodies_sleeve_both_name_and_mef_raises():
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "exo_sleeve",
+ "params": {"sleeve": "SD_4_7K", "MEF": 1.49}},
+ ]
+ with pytest.raises(ValueError, match="not both"):
+ validate_bodies(bodies)
+
+
+def test_validate_bodies_sleeve_neither_name_nor_mef_raises():
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "exo_sleeve", "params": {}},
+ ]
+ with pytest.raises(ValueError, match="'sleeve'.*'MEF'|'MEF'.*'sleeve'"):
+ validate_bodies(bodies)
+
+
+def test_validate_bodies_exo_role_with_insulating_typed_sleeve_raises():
+ # FOSECO_V88E is registry type=insulating, MEF=1.45 -- collides with an
+ # exo_calibration.csv row at MEF=1.45, so without the type cross-check
+ # this would silently resolve a confident nonzero exo_q0 for a
+ # non-exothermic product.
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "exo_sleeve", "params": {"sleeve": "FOSECO_V88E"}},
+ ]
+ with pytest.raises(ValueError, match="FOSECO_V88E.*type=insulating"):
+ validate_bodies(bodies)
+
+
+def test_validate_bodies_insulating_role_with_exo_typed_sleeve_raises():
+ # SD_4_7K is registry type=exo -- must not be accepted under
+ # insulating_sleeve.
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "insulating_sleeve", "params": {"sleeve": "SD_4_7K"}},
+ ]
+ with pytest.raises(ValueError, match="SD_4_7K.*type=exo"):
+ validate_bodies(bodies)
+
+
+def test_validate_bodies_sleeve_type_matches_role_ok():
+ # exo role + exo-typed sleeve, and insulating role + insulating-typed
+ # sleeve, both resolve cleanly (matched pairs).
+ exo_bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "exo_sleeve", "params": {"sleeve": "SD_4_7K"}},
+ ]
+ plan = validate_bodies(exo_bodies)
+ assert plan.bc_config == "exo_sleeve"
+ assert plan.sleeve_MEF == pytest.approx(1.547)
+
+ insulating_bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "insulating_sleeve", "params": {"sleeve": "FOSECO_V88E"}},
+ ]
+ plan2 = validate_bodies(insulating_bodies)
+ assert plan2.bc_config == "insulating_sleeve"
+ assert plan2.sleeve_MEF == pytest.approx(1.45)
+
+
+def test_validate_bodies_manual_mef_bypasses_type_check_either_role():
+ # Manual MEF override has no registry row to check against -- expert
+ # mode, allowed under either sleeve role.
+ exo_bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "exo_sleeve", "params": {"MEF": 1.45}},
+ ]
+ plan = validate_bodies(exo_bodies)
+ assert plan.bc_config == "exo_sleeve"
+ assert plan.sleeve_MEF == pytest.approx(1.45)
+
+ insulating_bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "insulating_sleeve", "params": {"MEF": 1.49}},
+ ]
+ plan2 = validate_bodies(insulating_bodies)
+ assert plan2.bc_config == "insulating_sleeve"
+ assert plan2.sleeve_MEF == pytest.approx(1.49)
+ assert plan2.exo_q0 == 0.0
+
+
+def test_validate_bodies_unknown_sleeve_name_raises_listing_options():
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "exo_sleeve", "params": {"sleeve": "NOT_A_SLEEVE"}},
+ ]
+ with pytest.raises(ValueError) as exc_info:
+ validate_bodies(bodies)
+ msg = str(exc_info.value)
+ assert "NOT_A_SLEEVE" in msg
+ assert "SD_4_7K" in msg
+
+
+def test_validate_bodies_casting_only_bc_none():
+ bodies = [{"solid_id": 1, "role": "casting"}]
+ plan = validate_bodies(bodies)
+ assert plan.bc_config == "none"
+ assert plan.casting_solid_id == 1
+ assert plan.role_solid_ids == []
+ assert plan.n_chills == 0
+ assert plan.chill_effusivity == 0.0
+ assert plan.sleeve_MEF == 0.0
+ assert plan.exo_q0 == 0.0
+
+
+def test_validate_bodies_ignore_bodies_dropped():
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "ignore"},
+ {"solid_id": 3, "role": "ignore", "params": {"whatever": 1}},
+ ]
+ plan = validate_bodies(bodies)
+ assert plan.bc_config == "none"
+ assert plan.casting_solid_id == 1
+
+
+def test_validate_bodies_zero_castings_raises():
+ bodies = [{"solid_id": 2, "role": "chill"}]
+ with pytest.raises(ValueError, match="exactly one"):
+ validate_bodies(bodies)
+
+
+def test_validate_bodies_two_castings_raises():
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "casting"},
+ ]
+ with pytest.raises(ValueError, match="exactly one"):
+ validate_bodies(bodies)
+
+
+def test_validate_bodies_chills_plus_sleeve_mixed_raises():
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "chill"},
+ {"solid_id": 3, "role": "exo_sleeve", "params": {"sleeve": "SD_4_7K"}},
+ ]
+ with pytest.raises(ValueError, match="chills and a sleeve cannot be combined"):
+ validate_bodies(bodies)
+
+
+def test_validate_bodies_mixed_sleeve_roles_raises():
+ # SP4 ADDENDUM (2026-08-08): N sleeve bodies are accepted now, but only
+ # when they share one role -- an exo_sleeve + an insulating_sleeve still
+ # raises, because the model has a single sleeve BC channel and cannot
+ # represent both kinds of sleeve heat exchange at once.
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "exo_sleeve", "params": {"sleeve": "SD_4_7K"}},
+ {"solid_id": 3, "role": "insulating_sleeve", "params": {"sleeve": "FOSECO_V90L"}},
+ ]
+ with pytest.raises(ValueError, match="one role"):
+ validate_bodies(bodies)
+
+
+def test_validate_bodies_two_same_role_sleeves_resolve():
+ # SP4 ADDENDUM (2026-08-08): two sleeve bodies of the SAME role and the
+ # SAME registry-resolved MEF must now be accepted (this used to raise
+ # "at most one sleeve body is supported"), resolving to that single MEF
+ # exactly and covering both solid_ids in role_solid_ids.
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "exo_sleeve", "params": {"sleeve": "SD_4_7K"}},
+ {"solid_id": 3, "role": "exo_sleeve", "params": {"sleeve": "SD_4_7K"}},
+ ]
+ plan = validate_bodies(bodies)
+ assert plan.bc_config == "exo_sleeve"
+ assert sorted(plan.role_solid_ids) == [2, 3]
+ assert plan.sleeve_MEF == pytest.approx(1.547)
+ assert plan.mef_weighted is False
+ assert plan.mef_spread == pytest.approx(0.0)
+
+
+def test_validate_bodies_unknown_role_raises():
+ bodies = [
+ {"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "typo_role"},
+ ]
+ with pytest.raises(ValueError, match="unknown role"):
+ validate_bodies(bodies)
+
+
+def test_bcplan_feature_spec_keys_and_values():
+ plan = BcPlan(bc_config="chills", casting_solid_id=1, role_solid_ids=[2],
+ n_chills=1, chill_effusivity=11968.0)
+ spec = plan.feature_spec(mold_effusivity=1250.0)
+ assert spec == {
+ "mold_effusivity": 1250.0,
+ "bc_config": "chills",
+ "chill_effusivity": 11968.0,
+ "sleeve_MEF": 0.0,
+ "exo_q0": 0.0,
+ }
+
+
+def test_bcplan_gbm_bc_spec_keys_and_exo_q0_mw_conversion():
+ plan = BcPlan(bc_config="exo_sleeve", casting_solid_id=1, role_solid_ids=[2],
+ sleeve_MEF=1.49, exo_q0=960000.0)
+ spec = plan.gbm_bc_spec()
+ assert spec == {
+ "bc_config": "exo_sleeve",
+ "n_chills": 0,
+ "sleeve_MEF": 1.49,
+ "exo_q0_MW": pytest.approx(0.96),
+ }
+
+
+# ---------------------------------------------------------------------- #
+# Part B Task 3 -- tag_contact_patches: mesh the user-designated casting
+# solid (via prep_parts' new `keep_solid_id`), classify its "Cooling"
+# triangles by proximity to role bodies, and rewrite the mesh's Adiabatic
+# facet group. Real gmsh throughout -- these build a fresh composer fixture
+# each time (see `src.assembly.compose_assembly`'s module docstring for why
+# a few mm of intentional overlap is the fixture convention, not a bug).
+# ---------------------------------------------------------------------- #
+def _elem_est_mm(mesh) -> float:
+ """Same `bbox_diag / n_nodes**(1/3)` local element-size estimate
+ `tag_contact_patches` uses for its default `tol_mm` -- used here only to
+ size test assertion margins, not to drive the function itself."""
+ bbox_diag = float(np.linalg.norm(mesh.node_xyz.max(axis=0) - mesh.node_xyz.min(axis=0)))
+ return bbox_diag / (mesh.node_xyz.shape[0] ** (1.0 / 3.0))
+
+
+def _build_chill_assembly(tmp_path, base_size=(30.0, 20.0, 10.0), chill_size=(10.0, 10.0, 6.0),
+ overlap_mm=2.0, target_nodes=500):
+ """Base box (the "casting") + a smaller box (the "chill") attached at
+ the "+z" face, exactly the `compose_assembly` fixture pattern used
+ throughout this file. Meshes ONLY the casting (`keep_solid_id`, Task 3's
+ other new capability -- the casting need not be the largest solid in
+ general, though it is here) and returns
+ `(casting_msh_path, assembly_step_path, casting_id, chill_id,
+ heal_and_mesh_result)`."""
+ base_step = tmp_path / "base.step"
+ _make_one_box_step(base_step, *base_size)
+
+ assembly_step = tmp_path / "assembly.step"
+ attachments = [{"shape": "box", "size": chill_size, "at_face": "+z", "overlap_mm": overlap_mm}]
+ compose_assembly(str(base_step), attachments, str(assembly_step))
+
+ rows = list_solids(str(assembly_step))
+ casting_id = max(rows, key=lambda r: r["volume_mm3"])["solid_id"]
+ chill_id = min(rows, key=lambda r: r["volume_mm3"])["solid_id"]
+
+ out_msh = tmp_path / "casting.msh"
+ result = heal_and_mesh(str(assembly_step), str(out_msh), target_nodes=target_nodes,
+ keep_solid_id=casting_id)
+ assert result["ok"] is True
+ return str(out_msh), str(assembly_step), casting_id, chill_id, result
+
+
+def test_tag_contact_patches_exact_contact_matches_chill_footprint(tmp_path):
+ """Overlap 2mm box chill at "+z" center (compose_assembly's own default
+ overlap) -- the Adiabatic patch must land on the chill's footprint: every
+ Adiabatic centroid within the footprint's xy extent (+- a couple of
+ element widths, since a nonzero tol_mm necessarily dilates the detected
+ region a little past the true edge) and at the casting's top z; total
+ patch area within 30% of the chill's analytic contact-face area. A
+ large-relative-to-tol footprint (40x40mm chill on a 60x50x15mm base) is
+ used so the fixed ~2mm overlap gap doesn't force a tol_mm whose
+ Minkowski-dilation ring is a big fraction of a small footprint's area
+ (verified numerically while writing this test: a 10x10mm footprint blows
+ well past 30% at any tol_mm large enough to bridge a 2mm gap)."""
+ base_size = (60.0, 50.0, 15.0)
+ chill_size = (40.0, 40.0, 6.0)
+ out_msh, assembly_step, casting_id, chill_id, _result = _build_chill_assembly(
+ tmp_path, base_size=base_size, chill_size=chill_size, overlap_mm=2.0,
+ target_nodes=3000)
+
+ mesh_before = parse_msh(out_msh)
+ elem_est = _elem_est_mm(mesh_before)
+ tol_mm = 2.5 # > the 2mm overlap gap, small relative to the 40x40 footprint
+
+ tag_result = tag_contact_patches(out_msh, assembly_step, [chill_id], tol_mm=tol_mm)
+
+ assert tag_result["n_adiab_tris"] > 0
+ chill_contact_area = chill_size[0] * chill_size[1] # dx * dy of the chill footprint
+ assert tag_result["per_body_patch_area_mm2"][chill_id] == pytest.approx(
+ chill_contact_area, rel=0.30)
+
+ mesh_after = parse_msh(out_msh)
+ adiab = mesh_after.tris["Adiabatic"]
+ assert adiab.shape[0] == tag_result["n_adiab_tris"]
+ centroids = mesh_after.node_xyz[adiab].mean(axis=1)
+
+ # Footprint is centered on the base's top face: x in [10, 50], y in [5, 45]
+ # (base 60x50, chill 40x40 centered -> (60-40)/2=10, (50-40)/2=5).
+ margin = tol_mm + elem_est
+ fx_lo, fx_hi = (base_size[0] - chill_size[0]) / 2.0, (base_size[0] + chill_size[0]) / 2.0
+ fy_lo, fy_hi = (base_size[1] - chill_size[1]) / 2.0, (base_size[1] + chill_size[1]) / 2.0
+ assert np.all(centroids[:, 0] >= fx_lo - margin)
+ assert np.all(centroids[:, 0] <= fx_hi + margin)
+ assert np.all(centroids[:, 1] >= fy_lo - margin)
+ assert np.all(centroids[:, 1] <= fy_hi + margin)
+ # at the casting's top face
+ assert np.allclose(centroids[:, 2], base_size[2], atol=margin)
+
+
+def test_tag_contact_patches_nodes_by_solid_matches_adiabatic_group(tmp_path):
+ """Part-B roles task (2026-08-03): `nodes_by_solid` -- the in-memory
+ per-solid contact-patch node-id mapping `tag_contact_patches` now also
+ returns -- must, for a single role body, exactly equal the node ids of
+ the "Adiabatic" group the SAME call wrote to `out_msh` (with only one
+ body, that body's own `close` mask IS the Adiabatic union)."""
+ out_msh, assembly_step, casting_id, chill_id, _result = _build_chill_assembly(
+ tmp_path, overlap_mm=2.0, target_nodes=500)
+
+ tag_result = tag_contact_patches(out_msh, assembly_step, [chill_id], tol_mm=3.0)
+
+ assert set(tag_result["nodes_by_solid"]) == {chill_id}
+ mesh_after = parse_msh(out_msh)
+ adiab_nodes = sorted(set(mesh_after.tris["Adiabatic"].reshape(-1).tolist()))
+ assert tag_result["nodes_by_solid"][chill_id] == adiab_nodes
+ # sorted, unique, and does not mutate the .msh/node-order contract --
+ # covered independently by test_tag_contact_patches_preserves_node_
+ # order_and_tets; here just confirm it's genuinely sorted+unique.
+ assert tag_result["nodes_by_solid"][chill_id] == sorted(
+ set(tag_result["nodes_by_solid"][chill_id]))
+
+
+def test_tag_contact_patches_small_clearance_still_found(tmp_path):
+ """0.1mm CLEARANCE (overlap_mm=-0.1, a hair short of touching) with
+ tol_mm >= 0.5mm must still find the patch."""
+ out_msh, assembly_step, casting_id, chill_id, _result = _build_chill_assembly(
+ tmp_path, overlap_mm=-0.1, target_nodes=500)
+
+ tag_result = tag_contact_patches(out_msh, assembly_step, [chill_id], tol_mm=0.5)
+
+ assert tag_result["n_adiab_tris"] > 0
+ assert tag_result["per_body_patch_area_mm2"][chill_id] > 0.0
+
+
+def test_tag_contact_patches_far_body_raises(tmp_path):
+ """A body 10mm away with tol_mm=1mm must raise -- never a silent
+ all-sand prediction for a body the caller explicitly assigned a role
+ to."""
+ out_msh, assembly_step, casting_id, chill_id, _result = _build_chill_assembly(
+ tmp_path, overlap_mm=-10.0, target_nodes=500)
+
+ with pytest.raises(ValueError, match=f"body {chill_id} does not touch"):
+ tag_contact_patches(out_msh, assembly_step, [chill_id], tol_mm=1.0)
+
+
+def test_tag_contact_patches_preserves_node_order_and_tets(tmp_path):
+ """The rewritten mesh's node ids/coords and tet connectivity must be
+ EXACTLY unchanged -- features/geodesics depend on it (see the module
+ docstring's node-order caution). Parses before AND after, independently
+ of `tag_contact_patches`'s own internal round-trip assertion."""
+ out_msh, assembly_step, casting_id, chill_id, _result = _build_chill_assembly(
+ tmp_path, overlap_mm=2.0, target_nodes=500)
+
+ mesh_before = parse_msh(out_msh)
+
+ tag_contact_patches(out_msh, assembly_step, [chill_id], tol_mm=3.0)
+
+ mesh_after = parse_msh(out_msh)
+ assert np.allclose(mesh_after.node_xyz, mesh_before.node_xyz)
+ assert mesh_after.tets.shape == mesh_before.tets.shape
+ assert np.array_equal(mesh_after.tets, mesh_before.tets)
+ # sanity: some Cooling triangles actually migrated to Adiabatic
+ assert "Adiabatic" in mesh_after.tris
+ assert mesh_after.tris["Cooling"].shape[0] < mesh_before.tris["Cooling"].shape[0]
+ assert (mesh_before.tris["Cooling"].shape[0]
+ == mesh_after.tris["Cooling"].shape[0] + mesh_after.tris["Adiabatic"].shape[0])
+
+
+def test_tag_contact_patches_unknown_role_solid_id_raises(tmp_path):
+ """A role `solid_id` that doesn't exist in `step_path` at all must fail
+ loudly (whatever form that takes -- gmsh itself raises when asked for
+ the boundary of a nonexistent entity), never silently produce an empty
+ or wrong patch."""
+ out_msh, assembly_step, casting_id, chill_id, _result = _build_chill_assembly(tmp_path)
+
+ with pytest.raises(Exception):
+ tag_contact_patches(out_msh, assembly_step, [9999], tol_mm=3.0)
+
+
+def test_heal_and_mesh_keep_solid_id_meshes_small_chill_from_assembly(tmp_path):
+ """Part B Task 3's other change: `heal_and_mesh(..., keep_solid_id=...)`
+ meshes the SMALL (chill) solid out of a composed multi-body assembly
+ (n_nodes/volume consistent with its own size), while the largest-solid
+ default on the SAME input is untouched (still meshes the big casting)."""
+ base_step = tmp_path / "base.step"
+ _make_one_box_step(base_step, 30.0, 20.0, 10.0)
+ assembly_step = tmp_path / "assembly.step"
+ compose_assembly(str(base_step),
+ [{"shape": "box", "size": (10.0, 10.0, 6.0), "at_face": "+z"}],
+ str(assembly_step))
+
+ rows = list_solids(str(assembly_step))
+ casting_row = max(rows, key=lambda r: r["volume_mm3"])
+ chill_row = min(rows, key=lambda r: r["volume_mm3"])
+
+ out_msh_chill = tmp_path / "chill.msh"
+ result_chill = heal_and_mesh(str(assembly_step), str(out_msh_chill), target_nodes=300,
+ keep_solid_id=chill_row["solid_id"])
+ assert result_chill["ok"] is True
+ assert result_chill["volume_mm3"] == pytest.approx(chill_row["volume_mm3"], rel=0.05)
+
+ out_msh_default = tmp_path / "default.msh"
+ result_default = heal_and_mesh(str(assembly_step), str(out_msh_default), target_nodes=300)
+ assert result_default["ok"] is True
+ assert result_default["volume_mm3"] == pytest.approx(casting_row["volume_mm3"], rel=0.05)
+
+
+# ---------------------------------------------------------------------- #
+# Code-review follow-ups (2026-07-29): total-Cooling-consumption guard,
+# per_body_patch_area_mm2's documented double-counting behavior, and the
+# clean (no bare StopIteration) too-few-flags path in
+# `_rewrite_msh_adiabatic`.
+# ---------------------------------------------------------------------- #
+def test_tag_contact_patches_enveloping_body_raises_no_cooling_remains(tmp_path):
+ """A role body that tightly ENVELOPES the entire casting (every Cooling
+ triangle ends up within tol_mm of it) must raise loudly -- if every
+ triangle were silently relabeled Adiabatic, `geodesic_d_cool` (see
+ `src.features`) would have no "Cooling" surface left to measure
+ distance to and would silently degrade to bbox_diag-derived constants
+ (confidently wrong physics, not a loud failure)."""
+ base_step = tmp_path / "base.step"
+ _make_one_box_step(base_step, 30.0, 20.0, 10.0)
+
+ assembly_step = tmp_path / "assembly.step"
+ # A box attached at "+z" with overlap_mm >= its own along-normal size
+ # sinks past the base's bottom too, and a face-plane size well past the
+ # base's own x/y extent -- the result envelopes the base on all 6 sides
+ # with a uniform ~5mm margin (worked out analytically: dz=20,
+ # overlap_mm=15 -> z in [-5, 15] vs the base's [0, 10]; dx=40, dy=30
+ # centered on the base's own center -> x in [-5, 35] vs [0, 30], y in
+ # [-5, 25] vs [0, 20]).
+ envelope = {"shape": "box", "size": (40.0, 30.0, 20.0), "at_face": "+z", "overlap_mm": 15.0}
+ compose_assembly(str(base_step), [envelope], str(assembly_step))
+
+ rows = list_solids(str(assembly_step))
+ casting_id = max(rows, key=lambda r: r["volume_mm3"])["solid_id"]
+ envelope_id = min(rows, key=lambda r: r["volume_mm3"])["solid_id"]
+
+ out_msh = tmp_path / "casting.msh"
+ result = heal_and_mesh(str(assembly_step), str(out_msh), target_nodes=500,
+ keep_solid_id=casting_id)
+ assert result["ok"] is True
+
+ with pytest.raises(ValueError, match="cover the entire casting surface"):
+ tag_contact_patches(str(out_msh), str(assembly_step), [envelope_id], tol_mm=8.0)
+
+
+def test_tag_contact_patches_two_bodies_double_count_area_single_count_tris(tmp_path):
+ """Two role bodies whose tolerance zones overlap on the casting surface:
+ `per_body_patch_area_mm2` sums each body's own area INDEPENDENTLY (the
+ shared triangles counted once per body -- documented double-counting),
+ while `n_adiab_tris` (and the actual rewritten "Adiabatic" group) counts
+ each such triangle exactly once, via the union mask."""
+ base_step = tmp_path / "base.step"
+ _make_one_box_step(base_step, 60.0, 50.0, 15.0)
+
+ assembly_step = tmp_path / "assembly.step"
+ # Two 40x40x6 chills at "+z", offset +-10mm in x -- footprints
+ # [0,40] and [20,60] in x (both y in [5,45]) overlap in a 20mm-wide
+ # x in [20,40] band.
+ attachments = [
+ {"shape": "box", "size": (40.0, 40.0, 6.0), "at_face": "+z", "offset": (-10.0, 0.0)},
+ {"shape": "box", "size": (40.0, 40.0, 6.0), "at_face": "+z", "offset": (10.0, 0.0)},
+ ]
+ compose_assembly(str(base_step), attachments, str(assembly_step))
+
+ rows = list_solids(str(assembly_step))
+ casting_id = max(rows, key=lambda r: r["volume_mm3"])["solid_id"]
+ role_ids = sorted(r["solid_id"] for r in rows if r["solid_id"] != casting_id)
+ assert len(role_ids) == 2
+
+ out_msh = tmp_path / "casting.msh"
+ result = heal_and_mesh(str(assembly_step), str(out_msh), target_nodes=3000,
+ keep_solid_id=casting_id)
+ assert result["ok"] is True
+
+ tag_result = tag_contact_patches(str(out_msh), str(assembly_step), role_ids, tol_mm=2.5)
+
+ areas = tag_result["per_body_patch_area_mm2"]
+ assert set(areas) == set(role_ids)
+ for solid_id in role_ids:
+ assert areas[solid_id] > 0.0
+ # symmetric footprints (mirrored offsets) -> roughly equal per-body areas
+ assert areas[role_ids[0]] == pytest.approx(areas[role_ids[1]], rel=0.05)
+
+ mesh_after = parse_msh(str(out_msh))
+ adiab = mesh_after.tris["Adiabatic"]
+ assert adiab.shape[0] == tag_result["n_adiab_tris"]
+
+ true_union_area = float(_triangle_areas(mesh_after.node_xyz, adiab).sum())
+ # The overlap region's triangles are counted in BOTH bodies' areas, so
+ # the sum across bodies must exceed the true (single-counted) union
+ # area -- pins the documented double-counting behavior.
+ assert sum(areas.values()) > true_union_area * 1.10
+
+ # nodes_by_solid (Part-B roles task): each body's own node set may share
+ # nodes with the other's (the overlap band's shared triangles), but the
+ # UNION across both bodies' node sets must equal exactly the true
+ # Adiabatic node set -- no node is dropped or invented.
+ assert set(tag_result["nodes_by_solid"]) == set(role_ids)
+ nodes_union = set()
+ for solid_id in role_ids:
+ nodes_union.update(tag_result["nodes_by_solid"][solid_id])
+ assert nodes_union == set(adiab.reshape(-1).tolist())
+
+
+def test_rewrite_msh_adiabatic_too_few_flags_raises_clean_assertion_error(tmp_path):
+ """A flag array SHORTER than the number of Cooling-tagged triangle
+ elements in the file must raise a clean `AssertionError` -- not a bare
+ `StopIteration` (the old `next(flags)`-inside-the-loop bug: exhausting
+ an iterator this way raises `StopIteration`, which is not a deliberate,
+ documented failure mode)."""
+ step = tmp_path / "box.step"
+ _make_one_box_step(step, 30.0, 20.0, 10.0)
+ out_msh = tmp_path / "box.msh"
+ result = heal_and_mesh(str(step), str(out_msh), target_nodes=300)
+ assert result["ok"] is True
+
+ n_cooling = parse_msh(str(out_msh)).tris["Cooling"].shape[0]
+ too_few = np.zeros(n_cooling - 1, dtype=bool) # one short
+
+ with pytest.raises(AssertionError, match="fewer"):
+ _rewrite_msh_adiabatic(str(out_msh), too_few)
+
+
+def test_rewrite_msh_adiabatic_too_many_flags_raises_clean_assertion_error(tmp_path):
+ step = tmp_path / "box.step"
+ _make_one_box_step(step, 30.0, 20.0, 10.0)
+ out_msh = tmp_path / "box.msh"
+ result = heal_and_mesh(str(step), str(out_msh), target_nodes=300)
+ assert result["ok"] is True
+
+ n_cooling = parse_msh(str(out_msh)).tris["Cooling"].shape[0]
+ too_many = np.zeros(n_cooling + 1, dtype=bool) # one extra
+
+ with pytest.raises(AssertionError, match="more flag"):
+ _rewrite_msh_adiabatic(str(out_msh), too_many)
diff --git a/tests/test_assembly_multisleeve.py b/tests/test_assembly_multisleeve.py
new file mode 100644
index 0000000000000000000000000000000000000000..d00cbd159a58b83bc650a16a26bf54448db38df2
--- /dev/null
+++ b/tests/test_assembly_multisleeve.py
@@ -0,0 +1,105 @@
+"""Multi-sleeve resolution (SP4 spec ADDENDUM).
+
+Real castings carry several sleeves; the pre-SP4 code accepted exactly one.
+These tests pin the three cases: identical sleeves resolve exactly, differing
+sleeves need explicit weights, and mixed sleeve ROLES still raise.
+"""
+import pytest
+
+from src.assembly import validate_bodies
+
+CASTING = {"solid_id": 1, "role": "casting"}
+
+
+def _exo(solid_id, sleeve):
+ return {"solid_id": solid_id, "role": "exo_sleeve", "params": {"sleeve": sleeve}}
+
+
+def test_three_identical_sleeves_resolve_exactly():
+ plan = validate_bodies([CASTING, _exo(2, "SD_4_7K"), _exo(3, "SD_4_7K"),
+ _exo(4, "SD_4_7K")])
+ assert plan.bc_config == "exo_sleeve"
+ assert sorted(plan.role_solid_ids) == [2, 3, 4]
+ assert plan.sleeve_MEF == pytest.approx(1.547)
+ assert plan.mef_weighted is False
+ assert plan.mef_spread == pytest.approx(0.0)
+ assert plan.exo_q0 > 0.0
+ assert plan.role_by_solid == {2: "exo_sleeve", 3: "exo_sleeve", 4: "exo_sleeve"}
+
+
+def test_differing_mef_without_weights_raises():
+ with pytest.raises(ValueError, match="differing MEF"):
+ validate_bodies([CASTING, _exo(2, "SD_4_7K"), _exo(3, "SD_5_8K")])
+
+
+def test_differing_mef_error_points_at_the_two_pass_path():
+ """Review I3: `predict_casting` cannot supply `mef_weights` (contact-
+ patch areas don't exist until after meshing/tagging), so the message a
+ user actually sees must not just name the missing parameter -- it must
+ say where areas come from and which callers already do the two-pass
+ resolution correctly."""
+ with pytest.raises(ValueError, match="tag_contact_patches") as exc:
+ validate_bodies([CASTING, _exo(2, "SD_4_7K"), _exo(3, "SD_5_8K")])
+ msg = str(exc.value)
+ assert "tools/sp4/prep.py" in msg
+ assert "tools/sp4/score.py" in msg
+
+
+def test_differing_mef_area_weighted():
+ # SD_4_7K MEF 1.547 over 300 mm2, SD_5_8K MEF 1.579 over 100 mm2
+ # (2026-08-09 sleeve-mef-correction: catalog_klx/catalog_geom, was
+ # 1.49/1.63 before the fix) -> (1.547*300 + 1.579*100) / 400 = 1.555
+ plan = validate_bodies([CASTING, _exo(2, "SD_4_7K"), _exo(3, "SD_5_8K")],
+ mef_weights={2: 300.0, 3: 100.0})
+ assert plan.sleeve_MEF == pytest.approx(1.555)
+ assert plan.mef_weighted is True
+ assert plan.mef_spread == pytest.approx(1.579 - 1.547)
+ assert sorted(plan.role_solid_ids) == [2, 3]
+
+
+def test_weights_missing_a_sleeve_raises():
+ with pytest.raises(ValueError, match="mef_weights"):
+ validate_bodies([CASTING, _exo(2, "SD_4_7K"), _exo(3, "SD_5_8K")],
+ mef_weights={2: 300.0})
+
+
+def test_nonpositive_weight_raises():
+ with pytest.raises(ValueError, match="positive"):
+ validate_bodies([CASTING, _exo(2, "SD_4_7K"), _exo(3, "SD_5_8K")],
+ mef_weights={2: 300.0, 3: 0.0})
+
+
+def test_mixed_sleeve_roles_still_raise():
+ bodies = [CASTING, _exo(2, "SD_4_7K"),
+ {"solid_id": 3, "role": "insulating_sleeve",
+ "params": {"sleeve": "FOSECO_V90L"}}]
+ with pytest.raises(ValueError, match="one role"):
+ validate_bodies(bodies)
+
+
+def test_single_sleeve_unchanged():
+ """The pre-SP4 single-sleeve path must be byte-identical."""
+ plan = validate_bodies([CASTING, _exo(2, "SD_4_7K")])
+ assert plan.role_solid_ids == [2]
+ assert plan.sleeve_MEF == pytest.approx(1.547)
+ assert plan.mef_weighted is False
+ assert plan.mef_spread == pytest.approx(0.0)
+
+
+def test_manual_mef_override_across_several_sleeves():
+ bodies = [CASTING,
+ {"solid_id": 2, "role": "exo_sleeve", "params": {"MEF": 1.49}},
+ {"solid_id": 3, "role": "exo_sleeve", "params": {"MEF": 1.49}}]
+ plan = validate_bodies(bodies)
+ assert plan.sleeve_MEF == pytest.approx(1.49)
+ assert plan.mef_weighted is False
+
+
+def test_chills_path_unchanged():
+ bodies = [CASTING,
+ {"solid_id": 2, "role": "chill", "params": {"material": "cast_iron"}},
+ {"solid_id": 3, "role": "chill", "params": {"material": "cast_iron"}}]
+ plan = validate_bodies(bodies)
+ assert plan.bc_config == "chills"
+ assert plan.n_chills == 2
+ assert plan.mef_weighted is False
diff --git a/tests/test_features_inference.py b/tests/test_features_inference.py
new file mode 100644
index 0000000000000000000000000000000000000000..055e7a912414b25e2c8935c7cf9626afc3711cf7
--- /dev/null
+++ b/tests/test_features_inference.py
@@ -0,0 +1,16 @@
+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
new file mode 100644
index 0000000000000000000000000000000000000000..dfedd289085af93032f7174e5cecfe6802288c66
--- /dev/null
+++ b/tests/test_features_v3.py
@@ -0,0 +1,36 @@
+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
new file mode 100644
index 0000000000000000000000000000000000000000..a508d4eeb9cd80e47a6e48d50ff6d3cdb109660c
--- /dev/null
+++ b/tests/test_features_v4.py
@@ -0,0 +1,275 @@
+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
new file mode 100644
index 0000000000000000000000000000000000000000..53c6155c7f9277d2502691779ba649254401d2c3
--- /dev/null
+++ b/tests/test_film.py
@@ -0,0 +1,46 @@
+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_geometry.py b/tests/test_geometry.py
new file mode 100644
index 0000000000000000000000000000000000000000..1ad7ffc04393d9e700c0d4cbc5abd399d9f57b0a
--- /dev/null
+++ b/tests/test_geometry.py
@@ -0,0 +1,104 @@
+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)
+
+
+def test_triangle_area_right_triangle():
+ xyz = np.array([[0, 0, 0], [4, 0, 0], [0, 3, 0]], float)
+ tri = np.array([[0, 1, 2]])
+ np.testing.assert_allclose(triangle_areas(xyz, tri), [6.0]) # 0.5*4*3
+ np.testing.assert_allclose(total_area(xyz, tri), 6.0)
+
+
+def test_tet_volume_unit():
+ # tetrahedron with legs 3,3,3 along axes from origin -> vol = 3^3/6 = 4.5
+ xyz = np.array([[0, 0, 0], [3, 0, 0], [0, 3, 0], [0, 0, 3]], float)
+ tets = np.array([[0, 1, 2, 3]])
+ np.testing.assert_allclose(mesh_volume(xyz, tets), 4.5)
+
+
+def test_nearest_distance():
+ pts = np.array([[0, 0, 0], [10, 0, 0]], float)
+ tgt = np.array([[0, 0, 0], [1, 0, 0]], float)
+ np.testing.assert_allclose(nearest_distance(pts, tgt), [0.0, 9.0])
+
+
+def test_nearest_distance_empty_targets():
+ pts = np.array([[0, 0, 0]], float)
+ tgt = np.empty((0, 3), float)
+ d = nearest_distance(pts, tgt)
+ assert np.isinf(d).all()
+
+
+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
new file mode 100644
index 0000000000000000000000000000000000000000..a2c2a93cbd91f2dfc6012e057c124ae796cf5b86
--- /dev/null
+++ b/tests/test_graph_build.py
@@ -0,0 +1,52 @@
+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
new file mode 100644
index 0000000000000000000000000000000000000000..0745f9512394191bf1d676d46e13eb4db960828e
--- /dev/null
+++ b/tests/test_hybrid.py
@@ -0,0 +1,516 @@
+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)
+
+
+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():
+ # gnn std is exactly 10 (two points, mean 10, deviations -+10)
+ gnn = np.array([0.0, 20.0])
+ total = 1000.0 # large enough that nothing clips
+ pred, info = hybrid_field_affine(gnn, total, std_seconds=20.0, q=0.995)
+ assert info["a"] == 2.0
+ np.testing.assert_allclose(np.std(pred), 20.0, atol=1e-9)
+ assert info["clip_frac"] == 0.0
+
+
+def test_hybrid_affine_total_anchor():
+ gnn = np.array([0.0, 20.0])
+ total = 1000.0 # large enough that nothing clips
+ pred, info = hybrid_field_affine(gnn, total, std_seconds=20.0, q=0.995)
+ assert np.quantile(pred, 0.995) == pytest.approx(total, abs=1e-9)
+ assert info["clip_frac"] == 0.0
+
+
+def test_hybrid_affine_preserves_ranking():
+ gnn = np.array([3.0, 1.0, 9.0, 5.0])
+ total = 1000.0
+ pred, info = hybrid_field_affine(gnn, total, std_seconds=5.0)
+ assert info["a"] > 0.0
+ assert list(np.argsort(pred)) == list(np.argsort(gnn))
+
+
+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.
+ gnn = np.array([0.0, 20.0])
+ total = 1000.0
+ pred, info = hybrid_field_affine(gnn, total, std_seconds=100.0)
+ assert info["a"] == 5.0
+ assert info["a_clamped"] is True
+
+
+def test_hybrid_affine_degenerate_fallback(caplog):
+ gnn = np.array([5.0, 5.0, 5.0])
+ total = 100.0
+ with caplog.at_level(logging.WARNING):
+ pred, info = hybrid_field_affine(gnn, total, std_seconds=20.0, q=0.995)
+ expected = hybrid_field_max_anchor(gnn, total, q=0.995)
+ np.testing.assert_allclose(pred, expected)
+ assert info["degenerate_fallback"] is True
+ assert any("degenerate" in rec.message.lower() for rec in caplog.records)
+
+
+def test_hybrid_affine_clip_negative():
+ # gnn std is 50; the naive (pre-amendment) closed form a=std_seconds/s=1.0
+ # already clips the low node (raw would be [-89.5, 10.5]), so the
+ # clip-aware solver kicks in. With only 2 nodes the post-clip std is capped
+ # well below 50 everywhere in the clamp range (max achievable is at
+ # a=5.0), so this lands on the a_clamp ceiling.
+ gnn = np.array([0.0, 100.0])
+ total = 10.0
+ pred, info = hybrid_field_affine(gnn, total, std_seconds=50.0)
+ np.testing.assert_allclose(pred, [0.0, 12.5])
+ assert info["a"] == 5.0
+ assert info["a_clamped"] is True
+ assert info["clip_frac"] == pytest.approx(0.5)
+
+
+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
+ # 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
+ # still lands on-target.
+ g = np.random.default_rng(1).exponential(5.0, size=200)
+ s = np.std(g)
+ total = 8.0
+ target = 0.3 * s
+
+ # sanity: the pre-amendment closed form (a = target/s, unclamped) clips
+ # ~30% of nodes here -- this is the regime the amendment targets.
+ a_naive = target / s
+ centered_naive = a_naive * (g - g.mean())
+ m_naive = total - np.quantile(centered_naive, 0.995)
+ naive_clip_frac = float((centered_naive + m_naive < 0.0).mean())
+ assert 0.25 <= naive_clip_frac <= 0.35
+
+ pred, info = hybrid_field_affine(g, total, target)
+ assert info["a_clamped"] is False
+ assert np.std(pred) == pytest.approx(target, rel=1e-3)
+
+
+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
+ # 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).
+ gnn = np.arange(10.0)
+ total = 10.5
+ pred, info = hybrid_field_affine(gnn, total, std_seconds=8.0)
+ assert info["a"] != 5.0 # must NOT blindly clamp to the ceiling
+ assert info["a_clamped"] is True
+ achieved = np.std(pred)
+ assert achieved == pytest.approx(3.816200508029165, rel=1e-6)
+ assert achieved < 8.0
+ # the peak (this test's outcome) beats the naive hi=5.0 clamp on fit:
+ at_hi = np.clip(5.0 * (gnn - gnn.mean() - np.quantile(gnn - gnn.mean(), 0.995)) + total,
+ 0.0, None)
+ 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():
+ # Invariant: when nothing clips at the closed-form solution, the amended
+ # (clip-aware) function must reproduce the pre-amendment closed form
+ # exactly (it takes it as a fast path; bisection is never invoked).
+ def pre_amendment_closed_form(gnn_seconds, total_seconds, std_seconds,
+ q=0.995, a_clamp=(0.3, 5.0)):
+ g = np.asarray(gnn_seconds, dtype=np.float64)
+ s = float(np.std(g))
+ a_raw = float(std_seconds) / s
+ a = float(np.clip(a_raw, a_clamp[0], a_clamp[1]))
+ centered = a * (g - g.mean())
+ m = float(total_seconds) - float(np.quantile(centered, q))
+ out = centered + m
+ return np.clip(out, 0.0, None), a
+
+ gnn = np.array([3.0, 1.0, 9.0, 5.0])
+ total = 1000.0
+ std_seconds = 5.0
+ pred_new, info_new = hybrid_field_affine(gnn, total, std_seconds)
+ pred_old, a_old = pre_amendment_closed_form(gnn, total, std_seconds)
+ assert info_new["a_clamped"] is False
+ assert info_new["a"] == pytest.approx(a_old, rel=1e-6)
+ 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)
+# --------------------------------------------------------------------------- #
+def _squeezed_case():
+ # A field with a small (3%) group of low-value outlier nodes and a tight
+ # high-value cluster: the affine solve's hard quantile-anchor pins the
+ # top near `total`, so scaling enough to hit a large `std_seconds` clips
+ # the low outliers hard -- `a_clamped` fires (the two-target problem is
+ # unreachable via a pure affine map). Verified interactively: at these
+ # exact parameters the affine solve lands short of the std target
+ # (std ~169.8 vs target ~330.1, clip_frac 0.03, a_clamped True) while the
+ # soft solve's default-weight optimum (~203.2) is strictly closer, by
+ # trading a ~20% total-anchor deviation for less clipping.
+ rng = np.random.default_rng(7)
+ n = 200
+ n_low = int(0.03 * n)
+ n_high = n - n_low
+ high_vals = 1000.0 + rng.normal(0.0, 2.0, n_high)
+ low_vals = rng.uniform(0.0, 50.0, n_low)
+ g = np.concatenate([high_vals, low_vals])
+ total = 1000.0
+ target = 2.0 * float(np.std(g))
+ return g, total, target
+
+
+def test_hybrid_soft_exact_degeneration_matches_affine():
+ # No-clip case (from test_hybrid_affine_spread_match/total_anchor): the
+ # affine solve already hits both targets essentially exactly, so the soft
+ # solve must return that seed unchanged.
+ gnn = np.array([0.0, 20.0])
+ total = 1000.0
+ pred_affine, info_affine = hybrid_field_affine(gnn, total, std_seconds=20.0, q=0.995)
+ pred_soft, info_soft = hybrid_field_soft(gnn, total, std_seconds=20.0, q=0.995)
+ np.testing.assert_allclose(pred_soft, pred_affine, rtol=1e-9)
+ assert abs(info_soft["total_deviation_frac"]) < 1e-6
+ assert info_soft["degenerate_fallback"] is False
+
+
+def test_hybrid_soft_trade_off_engages_on_squeezed_case():
+ g, total, target = _squeezed_case()
+ pred_affine, info_affine = hybrid_field_affine(g, total, target)
+ assert info_affine["a_clamped"] is True # affine truly cannot reach std_t
+
+ pred_soft, info_soft = hybrid_field_soft(g, total, target)
+ d_affine = abs(np.std(pred_affine) - target)
+ d_soft = abs(np.std(pred_soft) - target)
+ assert d_soft < d_affine # soft's std is strictly closer to the target
+ assert info_soft["total_deviation_frac"] > 0 # the top rose above total
+
+
+def test_hybrid_soft_weight_monotonicity_on_squeezed_case():
+ g, total, target = _squeezed_case()
+ pred_hi, info_hi = hybrid_field_soft(g, total, target, w_total=100.0, w_std=1.0)
+ pred_lo, info_lo = hybrid_field_soft(g, total, target, w_total=0.01, w_std=1.0)
+ assert abs(info_hi["total_deviation_frac"]) <= abs(info_lo["total_deviation_frac"])
+ assert np.std(pred_lo) >= np.std(pred_hi)
+
+
+def test_hybrid_soft_determinism():
+ g, total, target = _squeezed_case()
+ pred1, info1 = hybrid_field_soft(g, total, target)
+ pred2, info2 = hybrid_field_soft(g, total, target)
+ np.testing.assert_array_equal(pred1, pred2)
+ assert info1 == info2
+
+
+def test_hybrid_soft_degenerate_flat_field_falls_back(caplog):
+ gnn = np.array([5.0, 5.0, 5.0])
+ total = 100.0
+ with caplog.at_level(logging.WARNING):
+ pred, info = hybrid_field_soft(gnn, total, std_seconds=20.0, q=0.995)
+ expected = hybrid_field_max_anchor(gnn, total, q=0.995)
+ np.testing.assert_allclose(pred, expected)
+ assert info["degenerate_fallback"] is True
+ 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).
+# --------------------------------------------------------------------------- #
+_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;
+ # 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
+ # exactly (via the j_seed < 1e-12 early return, not by chance).
+ rng = np.random.default_rng(11)
+ g = rng.lognormal(mean=5.0, sigma=0.4, size=5000)
+ total = 5000.0 # large relative to g's own scale -> no clipping at a~1
+ target = 0.8 * float(np.std(g)) # well within reach, a_raw < 1 < a_clamp[1]
+
+ pred_affine, info_affine = hybrid_field_affine(g, total, target)
+ assert info_affine["a_clamped"] is False
+ pred_soft, info_soft = hybrid_field_soft(g, total, target)
+
+ np.testing.assert_allclose(pred_soft, pred_affine, rtol=1e-9)
+ assert info_soft["target_unreachable"] is False
+ assert abs(info_soft["total_deviation_frac"]) < 1e-6
+
+
+def test_hybrid_soft_a_clamp_lo_nonpositive_raises():
+ gnn = np.array([1.0, 2.0, 3.0])
+ with pytest.raises(ValueError, match="a_clamp"):
+ hybrid_field_soft(gnn, 10.0, 1.0, a_clamp=(0.0, 5.0))
+ with pytest.raises(ValueError, match="a_clamp"):
+ hybrid_field_soft(gnn, 10.0, 1.0, a_clamp=(-1.0, 5.0))
+
+
+def test_hybrid_soft_negative_weight_raises():
+ # Review round 1, 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])
+ with pytest.raises(ValueError, match="w_total|w_std"):
+ hybrid_field_soft(gnn, 10.0, 1.0, w_total=-1.0)
+ with pytest.raises(ValueError, match="w_total|w_std"):
+ hybrid_field_soft(gnn, 10.0, 1.0, w_std=-1.0)
+
+
+def test_hybrid_soft_both_weights_zero_raises():
+ # Review round 1, 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"):
+ hybrid_field_soft(gnn, 10.0, 1.0, w_total=0.0, w_std=0.0)
+
+
+def test_hybrid_soft_one_weight_zero_is_allowed():
+ # Sanity: only BOTH-zero is forbidden -- a single zero weight (pure
+ # single-target optimization) is a legitimate, deliberate use case.
+ gnn = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
+ pred, info = hybrid_field_soft(gnn, 10.0, 1.0, w_total=0.0, w_std=1.0)
+ assert np.all(np.isfinite(pred))
+ pred, info = hybrid_field_soft(gnn, 10.0, 1.0, w_total=1.0, w_std=0.0)
+ assert np.all(np.isfinite(pred))
+
+
+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
+ # 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.
+ rng = np.random.default_rng(3)
+ g = rng.exponential(1.0, size=500) ** 2
+ g = g - g.min()
+ total = 3.0
+ target = 3.0 * float(np.std(g))
+ with caplog.at_level(logging.WARNING):
+ pred, info = hybrid_field_soft(g, total, target, w_total=0.1, w_std=10.0)
+ assert info["clip_frac"] > 0.5
+ assert any("hybrid_field_soft" in rec.message and "majority" in rec.message.lower()
+ 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"
diff --git a/tests/test_hybrid_max_anchor.py b/tests/test_hybrid_max_anchor.py
new file mode 100644
index 0000000000000000000000000000000000000000..094d77d5bce6614a06f579d8170c083ad640c461
--- /dev/null
+++ b/tests/test_hybrid_max_anchor.py
@@ -0,0 +1,94 @@
+"""Task 2/7: 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
+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
+exact-max behavior bit-for-bit (`np.quantile(x, 1.0) == x.max()`).
+"""
+import numpy as np
+
+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."""
+ 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
+ assert np.allclose(np.diff(out), np.diff(g)) # additive shift only
+
+
+def test_max_anchor_default_q_preserves_shape_and_meets_or_exceeds_total():
+ """Default q<1: still an additive shift; max(out) >= total (not ==)."""
+ g = np.array([1.0, 4.0, 2.5])
+ out = hybrid_field_max_anchor(g, 10.0)
+ assert out.max() >= 10.0
+ assert out.shape == g.shape
+ assert np.allclose(np.diff(out), np.diff(g)) # additive shift only
+
+
+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
+ `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)
+ total = 37.5
+
+ out = hybrid_field_max_anchor(g, total, q=1.0)
+ expected = np.clip(g - g.max() + total, 0.0, None)
+
+ assert np.array_equal(out, expected)
+
+
+def test_max_anchor_default_q_is_robust_to_a_single_outlier_node():
+ """The failure mode the promotion gate exposed: one wildly-inflated
+ GNN prediction anchored to the exact max drags ~all other nodes below
+ the zero clip. With the default quantile anchor (q=0.995), >=99% of a
+ 1000-node field stays > 0 for the same outlier + total."""
+ rng = np.random.default_rng(0)
+ bulk = rng.uniform(0.0, 100.0, size=1000)
+ g = np.concatenate([bulk, [1e6]]) # one extreme outlier node
+ total = 500.0
+
+ out_exact = hybrid_field_max_anchor(g, total, q=1.0)
+ out_default = hybrid_field_max_anchor(g, total)
+
+ # Old exact-max behavior: the outlier alone sets the anchor, dragging
+ # essentially the whole bulk to the zero clip.
+ assert (out_exact[:-1] > 0).mean() < 0.01
+
+ # New default-q behavior: the bulk survives the clip.
+ assert (out_default[:-1] > 0).mean() >= 0.99
+ assert out_default.shape == g.shape
+
+
+def test_max_anchor_argsort_unchanged_between_q1_and_default_q():
+ """Hot-spot ranking must be unaffected by the anchor choice: both are
+ per-graph additive shifts of the same field, so (absent clipping
+ collisions) node order is identical regardless of q."""
+ rng = np.random.default_rng(7)
+ g = rng.uniform(10.0, 200.0, size=500)
+ # total large enough that neither anchor choice clips any node at zero
+ # (min(g) - anchor + total > 0 for anchor <= max(g)): isolates ordering
+ # from the (expected, separately-tested) clipping behavior.
+ total = 300.0
+
+ out_q1 = hybrid_field_max_anchor(g, total, q=1.0)
+ out_default = hybrid_field_max_anchor(g, total)
+
+ assert np.array_equal(np.argsort(out_q1), np.argsort(out_default))
diff --git a/tests/test_inference_props.py b/tests/test_inference_props.py
new file mode 100644
index 0000000000000000000000000000000000000000..323cac45b91a4cf4d988b192d218481c8e1a1923
--- /dev/null
+++ b/tests/test_inference_props.py
@@ -0,0 +1,125 @@
+import json, pytest
+from src.inference_props import resolve_props, list_alloys, list_molds, DEFAULT_ALLOY, \
+ resolve_alloy_family
+from src.shrinkage import SHRINKAGE_PCT
+
+REQ = {"k", "rho", "cp", "L_kJ_kg", "TL", "TS", "Tpour", "mold_effusivity"}
+
+def test_default_is_gjl250_green_sand():
+ p = resolve_props()
+ assert set(p) == REQ and all(isinstance(v, float) for v in p.values())
+ assert p["TL"] == 1220.0 and p["TS"] == 1150.0 # EN-GJL-250 row
+ assert p["Tpour"] == pytest.approx((1310 + 1330) / 2) # midpoint rule
+ assert p["mold_effusivity"] == pytest.approx(1018.0) # molds.csv, NOT 1250
+
+def test_named_alloy_and_lists():
+ assert DEFAULT_ALLOY in list_alloys() and "green_sand" in list_molds()
+ p = resolve_props(alloy="EN-GJS-400-15")
+ assert p["TS"] == 1120.0
+
+def test_unknown_names_raise_listing_options():
+ with pytest.raises(ValueError, match="EN-GJL-250"):
+ resolve_props(alloy="unobtainium")
+ with pytest.raises(ValueError, match="green_sand"):
+ resolve_props(mold="vacuum")
+
+def test_props_json_override_and_missing_keys(tmp_path):
+ full = dict(k=30.0, rho=7000.0, cp=500.0, L_kJ_kg=220.0, TL=1200.0, TS=1120.0,
+ Tpour=1350.0, mold_effusivity=900.0)
+ f = tmp_path / "p.json"; f.write_text(json.dumps(full))
+ assert resolve_props(props_json=str(f)) == full
+ f2 = tmp_path / "p2.json"; f2.write_text(json.dumps({"k": 30.0}))
+ with pytest.raises(ValueError, match="L_kJ_kg"):
+ resolve_props(props_json=str(f2))
+
+
+# --------------------------------------------------------------------------- #
+# Review finding I2 (final whole-branch review, 2026-08-01): 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):
+ inverted = dict(k=30.0, rho=7000.0, cp=500.0, L_kJ_kg=220.0, TL=1100.0, TS=1200.0,
+ Tpour=1350.0, mold_effusivity=900.0)
+ f = tmp_path / "inverted.json"; f.write_text(json.dumps(inverted))
+ with pytest.raises(ValueError, match=r"TL.*1100.*TS.*1200|TL.*must be strictly greater"):
+ resolve_props(props_json=str(f))
+
+
+def test_props_json_tl_equal_ts_raises(tmp_path):
+ # degenerate (TL == TS), not just inverted -- niyama_proxy's tdot_proxy
+ # would be exactly (TL-TS)/t_s == 0 for every node either way.
+ degenerate = dict(k=30.0, rho=7000.0, cp=500.0, L_kJ_kg=220.0, TL=1150.0, TS=1150.0,
+ Tpour=1350.0, mold_effusivity=900.0)
+ f = tmp_path / "degenerate.json"; f.write_text(json.dumps(degenerate))
+ with pytest.raises(ValueError, match="TL"):
+ resolve_props(props_json=str(f))
+
+
+def test_named_alloy_tl_ts_still_valid_after_i2_guard():
+ # sanity: the new guard must not break any real named grade (already
+ # validated by tools.doe.load_properties.load_materials's own TS 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():
+ # None -> DEFAULT_ALLOY (EN-GJL-250) -> grey_iron, not assumed.
+ assert resolve_alloy_family(None) == ("grey_iron", False)
+
+
+def test_resolve_alloy_family_unknown_prefix_falls_back_loud():
+ family, assumed = resolve_alloy_family("some_unknown_grade")
+ 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
+ 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
+ FUTURE grade added to the catalog under a name none of
+ `_ALLOY_FAMILY_PREFIXES` recognizes: that grade would otherwise silently
+ default to "steel" (a plausible-looking but wrong family for, say, a new
+ aluminum or a differently-prefixed steel grade) with no test ever
+ catching it."""
+ for grade in list_alloys():
+ family, assumed = resolve_alloy_family(grade)
+ assert assumed is False, (
+ f"catalog grade {grade!r} fell back to the 'assumed' family map "
+ f"(resolved family={family!r}) -- add a matching prefix to "
+ f"src.inference_props._ALLOY_FAMILY_PREFIXES")
+ assert family in SHRINKAGE_PCT
diff --git a/tests/test_input_encoder_extension.py b/tests/test_input_encoder_extension.py
new file mode 100644
index 0000000000000000000000000000000000000000..1b19ee257c6e49c0926af0ff1916614405cb873b
--- /dev/null
+++ b/tests/test_input_encoder_extension.py
@@ -0,0 +1,9 @@
+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
new file mode 100644
index 0000000000000000000000000000000000000000..7f6ae6a7e24374b76009f450a64b760e947aacb0
--- /dev/null
+++ b/tests/test_local_modulus.py
@@ -0,0 +1,247 @@
+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):
+ pts = np.array([[x, y, z] for x in (0, a) for y in (0, a) for z in (0, a)], float)
+ dt = Delaunay(pts)
+ tets = dt.simplices.astype(np.int64)
+ cool = dt.convex_hull.astype(np.int64) # all 12 boundary triangles
+ edges = edges_from_tets(tets)
+ elen = np.linalg.norm(pts[edges[:, 0]] - pts[edges[:, 1]], axis=1)
+ return pts, tets, cool, edges, elen
+
+
+def _box_grid(n=4, a=10.0):
+ """A structured grid of n^3 points spanning a cube of side `a`, Delaunay
+ tetrahedralized, with the convex-hull triangles as the cooling surface.
+ Denser than `_cube` (n=2) -> stands in for a training-density mesh."""
+ 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)
+ 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
+
+
+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))
+ assert np.isclose(mesh_volume(xyz, tets), 1000.0)
+
+
+def test_large_radius_equals_global_v_over_a():
+ xyz, tets, cool, edges, elen = _cube()
+ m = local_modulus(xyz, tets, cool, edges, elen, radii=[1e6])
+ expected = mesh_volume(xyz, tets) / total_area(xyz, cool) # 1000/600
+ assert m.shape == (8, 1) and np.allclose(m[:, 0], expected, rtol=1e-6)
+
+
+def test_no_cooling_gives_bbox_sentinel():
+ xyz, tets, _, edges, elen = _cube()
+ m = local_modulus(xyz, tets, np.empty((0, 3), np.int64), edges, elen, radii=[1e6])
+ assert np.allclose(m[:, 0], np.linalg.norm(xyz.max(0) - xyz.min(0)))
+
+
+def test_tiny_radius_gives_sentinel():
+ # ball = the node itself: no complete tet, no complete cooling tri -> sentinel
+ xyz, tets, cool, edges, elen = _cube()
+ m = local_modulus(xyz, tets, cool, edges, elen, radii=[0.1])
+ assert np.allclose(m[:, 0], np.linalg.norm(xyz.max(0) - xyz.min(0)))
+
+
+def test_multiple_radii_columns_ordered():
+ xyz, tets, cool, edges, elen = _cube()
+ m = local_modulus(xyz, tets, cool, edges, elen, radii=[0.1, 1e6])
+ 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
new file mode 100644
index 0000000000000000000000000000000000000000..470e4140e9f9b726fc8e31b8850043a91dc6ffa7
--- /dev/null
+++ b/tests/test_mesh_cache.py
@@ -0,0 +1,312 @@
+import json, os
+import pytest
+import src.mesh_cache as mc
+
+def _fake_heal(calls, ok=True, n_nodes=5000):
+ def fake(step_path, out_msh, target_nodes=8500, **kw):
+ calls.append((step_path, target_nodes))
+ if ok:
+ with open(out_msh, "w") as f:
+ f.write("MESH " + str(n_nodes))
+ return {"ok": ok, "n_nodes": n_nodes if ok else 0, "notes": "" if ok else "boom"}
+ return fake
+
+
+def _recording_heal(calls, results):
+ """A monkeypatch fake that records the `defeature` kwarg passed on each
+ call (`calls` accumulates dicts `{"defeature": bool}`) and returns
+ `results[i]` for the i-th call (clamped to the last entry once
+ exhausted) -- lets a single fake script a "fails plain, succeeds on
+ defeature retry" sequence for the trigger-policy tests below."""
+ def fake(step_path, out_msh, target_nodes=8500, defeature=False, **kw):
+ idx = min(len(calls), len(results) - 1)
+ calls.append({"defeature": defeature})
+ r = results[idx]
+ if r["ok"]:
+ with open(out_msh, "w") as f:
+ f.write("MESH " + str(r["n_nodes"]))
+ return dict(r)
+ return fake
+
+def test_miss_then_hit(tmp_path, monkeypatch):
+ step = tmp_path / "part.step"; step.write_bytes(b"solid-bytes-1")
+ 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 len(calls) == 1 # second call never meshed
+ assert open(tmp_path / "b.msh").read() == open(tmp_path / "a.msh").read()
+ assert r2["n_nodes"] == 5000 and r2["pipeline_version"] == mc.MESH_PIPELINE_VERSION
+
+def test_key_changes_with_content_target_and_version(tmp_path, monkeypatch):
+ step = tmp_path / "p.step"; step.write_bytes(b"AAA")
+ k1 = mc.mesh_cache_key(str(step), 8500)
+ step.write_bytes(b"BBB")
+ k2 = mc.mesh_cache_key(str(step), 8500)
+ k3 = mc.mesh_cache_key(str(step), 5000)
+ monkeypatch.setattr(mc, "MESH_PIPELINE_VERSION", mc.MESH_PIPELINE_VERSION + 1)
+ k4 = mc.mesh_cache_key(str(step), 8500)
+ assert len({k1, k2, k3, k4}) == 4
+
+def test_failure_not_cached(tmp_path, monkeypatch):
+ step = tmp_path / "p.step"; step.write_bytes(b"X")
+ calls = []
+ 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).
+ 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
+# 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
+# any of this (the policy is a deterministic function of the input file,
+# so a given STEP always resolves to the same `defeature` choice).
+# --------------------------------------------------------------------------
+_BIG_FILE_THRESHOLD_BYTES = 8 * 1024 * 1024
+
+
+def test_trigger_policy_big_file_single_defeature_call(tmp_path, monkeypatch):
+ step = tmp_path / "big.step"
+ step.write_bytes(b"x" * (_BIG_FILE_THRESHOLD_BYTES + 1))
+ calls = []
+ monkeypatch.setattr(mc, "heal_and_mesh",
+ _recording_heal(calls, [{"ok": True, "n_nodes": 999, "notes": ""}]))
+ r = mc.get_or_mesh(str(step), str(tmp_path / "a.msh"), 8500, cache_dir=str(tmp_path / "c"))
+ assert len(calls) == 1
+ assert calls[0]["defeature"] is True
+ assert r["ok"] is True and r["defeature"] is True and r["retry"] is False
+
+
+def test_trigger_policy_small_file_ok_plain_no_retry(tmp_path, monkeypatch):
+ step = tmp_path / "small.step"
+ step.write_bytes(b"tiny")
+ calls = []
+ monkeypatch.setattr(mc, "heal_and_mesh",
+ _recording_heal(calls, [{"ok": True, "n_nodes": 111, "notes": ""}]))
+ r = mc.get_or_mesh(str(step), str(tmp_path / "a.msh"), 8500, cache_dir=str(tmp_path / "c"))
+ assert len(calls) == 1
+ assert calls[0]["defeature"] is False
+ assert r["ok"] is True and r["defeature"] is False and r["retry"] is False
+
+
+def test_trigger_policy_small_file_retries_with_defeature_on_failure(tmp_path, monkeypatch):
+ step = tmp_path / "small.step"
+ step.write_bytes(b"tiny")
+ calls = []
+ monkeypatch.setattr(mc, "heal_and_mesh", _recording_heal(calls, [
+ {"ok": False, "n_nodes": 0, "notes": "boom"},
+ {"ok": True, "n_nodes": 777, "notes": ""},
+ ]))
+ r = mc.get_or_mesh(str(step), str(tmp_path / "a.msh"), 8500, cache_dir=str(tmp_path / "c"))
+ assert len(calls) == 2
+ assert calls[0]["defeature"] is False
+ assert calls[1]["defeature"] is True
+ assert r["ok"] is True and r["defeature"] is True and r["retry"] is True
+ assert r["n_nodes"] == 777
+
+
+def test_trigger_policy_big_file_defeature_fails_plain_retry_succeeds(tmp_path, monkeypatch):
+ """Symmetric with the small-file retry: a big file's first (defeature=True)
+ attempt can also fail outright, and must then retry ONCE with
+ defeature=False -- not be left with a single, permanently-failed attempt.
+ This must fail on pre-fix code (which had no big-file fallback at all)."""
+ step = tmp_path / "big.step"
+ step.write_bytes(b"x" * (_BIG_FILE_THRESHOLD_BYTES + 1))
+ calls = []
+ monkeypatch.setattr(mc, "heal_and_mesh", _recording_heal(calls, [
+ {"ok": False, "n_nodes": 0, "notes": "boom"},
+ {"ok": True, "n_nodes": 321, "notes": ""},
+ ]))
+ r = mc.get_or_mesh(str(step), str(tmp_path / "a.msh"), 8500, cache_dir=str(tmp_path / "c"))
+ assert len(calls) == 2
+ assert calls[0]["defeature"] is True
+ assert calls[1]["defeature"] is False
+ assert r["ok"] is True and r["defeature"] is False and r["retry"] is True
+ assert r["n_nodes"] == 321
+ assert r["cache"] == "miss"
+ # And now it must be cached: a second call must hit, with no further heal_and_mesh calls.
+ r2 = mc.get_or_mesh(str(step), str(tmp_path / "b.msh"), 8500, cache_dir=str(tmp_path / "c"))
+ assert r2["cache"] == "hit" and len(calls) == 2
+ assert r2["defeature"] is False and r2["retry"] is True
+
+
+def test_trigger_policy_small_file_fails_both_attempts_no_extra_retry(tmp_path, monkeypatch):
+ """Even if the defeature retry also fails, `get_or_mesh` must not retry
+ a third time -- exactly one retry, ever."""
+ step = tmp_path / "small.step"
+ step.write_bytes(b"tiny")
+ calls = []
+ monkeypatch.setattr(mc, "heal_and_mesh", _recording_heal(calls, [
+ {"ok": False, "n_nodes": 0, "notes": "boom1"},
+ {"ok": False, "n_nodes": 0, "notes": "boom2"},
+ ]))
+ r = mc.get_or_mesh(str(step), str(tmp_path / "a.msh"), 8500, cache_dir=str(tmp_path / "c"))
+ assert len(calls) == 2
+ assert r["ok"] is False and r["defeature"] is True and r["retry"] is True
+ assert r["cache"] == "miss"
+
+
+def test_tmp_names_are_unique_per_writer(tmp_path):
+ # Two concurrent workers racing on the SAME cache miss must never derive
+ # the same tmp path from the same final path -- otherwise their writes
+ # could interleave on one inode before either os.replace runs, and
+ # os.replace would then install a corrupted mesh for every later hit.
+ final = str(tmp_path / "c" / "somekey.msh")
+ tmp_a = mc._unique_tmp_path(final)
+ tmp_b = mc._unique_tmp_path(final)
+ assert tmp_a != tmp_b
+ 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_metrics.py b/tests/test_metrics.py
new file mode 100644
index 0000000000000000000000000000000000000000..73c63df6f99155929a8a861a4eb855901bc720c7
--- /dev/null
+++ b/tests/test_metrics.py
@@ -0,0 +1,34 @@
+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
new file mode 100644
index 0000000000000000000000000000000000000000..6b6d507331fb14fcecfe281367952494f44dcfd7
--- /dev/null
+++ b/tests/test_model_gnn.py
@@ -0,0 +1,187 @@
+import pytest
+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",
+}
+
+
+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 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)
+
+ # 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_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()
+
+ 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)
+
+ 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)
diff --git a/tests/test_msh_parse.py b/tests/test_msh_parse.py
new file mode 100644
index 0000000000000000000000000000000000000000..8d2bd93fceaf4edf2f860e0390a43a0d126a8674
--- /dev/null
+++ b/tests/test_msh_parse.py
@@ -0,0 +1,22 @@
+import numpy as np
+from src.msh_parse import parse_msh
+
+
+def test_parse_mini(mini_msh_file):
+ mesh = parse_msh(mini_msh_file)
+
+ # 5 nodes, in id order, in mm
+ assert mesh.node_xyz.shape == (5, 3)
+ np.testing.assert_allclose(mesh.node_xyz[0], [0, 0, 0])
+ np.testing.assert_allclose(mesh.node_xyz[4], [3, 3, 3])
+
+ # physical groups present, 0-based indices
+ assert set(mesh.tris.keys()) == {"Cooling", "Adiabatic"}
+ assert mesh.tris["Cooling"].shape == (2, 3)
+ assert mesh.tris["Adiabatic"].shape == (1, 3)
+ # cooling triangle 1 = nodes 1,2,3 (ids) -> indices 0,1,2
+ np.testing.assert_array_equal(sorted(mesh.tris["Cooling"][0]), [0, 1, 2])
+
+ # one tetrahedron, 0-based
+ assert mesh.tets.shape == (1, 4)
+ np.testing.assert_array_equal(sorted(mesh.tets[0]), [0, 1, 2, 3])
diff --git a/tests/test_ood.py b/tests/test_ood.py
new file mode 100644
index 0000000000000000000000000000000000000000..9a9c78b1336b06a67c508547a8ab9480a321f053
--- /dev/null
+++ b/tests/test_ood.py
@@ -0,0 +1,33 @@
+import numpy as np
+from src.ood import build_ood_ranges, ood_flags, save_ranges, load_ranges
+
+
+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
+
+
+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
+
+ # 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
diff --git a/tests/test_predict_casting_v4.py b/tests/test_predict_casting_v4.py
new file mode 100644
index 0000000000000000000000000000000000000000..b4e4dd0079845d8b7ad2ff8f5d8a86238524b472
--- /dev/null
+++ b/tests/test_predict_casting_v4.py
@@ -0,0 +1,1276 @@
+"""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.
+
+Part B Task 4: `predict_casting(bodies=...)` wiring -- `validate_bodies` ->
+mesh via the cache with `keep_solid_id`/`bodies_key` -> `tag_contact_patches`
+(cached as a `post_mesh` hook) -> the film spec/gbm bc_spec carry the
+resolved `BcPlan`. All of it is exercised with monkeypatched recorders
+(`validate_bodies`/`tag_contact_patches`/`get_or_mesh`/`predict_field`) --
+no real gmsh/model checkpoints are touched here either.
+"""
+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-GJL-250")
+ props_b = resolve_props(alloy="EN-GJS-500-7")
+ 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_with_bc_plan_uses_plan_feature_spec_and_gbm_bc_spec(monkeypatch):
+ """Part B Task 4: `predict_field(bc_plan=...)` feeds the film branch
+ `plan.feature_spec(...)` (the chill/sleeve spec) instead of the
+ all-sand literal, and `plan.gbm_bc_spec()` into `gbm_v2_total_seconds`.
+ Spies on the module-level `compute_features_inference` / `gbm_v2_total_seconds`
+ (both real, just observed) to capture exactly what spec/bc_spec were
+ built, without needing a real multi-body mesh."""
+ from src.assembly import BcPlan
+
+ mesh = _all_cooling_cube()
+ gnn = _make_film_gnn(seed=0)
+ hgb = _StubHGB(n_features_in_=16, total_seconds=50.0)
+ props = resolve_props()
+
+ plan = BcPlan(bc_config="chills", casting_solid_id=1, role_solid_ids=[2],
+ n_chills=2, chill_effusivity=9876.5)
+
+ 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, bc_plan=plan)
+
+ assert captured_spec["spec"]["bc_config"] == "chills"
+ assert captured_spec["spec"]["chill_effusivity"] == pytest.approx(9876.5)
+ assert captured_spec["spec"]["mold_effusivity"] == pytest.approx(props["mold_effusivity"])
+ assert captured_bc_spec["bc_spec"]["bc_config"] == "chills"
+ assert captured_bc_spec["bc_spec"]["n_chills"] == 2
+
+
+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 _cube_with_bc_patch(n=5, a=20.0):
+ """Same fixture as tests/test_features_v5.py: a Delaunay cube with the
+ z=0 face tagged 'Adiabatic' (BC patch) and every other hull face
+ 'Cooling' -- needed (unlike `_all_cooling_cube`) so a bc_plan's
+ chill/exo source patch has somewhere to sit."""
+ 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]
+ return Mesh(node_xyz=pts, tris={"Cooling": cool_tri, "Adiabatic": adiab_tri}, tets=tets)
+
+
+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()
+
+
+def test_v5_cfg_bodies_exo_patch_populates_exo_channels_not_chill(monkeypatch):
+ """(c) A bodies/BcPlan exo_sleeve patch resolved EXACTLY the way the
+ training builder does (`bc_plan.feature_spec(...)`, same spec keys)
+ must populate exo_mid/exo_long nonzero while leaving every chill
+ channel exactly zero -- the SAME spec resolution `compute_features_
+ inference`'s v4 branch already uses, now flowing into
+ `compute_bc_reach_channels` via the version=5 dispatch."""
+ from src.assembly import BcPlan
+
+ mesh = _cube_with_bc_patch()
+ gnn = _make_film_gnn_v5(seed=0)
+ hgb = _StubHGB(n_features_in_=16, total_seconds=50.0)
+ props = resolve_props()
+ plan = BcPlan(bc_config="exo_sleeve", casting_solid_id=1, role_solid_ids=[2],
+ sleeve_MEF=2.0, exo_q0=5000.0)
+
+ out, feat = predict_field(mesh, gnn, hgb, _norm20(), cfg=_film_cfg_v5(),
+ device="cpu", props=props, bc_plan=plan)
+ X = feat["X"]
+ idx = {n: i for i, n in enumerate(FEATURE_NAMES_V5)}
+ assert X[:, idx["exo_mid"]].max() > 0.0
+ assert X[:, idx["exo_long"]].max() > 0.0
+ assert (X[:, idx["chill_short"]] == 0.0).all()
+ assert (X[:, idx["chill_mid"]] == 0.0).all()
+ assert (X[:, idx["chill_long"]] == 0.0).all()
+
+
+def test_v5_cfg_bodies_chill_patch_populates_chill_channels_not_exo(monkeypatch):
+ """(c), vice versa: a chills BcPlan populates chill_short/mid/long
+ nonzero while exo_mid/exo_long stay exactly zero."""
+ from src.assembly import BcPlan
+
+ mesh = _cube_with_bc_patch()
+ gnn = _make_film_gnn_v5(seed=0)
+ hgb = _StubHGB(n_features_in_=16, total_seconds=50.0)
+ props = resolve_props()
+ plan = BcPlan(bc_config="chills", casting_solid_id=1, role_solid_ids=[2],
+ n_chills=2, chill_effusivity=9876.5)
+
+ out, feat = predict_field(mesh, gnn, hgb, _norm20(), cfg=_film_cfg_v5(),
+ device="cpu", props=props, bc_plan=plan)
+ X = feat["X"]
+ idx = {n: i for i, n in enumerate(FEATURE_NAMES_V5)}
+ assert X[:, idx["chill_short"]].max() > 0.0
+ assert X[:, idx["chill_mid"]].max() > 0.0
+ assert X[:, idx["chill_long"]].max() > 0.0
+ assert (X[:, idx["exo_mid"]] == 0.0).all()
+ assert (X[:, idx["exo_long"]] == 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_tag_contact_patches(calls, per_body_patch_area_mm2=None, n_adiab_tris=7,
+ nodes_by_solid=None):
+ per_body_patch_area_mm2 = per_body_patch_area_mm2 or {2: 42.0}
+
+ def _fn(msh_path, step_path, role_solid_ids, tol_mm=None):
+ calls.setdefault("tag_contact_patches", []).append(
+ dict(msh_path=msh_path, step_path=step_path,
+ role_solid_ids=list(role_solid_ids), tol_mm=tol_mm))
+ result = {"per_body_patch_area_mm2": dict(per_body_patch_area_mm2),
+ "n_adiab_tris": n_adiab_tris}
+ if nodes_by_solid is not None:
+ result["nodes_by_solid"] = {k: list(v) for k, v in nodes_by_solid.items()}
+ 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"
+
+
+# --------------------------------------------------------------------------- #
+# Part B Task 4: predict_casting(bodies=...) wiring -- monkeypatched
+# recorders for validate_bodies/tag_contact_patches/get_or_mesh/predict_field,
+# no real gmsh/model checkpoints touched.
+# --------------------------------------------------------------------------- #
+def _fake_validate_bodies(calls, plan):
+ def _fn(bodies, **kw):
+ calls["validate_bodies_arg"] = bodies
+ return plan
+ return _fn
+
+
+def test_bodies_given_wires_keep_solid_id_bodies_key_and_tagging(
+ monkeypatch, mini_msh_file, tmp_path):
+ from src.assembly import BcPlan
+
+ calls = {}
+ plan = BcPlan(bc_config="chills", casting_solid_id=1, role_solid_ids=[2],
+ n_chills=1, chill_effusivity=12345.0)
+
+ 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))
+ monkeypatch.setattr(pc_mod, "validate_bodies", _fake_validate_bodies(calls, plan))
+ monkeypatch.setattr(pc_mod, "tag_contact_patches", _fake_tag_contact_patches(calls))
+
+ captured = {}
+
+ def fake_predict_field(mesh, gnn_model, hgb_model, norm, cfg=None, device=None, tta=0,
+ baseline_norm=None, props=None, bc_plan=None, hgb_std=None,
+ lambda_shrink=1.0, w_total=1.0, w_std=1.0):
+ captured["bc_plan"] = bc_plan
+ return np.arange(5, dtype=np.float64) + 1.0, {"X": np.zeros((5, 9))}
+
+ monkeypatch.setattr(pc_mod, "load_hybrid",
+ lambda *a, **kw: (None, None, {}, {"film": True}, None, None))
+ monkeypatch.setattr(pc_mod, "predict_field", fake_predict_field)
+
+ bodies = [{"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "chill", "params": {"material": "cast_iron"}}]
+
+ out = pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path),
+ ood_json="does-not-exist.json", bodies=bodies)
+
+ assert calls["validate_bodies_arg"] == bodies
+ assert calls["get_or_mesh"]["keep_solid_id"] == 1
+ assert calls["get_or_mesh"]["bodies_key"] != "none"
+ assert "tag_contact_patches" in calls
+ assert calls["tag_contact_patches"][0]["role_solid_ids"] == [2]
+ assert captured["bc_plan"] is plan
+ assert out["bc_config"] == "chills"
+ assert out["per_body_patch_area_mm2"] == {2: 42.0}
+
+
+# --------------------------------------------------------------------------- #
+# Shrinkage-screening spec ADDENDUM 2 (2026-08-01, owner-approved SCOPED
+# SHIP): Gate C passed on all-sand/sleeve-only inputs but FAILED on the
+# chill-bearing slice (rho 0.313) -- a chill-bearing BC plan
+# (`bc_config == "chills"`) must mark "niyama_confidence": "low" + a summary
+# caveat; sleeve-only/bare inputs stay "full". Follows the same
+# validate_bodies/BcPlan monkeypatch pattern as the Part B tests above.
+# --------------------------------------------------------------------------- #
+def test_bodies_chill_role_marks_niyama_confidence_low_with_summary_caveat(
+ monkeypatch, mini_msh_file, tmp_path):
+ from src.assembly import BcPlan
+
+ calls = {}
+ plan = BcPlan(bc_config="chills", casting_solid_id=1, role_solid_ids=[2],
+ n_chills=1, chill_effusivity=12345.0)
+
+ 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))
+ monkeypatch.setattr(pc_mod, "validate_bodies", _fake_validate_bodies(calls, plan))
+ monkeypatch.setattr(pc_mod, "tag_contact_patches", _fake_tag_contact_patches(calls))
+ monkeypatch.setattr(pc_mod, "load_hybrid",
+ lambda *a, **kw: (None, None, {}, {"film": True}, None, None))
+ monkeypatch.setattr(pc_mod, "predict_field",
+ lambda mesh, *a, **kw: (np.arange(5, dtype=np.float64) + 1.0,
+ {"X": np.zeros((5, 9))}))
+
+ bodies = [{"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "chill", "params": {"material": "cast_iron"}}]
+
+ out = pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path),
+ ood_json="does-not-exist.json", bodies=bodies)
+
+ assert out["shrinkage"] is not None
+ assert out["shrinkage"]["niyama_confidence"] == "low"
+ with open(out["summary"]) as f:
+ summary_text = f.read()
+ assert "niyama confidence: LOW" in summary_text
+ assert "rho 0.313" in summary_text
+
+
+def test_bodies_sleeve_only_role_keeps_niyama_confidence_full_no_caveat(
+ monkeypatch, mini_msh_file, tmp_path):
+ """A sleeve-only BC plan (bc_config="insulating_sleeve", no chill role
+ at all) is UNAFFECTED by the chills-only Gate C gap -- stays "full",
+ no caveat line."""
+ from src.assembly import BcPlan
+
+ calls = {}
+ plan = BcPlan(bc_config="insulating_sleeve", casting_solid_id=1, role_solid_ids=[2])
+
+ 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))
+ monkeypatch.setattr(pc_mod, "validate_bodies", _fake_validate_bodies(calls, plan))
+ monkeypatch.setattr(pc_mod, "tag_contact_patches", _fake_tag_contact_patches(calls))
+ monkeypatch.setattr(pc_mod, "load_hybrid",
+ lambda *a, **kw: (None, None, {}, {"film": True}, None, None))
+ monkeypatch.setattr(pc_mod, "predict_field",
+ lambda mesh, *a, **kw: (np.arange(5, dtype=np.float64) + 1.0,
+ {"X": np.zeros((5, 9))}))
+
+ bodies = [{"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "insulating_sleeve"}]
+
+ out = pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path),
+ ood_json="does-not-exist.json", bodies=bodies)
+
+ assert out["shrinkage"] is not None
+ assert out["shrinkage"]["niyama_confidence"] == "full"
+ with open(out["summary"]) as f:
+ summary_text = f.read()
+ assert "niyama confidence: LOW" not in summary_text
+ assert "rho 0.313" not in summary_text
+
+
+def test_bodies_chill_role_learned_niyama_source_uses_neutral_caveat(
+ monkeypatch, mini_msh_file, tmp_path):
+ """Learned-heads Task 3 review round 1, F3: when `predict_field`'s
+ `feat` carries learned `gradT`/`dTdt` (so `niyama_source` is "learned",
+ not "proxy"), the chill-bearing low-confidence caveat text must be the
+ NEUTRAL "learned-head Niyama pending slice validation" -- NOT the
+ proxy-specific "rho 0.313" text, a claim about the PROXY identity's own
+ measured Gate C shortfall that has nothing to do with a learned
+ checkpoint's heads. The marking itself ("low") is UNCHANGED -- still
+ keyed only on `plan.bc_config == "chills"`, per the brief ("stays until
+ the promote step")."""
+ from src.assembly import BcPlan
+
+ calls = {}
+ plan = BcPlan(bc_config="chills", casting_solid_id=1, role_solid_ids=[2],
+ n_chills=1, chill_effusivity=12345.0)
+
+ 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))
+ monkeypatch.setattr(pc_mod, "validate_bodies", _fake_validate_bodies(calls, plan))
+ monkeypatch.setattr(pc_mod, "tag_contact_patches", _fake_tag_contact_patches(calls))
+ monkeypatch.setattr(pc_mod, "load_hybrid",
+ lambda *a, **kw: (None, None, {}, {"film": True}, None, None))
+
+ def fake_predict_field(mesh, *a, **kw):
+ n = 5
+ seconds = np.arange(n, dtype=np.float64) + 1.0
+ feat = {"X": np.zeros((n, 9)), "gradT": np.full(n, 500.0), "dTdt": np.full(n, 2.0)}
+ return seconds, feat
+
+ monkeypatch.setattr(pc_mod, "predict_field", fake_predict_field)
+
+ bodies = [{"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "chill", "params": {"material": "cast_iron"}}]
+
+ out = pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path),
+ ood_json="does-not-exist.json", bodies=bodies)
+
+ assert out["shrinkage"] is not None
+ assert out["shrinkage"]["niyama_confidence"] == "low"
+ assert out["shrinkage"]["niyama_source"] == "learned"
+ with open(out["summary"]) as f:
+ summary_text = f.read()
+ assert "niyama confidence: LOW" in summary_text
+ assert "learned-head Niyama pending slice validation" in summary_text
+ assert "rho 0.313" not in summary_text
+
+
+# --------------------------------------------------------------------------- #
+# Part-B roles task (2026-08-03): role-aware feed-node selection, replacing
+# review finding I1's coarse `plan.bc_config == "chills"` special case (which
+# forced feed_nodes empty for the WHOLE input, rather than per body role).
+# `tag_contact_patches` relabels EVERY role body's contact patch "Adiabatic"
+# -- chills included -- so on a chill-bearing input every "Adiabatic"
+# triangle is a chill footprint, not a riser/sleeve; on a sleeve-bearing
+# input it genuinely is a feed patch. `mini_msh_file` already carries one
+# real "Adiabatic" tri (0-indexed nodes 1,2,3 -- file node ids 2,3,4), so
+# both branches are exercised against the SAME underlying mesh, with the
+# fake `tag_contact_patches` now also returning `nodes_by_solid={2: [1, 2,
+# 3]}` (what the real tagging pass would have found for that one role body)
+# and each `BcPlan` carrying the matching `role_by_solid` -- only the ROLE
+# differs between the two tests, not merely `bc_config`.
+# --------------------------------------------------------------------------- #
+def _spy_screen_casting(monkeypatch, captured):
+ real_screen_casting = pc_mod.screen_casting
+
+ def _fn(*args, **kwargs):
+ captured["feed_nodes"] = kwargs.get("feed_nodes")
+ return real_screen_casting(*args, **kwargs)
+
+ monkeypatch.setattr(pc_mod, "screen_casting", _fn)
+
+
+def test_bodies_chill_role_forces_feed_nodes_empty(monkeypatch, mini_msh_file, tmp_path):
+ from src.assembly import BcPlan
+
+ calls = {}
+ plan = BcPlan(bc_config="chills", casting_solid_id=1, role_solid_ids=[2],
+ n_chills=1, chill_effusivity=12345.0, role_by_solid={2: "chill"})
+
+ 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))
+ monkeypatch.setattr(pc_mod, "validate_bodies", _fake_validate_bodies(calls, plan))
+ monkeypatch.setattr(pc_mod, "tag_contact_patches",
+ _fake_tag_contact_patches(calls, nodes_by_solid={2: [1, 2, 3]}))
+ monkeypatch.setattr(pc_mod, "load_hybrid",
+ lambda *a, **kw: (None, None, {}, {"film": True}, None, None))
+ monkeypatch.setattr(pc_mod, "predict_field",
+ lambda mesh, *a, **kw: (np.arange(5, dtype=np.float64) + 1.0,
+ {"X": np.zeros((5, 9))}))
+ captured = {}
+ _spy_screen_casting(monkeypatch, captured)
+
+ bodies = [{"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "chill", "params": {"material": "cast_iron"}}]
+
+ out = pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path),
+ ood_json="does-not-exist.json", bodies=bodies)
+
+ # the chill body's own patch DID find nodes (1,2,3) -- they are simply
+ # never fed nodes, proving this is a real per-role exclusion and not an
+ # accident of an empty/missing nodes_by_solid.
+ assert captured["feed_nodes"].size == 0
+ # with feed_nodes forced empty, is_fed is always False -- EVERY hot
+ # region hot_regions itself finds must survive into macro_shrinkage
+ # (none silently dropped as "fed" by a chill contact patch).
+ from src.shrinkage import hot_regions
+ mesh = parse_msh(osp.join(str(tmp_path), f"{out['name']}.msh"))
+ seconds = np.arange(mesh.node_xyz.shape[0], dtype=np.float64) + 1.0
+ all_hot_regions = hot_regions(mesh.node_xyz, mesh.tets, seconds)
+ assert out["shrinkage"]["n_macro_regions"] == len(all_hot_regions)
+
+
+def test_bodies_sleeve_role_keeps_adiabatic_feed_nodes(monkeypatch, mini_msh_file, tmp_path):
+ from src.assembly import BcPlan
+
+ calls = {}
+ plan = BcPlan(bc_config="insulating_sleeve", casting_solid_id=1, role_solid_ids=[2],
+ role_by_solid={2: "insulating_sleeve"})
+
+ 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))
+ monkeypatch.setattr(pc_mod, "validate_bodies", _fake_validate_bodies(calls, plan))
+ monkeypatch.setattr(pc_mod, "tag_contact_patches",
+ _fake_tag_contact_patches(calls, nodes_by_solid={2: [1, 2, 3]}))
+ monkeypatch.setattr(pc_mod, "load_hybrid",
+ lambda *a, **kw: (None, None, {}, {"film": True}, None, None))
+ monkeypatch.setattr(pc_mod, "predict_field",
+ lambda mesh, *a, **kw: (np.arange(5, dtype=np.float64) + 1.0,
+ {"X": np.zeros((5, 9))}))
+ captured = {}
+ _spy_screen_casting(monkeypatch, captured)
+
+ bodies = [{"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "insulating_sleeve"}]
+
+ pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path),
+ ood_json="does-not-exist.json", bodies=bodies)
+
+ # sleeve assembly -> feed == that patch's nodes EXACTLY (spec §3).
+ assert captured["feed_nodes"].tolist() == [1, 2, 3]
+
+
+def test_bodies_sleeve_role_feed_nodes_byte_identical_to_pre_change_adiabatic_union(
+ monkeypatch, mini_msh_file, tmp_path):
+ """Screening byte-identity (spec §3): on a legal (single-role, non-mixed)
+ sleeve input, the NEW per-role feed_nodes selection must produce EXACTLY
+ the same result as the pre-change code (`_shrinkage_surface_and_feed_
+ nodes`'s plain "Adiabatic"-tag node union, which is what `predict_
+ casting` used unconditionally for any non-chills bc_config before this
+ task)."""
+ from src.assembly import BcPlan
+ from src.msh_parse import parse_msh as _parse_msh
+
+ calls = {}
+ plan = BcPlan(bc_config="exo_sleeve", casting_solid_id=1, role_solid_ids=[2],
+ role_by_solid={2: "exo_sleeve"})
+
+ 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))
+ monkeypatch.setattr(pc_mod, "validate_bodies", _fake_validate_bodies(calls, plan))
+ monkeypatch.setattr(pc_mod, "tag_contact_patches",
+ _fake_tag_contact_patches(calls, nodes_by_solid={2: [1, 2, 3]}))
+ monkeypatch.setattr(pc_mod, "load_hybrid",
+ lambda *a, **kw: (None, None, {}, {"film": True}, None, None))
+ monkeypatch.setattr(pc_mod, "predict_field",
+ lambda mesh, *a, **kw: (np.arange(5, dtype=np.float64) + 1.0,
+ {"X": np.zeros((5, 9))}))
+ captured = {}
+ _spy_screen_casting(monkeypatch, captured)
+
+ bodies = [{"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "exo_sleeve"}]
+
+ out = pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path),
+ ood_json="does-not-exist.json", bodies=bodies)
+
+ old_mesh = _parse_msh(osp.join(str(tmp_path), f"{out['name']}.msh"))
+ _old_surface, old_feed_nodes = pc_mod._shrinkage_surface_and_feed_nodes(old_mesh)
+
+ assert np.array_equal(captured["feed_nodes"], old_feed_nodes)
+
+
+def test_bodies_legacy_cache_hit_sleeve_falls_back_to_adiabatic_union_with_warning(
+ monkeypatch, mini_msh_file, tmp_path, caplog):
+ """Review round 1, H1: a mesh-cache HIT served from a side-car written
+ BEFORE the Part-B roles task never had `nodes_by_solid` at all --
+ `runs/mesh_cache/` holds exactly such legacy entries today (`per_body_
+ patch_area_mm2` present, `nodes_by_solid` absent), incl. both gold
+ parts. On a sleeve plan, feed_nodes must fall back to the PRE-CHANGE
+ rule (the plain "Adiabatic"-tag node union) rather than silently going
+ empty, and a loud warning naming the legacy situation must fire."""
+ from src.assembly import BcPlan
+
+ calls = {}
+ plan = BcPlan(bc_config="insulating_sleeve", casting_solid_id=1, role_solid_ids=[2],
+ role_by_solid={2: "insulating_sleeve"})
+
+ monkeypatch.setattr(pc_mod, "get_or_mesh",
+ _fake_get_or_mesh_hit_with_str_keys(calls, mini_msh_file, {"2": 42.0}))
+ monkeypatch.setattr(pc_mod, "mesh_step", _fake_mesh_step(calls, mini_msh_file))
+ monkeypatch.setattr(pc_mod, "validate_bodies", _fake_validate_bodies(calls, plan))
+ monkeypatch.setattr(pc_mod, "tag_contact_patches", _fake_tag_contact_patches(calls))
+ monkeypatch.setattr(pc_mod, "load_hybrid",
+ lambda *a, **kw: (None, None, {}, {"film": True}, None, None))
+ monkeypatch.setattr(pc_mod, "predict_field",
+ lambda mesh, *a, **kw: (np.arange(5, dtype=np.float64) + 1.0,
+ {"X": np.zeros((5, 9))}))
+ captured = {}
+ _spy_screen_casting(monkeypatch, captured)
+
+ bodies = [{"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "insulating_sleeve"}]
+
+ with caplog.at_level(logging.WARNING):
+ out = pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path),
+ ood_json="does-not-exist.json", bodies=bodies)
+
+ mesh = parse_msh(osp.join(str(tmp_path), f"{out['name']}.msh"))
+ _surface, expected_adiab_union = pc_mod._shrinkage_surface_and_feed_nodes(mesh)
+ assert expected_adiab_union.size > 0 # sanity: mini_msh_file's real Adiabatic tri
+ assert np.array_equal(captured["feed_nodes"], expected_adiab_union)
+ assert any("legacy" in rec.message.lower() for rec in caplog.records)
+
+
+def test_bodies_legacy_cache_hit_chills_falls_back_to_empty_with_warning(
+ monkeypatch, mini_msh_file, tmp_path, caplog):
+ """Same legacy-cache-HIT scenario as above, chills plan: falls back to
+ empty feed_nodes (the pre-change chills rule), with the same loud
+ warning -- and the warning fires regardless of `bc_config`."""
+ from src.assembly import BcPlan
+
+ calls = {}
+ plan = BcPlan(bc_config="chills", casting_solid_id=1, role_solid_ids=[2],
+ n_chills=1, chill_effusivity=12345.0, role_by_solid={2: "chill"})
+
+ monkeypatch.setattr(pc_mod, "get_or_mesh",
+ _fake_get_or_mesh_hit_with_str_keys(calls, mini_msh_file, {"2": 42.0}))
+ monkeypatch.setattr(pc_mod, "mesh_step", _fake_mesh_step(calls, mini_msh_file))
+ monkeypatch.setattr(pc_mod, "validate_bodies", _fake_validate_bodies(calls, plan))
+ monkeypatch.setattr(pc_mod, "tag_contact_patches", _fake_tag_contact_patches(calls))
+ monkeypatch.setattr(pc_mod, "load_hybrid",
+ lambda *a, **kw: (None, None, {}, {"film": True}, None, None))
+ monkeypatch.setattr(pc_mod, "predict_field",
+ lambda mesh, *a, **kw: (np.arange(5, dtype=np.float64) + 1.0,
+ {"X": np.zeros((5, 9))}))
+ captured = {}
+ _spy_screen_casting(monkeypatch, captured)
+
+ bodies = [{"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "chill", "params": {"material": "cast_iron"}}]
+
+ with caplog.at_level(logging.WARNING):
+ pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path),
+ ood_json="does-not-exist.json", bodies=bodies)
+
+ assert captured["feed_nodes"].size == 0
+ assert any("legacy" in rec.message.lower() for rec in caplog.records)
+
+
+def test_nodes_by_role_from_tagging_incomplete_role_by_solid_warns_and_drops(caplog):
+ """Review round 1, M1: a solid_id present in `nodes_by_solid` but absent
+ from `role_by_solid` must be dropped LOUDLY (not silently) -- defense in
+ depth for any caller other than `_role_aware_feed_nodes` (which, in the
+ normal `predict_casting` path, never lets an incomplete mapping reach
+ this function at all -- see its own docstring)."""
+ with caplog.at_level(logging.WARNING):
+ grouped = pc_mod._nodes_by_role_from_tagging(
+ {2: [1, 2, 3], 5: [7, 8]}, {2: "insulating_sleeve"})
+
+ assert grouped == {"insulating_sleeve": [1, 2, 3]}
+ assert any("solid_id=5" in rec.message for rec in caplog.records)
+
+
+def test_feed_nodes_from_role_patches_mixed_mapping_bypassing_validate_bodies():
+ """A directly-constructed mixed role -> nodes mapping (spec §3) --
+ `validate_bodies` itself forbids combining chills with a sleeve, but
+ `_feed_nodes_from_role_patches` takes no `BcPlan`/`bc_config` at all, so
+ it must still exclude the chill nodes and keep only the sleeve nodes
+ even when both roles co-occur in the same mapping."""
+ nodes_by_role = {"chill": [10, 11, 12], "insulating_sleeve": [11, 20],
+ "exo_sleeve": [30]}
+
+ feed = pc_mod._feed_nodes_from_role_patches(nodes_by_role)
+
+ assert feed.tolist() == [11, 20, 30]
+ assert 10 not in feed.tolist() and 12 not in feed.tolist()
+
+
+def test_nodes_by_role_from_tagging_bare_mesh_empty_mapping_is_empty():
+ """Bare mesh (spec §3): no role bodies at all -> an empty `nodes_by_solid`
+ (or `role_by_solid`) grouping produces an empty per-role mapping, and
+ `_feed_nodes_from_role_patches` on that empty mapping stays empty --
+ matching `predict_casting`'s `if plan is not None and plan.role_solid_
+ ids:` guard, which never even enters this code path for a bare/bodies=
+ None input."""
+ assert pc_mod._nodes_by_role_from_tagging({}, {}) == {}
+ assert pc_mod._feed_nodes_from_role_patches({}).size == 0
+
+
+def test_nodes_by_role_from_tagging_normalizes_string_solid_id_keys():
+ """A mesh-cache HIT round-trips `nodes_by_solid` through JSON, turning
+ int solid_id keys into str -- `_nodes_by_role_from_tagging` must still
+ match them against `role_by_solid`'s int keys."""
+ grouped = pc_mod._nodes_by_role_from_tagging(
+ {"2": [1, 2, 3]}, {2: "insulating_sleeve"})
+ assert grouped == {"insulating_sleeve": [1, 2, 3]}
+
+
+def _fake_get_or_mesh_hit_with_str_keys(calls, mini_msh_file, per_body_patch_area_mm2):
+ """Simulates a real cache HIT: `per_body_patch_area_mm2` comes back with
+ STRING solid_id keys (the JSON-round-trip artifact -- see
+ tests/test_mesh_cache.py), and `post_mesh` is never called (a HIT never
+ re-tags)."""
+ 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)
+ return dict(n_solids=1, volume_mm3=1.0, bbox_diag_mm=17.3, n_nodes=5,
+ ok=True, healed=False, notes="", cache="hit",
+ per_body_patch_area_mm2=dict(per_body_patch_area_mm2), n_adiab_tris=3)
+ return _fn
+
+
+def test_bodies_cache_hit_normalizes_string_keys_to_int(
+ monkeypatch, mini_msh_file, tmp_path):
+ """Review fix (LOW): a cache HIT's `per_body_patch_area_mm2` comes back
+ with STRING dict keys (json.dump has no int-key mode) -- predict_casting
+ must normalize these back to int in its result dict, and must NOT call
+ tag_contact_patches again on a HIT."""
+ from src.assembly import BcPlan
+
+ calls = {}
+ plan = BcPlan(bc_config="chills", casting_solid_id=1, role_solid_ids=[2],
+ n_chills=1, chill_effusivity=12345.0)
+
+ monkeypatch.setattr(pc_mod, "get_or_mesh",
+ _fake_get_or_mesh_hit_with_str_keys(calls, mini_msh_file, {"2": 99.5}))
+ monkeypatch.setattr(pc_mod, "mesh_step", _fake_mesh_step(calls, mini_msh_file))
+ monkeypatch.setattr(pc_mod, "validate_bodies", _fake_validate_bodies(calls, plan))
+ monkeypatch.setattr(pc_mod, "tag_contact_patches", _fake_tag_contact_patches(calls))
+ monkeypatch.setattr(pc_mod, "load_hybrid",
+ lambda *a, **kw: (None, None, {}, {"film": True}, None, None))
+ monkeypatch.setattr(pc_mod, "predict_field",
+ lambda mesh, *a, **kw: (np.arange(5, dtype=np.float64) + 1.0,
+ {"X": np.zeros((5, 9))}))
+
+ bodies = [{"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "chill", "params": {"material": "cast_iron"}}]
+
+ out = pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path),
+ ood_json="does-not-exist.json", bodies=bodies)
+
+ assert "tag_contact_patches" not in calls # HIT never re-tags
+ assert out["cache"] == "hit"
+ assert out["per_body_patch_area_mm2"] == {2: 99.5}
+ assert all(isinstance(k, int) for k in out["per_body_patch_area_mm2"])
+
+
+def test_bodies_none_never_calls_tag_contact_patches_and_default_bc_config(
+ 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))
+ monkeypatch.setattr(pc_mod, "tag_contact_patches", _fake_tag_contact_patches(calls))
+ _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 "tag_contact_patches" not in calls
+ assert calls["get_or_mesh"]["keep_solid_id"] is None
+ assert calls["get_or_mesh"]["bodies_key"] == "none"
+ assert out["bc_config"] == "none"
+ assert out["per_body_patch_area_mm2"] == {}
+
+
+def test_bodies_as_json_path_loads(monkeypatch, mini_msh_file, tmp_path):
+ from src.assembly import BcPlan
+
+ calls = {}
+ plan = BcPlan(bc_config="none", casting_solid_id=1, role_solid_ids=[])
+
+ 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))
+ monkeypatch.setattr(pc_mod, "validate_bodies", _fake_validate_bodies(calls, plan))
+ # bodies requires a film-capable checkpoint (see the guard test below) --
+ # this test is only about JSON-path loading, so stub a film-capable cfg.
+ monkeypatch.setattr(pc_mod, "load_hybrid",
+ lambda *a, **kw: (None, None, {}, {"film": True}, None, None))
+ monkeypatch.setattr(pc_mod, "predict_field",
+ lambda mesh, *a, **kw: (np.arange(5, dtype=np.float64) + 1.0,
+ {"X": np.zeros((5, 9))}))
+
+ bodies = [{"solid_id": 1, "role": "casting"}]
+ bodies_path = tmp_path / "bodies.json"
+ bodies_path.write_text(json.dumps(bodies))
+
+ out = pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path),
+ ood_json="does-not-exist.json", bodies=str(bodies_path))
+
+ assert calls["validate_bodies_arg"] == bodies
+ assert calls["get_or_mesh"]["keep_solid_id"] == 1
+
+
+def test_bodies_with_raw_mesh_mode_raises(monkeypatch, mini_msh_file, tmp_path):
+ """`bodies` requires the prep/get_or_mesh pipeline (role-based casting
+ solid selection + tagging) -- combining it with mesh_mode='raw' (or an
+ explicit elem_size, which forces raw) must fail loudly rather than
+ silently ignore the role assignment."""
+ from src.assembly import BcPlan
+
+ calls = {}
+ plan = BcPlan(bc_config="none", casting_solid_id=1, role_solid_ids=[])
+ 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))
+ monkeypatch.setattr(pc_mod, "validate_bodies", _fake_validate_bodies(calls, plan))
+ _stub_prediction_stage(monkeypatch, n_nodes=5)
+
+ bodies = [{"solid_id": 1, "role": "casting"}]
+ with pytest.raises(ValueError):
+ pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path),
+ ood_json="does-not-exist.json", bodies=bodies,
+ mesh_mode="raw")
+
+
+def test_bodies_with_non_film_checkpoint_raises_and_never_tags(
+ monkeypatch, mini_msh_file, tmp_path):
+ """Review fix (MEDIUM): `bodies` with chill/sleeve roles + a NON-film
+ checkpoint (e.g. SP1 via --gnn) must fail loudly, not silently proceed.
+ A resolved BcPlan means tag_contact_patches would relabel the contact
+ patch "Adiabatic" -- for a chill that's backwards physics (a chill
+ COOLS, Adiabatic INSULATES) -- and SP1 has no BC channel to honor the
+ role assignment at all. The check runs BEFORE meshing (load_hybrid is
+ called first), so tag_contact_patches must never even be invoked."""
+ from src.assembly import BcPlan
+
+ calls = {}
+ plan = BcPlan(bc_config="chills", casting_solid_id=1, role_solid_ids=[2],
+ n_chills=1, chill_effusivity=12345.0)
+
+ 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))
+ monkeypatch.setattr(pc_mod, "validate_bodies", _fake_validate_bodies(calls, plan))
+ monkeypatch.setattr(pc_mod, "tag_contact_patches", _fake_tag_contact_patches(calls))
+ # A non-film (SP1-shaped) cfg -- no "film" key at all, matching a real
+ # SP1 checkpoint's cfg.
+ monkeypatch.setattr(pc_mod, "load_hybrid",
+ lambda *a, **kw: (None, None, {}, {"hidden": 16}, None, None))
+
+ bodies = [{"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "chill", "params": {"material": "cast_iron"}}]
+
+ with pytest.raises(ValueError, match="film-conditioned"):
+ pc_mod.predict_casting("fake_part.step", out_dir=str(tmp_path),
+ ood_json="does-not-exist.json", bodies=bodies)
+
+ assert "tag_contact_patches" not in calls
+ assert "get_or_mesh" not in calls
+
+
+def test_bodies_cache_key_order_invariant_and_role_sensitive():
+ from src.assembly import BcPlan
+
+ plan = BcPlan(bc_config="chills", casting_solid_id=1, role_solid_ids=[2, 3],
+ n_chills=2, chill_effusivity=12345.0)
+ bodies = [{"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "chill"},
+ {"solid_id": 3, "role": "chill"}]
+ bodies_reordered = [bodies[2], bodies[0], bodies[1]]
+
+ k1 = pc_mod._bodies_cache_key(bodies, plan)
+ k2 = pc_mod._bodies_cache_key(bodies_reordered, plan)
+ assert k1 == k2
+
+ plan2 = BcPlan(bc_config="chills", casting_solid_id=1, role_solid_ids=[3],
+ n_chills=1, chill_effusivity=12345.0)
+ bodies_role_changed = [{"solid_id": 1, "role": "casting"},
+ {"solid_id": 2, "role": "ignore"},
+ {"solid_id": 3, "role": "chill"}]
+ k3 = pc_mod._bodies_cache_key(bodies_role_changed, plan2)
+ assert k3 != k1
+
+
+# --------------------------------------------------------------------------
+# 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` -- which in turn is refused when `bodies` is present.
+# Density and per-body roles were therefore mutually exclusive from outside.
+# --------------------------------------------------------------------------
+
+
+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. Same reasoning as the
+ existing `bodies`-with-`raw` guard, and the same loud failure.
+ """
+ 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_prep_parts.py b/tests/test_prep_parts.py
new file mode 100644
index 0000000000000000000000000000000000000000..123a4b50659d118c22afe6acb89d3544b8bcbc71
--- /dev/null
+++ b/tests/test_prep_parts.py
@@ -0,0 +1,622 @@
+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,
+)
+from src.msh_parse import parse_msh
+
+
+def _make_one_box_step(path, dx=30.0, dy=20.0, dz=10.0):
+ import gmsh
+ gmsh.initialize()
+ try:
+ gmsh.option.setNumber("General.Terminal", 0)
+ gmsh.model.add("box")
+ gmsh.model.occ.addBox(0, 0, 0, dx, dy, dz)
+ gmsh.model.occ.synchronize()
+ gmsh.write(str(path))
+ finally:
+ gmsh.finalize()
+
+
+def _make_two_box_step(path, big=40.0, small=10.0):
+ """A synthetic 2-solid STEP: one `big`^3 box and one `small`^3 box,
+ positioned so they don't touch (kept as two separate OCC solids)."""
+ import gmsh
+ gmsh.initialize()
+ try:
+ gmsh.option.setNumber("General.Terminal", 0)
+ gmsh.model.add("two_box")
+ gmsh.model.occ.addBox(0, 0, 0, big, big, big)
+ gmsh.model.occ.addBox(big + 100, 0, 0, small, small, small)
+ gmsh.model.occ.synchronize()
+ gmsh.write(str(path))
+ finally:
+ 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).
+# --------------------------------------------------------------------------
+def test_elem_size_from_area_matches_formula():
+ area_mm2, target_nodes = 2200.0, 500
+ expected = np.sqrt(K_SURF * area_mm2 / target_nodes)
+ assert _elem_size_from_area(area_mm2, target_nodes) == pytest.approx(expected)
+
+
+def test_heal_and_mesh_area_seed_converges_on_first_attempt(tmp_path):
+ """The calibrated K_SURF must make the 30x20x10mm box @ target_nodes=500
+ converge within the +-35% band on the FIRST density-search attempt (no
+ refinement rounds needed) -- this is the whole point of seeding from
+ the OCC surface area instead of bbox_diag/28, which missed this same
+ fixture badly (see the multi-round test below, forced via a bad
+ monkeypatched K_SURF)."""
+ 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)
+
+ assert result["ok"] is True
+ assert result["attempts"] == 1
+ assert 500 * 0.65 <= result["n_nodes"] <= 500 * 1.35
+
+
+def test_first_guess_elem_size_falls_back_when_area_query_raises(monkeypatch):
+ """If the OCC surface-area query raises, `_first_guess_elem_size` must
+ fall back to the bbox_diag/28 heuristic rather than propagating the
+ exception."""
+ def _boom():
+ raise RuntimeError("simulated gmsh area-query failure")
+
+ monkeypatch.setattr(pp, "_total_surface_area_mm2", _boom)
+
+ bbox_diag_mm = 37.4
+ assert _first_guess_elem_size(bbox_diag_mm, 500) == pytest.approx(_elem_size(bbox_diag_mm))
+
+
+def test_first_guess_elem_size_falls_back_when_area_nonpositive(monkeypatch):
+ """If the OCC surface-area query returns <= 0 (no exception, just a
+ degenerate/empty answer), `_first_guess_elem_size` must fall back to
+ the bbox_diag/28 heuristic rather than dividing by (or sqrt-ing) a
+ non-positive area."""
+ monkeypatch.setattr(pp, "_total_surface_area_mm2", lambda: 0.0)
+
+ bbox_diag_mm = 37.4
+ assert _first_guess_elem_size(bbox_diag_mm, 500) == pytest.approx(_elem_size(bbox_diag_mm))
+
+
+# --------------------------------------------------------------------------
+# Heal-only-as-fallback regression tests (guards the bug fixed in 5220dbc:
+# `healShapes()` can run without raising and still silently destroy the
+# solid -- monkeypatched, deterministic, no real-part meshing).
+# --------------------------------------------------------------------------
+def test_heal_and_mesh_clean_part_never_calls_heal(tmp_path, monkeypatch):
+ """A part that meshes fine on pass 1 must never invoke `_apply_heal` --
+ this is the exact regression found in commit 5220dbc, where
+ `healShapes()` was called unconditionally and silently reduced a
+ perfectly meshable solid to 0 volume entities."""
+ step = tmp_path / "box.step"
+ _make_one_box_step(step, 30, 20, 10)
+
+ def _boom():
+ raise AssertionError("_apply_heal must not be called when pass 1 succeeds")
+
+ monkeypatch.setattr(pp, "_apply_heal", _boom)
+
+ 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["healed"] is False
+
+
+def test_heal_and_mesh_ineffective_heal_reported_as_failed(tmp_path, monkeypatch):
+ """If pass 1's mesh attempt fails and the pass-2 heal runs but does NOT
+ yield a volume (the historical bug: `healShapes()` silently zeroing out
+ the solid), `heal_and_mesh` must (a) actually attempt the heal, and
+ (b) report ok=False / healed=False -- never silently accept a heal
+ that produced no usable geometry."""
+ step = tmp_path / "box.step"
+ _make_one_box_step(step, 30, 20, 10)
+
+ monkeypatch.setattr(pp, "_tag_and_mesh",
+ lambda size, out_msh: (False, 0, "simulated mesh failure"))
+
+ heal_calls = []
+
+ def _fake_heal():
+ heal_calls.append(1)
+ return False, "simulated heal ineffective (0 volumes)"
+
+ monkeypatch.setattr(pp, "_apply_heal", _fake_heal)
+
+ out_msh = tmp_path / "box.msh"
+ result = heal_and_mesh(str(step), str(out_msh), target_nodes=500)
+
+ assert len(heal_calls) == 1 # (a) heal WAS attempted after pass-1 failure
+ assert result["ok"] is False # (b) not silently accepted
+ assert result["healed"] is False
+ assert "simulated heal ineffective" in result["notes"]
+ assert not out_msh.exists()
+
+
+def test_heal_and_mesh_effective_heal_then_succeeds(tmp_path, monkeypatch):
+ """If pass 1 fails but the pass-2 heal IS effective, pass 2 should
+ retry meshing and succeed (ok=True, healed=True)."""
+ step = tmp_path / "box.step"
+ _make_one_box_step(step, 30, 20, 10)
+
+ calls = {"n": 0}
+
+ def _fake_tag_and_mesh(size, out_msh):
+ calls["n"] += 1
+ if calls["n"] == 1:
+ return False, 0, "simulated pass-1 failure"
+ Path(out_msh).write_text("fake mesh")
+ return True, 777, ""
+
+ monkeypatch.setattr(pp, "_tag_and_mesh", _fake_tag_and_mesh)
+ monkeypatch.setattr(pp, "_apply_heal", lambda: (True, ""))
+
+ out_msh = tmp_path / "box.msh"
+ result = heal_and_mesh(str(step), str(out_msh), target_nodes=777)
+
+ assert calls["n"] >= 2
+ assert result["ok"] is True
+ assert result["healed"] is True
+ assert result["n_nodes"] == 777
+
+
+# --------------------------------------------------------------------------
+# 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
+# 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
+# tag numbers (`_VOLUME_PHYS_TAG`/`_COOLING_PHYS_TAG`) across every call in
+# a session, plus a loud-fail guard in case the binding fails for any other
+# reason.
+# --------------------------------------------------------------------------
+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
+ 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
+ physical names actually bind). This test's purpose is that TAG
+ behavior, not seed quality -- since the calibrated K_SURF now makes
+ 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
+ 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."""
+ step = tmp_path / "box.step"
+ _make_one_box_step(step, 30, 20, 10)
+
+ monkeypatch.setattr(pp, "K_SURF", 50.0)
+
+ calls = []
+ orig_tag_and_mesh = pp._tag_and_mesh
+
+ def _counting(size, out_msh):
+ calls.append(size)
+ return orig_tag_and_mesh(size, out_msh)
+
+ monkeypatch.setattr(pp, "_tag_and_mesh", _counting)
+
+ out_msh = tmp_path / "box.msh"
+ result = heal_and_mesh(str(step), str(out_msh), target_nodes=500)
+
+ assert len(calls) >= 2, "test setup must force a multi-round density search"
+ assert result["ok"] is True
+
+ mesh = parse_msh(str(out_msh))
+ assert "Cooling" in mesh.tris
+ assert mesh.tris["Cooling"].shape[0] > 0
+ 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():
+ """A plain fresh gmsh session (nothing fast-related called) must show
+ gmsh's own, un-touched default for Mesh.Algorithm3D -- never HXT (10).
+ Read BEFORE any fast call so this session can't have been polluted by
+ another test in the same process."""
+ import gmsh
+
+ gmsh.initialize()
+ try:
+ gmsh.option.setNumber("General.Terminal", 0)
+ algo = gmsh.option.getNumber("Mesh.Algorithm3D")
+ assert algo != 10
+ assert algo == 1 # gmsh's documented default (classical Delaunay)
+ finally:
+ 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).
+# --------------------------------------------------------------------------
+def test_heal_and_mesh_defeature_true_synthetic_box(tmp_path):
+ """defeature=True on a clean box must still succeed with a nonempty
+ "Cooling" surface group and tet volume -- defeaturing a clean part must
+ not break the Volume/Cooling tag guarantee."""
+ 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, defeature=True)
+
+ assert result["ok"] is True
+ assert result["n_nodes"] > 0
+
+ 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_heal_and_mesh_defeature_true_calls_apply_defeature(tmp_path, monkeypatch):
+ """defeature=True must actually invoke `_apply_defeature` (in pass 1)."""
+ step = tmp_path / "box.step"
+ _make_one_box_step(step, 30, 20, 10)
+
+ calls = []
+ orig = pp._apply_defeature
+
+ def _wrap():
+ calls.append(1)
+ return orig()
+
+ monkeypatch.setattr(pp, "_apply_defeature", _wrap)
+
+ out_msh = tmp_path / "box.msh"
+ result = heal_and_mesh(str(step), str(out_msh), target_nodes=500, defeature=True)
+
+ assert len(calls) >= 1
+ assert result["ok"] is True
+
+
+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`
+ 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
+ model, which raises an uncaught `Exception: Empty bounding box`."""
+ step = tmp_path / "box.step"
+ _make_one_box_step(step, 30, 20, 10)
+
+ monkeypatch.setattr(pp, "_apply_defeature",
+ lambda: (False, "defeature (removeAllDuplicates) yielded 0 volume entities"))
+
+ out_msh = tmp_path / "box.msh"
+ result = heal_and_mesh(str(step), str(out_msh), target_nodes=500, defeature=True)
+
+ assert result["ok"] is False
+ assert "0 volume entities" in result["notes"]
+ assert not out_msh.exists()
+
+
+def test_heal_and_mesh_defeature_default_never_calls_apply_defeature(tmp_path, monkeypatch):
+ """`defeature` defaults to False -- `heal_and_mesh` must not invoke
+ `_apply_defeature` 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_defeature must not be called when defeature=False (default)")
+
+ monkeypatch.setattr(pp, "_apply_defeature", _boom)
+
+ out_msh = tmp_path / "box.msh"
+ result = heal_and_mesh(str(step), str(out_msh), target_nodes=500)
+ assert result["ok"] is True
+
+
+def test_apply_defeature_removes_duplicates_without_losing_solid():
+ """Direct test of `_apply_defeature()` inside an initialized gmsh
+ session: runs cleanly on a plain box and does not destroy the solid."""
+ 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()
+ ok, note = pp._apply_defeature()
+ assert ok is True
+ assert note == ""
+ assert len(gmsh.model.getEntities(3)) == 1
+ finally:
+ 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`
+ must return ok=False and must NOT write `out_msh` -- never silently
+ hand back an untagged mesh."""
+ import gmsh
+
+ step = tmp_path / "box.step"
+ _make_one_box_step(step, 30, 20, 10)
+
+ gmsh.initialize()
+ try:
+ gmsh.option.setNumber("General.Terminal", 0)
+ gmsh.model.add("box")
+ gmsh.model.occ.importShapes(str(step))
+ gmsh.model.occ.synchronize()
+
+ monkeypatch.setattr(gmsh.model, "getPhysicalName", lambda dim, tag: "")
+
+ out_msh = tmp_path / "box.msh"
+ ok, n_nodes, note = pp._tag_and_mesh(2.0, str(out_msh))
+
+ assert ok is False
+ assert n_nodes == 0
+ assert "name binding missing" in note
+ assert not out_msh.exists()
+ finally:
+ gmsh.finalize()
diff --git a/tests/test_release_io.py b/tests/test_release_io.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f3e5c820dc8e516561ae31cd3541c86e8cf5057
--- /dev/null
+++ b/tests/test_release_io.py
@@ -0,0 +1,82 @@
+"""Public-release artifact formats round-trip exactly (src/release_io.py)."""
+import os
+
+import numpy as np
+import pytest
+import torch
+
+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 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_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)
+ with pytest.raises(ValueError, match="not a"):
+ rio.load_gnn_state(out)
+
+
+def test_gbm_skops_roundtrip(tmp_path):
+ pytest.importorskip("skops")
+ 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))
+
+
+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
diff --git a/tests/test_shrinkage.py b/tests/test_shrinkage.py
new file mode 100644
index 0000000000000000000000000000000000000000..6adb90440a0c19d47da2f1d1d3caf4fbd87941c9
--- /dev/null
+++ b/tests/test_shrinkage.py
@@ -0,0 +1,1305 @@
+import numpy as np
+import pytest
+
+from src.shrinkage import (
+ node_gradient, niyama_exact, niyama_proxy, porosity_risk, NY_CRIT,
+ 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
+# 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
+# between adjacent cubes cancel out of the boundary-surface-area count),
+# reused across the region/sink/feeder tests below.
+# ---------------------------------------------------------------------------
+
+_KUHN_TETS = [(0, 1, 3, 7), (0, 1, 5, 7), (0, 4, 5, 7),
+ (0, 2, 3, 7), (0, 2, 6, 7), (0, 4, 6, 7)]
+
+
+def _node_idx(nx, ny):
+ return lambda i, j, k: i + j * nx + k * nx * ny
+
+
+def _grid_mesh(nx, ny, nz, spacing=1.0):
+ idx = _node_idx(nx, ny)
+ pos = np.zeros((nx * ny * nz, 3))
+ for k in range(nz):
+ for j in range(ny):
+ for i in range(nx):
+ pos[idx(i, j, k)] = [i * spacing, j * spacing, k * spacing]
+
+ def local(i, j, k, bit):
+ dx, dy, dz = bit & 1, (bit >> 1) & 1, (bit >> 2) & 1
+ return idx(i + dx, j + dy, k + dz)
+
+ tets = []
+ for k in range(nz - 1):
+ for j in range(ny - 1):
+ for i in range(nx - 1):
+ for t in _KUHN_TETS:
+ tets.append([local(i, j, k, b) for b in t])
+ return pos, np.array(tets, dtype=np.int64)
+
+
+def _surface_nodes(nx, ny, nz):
+ idx = _node_idx(nx, ny)
+ nodes = []
+ for k in range(nz):
+ for j in range(ny):
+ for i in range(nx):
+ if i in (0, nx - 1) or j in (0, ny - 1) or k in (0, nz - 1):
+ nodes.append(idx(i, j, k))
+ return np.array(sorted(nodes), dtype=np.int64)
+
+
+# ---------------------------------------------------------------------------
+# (a) node_gradient
+# ---------------------------------------------------------------------------
+
+def test_node_gradient_single_regular_tet():
+ # unit tet at the origin: (0,0,0), (1,0,0), (0,1,0), (0,0,1)
+ pos = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]], float)
+ tets = np.array([[0, 1, 2, 3]])
+ field = 2.0 * pos[:, 0] # f = 2x -> grad = (2, 0, 0) everywhere
+ g = node_gradient(pos, tets, field)
+ assert g.shape == (4, 3)
+ np.testing.assert_allclose(g, np.tile([2.0, 0.0, 0.0], (4, 1)), atol=1e-10)
+
+
+def test_node_gradient_two_tet_bipyramid():
+ # two tets sharing triangular face ABC, apexes D (above) and E (below)
+ # A, B, C, D, E -> indices 0..4
+ pos = np.array([
+ [0, 0, 0], # A
+ [1, 0, 0], # B
+ [0, 1, 0], # C
+ [0, 0, 1], # D (apex above)
+ [0, 0, -1], # E (apex below)
+ ], float)
+ tets = np.array([
+ [0, 1, 2, 3], # A B C D
+ [0, 1, 2, 4], # A B C E
+ ])
+ field = pos[:, 0] + pos[:, 1] # f = x + y -> grad = (1, 1, 0) everywhere (exact, linear)
+ g = node_gradient(pos, tets, field)
+ assert g.shape == (5, 3)
+ np.testing.assert_allclose(g, np.tile([1.0, 1.0, 0.0], (5, 1)), atol=1e-10)
+
+
+def test_node_gradient_no_tets_returns_zeros():
+ pos = np.zeros((3, 3))
+ tets = np.empty((0, 4), dtype=np.int64)
+ field = np.zeros(3)
+ g = node_gradient(pos, tets, field)
+ np.testing.assert_allclose(g, np.zeros((3, 3)))
+
+
+# ---------------------------------------------------------------------------
+# (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])
+
+
+# ---------------------------------------------------------------------------
+# (c) niyama_proxy
+# ---------------------------------------------------------------------------
+
+def test_niyama_proxy_linear_slab_hand_computed():
+ # single regular tet, t_s = 2 + 2x (linear ramp along x) -> |grad t_s| = 2 exactly
+ pos = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]], float)
+ tets = np.array([[0, 1, 2, 3]])
+ t_s = np.array([2.0, 4.0, 2.0, 2.0]) # node0=2 (x=0), node1=4 (x=1), node2/3 at x=0 -> 2
+ TL, TS = 1400.0, 1100.0 # TL - TS = 300
+
+ ny = niyama_proxy(pos, tets, t_s, TL, TS)
+ assert ny.shape == (4,)
+
+ # hand computation: sqrt((TL-TS)/t_s) * |grad t_s|, |grad t_s| = 2 at every node
+ expected_node0 = np.sqrt(300.0 / 2.0) * 2.0 # ~24.494897...
+ expected_node1 = np.sqrt(300.0 / 4.0) * 2.0 # ~17.320508...
+ np.testing.assert_allclose(ny[0], expected_node0, rtol=1e-10)
+ np.testing.assert_allclose(ny[1], expected_node1, rtol=1e-10)
+
+
+def test_niyama_proxy_eps_guarded_at_zero_t_s():
+ pos = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]], float)
+ tets = np.array([[0, 1, 2, 3]])
+ t_s = np.array([0.0, 2.0, 0.0, 0.0])
+ ny = niyama_proxy(pos, tets, t_s, 1400.0, 1100.0)
+ assert np.isfinite(ny).all()
+
+
+# ---------------------------------------------------------------------------
+# (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,
+ }
+
+
+# ---------------------------------------------------------------------------
+# (e) porosity_risk -- percentile mode
+# ---------------------------------------------------------------------------
+
+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)
+ 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)
+ assert mask.sum() == 1
+ np.testing.assert_array_equal(np.where(mask)[0], [0])
+
+
+def test_porosity_risk_percentile_tie_handling_documented():
+ # duplicate values straddling the percentile threshold: both ties are
+ # flagged (mask is defined as ny <= percentile(ny, pct), so a tie AT the
+ # threshold is always included) -- this can flag slightly MORE than the
+ # 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)
+ 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")
+
+
+# ---------------------------------------------------------------------------
+# (a) hot_regions + is_fed -- isolated boss (unfed) vs. surface-connected
+# band (fed), on one 7x7x7 mesh with two disjoint elevated node clusters.
+# ---------------------------------------------------------------------------
+
+def _boss_and_band_mesh():
+ nx = ny = nz = 7
+ pos, tets = _grid_mesh(nx, ny, nz)
+ idx = _node_idx(nx, ny)
+ n_nodes = pos.shape[0]
+ t_s = np.zeros(n_nodes)
+
+ # boss: a fully-interior unit cube {2,3}x{2,3}x{2,3} -- 8 nodes, none on
+ # the (0 or 6) domain boundary, and not adjacent to the band below.
+ boss_nodes = [idx(i, j, k) for i in (2, 3) for j in (2, 3) for k in (2, 3)]
+ # band: a column at x=5,y=5, z=3..6 -- z=6 sits on the top boundary face,
+ # so the band directly touches a free surface.
+ band_nodes = [idx(5, 5, z) for z in (3, 4, 5, 6)]
+
+ t_s[boss_nodes] = 100.0
+ t_s[band_nodes] = 100.0
+ n_hot = len(boss_nodes) + len(band_nodes)
+ top_frac = n_hot / n_nodes # sized so the percentile cut lands exactly on {0, 100}
+
+ surface_nodes = _surface_nodes(nx, ny, nz)
+ return pos, tets, t_s, top_frac, surface_nodes, set(boss_nodes), set(band_nodes)
+
+
+def test_hot_regions_finds_boss_and_band_as_separate_components():
+ pos, tets, t_s, top_frac, surface_nodes, boss_nodes, band_nodes = _boss_and_band_mesh()
+ regions = hot_regions(pos, tets, t_s, top_frac=top_frac)
+ assert len(regions) == 2
+ node_sets = [set(r["nodes"].tolist()) for r in regions]
+ assert boss_nodes in node_sets
+ assert band_nodes in node_sets
+ for r in regions:
+ assert r["t_max"] == 100.0
+ np.testing.assert_allclose(r["centroid"], pos[np.array(sorted(r["nodes"]))].mean(axis=0))
+
+
+def test_hot_regions_sorted_by_t_max_desc():
+ pos, tets, t_s, top_frac, surface_nodes, boss_nodes, band_nodes = _boss_and_band_mesh()
+ # bump one region's peak so the sort order is unambiguous.
+ idx = _node_idx(7, 7)
+ t_s = t_s.copy()
+ t_s[idx(5, 5, 6)] = 150.0
+ regions = hot_regions(pos, tets, t_s, top_frac=top_frac)
+ t_maxes = [r["t_max"] for r in regions]
+ assert t_maxes == sorted(t_maxes, reverse=True)
+ assert t_maxes[0] == 150.0
+
+
+def test_is_fed_isolated_boss_is_false():
+ # feed_nodes = surface_nodes here (this mesh's domain boundary stands in
+ # for a riser/feed location) -- the boss never reaches the corridor
+ # threshold (0.95*t_max) outside its own 8 nodes, so it can't connect.
+ pos, tets, t_s, top_frac, surface_nodes, boss_nodes, band_nodes = _boss_and_band_mesh()
+ regions = hot_regions(pos, tets, t_s, top_frac=top_frac)
+ boss_region = next(r for r in regions if set(r["nodes"].tolist()) == boss_nodes)
+ assert is_fed(boss_region, t_s, tets, feed_nodes=surface_nodes) is False
+
+
+def test_is_fed_surface_connected_band_is_true():
+ # the band's own top node (z=6) is itself in feed_nodes=surface_nodes,
+ # so the region trivially self-connects to a feed node.
+ pos, tets, t_s, top_frac, surface_nodes, boss_nodes, band_nodes = _boss_and_band_mesh()
+ regions = hot_regions(pos, tets, t_s, top_frac=top_frac)
+ band_region = next(r for r in regions if set(r["nodes"].tolist()) == band_nodes)
+ assert is_fed(band_region, t_s, tets, feed_nodes=surface_nodes) is True
+
+
+def test_is_fed_empty_feed_nodes_is_always_false():
+ pos, tets, t_s, top_frac, surface_nodes, boss_nodes, band_nodes = _boss_and_band_mesh()
+ regions = hot_regions(pos, tets, t_s, top_frac=top_frac)
+ band_region = next(r for r in regions if set(r["nodes"].tolist()) == band_nodes)
+ # bare sand casting convention: no risers -> empty feed_nodes -> unfed,
+ # even for a region that DOES touch the domain surface.
+ assert is_fed(band_region, t_s, tets, feed_nodes=np.array([], dtype=np.int64)) is 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
+# "sub_resolution": True; `feeder_recommendations` skips them (logging the
+# count once); `macro_shrinkage` still lists them (volume/severity remain
+# meaningful, only the modulus diverges).
+# --------------------------------------------------------------------------- #
+def _single_isolated_hot_node_mesh():
+ """A 7x7x7 grid mesh with exactly ONE hot node (fully interior, no hot
+ neighbors) -- the minimal fixture for a 1-node `hot_regions` component."""
+ nx = ny = nz = 7
+ pos, tets = _grid_mesh(nx, ny, nz)
+ idx = _node_idx(nx, ny)
+ n_nodes = pos.shape[0]
+ t_s = np.zeros(n_nodes)
+ hot_node = idx(3, 3, 3)
+ t_s[hot_node] = 100.0
+ top_frac = 1.0 / n_nodes
+ surface_nodes = _surface_nodes(nx, ny, nz)
+ return pos, tets, t_s, top_frac, surface_nodes, hot_node
+
+
+def test_hot_regions_marks_sub_resolution_for_regions_under_4_nodes():
+ pos, tets, t_s, top_frac, surface_nodes, hot_node = _single_isolated_hot_node_mesh()
+ regions = hot_regions(pos, tets, t_s, top_frac=top_frac)
+ assert len(regions) == 1
+ assert regions[0]["nodes"].tolist() == [hot_node]
+ assert regions[0]["sub_resolution"] is True
+
+
+def test_hot_regions_does_not_mark_well_resolved_regions_sub_resolution():
+ # the boss (8 nodes) and band (4 nodes) fixture -- both >= the 4-node floor.
+ pos, tets, t_s, top_frac, surface_nodes, boss_nodes, band_nodes = _boss_and_band_mesh()
+ regions = hot_regions(pos, tets, t_s, top_frac=top_frac)
+ assert len(regions) == 2
+ for r in regions:
+ assert r["sub_resolution"] is False
+
+
+def test_feeder_recommendations_skips_sub_resolution_region_and_logs_count(caplog):
+ """A 1-node region must get NO feeder recommendation at all, and the
+ skip must be logged (once, with the count) rather than silent."""
+ import logging as _logging
+ pos, tets, t_s, top_frac, surface_nodes, hot_node = _single_isolated_hot_node_mesh()
+ regions = hot_regions(pos, tets, t_s, top_frac=top_frac)
+ assert regions[0]["sub_resolution"] is True
+
+ with caplog.at_level(_logging.INFO, logger="src.shrinkage"):
+ recs = feeder_recommendations(regions, pos, tets, surface_nodes,
+ gravity_vec=(0, 0, -1), t_s=t_s)
+
+ assert recs == []
+ assert any("skipped 1 sub-resolution region" in rec.message for rec in caplog.records)
+
+
+def test_feeder_recommendations_mixed_regions_skips_only_sub_resolution_ones(caplog):
+ """A regions list mixing well-resolved regions (get a feeder rec each)
+ and a sub-resolution one (skipped, not just de-prioritized) -- the
+ well-resolved regions' recs must be unaffected by the mix."""
+ import logging as _logging
+ pos, tets, t_s, top_frac, surface_nodes, boss_nodes, band_nodes = _boss_and_band_mesh()
+ regions = hot_regions(pos, tets, t_s, top_frac=top_frac)
+ assert all(not r["sub_resolution"] for r in regions)
+
+ sub_region = {
+ "nodes": np.array([0]), "centroid": pos[0], "volume_mm3": 1.0,
+ "t_max": float(t_s.max()), "region_id": 99,
+ "hot_threshold": regions[0]["hot_threshold"], "sub_resolution": True,
+ }
+ mixed_regions = regions + [sub_region]
+
+ with caplog.at_level(_logging.INFO, logger="src.shrinkage"):
+ recs = feeder_recommendations(mixed_regions, pos, tets, surface_nodes,
+ gravity_vec=(0, 0, -1), t_s=t_s)
+
+ assert len(recs) == len(regions)
+ assert {r["region_id"] for r in recs} == {r["region_id"] for r in regions}
+ assert any("skipped 1 sub-resolution region" in rec.message for rec in caplog.records)
+
+
+# --------------------------------------------------------------------------- #
+# Review finding I4 (performance): 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
+# of the helper on a hand-countable 2-tet mesh (one shared interior face).
+# --------------------------------------------------------------------------- #
+def test_boundary_faces_excludes_shared_interior_face():
+ # two tets sharing face (0,1,2) -- a triangular bipyramid.
+ tets = np.array([[0, 1, 2, 3], [0, 1, 2, 4]], dtype=np.int64)
+ bf = _boundary_faces(tets)
+ assert bf.shape == (6, 3)
+ boundary_keys = {tuple(sorted(row.tolist())) for row in bf}
+ assert (0, 1, 2) not in boundary_keys # the shared interior face is excluded
+ expected = {(0, 1, 3), (0, 2, 3), (1, 2, 3),
+ (0, 1, 4), (0, 2, 4), (1, 2, 4)}
+ assert boundary_keys == expected
+
+
+def test_boundary_faces_empty_tets_gives_empty_array():
+ bf = _boundary_faces(np.zeros((0, 4), dtype=np.int64))
+ assert bf.shape == (0, 3)
+
+
+def test_hot_regions_top_frac_out_of_range_raises():
+ pos, tets, t_s, top_frac, surface_nodes, boss_nodes, band_nodes = _boss_and_band_mesh()
+ with pytest.raises(ValueError):
+ hot_regions(pos, tets, t_s, top_frac=0.0)
+ with pytest.raises(ValueError):
+ hot_regions(pos, tets, t_s, top_frac=1.5)
+
+
+# ---------------------------------------------------------------------------
+# (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
+# 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
+# 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
+# hops from the region, connected only via the intermediate warm corridor).
+# ---------------------------------------------------------------------------
+
+def _corridor_chain_mesh():
+ """8x2x2 lattice (32 nodes): an x-adjacent chain of nodes (i,0,0),
+ i=0..7 (verified consecutively tet-adjacent). Node 6 will be the sole
+ hot region (t_max=100); nodes 0-5 are the candidate corridor (0 doubles
+ as the feed node); node 7 is the "far cold side" used as a decoy feed
+ node in variant 3. Every other mesh node stays at background 0 and
+ never participates (present only so this is a valid tet mesh)."""
+ nx, ny, nz = 8, 2, 2
+ pos, tets = _grid_mesh(nx, ny, nz)
+ idx = _node_idx(nx, ny)
+ chain = [idx(i, 0, 0) for i in range(nx)]
+ n_nodes = pos.shape[0]
+ return pos, tets, n_nodes, chain
+
+
+def test_is_fed_warm_corridor_bridges_region_to_distant_feed_node():
+ pos, tets, n_nodes, chain = _corridor_chain_mesh()
+ t_s = np.zeros(n_nodes)
+ t_s[chain[6]] = 100.0 # region peak
+ t_s[chain[0:6]] = 96.0 # corridor >= 0.95*100 = 95 -- stays warm
+ region = {"nodes": np.array([chain[6]]), "t_max": 100.0}
+ assert is_fed(region, t_s, tets, feed_nodes=np.array([chain[0]])) is True
+
+
+def test_is_fed_false_when_corridor_cools_below_corridor_frac():
+ pos, tets, n_nodes, chain = _corridor_chain_mesh()
+ t_s = np.zeros(n_nodes)
+ t_s[chain[6]] = 100.0
+ t_s[chain[0:6]] = 90.0 # < 95 -- corridor no longer warm enough
+ region = {"nodes": np.array([chain[6]]), "t_max": 100.0}
+ assert is_fed(region, t_s, tets, feed_nodes=np.array([chain[0]])) is False
+
+
+def test_is_fed_false_when_feed_node_is_on_far_cold_side():
+ pos, tets, n_nodes, chain = _corridor_chain_mesh()
+ t_s = np.zeros(n_nodes)
+ t_s[chain[6]] = 100.0
+ t_s[chain[0:6]] = 96.0 # corridor stays warm...
+ region = {"nodes": np.array([chain[6]]), "t_max": 100.0}
+ # ...but feed_nodes points to the far (cold, disconnected) side instead.
+ assert is_fed(region, t_s, tets, feed_nodes=np.array([chain[7]])) is False
+
+
+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
+ 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
+ connectivity. Guarded to report unfed instead."""
+ pos, tets, n_nodes, chain = _corridor_chain_mesh()
+ t_s = np.zeros(n_nodes) # every node, including all feed candidates, is 0
+ region = {"nodes": np.array([chain[6]]), "t_max": 0.0}
+ # without the guard, corridor threshold = 0.95*0 = 0 <= every t_s -> the
+ # whole mesh (one connected component) would trivially "fed" this region.
+ assert is_fed(region, t_s, tets, feed_nodes=np.array([chain[0]])) is False
+
+
+# ---------------------------------------------------------------------------
+# (b) macro_shrinkage -- severity = pct/100 * volume, hand-checked; grey_iron
+# (pct=0) still reports the region with zero severity; fed regions excluded.
+# ---------------------------------------------------------------------------
+
+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
+ 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
+ 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
+ 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
+ 27/270 = 10% of nodes, matching hot_regions'/macro_shrinkage's default
+ top_frac."""
+ nx, ny, nz = 9, 6, 5
+ pos, tets = _grid_mesh(nx, ny, nz)
+ idx = _node_idx(nx, ny)
+ n_nodes = pos.shape[0]
+ t_s = np.zeros(n_nodes)
+ block_nodes = [idx(i, j, k) for i in (3, 4, 5) for j in (1, 2, 3) for k in (1, 2, 3)]
+ t_s[block_nodes] = 100.0
+ surface_nodes = _surface_nodes(nx, ny, nz)
+ assert len(block_nodes) / n_nodes == pytest.approx(0.10)
+ assert not any(n in set(surface_nodes.tolist()) for n in block_nodes)
+ return pos, tets, t_s, surface_nodes, set(block_nodes)
+
+
+def test_macro_shrinkage_severity_hand_checked():
+ pos, tets, t_s, surface_nodes, block_nodes = _interior_cube_block_mesh()
+ # bare-casting convention: no risers -> empty feed_nodes -> conservative
+ # (this fully-interior block is unfed regardless, since it never reaches
+ # the domain surface anyway -- see is_fed's docstring).
+ no_feed = np.array([], dtype=np.int64)
+ regions = macro_shrinkage(pos, tets, t_s, no_feed, shrinkage_pct=2.5)
+ assert len(regions) == 1
+ r = regions[0]
+ assert set(r["nodes"].tolist()) == block_nodes
+ assert r["volume_mm3"] == pytest.approx(27.0)
+ assert r["severity"] == pytest.approx(2.5 / 100.0 * 27.0)
+
+
+def test_macro_shrinkage_grey_iron_zero_severity_still_reported():
+ pos, tets, t_s, surface_nodes, block_nodes = _interior_cube_block_mesh()
+ no_feed = np.array([], dtype=np.int64)
+ regions = macro_shrinkage(pos, tets, t_s, no_feed,
+ shrinkage_pct=SHRINKAGE_PCT["grey_iron"])
+ assert len(regions) == 1
+ assert regions[0]["severity"] == pytest.approx(0.0)
+ assert regions[0]["volume_mm3"] == pytest.approx(27.0) # region still fully reported
+
+
+def _unfed_boss_and_fed_band_10pct_mesh():
+ """10x5x4 lattice (200 nodes) with an interior unit-cube boss (8 nodes,
+ t_max=100, unfed) and a corner block band (12 nodes, t_max=150, touching
+ x=0/y=0/z=0 boundary faces) -- 20/200 = exactly 10%, matching
+ macro_shrinkage's (un-parametrized) default top_frac so BOTH regions are
+ picked up by its internal hot_regions() call. `feed_nodes=surface_nodes`
+ is used by the tests below as a stand-in "this domain boundary carries
+ risers" convention: the band's own boundary-touching nodes are literally
+ IN `surface_nodes`, so it self-connects trivially (fed); the boss's
+ warm corridor (0.95*its own t_max=100 -> threshold 95) never leaves its
+ own 8 nodes (all neighbors are background 0), so it cannot reach the
+ 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."""
+ nx, ny, nz = 10, 5, 4
+ pos, tets = _grid_mesh(nx, ny, nz)
+ idx = _node_idx(nx, ny)
+ n_nodes = pos.shape[0]
+ t_s = np.zeros(n_nodes)
+
+ boss_nodes = [idx(i, j, k) for i in (4, 5) for j in (1, 2) for k in (1, 2)]
+ band_nodes = [idx(i, j, k) for i in (0, 1) for j in (0, 1) for k in (0, 1, 2)]
+ t_s[boss_nodes] = 100.0
+ t_s[band_nodes] = 150.0
+ assert (len(boss_nodes) + len(band_nodes)) / n_nodes == pytest.approx(0.10)
+
+ surface_nodes = _surface_nodes(nx, ny, nz)
+ surface_set = set(surface_nodes.tolist())
+ assert not any(n in surface_set for n in boss_nodes)
+ assert any(n in surface_set for n in band_nodes)
+ return pos, tets, t_s, surface_nodes, set(boss_nodes), set(band_nodes)
+
+
+def test_macro_shrinkage_excludes_fed_regions():
+ pos, tets, t_s, surface_nodes, boss_nodes, band_nodes = _unfed_boss_and_fed_band_10pct_mesh()
+ regions = macro_shrinkage(pos, tets, t_s, feed_nodes=surface_nodes, shrinkage_pct=2.5)
+ assert len(regions) == 1
+ assert set(regions[0]["nodes"].tolist()) == boss_nodes
+
+
+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)."""
+ 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)
+ assert len(raw_regions) == 2
+ boss_raw = next(r for r in raw_regions if set(r["nodes"].tolist()) == boss_nodes)
+ band_raw = next(r for r in raw_regions if set(r["nodes"].tolist()) == band_nodes)
+ # band's t_max (150) > boss's (100) -> band sorts first.
+ assert band_raw["region_id"] == 0
+ assert boss_raw["region_id"] == 1
+
+ macro = macro_shrinkage(pos, tets, t_s, feed_nodes=surface_nodes, shrinkage_pct=2.5)
+ assert len(macro) == 1 # band filtered out (fed); a naive positional
+ # re-index would mislabel the sole survivor as region_id 0.
+ assert macro[0]["region_id"] == 1
+
+ sinks = top_sinks(macro, pos, surface_nodes, gravity_vec=(0, 0, -1))
+ feeders = feeder_recommendations(macro, pos, tets, surface_nodes, gravity_vec=(0, 0, -1),
+ t_s=t_s)
+ assert all(f["region_id"] == 1 for f in feeders)
+ assert all(s["region_id"] == 1 for s in sinks)
+
+
+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)
+ 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
+ `0.95*t_max` and the hot cut and bridges to a feed node -- i.e. a
+ genuine warm-corridor effect, not the trivial "feed node happens to be
+ inside the region" case exercised by `test_macro_shrinkage_excludes_
+ fed_regions`. Peak (node 1), bridge (node 3), and feed (node 2) are
+ chosen so the feed node is NOT itself tet-adjacent to the region peak
+ (hand-verified: nodes 1 and 2 share no tet in this mesh's Kuhn
+ decomposition) -- the bridge node is structurally required to connect
+ them, unlike a single-hop setup where a direct peak-feed edge could
+ make the test pass for the wrong reason. A single unit-cube mesh (N=8)
+ is used because hot_regions' un-parametrized default top_frac=0.10 only
+ isolates a lone peak node cleanly at small N (verified: at N=12+ the
+ corridor value gets pulled into the "hot" selection itself by
+ percentile interpolation, defeating the "external, not region-touching"
+ setup)."""
+ pos, tets = _grid_mesh(2, 2, 2)
+ n_nodes = pos.shape[0]
+ t_s = np.zeros(n_nodes)
+ t_s[1] = 100.0 # sole region peak
+ t_s[3] = 96.0 # bridge node: >= 0.95*100=95, but < the 97.2 hot cut
+ t_s[2] = 96.0 # feed node: same band; NOT tet-adjacent to node 1 directly
+
+ regions = hot_regions(pos, tets, t_s) # default top_frac=0.10
+ assert len(regions) == 1
+ assert regions[0]["nodes"].tolist() == [1] # bridge/feed excluded from the region itself
+
+ fed_via_corridor = macro_shrinkage(pos, tets, t_s, feed_nodes=np.array([2]),
+ shrinkage_pct=2.5)
+ assert fed_via_corridor == [] # fed -> not reported as macro-shrinkage
+
+ unfed_without_corridor = macro_shrinkage(pos, tets, t_s, feed_nodes=np.array([], dtype=np.int64),
+ shrinkage_pct=2.5)
+ assert len(unfed_without_corridor) == 1
+ assert unfed_without_corridor[0]["nodes"].tolist() == [1]
+
+ # cooling the bridge/feed band below 0.95*t_max flips the verdict back
+ # to unfed even with the SAME feed_nodes -- confirms the flip is driven
+ # by the corridor threshold, not merely feed_nodes' presence.
+ t_s_cold_corridor = t_s.copy()
+ t_s_cold_corridor[[2, 3]] = 90.0
+ cooled = macro_shrinkage(pos, tets, t_s_cold_corridor, feed_nodes=np.array([2]),
+ shrinkage_pct=2.5)
+ assert len(cooled) == 1
+
+
+def test_shrinkage_pct_table():
+ assert SHRINKAGE_PCT == {
+ "steel": 2.5, "low_alloy_steel": 2.5, "ductile_iron": 0.75, "grey_iron": 0.0,
+ }
+
+
+# ---------------------------------------------------------------------------
+# (c) top_sinks -- gravity -z -> sink on the top (z-max) face; rotating
+# gravity to -x/-y relocates the sink to the +x/+y face. Single 3x3x3 mesh,
+# 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():
+ nx = ny = nz = 3
+ pos, tets = _grid_mesh(nx, ny, nz)
+ idx = _node_idx(nx, ny)
+ n_nodes = pos.shape[0]
+ t_s = np.zeros(n_nodes)
+
+ node_top = idx(1, 1, 2) # z=2 -> top (z-max) surface face
+ node_bridge = idx(1, 1, 1) # interior
+ node_xface = idx(2, 1, 1) # x=2 -> +x surface face
+ node_yface = idx(1, 2, 1) # y=2 -> +y surface face
+
+ hot = [node_top, node_bridge, node_xface, node_yface]
+ t_s[hot] = 100.0
+ top_frac = len(hot) / n_nodes
+ surface_nodes = _surface_nodes(nx, ny, nz)
+ regions = hot_regions(pos, tets, t_s, top_frac=top_frac)
+ assert len(regions) == 1
+ return pos, tets, t_s, regions, surface_nodes, node_top, node_xface, node_yface
+
+
+def test_top_sinks_default_gravity_hits_top_face():
+ pos, tets, t_s, regions, surface_nodes, node_top, node_xface, node_yface = _rotatable_sink_mesh()
+ sinks = top_sinks(regions, pos, surface_nodes, gravity_vec=(0, 0, -1))
+ assert len(sinks) == 1
+ np.testing.assert_allclose(sinks[0]["location"], pos[node_top])
+
+
+def test_top_sinks_rotated_gravity_relocates_to_x_face():
+ pos, tets, t_s, regions, surface_nodes, node_top, node_xface, node_yface = _rotatable_sink_mesh()
+ sinks = top_sinks(regions, pos, surface_nodes, gravity_vec=(-1, 0, 0))
+ assert len(sinks) == 1
+ np.testing.assert_allclose(sinks[0]["location"], pos[node_xface])
+
+
+# ---------------------------------------------------------------------------
+# (d) feeder_recommendations -- min_modulus == 1.2*(V/A); location on the
+# 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
+# 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
+# field crosses its OWN selection threshold, which is an exact calculus
+# identity for the (artificial) binary step but measures the anisotropic
+# discrete perimeter of that step, not the true field's level-surface area:
+# it plateaus ~24-29% high on curved/diagonal boundaries (only looking exact
+# on an axis-aligned cut, where a midpoint crossing happens to coincide with
+# the true one). `feeder_recommendations` gained a required `t_s` parameter
+# in this round since the marching-tetrahedra computation needs the
+# continuous field, not just region-node membership.
+# ---------------------------------------------------------------------------
+
+def _reference_region_area_marching_tets(pos, tets, nodes, t_s, threshold):
+ """Independent reimplementation of the marching-tetrahedra + domain-
+ boundary-face area formula (mirrors, but does not call,
+ `src.shrinkage._region_surface_area` -- a genuine hand-check against the
+ same algorithm, not a tautological re-invocation of the code under
+ test)."""
+ n_nodes = pos.shape[0]
+ is_member = np.zeros(n_nodes, dtype=bool)
+ is_member[np.asarray(nodes, dtype=np.int64)] = True
+ hot_mask = t_s >= threshold
+
+ area = 0.0
+ member_per_tet = hot_mask[tets]
+ n_above = member_per_tet.sum(axis=1)
+ for ti in np.flatnonzero((n_above >= 1) & (n_above <= 3)):
+ tet = tets[ti]
+ if not is_member[tet].any():
+ continue
+ verts = pos[tet]
+ vals = t_s[tet]
+ above = member_per_tet[ti]
+ k = int(n_above[ti])
+ if k in (1, 3):
+ odd = int(np.flatnonzero(above)[0]) if k == 1 else int(np.flatnonzero(~above)[0])
+ others = [i for i in range(4) if i != odd]
+ pts = []
+ for o in others:
+ frac = (threshold - vals[odd]) / (vals[o] - vals[odd])
+ pts.append(verts[odd] + frac * (verts[o] - verts[odd]))
+ area += 0.5 * np.linalg.norm(np.cross(pts[1] - pts[0], pts[2] - pts[0]))
+ else:
+ hi = np.flatnonzero(above)
+ lo = np.flatnonzero(~above)
+
+ def _cp(h, c, verts=verts, vals=vals):
+ frac = (threshold - vals[h]) / (vals[c] - vals[h])
+ return verts[h] + frac * (verts[c] - verts[h])
+
+ a_pt, b_pt = _cp(hi[0], lo[0]), _cp(hi[0], lo[1])
+ c_pt, d_pt = _cp(hi[1], lo[1]), _cp(hi[1], lo[0])
+ area += 0.5 * np.linalg.norm(np.cross(c_pt - a_pt, d_pt - b_pt))
+
+ face_count = {}
+ for tet in tets:
+ for f in [(0, 1, 2), (0, 1, 3), (0, 2, 3), (1, 2, 3)]:
+ key = tuple(sorted(tet[list(f)].tolist()))
+ face_count[key] = face_count.get(key, 0) + 1
+ for key, cnt in face_count.items():
+ if cnt == 1:
+ v = np.array(key)
+ frac = is_member[v].sum() / 3.0
+ if frac > 0:
+ q0, q1, q2 = pos[v[0]], pos[v[1]], pos[v[2]]
+ area += frac * 0.5 * np.linalg.norm(np.cross(q1 - q0, q2 - q0))
+ return area
+
+
+def _interior_chebyshev_cube_mesh():
+ """9x6x5 lattice (same domain as `_interior_cube_block_mesh`) with a
+ CONTINUOUS Chebyshev-distance field `t_s = -max(|x-cx|,|y-cy|,|z-cz|)`
+ centered at (4,2,2) -- the level sets of this field are EXACT axis-
+ aligned cubes at any threshold (unlike the binary 0/100 step used
+ elsewhere in this file, whose arbitrary node-count-derived threshold
+ has no meaningful relationship to a genuine half-integer offset). At
+ threshold=-1.5 (exactly bisecting the "distance 1" hot shell and the
+ "distance 2" next shell), `t_s >= -1.5` selects the SAME 27-node block
+ as `_interior_cube_block_mesh` (hand-verified) -- reusing `hot_regions`
+ with top_frac=27/270 reproduces this node set and the correct volume/
+ centroid/t_max/region_id, but its OWN `hot_threshold` lands at -1.9 (a
+ percentile-interpolation artifact of this exact node-count split, not a
+ physically meaningful cut) -- overridden to the analytically clean -1.5
+ for this specific hand-check, isolating the area ALGORITHM's behavior at
+ a genuinely representative threshold position from `hot_regions`' own
+ (separately tested) node-selection quirks."""
+ nx, ny, nz = 9, 6, 5
+ pos, tets = _grid_mesh(nx, ny, nz)
+ center = np.array([4.0, 2.0, 2.0])
+ t_s = -np.max(np.abs(pos - center), axis=1)
+ surface_nodes = _surface_nodes(nx, ny, nz)
+
+ regions = hot_regions(pos, tets, t_s, top_frac=27.0 / 270.0)
+ assert len(regions) == 1
+ region = dict(regions[0])
+ assert region["volume_mm3"] == pytest.approx(27.0)
+ region["hot_threshold"] = -1.5 # see docstring: the clean, symmetric cut
+ return pos, tets, t_s, surface_nodes, region
+
+
+def test_feeder_recommendations_modulus_level_set_consistent():
+ """review round 3: 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."""
+ 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)
+ assert len(recs) == 1
+
+ expected_area = _reference_region_area_marching_tets(pos, tets, region["nodes"], t_s,
+ region["hot_threshold"])
+ expected_modulus = region["volume_mm3"] / expected_area
+ assert recs[0]["min_modulus_mm"] == pytest.approx(1.2 * expected_modulus, rel=1e-9)
+ assert recs[0]["region_id"] == region["region_id"]
+
+ M = recs[0]["min_modulus_mm"] / 1.2
+ assert abs(M - 0.5) / 0.5 < 0.15, f"M={M} not within 15% of the Chvorinov M=0.5 cube reference"
+
+
+def test_feeder_recommendations_location_on_top_surface_above_centroid():
+ pos, tets, t_s, surface_nodes, block_nodes = _interior_cube_block_mesh()
+ regions = hot_regions(pos, tets, t_s)
+ recs = feeder_recommendations(regions, pos, tets, surface_nodes, gravity_vec=(0, 0, -1),
+ t_s=t_s)
+ loc = recs[0]["location"]
+ centroid = regions[0]["centroid"]
+
+ # the recommended point must itself be a surface node...
+ surface_pos = pos[surface_nodes]
+ assert np.any(np.all(np.isclose(surface_pos, loc), axis=1))
+ # ...directly above the centroid (matching x,y; z at the domain's top face)
+ np.testing.assert_allclose(loc[:2], centroid[:2])
+ assert loc[2] == pos[:, 2].max()
+
+
+def test_feeder_recommendations_gravity_rotation_relocates_location():
+ pos, tets, t_s, regions, surface_nodes, node_top, node_xface, node_yface = _rotatable_sink_mesh()
+
+ recs_z = feeder_recommendations(regions, pos, tets, surface_nodes, gravity_vec=(0, 0, -1),
+ t_s=t_s)
+ np.testing.assert_allclose(recs_z[0]["location"], pos[node_top])
+
+ recs_x = feeder_recommendations(regions, pos, tets, surface_nodes, gravity_vec=(-1, 0, 0),
+ t_s=t_s)
+ np.testing.assert_allclose(recs_x[0]["location"], pos[node_xface])
+
+ recs_y = feeder_recommendations(regions, pos, tets, surface_nodes, gravity_vec=(0, -1, 0),
+ t_s=t_s)
+ np.testing.assert_allclose(recs_y[0]["location"], pos[node_yface])
+
+
+def test_feeder_recommendations_empty_surface_nodes_raises_value_error():
+ pos, tets, t_s, surface_nodes, block_nodes = _interior_cube_block_mesh()
+ regions = hot_regions(pos, tets, t_s)
+ with pytest.raises(ValueError):
+ feeder_recommendations(regions, pos, tets, np.array([], dtype=np.int64),
+ gravity_vec=(0, 0, -1), t_s=t_s)
+
+
+def test_feeder_recommendations_degenerate_region_warns_instead_of_silent_zero(caplog):
+ import logging as _logging
+ pos, tets, t_s, surface_nodes, block_nodes = _interior_cube_block_mesh()
+ # a hand-built region touching no tet at all (an isolated single node
+ # from the fixture's background, never part of any hot cluster) has
+ # volume_mm3 == 0 and boundary area == 0 -> degenerate modulus.
+ degenerate_region = {
+ "nodes": np.array([0]),
+ "centroid": pos[0],
+ "volume_mm3": 0.0,
+ "t_max": 0.0,
+ "region_id": 0,
+ "hot_threshold": t_s.max() + 1.0, # nothing in the mesh qualifies as "hot"
+ }
+ with caplog.at_level(_logging.WARNING, logger="src.shrinkage"):
+ recs = feeder_recommendations([degenerate_region], pos, tets, surface_nodes,
+ gravity_vec=(0, 0, -1), t_s=t_s)
+ assert recs[0]["min_modulus_mm"] == 0.0
+ assert any("degenerate" in rec.message for rec in caplog.records)
+
+
+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
+ 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-
+ rotation tests above already use a skin-touching region but only assert
+ `location`, never `min_modulus_mm`. Reuses the boss+band fixture (task 2
+ round 1): the band directly touches 3 domain-boundary faces."""
+ pos, tets, t_s, top_frac, surface_nodes, boss_nodes, band_nodes = _boss_and_band_mesh()
+ regions = hot_regions(pos, tets, t_s, top_frac=top_frac)
+ band_region = next(r for r in regions if set(r["nodes"].tolist()) == band_nodes)
+
+ with caplog.at_level("WARNING", logger="src.shrinkage"):
+ recs = feeder_recommendations([band_region], pos, tets, surface_nodes,
+ gravity_vec=(0, 0, -1), t_s=t_s)
+ assert not any("degenerate" in rec.message for rec in caplog.records)
+ assert recs[0]["min_modulus_mm"] > 0.0
+ assert np.isfinite(recs[0]["min_modulus_mm"])
+
+ # independent hand-check: the reference reimplementation must agree.
+ expected_area = _reference_region_area_marching_tets(
+ pos, tets, band_region["nodes"], t_s, band_region["hot_threshold"])
+ expected_modulus = band_region["volume_mm3"] / expected_area
+ assert recs[0]["min_modulus_mm"] == pytest.approx(1.2 * expected_modulus, rel=1e-9)
+
+
+def _smooth_blob_mesh(spacing, extent=10.0, top_frac=0.05, center_offset=(0.13, 0.29, 0.07)):
+ """A cubic domain of side `extent` mm with a smooth (non-binary-block)
+ field `t_s = -|x - center|^2` -- hottest at the center, smoothly falling
+ off outward, so the top-`top_frac` hot selection is a roughly spherical
+ blob whose boundary genuinely straddles many partial tets at every mesh
+ 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)."""
+ 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)
+ t_s = -((pos - center) ** 2).sum(axis=1)
+ surface_nodes = _surface_nodes(n, n, n)
+ return pos, tets, t_s, top_frac, surface_nodes
+
+
+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-
+ 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
+ grid position at both tested densities. Under the OLD (round-1-shipped)
+ 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."""
+ import logging as _logging
+
+ def modulus_at(spacing):
+ pos, tets, t_s, top_frac, surface_nodes = _smooth_blob_mesh(spacing)
+ regions = hot_regions(pos, tets, t_s, top_frac=top_frac)
+ assert len(regions) == 1
+ with caplog.at_level(_logging.WARNING, logger="src.shrinkage"):
+ recs = feeder_recommendations(regions, pos, tets, surface_nodes,
+ gravity_vec=(0, 0, -1), t_s=t_s)
+ return recs[0]["min_modulus_mm"] / 1.2
+
+ m_coarse = modulus_at(1.0)
+ m_fine = modulus_at(0.5)
+ drift = abs(m_fine - m_coarse) / m_coarse
+ assert drift < 0.15, f"M drift {drift:.1%} exceeds the 15% bound (coarse={m_coarse}, fine={m_fine})"
+
+
+# ---------------------------------------------------------------------------
+# (g) review round 3 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
+# orientation sweep (axis/face-diagonal/space-diagonal/generic normals)
+# showing the error stays small and roughly ORIENTATION-INDEPENDENT.
+# ---------------------------------------------------------------------------
+
+def _sphere_region(spacing, extent=12.0, top_frac=0.03, center_offset=(0.13, 0.29, 0.07)):
+ """A spherical hot blob (off-grid center, see `_smooth_blob_mesh`) with
+ its ANALYTIC area/radius derived directly from the threshold `hot_
+ regions` actually used (R = sqrt(-threshold) exactly, since the
+ continuous field t_s = -|pos-center|^2 has a perfectly spherical level
+ set at every threshold, independent of mesh discretization)."""
+ 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)
+ t_s = -((pos - center) ** 2).sum(axis=1)
+ threshold = np.percentile(t_s, 100 * (1 - top_frac))
+ nodes = np.flatnonzero(t_s >= threshold)
+ radius = np.sqrt(-threshold)
+ true_area = 4.0 * np.pi * radius ** 2
+ return pos, tets, nodes, t_s, threshold, true_area
+
+
+def test_region_area_sphere_converges_not_plateaus():
+ """review round 3 acceptance check #1: 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
+ co-area formula."""
+ errors = []
+ for spacing in (1.0, 0.5, 0.25):
+ pos, tets, nodes, t_s, threshold, true_area = _sphere_region(spacing)
+ area = _reference_region_area_marching_tets(pos, tets, nodes, t_s, threshold)
+ errors.append(abs(area - true_area) / true_area)
+
+ assert errors[0] < 0.10, f"coarsest-mesh sphere error too high: {errors[0]:.1%}"
+ assert errors[-1] < errors[0], f"error did not decrease with refinement: {errors}"
+ assert errors[-1] < 0.02, f"finest-mesh sphere error too high: {errors[-1]:.1%}"
+
+
+def _halfspace_box_total_area(box_extent, normal, offset):
+ """Analytic total boundary area of `box ∩ {p : normal.p <= offset}` for
+ an axis-aligned box `[0,box_extent]^3` -- exact via `scipy.spatial.
+ ConvexHull` on the (convex) polytope's vertices: the box corners already
+ inside the half-space, plus the half-space plane's intersection points
+ with each of the box's 12 edges. Serves as the analytic reference for
+ the diagonal-cut and orientation-sweep tests below (hand-verified
+ against known closed forms: an axis-aligned cut of a 10mm cube at x=5
+ gives exactly 400mm^2 = surface area of a 5x10x10 box; a cut through the
+ center along the space diagonal gives 300 [half the original 600mm^2
+ surface, split evenly by symmetry] + 3*sqrt(3)/4*10^2 [=129.90, the
+ well-known hexagonal cross-section area of a cube bisected
+ perpendicular to its space diagonal] = 429.90)."""
+ from scipy.spatial import ConvexHull
+
+ normal = np.asarray(normal, dtype=np.float64)
+ normal = normal / np.linalg.norm(normal)
+ lo, hi = 0.0, float(box_extent)
+ corners = np.array([[x, y, z] for x in (lo, hi) for y in (lo, hi) for z in (lo, hi)])
+ edges = []
+ for xi in range(2):
+ for yi in range(2):
+ for zi in range(2):
+ idx0 = 4 * xi + 2 * yi + zi
+ if xi == 0:
+ edges.append((idx0, 4 * 1 + 2 * yi + zi))
+ if yi == 0:
+ edges.append((idx0, 4 * xi + 2 * 1 + zi))
+ if zi == 0:
+ edges.append((idx0, 4 * xi + 2 * yi + 1))
+
+ pts = [c for c in corners if normal @ c - offset <= 1e-9]
+ for i, j in edges:
+ vi, vj = corners[i], corners[j]
+ si, sj = normal @ vi - offset, normal @ vj - offset
+ if si * sj < 0:
+ t = si / (si - sj)
+ pts.append(vi + t * (vj - vi))
+ return ConvexHull(np.array(pts)).area
+
+
+def _diagonal_cut_region(spacing, normal, extent=10.0, top_frac=0.3):
+ pos, tets = _grid_mesh(int(round(extent / spacing)) + 1,
+ int(round(extent / spacing)) + 1,
+ int(round(extent / spacing)) + 1, spacing=spacing)
+ normal = np.asarray(normal, dtype=np.float64)
+ normal = normal / np.linalg.norm(normal)
+ t_s = -(pos @ normal)
+ threshold = np.percentile(t_s, 100 * (1 - top_frac))
+ nodes = np.flatnonzero(t_s >= threshold)
+ true_area = _halfspace_box_total_area(extent, normal, offset=-threshold)
+ return pos, tets, nodes, t_s, threshold, true_area
+
+
+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
+ 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
+ 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))
+ area = _reference_region_area_marching_tets(pos, tets, nodes, t_s, threshold)
+ err = abs(area - true_area) / true_area
+ 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):
+# each tet contributes tet_volume * n_hot_vertices/4, replacing an earlier
+# all-4-vertices-hot rule that discarded (undercounted) every boundary tet.
+# ---------------------------------------------------------------------------
+
+def test_hot_regions_volume_fraction_weighted_partial_tets_hand_checked():
+ """An isolated single unit-cube mesh (2x2x2 nodes, spacing 1mm, no
+ neighboring cubes -- nothing outside this mesh to contaminate the
+ count): corner 7 (the (1,1,1) corner) is background/cold, the other 7
+ corners are hot. Every one of the cube's 6 Kuhn tets contains corner 7
+ (hand-checked from the tet list), so every tet has exactly 3 of its 4
+ vertices hot; each tet has volume 1/6 mm^3 (regular unit cube / 6). Total
+ fraction-weighted volume = 6 * (1/6) * (3/4) = 0.75 mm^3 exactly. Under
+ the old all-4-vertices rule this region would score 0.0 mm^3 (no tet has
+ all 4 vertices hot) -- the undercount this fix targets."""
+ pos, tets = _grid_mesh(2, 2, 2)
+ t_s = np.full(8, 100.0)
+ t_s[7] = 0.0
+ top_frac = 7.0 / 8.0
+ regions = hot_regions(pos, tets, t_s, top_frac=top_frac)
+ assert len(regions) == 1
+ assert set(regions[0]["nodes"].tolist()) == {0, 1, 2, 3, 4, 5, 6}
+ assert regions[0]["volume_mm3"] == pytest.approx(0.75)
+
+
+def test_hot_regions_volume_full_mesh_matches_analytic_total():
+ """Sanity/regression check: when EVERY node is hot (top_frac=1.0), no
+ tet straddles the region boundary (n_hot_vertices == 4 everywhere), so
+ the fraction-weighted rule reduces to the plain tet-volume sum, which
+ for a unit-spacing (nx-1)x(ny-1)x(nz-1) lattice equals that product
+ exactly."""
+ nx, ny, nz = 4, 3, 3
+ pos, tets = _grid_mesh(nx, ny, nz)
+ t_s = np.ones(pos.shape[0]) # uniform -> every node is "hot" at top_frac=1.0
+ regions = hot_regions(pos, tets, t_s, top_frac=1.0)
+ assert len(regions) == 1
+ assert regions[0]["volume_mm3"] == pytest.approx((nx - 1) * (ny - 1) * (nz - 1))
+
+
+# ---------------------------------------------------------------------------
+# (e) determinism -- two runs on the same inputs produce identical results.
+# ---------------------------------------------------------------------------
+
+def test_determinism_hot_regions_and_downstream():
+ pos, tets, t_s, surface_nodes, block_nodes = _interior_cube_block_mesh()
+ no_feed = np.array([], dtype=np.int64)
+
+ regions1 = hot_regions(pos, tets, t_s)
+ regions2 = hot_regions(pos, tets, t_s)
+ assert len(regions1) == len(regions2) == 1
+ np.testing.assert_array_equal(regions1[0]["nodes"], regions2[0]["nodes"])
+ np.testing.assert_allclose(regions1[0]["centroid"], regions2[0]["centroid"])
+ assert regions1[0]["volume_mm3"] == regions2[0]["volume_mm3"]
+ assert regions1[0]["t_max"] == regions2[0]["t_max"]
+ assert regions1[0]["region_id"] == regions2[0]["region_id"]
+
+ macro1 = macro_shrinkage(pos, tets, t_s, no_feed, shrinkage_pct=2.5)
+ macro2 = macro_shrinkage(pos, tets, t_s, no_feed, shrinkage_pct=2.5)
+ assert macro1[0]["severity"] == macro2[0]["severity"]
+
+ feeders1 = feeder_recommendations(regions1, pos, tets, surface_nodes, gravity_vec=(0, 0, -1),
+ t_s=t_s)
+ feeders2 = feeder_recommendations(regions2, pos, tets, surface_nodes, gravity_vec=(0, 0, -1),
+ t_s=t_s)
+ assert feeders1[0]["min_modulus_mm"] == feeders2[0]["min_modulus_mm"]
+ np.testing.assert_array_equal(feeders1[0]["location"], feeders2[0]["location"])
+
+
+# ---------------------------------------------------------------------------
+# 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.
+# ---------------------------------------------------------------------------
+
+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():
+ nx = ny = nz = 4
+ pos, tets = _grid_mesh(nx, ny, nz)
+ t_s = pos[:, 0].copy() + 1.0
+ surface_nodes = _surface_nodes(nx, ny, nz)
+ TL, TS = 1400.0, 1100.0
+
+ result = screen_casting(pos, tets, surface_nodes, t_s, TL, TS, "grey_iron")
+
+ 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")
+ 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():
+ pos, tets, t_s, surface_nodes, block_nodes = _interior_cube_block_mesh()
+ with pytest.raises(KeyError) as exc_info:
+ screen_casting(pos, tets, surface_nodes, t_s, 1400.0, 1100.0, "unobtainium")
+ msg = str(exc_info.value)
+ assert "unobtainium" in msg
+ for family in SHRINKAGE_PCT:
+ assert family in msg
+
+
+def test_screen_casting_feed_nodes_default_is_empty_bare_casting_convention():
+ pos, tets, t_s, surface_nodes, boss_nodes, band_nodes = _unfed_boss_and_fed_band_10pct_mesh()
+ TL, TS = 1400.0, 1100.0
+
+ default_call = screen_casting(pos, tets, surface_nodes, t_s, TL, TS, "steel")
+ explicit_empty = screen_casting(pos, tets, surface_nodes, t_s, TL, TS, "steel",
+ feed_nodes=np.array([], dtype=np.int64))
+ assert len(default_call["macro_shrinkage"]) == len(explicit_empty["macro_shrinkage"]) == 2
+
+ 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
diff --git a/tools/__init__.py b/tools/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/tools/doe/__init__.py b/tools/doe/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/tools/doe/load_properties.py b/tools/doe/load_properties.py
new file mode 100644
index 0000000000000000000000000000000000000000..f4e4038767ba57cf6e8dd52316646b419cdd4209
--- /dev/null
+++ b/tools/doe/load_properties.py
@@ -0,0 +1,114 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..2dbaf86fc7846d6948f42531dc162439e6d55bc2
--- /dev/null
+++ b/tools/doe/mesh_study.py
@@ -0,0 +1,18 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..f2b51789d48f9b2517a33c824251bad4be8a3c85
--- /dev/null
+++ b/tools/doe/prep_parts.py
@@ -0,0 +1,890 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..462a540c855a912a93b2327710cc6ab8091bf817
--- /dev/null
+++ b/tools/export_release.py
@@ -0,0 +1,111 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..e5f86bf48acacd984eb9e36b1c1b2186fb035cc1
--- /dev/null
+++ b/tools/recovery/README.md
@@ -0,0 +1,15 @@
+# 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
new file mode 100644
index 0000000000000000000000000000000000000000..2eedcede5f0efe9120bd82677676808feb1ba1ba
--- /dev/null
+++ b/tools/recovery/compute_perms.py
@@ -0,0 +1,97 @@
+"""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
new file mode 100644
index 0000000000000000000000000000000000000000..31b48e840fc8c89fc0fdeb17672f2c0a62553480
--- /dev/null
+++ b/tools/recovery/reprocess_failures.py
@@ -0,0 +1,99 @@
+"""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)
diff --git a/tools/sleeves/__init__.py b/tools/sleeves/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/tools/sleeves/sleeve_registry.csv b/tools/sleeves/sleeve_registry.csv
new file mode 100644
index 0000000000000000000000000000000000000000..21c5a7e2faf1745e2049cd68444e3d6ac1ac57a9
--- /dev/null
+++ b/tools/sleeves/sleeve_registry.csv
@@ -0,0 +1,15 @@
+name,source,type,d1_mm,d2_mm,H_mm,t_mm,V_cm3,A_cm2,M_geom_cm,MEF,M_eff_cm,mef_source
+SYNTH_M050,synthetic,exo,27.952,24.202,42.950,5.710,22.978,45.956,0.5000,1.400,0.7000,synthetic_spread
+SYNTH_M065,synthetic,exo,36.337,31.463,55.835,7.423,50.483,77.666,0.6500,1.650,1.0725,synthetic_spread
+SD_4_7K,catalog,exo,41.000,35.500,63.000,8.375,72.517,98.877,0.7334,1.547,1.1343,catalog_ratio
+SD_4_95K,catalog,exo,42.000,36.000,84.000,8.750,100.544,127.017,0.7916,1.562,1.2368,catalog_ratio
+SD_5_8K,catalog,exo,51.000,47.500,69.500,10.250,132.455,145.716,0.9090,1.579,1.4353,catalog_ratio
+FOSECO_P5_8K,foseco_step,exo,,,,,,,0.9500,1.580,1.5000,datasheet
+SYNTH_M080,synthetic,exo,44.723,38.723,68.720,9.135,94.119,117.648,0.8000,1.900,1.5200,synthetic_spread
+FOSECO_P7_10K,foseco_step,exo,,,,,,,1.2500,1.560,1.9500,datasheet
+FOSECO_P8_11K,foseco_step,exo,,,,,,,1.4000,1.570,2.2000,datasheet
+FOSECO_ZP912K14QC,foseco_step,exo,,,,,,,1.5500,1.550,2.4000,datasheet
+FOSECO_P10_13K,foseco_step,exo,,,,,,,1.7500,1.540,2.7000,datasheet
+FOSECO_P12_15K,foseco_step,exo,,,,,,,2.0000,1.550,3.1000,datasheet
+FOSECO_V90L,foseco_step,insulating,,,,,,,,1.150,1.7000,unverified
+FOSECO_V88E,foseco_step,insulating,,,,,,,,1.450,1.7000,unverified
diff --git a/weights/SHA256SUMS b/weights/SHA256SUMS
new file mode 100644
index 0000000000000000000000000000000000000000..279783e960855d6afeb3945f4c87bf6a1fd8fd87
--- /dev/null
+++ b/weights/SHA256SUMS
@@ -0,0 +1,4 @@
+e1dd5d9227a9495af51177f52f78bd28507134c19e1a82d9341fbf49733a820a gnn_v5.1.safetensors
+0a0b9caa7f1ebcea94e220dd4704ef63309c5f1619fbc20427f90ae1e6727b83 gbm_total_v3.skops
+bd4469c06490c4eb8b3d90b24144b2fd7059ae8faab7012e78b12cfaaf25239a gbm_std_v2.skops
+5c24d0c99064ea264e17b44c9c419f73e6793e509215b7099cf5a5902c9a12ae ood_ranges.json
diff --git a/weights/gbm_std_v2.skops b/weights/gbm_std_v2.skops
new file mode 100644
index 0000000000000000000000000000000000000000..1b974ba3612d6d30c40b9f7ab5fd2740d4b31591
--- /dev/null
+++ b/weights/gbm_std_v2.skops
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd4469c06490c4eb8b3d90b24144b2fd7059ae8faab7012e78b12cfaaf25239a
+size 849155
diff --git a/weights/gbm_total_v3.skops b/weights/gbm_total_v3.skops
new file mode 100644
index 0000000000000000000000000000000000000000..10f84e249ce2a944c46010401d9ab163726e94d1
--- /dev/null
+++ b/weights/gbm_total_v3.skops
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a0b9caa7f1ebcea94e220dd4704ef63309c5f1619fbc20427f90ae1e6727b83
+size 615362
diff --git a/weights/gnn_v5.1.safetensors b/weights/gnn_v5.1.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..8031e1d0776249afb6ff37e5852c60e85265cd86
--- /dev/null
+++ b/weights/gnn_v5.1.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e1dd5d9227a9495af51177f52f78bd28507134c19e1a82d9341fbf49733a820a
+size 7409448
diff --git a/weights/ood_ranges.json b/weights/ood_ranges.json
new file mode 100644
index 0000000000000000000000000000000000000000..03132174bed53ff50edb71875bd7857e6be7620c
--- /dev/null
+++ b/weights/ood_ranges.json
@@ -0,0 +1,5 @@
+{
+ "volume_mm3_max": 390663008.0,
+ "bbox_diag_max": 2177.705322265625,
+ "geodesic_d_cool_max": 125.68257904052734
+}
\ No newline at end of file