feat: implement reproducible PROTEAN baseline pipeline

tests/test_baseline_aggregation.py

@@ -0,0 +1,76 @@
+"""Tests for the reproduced original-PROTEAN prototype aggregation.
+
+These tests cover only the baseline aggregation present in this checkpoint. The
+participant rule (``S_j`` with equal weights) is a documented reproduction
+assumption; these tests check its arithmetic and its explicit handling of
+missing classes, not its equivalence to the original method.
+"""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+
+from protean_research.aggregation import (
+    AggregationMode,
+    aggregate_class_prototypes,
+    aggregate_prototypes,
+)
+
+PROTOTYPES = [[0.0, 2.0], [2.0, 0.0], [4.0, 4.0]]
+
+
+def mean(vectors):
+    dim = len(vectors[0])
+    return [sum(v[d] for v in vectors) / len(vectors) for d in range(dim)]
+
+
+def test_original_protean_is_equal_mean_over_contributors():
+    out = aggregate_prototypes(PROTOTYPES, AggregationMode.ORIGINAL_PROTEAN)
+    assert out == pytest.approx(mean(PROTOTYPES))
+
+
+def test_original_protean_is_default_mode():
+    assert aggregate_prototypes(PROTOTYPES) == pytest.approx(mean(PROTOTYPES))
+
+
+def test_single_contributor_returns_itself():
+    out = aggregate_prototypes([[1.0, -2.0, 3.0]])
+    assert out == pytest.approx([1.0, -2.0, 3.0])
+
+
+def test_output_is_finite_for_valid_inputs():
+    out = aggregate_prototypes(PROTOTYPES)
+    assert all(math.isfinite(v) for v in out)
+
+
+def test_aggregate_class_prototypes_builds_s_j_and_drops_missing():
+    contributions = {
+        0: [([1.0, 1.0], 5), (None, 0), ([3.0, 3.0], 5)],
+        1: [([2.0, 2.0], 0)],  # no positive support -> class omitted entirely
+    }
+    result = aggregate_class_prototypes(contributions)
+    assert set(result) == {0}
+    assert result[0] == pytest.approx([2.0, 2.0])
+
+
+def test_missing_class_is_not_fabricated():
+    # A client with N_ij == 0 contributes no prototype and must not enter S_j.
+    result = aggregate_class_prototypes({0: [(None, 0), ([1.0], 0)]})
+    assert result == {}
+
+
+def test_proposed_mode_is_not_available():
+    # The support-weighted mode belongs to the later beta/controls phase.
+    with pytest.raises(ValueError):
+        aggregate_prototypes(PROTOTYPES, "support_weighted")
+
+
+def test_validation_errors():
+    with pytest.raises(ValueError):
+        aggregate_prototypes([[1.0, 2.0], [1.0]])  # mismatched dimensions
+    with pytest.raises(ValueError):
+        aggregate_prototypes([])  # no prototypes
+    with pytest.raises(ValueError):
+        aggregate_prototypes([[1.0, float("nan")]])  # non-finite values

tests/test_data_pipeline.py

