groupoid 0.1.0.dev1__py3-none-any.whl

groupoid/__init__.py

@@ -0,0 +1,31 @@
+"""GROUPOID: Groupoid-based federated learning with Riemannian geometry."""
+
+__version__ = "0.1.0.dev1"
+
+from groupoid.aggregation import (
+    DisconnectedClientGraphError,
+    FederatedRound,
+    TransportGroupoidAggregator,
+)
+from groupoid.cohomology import IncompleteCocycleError, compute_h1
+from groupoid.groupoid import CompositionError, Morphism, compose, inverse
+from groupoid.laplacian import SpectralSummary, spectral_analysis
+from groupoid.manifold import karcher_mean
+from groupoid.sheaf import Sheaf
+
+__all__ = [
+    "__version__",
+    "CompositionError",
+    "DisconnectedClientGraphError",
+    "FederatedRound",
+    "IncompleteCocycleError",
+    "Morphism",
+    "Sheaf",
+    "SpectralSummary",
+    "TransportGroupoidAggregator",
+    "compose",
+    "compute_h1",
+    "inverse",
+    "karcher_mean",
+    "spectral_analysis",
+]

groupoid/aggregation.py

@@ -0,0 +1,234 @@
+"""Federated aggregation via groupoid transport and Karcher mean.
+
+This module ties together the transport groupoid, sheaf consistency,
+and Riemannian aggregation into a complete federated learning round.
+The key insight is that naive parameter averaging fails when client
+models live on different tangent spaces. Instead, we:
+
+1. Transport all client parameters to a common base point via the
+   groupoid morphisms.
+2. Check cohomological consistency (H^1) to detect irreconcilable
+   disagreements before aggregation.
+3. Compute the Karcher mean on the manifold of the transported
+   parameters.
+4. Transport the aggregated model back to each client's local frame.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import networkx as nx
+import numpy as np
+import numpy.typing as npt
+from loguru import logger
+
+from groupoid.cohomology import compute_h1
+from groupoid.groupoid import Morphism, compose, inverse
+from groupoid.manifold import karcher_mean
+
+
+class DisconnectedClientGraphError(Exception):
+    """Raised when the client graph has no transport path to the base node.
+
+    Aggregation transports every client's parameters to the base node along
+    a path in the client graph. If a client lies in a different connected
+    component from the base node, no such path exists and the parameters
+    cannot be aligned. This error replaces the raw networkx
+    ``NetworkXNoPath`` with a domain-meaningful message naming the
+    unreachable node and the base.
+    """
+
+
+@dataclass
+class FederatedRound:
+    """Result of a single federated aggregation round."""
+
+    global_params: npt.NDArray[np.float64]
+    local_updates: dict[str, npt.NDArray[np.float64]]
+    h1_norm: float
+    is_consistent: bool
+    transport_residuals: dict[str, float]
+    round_idx: int = 0
+
+
+@dataclass
+class TransportGroupoidAggregator:
+    """Federated aggregator using groupoid transport and Riemannian geometry.
+
+    Given a network of clients with local model parameters on a Riemannian
+    manifold, this aggregator:
+    - Uses parallel transport (groupoid morphisms) to align parameters
+    - Checks cohomological consistency before aggregation
+    - Computes intrinsic Karcher mean on the parameter manifold
+    - Distributes the result back via inverse transport
+
+    Parameters
+    ----------
+    manifold
+        A geomstats manifold instance for the parameter space.
+    graph
+        Directed graph of the client network.
+    base_node
+        The node to which all parameters are transported for aggregation.
+    consistency_threshold
+        Maximum H^1 norm before raising a consistency warning.
+    """
+
+    manifold: object
+    graph: nx.DiGraph
+    base_node: str
+    consistency_threshold: float = 1e-6
+    morphisms: dict[tuple[str, str], Morphism] = field(default_factory=dict)
+    _round_idx: int = field(default=0, init=False)
+
+    def register_transport(
+        self, source: str, target: str, matrix: npt.NDArray[np.float64]
+    ) -> None:
+        """Register a transport map between two clients."""
+        self.morphisms[(source, target)] = Morphism(
+            source=source,
+            target=target,
+            transport_map=matrix,
+        )
+        logger.debug("Registered transport {} -> {}", source, target)
+
+    def _get_transport_to_base(self, node: str) -> npt.NDArray[np.float64] | None:
+        """Compute the composite transport map from node to base_node.
+
+        Uses shortest path in the graph and composes morphisms along it.
+        Returns None if node == base_node (identity transport).
+        """
+        if node == self.base_node:
+            return None
+
+        try:
+            path = nx.shortest_path(self.graph.to_undirected(), node, self.base_node)
+        except nx.NetworkXNoPath as e:
+            raise DisconnectedClientGraphError(
+                f"client graph is disconnected: no transport path from "
+                f"{node} to base {self.base_node}"
+            ) from e
+        composite: Morphism | None = None
+
+        for i in range(len(path) - 1):
+            src, tgt = path[i], path[i + 1]
+            if (src, tgt) in self.morphisms:
+                m = self.morphisms[(src, tgt)]
+            elif (tgt, src) in self.morphisms:
+                m = inverse(self.morphisms[(tgt, src)])
+            else:
+                raise ValueError(f"No transport map for edge ({src}, {tgt})")
+
+            composite = m if composite is None else compose(composite, m)
+
+        return composite.transport_map if composite is not None else None
+
+    def check_consistency(self, client_params: dict[str, npt.NDArray[np.float64]]) -> float:
+        """Check cohomological consistency of current transport maps.
+
+        Returns the H^1 norm. A value near zero indicates that the
+        local models can be consistently aggregated.
+        """
+        transport_maps = {(m.source, m.target): m.transport_map for m in self.morphisms.values()}
+        h1 = compute_h1(self.graph, transport_maps)
+        logger.info("Cohomological consistency check: H^1 = {:.2e}", h1)
+        return h1
+
+    def aggregate(
+        self,
+        client_params: dict[str, npt.NDArray[np.float64]],
+        weights: dict[str, float] | None = None,
+    ) -> FederatedRound:
+        """Run one round of groupoid-aware federated aggregation.
+
+        Parameters
+        ----------
+        client_params
+            Map from client node ID to local model parameters.
+        weights
+            Optional per-client weights for the Karcher mean.
+
+        Returns
+        -------
+        FederatedRound
+            The aggregation result including global params, consistency
+            metrics, and per-client local updates.
+        """
+        self._round_idx += 1
+        logger.info("Starting aggregation round {}", self._round_idx)
+
+        h1 = self.check_consistency(client_params)
+        is_consistent = h1 < self.consistency_threshold
+
+        if not is_consistent:
+            logger.warning(
+                "H^1 = {:.2e} exceeds threshold {:.2e}, "
+                "aggregation may produce inconsistent results",
+                h1,
+                self.consistency_threshold,
+            )
+
+        # Transport all client params to base node frame
+        transported = {}
+        transport_residuals = {}
+        for node, params in client_params.items():
+            if node == self.base_node:
+                transported[node] = params
+                transport_residuals[node] = 0.0
+            else:
+                # _get_transport_to_base returns None only for node == base_node
+                # (excluded by this else); a disconnected graph raises
+                # NetworkXNoPath. This guard is therefore defensively
+                # unreachable.
+                T = self._get_transport_to_base(node)
+                if T is None:  # pragma: no cover - unreachable defensive guard (see above)
+                    raise ValueError(f"No transport path from {node} to {self.base_node}")
+                transported_params = T @ params
+                transported[node] = transported_params
+                transport_residuals[node] = float(
+                    np.linalg.norm(T @ T.T - np.eye(T.shape[0]), "fro")
+                )
+
+        # Stack transported params and compute Karcher mean
+        nodes = sorted(transported.keys())
+        param_stack = np.stack([transported[n] for n in nodes])
+
+        if weights is not None:
+            w = np.array([weights.get(n, 1.0) for n in nodes])
+            w = w / w.sum()
+        else:
+            w = None
+
+        global_params = karcher_mean(self.manifold, param_stack, weights=w)
+
+        # Transport global params back to each client's local frame
+        local_updates = {}
+        for node in client_params:
+            if node == self.base_node:
+                local_updates[node] = global_params
+            else:
+                # Same defensively-unreachable guard as in the forward
+                # transport loop above.
+                T = self._get_transport_to_base(node)
+                if T is None:  # pragma: no cover - unreachable defensive guard (see above)
+                    raise ValueError(f"No transport path from {node} to {self.base_node}")
+                T_inv = np.linalg.inv(T)
+                local_updates[node] = T_inv @ global_params
+
+        result = FederatedRound(
+            global_params=global_params,
+            local_updates=local_updates,
+            h1_norm=h1,
+            is_consistent=is_consistent,
+            transport_residuals=transport_residuals,
+            round_idx=self._round_idx,
+        )
+
+        logger.info(
+            "Round {} complete: H^1={:.2e}, consistent={}",
+            self._round_idx,
+            h1,
+            is_consistent,
+        )
+        return result

groupoid/cohomology.py

@@ -0,0 +1,99 @@
+"""First cohomology of transport groupoids."""
+
+from __future__ import annotations
+
+import networkx as nx
+import numpy as np
+import numpy.typing as npt
+from loguru import logger
+
+
+class IncompleteCocycleError(Exception):
+    """Raised when a cocycle is incompletely specified.
+
+    The holonomy around a cycle is the ordered product of the transport
+    maps on *every* edge of that cycle. If any edge on a cycle has no
+    transport map (in either direction), the holonomy is undefined: a
+    product over a strict subset of the cycle's edges is mathematically
+    meaningless and must not be reported as an H^1 (in)consistency. This
+    error names the missing edge so the caller can supply it.
+    """
+
+
+def compute_h1(
+    graph: nx.DiGraph,
+    transport_maps: dict[tuple[str, str], npt.NDArray[np.float64]],
+) -> float:
+    """Compute the H^1 cohomology norm of a transport cocycle.
+
+    Returns 0.0 (up to numerical tolerance) when the cocycle is a
+    coboundary, meaning the local data is globally consistent.
+
+    The H^1 obstruction on a cycle is the deviation from the identity of
+    the holonomy, i.e. the ordered product of the restriction/transport
+    maps around the *entire* cycle. The cocycle must therefore be fully
+    specified: every edge of every basis cycle needs a transport map in
+    one direction or the other (the reverse direction is inverted). A
+    missing edge map leaves the holonomy undefined, so this function
+    raises :class:`IncompleteCocycleError` rather than silently forming a
+    partial product.
+
+    Parameters
+    ----------
+    graph : nx.DiGraph
+        The groupoid nerve (directed graph of nodes/clients).
+    transport_maps : dict
+        Maps (source, target) edge tuples to transport matrices.
+
+    Returns
+    -------
+    float
+        The cohomology norm. Zero indicates global consistency.
+
+    Raises
+    ------
+    IncompleteCocycleError
+        If any edge on a basis cycle has no transport map in either
+        direction, so the cycle's holonomy is undefined.
+    """
+    undirected = graph.to_undirected()
+    cycles = nx.cycle_basis(undirected)
+
+    if not cycles:
+        logger.debug("No cycles in graph, H^1 = 0 trivially")
+        return 0.0
+
+    max_holonomy_norm = 0.0
+
+    for cycle in cycles:
+        n = len(cycle)
+        # Resolve every edge's transport map first, raising on the first
+        # missing one so a partial product is never formed. The reverse
+        # orientation is inverted. Every cycle-basis cycle has length >= 3, so
+        # the resulting product is always fully formed.
+        edge_maps = []
+        for i in range(n):
+            u = cycle[i]
+            v = cycle[(i + 1) % n]
+
+            if (u, v) in transport_maps:
+                edge_maps.append(transport_maps[(u, v)])
+            elif (v, u) in transport_maps:
+                edge_maps.append(np.linalg.inv(transport_maps[(v, u)]))
+            else:
+                raise IncompleteCocycleError(
+                    f"Incomplete cocycle: no transport map for edge ({u}, {v}) "
+                    f"on cycle {cycle}; holonomy is undefined. Supply the edge "
+                    f"map in either direction before computing H^1."
+                )
+
+        holonomy = edge_maps[0]
+        for T in edge_maps[1:]:
+            holonomy = T @ holonomy
+
+        dim = holonomy.shape[0]
+        deviation = float(np.linalg.norm(holonomy - np.eye(dim), ord="fro"))
+        max_holonomy_norm = max(max_holonomy_norm, deviation)
+
+    logger.debug("H^1 norm = {:.6e}", max_holonomy_norm)
+    return max_holonomy_norm

groupoid/groupoid.py

@@ -0,0 +1,85 @@
+"""Groupoid structure for federated learning transport maps."""
+
+from __future__ import annotations
+
+import numpy as np
+import numpy.typing as npt
+from loguru import logger
+from pydantic import BaseModel, ConfigDict
+
+
+class Morphism(BaseModel):
+    """A morphism in the transport groupoid.
+
+    Represents a transport map between two nodes (clients) in the
+    federated learning network.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    source: str
+    target: str
+    transport_map: npt.NDArray[np.float64]
+
+    def __repr__(self) -> str:
+        return f"Morphism({self.source} -> {self.target})"
+
+
+class CompositionError(Exception):
+    """Raised when morphism composition is not defined."""
+
+
+def compose(f: Morphism, g: Morphism) -> Morphism:
+    """Compose two morphisms f and g (f followed by g).
+
+    Requires f.target == g.source (category composition law).
+
+    Parameters
+    ----------
+    f : Morphism
+        First morphism (applied first).
+    g : Morphism
+        Second morphism (applied second).
+
+    Returns
+    -------
+    Morphism
+        The composed morphism from f.source to g.target.
+
+    Raises
+    ------
+    CompositionError
+        If f.target != g.source.
+    """
+    if f.target != g.source:
+        raise CompositionError(f"Cannot compose: {f.target} != {g.source}")
+
+    logger.debug("Composing {} with {}", f, g)
+    composed_map = g.transport_map @ f.transport_map
+    return Morphism(
+        source=f.source,
+        target=g.target,
+        transport_map=composed_map,
+    )
+
+
+def inverse(f: Morphism) -> Morphism:
+    """Compute the inverse of a morphism.
+
+    Parameters
+    ----------
+    f : Morphism
+        The morphism to invert.
+
+    Returns
+    -------
+    Morphism
+        The inverse morphism from f.target to f.source.
+    """
+    logger.debug("Inverting {}", f)
+    inv_map = np.linalg.inv(f.transport_map)
+    return Morphism(
+        source=f.target,
+        target=f.source,
+        transport_map=inv_map,
+    )

