59 lines
2.0 KiB
Python
59 lines
2.0 KiB
Python
"""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]
|