| |
| """Compare task loss/score/metric semantics for two forward artifacts.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
| from types import SimpleNamespace |
|
|
| import torch |
| from sklearn.metrics import roc_auc_score |
|
|
| from validation.artifacts import load_artifact |
|
|
|
|
| def _evaluate(artifact, *, rewrite: bool): |
| logits = torch.from_numpy(artifact.logits).to(torch.float32) |
| labels = torch.from_numpy(artifact.labels).to(torch.long) |
| weights = torch.from_numpy(artifact.weights).to(torch.float32) |
| elementwise = torch.nn.functional.cross_entropy(logits, labels, reduction="none") |
| loss = torch.zeros(()) |
| for label in torch.unique(labels): |
| mask = labels == label |
| loss = loss + (weights[mask] * elementwise[mask]).sum() / weights[mask].sum() |
| loss = loss / len(torch.unique(labels)) |
| scores = torch.softmax(logits, dim=1) |
| predictions = scores.argmax(dim=1) |
| positive = weights > 0 |
| one_hot = torch.nn.functional.one_hot(labels, num_classes=scores.shape[1]).numpy() |
| positive_numpy = positive.numpy() |
| legacy_auc = float( |
| roc_auc_score( |
| one_hot[positive_numpy], |
| scores.detach().numpy()[positive_numpy], |
| multi_class="ovr", |
| sample_weight=weights.numpy()[positive_numpy], |
| ) |
| ) |
| if rewrite: |
| from gnn4colliders.tasks import MulticlassClassificationTask |
|
|
| batch = SimpleNamespace( |
| labels=labels, |
| metadata=SimpleNamespace(weight=weights), |
| ) |
| auc = MulticlassClassificationTask().metrics(logits, batch)["roc_auc"] |
| else: |
| auc = legacy_auc |
| return { |
| "loss": float(loss), |
| "accuracy": float((predictions == labels).float().mean()), |
| "roc_auc": auc, |
| "positive_weight_events": int(positive.sum()), |
| "negative_weight_events": int((weights < 0).sum()), |
| "zero_weight_events": int((weights == 0).sum()), |
| } |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("legacy", type=Path) |
| parser.add_argument("rewrite", type=Path) |
| parser.add_argument("output", type=Path) |
| args = parser.parse_args() |
| left = _evaluate(load_artifact(args.legacy), rewrite=False) |
| right = _evaluate(load_artifact(args.rewrite), rewrite=True) |
| differences = { |
| name: abs(left[name] - right[name]) for name in ("loss", "accuracy", "roc_auc") |
| } |
| report = { |
| "schema_version": 1, |
| "legacy": left, |
| "rewrite": right, |
| "absolute_differences": differences, |
| "overall_status": "PASS" |
| if all(value <= 1e-5 for value in differences.values()) |
| else "FAIL", |
| "note": ( |
| "Legacy uses sklearn weighted OVR AUC; rewrite uses its task-owned " |
| "weighted OVR implementation." |
| ), |
| } |
| args.output.mkdir(parents=True, exist_ok=True) |
| (args.output / "task_metrics.json").write_text( |
| json.dumps(report, indent=2, sort_keys=True, allow_nan=True) + "\n" |
| ) |
| print(json.dumps({"overall_status": report["overall_status"]}, indent=2)) |
| return 0 if report["overall_status"] == "PASS" else 1 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