@@ -0,0 +1,156 @@
+"""Tests for data configuration, loading, splitting and preprocessing."""
+
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from protean_research.data import loading, splits
+from protean_research.data.preprocessing import Preprocessor, fit_preprocessor, transform
+
+
+@pytest.fixture()
+def toy_csv(tmp_path):
+    frame = pd.DataFrame(
+        {
+            "id": list(range(12)),
+            "num": ["1", "-", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12"],
+            "cat": ["a", "b", "a", "b", "a", "b", "a", "b", "a", "b", "a", "b"],
+            "label": ["Normal"] * 4
+            + ["Attack"] * 4
+            + ["Lateral _movement"] * 2
+            + ["Attack"] * 2,
+        }
+    )
+    # Duplicate one row to exercise the duplicate-drop policy.
+    frame = pd.concat([frame, frame.iloc[[0]]], ignore_index=True)
+    path = tmp_path / "toy.csv"
+    frame.to_csv(path, index=False)
+    return path
+
+
+@pytest.fixture()
+def toy_cfg(toy_csv):
+    return {
+        "datasets": {
+            "toy": {
+                "name": "Toy",
+                "raw_path": str(toy_csv),
+                "label_column": "label",
+                "benign_label": "Normal",
+                "drop_labels": [],
+                "missing_tokens": ["-"],
+                "drop_duplicate_rows": True,
+                "categorical_columns": ["cat"],
+                "non_feature_columns": ["id"],
+            }
+        }
+    }
+
+
+def test_spec_resolution_and_feature_columns(toy_cfg):
+    spec = loading.load_spec("toy", toy_cfg)
+    assert spec.label_column == "label"
+    assert spec.non_feature_set == {"id", "label"}
+    cleaned = loading.clean_frame(loading.read_raw(spec), spec)
+    # 13 written rows - 1 duplicate = 12.
+    assert cleaned.n_rows_raw == 13
+    assert cleaned.n_duplicates_dropped == 1
+    assert len(cleaned.labels) == 12
+    assert set(cleaned.feature_columns) == {"num", "cat"}
+    # Non-feature and alias columns are excluded from features.
+    assert "id" not in cleaned.feature_columns
+    assert "label" not in cleaned.feature_columns
+
+
+def test_label_mapping_deterministic_and_canonical(toy_cfg):
+    spec = loading.load_spec("toy", toy_cfg)
+    cleaned = loading.clean_frame(loading.read_raw(spec), spec)
+    assert cleaned.class_names == ["Attack", "Lateral_movement", "Normal"]
+    assert cleaned.class_names[cleaned.benign_index] == "Normal"
+    assert loading.canonical_label("Lateral _movement") == "Lateral_movement"
+    # Re-running produces the identical encoding.
+    again = loading.clean_frame(loading.read_raw(spec), spec)
+    assert np.array_equal(cleaned.labels, again.labels)
+
+
+def test_unknown_dataset_raises():
+    with pytest.raises(KeyError):
+        loading.load_spec("does-not-exist", {"datasets": {}})
+
+
+def test_real_dataset_config_resolves_without_reading():
+    # The tracked dataset config must resolve both research datasets and their
+    # external raw paths without touching the raw files.
+    for key in ("xiiotid", "5gnidd"):
+        spec = loading.load_spec(key)
+        assert spec.key == key
+        assert spec.label_column
+        assert spec.benign_label
+        assert spec.raw_path.is_absolute()
+        assert spec.non_feature_columns
+    spec = loading.load_spec("5gnidd")
+    assert "Unnamed: 0" in spec.non_feature_columns
+    assert spec.missing_tokens == ("",)
+
+
+def test_split_is_deterministic_and_disjoint():
+    labels = np.repeat([0, 1, 2, 3], 25)
+    a = splits.stratified_split(labels, "toy", seed=0)
+    b = splits.stratified_split(labels, "toy", seed=0)
+    assert np.array_equal(a.train, b.train)
+    assert np.array_equal(a.validation, b.validation)
+    assert np.array_equal(a.test, b.test)
+    splits.check_disjoint(a, n_rows=len(labels))
+    # Every class appears in validation and test.
+    for split in (a.validation, a.test):
+        assert set(np.unique(labels[split])) == {0, 1, 2, 3}
+
+
+def test_different_seed_changes_split():
+    labels = np.repeat([0, 1, 2, 3], 25)
+    a = splits.stratified_split(labels, "toy", seed=0)
+    b = splits.stratified_split(labels, "toy", seed=1)
+    assert not np.array_equal(a.train, b.train)
+
+
+def test_preprocessing_is_deterministic(toy_cfg):
+    spec = loading.load_spec("toy", toy_cfg)
+    cleaned = loading.clean_frame(loading.read_raw(spec), spec)
+    split = splits.stratified_split(cleaned.labels, "toy", seed=0)
+    pre1 = fit_preprocessor(cleaned, split.train)
+    pre2 = fit_preprocessor(cleaned, split.train)
+    assert pre1.to_dict() == pre2.to_dict()
+    x1 = transform(pre1, cleaned)
+    x2 = transform(pre1, cleaned)
+    assert np.array_equal(x1, x2)
+    assert x1.dtype == np.float32
+    assert np.isfinite(x1).all()
+
+
+def test_preprocessing_fits_on_train_only(toy_cfg):
+    spec = loading.load_spec("toy", toy_cfg)
+    cleaned = loading.clean_frame(loading.read_raw(spec), spec)
+    train_idx = np.arange(0, 8)
+    pre = fit_preprocessor(cleaned, train_idx)
+    train_only = cleaned.features.iloc[train_idx]
+    values = pd.to_numeric(
+        train_only["num"].replace({"-": np.nan}), errors="coerce"
+    )
+    expected_median = float(values.median())
+    expected_mean = float(values.fillna(expected_median).mean())
+    assert pre.medians["num"] == pytest.approx(expected_median)
+    assert pre.means["num"] == pytest.approx(expected_mean)
+
+
+def test_preprocessor_roundtrip(toy_cfg, tmp_path):
+    spec = loading.load_spec("toy", toy_cfg)
+    cleaned = loading.clean_frame(loading.read_raw(spec), spec)
+    split = splits.stratified_split(cleaned.labels, "toy", seed=0)
+    pre = fit_preprocessor(cleaned, split.train)
+    path = tmp_path / "pre.json"
+    pre.save(path)
+    loaded = Preprocessor.load(path)
+    assert loaded.to_dict() == pre.to_dict()
+    assert np.array_equal(transform(loaded, cleaned), transform(pre, cleaned))

tests/test_federated_smoke.py

@@ -0,0 +1,50 @@
+"""A tiny end-to-end federated smoke test on synthetic, separable data.
+
+Guards the baseline loop (local update -> model averaging -> prototype
+aggregation -> redistribution -> validation metrics) without needing the raw
+datasets. This is a pipeline test, not a research result.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from protean_research.federated import TrainingConfig, run_federated  # noqa: E402
+
+
+def test_federated_smoke_runs_and_produces_prototypes():
+    rng = np.random.default_rng(0)
+    n_per_class = 20
+    num_classes = 3
+    dim = 6
+    features = rng.normal(size=(n_per_class * num_classes, dim)).astype(np.float32)
+    labels = np.repeat(np.arange(num_classes), n_per_class).astype(np.int64)
+    features = features + labels[:, None] * 2.0  # make classes separable
+    assignments = [np.arange(0, 20), np.arange(20, 40), np.arange(40, 60)]
+
+    config = TrainingConfig(
+        rounds=1,
+        local_epochs=1,
+        batch_size=16,
+        num_clients=3,
+        embedding_dim=8,
+        dense_units=8,
+    )
+    result = run_federated(
+        train_features=features,
+        train_labels=labels,
+        assignments=assignments,
+        val_features=features,
+        val_labels=labels,
+        num_classes=num_classes,
+        input_dim=dim,
+        config=config,
+        init_seed_parts=("toy", "seed=0"),
+    )
+    assert set(result["global_prototypes"]) == {0, 1, 2}
+    assert len(result["parameters"]) > 0
+    assert np.isfinite(result["history"][-1]["accuracy"])
+    assert result["snapshot"] is not None

tests/test_metrics.py

@@ -0,0 +1,44 @@
+"""Tests for the evaluation metrics on small known fixtures."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from protean_research.evaluation.metrics import (
+    classification_metrics,
+    confusion_matrix,
+    rare_class_metrics,
+)
+
+
+def test_confusion_matrix_counts():
+    matrix = confusion_matrix([0, 0, 1, 1], [0, 1, 1, 1], num_classes=2)
+    assert matrix.tolist() == [[1, 1], [0, 2]]
+
+
+def test_classification_metrics_known_fixture():
+    y_true = [0, 0, 1, 1, 2, 2]
+    y_pred = [0, 0, 1, 2, 2, 1]
+    metrics = classification_metrics(y_true, y_pred, num_classes=3)
+    assert metrics["accuracy"] == pytest.approx(4 / 6)
+    assert metrics["macro_f1"] == pytest.approx((1.0 + 0.5 + 0.5) / 3)
+    per_class = {row["class_id"]: row for row in metrics["per_class"]}
+    assert per_class[0]["f1"] == pytest.approx(1.0)
+    assert per_class[1]["precision"] == pytest.approx(0.5)
+    assert per_class[2]["recall"] == pytest.approx(0.5)
+
+
+def test_rare_class_metrics_restricts_to_given_classes():
+    y_true = [0, 0, 1, 1, 2, 2]
+    y_pred = [0, 0, 1, 2, 2, 1]
+    rare = rare_class_metrics(y_true, y_pred, num_classes=3, rare_class_ids=[1, 2])
+    assert rare["rare_class_macro_f1"] == pytest.approx(0.5)
+    assert rare["rare_class_recall"] == pytest.approx(0.5)
+    assert rare["rare_support"] == 4
+
+
+def test_rare_class_metrics_skips_absent_classes():
+    rare = rare_class_metrics([0, 0], [0, 0], num_classes=3, rare_class_ids=[2])
+    assert rare["rare_class_macro_f1"] == 0.0
+    assert rare["rare_support"] == 0

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) == []

