feat: implement reproducible PROTEAN baseline pipeline

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]