feat: implement support-weighted prototype aggregation

src/protean_research/aggregation/__init__.py

@@ -1,18 +1,26 @@
 """Server-side prototype aggregation.
 
-This package is the single locus of the proposed modification. In the baseline
-checkpoint it implements only the reproduced original-PROTEAN aggregation. The
-proposed support-weighted aggregation is deferred to the beta/controls phase.
+This package is the single locus of the proposed modification. It implements
+the reproduced original-PROTEAN aggregation (``original_protean``) and the
+proposed class-specific sample-support weighting (``support_weighted``) behind an
+explicit mode/beta interface. Missing-class handling (``S_j``) stays explicit in
+the aggregation functions; no zero-vector padding is used.
 """
 
 from .prototype_aggregation import (
+    BETA_CANDIDATES,
     AggregationMode,
     aggregate_class_prototypes,
     aggregate_prototypes,
+    normalized_support_weights,
+    validate_beta,
 )
 
 __all__ = [
+    "BETA_CANDIDATES",
     "AggregationMode",
     "aggregate_class_prototypes",
     "aggregate_prototypes",
+    "normalized_support_weights",
+    "validate_beta",
 ]

src/protean_research/aggregation/prototype_aggregation.py

@@ -1,25 +1,49 @@
-"""Baseline server-side prototype aggregation.
+"""Server-side prototype aggregation.
 
-This module implements **only the reproduced original-PROTEAN prototype
-aggregation** used by the baseline. The proposed support-weighted aggregation
-(``N_ij ** beta``) is deliberately NOT present in this checkpoint; it belongs to
-the later beta/controls phase and must not be introduced before the baseline is
-verified.
+This module implements the reproduced original-PROTEAN prototype aggregation
+(``original_protean``) and the proposed class-specific sample-support weighting
+(``support_weighted``). Only this server-side step differs between the two modes;
+the data pipeline, model, loss, model averaging and inference stay frozen.
 
-Participant rule (reproduction assumption, not verified against code)
----------------------------------------------------------------------
-The published Equation (2) / thesis Eq. (5.5) writes ``C_bar_j = (1/M) * sum_i
-C_ij`` over all ``M`` clients, but never states how a class that is missing on a
-client (``N_ij = 0``) is handled. The thesis separately states that "prototypes
-are shared only for locally observed classes", so this baseline assumes the
-server receives prototypes only from the contributor set::
+Supported modes
+---------------
+``original_protean``
+    Equal weights over the contributor set ``S_j`` (reproduction baseline)::
+
+        Cbar_j = (1 / |S_j|) * sum_{i in S_j} C_ij
+
+``support_weighted``
+    Class-specific sample-support weighting (the proposed method)::
+
+        Cbar_j(beta) = ( sum_{i in S_j} N_ij^beta * C_ij )
+                       / ( sum_{i in S_j} N_ij^beta )
+
+    with ``0 <= beta <= 1``.
+
+Participant rule (``S_j``)
+--------------------------
+Both modes share the same participant set::
 
     S_j = { i | N_ij > 0 and a valid local prototype C_ij exists }
 
