#!/usr/bin/env python """Baseline PROTEAN smoke/run entry point. Runs the frozen baseline pipeline end to end for one ``(dataset, alpha, seed)``: load -> clean -> split -> preprocess -> Dirichlet partition -> federated training -> local prototypes -> server aggregation -> global redistribution -> evaluation -> structured artifacts. This runner defaults to a small configuration suitable for a smoke test. A full scientific sweep must be launched separately once the baseline is verified. Usage (smoke):: python scripts/run_baseline.py --dataset xiiotid --alpha 0.5 --seed 0 \ --max-rows 4000 --clients 4 --rounds 2 --local-epochs 1 """ 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.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.data.characterization import characterize_support # noqa: E402 from protean_research.data.pipeline import ( # noqa: E402 create_partition, prepare_dataset, preprocessing_content_hash, ) from protean_research.evaluation.metrics import ( # noqa: E402 classification_metrics, rare_class_metrics, ) from protean_research.federated import TrainingConfig, run_federated # noqa: E402 from protean_research.models.protean_cnn import extract_embeddings # noqa: E402 from protean_research.models.prototypes import nearest_prototype_predict # noqa: E402 from protean_research.utils.config import ( # noqa: E402 expand_path, load_datasets_config, load_yaml, ) def parse_args(argv=None) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--dataset", required=True) parser.add_argument("--alpha", type=float, required=True) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--clients", type=int, default=10) parser.add_argument("--rounds", type=int, default=10) parser.add_argument("--local-epochs", type=int, default=3) parser.add_argument("--batch-size", type=int, default=256) 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( "--max-rows", type=int, default=None, help="SMOKE ONLY: limit how many raw rows are read; never use for science.", ) parser.add_argument("--sample-from", type=int, default=200000, help="SMOKE ONLY: rows to read before random subsampling.") parser.add_argument("--partitions-dir", default=None) parser.add_argument("--processed-dir", default=None) parser.add_argument("--results-dir", default=None) parser.add_argument("--analysis-dir", default=None) return parser.parse_args(argv) def write_json(path: Path, payload) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as handle: json.dump(payload, handle, indent=2, sort_keys=True) handle.write("\n") 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) datasets_cfg = load_datasets_config() protocol = load_yaml(REPO_ROOT / "configs" / "experiments" / "experiment_protocol.yaml") artifacts_cfg = protocol.get("artifacts", {}) partitions_dir = expand_path(args.partitions_dir or artifacts_cfg.get("partition_manifest_dir", "artifacts/partitions")) processed_dir = expand_path(args.processed_dir or artifacts_cfg.get("processed_dir", "artifacts/processed")) analysis_dir = expand_path(args.analysis_dir or artifacts_cfg.get("analysis_dir", "artifacts/analysis")) results_root = expand_path(args.results_dir or artifacts_cfg.get("results_dir", "results")) # Smoke runs must never pollute the tracked partition/analysis directories. if args.max_rows is not None: if args.partitions_dir is None: partitions_dir = processed_dir / "smoke" / "partitions" if args.analysis_dir is None: analysis_dir = processed_dir / "smoke" / "analysis" if args.results_dir is None: results_root = results_root / "smoke" if args.max_rows is not None: from protean_research.data import loading as loading_mod import numpy as np spec = loading_mod.load_spec(args.dataset, datasets_cfg) pool = loading_mod.read_raw(spec, nrows=args.sample_from) rng = np.random.default_rng(args.seed) take = min(args.max_rows, len(pool)) chosen = np.sort(rng.choice(len(pool), size=take, replace=False)) pool = pool.iloc[chosen].reset_index(drop=True) bundle = prepare_dataset( args.dataset, seed=args.seed, datasets_cfg=datasets_cfg, raw_frame=pool ) else: bundle = prepare_dataset(args.dataset, seed=args.seed, datasets_cfg=datasets_cfg) # ---- partition (create or reload) ------------------------------------- pre_hash = preprocessing_content_hash(bundle.preprocessor) manifest_file = manifest_mod.manifest_path(partitions_dir, args.dataset, args.alpha, args.seed) if manifest_file.exists(): manifest = manifest_mod.load_manifest(manifest_file) split, assignments, support = manifest_mod.reproduce_train_assignments(manifest, bundle.labels) # Reuse the manifest's split for pairing; verify it matches the bundle. bundle.split = split else: manifest, assignments, support = create_partition( bundle, alpha=args.alpha, seed=args.seed, num_clients=args.clients, partition_dir=partitions_dir, preprocessing_hash=pre_hash, ) manifest_reload_ok = True try: manifest_mod.reproduce_train_assignments(manifest, bundle.labels) except Exception: # pragma: no cover - surfaced in metadata manifest_reload_ok = False labels = bundle.labels train_features = bundle.features[bundle.split.train] train_labels = labels[bundle.split.train] val_features = bundle.features[bundle.split.validation] val_labels = labels[bundle.split.validation] test_features = bundle.features[bundle.split.test] test_labels = labels[bundle.split.test] benign_index = bundle.cleaned.benign_index per_client_rare = rare_unseen.rare_class_ids(support, benign_index) union_rare = sorted({c for row in per_client_rare for c in row}) cfg = TrainingConfig( rounds=args.rounds, local_epochs=args.local_epochs, batch_size=args.batch_size, learning_rate=args.lr, mu=args.mu, lam=args.lam, num_clients=args.clients, ) result = run_federated( train_features=train_features, train_labels=train_labels, assignments=assignments, val_features=val_features, val_labels=val_labels, num_classes=bundle.num_classes, input_dim=bundle.features.shape[1], config=cfg, init_seed_parts=(args.dataset, f"seed={args.seed}"), rare_class_ids=union_rare, ) model = result["model"] global_prototypes = result["global_prototypes"] # ---- global test evaluation ------------------------------------------ test_embeddings = extract_embeddings(model, test_features) test_pred = nearest_prototype_predict(test_embeddings, global_prototypes) test_metrics = classification_metrics(test_labels, test_pred, bundle.num_classes) test_rare = rare_class_metrics(test_labels, test_pred, bundle.num_classes, union_rare) # ---- per-client evaluation ------------------------------------------- eval_assignments, eval_rng = partition_mod.partition_eval_split( test_labels, args.clients, args.alpha, args.dataset, args.seed, "test" ) per_client_rows = [] for client_id in range(args.clients): idx = eval_assignments[client_id] if len(idx) == 0: per_client_rows.append({ "client_id": client_id, "n_samples": 0, "accuracy": None, "macro_f1": None, "rare_class_macro_f1": None, "rare_class_recall": None, }) continue emb = extract_embeddings(model, test_features[idx]) pred = nearest_prototype_predict(emb, global_prototypes) metrics = classification_metrics(test_labels[idx], pred, bundle.num_classes) rare_ids = per_client_rare[client_id] if rare_ids: rare = rare_class_metrics(test_labels[idx], pred, bundle.num_classes, rare_ids) else: rare = {"rare_class_macro_f1": None, "rare_class_recall": None} per_client_rows.append({ "client_id": client_id, "n_samples": int(len(idx)), "accuracy": metrics["accuracy"], "macro_f1": metrics["macro_f1"], "rare_class_macro_f1": rare["rare_class_macro_f1"], "rare_class_recall": rare["rare_class_recall"], }) # ---- structured artifacts -------------------------------------------- run_name = f"{args.dataset}__alpha{args.alpha:g}__seed{args.seed}__original_protean" run_dir = results_root / run_name run_dir.mkdir(parents=True, exist_ok=True) metadata = { "run_name": run_name, "dataset": args.dataset, "alpha": args.alpha, "seed": args.seed, "num_clients": args.clients, "num_classes": bundle.num_classes, "class_names": bundle.class_names, "benign_class": bundle.spec.benign_label, "aggregation_mode": "original_protean", "training": { "rounds": args.rounds, "local_epochs": args.local_epochs, "batch_size": args.batch_size, "learning_rate": args.lr, "mu": args.mu, "lambda": args.lam, }, "smoke": args.max_rows is not None, "max_rows": args.max_rows, "manifest_path": str(manifest_file), "manifest_content_hash": manifest_mod.manifest_content_hash(manifest), "preprocessing_hash": pre_hash, "manifest_reload_ok": manifest_reload_ok, "repo_root": str(REPO_ROOT), } write_json(run_dir / "metadata.json", metadata) write_json(run_dir / "history.json", result["history"]) write_json(run_dir / "test_global_metrics.json", { "accuracy": test_metrics["accuracy"], "macro_f1": test_metrics["macro_f1"], "weighted_f1": test_metrics["weighted_f1"], **test_rare, }) write_json(run_dir / "test_confusion_matrix.json", test_metrics["confusion_matrix"]) write_csv(run_dir / "test_per_class.csv", [ {**row, "class_name": bundle.class_names[row["class_id"]]} for row in test_metrics["per_class"] ]) write_csv(run_dir / "per_client.csv", per_client_rows) write_csv(run_dir / "train_support_Nij.csv", [ {"client_id": i, **{bundle.class_names[c]: int(support[i, c]) for c in range(bundle.num_classes)}} for i in range(args.clients) ]) write_csv(run_dir / "characterization.csv", characterize_support(support, bundle.class_names)) write_csv(run_dir / "rare_unseen.csv", rare_unseen.rare_unseen_rows( support, bundle.class_names, benign_index, args.dataset, args.alpha, args.seed )) processed_dir.mkdir(parents=True, exist_ok=True) bundle.preprocessor.save(processed_dir / f"{args.dataset}__seed{args.seed}__preprocessor.json") snapshot = result["snapshot"] if snapshot is not None: write_json(analysis_dir / f"{run_name}__prototype_snapshot.json", snapshot) summary = { "run": run_name, "test_accuracy": test_metrics["accuracy"], "test_macro_f1": test_metrics["macro_f1"], "test_rare_class_macro_f1": test_rare["rare_class_macro_f1"], "n_global_prototypes": len(global_prototypes), "history_last": result["history"][-1] if result["history"] else None, "results_dir": str(run_dir), } print(json.dumps(summary, indent=2, sort_keys=True)) return 0 if __name__ == "__main__": raise SystemExit(main())