#!/usr/bin/env python3 """Weak training diagnostic at every AC28 and ME1a mesh-graph node budget.""" from __future__ import annotations import csv import argparse import json import sys 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 sys.path.insert(0,str(Path(__file__).resolve().parents[1])) from match_targets import geodesic_distance from models.architectures import MeshGraphProcessor from models.compare_cad_predictions import fit, metrics from models.data import fixed_physical_normalize, split_embedded_units from models.training import seed_everything PROJECT=Path(__file__).resolve().parents[1] PARTS={"AC28":(PROJECT,"18h01"),"ME1a_DT":(PROJECT/"experiments/ME1a_DT","14h06")} OUT=PROJECT/"results/ME1a_DT_comparison/graph_fit" BUDGETS=(12000,6000,5000,3000,1500,750,300,100) def process_vector(root: Path, run: str) -> torch.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)) labels=np.r_[names,d["geometry_feature_names"].astype(str)] unit_labels=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) normalized=fixed_physical_normalize(np.where(valid,raw,0),labels,unit_labels) return torch.tensor(np.concatenate([normalized,valid.astype(np.float32)],1)) def main() -> None: parser=argparse.ArgumentParser(description=__doc__) parser.add_argument("--budgets",type=int,nargs="+",default=list(BUDGETS)) parser.add_argument("--append",action="store_true",help="merge these budgets with prior report records") args=parser.parse_args() OUT.mkdir(parents=True,exist_ok=True) torch.set_num_threads(4) seed_everything() records=[] for part,(root,run) in PARTS.items(): paired=root/f"data_processed/paired/{run}_cad_ply_targets.npz" with np.load(paired,allow_pickle=False) as d: fine_pos=np.asarray(d["cad_pos_mm"],np.float32) fine_mask=np.asarray(d["weak_defect_mask"][:,1],bool) anchors=np.asarray(d["defect_anchor_cad_mm"],np.float32).reshape(-1,3) condition=process_vector(root,run) for budget in args.budgets: path=root/f"data_processed/graphs/graph_{budget}.npz" with np.load(path,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) names=np.asarray(d["feature_names"]).astype(str) edge=np.asarray(d["edge_index"],np.int64) edge_attr=np.asarray(d["edge_attr"],np.float32).copy() edge_attr[:,0]/=1000 node_values=fixed_physical_normalize(np.where(valid,raw,0),names) node=torch.tensor(np.concatenate([node_values,valid.astype(np.float32)],1)) edge_t=torch.tensor(edge,dtype=torch.long) attr_t=torch.tensor(edge_attr) batch=torch.zeros(len(pos),dtype=torch.long) tree=cKDTree(pos) anchor_dist,seed_node=tree.query(anchors,workers=-1) geodesic=np.min(np.stack([geodesic_distance(len(pos),edge, np.asarray(edge_attr[:,0]*1000,np.float32),int(i)) for i in seed_node]),axis=0) native_soft=np.exp(-.5*(geodesic/10)**2).astype(np.float32) native_mask=geodesic<=10 fine_to_coarse=tree.query(fine_pos,workers=-1)[1] transfer=native_mask[fine_to_coarse] label_iou=int((transfer&fine_mask).sum())/max(1,int((transfer|fine_mask).sum())) model=MeshGraphProcessor(node.shape[1],attr_t.shape[1],condition.shape[1],width=48,steps=3) prediction,history=fit(model,lambda:model(node,edge_t,attr_t,condition,batch)[:,0], torch.tensor(native_soft),torch.ones(len(pos),dtype=torch.bool),80) native_metrics=metrics(prediction,native_mask) transferred=prediction[fine_to_coarse] fine_metrics=metrics(transferred,fine_mask) checkpoint=OUT/f"{part}_{budget}.pt" torch.save({"state_dict":model.state_dict(),"part":part,"requested_nodes":budget, "actual_nodes":len(pos),"production_valid":False, "generalization_measured":False,"target":"diagnostic weak CSV seeds"},checkpoint) records.append({"part":part,"requested_nodes":budget,"actual_nodes":len(pos), "anchor_nearest_node_median_mm":float(np.median(anchor_dist)), "anchor_nearest_node_max_mm":float(anchor_dist.max()), "native_positive_nodes":int(native_mask.sum()), "weak_label_transfer_iou_on_finest_graph":label_iou, "training_loss_initial":history["loss_initial"],"training_loss_final":history["loss_final"], "native_training_iou":native_metrics["threshold_iou"], "mapped_training_iou_on_finest_graph":fine_metrics["threshold_iou"], "mapped_training_ap_on_finest_graph":fine_metrics["average_precision"], "checkpoint":str(checkpoint.relative_to(PROJECT))}) print(part,budget,len(pos),"label transfer IoU",round(label_iou,3), "mapped training IoU",round(fine_metrics["threshold_iou"],3),flush=True) if args.append and (OUT/"report.json").exists(): previous=json.loads((OUT/"report.json").read_text())["records"] changed={(r["part"],r["requested_nodes"]) for r in records} records=[r for r in previous if (r["part"],r["requested_nodes"]) not in changed]+records records.sort(key=lambda r:(list(PARTS).index(r["part"]),BUDGETS.index(r["requested_nodes"]))) (OUT/"report.json").write_text(json.dumps({"status":"weak_training_resolution_diagnostic", "epochs_per_fit":80,"production_valid":False,"generalization_measured":False, "label_policy":"Each graph receives a fresh 10 mm geodesic weak mask from the same CSV-derived anchors; transferred labels and predictions are compared to the finest graph. All metrics are same-run diagnostics.", "records":records},indent=2)) with (OUT/"report.csv").open("w",newline="") as fh: writer=csv.DictWriter(fh,fieldnames=records[0].keys());writer.writeheader();writer.writerows(records) fig,axes=plt.subplots(1,3,figsize=(14,4),constrained_layout=True) for part in PARTS: part_rows=[r for r in records if r["part"]==part] x=[r["actual_nodes"] for r in part_rows] for ax,key,label in zip(axes,("weak_label_transfer_iou_on_finest_graph", "mapped_training_iou_on_finest_graph","anchor_nearest_node_max_mm"), ("Weak-label transfer IoU","Mapped training IoU","Farthest CSV seed to graph node (mm)")): ax.plot(x,[r[key] for r in part_rows],"o-",label=part) ax.set_xscale("log");ax.invert_xaxis();ax.set_xlabel("Actual graph nodes");ax.set_ylabel(label);ax.grid(alpha=.25) axes[0].legend();fig.suptitle("Same-run weak-label fits; changed graph budgets alter target support") fig.savefig(OUT/"resolution_training.png",dpi=160);plt.close(fig) if __name__=="__main__":main()