holonomy-lib 0.8.0__py3-none-any.whl

holonomy_lib/info_geometry/__init__.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: FSL-1.1-ALv2
+# Copyright 2026 Synoros
+
+"""holonomy_lib.info_geometry: information-geometric primitives.
+
+Information geometry studies probability distributions as points on a
+Riemannian manifold whose metric is the Fisher information. Three
+families of primitives appear here:
+
+  Bregman divergences: `bregman_divergence(p, q, potential)`.
+    The asymmetric "distance" induced by a convex potential function.
+    Generalizes squared-Euclidean (potential = ½‖·‖²) and Itakura-
+    Saito (potential = −log(·)) and Kullback-Leibler (potential =
+    Σ p log p, the negative entropy). Banerjee et al. (2005) is the
+    canonical reference.
+
+  KL divergences: closed forms for the common exponential
+    families:
+      `kl_divergence_categorical(p, q)`: discrete distributions.
+      `kl_divergence_gaussian(mu_p, Sigma_p, mu_q, Sigma_q)`: multi-
+        variate Gaussians.
+
+  Fisher metric: the Riemannian inner product on the statistical
+    manifold. For q near p, KL(p || q) is half the squared length of
+    q - p in this metric. Plus `natural_gradient`, the steepest-descent
+    direction under this metric. See Amari (1998).
+      `fisher_information_categorical(p)`: diagonal Fisher on the simplex.
+      `fisher_information_gaussian_mean(Sigma)`: Fisher for the mean parameter.
+      `natural_gradient(grad, fisher_matrix)`: F^{-1} · grad.
+
+Each primitive takes a leading batch dimension and cites its source.
+
+References:
+  Amari, S.-I. (1998). Natural gradient works efficiently in learning.
+    Neural Computation 10(2):251–276.
+  Amari, S.-I. (2016). Information Geometry and Its Applications.
+    Applied Mathematical Sciences 194, Springer. The textbook.
+  Banerjee, A., Merugu, S., Dhillon, I. S., Ghosh, J. (2005).
+    Clustering with Bregman divergences. JMLR 6:1705–1749.
+  Martens, J. (2020). New insights and perspectives on the natural
+    gradient method. JMLR 21(146):1–76.
+  Nielsen, F. (2020). An elementary introduction to information
+    geometry. Entropy 22(10):1100.
+"""
+
+from holonomy_lib.info_geometry.divergences import (
+    bregman_divergence,
+    kl_divergence_categorical,
+    kl_divergence_gaussian,
+)
+from holonomy_lib.info_geometry.fisher import (
+    fisher_information_categorical,
+    fisher_information_gaussian_mean,
+    natural_gradient,
+)
+
+__all__ = [
+    "bregman_divergence",
+    "fisher_information_categorical",
+    "fisher_information_gaussian_mean",
+    "kl_divergence_categorical",
+    "kl_divergence_gaussian",
+    "natural_gradient",
+]

holonomy_lib/info_geometry/divergences.py