-and averages them with **equal weights ``1/|S_j|``**. This is an explicit
-REPRODUCTION ASSUMPTION of this project, not a source-confirmed fact. It must be
-validated in the baseline-verification phase; see
-``docs/baseline/BASELINE_IMPLEMENTATION.md``.
+Missing classes (``N_ij = 0`` or no prototype) never enter ``S_j`` and are never
+fabricated as zero vectors. For ``original_protean`` this participant rule is an
+explicit REPRODUCTION ASSUMPTION, not a source-confirmed fact (the published
+Eq. (2) / thesis Eq. (5.5) writes ``1/M`` over all clients but never defines
+missing-class handling). The same documented rule is reused for
+``support_weighted`` so that beta candidates are compared on an identical
+contributor set.
+
+Beta semantics
+--------------
+- ``beta = 0``: every weight is ``N_ij ** 0 = 1``, so for a fixed ``S_j`` the
+  aggregate is the equal (arithmetic) mean over ``S_j``. This is *not* asserted to
+  be the original PROTEAN baseline beyond the shared participant rule.
+- ``beta = 1``: weights are proportional to the class-specific support; because
+  each ``C_ij`` is a local class-embedding mean, the aggregate corresponds to the
+  combined mean of the included embeddings when all assumptions match.
+- A single contributor and an all-equal support vector make the aggregate
+  independent of beta.
 """
 
 from __future__ import annotations
@@ -30,11 +54,15 @@ from typing import Mapping, Sequence
 
 Vector = Sequence[float]
 
+#: Frozen beta candidates from the research specification.
+BETA_CANDIDATES = (0.0, 0.25, 0.50, 0.75, 1.00)
+
 
 class AggregationMode(str, Enum):
-    """Prototype-aggregation modes implemented in the baseline checkpoint."""
+    """Prototype-aggregation modes implemented in this checkpoint."""
 
     ORIGINAL_PROTEAN = "original_protean"
+    SUPPORT_WEIGHTED = "support_weighted"
 
 
 def _validate_vectors(prototypes: Sequence[Vector]) -> int:
@@ -56,9 +84,53 @@ def _validate_vectors(prototypes: Sequence[Vector]) -> int:
     return dim
 
 
+def validate_beta(beta) -> float:
+    """Return ``beta`` as a float in ``[0, 1]`` or raise ``ValueError``.
+
+    Explicit and strict: ``None``, non-numeric, non-finite and out-of-range
+    values are all rejected rather than silently defaulted.
+    """
+    if beta is None:
+        raise ValueError("beta is required for support_weighted aggregation")
+    if isinstance(beta, bool) or not isinstance(beta, (int, float)):
+        raise ValueError(f"beta must be a real number, got {beta!r}")
+    value = float(beta)
+    if not math.isfinite(value):
+        raise ValueError(f"beta must be finite, got {beta!r}")
+    if value < 0.0 or value > 1.0:
+        raise ValueError(f"beta must lie in the closed interval [0, 1], got {value}")
+    return value
+
+
+def normalized_support_weights(supports: Sequence[float], beta) -> list[float]:
+    """Normalised ``N_ij ** beta`` weights that sum to one.
+
+    ``supports`` are the class-specific supports ``N_ij`` of the clients in
+    ``S_j``; they must be strictly positive and finite. Used both by the
+    aggregation and by the direct prototype-effect instrumentation.
+    """
+    value = validate_beta(beta)
+    if supports is None:
+        raise ValueError("supports are required for support_weighted aggregation")
+    weights = []
+    for support in supports:
+        n = float(support)
+        if not math.isfinite(n):
+            raise ValueError(f"N_ij must be finite, got {support!r}")
+        if n <= 0.0:
+            raise ValueError("support_weighted requires strictly positive N_ij")
+        weights.append(n ** value)
+    total = math.fsum(weights)
+    if not math.isfinite(total) or total <= 0.0:
+        raise ValueError("support weights must sum to a positive finite value")
+    return [w / total for w in weights]
+
+
 def aggregate_prototypes(
     prototypes: Sequence[Vector],
     mode: AggregationMode | str = AggregationMode.ORIGINAL_PROTEAN,
+    supports: Sequence[float] | None = None,
+    beta=None,
 ) -> list[float]:
     """Aggregate the local prototypes ``C_ij`` of one class into ``C_bar_j``.
 
