API Reference

Aggregation

groupoid.aggregation.TransportGroupoidAggregator dataclass

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. track_divergence When True, each :meth:aggregate round computes the persistent homology of the transported client parameters and the bottleneck distance to the previous round (H0 against H0), populating FederatedRound.divergence.

Source code in groupoid/aggregation.py

@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.
    track_divergence
        When True, each :meth:`aggregate` round computes the persistent
        homology of the transported client parameters and the bottleneck
        distance to the previous round (H0 against H0), populating
        ``FederatedRound.divergence``.
    """

    manifold: object
    graph: nx.DiGraph
    base_node: str
    consistency_threshold: float = 1e-6
    track_divergence: bool = False
    morphisms: dict[tuple[str, str], Morphism] = field(default_factory=dict)
    _round_idx: int = field(default=0, init=False)
    _prev_divergence: _persistence.PersistenceSummary | None = field(
        default=None, init=False, repr=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 register_transport_from_points(
        self,
        source: str,
        target: str,
        source_point: npt.NDArray[np.float64],
        target_point: npt.NDArray[np.float64],
        method: str = "pole",
        n_rungs: int = 2,
    ) -> npt.NDArray[np.float64]:
        """Compute and register a transport map from two manifold points.

        Builds the parallel-transport matrix along the geodesic from
        ``source_point`` to ``target_point`` via
        :func:`groupoid.transport.compute_transport_matrix` (pole or
        Schild's ladder) and registers it for the ``(source, target)``
        edge. This wires the ladder approximations into the aggregation
        pipeline for the case where per-client base points are known but
        explicit transport matrices are not. The ladders are
        approximations with a small residual; see LIMITATIONS.md for the
        measured behavior.

        Returns
        -------
        np.ndarray
            The registered transport matrix.
        """
        matrix = compute_transport_matrix(
            self.manifold,
            source_point,
            target_point,
            method=method,
            n_rungs=n_rungs,
        )
        self.register_transport(source, target, matrix)
        return matrix

    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])

        # Optional topological divergence tracking on the transported
        # parameters (all in the base frame, so rounds are comparable).
        divergence: _persistence.PersistenceSummary | None = None
        if self.track_divergence:
            divergence = _persistence.track_divergence(
                param_stack, previous_summary=self._prev_divergence
            )
            self._prev_divergence = divergence

        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,
            divergence=divergence,
        )

        logger.info(
            "Round {} complete: H^1={:.2e}, consistent={}",
            self._round_idx,
            h1,
            is_consistent,
        )
        return result

aggregate(client_params, weights=None)

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.

Source code in groupoid/aggregation.py

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])

    # Optional topological divergence tracking on the transported
    # parameters (all in the base frame, so rounds are comparable).
    divergence: _persistence.PersistenceSummary | None = None
    if self.track_divergence:
        divergence = _persistence.track_divergence(
            param_stack, previous_summary=self._prev_divergence
        )
        self._prev_divergence = divergence

    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,
        divergence=divergence,
    )

    logger.info(
        "Round {} complete: H^1={:.2e}, consistent={}",
        self._round_idx,
        h1,
        is_consistent,
    )
    return result

check_consistency(client_params)

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.

Source code in groupoid/aggregation.py

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

register_transport(source, target, matrix)

Register a transport map between two clients.

Source code in groupoid/aggregation.py

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)

register_transport_from_points(source, target, source_point, target_point, method='pole', n_rungs=2)

Compute and register a transport map from two manifold points.

Builds the parallel-transport matrix along the geodesic from source_point to target_point via :func:groupoid.transport.compute_transport_matrix (pole or Schild's ladder) and registers it for the (source, target) edge. This wires the ladder approximations into the aggregation pipeline for the case where per-client base points are known but explicit transport matrices are not. The ladders are approximations with a small residual; see LIMITATIONS.md for the measured behavior.

Returns

np.ndarray The registered transport matrix.

Source code in groupoid/aggregation.py

def register_transport_from_points(
    self,
    source: str,
    target: str,
    source_point: npt.NDArray[np.float64],
    target_point: npt.NDArray[np.float64],
    method: str = "pole",
    n_rungs: int = 2,
) -> npt.NDArray[np.float64]:
    """Compute and register a transport map from two manifold points.

    Builds the parallel-transport matrix along the geodesic from
    ``source_point`` to ``target_point`` via
    :func:`groupoid.transport.compute_transport_matrix` (pole or
    Schild's ladder) and registers it for the ``(source, target)``
    edge. This wires the ladder approximations into the aggregation
    pipeline for the case where per-client base points are known but
    explicit transport matrices are not. The ladders are
    approximations with a small residual; see LIMITATIONS.md for the
    measured behavior.

    Returns
    -------
    np.ndarray
        The registered transport matrix.
    """
    matrix = compute_transport_matrix(
        self.manifold,
        source_point,
        target_point,
        method=method,
        n_rungs=n_rungs,
    )
    self.register_transport(source, target, matrix)
    return matrix

groupoid.aggregation.FederatedRound dataclass

Result of a single federated aggregation round.

divergence is populated only when the aggregator was constructed with track_divergence=True: it is the persistent-homology summary of the transported client parameters, with the bottleneck distance to the previous round's summary (None on the first round with tracking off; see :func:groupoid.persistence.track_divergence).

Source code in groupoid/aggregation.py

@dataclass
class FederatedRound:
    """Result of a single federated aggregation round.

    ``divergence`` is populated only when the aggregator was constructed
    with ``track_divergence=True``: it is the persistent-homology summary
    of the transported client parameters, with the bottleneck distance to
    the previous round's summary (None on the first round with tracking
    off; see :func:`groupoid.persistence.track_divergence`).
    """

    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
    divergence: _persistence.PersistenceSummary | None = None

Manifold Operations

groupoid.manifold.karcher_mean(manifold, points, weights=None, max_iter=100, tol=1e-06)

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.

Source code in groupoid/manifold.py

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

groupoid.groupoid.Morphism

Bases: BaseModel

A morphism in the transport groupoid.

Represents a transport map between two nodes (clients) in the federated learning network.

Source code in groupoid/groupoid.py

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})"

    # pydantic's BaseModel defines __str__ separately (a field dump including
    # the full transport matrix), which would bypass the compact __repr__ in
    # print(), f-strings, and loguru's {} formatting. Alias it explicitly.
    __str__ = __repr__

groupoid.groupoid.compose(f, g)

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.

Source code in groupoid/groupoid.py

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,
    )

groupoid.groupoid.inverse(f)

Compute the inverse of a morphism.

Parameters

f : Morphism The morphism to invert.

Returns

Morphism The inverse morphism from f.target to f.source.

Source code in groupoid/groupoid.py

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,
    )

Cohomology

groupoid.cohomology.compute_h1(graph, transport_maps)

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.

Source code in groupoid/cohomology.py

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

Sheaf

groupoid.sheaf.Sheaf

A cellular sheaf on a graph.

Assigns vector spaces to nodes and linear restriction maps to edges.

Source code in groupoid/sheaf.py

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

get_restriction_map(source, target)

Get the restriction map for an edge.

Source code in groupoid/sheaf.py

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)]

get_section(node)

Get the section value at a node.

Source code in groupoid/sheaf.py

def get_section(self, node: str) -> npt.NDArray[np.float64]:
    """Get the section value at a node."""
    return self._sections[node]

restrict(section, source, target)

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.

Source code in groupoid/sheaf.py

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

restrict_along_path(section, path)

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].

Source code in groupoid/sheaf.py

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

set_restriction_map(source, target, matrix)

Set the restriction map for an edge.

Source code in groupoid/sheaf.py

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)

set_section(node, value)

Set a section value at a node.

Source code in groupoid/sheaf.py

def set_section(self, node: str, value: npt.NDArray[np.float64]) -> None:
    """Set a section value at a node."""
    self._sections[node] = value

Sheaf Laplacian

groupoid.laplacian.build_sheaf_laplacian(sheaf, stalk_dim)

Build the sheaf Laplacian matrix.

For a sheaf F on graph G with n nodes and stalk dimension d, the sheaf Laplacian is a (nd) x (nd) 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 (nd, nd).

Source code in groupoid/laplacian.py

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

groupoid.laplacian.spectral_analysis(sheaf, stalk_dim, tol=1e-10)

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.

Source code in groupoid/laplacian.py

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,
    )

groupoid.laplacian.sheaf_diffusion_step(sheaf, sections, stalk_dim, step_size=0.1)

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.

Source code in groupoid/laplacian.py

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

Parallel Transport

groupoid.transport.schild_ladder(manifold, tangent_vec, base_point, end_point, n_rungs=1)

Parallel transport via Schild's ladder.

An iterative approximation of parallel transport that only requires the exponential and logarithmic maps. This is the coarser of the two ladders: on S^2 its direction agrees with the analytic parallel transport only to cosine ~0.98 and does not converge to the exact value as n_rungs increases (see pole_ladder for the accurate variant and LIMITATIONS.md for the measured behavior).

Parameters

manifold A geomstats manifold with exp and log maps. tangent_vec Tangent vector at base_point to transport. base_point Starting point on the manifold. end_point Destination point on the manifold. n_rungs Number of ladder rungs (higher = more accurate).

Returns

np.ndarray The transported tangent vector at end_point.

Source code in groupoid/transport.py

def schild_ladder(
    manifold: Any,  # geomstats manifold; no upstream type stubs
    tangent_vec: npt.NDArray[np.float64],
    base_point: npt.NDArray[np.float64],
    end_point: npt.NDArray[np.float64],
    n_rungs: int = 1,
) -> npt.NDArray[np.float64]:
    """Parallel transport via Schild's ladder.

    An iterative approximation of parallel transport that only requires
    the exponential and logarithmic maps. This is the coarser of the two
    ladders: on S^2 its direction agrees with the analytic parallel
    transport only to cosine ~0.98 and does not converge to the exact
    value as ``n_rungs`` increases (see ``pole_ladder`` for the accurate
    variant and LIMITATIONS.md for the measured behavior).

    Parameters
    ----------
    manifold
        A geomstats manifold with exp and log maps.
    tangent_vec
        Tangent vector at base_point to transport.
    base_point
        Starting point on the manifold.
    end_point
        Destination point on the manifold.
    n_rungs
        Number of ladder rungs (higher = more accurate).

    Returns
    -------
    np.ndarray
        The transported tangent vector at end_point.
    """
    metric = manifold.metric
    direction = metric.log(end_point, base_point)

    current_base = base_point
    current_vec = tangent_vec

    for _k in range(n_rungs):
        step = direction * (1.0 / n_rungs)
        next_base = metric.exp(step, current_base)

        # Schild's ladder: reflect the antipodal endpoint through
        # the midpoint of the base-to-next geodesic
        u = metric.exp(-current_vec, current_base)
        midpoint = metric.exp(step / 2.0, current_base)
        log_mid_u = metric.log(u, midpoint)
        u_prime = metric.exp(-log_mid_u, midpoint)
        current_vec = metric.log(u_prime, next_base)

        current_base = next_base
        direction = metric.log(end_point, current_base)

    result: npt.NDArray[np.float64] = current_vec
    return result

groupoid.transport.pole_ladder(manifold, tangent_vec, base_point, end_point, n_rungs=1)

Parallel transport via pole ladder.

A variant of Schild's ladder that uses geodesic symmetry instead of midpoint construction. More accurate than Schild's ladder for the same number of rungs, and exactly preserves the norm of the transported vector on symmetric spaces.

Parameters

manifold A geomstats manifold with exp and log maps. tangent_vec Tangent vector at base_point to transport. base_point Starting point on the manifold. end_point Destination point on the manifold. n_rungs Number of ladder rungs.

Returns

np.ndarray The transported tangent vector at end_point.

Source code in groupoid/transport.py

def pole_ladder(
    manifold: Any,  # geomstats manifold; no upstream type stubs
    tangent_vec: npt.NDArray[np.float64],
    base_point: npt.NDArray[np.float64],
    end_point: npt.NDArray[np.float64],
    n_rungs: int = 1,
) -> npt.NDArray[np.float64]:
    """Parallel transport via pole ladder.

    A variant of Schild's ladder that uses geodesic symmetry instead
    of midpoint construction. More accurate than Schild's ladder for
    the same number of rungs, and exactly preserves the norm of the
    transported vector on symmetric spaces.

    Parameters
    ----------
    manifold
        A geomstats manifold with exp and log maps.
    tangent_vec
        Tangent vector at base_point to transport.
    base_point
        Starting point on the manifold.
    end_point
        Destination point on the manifold.
    n_rungs
        Number of ladder rungs.

    Returns
    -------
    np.ndarray
        The transported tangent vector at end_point.
    """
    metric = manifold.metric
    direction = metric.log(end_point, base_point)

    current_base = base_point
    current_vec = tangent_vec

    for _k in range(n_rungs):
        step = direction * (1.0 / n_rungs)
        next_base = metric.exp(step, current_base)

        # Pole construction: reflect through the midpoint of the geodesic
        pole = metric.exp(-current_vec, current_base)
        mid_of_geodesic = metric.exp(step / 2.0, current_base)
        log_pole_from_mid = metric.log(pole, mid_of_geodesic)
        reflected = metric.exp(-log_pole_from_mid, mid_of_geodesic)
        current_vec = metric.log(reflected, next_base)

        current_base = next_base

    return current_vec

groupoid.transport.compute_transport_matrix(manifold, base_point, end_point, method='pole', n_rungs=2)

Compute the parallel transport matrix between two points.

Constructs the full linear map T such that for any tangent vector v at base_point, T @ v gives the parallel transport of v to end_point.

Parameters

manifold A geomstats manifold. base_point Starting point on the manifold. end_point Destination point on the manifold. method Transport method: "pole" or "schild". n_rungs Number of ladder rungs for the approximation.

Returns

np.ndarray Transport matrix of shape (dim, dim).

Source code in groupoid/transport.py

def compute_transport_matrix(
    manifold: Any,  # geomstats manifold; no upstream type stubs
    base_point: npt.NDArray[np.float64],
    end_point: npt.NDArray[np.float64],
    method: str = "pole",
    n_rungs: int = 2,
) -> npt.NDArray[np.float64]:
    """Compute the parallel transport matrix between two points.

    Constructs the full linear map T such that for any tangent vector v
    at base_point, T @ v gives the parallel transport of v to end_point.

    Parameters
    ----------
    manifold
        A geomstats manifold.
    base_point
        Starting point on the manifold.
    end_point
        Destination point on the manifold.
    method
        Transport method: "pole" or "schild".
    n_rungs
        Number of ladder rungs for the approximation.

    Returns
    -------
    np.ndarray
        Transport matrix of shape (dim, dim).
    """
    transport_fn = pole_ladder if method == "pole" else schild_ladder
    dim = base_point.shape[-1]
    T = np.zeros((dim, dim))

    # Transport each basis vector
    for i in range(dim):
        e_i = np.zeros(dim)
        e_i[i] = 1.0

        # Project onto tangent space at base_point
        tangent = manifold.to_tangent(e_i, base_point)

        transported = transport_fn(manifold, tangent, base_point, end_point, n_rungs=n_rungs)
        T[:, i] = transported

    logger.debug(
        "Transport matrix computed ({} method, {} rungs): det={:.4f}",
        method,
        n_rungs,
        np.linalg.det(T),
    )
    return T

Riemannian Optimizers

groupoid.optimizer.RiemannianSGD dataclass

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. With momentum, the velocity is parallel-transported into each new iterate's tangent space (see the module docstring), so it accumulates geometry-consistently across steps.

Parameters

manifold A geomstats manifold instance. lr Learning rate. momentum Momentum coefficient (0 = no momentum).

Source code in groupoid/optimizer.py

@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. With momentum, the velocity is
    parallel-transported into each new iterate's tangent space (see the
    module docstring), so it accumulates geometry-consistently across
    steps.

    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. _velocity was parallel-transported into `point`'s
        # tangent space at the end of the previous step; the to_tangent wrap
        # keeps the combination numerically tangent (and tolerates callers
        # stepping from a point other than the previous iterate).
        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
                )
            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)

        # Carry the velocity to the new iterate's tangent space.
        if self.momentum > 0:
            self._velocity = _transport_moment(self.manifold, vel, point, new_point)

        return new_point

step(point, euclidean_grad)

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.

Source code in groupoid/optimizer.py

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. _velocity was parallel-transported into `point`'s
    # tangent space at the end of the previous step; the to_tangent wrap
    # keeps the combination numerically tangent (and tolerates callers
    # stepping from a point other than the previous iterate).
    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
            )
        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)

    # Carry the velocity to the new iterate's tangent space.
    if self.momentum > 0:
        self._velocity = _transport_moment(self.manifold, vel, point, new_point)

    return new_point

groupoid.optimizer.RiemannianAdam dataclass

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. The first moment is parallel-transported into each new iterate's tangent space (see the module docstring); the second moment is a scalar gradient-norm average and needs no transport because parallel transport is an isometry (norms are invariant).

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.

Source code in groupoid/optimizer.py

@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. The first moment
    is parallel-transported into each new iterate's tangent space (see
    the module docstring); the second moment is a scalar gradient-norm
    average and needs no transport because parallel transport is an
    isometry (norms are invariant).

    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). Seed from zero like the
        # second moment below: m_1 = (1 - beta1) * grad, so the bias correction
        # m_hat = m_1 / (1 - beta1) recovers `grad` on the first step. Seeding
        # m_1 = grad directly would leave the 1/(1-beta1) factor uncancelled and
        # inflate the first update by ~10x (beta1=0.9).
        # _m was parallel-transported into `point`'s tangent space at the end
        # of the previous step; the to_tangent wrap keeps the combination
        # numerically tangent (and tolerates callers stepping from a point
        # other than the previous iterate).
        first_moment: npt.NDArray[np.float64]
        if self._m is None:
            first_moment = (1 - self.beta1) * grad
        else:
            first_moment = self.manifold.to_tangent(
                self.beta1 * self._m + (1 - self.beta1) * grad, point
            )

        # 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)

        # Carry the first moment to the new iterate's tangent space.
        self._m = _transport_moment(self.manifold, first_moment, point, new_point)

        return new_point

step(point, euclidean_grad)

Perform one optimization step.

Parameters

point Current point on the manifold. euclidean_grad Euclidean gradient.

Returns

np.ndarray Updated point on the manifold.

Source code in groupoid/optimizer.py

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). Seed from zero like the
    # second moment below: m_1 = (1 - beta1) * grad, so the bias correction
    # m_hat = m_1 / (1 - beta1) recovers `grad` on the first step. Seeding
    # m_1 = grad directly would leave the 1/(1-beta1) factor uncancelled and
    # inflate the first update by ~10x (beta1=0.9).
    # _m was parallel-transported into `point`'s tangent space at the end
    # of the previous step; the to_tangent wrap keeps the combination
    # numerically tangent (and tolerates callers stepping from a point
    # other than the previous iterate).
    first_moment: npt.NDArray[np.float64]
    if self._m is None:
        first_moment = (1 - self.beta1) * grad
    else:
        first_moment = self.manifold.to_tangent(
            self.beta1 * self._m + (1 - self.beta1) * grad, point
        )

    # 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)

    # Carry the first moment to the new iterate's tangent space.
    self._m = _transport_moment(self.manifold, first_moment, point, new_point)

    return new_point

groupoid.optimizer.curvature_adaptive_lr(manifold, point, base_lr, tangent_vec)

Adapt learning rate based on local sectional curvature.

In regions of high positive curvature, geodesics converge and we should take smaller steps, so the base rate is damped by 1 / (1 + kappa). In flat or negatively curved regions the base rate is returned unchanged; no enlargement is applied.

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.

Source code in groupoid/optimizer.py

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, so the base rate is damped by
    ``1 / (1 + kappa)``. In flat or negatively curved regions the base
    rate is returned unchanged; no enlargement is applied.

    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

Persistent Homology

groupoid.persistence.compute_persistence(points, max_dim=1, max_edge_length=np.inf)

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.

Source code in groupoid/persistence.py

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,
    )

groupoid.persistence.track_divergence(current_params, previous_summary=None, max_dim=1)

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.

Source code in groupoid/persistence.py

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