groupoid 0.1.0.dev1__py3-none-any.whl

groupoid/manifold.py

@@ -0,0 +1,71 @@
+"""Riemannian manifold operations for GROUPOID."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import numpy.typing as npt
+from loguru import logger
+
+# TYPE_CHECKING is always False at runtime; import is for static type checkers.
+if TYPE_CHECKING:  # pragma: no cover
+    from geomstats.geometry.base import LevelSet
+
+
+def karcher_mean(
+    manifold: LevelSet,
+    points: npt.NDArray[np.float64],
+    weights: npt.NDArray[np.float64] | None = None,
+    max_iter: int = 100,
+    tol: float = 1e-6,
+) -> npt.NDArray[np.float64]:
+    """Compute the Karcher (Frechet) mean on a Riemannian manifold.
+
+    Parameters
+    ----------
+    manifold : geomstats manifold
+        A geomstats manifold instance with a metric.
+    points : np.ndarray
+        Array of shape (n_points, *point_shape) on the manifold.
+    weights : np.ndarray or None
+        Optional weights for the mean computation.
+    max_iter : int
+        Maximum iterations for gradient descent.
+    tol : float
+        Convergence tolerance.
+
+    Returns
+    -------
+    np.ndarray
+        The Karcher mean point on the manifold.
+    """
+    from geomstats.learning.frechet_mean import FrechetMean
+
+    n_points = points.shape[0]
+    logger.debug(
+        "Computing Karcher mean of {} points on {}",
+        n_points,
+        type(manifold).__name__,
+    )
+
+    estimator = FrechetMean(manifold)
+
+    # Forward convergence controls to the underlying gradient-descent
+    # optimizer when the installed geomstats version exposes it. Older
+    # versions without an `optimizer` attribute fall back to their defaults.
+    optimizer = getattr(estimator, "optimizer", None)
+    if optimizer is not None:
+        if hasattr(optimizer, "max_iter"):
+            optimizer.max_iter = max_iter
+        if hasattr(optimizer, "epsilon"):
+            optimizer.epsilon = tol
+
+    if weights is not None:
+        estimator.fit(points, weights=weights)
+    else:
+        estimator.fit(points)
+
+    mean: npt.NDArray[np.float64] = estimator.estimate_
+    logger.debug("Karcher mean converged")
+    return mean

groupoid/optimizer.py

@@ -0,0 +1,208 @@
+"""Riemannian optimizers for federated learning on manifolds.
+
+Standard SGD and Adam assume Euclidean parameter spaces. When model
+parameters live on a Riemannian manifold (e.g., Stiefel manifold for
+orthogonal layers, SPD manifold for covariance matrices), we need
+Riemannian counterparts that respect the geometry.
+
+This module provides:
+- Riemannian SGD with retraction
+- Riemannian Adam with exponential map updates
+- Adaptive learning rate based on sectional curvature
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+from loguru import logger
+
+
+@dataclass
+class RiemannianSGD:
+    """Riemannian stochastic gradient descent.
+
+    Updates parameters by computing the Riemannian gradient (projection
+    of Euclidean gradient onto tangent space) and retracting back to
+    the manifold via the exponential map.
+
+    Parameters
+    ----------
+    manifold
+        A geomstats manifold instance.
+    lr
+        Learning rate.
+    momentum
+        Momentum coefficient (0 = no momentum).
+    """
+
+    manifold: Any
+    lr: float = 0.01
+    momentum: float = 0.0
+    _velocity: npt.NDArray[np.float64] | None = field(default=None, init=False, repr=False)
+
+    def step(
+        self, point: npt.NDArray[np.float64], euclidean_grad: npt.NDArray[np.float64]
+    ) -> npt.NDArray[np.float64]:
+        """Perform one optimization step.
+
+        Parameters
+        ----------
+        point
+            Current point on the manifold.
+        euclidean_grad
+            Euclidean gradient (will be projected to tangent space).
+
+        Returns
+        -------
+        np.ndarray
+            Updated point on the manifold.
+        """
+        # Project gradient onto tangent space
+        riemannian_grad = self.manifold.to_tangent(euclidean_grad, point)
+
+        # Apply momentum
+        vel: npt.NDArray[np.float64]
+        if self.momentum > 0:
+            if self._velocity is None:
+                vel = riemannian_grad
+            else:
+                vel = self.manifold.to_tangent(
+                    self.momentum * self._velocity + riemannian_grad, point
+                )
+            self._velocity = vel
+            update = -self.lr * vel
+        else:
+            update = -self.lr * riemannian_grad
+
+        # Retract to manifold via exponential map
+        new_point: npt.NDArray[np.float64] = self.manifold.metric.exp(update, point)
+
+        return new_point
+
+
+@dataclass
+class RiemannianAdam:
+    """Riemannian Adam optimizer.
+
+    Adapts the Adam optimizer to Riemannian manifolds by maintaining
+    exponential moving averages of the Riemannian gradient and its
+    squared norm, with updates via the exponential map.
+
+    Parameters
+    ----------
+    manifold
+        A geomstats manifold instance.
+    lr
+        Learning rate.
+    beta1
+        Exponential decay rate for first moment.
+    beta2
+        Exponential decay rate for second moment.
+    eps
+        Small constant for numerical stability.
+    """
+
+    manifold: Any
+    lr: float = 0.001
+    beta1: float = 0.9
+    beta2: float = 0.999
+    eps: float = 1e-8
+    _m: npt.NDArray[np.float64] | None = field(default=None, init=False, repr=False)
+    _v: float = field(default=0.0, init=False, repr=False)
+    _t: int = field(default=0, init=False, repr=False)
+
+    def step(
+        self, point: npt.NDArray[np.float64], euclidean_grad: npt.NDArray[np.float64]
+    ) -> npt.NDArray[np.float64]:
+        """Perform one optimization step.
+
+        Parameters
+        ----------
+        point
+            Current point on the manifold.
+        euclidean_grad
+            Euclidean gradient.
+
+        Returns
+        -------
+        np.ndarray
+            Updated point on the manifold.
+        """
+        self._t += 1
+
+        # Riemannian gradient
+        grad = self.manifold.to_tangent(euclidean_grad, point)
+        grad_norm_sq = float(np.sum(grad**2))
+
+        # Update biased first moment (tangent vector)
+        first_moment: npt.NDArray[np.float64]
+        if self._m is None:
+            first_moment = grad.copy()
+        else:
+            first_moment = self.manifold.to_tangent(
+                self.beta1 * self._m + (1 - self.beta1) * grad, point
+            )
+        self._m = first_moment
+
+        # Update biased second moment (scalar, norm-based)
+        self._v = self.beta2 * self._v + (1 - self.beta2) * grad_norm_sq
+
+        # Bias correction
+        m_hat = first_moment / (1 - self.beta1**self._t)
+        v_hat = self._v / (1 - self.beta2**self._t)
+
+        # Adaptive update
+        update = -self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
+        update = self.manifold.to_tangent(update, point)
+
+        new_point: npt.NDArray[np.float64] = self.manifold.metric.exp(update, point)
+        return new_point
+
+
+def curvature_adaptive_lr(
+    manifold: Any,  # geomstats manifold; no upstream type stubs
+    point: npt.NDArray[np.float64],
+    base_lr: float,
+    tangent_vec: npt.NDArray[np.float64],
+) -> float:
+    """Adapt learning rate based on local sectional curvature.
+
+    In regions of high positive curvature, geodesics converge and
+    we should take smaller steps. In regions of negative curvature,
+    geodesics diverge and we can take larger steps.
+
+    Parameters
+    ----------
+    manifold
+        A geomstats manifold with a curvature method.
+    point
+        Current point on the manifold.
+    base_lr
+        Base learning rate to adapt.
+    tangent_vec
+        Direction of the update.
+
+    Returns
+    -------
+    float
+        Adapted learning rate.
+    """
+    try:
+        if hasattr(manifold.metric, "sectional_curvature"):
+            # Use a random orthogonal vector for the plane
+            random_vec = manifold.to_tangent(np.random.randn(*tangent_vec.shape), point)
+            kappa = manifold.metric.sectional_curvature(tangent_vec, random_vec, point)
+            kappa = float(np.mean(kappa)) if hasattr(kappa, "__len__") else float(kappa)
+
+            # Scale: lr / (1 + max(kappa, 0)) damps in positive curvature
+            adapted: float = base_lr / (1.0 + max(kappa, 0.0))
+            logger.debug("Curvature-adapted LR: {:.6f} (kappa={:.4f})", adapted, kappa)
+            return adapted
+    except (AttributeError, NotImplementedError):
+        pass
+
+    return base_lr

