feat: implement support-weighted prototype aggregation

tests/test_baseline_aggregation.py

@@ -61,10 +61,19 @@ def test_missing_class_is_not_fabricated():
     assert result == {}
 
 
-def test_proposed_mode_is_not_available():
-    # The support-weighted mode belongs to the later beta/controls phase.
+def test_support_weighted_requires_supports_and_beta():
+    # The proposed mode now exists but must be fully parameterised.
     with pytest.raises(ValueError):
-        aggregate_prototypes(PROTOTYPES, "support_weighted")
+        aggregate_prototypes(PROTOTYPES, "support_weighted")  # no supports/beta
+    with pytest.raises(ValueError):
+        aggregate_prototypes(PROTOTYPES, "support_weighted", supports=[1, 2, 3])
+    with pytest.raises(ValueError):
+        aggregate_prototypes(PROTOTYPES, "not_a_mode")
+
+
+def test_beta_is_rejected_for_original_protean():
+    with pytest.raises(ValueError):
+        aggregate_prototypes(PROTOTYPES, "original_protean", beta=0.5)
 
 
 def test_validation_errors():

tests/test_prototype_effect.py

@@ -0,0 +1,142 @@
+"""Direct prototype-effect instrumentation and support_weighted end-to-end wiring."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from protean_research.analysis.direct_prototype import (
+    aggregate_snapshot,
+    client_weight_rows,
+    prototype_effect_rows,
+    snapshot_class_members,
+)
+
+
+def make_snapshot():
+    # Two classes. Class 0 has strongly imbalanced supports across 3 clients;
+    # class 1 has a single contributor (beta-invariant).
+    return {
+        "round": 0,
+        "client_prototypes": [
+            {0: [0.0, 0.0], 1: [1.0, 1.0]},
+            {0: [4.0, 8.0]},
+            {0: [10.0, 2.0]},
+        ],
+        "client_supports": [
+            {0: 1, 1: 5},
+            {0: 9},
+            {0: 4},
+        ],
+    }
+
+
+def test_snapshot_class_members_aligns_supports():
+    members = snapshot_class_members(make_snapshot())
+    assert sorted(members) == [0, 1]
+    assert [cid for cid, _, _ in members[0]] == [0, 1, 2]
+    assert [n for _, _, n in members[0]] == [1, 9, 4]
+    assert members[1] == [(0, [1.0, 1.0], 5)]
+
+
+def test_beta0_effect_is_zero_and_beta1_displaces():
+    rows = prototype_effect_rows(make_snapshot(), betas=[0.0, 1.0])
+    by_key = {(r["class_id"], r["beta"]): r for r in rows}
+    assert by_key[(0, 0.0)]["displacement_from_beta0"] == pytest.approx(0.0)
+    assert by_key[(0, 0.0)]["relative_displacement"] == pytest.approx(0.0)
+    assert by_key[(0, 1.0)]["displacement_from_beta0"] > 0.0
+    # Single contributor -> no displacement for any beta.
+    assert by_key[(1, 1.0)]["displacement_from_beta0"] == pytest.approx(0.0)
+    assert by_key[(1, 1.0)]["n_contributors"] == 1
+
+
+def test_effect_beta0_matches_arithmetic_mean_aggregate():
+    snapshot = make_snapshot()
+    aggregated = aggregate_snapshot(snapshot, 0.0)
+    assert aggregated[0] == pytest.approx([14.0 / 3.0, 10.0 / 3.0])
+    assert aggregated[1] == pytest.approx([1.0, 1.0])
+
+
+def test_client_weights_sum_to_one_per_class_beta():
+    rows = client_weight_rows(make_snapshot(), betas=[0.0, 0.5, 1.0])
+    groups: dict[tuple, float] = {}
+    for row in rows:
+        key = (row["class_id"], row["beta"])
+        groups[key] = groups.get(key, 0.0) + row["weight"]
+    assert groups
+    for total in groups.values():
+        assert total == pytest.approx(1.0)
+
+
+def test_json_style_string_class_ids_are_supported():
+    snapshot = make_snapshot()
+    json_like = {
+        "round": 0,
+        "client_prototypes": [
+            {str(c): v for c, v in prototypes.items()}
+            for prototypes in snapshot["client_prototypes"]
+        ],
+        "client_supports": [
+            {str(c): v for c, v in supports.items()}
+            for supports in snapshot["client_supports"]
+        ],
+    }
+    rows = prototype_effect_rows(json_like, betas=[0.0, 1.0])
+    assert {(r["class_id"], r["beta"]) for r in rows} == {(0, 0.0), (0, 1.0), (1, 0.0), (1, 1.0)}
+
+
+# ---------------------------------------------------------------------------
+# End-to-end wiring: support_weighted must run through the federated loop.
+# ---------------------------------------------------------------------------
+
+torch = pytest.importorskip("torch")
+
+from protean_research.federated import TrainingConfig, run_federated  # noqa: E402
+
+
+def _run_support_weighted(beta, seed_parts):
+    rng = np.random.default_rng(0)
+    num_classes, per_class, dim = 3, 10, 6
+    features = rng.normal(size=(per_class * num_classes, dim)).astype(np.float32)
+    labels = np.repeat(np.arange(num_classes), per_class).astype(np.int64)
+    features = features + labels[:, None] * 2.0
+    assignments = [np.arange(0, 10), np.arange(10, 20), np.arange(20, 30)]
+    config = TrainingConfig(
+        rounds=2,
+        local_epochs=1,
+        batch_size=10,
+        num_clients=3,
+        embedding_dim=8,
+        dense_units=8,
+        aggregation_mode="support_weighted",
+        beta=beta,
+    )
+    return 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=seed_parts,
+    )
+
+
+def test_support_weighted_end_to_end_is_deterministic_and_finite():
+    first = _run_support_weighted(0.5, ("sw", "seed=0"))
+    second = _run_support_weighted(0.5, ("sw", "seed=0"))
+    assert set(first["global_prototypes"]) == {0, 1, 2}
+    for c in first["global_prototypes"]:
+        assert np.isfinite(np.asarray(first["global_prototypes"][c])).all()
+        assert np.array_equal(
+            np.asarray(first["global_prototypes"][c]),
+            np.asarray(second["global_prototypes"][c]),
+        )
+    assert np.isfinite(first["history"][-1]["accuracy"])
+
+
+def test_support_weighted_rejects_bad_beta_end_to_end():
+    with pytest.raises(ValueError):
+        _run_support_weighted(1.5, ("sw", "bad"))

tests/test_support_weighted_aggregation.py

@@ -0,0 +1,200 @@
+"""Invariant tests for the proposed support-weighted prototype aggregation.
+
+These tests pin the required arithmetic properties of ``support_weighted`` with
+``w_ij = N_ij ** beta`` over the same contributor set ``S_j`` as the baseline:
+
+- ``beta = 0`` is the equal (arithmetic) mean over ``S_j``;
+- ``beta = 1`` is the support-weighted mean;
+- a hand-computed intermediate-beta fixture;
+- single-contributor invariance;
+- equal-support invariance;
+- missing-class exclusion (``S_j`` construction);
+- normalised weights sum to one;
+- invalid-beta rejection.
+
+They do not assert that ``beta = 0`` *is* original PROTEAN; that depends on the
+separately documented participant rule.
+"""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+
+from protean_research.aggregation import (
+    BETA_CANDIDATES,
+    AggregationMode,
+    aggregate_class_prototypes,
+    aggregate_prototypes,
+    normalized_support_weights,
+    validate_beta,
+)
+
+SUPPORT_WEIGHTED = AggregationMode.SUPPORT_WEIGHTED
+ORIGINAL = AggregationMode.ORIGINAL_PROTEAN
+
+# Prototypes with strongly different supports to make weighting observable.
+PROTOTYPES = [[0.0, 0.0], [4.0, 8.0], [10.0, 2.0]]
+SUPPORTS = [1.0, 9.0, 4.0]
+
+
+def arithmetic_mean(vectors):
+    dim = len(vectors[0])
+    return [sum(v[d] for v in vectors) / len(vectors) for d in range(dim)]
+
+
+def support_weighted_mean(vectors, supports, beta):
+    weights = [float(n) ** beta for n in supports]
+    total = sum(weights)
+    dim = len(vectors[0])
+    return [sum(w * v[d] for w, v in zip(weights, vectors)) / total for d in range(dim)]
+
+
+# ---------------------------------------------------------------------------
+# beta = 0
+# ---------------------------------------------------------------------------
+
+
+def test_beta0_is_equal_arithmetic_mean_over_s_j():
+    out = aggregate_prototypes(
+        PROTOTYPES, SUPPORT_WEIGHTED, supports=SUPPORTS, beta=0.0
+    )
+    assert out == pytest.approx(arithmetic_mean(PROTOTYPES))
+
+
+def test_beta0_matches_original_protean_on_identical_s_j():
+    weighted = aggregate_prototypes(
+        PROTOTYPES, SUPPORT_WEIGHTED, supports=SUPPORTS, beta=0.0
+    )
+    baseline = aggregate_prototypes(PROTOTYPES, ORIGINAL)
+    assert weighted == pytest.approx(baseline, rel=1e-6, abs=1e-6)
+
+
+# ---------------------------------------------------------------------------
+# beta = 1 and intermediate beta
+# ---------------------------------------------------------------------------
+
+
+def test_beta1_is_support_weighted_mean():
+    out = aggregate_prototypes(
+        PROTOTYPES, SUPPORT_WEIGHTED, supports=SUPPORTS, beta=1.0
+    )
+    assert out == pytest.approx(support_weighted_mean(PROTOTYPES, SUPPORTS, 1.0))
+
+
+def test_intermediate_beta_fixture():
+    # two contributors: C = [0, 4], N = 1 and 9; beta = 0.5 -> weights 1, 3.
+    out = aggregate_prototypes(
+        [[0.0, 0.0], [4.0, 4.0]], SUPPORT_WEIGHTED, supports=[1.0, 9.0], beta=0.5
+    )
+    expected = [3.0, 3.0]  # (1*0 + 3*4) / 4
+    assert out == pytest.approx(expected)
+
+
+def test_intermediate_beta_matches_closed_form():
+    for beta in (0.25, 0.5, 0.75):
+        out = aggregate_prototypes(
+            PROTOTYPES, SUPPORT_WEIGHTED, supports=SUPPORTS, beta=beta
+        )
+        assert out == pytest.approx(support_weighted_mean(PROTOTYPES, SUPPORTS, beta))
+
+
+# ---------------------------------------------------------------------------
+# invariance
+# ---------------------------------------------------------------------------
+
+
+def test_single_contributor_is_invariant_to_beta():
+    single = [[1.5, -2.0, 3.0]]
+    for beta in BETA_CANDIDATES:
+        out = aggregate_prototypes(single, SUPPORT_WEIGHTED, supports=[7.0], beta=beta)
+        assert out == pytest.approx(single[0])
+
+
+def test_equal_support_is_invariant_to_beta():
+    supports = [5.0, 5.0, 5.0]
+    reference = arithmetic_mean(PROTOTYPES)
+    for beta in BETA_CANDIDATES:
+        out = aggregate_prototypes(
+            PROTOTYPES, SUPPORT_WEIGHTED, supports=supports, beta=beta
+        )
+        assert out == pytest.approx(reference, rel=1e-9, abs=1e-9)
+
+
+# ---------------------------------------------------------------------------
+# S_j construction / missing classes
+# ---------------------------------------------------------------------------
+
+
+def test_missing_class_is_excluded_from_s_j():
+    contributions = {
+        0: [([1.0, 1.0], 5), (None, 5), ([3.0, 3.0], 0), ([5.0, 5.0], 5)],
+    }
+    result = aggregate_class_prototypes(contributions, SUPPORT_WEIGHTED, beta=1.0)
+    # Only the two positive-support entries with prototypes enter: mean = 3.
+    assert result[0] == pytest.approx([3.0, 3.0])
+
+
+def test_class_with_only_missing_contributors_is_omitted():
+    result = aggregate_class_prototypes(
+        {0: [(None, 3), ([1.0], 0)]}, SUPPORT_WEIGHTED, beta=0.5
+    )
+    assert result == {}
+
+
+# ---------------------------------------------------------------------------
+# weights
+# ---------------------------------------------------------------------------
+
+
+def test_normalized_weights_sum_to_one():
+    for beta in BETA_CANDIDATES:
+        weights = normalized_support_weights(SUPPORTS, beta)
+        assert sum(weights) == pytest.approx(1.0)
+        assert all(w > 0.0 for w in weights)
+
+
+def test_beta0_weights_are_uniform():
+    weights = normalized_support_weights(SUPPORTS, 0.0)
+    assert weights == pytest.approx([1 / 3, 1 / 3, 1 / 3])
+
+
+def test_output_is_finite_for_all_candidates():
+    for beta in BETA_CANDIDATES:
+        out = aggregate_prototypes(
+            PROTOTYPES, SUPPORT_WEIGHTED, supports=SUPPORTS, beta=beta
+        )
+        assert all(math.isfinite(v) for v in out)
+
+
+# ---------------------------------------------------------------------------
+# validation / rejection
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("beta", [-0.01, 1.01, -1.0, 2.0, float("nan"), float("inf")])
+def test_invalid_beta_is_rejected(beta):
+    with pytest.raises(ValueError):
+        validate_beta(beta)
+    with pytest.raises(ValueError):
+        aggregate_prototypes(PROTOTYPES, SUPPORT_WEIGHTED, supports=SUPPORTS, beta=beta)
+
+
+def test_none_and_non_numeric_beta_are_rejected():
+    for beta in (None, "0.5", True, [0.5]):
+        with pytest.raises(ValueError):
+            validate_beta(beta)
+
+
+def test_supports_must_be_strictly_positive_and_aligned():
+    with pytest.raises(ValueError):
+        aggregate_prototypes(PROTOTYPES, SUPPORT_WEIGHTED, supports=[1, 2], beta=1.0)
+    with pytest.raises(ValueError):
+        aggregate_prototypes(PROTOTYPES, SUPPORT_WEIGHTED, supports=[1, 0, 3], beta=1.0)
+    with pytest.raises(ValueError):
+        aggregate_prototypes(PROTOTYPES, SUPPORT_WEIGHTED, supports=None, beta=1.0)
+
+
+def test_beta_candidates_are_the_frozen_protocol_values():
+    assert BETA_CANDIDATES == (0.0, 0.25, 0.50, 0.75, 1.00)