feat: implement support-weighted prototype aggregation

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,