125 lines
5.4 KiB
Python
125 lines
5.4 KiB
Python
#!/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())
|