tests/test_prototype_aggregation.py

@@ -1,126 +0,0 @@
-"""Unit tests for the server-side prototype aggregation math.
-
-Covers the dependency-free subset of the testing strategy in
-``docs/research/IMPLEMENTATION_PLAN.md`` (items 1-7). Items requiring the
-future data/federated modules (partition determinism, manifest reload,
-baseline invariance, validation/test separation) are deferred.
-"""
-
-from __future__ import annotations
-
-import math
-
-import pytest
-
-from protean_research.aggregation import (
-    AggregationMode,
-    aggregate_prototypes,
-    aggregation_weights,
-    beta_weights,
-)
-
-BETAS = [0.0, 0.25, 0.5, 0.75, 1.0]
-
-
-def mean(vectors):
-    dim = len(vectors[0])
-    return [sum(v[d] for v in vectors) / len(vectors) for d in range(dim)]
-
-
-def test_beta_zero_equals_arithmetic_mean() -> None:
-    prototypes = [[0.0, 2.0, -1.0], [2.0, 0.0, 1.0], [4.0, 4.0, 4.0]]
-    supports = [1.0, 3.0, 10.0]
-    out = aggregate_prototypes(
-        prototypes, AggregationMode.SUPPORT_WEIGHTED, beta=0.0, supports=supports
-    )
-    assert out == pytest.approx(mean(prototypes))
-
-
-def test_beta_one_gives_sample_count_weighting() -> None:
-    prototypes = [[0.0, 2.0], [2.0, 0.0], [4.0, 4.0]]
-    supports = [1.0, 3.0, 10.0]
-    out = aggregate_prototypes(
-        prototypes, AggregationMode.SUPPORT_WEIGHTED, beta=1.0, supports=supports
-    )
-    total = sum(supports)
-    expected = [
-        sum(s * v[d] for s, v in zip(supports, prototypes)) / total
-        for d in range(2)
-    ]
-    assert out == pytest.approx(expected)
-
-
-def test_single_contributor_invariant_across_betas() -> None:
-    prototypes = [[1.0, -2.0, 3.0]]
-    supports = [7.0]
-    for beta in BETAS:
-        out = aggregate_prototypes(
-            prototypes, AggregationMode.SUPPORT_WEIGHTED, beta=beta, supports=supports
-        )
-        assert out == pytest.approx(prototypes[0])
-
-
-def test_equal_supports_invariant_across_betas() -> None:
-    prototypes = [[0.0, 2.0], [2.0, 0.0], [4.0, 4.0]]
-    supports = [5.0, 5.0, 5.0]
-    baseline = aggregate_prototypes(
-        prototypes, AggregationMode.SUPPORT_WEIGHTED, beta=0.0, supports=supports
-    )
-    for beta in BETAS:
-        out = aggregate_prototypes(
-            prototypes, AggregationMode.SUPPORT_WEIGHTED, beta=beta, supports=supports
-        )
-        assert out == pytest.approx(baseline)
-
-
-def test_missing_class_not_in_s_j() -> None:
-    # A client with N_ij == 0 must never be passed into S_j.
-    with pytest.raises(ValueError):
-        beta_weights([1.0, 0.0], beta=0.5)
-    with pytest.raises(ValueError):
-        beta_weights([1.0, -3.0], beta=0.5)
-    with pytest.raises(ValueError):
-        aggregate_prototypes(
-            [[1.0], [2.0]], AggregationMode.SUPPORT_WEIGHTED, beta=0.5, supports=[1.0, 0.0]
-        )
-
-
-def test_normalized_weights_sum_to_one() -> None:
-    prototypes = [[0.0, 2.0], [2.0, 0.0], [4.0, 4.0]]
-    for beta in BETAS:
-        w = aggregation_weights(
-            prototypes, AggregationMode.SUPPORT_WEIGHTED, beta=beta, supports=[1.0, 3.0, 10.0]
-        )
-        assert sum(w) == pytest.approx(1.0)
-    w_eq = aggregation_weights(prototypes, AggregationMode.EQUAL_AVAILABLE)
-    assert sum(w_eq) == pytest.approx(1.0)
-    assert w_eq == pytest.approx([1 / 3, 1 / 3, 1 / 3])
-
-
-def test_valid_inputs_never_nan_or_inf() -> None:
-    prototypes = [[0.0, 2.0], [2.0, 0.0], [4.0, 4.0]]
-    for beta in BETAS:
-        out = aggregate_prototypes(
-            prototypes, AggregationMode.SUPPORT_WEIGHTED, beta=beta, supports=[1.0, 3.0, 10.0]
-        )
-        assert all(math.isfinite(v) for v in out)
-
-
-def test_original_protean_is_not_implemented() -> None:
-    with pytest.raises(NotImplementedError):
-        aggregate_prototypes([[1.0, 2.0]], AggregationMode.ORIGINAL_PROTEAN)
-
-
-def test_validation_errors() -> None:
-    with pytest.raises(ValueError):
-        beta_weights([1.0], beta=1.5)
-    with pytest.raises(ValueError):
-        beta_weights([1.0], beta=None)  # type: ignore[arg-type]
-    with pytest.raises(ValueError):
-        aggregate_prototypes([[1.0, 2.0], [1.0]], AggregationMode.SUPPORT_WEIGHTED, beta=0.5, supports=[1.0, 2.0])
-    with pytest.raises(ValueError):
-        aggregate_prototypes(
-            [[1.0, 2.0]], AggregationMode.SUPPORT_WEIGHTED, beta=0.5, supports=[1.0, 2.0]
-        )
-    with pytest.raises(ValueError):
-        aggregate_prototypes([[1.0, float("nan")]], AggregationMode.EQUAL_AVAILABLE)