@@ -0,0 +1,224 @@
+# SPDX-License-Identifier: FSL-1.1-ALv2
+# Copyright 2026 Synoros
+
+"""Bregman and KL divergences, batched over leading dimensions on any device.
+
+A Bregman divergence is the asymmetric "distance" induced by a strictly
+convex generator F : Ω ⊂ ℝᵈ → ℝ:
+
+    D_F(p ‖ q) = F(p) − F(q) − ⟨∇F(q), p − q⟩.
+
+It is the gap between the function value at `p` and its first-order
+Taylor approximation at `q`. Choices of `F` recover familiar
+divergences:
+
+  F(x) = ½ ⟨x, x⟩                        → squared Euclidean ½‖p−q‖²
+  F(x) = Σ x_i log x_i − x_i             → generalized KL (Csiszár 1991)
+                                            (= KL on the simplex)
+  F(x) = − Σ log x_i                     → Itakura-Saito
+  F(X) = −log det X (for SPD matrices)  → LogDet / Burg matrix divergence
+
+This module provides:
+
+  bregman_divergence(p, q, potential)
+      The fully general form, parameterized by a callable `potential`
+      that returns (F(x), ∇F(x)) given x. Caller supplies the convex
+      function; we evaluate the Bregman formula and respect the leading
+      batch dimensions of `p`/`q`.
+
+  kl_divergence_categorical(p, q)
+      KL(p ‖ q) for discrete distributions on the simplex, with the
+      0·log 0 = 0 convention. On the simplex this equals `bregman_divergence`
+      with `potential = (x · log x − x).sum()`, and it is cheaper. The
+      Bregman form also adds Σq − Σp, which vanishes only when both sum to one.
+
+      Autograd note: `xlogy` preserves zero-mass terms and returns
+      `+inf` when `p_i > 0` and `q_i = 0` (support mismatch).
+      Gradients through near-zero bins are not numerically reliable;
+      use only for inference unless you're sure about the support.
+
+  kl_divergence_gaussian(mu_p, Sigma_p, mu_q, Sigma_q)
+      KL between two multivariate Gaussians:
+
+          KL = ½ ( tr(Σ_q⁻¹ Σ_p) + (μ_q − μ_p)ᵀ Σ_q⁻¹ (μ_q − μ_p)
+                   − d + log det Σ_q − log det Σ_p ).
+
+      Uses Cholesky-based solves for stability.
+
+References:
+  Bregman, L. M. (1967). The relaxation method of finding the common
+    point of convex sets ... USSR Comp. Math. Math. Phys. 7(3):200-217.
+    The original definition.
+  Banerjee, A., Merugu, S., Dhillon, I. S., Ghosh, J. (2005).
+    Clustering with Bregman divergences. JMLR 6:1705-1749, §2.
+  Csiszár, I. (1991). Why least squares and maximum entropy?
+    An axiomatic approach to inference for linear inverse problems.
+    Ann. Stat. 19(4):2032-2066. Bregman ↔ generalized KL connection.
+  Amari, S.-I. (2016). Information Geometry and Its Applications, §1.5.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+
+
+def bregman_divergence(
+    p: torch.Tensor,
+    q: torch.Tensor,
+    potential: Callable[[torch.Tensor], tuple[torch.Tensor, torch.Tensor]],
+) -> torch.Tensor:
+    """Bregman divergence `D_F(p ‖ q) = F(p) − F(q) − ⟨∇F(q), p − q⟩`.
+
+    Args:
+      p, q: (B, d) tensors. Must broadcast against each other; output
+        has the broadcast batch shape.
+      potential: callable taking `x: (..., d)` and returning
+        `(F(x): (...,), grad_F(x): (..., d))`. Both outputs must
+        preserve the leading dims of the input. The caller is
+        responsible for choosing a strictly convex `F` and confirming
+        `p`, `q` live in its domain.
+
+    Returns:
+      D_F(p ‖ q): (B,) divergence values. Non-negative when F is
+      strictly convex (Bregman 1967); equals zero iff p = q.
+
+    Example (squared Euclidean):
+      >>> def half_norm_sq(x):
+      ...     return 0.5 * (x * x).sum(dim=-1), x
+      >>> bregman_divergence(p, q, half_norm_sq)
+      # = ½ ‖p − q‖²
+
+    References:
+      Bregman (1967).
+      Banerjee et al. (2005), §2.
+    """
+    if p.shape[-1] != q.shape[-1]:
+        raise ValueError(
+            f"p and q must have matching last dim; got p.shape[-1]="
+            f"{p.shape[-1]} vs q.shape[-1]={q.shape[-1]}"
+        )
+    # Only the gradient at q enters the formula, so the potential's gradient
+    # at p is discarded.
+    F_p, _ = potential(p)
+    F_q, grad_F_q = potential(q)
+    # ⟨∇F(q), p − q⟩, summed over the feature dim so it has the batch shape.
+    inner = (grad_F_q * (p - q)).sum(dim=-1)
+    return F_p - F_q - inner
+
+
+def kl_divergence_categorical(
+    p: torch.Tensor,
+    q: torch.Tensor,
+) -> torch.Tensor:
+    """KL divergence between two discrete distributions on the simplex.
+
+        KL(p ‖ q) = Σ_i p_i · (log p_i − log q_i).
+
+    Uses the convention 0 · log(0/x) = 0 (continuous limit). Both `p`
+    and `q` must be on the simplex (non-negative, sum to 1 along the
+    last dim); we do not enforce this, so it is a precondition.
+    Gradients at zero probabilities need not be finite; differentiate
+    only where both distributions have strictly positive entries.
+
+    Args:
+      p, q: (B, k) probability tensors. Last dim is the support size.
+
+    Returns:
+      KL: (B,) divergence values, non-negative (Gibbs inequality).
+
+    References:
+      Cover, T. M., Thomas, J. A. (2006). Elements of Information
+        Theory, 2nd ed., §2.3.
+      Amari (2016), §2.4.
+    """
+    if p.shape[-1] != q.shape[-1]:
+        raise ValueError(
+            f"p and q must have matching last dim; got {p.shape[-1]} vs "
+            f"{q.shape[-1]}"
+        )
+    # p log(p / q) is split into two xlogy terms. xlogy(p, x) is 0 wherever
+    # p is 0, so zero-mass bins drop out without a 0 * log 0 nan. No clamp
+    # is applied, so tiny positive p keep their exact p log p. A bin with
+    # p_i > 0 and q_i = 0 gives +inf, as the module note says.
+    return (torch.special.xlogy(p, p) - torch.special.xlogy(p, q)).sum(dim=-1)
+
+
+def kl_divergence_gaussian(
+    mu_p: torch.Tensor,
+    Sigma_p: torch.Tensor,
+    mu_q: torch.Tensor,
+    Sigma_q: torch.Tensor,
+) -> torch.Tensor:
+    """KL divergence between two multivariate Gaussians.
+
+        KL(N(μ_p, Σ_p) ‖ N(μ_q, Σ_q))
+          = ½ ( tr(Σ_q⁻¹ Σ_p)
+                + (μ_q − μ_p)ᵀ Σ_q⁻¹ (μ_q − μ_p)
+                − d
+                + log det Σ_q − log det Σ_p ).
+
+    Stability: we factorize Σ_q via Cholesky once and reuse the
+    factorization for both the trace term `tr(Σ_q⁻¹ Σ_p)` and the
+    Mahalanobis term `(μ_q − μ_p)ᵀ Σ_q⁻¹ (μ_q − μ_p)`. `log det Σ_q`
+    comes for free from the Cholesky diagonal. For `Σ_p` we just need
+    `log det Σ_p`, also via Cholesky. A covariance that is not SPD makes
+    `torch.linalg.cholesky` raise, so bad input fails instead of returning
+    a number.
+
+    Args:
+      mu_p, mu_q: (..., d) means.
+      Sigma_p, Sigma_q: (..., d, d) covariances. Must be SPD.
+
+    Returns:
+      KL: (...,) divergence values, non-negative.
+
+    References:
+      Amari (2016), §2.4; exponential families.
+      Petersen, K. B., Pedersen, M. S. (2012). The Matrix Cookbook,
+        eq. 380; Gaussian-Gaussian KL.
+    """
+    d = mu_p.shape[-1]
+    if mu_q.shape[-1] != d:
+        raise ValueError(
+            f"mu_q's last dim must match mu_p's d={d}; "
+            f"got mu_q.shape={tuple(mu_q.shape)}"
+        )
+    if Sigma_p.shape[-1] != d or Sigma_p.shape[-2] != d:
+        raise ValueError(
+            f"Sigma_p must be (..., d, d) with d={d}; "
+            f"got {tuple(Sigma_p.shape)}"
+        )
+    if Sigma_q.shape[-1] != d or Sigma_q.shape[-2] != d:
+        raise ValueError(
+            f"Sigma_q must be (..., d, d) with d={d}; "
+            f"got {tuple(Sigma_q.shape)}"
+        )
+
+    L_q = torch.linalg.cholesky(Sigma_q)
+    # Σ_q⁻¹ Σ_p through the Cholesky factor of Σ_q: two triangular solves,
+    # so no explicit inverse of Σ_q is formed.
+    Sigma_q_inv_Sigma_p = torch.cholesky_solve(Sigma_p, L_q)
+    trace_term = torch.diagonal(
+        Sigma_q_inv_Sigma_p, dim1=-2, dim2=-1,
+    ).sum(dim=-1)
+
+    diff = (mu_q - mu_p).unsqueeze(dim=-1)  # (..., d, 1)
+    # Σ_q⁻¹ Δμ via cholesky_solve; then ⟨Δμ, ·⟩ is a dot product.
+    Sigma_q_inv_diff = torch.cholesky_solve(diff, L_q).squeeze(dim=-1)
+    mahalanobis = ((mu_q - mu_p) * Sigma_q_inv_diff).sum(dim=-1)
+
+    # log det Σ = 2 Σ log diag(L) for the Cholesky factor.
+    log_det_Sigma_q = 2.0 * torch.log(
+        torch.diagonal(L_q, dim1=-2, dim2=-1)
+    ).sum(dim=-1)
+    L_p = torch.linalg.cholesky(Sigma_p)
+    log_det_Sigma_p = 2.0 * torch.log(
+        torch.diagonal(L_p, dim1=-2, dim2=-1)
+    ).sum(dim=-1)
+
+    return 0.5 * (
+        trace_term + mahalanobis - float(d)
+        + log_det_Sigma_q - log_det_Sigma_p
+    )

holonomy_lib/info_geometry/fisher.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: FSL-1.1-ALv2
+# Copyright 2026 Synoros
+
+"""Fisher information metric + natural-gradient primitives.
+
+The Fisher information matrix is the canonical Riemannian metric on a
+parametric statistical manifold:
+
+    g_{ij}(θ) = E_{p_θ}[ (∂_i log p_θ)(∂_j log p_θ) ]
+              = − E_{p_θ}[ ∂_i ∂_j log p_θ ].
+
+It is the second-order Taylor coefficient of KL near the diagonal:
+
+    KL(p_θ ‖ p_{θ + dθ}) = ½ dθᵀ F(θ) dθ + O(‖dθ‖³).
+
+Information geometry studies the geodesics of this metric. In
+optimization, the natural gradient `F(θ)^{−1} ∇L` is the steepest-descent
+direction under this metric. The metric does not depend on how θ is
+parameterized, so the direction does not either (Amari 1998).
+
+Two closed-form Fisher matrices are exposed for the most common
+families; users with a custom likelihood can compute their own Fisher
+either analytically or via `torch.autograd` over `log p`.
+
+  fisher_information_categorical(p)
+    Diagonal Fisher on the simplex, g_ii(p) = 1/p_i. Recovers the
+    Hellinger-Bhattacharyya geometry whose geodesic distance is twice
+    the Hellinger angle.
+
+  fisher_information_gaussian_mean(Sigma)
+    The Fisher information for the mean of N(μ, Σ) at fixed Σ. The
+    closed form is g(μ) = Σ^{−1}, which makes the Fisher-metric
+    geodesics between two means coincide with the Mahalanobis-distance
+    straight lines for that fixed Σ.
+
+  natural_gradient(grad, fisher_matrix)
+    Apply the Fisher inverse to a Euclidean gradient via a Cholesky
+    solve (no explicit inverse). The output is the steepest-descent
+    direction under the Fisher metric.
+
+References:
+  Amari, S.-I. (1998). Natural gradient works efficiently in learning.
+    Neural Computation 10(2):251–276. The natural-gradient paper.
+  Amari, S.-I. (2016). Information Geometry and Its Applications.
+    Applied Mathematical Sciences 194, Springer. §2.5, §4.
+  Martens, J. (2020). New insights and perspectives on the natural
+    gradient method. JMLR 21(146):1–76. Modern revisit.
+  Petersen, K. B., Pedersen, M. S. (2012). The Matrix Cookbook,
+    §6 (Statistics); Fisher information for the Gaussian family.
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def fisher_information_categorical(p: torch.Tensor) -> torch.Tensor:
+    """Fisher information metric on the simplex Δ^{k−1}.
+
+    For p in the open simplex, the Fisher information metric is the
+    diagonal matrix
+
+        g_ij(p) = δ_ij / p_i.
+
+    This is the metric induced by the KL divergence's second-order
+    Taylor expansion: KL(p ‖ p + δ) ≈ ½ Σ_i δ_i² / p_i.
+
+    Args:
+      p: (B, k) probability vectors on the open simplex. Entries should
+        be positive. Metric coefficients diverge as p_i approaches zero,
+        but the boundary is at finite Fisher distance. For example, the
+        Bernoulli distance from p=1/2 to p=0 is pi/2.
+        Entries below `1e-9` are clamped up to `1e-9` in the input dtype
+        (float16 is a special case, see the body).
+
+    Returns:
+      Fisher matrix: (B, k, k) diagonal SPD.
+
+    References:
+      Amari (2016), §2.4 example 4.
+      Nielsen, F. (2020). An elementary introduction to information
+        geometry. Entropy 22(10):1100.
+    """
+    if p.ndim < 1:
+        raise ValueError(f"p must have at least 1 dim; got p.shape={tuple(p.shape)}")
+    # The clamp is what keeps 1/p finite at the boundary. The floor is the
+    # library-wide convention `numerical_floor_convention` in
+    # notes/magic_numbers.md, not a per-call knob. Keeping p in the interior
+    # is the caller's job; a clamped entry gives a large but finite value.
+    # float16 cannot represent 1e-9 (it rounds to 0), so the floor is raised
+    # to the dtype's smallest normal number when that is larger. Every other
+    # float dtype keeps the 1e-9 convention.
+    p_safe = p.clamp(min=max(1e-9, torch.finfo(p.dtype).tiny))
+    inv_p = torch.reciprocal(p_safe)
+    return torch.diag_embed(inv_p)
+
+
+def fisher_information_gaussian_mean(Sigma: torch.Tensor) -> torch.Tensor:
+    """Fisher information for the mean μ of N(μ, Σ), at fixed Σ.
+
+    Standard result (Amari 2016, eq. 2.43): for the Gaussian family
+    parameterized by the mean with covariance held fixed,
+
+        g(μ) = Σ^{−1}.
+
+    The Fisher-geodesic distance between two means μ_1, μ_2 is then
+    `sqrt((μ_1 − μ_2)ᵀ Σ^{−1} (μ_1 − μ_2))`, i.e. the Mahalanobis
+    distance for the given Σ.
+
+    Computed from a Cholesky factor with `cholesky_inverse`, not
+    `linalg.inv`. A matrix that is not SPD makes `cholesky` raise.
+
+    Args:
+      Sigma: (B, d, d) SPD covariance matrices.
+
+    Returns:
+      Fisher matrix: (B, d, d) SPD.
+
+    References:
+      Amari (2016), §2.4 example 5.
+      Petersen-Pedersen (2012), §6.
+    """
+    if Sigma.ndim < 2 or Sigma.shape[-1] != Sigma.shape[-2]:
+        raise ValueError(
+            f"Sigma must be (..., d, d); got Sigma.shape={tuple(Sigma.shape)}"
+        )
+    # Σ is SPD, so factor it once and invert from the factor. This skips
+    # the pivoting an LU inverse does, and kl_divergence_gaussian uses the
+    # same factorization.
+    L = torch.linalg.cholesky(Sigma)
+    return torch.cholesky_inverse(L)
+
+
+def natural_gradient(
+    grad: torch.Tensor,
+    fisher_matrix: torch.Tensor,
+) -> torch.Tensor:
+    """Natural gradient: F(θ)^{−1} · ∇L.
+
+    The natural gradient (Amari 1998) is the steepest-descent direction
+    on a statistical manifold under the Fisher metric. Unlike the
+    Euclidean gradient, the direction does not change under a smooth
+    reparameterization of θ, since the metric is the same geometric
+    object in every coordinate system. That removes one source of
+    parameterization-dependent behavior. It does not make an
+    ill-conditioned F any better conditioned.
+
+    Implementation: solve `F · d = grad` through a Cholesky factor of F,
+    rather than forming `inv(F) @ grad`. The solve inherits whatever
+    conditioning F has, without the extra rounding from an explicit inverse.
+
+    Args:
+      grad: (B, d) Euclidean gradient at θ.
+      fisher_matrix: (B, d, d) Fisher information matrix at θ, SPD.
+
+    Returns:
+      (B, d) natural gradient.
+
+    References:
+      Amari (1998).
+      Martens (2020), §2; modern derivation.
+    """
+    if fisher_matrix.ndim < 2:
+        raise ValueError(
+            f"fisher_matrix must be (..., d, d) with at least 2 dims; "
+            f"got fisher_matrix.shape={tuple(fisher_matrix.shape)}"
+        )
+    if grad.shape[-1] != fisher_matrix.shape[-1]:
+        raise ValueError(
+            f"grad last dim {grad.shape[-1]} must match fisher_matrix "
+            f"last dim {fisher_matrix.shape[-1]}"
+        )
+    if fisher_matrix.shape[-1] != fisher_matrix.shape[-2]:
+        raise ValueError(
+            f"fisher_matrix must be (..., d, d); got "
+            f"fisher_matrix.shape={tuple(fisher_matrix.shape)}"
+        )
+    L = torch.linalg.cholesky(fisher_matrix)
+    return torch.cholesky_solve(grad.unsqueeze(dim=-1), L).squeeze(dim=-1)