@@ -67,28 +139,48 @@ def aggregate_prototypes(
     prototypes:
         Local prototypes ``C_ij`` for the contributing clients in ``S_j``.
     mode:
-        Only :attr:`AggregationMode.ORIGINAL_PROTEAN` is implemented in this
-        baseline checkpoint. Other modes raise ``NotImplementedError``.
-
-    The baseline uses equal weights over the supplied contributor set, so the
-    result is the arithmetic mean of ``prototypes``.
+        :attr:`AggregationMode.ORIGINAL_PROTEAN` (equal mean) or
+        :attr:`AggregationMode.SUPPORT_WEIGHTED` (``N_ij ** beta``).
+    supports:
+        Class-specific supports ``N_ij`` aligned with ``prototypes``. Required
+        for ``support_weighted``; ignored by ``original_protean``.
+    beta:
+        Support-weighting exponent in ``[0, 1]``. Only valid for
+        ``support_weighted``; passing it with ``original_protean`` is an error.
     """
-    # ``AggregationMode`` contains only the baseline mode, so any other value
-    # raises ``ValueError`` here. This prevents the proposed support-weighted
-    # aggregation from being used before the baseline is verified.
-    AggregationMode(mode)
+    resolved = AggregationMode(mode)
     dim = _validate_vectors(prototypes)
-    count = len(prototypes)
+
+    if resolved is AggregationMode.ORIGINAL_PROTEAN:
+        if beta is not None:
+            raise ValueError("beta is only valid for support_weighted aggregation")
+        count = len(prototypes)
+        result = [0.0] * dim
+        for vec in prototypes:
+            for d in range(dim):
+                result[d] += float(vec[d]) / count
+        return result
+
+    # support_weighted
+    if supports is None:
+        raise ValueError("supports are required for support_weighted aggregation")
+    if len(supports) != len(prototypes):
+        raise ValueError(
+            f"supports and prototypes must have equal length, got "
+            f"{len(supports)} vs {len(prototypes)}"
+        )
+    weights = normalized_support_weights(supports, beta)
     result = [0.0] * dim
-    for vec in prototypes:
+    for weight, vec in zip(weights, prototypes):
         for d in range(dim):
-            result[d] += float(vec[d]) / count
+            result[d] += weight * float(vec[d])
     return result
 
 
 def aggregate_class_prototypes(
     contributions,
     mode: AggregationMode | str = AggregationMode.ORIGINAL_PROTEAN,
+    beta=None,
 ):
     """Aggregate per-class local prototypes into global prototypes.
 
@@ -101,16 +193,21 @@ def aggregate_class_prototypes(
     """
     if not isinstance(contributions, Mapping):
         raise TypeError("contributions must be a mapping class_id -> [(C_ij, N_ij), ...]")
+    resolved = AggregationMode(mode)
     result = {}
     for class_id in sorted(contributions, key=lambda c: int(c)):
         member_prototypes: list[Vector] = []
+        member_supports: list[float] = []
         for proto, support in contributions[class_id]:
             if proto is None:
                 continue
             if support is None or float(support) <= 0.0:
                 continue
             member_prototypes.append(proto)
+            member_supports.append(float(support))
         if not member_prototypes:
             continue
-        result[int(class_id)] = aggregate_prototypes(member_prototypes, mode)
+        result[int(class_id)] = aggregate_prototypes(
+            member_prototypes, resolved, supports=member_supports, beta=beta
+        )
     return result

src/protean_research/analysis/__init__.py

@@ -0,0 +1,24 @@
+"""Analysis utilities that do not depend on training.
+
+Currently this package implements the direct prototype-effect analysis
+(``direct_prototype``), which recomputes global prototypes from a fixed round
+snapshot of ``{C_ij, N_ij}`` for several beta without retraining.
+"""
+
+from .direct_prototype import (
+    DEFAULT_BETAS,
+    aggregate_snapshot,
+    client_weight_rows,
+    prototype_effect_rows,
+    snapshot_class_members,
+    snapshot_contributions,
+)
+
+__all__ = [
+    "DEFAULT_BETAS",
+    "aggregate_snapshot",
+    "client_weight_rows",
+    "prototype_effect_rows",
+    "snapshot_class_members",
+    "snapshot_contributions",
+]

src/protean_research/analysis/direct_prototype.py

@@ -0,0 +1,184 @@
+"""Direct prototype-effect analysis on a fixed ``{C_ij, N_ij}`` snapshot.
+
+This separates the *direct aggregation effect* (what changing beta does to the
+global prototypes for a fixed set of local prototypes) from the *end-to-end
+training effect* (what happens when beta also changes training). It requires no
+retraining: the federated loop already stores a round snapshot of every client's
+local prototypes ``C_ij`` and class supports ``N_ij``.
+
+The snapshot format is the one produced by
+``protean_research.federated.round.run_federated``::
+
+    {
+      "round": 0,
+      "client_prototypes": [ {class_id: [floats], ...}, ... ],
+      "client_supports":   [ {class_id: int, ...}, ... ],
+    }
+
+Both JSON-loaded (string class ids) and in-memory (int class ids) snapshots are
+accepted. All analysis reuses the production aggregation functions, so the
+reported global prototypes are exactly those a beta run would have aggregated in
+that round; only the training feedback is removed.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Mapping, Sequence
+
+from ..aggregation import (
+    AggregationMode,
+    aggregate_class_prototypes,
+    normalized_support_weights,
+)
+
+DEFAULT_BETAS = (0.0, 0.25, 0.50, 0.75, 1.00)
+
+_SUPPORT_WEIGHTED = AggregationMode.SUPPORT_WEIGHTED
+
+
+def _lookup(mapping: Mapping, class_id: int):
+    """Support both JSON (str keys) and in-memory (int keys) class mappings."""
+    if class_id in mapping:
+        return mapping[class_id]
+    return mapping.get(str(class_id))
+
+
+def snapshot_class_members(snapshot: Mapping) -> dict[int, list[tuple[int, object, float]]]:
+    """Return ``class_id -> [(client_id, C_ij, N_ij), ...]`` from a snapshot."""
+    try:
+        client_prototypes = snapshot["client_prototypes"]
+        client_supports = snapshot["client_supports"]
+    except (KeyError, TypeError) as exc:  # pragma: no cover - defensive
+        raise ValueError("snapshot must contain client_prototypes and client_supports") from exc
+    if len(client_prototypes) != len(client_supports):
+        raise ValueError(
+            "snapshot client_prototypes and client_supports must have equal length"
+        )
+
+    members: dict[int, list[tuple[int, object, float]]] = {}
+    for client_id, (prototypes, supports) in enumerate(
+        zip(client_prototypes, client_supports)
+    ):
+        for raw_class_id, proto in prototypes.items():
+            class_id = int(raw_class_id)
+            support = _lookup(supports, class_id)
+            members.setdefault(class_id, []).append((client_id, proto, support))
+    return members
+
+
+def snapshot_contributions(snapshot: Mapping) -> dict[int, list[tuple[object, float]]]:
+    """Return ``class_id -> [(C_ij, N_ij), ...]`` for the aggregation module."""
+    members = snapshot_class_members(snapshot)
+    return {
+        class_id: [(proto, support) for _, proto, support in entries]
+        for class_id, entries in members.items()
+    }
+
+
+def aggregate_snapshot(snapshot: Mapping, beta) -> dict[int, list[float]]:
+    """Aggregate the snapshot's local prototypes with the proposed mode."""
+    return aggregate_class_prototypes(
+        snapshot_contributions(snapshot), _SUPPORT_WEIGHTED, beta=beta
+    )
+
+
+def _l2_distance(a: Sequence[float], b: Sequence[float]) -> float:
+    if len(a) != len(b):
+        raise ValueError("prototype vectors must share the same dimension")
+    return math.sqrt(math.fsum((float(x) - float(y)) ** 2 for x, y in zip(a, b)))
+
+
+def _norm(vec: Sequence[float]) -> float:
+    return math.sqrt(math.fsum(float(x) ** 2 for x in vec))
+
+
+def prototype_effect_rows(
+    snapshot: Mapping,
+    betas: Sequence[float] = DEFAULT_BETAS,
+    class_names: Mapping[int, str] | None = None,
+) -> list[dict]:
+    """Per-class rows of ``|| Cbar_j(beta) - Cbar_j(0) ||`` and prototype norms.
+
+    One row per ``(class_id, beta)``. ``displacement`` is the direct aggregation
+    effect relative to the equal-weight ``beta = 0`` aggregate on the *same*
+    contributor set. A large displacement is not evidence of improved detection.
+    """
+    if not betas:
+        raise ValueError("at least one beta is required")
+    contributions = snapshot_contributions(snapshot)
+    baseline = aggregate_class_prototypes(contributions, _SUPPORT_WEIGHTED, beta=0.0)
+    class_names = class_names or {}
+    members_by_class = snapshot_class_members(snapshot)
+
+    rows: list[dict] = []
+    for class_id in sorted(baseline):
+        members = [
+            support
+            for _, _, support in members_by_class.get(class_id, [])
+            if support is not None and float(support) > 0.0
+        ]
+        base_norm = _norm(baseline[class_id])
+        for beta in betas:
+            cbar = aggregate_snapshot(snapshot, beta)[class_id]
+            displacement = _l2_distance(cbar, baseline[class_id])
+            rows.append(
+                {
+                    "class_id": class_id,
+                    "class_name": class_names.get(class_id),
+                    "beta": float(beta),
+                    "n_contributors": len(members),
+                    "support_min": min(members) if members else None,
+                    "support_max": max(members) if members else None,
+                    "support_mean": (sum(members) / len(members)) if members else None,
+                    "prototype_norm": _norm(cbar),
+                    "reference_norm_beta0": base_norm,
+                    "displacement_from_beta0": displacement,
+                    "relative_displacement": (
+                        displacement / base_norm if base_norm > 0.0 else None
+                    ),
+                }
+            )
+    return rows
+
+
+def client_weight_rows(
+    snapshot: Mapping,
+    betas: Sequence[float] = DEFAULT_BETAS,
+    class_names: Mapping[int, str] | None = None,
+) -> list[dict]:
+    """Per-``(class, beta, client)`` normalised aggregation weights.
+
+    Only clients in ``S_j`` (positive support and a valid prototype) appear; the
+    reported weights sum to one within each ``(class, beta)`` group.
+    """
+    if not betas:
+        raise ValueError("at least one beta is required")
+    members = snapshot_class_members(snapshot)
+    class_names = class_names or {}
+
+    rows: list[dict] = []
+    for class_id in sorted(members):
+        valid = [
+            (client_id, support)
+            for client_id, _, support in members[class_id]
+            if support is not None and float(support) > 0.0
+        ]
+        if not valid:
+            continue
+        for beta in betas:
+            weights = normalized_support_weights(
+                [support for _, support in valid], beta
+            )
+            for (client_id, support), weight in zip(valid, weights):
+                rows.append(
+                    {
+                        "class_id": class_id,
+                        "class_name": class_names.get(class_id),
+                        "beta": float(beta),
+                        "client_id": client_id,
+                        "N_ij": float(support),
+                        "weight": weight,
+                    }
+                )
+    return rows

src/protean_research/federated/round.py

@@ -38,6 +38,8 @@ class TrainingConfig:
     dense_units: int = 128
     snapshot_round: Optional[int] = 0
     device: str = "cpu"
+    aggregation_mode: str = "original_protean"
+    beta: Optional[float] = None
 
 
 def _evaluate(model, features, labels, global_prototypes, num_classes, rare_class_ids=None):
@@ -122,6 +124,8 @@ def run_federated(
             [u.prototypes for u in updates],
             [u.supports for u in updates],
             num_classes,
+            mode=config.aggregation_mode,
+            beta=config.beta,
         )
 
         if config.snapshot_round is not None and round_id == config.snapshot_round:

src/protean_research/federated/server.py

@@ -25,10 +25,16 @@ def aggregate_prototypes(
     client_supports: Sequence[Mapping[int, int]],
     num_classes: int,
     mode: AggregationMode | str = AggregationMode.ORIGINAL_PROTEAN,
+    beta=None,
 ):
-    """Aggregate client prototypes into ``{class_id: Cbar_j}`` (baseline)."""
+    """Aggregate client prototypes into ``{class_id: Cbar_j}``.
+
+    ``mode`` selects equal ``original_protean`` averaging or the proposed
+    ``support_weighted`` aggregation with exponent ``beta``. The contributor set
+    ``S_j`` is identical for both modes.
+    """
     return collect_global_prototypes(
-        client_prototypes, client_supports, num_classes, mode=mode
+        client_prototypes, client_supports, num_classes, mode=mode, beta=beta
     )
 
 

src/protean_research/models/prototypes.py

@@ -41,18 +41,20 @@ def collect_global_prototypes(
     client_supports: Sequence[Mapping[int, int]],
     num_classes: int,
     mode: AggregationMode | str = AggregationMode.ORIGINAL_PROTEAN,
+    beta=None,
 ):
     """Server-side aggregation of per-client prototypes into global prototypes.
 
     Clients that did not observe a class simply do not appear in
     ``client_prototypes`` for that class, so ``S_j`` is built implicitly and
-    missing classes are never inserted as zero vectors.
+    missing classes are never inserted as zero vectors. ``beta`` is forwarded to
+    the aggregation module and is only valid for ``support_weighted``.
     """
     contributions: dict[int, list] = {c: [] for c in range(num_classes)}
     for prototypes, supports in zip(client_prototypes, client_supports):
         for class_id, proto in prototypes.items():
             contributions[int(class_id)].append((proto, supports.get(int(class_id), 0)))
-    return aggregate_class_prototypes(contributions, mode)
+    return aggregate_class_prototypes(contributions, mode, beta=beta)
 
 
 def nearest_prototype_predict(embeddings, global_prototypes: Mapping[int, object]):