tests/test_prototypes.py

@@ -0,0 +1,58 @@
+"""Tests for local prototype construction and global prototype collection."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from protean_research.aggregation import AggregationMode
+from protean_research.models.prototypes import (
+    collect_global_prototypes,
+    compute_local_prototypes,
+    nearest_prototype_predict,
+)
+
+
+def test_local_prototype_is_mean_embedding():
+    embeddings = np.array(
+        [[0.0, 0.0], [2.0, 0.0], [10.0, 10.0], [0.0, 4.0]], dtype=np.float32
+    )
+    labels = np.array([0, 0, 1, 1])
+    prototypes, supports = compute_local_prototypes(embeddings, labels, num_classes=3)
+    assert supports == {0: 2, 1: 2}
+    assert prototypes[0] == pytest.approx([1.0, 0.0])
+    assert prototypes[1] == pytest.approx([5.0, 7.0])
+
+
+def test_missing_class_has_no_prototype():
+    embeddings = np.array([[1.0, 1.0]], dtype=np.float32)
+    labels = np.array([2])
+    prototypes, supports = compute_local_prototypes(embeddings, labels, num_classes=4)
+    assert set(prototypes) == {2}
+    assert 0 not in prototypes
+    assert 1 not in prototypes
+
+
+def test_collect_global_original_is_equal_mean_over_contributors():
+    client_prototypes = [{0: np.array([0.0, 2.0])}, {0: np.array([4.0, 4.0])}]
+    client_supports = [{0: 1}, {0: 100}]
+    result = collect_global_prototypes(
+        client_prototypes, client_supports, num_classes=1, mode=AggregationMode.ORIGINAL_PROTEAN
+    )
+    assert result[0] == pytest.approx([2.0, 3.0])
+
+
+def test_collect_global_omits_missing_class():
+    client_prototypes = [{0: np.array([1.0, 1.0])}, {}]
+    client_supports = [{0: 3}, {}]
+    result = collect_global_prototypes(
+        client_prototypes, client_supports, num_classes=2, mode=AggregationMode.ORIGINAL_PROTEAN
+    )
+    assert set(result) == {0}
+
+
+def test_nearest_prototype_predict():
+    prototypes = {0: np.array([0.0, 0.0]), 1: np.array([10.0, 10.0])}
+    embeddings = np.array([[0.1, 0.1], [9.0, 9.0]])
+    predictions = nearest_prototype_predict(embeddings, prototypes)
+    assert predictions.tolist() == [0, 1]