holonomy_lib/lie/__init__.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: FSL-1.1-ALv2
+# Copyright 2026 Synoros
+
+"""holonomy_lib.lie : Lie group primitives for SO(3) and the 2-sphere.
+
+The Lie group `SO(3)` is the group of orientation-preserving
+rotations of R³; its Lie algebra `so(3)` is the 3-dim space of
+skew-symmetric matrices (equivalently, angular-velocity vectors).
+This module ships the standard exp/log/random/compose primitives
+plus real spherical harmonics, the standard basis for functions on the
+unit sphere that rotate in a known way under SO(3).
+
+Currently exposed:
+  so3.axis_angle_to_matrix(axis, angle)  : Rodrigues formula.
+  so3.matrix_to_axis_angle(R)            : inverse log map.
+  so3.so3_exp(omega)                     : matrix exp on so(3).
+  so3.so3_log(R)                         : matrix log on SO(3).
+  so3.random_so3(batch_size, generator)  : Haar-uniform via
+                                            quaternion sampling.
+  so3.compose(R1, R2)                    : group product (matmul).
+  real_spherical_harmonics(directions, l_max)
+                                          : closed-form Y_lm for
+                                            l ≤ 4 evaluated at
+                                            (B, 3) direction vectors
+                                            (normalized internally).
+
+Planned, not implemented yet:
+  Wigner-D matrices (real basis) for higher-l rotation actions on
+    spherical-harmonic features.
+  Clebsch-Gordan coefficients for SO(3) tensor products (l₁ ⊗ l₂ → L).
+  SE(3), SU(2), SL(n) primitives.
+
+References:
+  Hall, B. C. (2015). Lie Groups, Lie Algebras, and Representations,
+    2nd ed. Springer GTM 222. The canonical reference for matrix Lie
+    groups; §3.1 covers SO(3) and its Lie algebra.
+  Edmonds, A. R. (1957). Angular Momentum in Quantum Mechanics.
+    Princeton University Press. Real and complex spherical harmonics
+    + Wigner D-matrices.
+  Shoemake, K. (1992). Uniform random rotations. In Graphics Gems
+    III, pp. 124–132. Academic Press. The quaternion-from-3-uniforms
+    construction used by `random_so3`.
+  Cohen, T. S., Geiger, M., Köhler, J., Welling, M. (2018).
+    Spherical CNNs. ICLR. Modern application of these primitives in
+    SO(3)-equivariant deep learning.
+"""
+
+from holonomy_lib.lie import so3
+from holonomy_lib.lie.spherical_harmonics import real_spherical_harmonics
+
+__all__ = [
+    "real_spherical_harmonics",
+    "so3",
+]