groupoid/persistence.py

@@ -0,0 +1,195 @@
+"""Persistent homology for federated model divergence detection.
+
+Applies topological data analysis to track how the distribution of
+client model parameters evolves across federated rounds. The key
+invariant: if the persistence diagram is stable across rounds, the
+federation is converging; if new topological features appear, clients
+are diverging into distinct clusters.
+
+Uses Vietoris-Rips filtration on the space of model parameters to
+compute persistence diagrams, then measures stability via bottleneck
+and Wasserstein distances.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import numpy.typing as npt
+from loguru import logger
+
+
+@dataclass
+class PersistenceSummary:
+    """Summary of persistent homology computation.
+
+    The ``diagram`` array is dimension-labelled: it has shape ``(k, 3)``
+    where column 0 is birth, column 1 is death, and column 2 is the
+    homology dimension (0 for H0/connected components, 1 for H1/loops,
+    ...). Retaining the dimension label is what lets ``track_divergence``
+    compare like-with-like (H0 against H0) instead of pooling features
+    from different homology dimensions into one bottleneck computation.
+    Use :meth:`diagram_for_dim` to extract a single dimension's bars.
+    """
+
+    betti_0: int
+    betti_1: int
+    total_persistence: float
+    max_persistence: float
+    diagram: npt.NDArray[np.float64]
+    bottleneck_to_previous: float | None = None
+
+    def diagram_for_dim(self, dim: int, *, finite_only: bool = False) -> npt.NDArray[np.float64]:
+        """Return the ``(birth, death)`` bars of a single homology dimension.
+
+        Parameters
+        ----------
+        dim
+            Homology dimension to select (0 = H0, 1 = H1, ...).
+        finite_only
+            If True, drop bars that die at infinity (the essential
+            classes). H0 always carries exactly one infinite bar (the
+            whole space's single surviving component under ``thresh=inf``);
+            excluding it leaves the finite bars, which for H0 are exactly
+            the minimum-spanning-tree edge weights of the point cloud.
+
+        Returns
+        -------
+        np.ndarray
+            Array of shape ``(m, 2)`` with columns ``(birth, death)``.
+            Empty diagrams return shape ``(0, 2)``.
+        """
+        if self.diagram.shape[0] == 0:
+            return np.empty((0, 2))
+        mask = self.diagram[:, 2] == dim
+        if finite_only:
+            mask = mask & (self.diagram[:, 1] < np.inf)
+        # np.asarray pins the return type (numpy fancy-indexing is typed as
+        # Any, which would trip mypy's warn_return_any); it is a no-copy view
+        # here since the slice is already an ndarray.
+        return np.asarray(self.diagram[mask][:, :2])
+
+
+def compute_persistence(
+    points: npt.NDArray[np.float64],
+    max_dim: int = 1,
+    max_edge_length: float = np.inf,
+) -> PersistenceSummary:
+    """Compute persistent homology of a point cloud.
+
+    Parameters
+    ----------
+    points
+        Array of shape (n_points, n_features) representing model
+        parameters or their embeddings.
+    max_dim
+        Maximum homological dimension to compute.
+    max_edge_length
+        Maximum edge length for the Rips filtration.
+
+    Returns
+    -------
+    PersistenceSummary
+        Topological summary including Betti numbers and persistence.
+    """
+    from ripser import ripser
+
+    logger.debug("Computing persistence: {} points, max_dim={}", points.shape[0], max_dim)
+
+    result = ripser(points, maxdim=max_dim, thresh=max_edge_length)
+    diagrams = result["dgms"]
+
+    # Compute Betti numbers (count features that persist at infinity)
+    betti_0 = int(np.sum(diagrams[0][:, 1] == np.inf)) if len(diagrams) > 0 else 0
+    betti_1 = int(np.sum(diagrams[1][:, 1] == np.inf)) if len(diagrams) > 1 else 0
+
+    # Compute total and max persistence (excluding infinite features)
+    all_finite = []
+    for dgm in diagrams:
+        finite_mask = dgm[:, 1] < np.inf
+        if np.any(finite_mask):
+            lifetimes = dgm[finite_mask, 1] - dgm[finite_mask, 0]
+            all_finite.extend(lifetimes.tolist())
+
+    total_persistence = float(sum(all_finite)) if all_finite else 0.0
+    max_persistence = float(max(all_finite)) if all_finite else 0.0
+
+    # Concatenate all diagrams for storage, RETAINING the homology
+    # dimension as a third column. ripser returns a per-dimension list
+    # (diagrams[d] holds dimension d); a bare np.vstack would discard
+    # that label and make H0 and H1 bars indistinguishable. We append a
+    # dim column so downstream consumers (e.g. track_divergence) can
+    # compare like dimensions instead of pooling them.
+    labelled = []
+    for dim, dgm in enumerate(diagrams):
+        if dgm.shape[0] == 0:
+            continue
+        dim_col = np.full((dgm.shape[0], 1), float(dim))
+        labelled.append(np.hstack([dgm, dim_col]))
+    full_diagram = np.vstack(labelled) if labelled else np.empty((0, 3))
+
+    logger.debug(
+        "Persistence: beta_0={}, beta_1={}, total={:.4f}",
+        betti_0,
+        betti_1,
+        total_persistence,
+    )
+
+    return PersistenceSummary(
+        betti_0=betti_0,
+        betti_1=betti_1,
+        total_persistence=total_persistence,
+        max_persistence=max_persistence,
+        diagram=full_diagram,
+    )
+
+
+def track_divergence(
+    current_params: npt.NDArray[np.float64],
+    previous_summary: PersistenceSummary | None = None,
+    max_dim: int = 1,
+) -> PersistenceSummary:
+    """Track federation divergence across rounds.
+
+    Computes persistence of current parameter distribution and, if a
+    previous summary exists, computes the bottleneck distance to
+    measure how much the topological structure has changed.
+
+    Parameters
+    ----------
+    current_params
+        Array of shape (n_clients, n_features).
+    previous_summary
+        PersistenceSummary from the previous round, if available.
+    max_dim
+        Maximum homological dimension.
+
+    Returns
+    -------
+    PersistenceSummary
+        Updated summary with bottleneck distance to previous round.
+    """
+    summary = compute_persistence(current_params, max_dim=max_dim)
+
+    if previous_summary is not None:
+        from persim import bottleneck
+
+        # Compare H0 diagrams (connected components) ONLY. We select bars
+        # whose homology dimension is 0 and which are finite (the single
+        # infinite H0 bar -- the whole space's surviving component -- is
+        # excluded so persim never has to match infinities). Selecting on
+        # the dimension label is essential: without it, H1 (loop) bars
+        # leak into this pool and silently contaminate the "H0 divergence"
+        # with loop-structure changes. See PersistenceSummary.diagram_for_dim.
+        current_h0 = summary.diagram_for_dim(0, finite_only=True)
+        prev_h0 = previous_summary.diagram_for_dim(0, finite_only=True)
+
+        if len(current_h0) > 0 and len(prev_h0) > 0:
+            dist = bottleneck(current_h0, prev_h0)
+            summary.bottleneck_to_previous = float(dist)
+            logger.info("Bottleneck distance to previous round: {:.4f}", dist)
+        else:
+            summary.bottleneck_to_previous = 0.0
+
+    return summary

groupoid/sheaf.py

@@ -0,0 +1,84 @@
+"""Sheaf-theoretic data structures for federated learning."""
+
+from __future__ import annotations
+
+import networkx as nx
+import numpy as np
+import numpy.typing as npt
+from loguru import logger
+
+
+class Sheaf:
+    """A cellular sheaf on a graph.
+
+    Assigns vector spaces to nodes and linear restriction maps to edges.
+    """
+
+    def __init__(self, graph: nx.DiGraph) -> None:
+        self.graph = graph
+        self._restriction_maps: dict[tuple[str, str], npt.NDArray[np.float64]] = {}
+        self._sections: dict[str, npt.NDArray[np.float64]] = {}
+
+    def set_restriction_map(
+        self, source: str, target: str, matrix: npt.NDArray[np.float64]
+    ) -> None:
+        """Set the restriction map for an edge."""
+        self._restriction_maps[(source, target)] = matrix
+        logger.debug("Set restriction map {} -> {}", source, target)
+
+    def get_restriction_map(self, source: str, target: str) -> npt.NDArray[np.float64]:
+        """Get the restriction map for an edge."""
+        return self._restriction_maps[(source, target)]
+
+    def set_section(self, node: str, value: npt.NDArray[np.float64]) -> None:
+        """Set a section value at a node."""
+        self._sections[node] = value
+
+    def get_section(self, node: str) -> npt.NDArray[np.float64]:
+        """Get the section value at a node."""
+        return self._sections[node]
+
+    def restrict(
+        self, section: npt.NDArray[np.float64], source: str, target: str
+    ) -> npt.NDArray[np.float64]:
+        """Apply the restriction map to a section.
+
+        Parameters
+        ----------
+        section : np.ndarray
+            The section value at the source node.
+        source : str
+            Source node identifier.
+        target : str
+            Target node identifier.
+
+        Returns
+        -------
+        np.ndarray
+            The restricted section at the target node.
+        """
+        R = self._restriction_maps[(source, target)]
+        result: npt.NDArray[np.float64] = R @ section
+        return result
+
+    def restrict_along_path(
+        self, section: npt.NDArray[np.float64], path: list[str]
+    ) -> npt.NDArray[np.float64]:
+        """Restrict a section along a path of nodes.
+
+        Parameters
+        ----------
+        section : np.ndarray
+            The section value at path[0].
+        path : list[str]
+            Ordered list of nodes forming a path in the graph.
+
+        Returns
+        -------
+        np.ndarray
+            The section restricted to path[-1].
+        """
+        result = section
+        for i in range(len(path) - 1):
+            result = self.restrict(result, path[i], path[i + 1])
+        return result