"""Compare graph sizes and joint/individual weak-label GNN fits on AC28 and ME1a.
Every score is a diagnostic on provisional CSV-seeded masks. In particular,
holding out a run does not validate its registration or defect extent.
"""
from __future__ import annotations
import argparse
import csv
import html
import json
import time
from dataclasses import dataclass
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from scipy.spatial import cKDTree
import torch
from torch import Tensor, nn
import torch.nn.functional as F
from models.architectures import MeshGraphProcessor
from models.compare_cad_predictions import metrics
from models.data import fixed_physical_normalize, split_embedded_units
from models.training import seed_everything
from 전처리.match_targets import geodesic_distance
PROJECT = Path(__file__).resolve().parents[1]
OUT = PROJECT / "results/ME1a_DT_comparison/gnn_joint"
RUNS = (("AC28_16h", "AC28", PROJECT, "16h"),
("AC28_18h01", "AC28", PROJECT, "18h01"),
("ME1a_14h06", "ME1a_DT", PROJECT / "experiments/ME1a_DT", "14h06"))
@dataclass
class Sample:
name: str
part: str
node: Tensor
edge: Tensor
attr: Tensor
process: Tensor
target: Tensor
binary: np.ndarray
fine_binary: np.ndarray
fine_from_coarse: np.ndarray
transfer_iou: float
actual_nodes: int
def _process(root: Path, run: str) -> Tensor:
with np.load(root / "data_processed/tabular/process_and_geometry.npz", allow_pickle=False) as d:
idx = list(d["run_id"].astype(str)).index(run)
names, units = split_embedded_units(d["process_feature_names"].astype(str))
names = np.r_[names, d["geometry_feature_names"].astype(str)]
units = np.r_[units, np.repeat("", len(d["geometry_feature_names"]))]
raw = np.asarray(d["x"][idx:idx+1], np.float32)
valid = np.asarray(d["feature_valid"][idx:idx+1], bool)
scaled = fixed_physical_normalize(np.where(valid, raw, 0), names, units)
return torch.tensor(np.concatenate([scaled, valid.astype(np.float32)], 1)[0])
def load_samples(budget: int) -> list[Sample]:
out = []
for name, part, root, run in RUNS:
with np.load(root / f"data_processed/paired/{run}_cad_ply_targets.npz", allow_pickle=False) as d:
fine_pos = np.asarray(d["cad_pos_mm"], np.float32)
fine_binary = np.asarray(d["weak_defect_mask"][:, 1], bool)
anchors = np.asarray(d["defect_anchor_cad_mm"], np.float32).reshape(-1, 3)
with np.load(root / f"data_processed/graphs/graph_{budget}.npz", allow_pickle=False) as d:
pos = np.asarray(d["pos_mm"], np.float32)
raw = np.asarray(d["x"], np.float32)
valid = np.asarray(d["feature_valid"], bool)
feature_names = d["feature_names"].astype(str)
edge = np.asarray(d["edge_index"], np.int64)
attr = np.asarray(d["edge_attr"], np.float32).copy()
lengths_mm = attr[:, 0].copy()
attr[:, 0] /= 1000.0
node = fixed_physical_normalize(np.where(valid, raw, 0), feature_names)
node = np.concatenate([node, valid.astype(np.float32)], 1)
tree = cKDTree(pos)
seeds = np.unique(tree.query(anchors, workers=-1)[1])
distance = np.min(np.stack([geodesic_distance(len(pos), edge, lengths_mm, int(i))
for i in seeds]), axis=0)
target = np.exp(-.5 * (distance / 10.0) ** 2).astype(np.float32)
binary = distance <= 10.0
fine_from_coarse = tree.query(fine_pos, workers=-1)[1]
transferred = binary[fine_from_coarse]
transfer_iou = float((transferred & fine_binary).sum() / max(1, (transferred | fine_binary).sum()))
out.append(Sample(name, part, torch.tensor(node), torch.tensor(edge), torch.tensor(attr),
_process(root, run), torch.tensor(target), binary, fine_binary,
fine_from_coarse, transfer_iou, len(pos)))
return out
class ResidualGeoGNN(nn.Module):
"""Short residual edge messages plus multiscale positions for small regions."""
def __init__(self, node_dim: int, edge_dim: int, process_dim: int, width: int = 48, steps: int = 2):
super().__init__()
self.process = nn.Sequential(nn.Linear(process_dim, 24), nn.GELU(), nn.Linear(24, 24))
self.node = nn.Sequential(nn.Linear(node_dim + 24 + 24, width), nn.GELU(), nn.LayerNorm(width))
self.msg = nn.ModuleList([nn.Sequential(nn.Linear(width + edge_dim, width), nn.GELU(),
nn.Linear(width, width)) for _ in range(steps)])
self.update = nn.ModuleList([nn.Sequential(nn.Linear(width * 2, width), nn.GELU(),
nn.Linear(width, width)) for _ in range(steps)])
self.norm = nn.ModuleList([nn.LayerNorm(width) for _ in range(steps)])
self.head = nn.Sequential(nn.Linear(width * 2, width), nn.GELU(), nn.Linear(width, 1))
def forward(self, node_x: Tensor, edge_index: Tensor, edge_attr: Tensor,
global_x: Tensor, batch: Tensor) -> Tensor:
pos = node_x[:, :3]
angles = pos[:, :, None] * pos.new_tensor([np.pi, 2 * np.pi, 4 * np.pi, 8 * np.pi])
fourier = torch.cat([angles.sin(), angles.cos()], dim=2).flatten(1)
process = self.process(global_x)[batch]
h = self.node(torch.cat([node_x, fourier, process], 1))
initial = h
src, dst = edge_index
degree = torch.bincount(dst, minlength=len(h)).clamp_min(1).to(h.dtype)[:, None]
for message, update, norm in zip(self.msg, self.update, self.norm):
m = message(torch.cat([h[src], edge_attr], 1))
agg = torch.zeros_like(h).index_add_(0, dst, m) / degree
h = norm(h + update(torch.cat([h, agg], 1)))
return self.head(torch.cat([h, initial], 1)).squeeze(1)
class GeoMLP(nn.Module):
"""Control with the same positional and process information, no messages."""
def __init__(self, node_dim: int, edge_dim: int, process_dim: int, width: int = 48):
super().__init__()
self.process = nn.Sequential(nn.Linear(process_dim, 24), nn.GELU(), nn.Linear(24, 24))
self.head = nn.Sequential(nn.Linear(node_dim + 24 + 24, width), nn.GELU(),
nn.Linear(width, width), nn.GELU(), nn.Linear(width, 1))
def forward(self, node_x: Tensor, edge_index: Tensor, edge_attr: Tensor,
global_x: Tensor, batch: Tensor) -> Tensor:
pos = node_x[:, :3]
angles = pos[:, :, None] * pos.new_tensor([np.pi, 2 * np.pi, 4 * np.pi, 8 * np.pi])
fourier = torch.cat([angles.sin(), angles.cos()], dim=2).flatten(1)
return self.head(torch.cat([node_x, fourier, self.process(global_x)[batch]], 1)).squeeze(1)
def make_model(kind: str, sample: Sample) -> nn.Module:
dims = (sample.node.shape[1], sample.attr.shape[1], sample.process.shape[0])
if kind == "mesh": return MeshGraphProcessor(*dims, width=48, steps=3)
if kind == "residual_geo": return ResidualGeoGNN(*dims)
if kind == "geo_mlp": return GeoMLP(*dims)
raise ValueError(kind)
def combine(samples: list[Sample]) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, list[slice]]:
offset = 0
edges, batches, slices = [], [], []
for i, sample in enumerate(samples):
edges.append(sample.edge + offset)
batches.append(torch.full((sample.actual_nodes,), i, dtype=torch.long))
slices.append(slice(offset, offset + sample.actual_nodes))
offset += sample.actual_nodes
return (torch.cat([s.node for s in samples]), torch.cat(edges, 1),
torch.cat([s.attr for s in samples]), torch.stack([s.process for s in samples]),
torch.cat(batches), slices)
def forward(model: nn.Module, batch: tuple) -> Tensor:
output = model(*batch[:5])
return output[:, 0] if output.ndim == 2 else output
def risk_loss(logit: Tensor, samples: list[Sample], slices: list[slice], mode: str = "soft") -> Tensor:
losses = []
for sample, sl in zip(samples, slices):
target = sample.target if mode == "soft" else torch.as_tensor(sample.binary, dtype=logit.dtype)
positive = max(1, int(sample.binary.sum()))
weight = min(200.0 if mode == "soft" else 1000.0, (len(target) - positive) / positive)
bce = F.binary_cross_entropy_with_logits(logit[sl], target, pos_weight=target.new_tensor(weight))
prob = logit[sl].sigmoid()
dice = 1 - (2 * (prob * target).sum() + 1) / (prob.sum() + target.sum() + 1)
losses.append(bce + dice)
return torch.stack(losses).mean()
def train(kind: str, samples: list[Sample], epochs: int, seed: int, loss_mode: str = "soft") -> tuple[nn.Module, dict]:
torch.manual_seed(seed)
model = make_model(kind, samples[0])
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=1e-5)
batch = combine(samples)
best_loss = float("inf")
best_epoch = 0
best = None
start = time.monotonic()
for epoch in range(epochs):
model.train()
optimizer.zero_grad(set_to_none=True)
prediction = forward(model, batch)
loss = risk_loss(prediction, samples, batch[5], loss_mode)
loss.backward()
optimizer.step()
value = float(loss.detach())
if value < best_loss:
best_loss, best_epoch = value, epoch + 1
best = {k: v.detach().clone() for k, v in model.state_dict().items()}
model.load_state_dict(best)
return model, {"best_loss": best_loss, "best_epoch": best_epoch,
"seconds": round(time.monotonic() - start, 2)}
def evaluate(model: nn.Module, sample: Sample) -> dict:
model.eval()
with torch.no_grad():
probability = forward(model, combine([sample])).sigmoid().cpu().numpy()
native = metrics(probability, sample.binary)
mapped = metrics(probability[sample.fine_from_coarse], sample.fine_binary)
return {"native_ap": native["average_precision"], "native_iou": native["threshold_iou"],
"mapped_ap": mapped["average_precision"], "mapped_iou": mapped["threshold_iou"],
"mapped_dice": mapped["threshold_dice"], "mapped_top_k_recall": mapped["top_k_recall"],
"fine_prevalence": float(sample.fine_binary.mean()),
"predicted_positive_fine": mapped["predicted_positive_nodes"]}
def main() -> None:
global OUT
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--budgets", default="3000,5000,6000,12000")
parser.add_argument("--models", default="mesh,residual_geo,geo_mlp")
parser.add_argument("--epochs", type=int, default=80)
parser.add_argument("--loss", choices=("soft", "hard"), default="soft")
parser.add_argument("--out", type=Path, default=OUT)
args = parser.parse_args()
OUT = args.out.resolve()
budgets = [int(x) for x in args.budgets.split(",")]
kinds = args.models.split(",")
OUT.mkdir(parents=True, exist_ok=True)
torch.set_num_threads(4)
seed_everything()
rows = []
for budget in budgets:
samples = load_samples(budget)
by_name = {s.name: s for s in samples}
setups = [("joint", [s.name for s in samples], [s.name for s in samples])]
setups += [(f"holdout_{s.name}", [x.name for x in samples if x.name != s.name], [s.name])
for s in samples]
setups += [(f"single_{s.name}", [s.name], [s.name]) for s in samples]
setups += [("AC28_only_holdout_16h", ["AC28_18h01"], ["AC28_16h"]),
("AC28_only_holdout_18h01", ["AC28_16h"], ["AC28_18h01"])]
for kind in kinds:
for setup, train_names, test_names in setups:
model, fit = train(kind, [by_name[n] for n in train_names], args.epochs, seed=17,
loss_mode=args.loss)
checkpoint = OUT / f"{kind}_{budget}_{setup}.pt"
torch.save({"state_dict": model.state_dict(), "kind": kind, "budget": budget,
"train_runs": train_names, "loss_mode": args.loss, "weak_label_only": True,
"production_valid": False}, checkpoint)
for test_name in test_names:
sample = by_name[test_name]
row = {"budget": budget, "kind": kind, "loss_mode": args.loss, "setup": setup,
"train_runs": "+".join(train_names), "eval_run": test_name,
"seen_in_training": test_name in train_names,
"actual_nodes": sample.actual_nodes, "positive_native": int(sample.binary.sum()),
"positive_fine": int(sample.fine_binary.sum()),
"label_transfer_iou": sample.transfer_iou,
**fit, **evaluate(model, sample),
"checkpoint": str(checkpoint.relative_to(PROJECT))}
rows.append(row)
print(kind, budget, setup, test_name, "AP", round(row["mapped_ap"], 4),
"IoU", round(row["mapped_iou"], 4), "sec", fit["seconds"], flush=True)
report = {"status": "joint_and_individual_weak_label_diagnostic",
"production_valid": False, "independent_physical_runs": 3,
"label_policy": "10 mm geodesic Gaussian/area from provisional CSV anchors at each graph size",
"selection_caution": "The three runs and unvalidated labels cannot establish production accuracy; same-run fits are separate from leave-one-run-out predictions.",
"epochs_per_fit": args.epochs, "loss_mode": args.loss, "rows": rows}
(OUT / "report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
with (OUT / "report.csv").open("w", newline="") as fh:
writer = csv.DictWriter(fh, fieldnames=list(rows[0]))
writer.writeheader(); writer.writerows(rows)
make_figures(rows, budgets, kinds)
make_html(rows, budgets, kinds)
def make_figures(rows: list[dict], budgets: list[int], kinds: list[str]) -> None:
fig, axes = plt.subplots(1, 3, figsize=(15, 4.5), constrained_layout=True)
for kind in kinds:
for part, label in (("AC28_16h", "AC28 16h"), ("AC28_18h01", "AC28 18h01"),
("ME1a_14h06", "ME1a 14h06")):
selected = [next(r for r in rows if r["budget"] == budget and r["kind"] == kind
and r["setup"] == "joint" and r["eval_run"] == part) for budget in budgets]
axes[0].plot([r["actual_nodes"] for r in selected], [r["mapped_ap"] for r in selected], "o-", label=f"{kind} {label}")
selected = [next(r for r in rows if r["budget"] == budget and r["kind"] == kind
and r["setup"] == f"holdout_{part}") for budget in budgets]
axes[1].plot([r["actual_nodes"] for r in selected], [r["mapped_ap"] for r in selected], "o-", label=f"{kind} {label}")
for budget in budgets:
row = next(r for r in rows if r["budget"] == budget and r["kind"] == kinds[0] and r["eval_run"] == "ME1a_14h06")
axes[2].scatter(row["actual_nodes"], row["label_transfer_iou"], s=60, color="#d35400")
axes[2].annotate(str(budget), (row["actual_nodes"], row["label_transfer_iou"]), xytext=(3, 4), textcoords="offset points")
axes[0].set_title("Joint training: same-run weak AP")
axes[1].set_title("Leave-one-run-out weak AP")
axes[2].set_title("ME1a weak-label transfer")
for ax in axes:
ax.set_xscale("log"); ax.set_xlabel("actual graph nodes"); ax.grid(alpha=.25)
for ax in axes[:2]:ax.set_ylabel("mapped average precision")
axes[2].set_ylabel("IoU of labels mapped to finest graph")
axes[1].legend(fontsize=6, loc="best")
fig.savefig(OUT / "comparison.png", dpi=170); plt.close(fig)
def make_html(rows: list[dict], budgets: list[int], kinds: list[str]) -> None:
selected = [r for r in rows if r["setup"] == "joint" or r["setup"].startswith("holdout_")]
table = "".join(f"| {r['budget']:,} | {html.escape(r['kind'])} | "
f"{html.escape(r['setup'])} | {html.escape(r['eval_run'])} | "
f"{r['mapped_ap']:.4f} | {r['mapped_iou']:.4f} | "
f"{r['label_transfer_iou']:.4f} |
" for r in selected)
focus = "residual_geo" if "residual_geo" in kinds else kinds[0]
effect_rows = []
for budget in budgets:
for run in ("AC28_16h", "AC28_18h01", "ME1a_14h06"):
base = [r for r in rows if r["budget"] == budget and r["kind"] == focus
and r["eval_run"] == run]
find = lambda setup: next((r["mapped_ap"] for r in base if r["setup"] == setup), None)
single, joint, held = find(f"single_{run}"), find("joint"), find(f"holdout_{run}")
ac_only = next((r["mapped_ap"] for r in base if r["setup"].startswith("AC28_only_holdout")), None)
effect_rows.append(f"| {budget:,} | {run} | {single:.4f} | "
f"{joint:.4f} | {held:.4f} | "
f"{ac_only:.4f} |
" if ac_only is not None else
f"| {budget:,} | {run} | {single:.4f} | "
f"{joint:.4f} | {held:.4f} | — |
")
image = 'Where the model predicts
' if (OUT / "prediction_maps.png").exists() else ""
page = f'''
Joint AC28 and ME1a GNN diagnostics
AC28 + ME1a: GNN and node-size experiment
All labels are provisional CSV-seeded areas; registration is unvalidated. Same-run fit is separate from leave-one-run-out prediction. Three production runs cannot establish production accuracy.
Architecture and graph size
Mesh GNN uses the previous message processor. Residual Geo GNN adds short residual edge messages and position frequencies. Geo MLP has the same position and process input without graph messages. The experiment covers {len(budgets)} graph budgets and {len(kinds)} model families, using per-run balanced weak loss.
Does pooling AC28 and ME1a help?
Values below show {focus} average precision. Individual and joint values are evaluated on runs used for fitting. Held-out values evaluate an unseen run. The last column trains on the other AC28 run only, allowing a direct comparison with adding ME1a for AC28 holdouts.
{image}
Full fit and per-run metrics CSV · Metrics and checkpoints JSON · Main AC28 + ME1a report
'''
(OUT / "index.html").write_text(page, encoding="utf-8")
if __name__ == "__main__": main()