groupoid/laplacian.py

@@ -0,0 +1,201 @@
+"""Sheaf Laplacian for federated consensus and spectral analysis.
+
+The sheaf Laplacian generalizes the graph Laplacian by incorporating
+the restriction maps of a cellular sheaf. Its spectrum reveals the
+structure of the agreement space:
+
+- Kernel of L = space of global sections (consistent models)
+- Smallest nonzero eigenvalue = algebraic connectivity of the sheaf
+  (how fast consensus can be reached)
+- Spectral gap = robustness of the consensus to perturbation
+
+In federated learning, the sheaf Laplacian governs the diffusion
+process that drives local models toward global consistency.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import numpy.typing as npt
+from loguru import logger
+
+from groupoid.sheaf import Sheaf
+
+
+@dataclass
+class SpectralSummary:
+    """Spectral decomposition of the sheaf Laplacian."""
+
+    eigenvalues: npt.NDArray[np.float64]
+    eigenvectors: npt.NDArray[np.float64]
+    spectral_gap: float
+    algebraic_connectivity: float
+    kernel_dimension: int
+    consensus_rate: float
+
+
+def build_sheaf_laplacian(sheaf: Sheaf, stalk_dim: int) -> npt.NDArray[np.float64]:
+    """Build the sheaf Laplacian matrix.
+
+    For a sheaf F on graph G with n nodes and stalk dimension d,
+    the sheaf Laplacian is a (n*d) x (n*d) block matrix defined as:
+
+        L_F = delta^T @ delta
+
+    where delta is the connection coboundary. For each edge (u, v) with
+    restriction (transport) map R = R_{uv}: stalk(u) -> stalk(v), the
+    coboundary acts as (delta x)_{(u,v)} = x_v - R_{uv} x_u, so
+    L = delta^T @ delta has blocks (summed over incident edges):
+
+        L[u,u] += R_{uv}^T @ R_{uv}    (source diagonal)
+        L[v,v] += I                    (target diagonal)
+        L[u,v] += -R_{uv}^T            (off-diagonal)
+        L[v,u] += -R_{uv}              (off-diagonal)
+
+    L is symmetric positive semi-definite for ANY restriction maps (it is
+    delta^T delta); its kernel is the space of transport-consistent global
+    sections (x_v = R_{uv} x_u on every edge).
+
+    Parameters
+    ----------
+    sheaf
+        A Sheaf instance with restriction maps set.
+    stalk_dim
+        Dimension of each stalk (vector space at each node).
+
+    Returns
+    -------
+    np.ndarray
+        The sheaf Laplacian matrix of shape (n*d, n*d).
+    """
+    nodes = sorted(sheaf.graph.nodes())
+    n = len(nodes)
+    node_idx = {node: i for i, node in enumerate(nodes)}
+    N = n * stalk_dim
+
+    L = np.zeros((N, N))
+
+    for u, v in sheaf.graph.edges():
+        i, j = node_idx[u], node_idx[v]
+        R = sheaf.get_restriction_map(u, v)
+        i_slice = slice(i * stalk_dim, (i + 1) * stalk_dim)
+        j_slice = slice(j * stalk_dim, (j + 1) * stalk_dim)
+
+        # L = delta^T delta for coboundary (delta x)_(u,v) = x_v - R_uv x_u:
+        L[i_slice, i_slice] += R.T @ R  # source diagonal: R^T R
+        L[j_slice, j_slice] += np.eye(stalk_dim)  # target diagonal: I
+        L[i_slice, j_slice] += -R.T  # off-diagonal: -R^T
+        L[j_slice, i_slice] += -R  # off-diagonal: -R
+
+    logger.debug("Built sheaf Laplacian: {}x{} ({} nodes, stalk_dim={})", N, N, n, stalk_dim)
+    return L
+
+
+def spectral_analysis(
+    sheaf: Sheaf,
+    stalk_dim: int,
+    tol: float = 1e-10,
+) -> SpectralSummary:
+    """Compute spectral decomposition of the sheaf Laplacian.
+
+    Parameters
+    ----------
+    sheaf
+        A Sheaf instance with restriction maps.
+    stalk_dim
+        Dimension of each stalk.
+    tol
+        Tolerance for identifying zero eigenvalues.
+
+    Returns
+    -------
+    SpectralSummary
+        Full spectral summary including connectivity and consensus rate.
+    """
+    L = build_sheaf_laplacian(sheaf, stalk_dim)
+    eigenvalues, eigenvectors = np.linalg.eigh(L)
+
+    # Sort by magnitude
+    idx = np.argsort(eigenvalues)
+    eigenvalues = eigenvalues[idx]
+    eigenvectors = eigenvectors[:, idx]
+
+    # Kernel dimension (number of zero eigenvalues)
+    kernel_dim = int(np.sum(np.abs(eigenvalues) < tol))
+
+    # Spectral gap and algebraic connectivity
+    nonzero_eigs = eigenvalues[np.abs(eigenvalues) >= tol]
+    if len(nonzero_eigs) > 0:
+        algebraic_connectivity = float(nonzero_eigs[0])
+        spectral_gap = float(nonzero_eigs[0])
+    else:
+        algebraic_connectivity = 0.0
+        spectral_gap = 0.0
+
+    # Consensus rate: exponential convergence rate of sheaf diffusion
+    # x(t+1) = (I - epsilon * L) @ x(t), converges as exp(-lambda_1 * t)
+    consensus_rate = algebraic_connectivity
+
+    logger.info(
+        "Spectral analysis: kernel_dim={}, spectral_gap={:.4f}, connectivity={:.4f}",
+        kernel_dim,
+        spectral_gap,
+        algebraic_connectivity,
+    )
+
+    return SpectralSummary(
+        eigenvalues=eigenvalues,
+        eigenvectors=eigenvectors,
+        spectral_gap=spectral_gap,
+        algebraic_connectivity=algebraic_connectivity,
+        kernel_dimension=kernel_dim,
+        consensus_rate=consensus_rate,
+    )
+
+
+def sheaf_diffusion_step(
+    sheaf: Sheaf,
+    sections: dict[str, npt.NDArray[np.float64]],
+    stalk_dim: int,
+    step_size: float = 0.1,
+) -> dict[str, npt.NDArray[np.float64]]:
+    """One step of sheaf diffusion (Laplacian smoothing).
+
+    Drives local sections toward global consistency by flowing along
+    the negative gradient of the sheaf Laplacian energy:
+
+        E(x) = x^T L_F x = sum_{(i,j)} ||R_{ij} x_i - x_j||^2
+
+    Parameters
+    ----------
+    sheaf
+        Sheaf with restriction maps.
+    sections
+        Current section values at each node.
+    stalk_dim
+        Dimension of each stalk.
+    step_size
+        Diffusion step size (must be < 1/lambda_max for stability).
+
+    Returns
+    -------
+    dict[str, npt.NDArray[np.float64]]
+        Updated section values after one diffusion step.
+    """
+    L = build_sheaf_laplacian(sheaf, stalk_dim)
+    nodes = sorted(sheaf.graph.nodes())
+
+    # Stack sections into vector
+    x = np.concatenate([sections[n] for n in nodes])
+
+    # Diffusion step: x' = x - step_size * L @ x
+    x_new = x - step_size * L @ x
+
+    # Unstack
+    result = {}
+    for idx, node in enumerate(nodes):
+        result[node] = x_new[idx * stalk_dim : (idx + 1) * stalk_dim]
+
+    return result