feat: implement reproducible PROTEAN baseline pipeline

tests/test_partition.py

@@ -0,0 +1,126 @@
+"""Tests for deterministic Dirichlet partitioning and manifest reload."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from protean_research.data import manifest as manifest_mod
+from protean_research.data import partition as partition_mod
+from protean_research.data import rare_unseen, splits
+
+
+def make_labels(num_classes=5, per_class=40):
+    return np.repeat(np.arange(num_classes), per_class).astype(np.int64)
+
+
+def test_partition_same_seed_identical():
+    labels = make_labels()
+    idx = partition_mod.class_index_map(labels, 5)
+    a = partition_mod.dirichlet_partition(idx, 4, 0.3, rng_seed=123)
+    b = partition_mod.dirichlet_partition(idx, 4, 0.3, rng_seed=123)
+    for x, y in zip(a, b):
+        assert np.array_equal(x, y)
+
+
+def test_partition_different_alpha_differs():
+    labels = make_labels()
+    idx = partition_mod.class_index_map(labels, 5)
+    a = partition_mod.dirichlet_partition(idx, 4, 0.1, rng_seed=7)
+    b = partition_mod.dirichlet_partition(idx, 4, 5.0, rng_seed=7)
+    assert partition_mod.partition_hash(a) != partition_mod.partition_hash(b)
+
+
+def test_every_sample_assigned_exactly_once():
+    labels = make_labels()
+    idx = partition_mod.class_index_map(labels, 5)
+    assignments = partition_mod.dirichlet_partition(idx, 4, 0.4, rng_seed=99)
+    union = np.sort(np.concatenate(assignments))
+    assert np.array_equal(union, np.arange(len(labels)))
+    assert len(union) == len(set(union.tolist()))
+
+
+def test_support_matrix_matches_assignments():
+    labels = make_labels()
+    idx = partition_mod.class_index_map(labels, 5)
+    assignments = partition_mod.dirichlet_partition(idx, 4, 0.4, rng_seed=99)
+    support = partition_mod.support_matrix(assignments, labels, 5)
+    assert support.shape == (4, 5)
+    assert int(support.sum()) == len(labels)
+    # Column sums equal per-class counts.
+    assert np.array_equal(support.sum(axis=0), np.bincount(labels, minlength=5))
+
+
+def test_manifest_save_reload_identity(tmp_path):
+    labels = make_labels()
+    split = splits.stratified_split(labels, "toy", seed=0)
+    labels_train = labels[split.train]
+    idx = partition_mod.class_index_map(labels_train, 5)
+    rng_seed = partition_mod.partition_rng_seed("toy", 0.5, 0)
+    assignments = partition_mod.dirichlet_partition(idx, 4, 0.5, rng_seed)
+    manifest = manifest_mod.build_manifest(
+        dataset_key="toy",
+        dataset_name="Toy",
+        alpha=0.5,
+        seed=0,
+        num_clients=4,
+        class_names=[f"c{i}" for i in range(5)],
+        benign_class="c0",
+        labels=labels,
+        split=split,
+        assignments=assignments,
+        rng_seed=rng_seed,
+        train_fraction=0.7,
+        validation_fraction=0.1,
+        test_fraction=0.2,
+    )
+    path = tmp_path / "toy__alpha0.5__seed0.json"
+    manifest_mod.save_manifest(path, manifest)
+    loaded = manifest_mod.load_manifest(path)
+    split_reload, assignments_reload, support_reload = manifest_mod.reproduce_train_assignments(
+        loaded, labels
+    )
+    assert np.array_equal(split_reload.train, split.train)
+    assert partition_mod.partition_hash(assignments_reload) == manifest["partition_hash"]
+    assert np.array_equal(support_reload, np.asarray(manifest["support"]))
+
+
+def test_manifest_detects_corruption(tmp_path):
+    labels = make_labels()
+    split = splits.stratified_split(labels, "toy", seed=0)
+    labels_train = labels[split.train]
+    idx = partition_mod.class_index_map(labels_train, 5)
+    rng_seed = partition_mod.partition_rng_seed("toy", 0.5, 0)
+    assignments = partition_mod.dirichlet_partition(idx, 4, 0.5, rng_seed)
+    manifest = manifest_mod.build_manifest(
+        dataset_key="toy", dataset_name="Toy", alpha=0.5, seed=0, num_clients=4,
+        class_names=[f"c{i}" for i in range(5)], benign_class="c0", labels=labels,
+        split=split, assignments=assignments, rng_seed=rng_seed,
+        train_fraction=0.7, validation_fraction=0.1, test_fraction=0.2,
+    )
+    manifest["rng_seed"] = int(manifest["rng_seed"]) + 1
+    with pytest.raises(ValueError):
+        manifest_mod.reproduce_train_assignments(manifest, labels)
+
+
+def test_rare_unseen_semantics():
+    # Client 0: benign + classes 1,2,3 with supports 10, 5, 2 -> rare = {3,2}.
+    support = np.array(
+        [
+            [50, 10, 5, 2, 0],
+            [50, 1, 0, 0, 0],
+            [50, 0, 0, 0, 0],
+        ]
+    )
+    rows = rare_unseen.rare_unseen_rows(support, ["b", "c1", "c2", "c3", "c4"], 0, "toy", 0.5, 0)
+    client0 = [r for r in rows if r["client_id"] == 0]
+    rare = {r["class_id"] for r in client0 if r["rare_flag"]}
+    unseen = {r["class_id"] for r in client0 if r["unseen_flag"]}
+    assert rare == {2, 3}
+    assert unseen == {4}
+    # Benign is never rare or unseen.
+    assert not any(r["rare_flag"] or r["unseen_flag"] for r in rows if r["class_id"] == 0)
+    # Client with a single attack class uses that one.
+    assert rare_unseen.rare_class_ids_for_client([50, 1, 0, 0, 0], 0) == [1]
+    # Client with no attack class contributes nothing to the rare aggregate.
+    assert rare_unseen.rare_class_ids_for_client([50, 0, 0, 0, 0], 0) == []