feat: implement reproducible PROTEAN baseline pipeline

scripts/generate_partitions.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python
+"""Generate deterministic Dirichlet partition manifests for all conditions.
+
+This does NOT train anything. It loads each dataset once (labels only), computes
+the frozen stratified split, generates the Dirichlet partition for every
+``(alpha, seed)`` and persists a manifest plus characterization/rare-unseen
+metadata. Training later loads these manifests; it never regenerates a partition.
+
+Usage::
+
+    python scripts/generate_partitions.py --datasets xiiotid 5gnidd \
+        --alphas 0.75 0.5 0.25 --seeds 0 1 2 --clients 10
+"""
+
+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 loading, manifest as manifest_mod, partition as partition_mod  # noqa: E402
+from protean_research.data import rare_unseen, splits as splits_mod  # noqa: E402
+from protean_research.data.characterization import characterize_support  # noqa: E402
+from protean_research.data.pipeline import (  # noqa: E402
+    TEST_FRACTION,
+    TRAIN_FRACTION,
+    VALIDATION_FRACTION,
+)
+from protean_research.utils.config import expand_path, load_datasets_config, load_yaml  # noqa: E402
+
+
+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
+    with path.open("w", encoding="utf-8", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys()))
+        writer.writeheader()
+        writer.writerows(rows)
+
+
+def main(argv=None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--datasets", nargs="+", default=["xiiotid", "5gnidd"])
+    parser.add_argument("--alphas", nargs="+", type=float, default=[0.75, 0.5, 0.25])
+    parser.add_argument("--seeds", nargs="+", type=int, default=[0, 1, 2])
+    parser.add_argument("--clients", type=int, default=10)
+    parser.add_argument("--partitions-dir", default=None)
+    parser.add_argument("--analysis-dir", default=None)
+    args = parser.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"))
+    analysis_dir = expand_path(args.analysis_dir or artifacts_cfg.get("analysis_dir", "artifacts/analysis"))
+
+    created = []
+    for dataset_key in args.datasets:
+        spec = loading.load_spec(dataset_key, datasets_cfg)
+        cleaned = loading.clean_frame(loading.read_raw(spec), spec)
+        for alpha in args.alphas:
+            for seed in args.seeds:
+                split = splits_mod.stratified_split(
+                    cleaned.labels, dataset_key, seed,
+                    train_fraction=TRAIN_FRACTION,
+                    validation_fraction=VALIDATION_FRACTION,
+                    test_fraction=TEST_FRACTION,
+                )
+                labels_train = cleaned.labels[split.train]
+                class_indices = partition_mod.class_index_map(labels_train, len(cleaned.class_names))
+                rng_seed = partition_mod.partition_rng_seed(dataset_key, alpha, seed)
+                assignments = partition_mod.dirichlet_partition(
+                    class_indices, args.clients, alpha, rng_seed
+                )
+                manifest = manifest_mod.build_manifest(
+                    dataset_key=dataset_key,
+                    dataset_name=spec.name,
+                    alpha=alpha,
+                    seed=seed,
+                    num_clients=args.clients,
+                    class_names=cleaned.class_names,
+                    benign_class=spec.benign_label,
+                    labels=cleaned.labels,
+                    split=split,
+                    assignments=assignments,
+                    rng_seed=rng_seed,
+                    train_fraction=TRAIN_FRACTION,
+                    validation_fraction=VALIDATION_FRACTION,
+                    test_fraction=TEST_FRACTION,
+                )
+                path = manifest_mod.manifest_path(partitions_dir, dataset_key, alpha, seed)
+                manifest_mod.save_manifest(path, manifest)
+
+                # Verify exact reload before declaring success.
+                manifest_mod.reproduce_train_assignments(manifest, cleaned.labels)
+
+                support = manifest["support"]
+                import numpy as np
+
+                support_arr = np.asarray(support, dtype=np.int64)
+                tag = f"{dataset_key}__alpha{alpha:g}__seed{seed}"
+                write_csv(analysis_dir / f"{tag}__characterization.csv",
+                          characterize_support(support_arr, cleaned.class_names))
+                write_csv(analysis_dir / f"{tag}__rare_unseen.csv",
+                          rare_unseen.rare_unseen_rows(
+                              support_arr, cleaned.class_names, cleaned.benign_index,
+                              dataset_key, alpha, seed))
+                created.append({"dataset": dataset_key, "alpha": alpha, "seed": seed,
+                                "path": str(path), "partition_hash": manifest["partition_hash"]})
+                print(f"created {path.name}")
+
+    print(json.dumps({"created": len(created)}, indent=2))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())