feat: implement support-weighted prototype aggregation

scripts/analyze_prototype_effect.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python
+"""Direct prototype-effect analysis from a fixed ``{C_ij, N_ij}`` snapshot.
+
+Reads a round snapshot written by ``run_baseline`` (or ``run_federated``) and
+recomputes ``Cbar_j(beta)`` for the candidate betas **without retraining**,
+reporting per-class ``|| Cbar_j(beta) - Cbar_j(0) ||`` and the normalised client
+weights. This measures the direct aggregation effect only; it is not an
+end-to-end detection result.
+
+Usage::
+
+    python scripts/analyze_prototype_effect.py \
+        --snapshot artifacts/analysis/xiiotid__alpha0.5__seed0__original_protean__prototype_snapshot.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO_ROOT / "src"))
+
+from protean_research.aggregation import BETA_CANDIDATES  # noqa: E402
+from protean_research.analysis.direct_prototype import (  # noqa: E402
+    client_weight_rows,
+    prototype_effect_rows,
+)
+
+
+def parse_args(argv=None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--snapshot", required=True)
+    parser.add_argument(
+        "--betas",
+        type=float,
+        nargs="+",
+        default=list(BETA_CANDIDATES),
+        help="beta candidates to evaluate (default: frozen candidate set).",
+    )
+    parser.add_argument(
+        "--class-names",
+        default=None,
+        help="optional comma-separated class names indexed by class id.",
+    )
+    parser.add_argument("--output-dir", default="artifacts/analysis")
+    return parser.parse_args(argv)
+
+
+def write_csv(path: Path, rows: list[dict]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if not rows:
+        path.write_text("", encoding="utf-8")
+        return
+    fieldnames = list(rows[0].keys())
+    with path.open("w", encoding="utf-8", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=fieldnames)
+        writer.writeheader()
+        writer.writerows(rows)
+
+
+def main(argv=None) -> int:
+    args = parse_args(argv)
+    snapshot_path = Path(args.snapshot).expanduser()
+    if not snapshot_path.is_absolute():
+        snapshot_path = REPO_ROOT / snapshot_path
+    with snapshot_path.open("r", encoding="utf-8") as handle:
+        snapshot = json.load(handle)
+
+    class_names = None
+    if args.class_names:
+        class_names = {
+            idx: name.strip()
+            for idx, name in enumerate(args.class_names.split(","))
+        }
+
+    effect_rows = prototype_effect_rows(snapshot, args.betas, class_names)
+    weight_rows = client_weight_rows(snapshot, args.betas, class_names)
+
+    output_dir = Path(args.output_dir).expanduser()
+    if not output_dir.is_absolute():
+        output_dir = REPO_ROOT / output_dir
+    stem = snapshot_path.stem
+    effect_path = output_dir / f"{stem}__prototype_effect.csv"
+    weight_path = output_dir / f"{stem}__prototype_weights.csv"
+    write_csv(effect_path, effect_rows)
+    write_csv(weight_path, weight_rows)
+
+    summary = {
+        "snapshot": str(snapshot_path),
+        "round": snapshot.get("round"),
+        "betas": [float(b) for b in args.betas],
+        "n_classes": len({row["class_id"] for row in effect_rows}),
+        "max_displacement": max(
+            (row["displacement_from_beta0"] for row in effect_rows), default=0.0
+        ),
+        "effect_csv": str(effect_path),
+        "weights_csv": str(weight_path),
+    }
+    print(json.dumps(summary, indent=2, sort_keys=True))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/run_baseline.py

@@ -29,6 +29,7 @@ sys.path.insert(0, str(REPO_ROOT / "src"))
 from protean_research.data import manifest as manifest_mod  # noqa: E402
 from protean_research.data import partition as partition_mod  # noqa: E402
 from protean_research.data import rare_unseen  # noqa: E402
+from protean_research.aggregation import validate_beta  # noqa: E402
 from protean_research.data.characterization import characterize_support  # noqa: E402
 from protean_research.data.pipeline import (  # noqa: E402
     create_partition,
@@ -61,6 +62,18 @@ def parse_args(argv=None) -> argparse.Namespace:
     parser.add_argument("--lr", type=float, default=1e-3)
     parser.add_argument("--mu", type=float, default=0.1)
     parser.add_argument("--lam", type=float, default=1.0)
+    parser.add_argument(
+        "--aggregation-mode",
+        choices=["original_protean", "support_weighted"],
+        default="original_protean",
+        help="server-side prototype aggregation mode (default: frozen baseline).",
+    )
+    parser.add_argument(
+        "--beta",
+        type=float,
+        default=None,
+        help="support-weighting exponent in [0, 1]; required for support_weighted.",
+    )
     parser.add_argument(
         "--max-rows",
         type=int,
@@ -97,6 +110,20 @@ def write_csv(path: Path, rows: list[dict]) -> None:
 
 def main(argv=None) -> int:
     args = parse_args(argv)
+
+    if args.aggregation_mode == "support_weighted":
+        if args.beta is None:
+            raise SystemExit(
+                "--beta is required when --aggregation-mode support_weighted"
+            )
+        beta = validate_beta(args.beta)
+    else:
+        if args.beta is not None:
+            raise SystemExit(
+                "--beta is only valid with --aggregation-mode support_weighted"
+            )
+        beta = None
+
     datasets_cfg = load_datasets_config()
     protocol = load_yaml(REPO_ROOT / "configs" / "experiments" / "experiment_protocol.yaml")
 
@@ -170,6 +197,8 @@ def main(argv=None) -> int:
         mu=args.mu,
         lam=args.lam,
         num_clients=args.clients,
+        aggregation_mode=args.aggregation_mode,
+        beta=beta,
     )
     result = run_federated(
         train_features=train_features,
@@ -224,7 +253,10 @@ def main(argv=None) -> int:
         })
 
     # ---- structured artifacts --------------------------------------------
-    run_name = f"{args.dataset}__alpha{args.alpha:g}__seed{args.seed}__original_protean"
+    mode_tag = args.aggregation_mode
+    if args.aggregation_mode == "support_weighted":
+        mode_tag = f"support_weighted__beta{beta:g}"
+    run_name = f"{args.dataset}__alpha{args.alpha:g}__seed{args.seed}__{mode_tag}"
     run_dir = results_root / run_name
     run_dir.mkdir(parents=True, exist_ok=True)
 
@@ -237,7 +269,8 @@ def main(argv=None) -> int:
         "num_classes": bundle.num_classes,
         "class_names": bundle.class_names,
         "benign_class": bundle.spec.benign_label,
-        "aggregation_mode": "original_protean",
+        "aggregation_mode": args.aggregation_mode,
+        "beta": beta,
         "training": {
             "rounds": args.rounds,
             "local_epochs": args.local_epochs,