127 lines
5.0 KiB
Python
127 lines
5.0 KiB
Python
"""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) == []
|