calibrax.metrics.functional.manifold

Manifold distance metrics for structured matrix spaces. Includes the affine-invariant and log-Euclidean distances on symmetric positive definite (SPD) matrices, geodesic distance on Grassmann and Stiefel manifolds, and ultrahyperbolic distance on pseudo-hyperboloids.

Manifold distance metrics -- Riemannian and pseudo-Riemannian geometry.

Geometrically correct distances for structured mathematical objects that live on curved manifolds rather than flat Euclidean space. Using naive Euclidean distance on these objects produces geometrically meaningless results.

Manifold hierarchy (each level generalizes the previous): - Euclidean: flat space, standard L2 norm - Riemannian: curved space with positive-definite metric tensor - Pseudo-Riemannian: metric tensor may be indefinite (mixed signature)

Manifolds covered: - SPD(n): Symmetric Positive Definite matrices (covariance matrices, diffusion tensors). Two distances: affine-invariant (geometrically exact) and log-Euclidean (faster approximation). - Gr(p,n): Grassmann manifold of p-dimensional subspaces (PCA subspace comparison, feature space analysis). Basis-independent. - St(p,n): Stiefel manifold of orthonormal p-frames (orientation comparison, orthogonal constraints). Basis-dependent. - Pseudo-hyperboloid: Mixed-curvature space for knowledge graph embeddings. Generalizes Lorentz/hyperbolic to arbitrary signature.

All implemented using JAX linear algebra. No external manifold libraries. Registered with domain="manifold".

spd_affine_invariant_distance(a: Any, b: Any) -> Any

Affine-invariant Riemannian distance between SPD matrices.

Geodesic distance on the SPD manifold: d(A, B) = ||log(A^{-1/2} B A^{-1/2})||_F

Equivalent to sqrt(sum(log(lambda_i)^2)) where lambda_i are the generalized eigenvalues of (B, A).

Invariant under congruence: d(A, B) = d(MAM^T, MBM^T) for any invertible M. This is the natural distance for covariance matrices.

Complexity: O(n^3) for generalized eigenvalue decomposition.

Parameters:

Name	Type	Description	Default
a	Any	First SPD matrix, shape (n, n).	required
b	Any	Second SPD matrix, shape (n, n).	required

Returns:

Type	Description
Any	Affine-invariant distance. Lower is better. 0.0 for identical matrices.

Examples:

>>> import jax.numpy as jnp
>>> a = jnp.eye(3) * 2.0
>>> spd_affine_invariant_distance(a, a)
0.0

spd_log_euclidean_distance(a: Any, b: Any) -> Any

Log-Euclidean distance between SPD matrices.

Maps SPD matrices to Euclidean space via matrix logarithm: d(A, B) = ||log(A) - log(B)||_F

Faster than affine-invariant but less geometrically faithful. Invariant under orthogonal transformations but NOT affine-invariant.

Complexity: O(n^3) for eigendecomposition (no matrix inverse).

Parameters:

Name	Type	Description	Default
a	Any	First SPD matrix, shape (n, n).	required
b	Any	Second SPD matrix, shape (n, n).	required

Returns:

Type	Description
Any	Log-Euclidean distance. Lower is better. 0.0 for identical matrices.

Examples:

>>> import jax.numpy as jnp
>>> a = jnp.eye(3) * 2.0
>>> spd_log_euclidean_distance(a, a)
0.0

grassmann_distance(u: Any, v: Any) -> Any

Geodesic distance on the Grassmann manifold Gr(p, n).

Distance between two p-dimensional subspaces of R^n, based on principal angles. The principal angles theta_i = arccos(sigma_i) where sigma_i are singular values of U^T V.

Distance = sqrt(sum(theta_i^2)).

Basis-independent: same subspace with different orthonormal basis gives distance 0.

Complexity: O(n*p^2) for SVD of the p x p matrix U^T V.

Parameters:

Name	Type	Description	Default
u	Any	First orthonormal matrix, shape (n, p).	required
v	Any	Second orthonormal matrix, shape (n, p).	required

Returns:

Type	Description
Any	Grassmann distance. Lower is better. 0.0 for identical subspaces.
Any	Maximum pi/2 * sqrt(p) for fully orthogonal subspaces.

Examples:

>>> import jax.numpy as jnp
>>> u = jnp.eye(3, 2)
>>> grassmann_distance(u, u)
0.0

stiefel_distance(u: Any, v: Any) -> Any

Extrinsic distance on the Stiefel manifold St(p, n).

Frobenius distance between orthonormal p-frames: ||U - V||_F.

Unlike Grassmann distance, this is basis-DEPENDENT: the same subspace with a different orthonormal basis gives a nonzero distance.

Complexity: O(n*p) for element-wise difference and norm.

Parameters:

Name	Type	Description	Default
u	Any	First orthonormal matrix, shape (n, p).	required
v	Any	Second orthonormal matrix, shape (n, p).	required

Returns:

Type	Description
Any	Stiefel distance. Lower is better. 0.0 for identical frames.

Examples:

>>> import jax.numpy as jnp
>>> u = jnp.eye(3, 2)
>>> stiefel_distance(u, u)
0.0

ultrahyperbolic_distance(a: Any, b: Any, *, signature: tuple[int, int]) -> Any

Geodesic distance on the pseudo-hyperboloid with given signature.

Pseudo-Riemannian inner product: <a,b>_{p,q} = -sum(a[:p]*b[:p]) + sum(a[p:]*b[p:])

Points live on <x,x>_{p,q} = -1. Distance = arccosh(-<a,b>_{p,q}) for timelike-separated points.

Generalizes Lorentz distance (signature=(1,n)) to arbitrary mixed-curvature spaces (UltraE pattern, Xiong et al. KDD 2022).

Special cases: - signature=(1, n): Lorentz/hyperboloid (pure hyperbolic) - signature=(0, n): Spherical (all spacelike) - signature=(p, q) with p,q > 0: Mixed curvature

Complexity: O(n) for inner product and arccosh.

Parameters:

Name	Type	Description	Default
a	Any	First point on the pseudo-hyperboloid.	required
b	Any	Second point on the pseudo-hyperboloid.	required
signature	tuple[int, int]	Tuple (p, q) specifying p timelike and q spacelike dims.	required

Returns:

Type	Description
Any	Pseudo-Riemannian distance. Lower is better. 0.0 for identical points.

Examples:

>>> import jax.numpy as jnp
>>> a = jnp.array([jnp.sqrt(2.0), 1.0, 0.0])
>>> ultrahyperbolic_distance(a, a, signature=(1, 2))
0.0