"""Train weak-label diagnostics and transfer every model output to one CAD graph.
The experiment uses run 18h01 because its diagnostic registration has determinant
+1. Labels and registration are still unvalidated; all metrics are training-set
representation checks, never held-out performance.
"""
from __future__ import annotations
import json
import math
import os
import time
from pathlib import Path
from typing import Callable
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from scipy.spatial import cKDTree
from sklearn.metrics import average_precision_score
import torch
from torch import Tensor, nn
import torch.nn.functional as F
from .architectures import (CNN2DUNet, DGCNNSeg, FNO2d, MeshGraphProcessor,
PhysicsConstrainedRiskGNN, PointNetSeg,
SparseUNetOptional, TabularSurrogate, UNet3D,
UVNetFaceModel)
from .data import (ROOT, brep_input, graph_inputs, image_input, point_inputs,
process_conditions, sparse_batch_collate, sparse_voxel_input,
tabular_input, voxel_input)
from .train_weak_gnn import load_training_data
from .training import seed_everything
OUT = ROOT / "results/models/cad_predictions"
RUN = os.environ.get("MOLTEX_RUN", "18h01")
def weak_fit_checkpoint(model: nn.Module, name: str, *, scope: str | None = None) -> str:
"""Persist the exact selected state used for the displayed CAD field."""
slug = name.lower().replace("/", "-").replace(" ", "-")
path = OUT / "checkpoints" / f"{RUN}-{slug}.pt"
path.parent.mkdir(parents=True, exist_ok=True)
torch.save({
"state_dict": {key: value.detach().cpu().clone() for key, value in model.state_dict().items()},
"production_valid": False,
"generalization_measured": False,
"run": RUN,
"checkpoint_scope": scope or f"{RUN} weak-label diagnostic fit",
"warning": "Weak-label training-set diagnostic only; not a production model.",
}, path)
return str(path.relative_to(ROOT))
def tensor(value: np.ndarray, dtype: torch.dtype = torch.float32) -> Tensor:
return torch.as_tensor(value, dtype=dtype)
def risk_loss(logit: Tensor, target: Tensor, valid: Tensor) -> Tensor:
selected_logit, selected_target = logit[valid], target[valid]
positive = (selected_target >= .5).sum().clamp_min(1)
weight = ((len(selected_target) - positive) / positive).clamp(max=200)
bce = F.binary_cross_entropy_with_logits(selected_logit, selected_target, pos_weight=weight)
probability = torch.sigmoid(selected_logit)
dice = 1 - (2 * (probability * selected_target).sum() + 1) / (probability.sum() + selected_target.sum() + 1)
return bce + dice
def fit(model: nn.Module, forward: Callable[[], Tensor], target: Tensor, valid: Tensor,
epochs: int, lr: float = 3e-3) -> tuple[np.ndarray, dict]:
optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-5)
history = []
best_loss = float("inf"); best_epoch = 0; best_state = None
start = time.perf_counter()
for epoch in range(epochs):
model.train(); optimizer.zero_grad(set_to_none=True)
logit = forward(); loss = risk_loss(logit, target, valid)
loss.backward(); optimizer.step()
value = float(loss.detach()); history.append(value)
if value < best_loss:
best_loss, best_epoch = value, epoch + 1
best_state = {key: value.detach().clone() for key, value in model.state_dict().items()}
if best_state is None:
raise RuntimeError("No training epoch completed")
model.load_state_dict(best_state); model.eval()
with torch.no_grad(): probability = torch.sigmoid(forward()).cpu().numpy()
return probability, {"loss_initial": history[0], "loss_final": best_loss,
"selected_epoch": best_epoch, "epochs": epochs,
"training_seconds": time.perf_counter() - start}
def nearest_target(query: np.ndarray, cad_pos: np.ndarray, target: np.ndarray) -> np.ndarray:
return target[cKDTree(cad_pos).query(query, workers=-1)[1]]
def barycentric_target(target: np.ndarray, source_triangle: np.ndarray,
barycentric: np.ndarray, faces: np.ndarray) -> np.ndarray:
return np.einsum("ni,ni->n", target[faces[source_triangle]], barycentric)
def barycentric_to_cad(values: np.ndarray, source_triangle: np.ndarray,
barycentric: np.ndarray, faces: np.ndarray,
cad_pos: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
vertices = faces[source_triangle]
weighted = np.asarray(values, np.float64)[:, None] * barycentric
total = np.zeros(len(cad_pos), np.float64); weight = np.zeros(len(cad_pos), np.float64)
np.add.at(total, vertices.ravel(), weighted.ravel()); np.add.at(weight, vertices.ravel(), barycentric.ravel())
direct = weight > 1e-10; field = np.zeros(len(cad_pos), np.float64); field[direct] = total[direct] / weight[direct]
distance = np.zeros(len(cad_pos), np.float64)
if not direct.all():
fill_distance, fill_index = cKDTree(cad_pos[direct]).query(cad_pos[~direct], workers=-1)
field[~direct] = field[direct][fill_index]; distance[~direct] = fill_distance
return field.astype(np.float32), weight.astype(np.float32), distance.astype(np.float32)
def sparse_logits_on_prepared_coords(output: Tensor, prepared_coords: np.ndarray) -> Tensor:
"""Return sparse logits in the exact prepared ``[batch, x, y, z]`` order.
MinkowskiEngine is allowed to choose its own internal coordinate ordering.
The weak target and the saved ``surface_sample_voxel_index`` instead use the
explicit sparse NPZ row order, so coordinate equality is checked before
gathering logits back into that order.
"""
output_coords = np.asarray(output.C.detach().cpu(), np.int32)
expected = np.asarray(prepared_coords, np.int32)
if output_coords.shape != expected.shape:
raise RuntimeError(
f"Sparse U-Net coordinate shape changed: got {output_coords.shape}, expected {expected.shape}")
if len(np.unique(output_coords, axis=0)) != len(output_coords):
raise RuntimeError("Sparse U-Net returned duplicate output coordinates")
row_for_coordinate = {tuple(row): index for index, row in enumerate(output_coords.tolist())}
try:
order = [row_for_coordinate[tuple(row)] for row in expected.tolist()]
except KeyError as exc:
raise RuntimeError("Sparse U-Net output coordinates do not match the prepared coordinate set") from exc
if len(order) != len(output_coords):
raise RuntimeError("Sparse U-Net output coordinate set has unexpected extra coordinates")
if output.F.ndim != 2 or output.F.shape[1] != 1:
raise RuntimeError(f"Sparse U-Net diagnostic requires one logit per voxel, got {tuple(output.F.shape)}")
return output.F[torch.as_tensor(order, dtype=torch.long, device=output.F.device), 0]
def metrics(field: np.ndarray, binary: np.ndarray) -> dict:
field = np.clip(np.asarray(field, float), 0, 1)
truth = np.asarray(binary, bool)
pred = field >= .5
intersection = int((pred & truth).sum()); union = int((pred | truth).sum())
positives = int(truth.sum()); k = max(1, positives)
top = np.argpartition(field, -k)[-k:]
return {
"average_precision": float(average_precision_score(truth, field)),
"threshold_iou": intersection / max(1, union),
"threshold_dice": 2 * intersection / max(1, int(pred.sum()) + positives),
"top_k_recall": float(truth[top].sum() / k),
"predicted_positive_nodes": int(pred.sum()),
"target_positive_nodes": positives,
}
def train_all() -> tuple[np.ndarray, np.ndarray, dict[str, np.ndarray], dict[str, dict], dict[str, np.ndarray], dict]:
torch.set_num_threads(max(1, min(4, torch.get_num_threads())))
seed_everything()
paired_path = ROOT / f"data_processed/paired/{RUN}_cad_ply_targets.npz"
with np.load(paired_path, allow_pickle=False) as paired:
cad_pos = np.asarray(paired["cad_pos_mm"], np.float32)
target_soft = np.asarray(paired["weak_defect_probability"][:, 1], np.float32)
target_binary = np.asarray(paired["weak_defect_mask"][:, 1], bool)
process_x, process_mask, run_ids, _ = process_conditions()
run_index = run_ids.index(RUN)
process = tensor(np.concatenate([process_x, process_mask.astype(np.float32)], 1)[run_index:run_index+1])
fields: dict[str, np.ndarray] = {}
reports: dict[str, dict] = {}
display_valid: dict[str, np.ndarray] = {}
sparse_status: dict = {"status": "unavailable_optional_backend", "backend": "MinkowskiEngine",
"reason": "MinkowskiEngine not installed", "production_valid": False}
canonical = np.load(ROOT / "data_processed/graphs/graph_12000.npz", allow_pickle=False)
faces = np.asarray(canonical["face"], np.int64).T
# Image models: supervise only visible CAD pixels, then nearest-project all CAD nodes.
raw_image = np.load(ROOT / "data_processed/images/cad_multiview_128.npz", allow_pickle=False)
image = image_input(); image_x = tensor(np.concatenate([image["x"], image["mask"].astype(np.float32)], 1))
visible = np.asarray(raw_image["visibility"], bool); xyz = np.asarray(raw_image["xyz_mm"], np.float32)
image_target = np.zeros(visible.shape, np.float32)
image_source = np.asarray(raw_image["source_triangle"], np.int64); image_bary = np.asarray(raw_image["barycentric"], np.float32)
image_target[visible] = barycentric_target(target_soft, image_source[visible], image_bary[visible], faces)
image_valid = tensor(visible, torch.bool); image_target_t = tensor(image_target)
image_process = process.repeat(len(image_x), 1)
for name, model, epochs in (
("2D CNN/U-Net", CNN2DUNet(image_x.shape[1], global_dim=process.shape[1]), 80),
("FNO2d", FNO2d(image_x.shape[1], global_dim=process.shape[1]), 60),
):
probability, report = fit(model, lambda m=model: m(image_x, image_process)[:, 0], image_target_t, image_valid, epochs)
checkpoint = weak_fit_checkpoint(model, name)
fields[name], coverage, distance = barycentric_to_cad(probability[visible], image_source[visible], image_bary[visible], faces, cad_pos)
display_valid[name] = distance <= 10
reports[name] = report | {"checkpoint": checkpoint, "checkpoint_scope": f"{RUN} weak-label diagnostic fit",
"native_domain": "six visible CAD image planes", "process_conditioned": True,
"mapping_method": "stored source triangles and barycentric scatter; nearest direct-node fill",
"direct_node_coverage": float((coverage > 0).mean()), "fill_distance_p95_mm": float(np.quantile(distance, .95))}
# Dense surface voxel model; supervise occupied cells and query nearest occupied cell for CAD display.
dense = voxel_input(); voxel_x = tensor(np.concatenate([dense["x"], dense["mask"].astype(np.float32)], 0)[None])
sparse = np.load(ROOT / "data_processed/voxels/cad_surface_sparse_32.npz", allow_pickle=False)
coord = np.asarray(sparse["voxel_coords"], int); origin = np.asarray(sparse["grid_origin_mm"], float); pitch = float(sparse["voxel_pitch_mm"])
voxel_target = np.zeros(tuple(np.asarray(sparse["grid_dims"], int)), np.float32)
voxel_valid = np.zeros_like(voxel_target, bool)
sample_tri = np.asarray(sparse["surface_sample_source_triangle"], np.int64)
sample_bary = np.asarray(sparse["surface_sample_barycentric"], np.float32)
sample_voxel = np.asarray(sparse["surface_sample_voxel_index"], np.int64)
sample_target = barycentric_target(target_soft, sample_tri, sample_bary, faces)
voxel_values = np.zeros(len(coord), np.float32); np.maximum.at(voxel_values, sample_voxel, sample_target)
voxel_target[tuple(coord.T)] = voxel_values; voxel_valid[tuple(coord.T)] = True
model3d = UNet3D(voxel_x.shape[1], global_dim=process.shape[1])
probability, report = fit(model3d, lambda: model3d(voxel_x, process)[:, 0], tensor(voxel_target[None]), tensor(voxel_valid[None], torch.bool), 100)
checkpoint = weak_fit_checkpoint(model3d, "Dense 3D U-Net")
occupied_probability = probability[0][tuple(coord.T)]
sample_probability = occupied_probability[sample_voxel]
fields["Dense 3D U-Net"], coverage, distance = barycentric_to_cad(sample_probability, sample_tri, sample_bary, faces, cad_pos)
display_valid["Dense 3D U-Net"] = distance <= 10
reports["Dense 3D U-Net"] = report | {"checkpoint": checkpoint, "checkpoint_scope": f"{RUN} weak-label diagnostic fit",
"native_domain": "403 occupied cells in a dense 32×4×16 grid", "process_conditioned": True,
"mapping_method": "surface voxel → stored surface samples → triangle/barycentric scatter",
"direct_node_coverage": float((coverage > 0).mean()), "fill_distance_p95_mm": float(np.quantile(distance, .95))}
# Sparse surface voxel model. Its weak labels use the same canonical 403
# occupied coordinates as the sparse NPZ; output rows are coordinate-matched
# before being transferred through stored sample triangle/barycentric links.
try:
import MinkowskiEngine as ME # type: ignore
except ImportError as exc:
sparse_status["reason"] = f"backend import unavailable: {exc}"
else:
sparse_input = sparse_voxel_input()
sparse_coords_np, sparse_features_np = sparse_batch_collate(
sparse_input["coords"], sparse_input["features"])
if not np.array_equal(sparse_coords_np[:, 1:], coord):
raise RuntimeError("Sparse input coordinates differ from cad_surface_sparse_32 voxel_coords")
sparse_coords = tensor(sparse_coords_np, torch.int32)
sparse_features = tensor(sparse_features_np)
sparse_model = SparseUNetOptional(sparse_features.shape[1])
def sparse_forward() -> Tensor:
sparse_tensor = ME.SparseTensor(features=sparse_features, coordinates=sparse_coords)
output = sparse_model(sparse_tensor)
return sparse_logits_on_prepared_coords(output, sparse_coords_np)
probability, report = fit(
sparse_model, sparse_forward, tensor(voxel_values),
torch.ones(len(voxel_values), dtype=torch.bool), 100)
checkpoint = weak_fit_checkpoint(sparse_model, "Sparse 3D U-Net")
sparse_sample_probability = probability[sample_voxel]
fields["Sparse 3D U-Net"], coverage, distance = barycentric_to_cad(
sparse_sample_probability, sample_tri, sample_bary, faces, cad_pos)
display_valid["Sparse 3D U-Net"] = distance <= 10
reports["Sparse 3D U-Net"] = report | {"checkpoint": checkpoint, "checkpoint_scope": f"{RUN} weak-label diagnostic fit",
"native_domain": "403 canonical occupied sparse-surface voxels",
"process_conditioned": False,
"mapping_method": "coordinate-matched sparse output → stored surface samples → triangle/barycentric scatter",
"output_coordinate_alignment_verified": True,
"direct_node_coverage": float((coverage > 0).mean()),
"fill_distance_p95_mm": float(np.quantile(distance, .95)),
"production_valid": False,
}
sparse_status = {
"status": "diagnostic_weak_label_fit_completed", "backend": "MinkowskiEngine",
"production_valid": False, "generalization_measured": False,
"checkpoint": checkpoint, "loss_initial": report["loss_initial"],
"loss_final": report["loss_final"], "selected_epoch": report["selected_epoch"],
"epochs": report["epochs"], "training_seconds": report["training_seconds"],
"native_voxels": int(len(voxel_values)),
"output_coordinate_alignment_verified": True,
}
# CAD point models.
point_data = point_inputs()[0]; raw_point = np.load(ROOT / "data_processed/points/cad_surface_points.npz", allow_pickle=False)
point_pos = np.asarray(raw_point["pos_mm"], np.float32)
point_x = tensor(np.concatenate([point_data["x"], point_data["mask"].astype(np.float32)], 1)[None])
point_source = np.asarray(raw_point["source_triangle"], np.int64); point_bary = np.asarray(raw_point["barycentric"], np.float32)
point_target = barycentric_target(target_soft, point_source, point_bary, faces)
point_model = PointNetSeg(point_x.shape[-1])
probability, report = fit(point_model, lambda: point_model(point_x)[0, :, 0], tensor(point_target), torch.ones(len(point_target), dtype=torch.bool), 120)
checkpoint = weak_fit_checkpoint(point_model, "PointNet")
fields["PointNet"], coverage, distance = barycentric_to_cad(probability, point_source, point_bary, faces, cad_pos)
display_valid["PointNet"] = distance <= 10
reports["PointNet"] = report | {"checkpoint": checkpoint, "checkpoint_scope": f"{RUN} weak-label diagnostic fit",
"native_domain": "8,192 CAD surface points", "process_conditioned": False,
"mapping_method": "stored point triangle/barycentric scatter; nearest direct-node fill",
"direct_node_coverage": float((coverage > 0).mean()), "fill_distance_p95_mm": float(np.quantile(distance, .95))}
dcount = min(512, len(point_pos)); dg_x = point_x[:, :dcount]; dg_pos = point_pos[:dcount]; dg_target = point_target[:dcount]
dg_model = DGCNNSeg(dg_x.shape[-1], k=12)
probability, report = fit(dg_model, lambda: dg_model(dg_x)[0, :, 0], tensor(dg_target), torch.ones(dcount, dtype=torch.bool), 100)
checkpoint = weak_fit_checkpoint(dg_model, "DGCNN")
fields["DGCNN"], coverage, distance = barycentric_to_cad(probability, point_source[:dcount], point_bary[:dcount], faces, cad_pos)
display_valid["DGCNN"] = distance <= 10
reports["DGCNN"] = report | {"checkpoint": checkpoint, "checkpoint_scope": f"{RUN} weak-label diagnostic fit",
"native_domain": "512-point CPU subset", "process_conditioned": False,
"mapping_method": "stored point triangle/barycentric scatter; large nearest-node fill",
"direct_node_coverage": float((coverage > 0).mean()), "fill_distance_p95_mm": float(np.quantile(distance, .95))}
# Direct CAD graph risk field.
graph = max(graph_inputs(), key=lambda item: item["nodes"])
graph_x = tensor(np.concatenate([graph["x"], graph["mask"].astype(np.float32)], 1))
graph_edge = tensor(graph["edge_index"], torch.long); graph_attr_np = np.asarray(graph["edge_attr"], np.float32).copy(); graph_attr_np[:, 0] /= 1000
graph_attr = tensor(graph_attr_np); graph_batch = torch.zeros(len(graph_x), dtype=torch.long)
graph_model = MeshGraphProcessor(graph_x.shape[1], graph_attr.shape[1], process.shape[1], width=48, steps=3)
probability, report = fit(graph_model, lambda: graph_model(graph_x, graph_edge, graph_attr, process, graph_batch)[:, 0], tensor(target_soft), torch.ones(len(target_soft), dtype=torch.bool), 100)
checkpoint = weak_fit_checkpoint(graph_model, "Mesh GNN")
fields["Mesh GNN"] = probability
display_valid["Mesh GNN"] = np.ones(len(cad_pos), bool)
reports["Mesh GNN"] = report | {"checkpoint": checkpoint, "checkpoint_scope": f"{RUN} weak-label diagnostic fit",
"native_domain": f"{len(cad_pos):,} CAD graph nodes", "process_conditioned": True}
# The constrained-GNN fit uses all paired runs available in this experiment root.
weak_data = load_training_data(); weak_model = PhysicsConstrainedRiskGNN(
weak_data["node"].shape[1], weak_data["edge_attr"].shape[1], weak_data["process"].shape[1], max(1, len(weak_data["class_names"])))
checkpoint = torch.load(ROOT / "results/models/weak_gnn/diagnostic_weak_gnn.pt", map_location="cpu", weights_only=False)
weak_model.load_state_dict(checkpoint["state_dict"]); weak_model.eval()
risk_checkpoint = weak_fit_checkpoint(weak_model, "Risk GNN", scope=f"{len(weak_data['runs'])}-run weak-label diagnostic fit; displayed {RUN} condition")
with torch.no_grad():
out = weak_model(tensor(weak_data["node"]), tensor(weak_data["edge"], torch.long), tensor(weak_data["edge_attr"]),
tensor(weak_data["process"]), tensor(weak_data["batch"], torch.long))
count = weak_data["nodes_per_run"]; index = weak_data["runs"].index(RUN)
fields["Risk GNN"] = torch.sigmoid(out["risk_logit"])[index*count:(index+1)*count].numpy()
display_valid["Risk GNN"] = np.ones(len(cad_pos), bool)
weak_report = json.loads((ROOT / "results/models/weak_gnn/report.json").read_text())
reports["Risk GNN"] = {"checkpoint": risk_checkpoint, "checkpoint_scope": f"{len(weak_data['runs'])}-run weak-label diagnostic fit; displayed {RUN} condition",
"loss_initial": weak_report["loss_initial"], "loss_final": weak_report["loss_final"],
"selected_epoch": weak_report["selected_epoch"], "epochs": weak_report["epochs"],
"training_seconds": None, "native_domain": f"{len(cad_pos):,} CAD graph nodes; {len(weak_data['runs'])} run(s)",
"process_conditioned": True}
# UV face logits, then nearest valid face-grid sample back to CAD nodes.
brep = brep_input(); raw_brep = np.load(ROOT / "data_processed/brep/faces_uv.npz", allow_pickle=False)
face_pos = np.asarray(raw_brep["pos_mm"], np.float32); face_valid_grid = np.asarray(raw_brep["trim_valid"], bool)
valid_pos = face_pos[face_valid_grid]; valid_target = nearest_target(valid_pos, cad_pos, target_soft)
face_index = np.broadcast_to(np.arange(len(face_pos))[:, None, None], face_valid_grid.shape)[face_valid_grid]
face_target = np.zeros(len(face_pos), np.float32); face_valid = np.zeros(len(face_pos), bool)
np.maximum.at(face_target, face_index, valid_target); face_valid[np.unique(face_index)] = True
grids = tensor(brep["grids"]); adjacency = tensor(brep["adjacency"], torch.long); uv_model = UVNetFaceModel(grids.shape[1])
probability, report = fit(uv_model, lambda: uv_model(grids, adjacency)[:, 0], tensor(face_target), tensor(face_valid, torch.bool), 100)
checkpoint = weak_fit_checkpoint(uv_model, "UV-Net")
uv_distance, nearest_grid = cKDTree(valid_pos).query(cad_pos, workers=-1)
fields["UV-Net"] = probability[face_index[nearest_grid]]
display_valid["UV-Net"] = uv_distance <= 10
reports["UV-Net"] = report | {"checkpoint": checkpoint, "checkpoint_scope": f"{RUN} weak-label diagnostic fit",
"native_domain": "6,564 B-Rep faces", "process_conditioned": False,
"mapping_method": "nearest valid B-Rep UV sample; nodes beyond 10 mm grayed",
"direct_node_coverage": float((uv_distance <= 10).mean()), "fill_distance_p95_mm": float(np.quantile(uv_distance, .95)),
"fill_distance_max_mm": float(uv_distance.max())}
# A run-level model cannot localize; its scalar is deliberately uniform on CAD.
tab = tabular_input(); tab_x = tensor(np.concatenate([tab["x"], tab["mask"].astype(np.float32)], 1)[run_index:run_index+1])
prevalence = np.asarray([target_binary.mean()], np.float32); tab_model = TabularSurrogate(tab_x.shape[1])
optimizer = torch.optim.AdamW(tab_model.parameters(), lr=3e-3); history = []
for _ in range(120):
optimizer.zero_grad(); pred = torch.sigmoid(tab_model(tab_x)[:, 0]); loss = F.mse_loss(pred, tensor(prevalence)); loss.backward(); optimizer.step(); history.append(float(loss.detach()))
with torch.no_grad(): scalar = float(torch.sigmoid(tab_model(tab_x))[0, 0])
checkpoint = weak_fit_checkpoint(tab_model, "Tabular MLP")
fields["Tabular MLP"] = np.full(len(cad_pos), scalar, np.float32)
display_valid["Tabular MLP"] = np.ones(len(cad_pos), bool)
reports["Tabular MLP"] = {"checkpoint": checkpoint, "checkpoint_scope": f"{RUN} weak-label diagnostic fit",
"loss_initial": history[0], "loss_final": min(history), "selected_epoch": int(np.argmin(history)+1),
"epochs": len(history), "training_seconds": None, "native_domain": "one scalar repeated over CAD",
"process_conditioned": True, "uniform_field": True}
return cad_pos, target_binary, fields, reports, display_valid, sparse_status
def plot_field(path: Path, cad_pos: np.ndarray, field: np.ndarray, title: str, subtitle: str = "", valid: np.ndarray | None = None) -> None:
fig, ax = plt.subplots(figsize=(10, 5), constrained_layout=True)
if valid is None: valid = np.ones(len(cad_pos), bool)
ax.scatter(cad_pos[~valid, 0], cad_pos[~valid, 2], c="#cfd4dc", s=4, rasterized=True)
draw = ax.scatter(cad_pos[valid, 0], cad_pos[valid, 2], c=field[valid], s=5, cmap="magma", vmin=0, vmax=1, rasterized=True)
ax.set_aspect("equal"); ax.set_xlabel("CAD X (mm)"); ax.set_ylabel("CAD Z (mm)"); ax.set_title(title + (f"\n{subtitle}" if subtitle else ""))
fig.colorbar(draw, ax=ax, label="mapped risk probability")
fig.savefig(path, dpi=150); plt.close(fig)
def write_results(cad_pos: np.ndarray, target_binary: np.ndarray, fields: dict[str, np.ndarray], reports: dict[str, dict], display_valid: dict[str, np.ndarray], sparse_status: dict) -> None:
OUT.mkdir(parents=True, exist_ok=True)
results = {}
target_float = target_binary.astype(np.float32)
plot_field(OUT / "target.png", cad_pos, target_float, f"{RUN}: weak 10 mm CAD target", f"{int(target_binary.sum())} positive nodes; unvalidated training label")
order = ["2D CNN/U-Net", "FNO2d", "Dense 3D U-Net"]
if "Sparse 3D U-Net" in fields:
order.append("Sparse 3D U-Net")
order += ["PointNet", "DGCNN", "Mesh GNN", "Risk GNN", "UV-Net", "Tabular MLP"]
for name in order:
value = np.clip(fields[name], 0, 1).astype(np.float32)
score = metrics(value, target_binary)
slug = name.lower().replace("/", "-").replace(" ", "-")
plot_field(OUT / f"{slug}.png", cad_pos, value, f"{RUN}: {name} mapped to CAD",
f"training AP {score['average_precision']:.3f} · IoU@0.5 {score['threshold_iou']:.3f} · gray = mapping distance >10 mm", display_valid[name])
results[name] = reports[name] | score | {"image": f"{slug}.png", "display_coverage": float(display_valid[name].mean())}
np.savez_compressed(OUT / "cad_risk_fields.npz", cad_pos_mm=cad_pos, weak_target=target_float,
**{name.lower().replace("/", "_").replace(" ", "_"): fields[name] for name in order},
**{name.lower().replace("/", "_").replace(" ", "_")+"__display_valid": display_valid[name] for name in order})
panels = [("Weak target", target_float)] + [(name, fields[name]) for name in order]
if "Sparse 3D U-Net" not in fields:
panels.append(("Sparse U-Net unavailable", None))
cols = 4; rows = math.ceil(len(panels) / cols)
fig, axes = plt.subplots(rows, cols, figsize=(18, 4.5 * rows), constrained_layout=True); axes = np.asarray(axes).ravel()
for ax, (name, value) in zip(axes, panels):
if value is None:
ax.text(.5, .5, "MinkowskiEngine\nunavailable", ha="center", va="center", transform=ax.transAxes, fontsize=14); ax.axis("off"); continue
valid = np.ones(len(cad_pos), bool) if name not in display_valid else display_valid[name]
ax.scatter(cad_pos[~valid, 0], cad_pos[~valid, 2], c="#cfd4dc", s=2, rasterized=True)
draw = ax.scatter(cad_pos[valid, 0], cad_pos[valid, 2], c=np.asarray(value)[valid], s=2.5, cmap="magma", vmin=0, vmax=1, rasterized=True)
ax.set_aspect("equal"); ax.set_xlabel("X (mm)"); ax.set_ylabel("Z (mm)")
if name in results:
ax.set_title(f"{name}\nAP {results[name]['average_precision']:.3f} · IoU {results[name]['threshold_iou']:.3f}")
else: ax.set_title(name)
for ax in axes[len(panels):]: ax.axis("off")
fig.colorbar(draw, ax=axes.tolist(), label="mapped risk probability", fraction=.016, pad=.01)
fig.suptitle(f"{RUN} training-set CAD risk comparison — common nodes/view/scale; weak unvalidated targets")
fig.savefig(OUT / "comparison.png", dpi=160); plt.close(fig)
report = {"status": "weak_training_cad_mapping_comparison", "run": RUN, "production_valid": False,
"generalization_measured": False, "common_cad_nodes": int(len(cad_pos)),
"target": "10 mm weak geodesic mask from supplied CSV annotations",
"comparison_warning": "Every metric is evaluated on the training run. Native grids and transfer resolution differ. Do not rank production accuracy.",
"models": results, "sparse_3d_unet": sparse_status}
(OUT / "report.json").write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
rows_html = "".join(
f"| {name} | {results[name]['loss_initial']:.4g} → {results[name]['loss_final']:.4g} | "
f"{results[name]['average_precision']:.3f} | {results[name]['threshold_iou']:.3f} | "
f"{results[name]['threshold_dice']:.3f} | {results[name]['top_k_recall']:.3f} | {results[name]['display_coverage']:.1%} | "
f"{'yes' if results[name]['process_conditioned'] else 'no'} | CAD field |
|---|
"
for name in order)
sparse_html = "" if "Sparse 3D U-Net" in fields else (
f'| Sparse 3D U-Net | Unavailable: {sparse_status["reason"]} |
|---|
')
cards = "".join(f'{name}
' for name in order)
nav_href = "../dashboard/index.html" if (OUT.parent / "dashboard/index.html").exists() else "../../../../../results/ME1a_DT_comparison/index.html"
nav_label = "model structures and inputs" if (OUT.parent / "dashboard/index.html").exists() else "ME1a + AC28 results"
(OUT / "index.html").write_text(f"""All-model CAD prediction comparison
← {nav_label} · JSON report · mapped fields NPZ
All models mapped onto the same CAD
Training-set diagnostic only. Run {RUN}, CSV-seeded weak target and unvalidated registration. AP/IoU/Dice below measure fit to this same run, not generalization. Native resolution and transfer back to CAD differ by model. Gray CAD nodes are more than 10 mm from direct mapped support.
| Model | Native loss | Training AP | IoU@0.5 | Dice@0.5 | Top-K recall | Display coverage | Process | View |
|---|
{rows_html}{sparse_html}
{cards}
""", encoding="utf-8")
def main() -> None:
cad_pos, target, fields, reports, display_valid, sparse_status = train_all()
write_results(cad_pos, target, fields, reports, display_valid, sparse_status)
print(OUT / "index.html")
if __name__ == "__main__": main()