Manifolds API
This page documents the core manifold operations in Hyperbolix. Each manifold is a class that provides geometric operations and automatic dtype casting.
Overview
Hyperbolix provides four base manifold classes plus a composition class:
- Euclidean: Flat Euclidean space (baseline)
- Poincaré Ball: Conformal model of hyperbolic space
- Hyperboloid: Lorentz/Minkowski model of hyperbolic space
- Proper Velocity: Unconstrained \(\mathbb{R}^n\) model from special relativity (Chen et al. 2026)
- Product Manifold: Heterogeneous-curvature product spaces \(M_1 \times M_2 \times \dots \times M_n\) (Gu et al. 2019)
All manifolds share a common interface defined by the Manifold protocol and support:
- Automatic dtype casting: Pass
dtype=jnp.float64 for higher precision
- vmap-native methods: Methods operate on single points; use
jax.vmap for batching
- JIT compatibility: All methods are JIT-compilable
- Learnable curvature: Use the
LearnableCurvature module to add trainable curvature to any model (softplus or log/exp reparameterization, optional clamping)
Manifold Protocol
The Curvature type
Manifold methods take the curvature as a positional c: Curvature argument
(hyperbolix.manifolds.Curvature). It is the union
ScalarCurvature | Sequence[ScalarCurvature], where ScalarCurvature = float |
jax.Array: single manifolds (Poincare, Hyperboloid, ProperVelocity,
Euclidean) take a scalar c, while ProductManifold takes a sequence
of per-factor scalars. Passing a traced jax.Array (e.g. the value returned by a
LearnableCurvature call) makes the curvature differentiable.
hyperbolix.manifolds.protocol.Manifold
Bases: Protocol
Structural protocol for manifold classes.
All concrete manifold classes (Poincare, Hyperboloid,
ProperVelocity, Euclidean, ProductManifold) satisfy this
protocol without modification. For single manifolds, c is a scalar
(ScalarCurvature); for ProductManifold, c is a sequence of
length n_factors (one curvature per factor).
The method signatures use the minimal common interface so that
manifold-specific optional parameters (e.g. version_idx, atol)
do not break compatibility.
Euclidean
Flat Euclidean space (identity operations).
hyperbolix.manifolds.euclidean.Euclidean
Euclidean(dtype: dtype = jnp.float32)
Bases: ManifoldBase
Euclidean manifold with automatic dtype casting.
Provides all manifold operations with automatic casting of array inputs
to the specified dtype. Curvature is always fixed at 0.0.
Args:
dtype: Target JAX dtype for computations (default: jnp.float32)
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds.euclidean import Euclidean
>>>
>>> manifold = Euclidean(dtype=jnp.float64)
>>> x = jnp.array([1.0, 2.0])
>>> y = jnp.array([3.0, 4.0])
>>> d = manifold.dist(x, y)
Source code in hyperbolix/manifolds/euclidean.py
| def __init__(self, dtype: jnp.dtype = jnp.float32) -> None:
super().__init__(dtype, c=0.0)
|
proj
proj(
x: Float[Array, dim], c: Curvature = 0.0
) -> Float[Array, dim]
Project point onto Euclidean space (identity).
Source code in hyperbolix/manifolds/euclidean.py
| def proj(self, x: Float[Array, "dim"], c: Curvature = 0.0) -> Float[Array, "dim"]:
"""Project point onto Euclidean space (identity)."""
return _proj(self._cast(x), c)
|
addition
addition(
x: Float[Array, dim],
y: Float[Array, dim],
c: Curvature = 0.0,
) -> Float[Array, dim]
Add Euclidean points.
Source code in hyperbolix/manifolds/euclidean.py
| def addition(self, x: Float[Array, "dim"], y: Float[Array, "dim"], c: Curvature = 0.0) -> Float[Array, "dim"]:
"""Add Euclidean points."""
return _addition(self._cast(x), self._cast(y), c)
|
scalar_mul
scalar_mul(
r: float, x: Float[Array, dim], c: Curvature = 0.0
) -> Float[Array, dim]
Scalar multiplication.
Source code in hyperbolix/manifolds/euclidean.py
| def scalar_mul(self, r: float, x: Float[Array, "dim"], c: Curvature = 0.0) -> Float[Array, "dim"]:
"""Scalar multiplication."""
x = self._cast(x)
r_cast = jnp.asarray(r, dtype=x.dtype)
return _scalar_mul(r_cast, x, c) # type: ignore[arg-type]
|
dist
dist(
x: Float[Array, dim],
y: Float[Array, dim],
c: Curvature = 0.0,
) -> Float[Array, ""]
Compute distance.
Source code in hyperbolix/manifolds/euclidean.py
| def dist(self, x: Float[Array, "dim"], y: Float[Array, "dim"], c: Curvature = 0.0) -> Float[Array, ""]:
"""Compute distance."""
return _dist(self._cast(x), self._cast(y), c)
|
dist_0
dist_0(
x: Float[Array, dim], c: Curvature = 0.0
) -> Float[Array, ""]
Distance from origin.
Source code in hyperbolix/manifolds/euclidean.py
| def dist_0(self, x: Float[Array, "dim"], c: Curvature = 0.0) -> Float[Array, ""]:
"""Distance from origin."""
return _dist_0(self._cast(x), c)
|
expmap
expmap(
v: Float[Array, dim],
x: Float[Array, dim],
c: Curvature = 0.0,
) -> Float[Array, dim]
Exponential map.
Source code in hyperbolix/manifolds/euclidean.py
| def expmap(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature = 0.0) -> Float[Array, "dim"]:
"""Exponential map."""
return _expmap(self._cast(v), self._cast(x), c)
|
expmap_0
expmap_0(
v: Float[Array, dim], c: Curvature = 0.0
) -> Float[Array, dim]
Exponential map from origin.
Source code in hyperbolix/manifolds/euclidean.py
| def expmap_0(self, v: Float[Array, "dim"], c: Curvature = 0.0) -> Float[Array, "dim"]:
"""Exponential map from origin."""
return _expmap_0(self._cast(v), c)
|
retraction
retraction(
v: Float[Array, dim],
x: Float[Array, dim],
c: Curvature = 0.0,
) -> Float[Array, dim]
Retraction.
Source code in hyperbolix/manifolds/euclidean.py
| def retraction(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature = 0.0) -> Float[Array, "dim"]:
"""Retraction."""
return _retraction(self._cast(v), self._cast(x), c)
|
logmap
logmap(
y: Float[Array, dim],
x: Float[Array, dim],
c: Curvature = 0.0,
) -> Float[Array, dim]
Logarithmic map.
Source code in hyperbolix/manifolds/euclidean.py
| def logmap(self, y: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature = 0.0) -> Float[Array, "dim"]:
"""Logarithmic map."""
return _logmap(self._cast(y), self._cast(x), c)
|
logmap_0
logmap_0(
y: Float[Array, dim], c: Curvature = 0.0
) -> Float[Array, dim]
Logarithmic map from origin.
Source code in hyperbolix/manifolds/euclidean.py
| def logmap_0(self, y: Float[Array, "dim"], c: Curvature = 0.0) -> Float[Array, "dim"]:
"""Logarithmic map from origin."""
return _logmap_0(self._cast(y), c)
|
ptransp
ptransp(
v: Float[Array, dim],
x: Float[Array, dim],
y: Float[Array, dim],
c: Curvature = 0.0,
) -> Float[Array, dim]
Parallel transport.
Source code in hyperbolix/manifolds/euclidean.py
| def ptransp(
self, v: Float[Array, "dim"], x: Float[Array, "dim"], y: Float[Array, "dim"], c: Curvature = 0.0
) -> Float[Array, "dim"]:
"""Parallel transport."""
return _ptransp(self._cast(v), self._cast(x), self._cast(y), c)
|
ptransp_0
ptransp_0(
v: Float[Array, dim],
y: Float[Array, dim],
c: Curvature = 0.0,
) -> Float[Array, dim]
Parallel transport from origin.
Source code in hyperbolix/manifolds/euclidean.py
| def ptransp_0(self, v: Float[Array, "dim"], y: Float[Array, "dim"], c: Curvature = 0.0) -> Float[Array, "dim"]:
"""Parallel transport from origin."""
return _ptransp_0(self._cast(v), self._cast(y), c)
|
tangent_inner
tangent_inner(
u: Float[Array, dim],
v: Float[Array, dim],
x: Float[Array, dim],
c: Curvature = 0.0,
) -> Float[Array, ""]
Tangent inner product.
Source code in hyperbolix/manifolds/euclidean.py
| def tangent_inner(
self, u: Float[Array, "dim"], v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature = 0.0
) -> Float[Array, ""]:
"""Tangent inner product."""
return _tangent_inner(self._cast(u), self._cast(v), self._cast(x), c)
|
tangent_norm
tangent_norm(
v: Float[Array, dim],
x: Float[Array, dim],
c: Curvature = 0.0,
) -> Float[Array, ""]
Tangent norm.
Source code in hyperbolix/manifolds/euclidean.py
| def tangent_norm(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature = 0.0) -> Float[Array, ""]:
"""Tangent norm."""
return _tangent_norm(self._cast(v), self._cast(x), c)
|
egrad2rgrad
egrad2rgrad(
grad: Float[Array, dim],
x: Float[Array, dim],
c: Curvature = 0.0,
) -> Float[Array, dim]
Euclidean to Riemannian gradient.
Source code in hyperbolix/manifolds/euclidean.py
| def egrad2rgrad(self, grad: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature = 0.0) -> Float[Array, "dim"]:
"""Euclidean to Riemannian gradient."""
return _egrad2rgrad(self._cast(grad), self._cast(x), c)
|
tangent_proj
tangent_proj(
v: Float[Array, dim],
x: Float[Array, dim],
c: Curvature = 0.0,
) -> Float[Array, dim]
Project onto tangent space.
Source code in hyperbolix/manifolds/euclidean.py
| def tangent_proj(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature = 0.0) -> Float[Array, "dim"]:
"""Project onto tangent space."""
return _tangent_proj(self._cast(v), self._cast(x), c)
|
is_in_manifold
is_in_manifold(
x: Float[Array, dim], c: Curvature = 0.0
) -> Array
Check if on manifold.
Source code in hyperbolix/manifolds/euclidean.py
| def is_in_manifold(self, x: Float[Array, "dim"], c: Curvature = 0.0) -> Array:
"""Check if on manifold."""
return _is_in_manifold(self._cast(x), c)
|
is_in_tangent_space
is_in_tangent_space(
v: Float[Array, dim],
x: Float[Array, dim],
c: Curvature = 0.0,
) -> Array
Check if in tangent space.
Source code in hyperbolix/manifolds/euclidean.py
| def is_in_tangent_space(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature = 0.0) -> Array:
"""Check if in tangent space."""
return _is_in_tangent_space(self._cast(v), self._cast(x), c)
|
Poincaré Ball
The Poincaré ball model with Möbius operations.
Distance Versions
The Poincaré dist method has a version_idx parameter selecting between 3 formulations:
VERSION_MOBIUS_DIRECT (0): Möbius addition formula (default, fastest)VERSION_MOBIUS (1): Möbius via additionVERSION_METRIC_TENSOR (2): Direct metric tensor integration
Constants are available as poincare.VERSION_MOBIUS_DIRECT etc., or from
hyperbolix.manifolds.poincare.
Apollonian weak metric
apollonian_dist(x, y, c) is the non-symmetric Apollonian weak metric \(\delta\)
(Papadopoulos & Troyanov, Weak metrics on Euclidean domains, Thm 2) — a weak metric,
not a geodesic distance:
\[\delta_c(x,y) = \log\!\left(\frac{\sqrt{c}\,\lVert x-y\rVert + \sqrt{c^2\lVert x\rVert^2\lVert y\rVert^2 - 2c\langle x,y\rangle + 1}}{1-c\lVert y\rVert^2}\right)\]
It satisfies \(\delta(x,x)=0\), \(\delta\ge 0\) and the triangle inequality, but
\(\delta(x,y) \neq \delta(y,x)\) in general. Its symmetrization recovers the geodesic distance:
\(\delta(x,y) + \delta(y,x) = \sqrt{c}\cdot\) dist(x, y, c).
Warning
The antisymmetric part of \(\delta\) is an exact coboundary (a difference of a per-point
potential), so it carries no circulation and is useless as an asymmetric quasimetric energy.
For that, use the busemann coordinate below with an external quasimetric combinator.
Busemann function (Chen et al. 2026)
busemann(x, v, c) is the closed-form point-to-horosphere coordinate \(B^v(x)\) for a unit
ideal direction \(v\in\mathbb{S}^{n-1}\) — the horospherical analog of the point-to-hyperplane
compute_mlr/compute_mlr_pp. v must be unit-norm (not normalized internally). It is an
intrinsic quantity, so Poincare.busemann and Hyperboloid.busemann agree under
poincare_to_hyperboloid, and \(B^v(\text{origin})=0\). Backs the *Busemann MLR/FC layers.
\[\mathbb{P}^n:\ B^v(x) = \tfrac{1}{\sqrt c}\log\!\frac{\lVert v-\sqrt c\,x\rVert^2}{1-c\lVert x\rVert^2}
\qquad\quad \mathbb{L}^n:\ B^v(x) = \tfrac{1}{\sqrt c}\log\!\big(\sqrt c\,(x_t-\langle x_s,v\rangle)\big)\]
hyperbolix.manifolds.poincare.Poincare
Poincare(dtype: dtype = jnp.float32, *, c: float = 1.0)
Bases: ManifoldBase
Poincaré ball manifold with automatic dtype casting.
Provides all manifold operations with automatic casting of array inputs
to the specified dtype. This eliminates the need for manual casting and
provides better numerical stability control.
Args:
dtype: Target JAX dtype for computations (default: jnp.float32)
c: Curvature value (default: 1.0). Must be positive.
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds.poincare import Poincare, VERSION_MOBIUS_DIRECT
>>>
>>> # Create manifold with float64 for better precision
>>> manifold = Poincare(dtype=jnp.float64)
>>>
>>> # Static curvature
>>> manifold = Poincare(c=0.1)
>>> c = manifold.c # returns 0.1
Source code in hyperbolix/manifolds/_base.py
| def __init__(
self,
dtype: jnp.dtype = jnp.float32,
*,
c: float = 1.0,
) -> None:
self.dtype = dtype
self._c_val = c
|
proj
proj(
x: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Project point onto Poincaré ball by clipping norm.
Source code in hyperbolix/manifolds/poincare.py
| def proj(self, x: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Project point onto Poincaré ball by clipping norm."""
return _proj(self._cast(x), c)
|
gyration
gyration(
x: Float[Array, dim],
y: Float[Array, dim],
z: Float[Array, dim],
c: Curvature,
) -> Float[Array, dim]
Compute gyration gyr[x,y]z to restore commutativity.
Source code in hyperbolix/manifolds/poincare.py
| def gyration(
self, x: Float[Array, "dim"], y: Float[Array, "dim"], z: Float[Array, "dim"], c: Curvature
) -> Float[Array, "dim"]:
"""Compute gyration gyr[x,y]z to restore commutativity."""
return _gyration(self._cast(x), self._cast(y), self._cast(z), c)
|
addition
addition(
x: Float[Array, dim], y: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Möbius gyrovector addition x ⊕ y.
Source code in hyperbolix/manifolds/poincare.py
| def addition(self, x: Float[Array, "dim"], y: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Möbius gyrovector addition x ⊕ y."""
return _addition(self._cast(x), self._cast(y), c)
|
scalar_mul
scalar_mul(
r: float, x: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Scalar multiplication r ⊗ x on Poincaré ball.
Source code in hyperbolix/manifolds/poincare.py
| def scalar_mul(self, r: float, x: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Scalar multiplication r ⊗ x on Poincaré ball."""
x = self._cast(x)
r_cast = jnp.asarray(r, dtype=x.dtype)
return _scalar_mul(r_cast, x, c) # type: ignore[arg-type]
|
dist
dist(
x: Float[Array, dim],
y: Float[Array, dim],
c: Curvature,
version_idx: int = VERSION_MOBIUS_DIRECT,
) -> Float[Array, ""]
Compute geodesic distance between Poincaré ball points.
Source code in hyperbolix/manifolds/poincare.py
| def dist(
self,
x: Float[Array, "dim"],
y: Float[Array, "dim"],
c: Curvature,
version_idx: int = VERSION_MOBIUS_DIRECT,
) -> Float[Array, ""]:
"""Compute geodesic distance between Poincaré ball points."""
return _dist(self._cast(x), self._cast(y), c, version_idx)
|
dist_0
dist_0(
x: Float[Array, dim],
c: Curvature,
version_idx: int = VERSION_MOBIUS_DIRECT,
) -> Float[Array, ""]
Compute geodesic distance from Poincaré ball origin.
Source code in hyperbolix/manifolds/poincare.py
| def dist_0(self, x: Float[Array, "dim"], c: Curvature, version_idx: int = VERSION_MOBIUS_DIRECT) -> Float[Array, ""]:
"""Compute geodesic distance from Poincaré ball origin."""
return _dist_0(self._cast(x), c, version_idx)
|
apollonian_dist
apollonian_dist(
x: Float[Array, dim], y: Float[Array, dim], c: Curvature
) -> Float[Array, ""]
Apollonian weak metric δ(x, y) — non-symmetric; symmetrizes to √c·dist(x, y).
.. warning::
Although δ is non-symmetric, its antisymmetric part is an exact coboundary
(δ(x, y) - δ(y, x) is a difference of a per-point potential), so it carries
no circulation and is useless as a quasimetric energy. Do not reach for this
expecting genuine asymmetry — use a :meth:busemann coordinate fed to an external
quasimetric combinator (IQE/MRN) instead.
Source code in hyperbolix/manifolds/poincare.py
| def apollonian_dist(self, x: Float[Array, "dim"], y: Float[Array, "dim"], c: Curvature) -> Float[Array, ""]:
"""Apollonian weak metric δ(x, y) — non-symmetric; symmetrizes to √c·dist(x, y).
.. warning::
Although ``δ`` is non-symmetric, its antisymmetric part is an exact **coboundary**
(``δ(x, y) - δ(y, x)`` is a difference of a per-point potential), so it carries
**no circulation** and is useless as a quasimetric energy. Do not reach for this
expecting genuine asymmetry — use a :meth:`busemann` coordinate fed to an external
quasimetric combinator (IQE/MRN) instead.
"""
return _apollonian_dist(self._cast(x), self._cast(y), c)
|
expmap
expmap(
v: Float[Array, dim], x: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Exponential map: map tangent vector v at point x to manifold.
Source code in hyperbolix/manifolds/poincare.py
| def expmap(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Exponential map: map tangent vector v at point x to manifold."""
return _expmap(self._cast(v), self._cast(x), c)
|
expmap_0
expmap_0(
v: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Exponential map from origin: map tangent vector v at origin to manifold.
Source code in hyperbolix/manifolds/poincare.py
| def expmap_0(self, v: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Exponential map from origin: map tangent vector v at origin to manifold."""
return _expmap_0(self._cast(v), c)
|
retraction
retraction(
v: Float[Array, dim], x: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Retraction: first-order approximation of exponential map.
Source code in hyperbolix/manifolds/poincare.py
| def retraction(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Retraction: first-order approximation of exponential map."""
return _retraction(self._cast(v), self._cast(x), c)
|
logmap
logmap(
y: Float[Array, dim], x: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Logarithmic map: map point y to tangent space at point x.
Source code in hyperbolix/manifolds/poincare.py
| def logmap(self, y: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Logarithmic map: map point y to tangent space at point x."""
return _logmap(self._cast(y), self._cast(x), c)
|
logmap_0
logmap_0(
y: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Logarithmic map from origin: map point y to tangent space at origin.
Source code in hyperbolix/manifolds/poincare.py
| def logmap_0(self, y: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Logarithmic map from origin: map point y to tangent space at origin."""
return _logmap_0(self._cast(y), c)
|
ptransp
ptransp(
v: Float[Array, dim],
x: Float[Array, dim],
y: Float[Array, dim],
c: Curvature,
) -> Float[Array, dim]
Parallel transport tangent vector v from point x to point y.
Source code in hyperbolix/manifolds/poincare.py
| def ptransp(
self, v: Float[Array, "dim"], x: Float[Array, "dim"], y: Float[Array, "dim"], c: Curvature
) -> Float[Array, "dim"]:
"""Parallel transport tangent vector v from point x to point y."""
return _ptransp(self._cast(v), self._cast(x), self._cast(y), c)
|
ptransp_0
ptransp_0(
v: Float[Array, dim], y: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Parallel transport tangent vector v from origin to point y.
Source code in hyperbolix/manifolds/poincare.py
| def ptransp_0(self, v: Float[Array, "dim"], y: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Parallel transport tangent vector v from origin to point y."""
return _ptransp_0(self._cast(v), self._cast(y), c)
|
tangent_inner
tangent_inner(
u: Float[Array, dim],
v: Float[Array, dim],
x: Float[Array, dim],
c: Curvature,
) -> Float[Array, ""]
Compute inner product of tangent vectors u and v at point x.
Source code in hyperbolix/manifolds/poincare.py
| def tangent_inner(
self, u: Float[Array, "dim"], v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature
) -> Float[Array, ""]:
"""Compute inner product of tangent vectors u and v at point x."""
return _tangent_inner(self._cast(u), self._cast(v), self._cast(x), c)
|
tangent_norm
tangent_norm(
v: Float[Array, dim], x: Float[Array, dim], c: Curvature
) -> Float[Array, ""]
Compute norm of tangent vector v at point x.
Source code in hyperbolix/manifolds/poincare.py
| def tangent_norm(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature) -> Float[Array, ""]:
"""Compute norm of tangent vector v at point x."""
return _tangent_norm(self._cast(v), self._cast(x), c)
|
egrad2rgrad
egrad2rgrad(
grad: Float[Array, dim],
x: Float[Array, dim],
c: Curvature,
) -> Float[Array, dim]
Convert Euclidean gradient to Riemannian gradient.
Source code in hyperbolix/manifolds/poincare.py
| def egrad2rgrad(self, grad: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Convert Euclidean gradient to Riemannian gradient."""
return _egrad2rgrad(self._cast(grad), self._cast(x), c)
|
tangent_proj
tangent_proj(
v: Float[Array, dim], x: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Project vector v onto tangent space at point x.
Source code in hyperbolix/manifolds/poincare.py
| def tangent_proj(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Project vector v onto tangent space at point x."""
return _tangent_proj(self._cast(v), self._cast(x), c)
|
is_in_manifold
is_in_manifold(
x: Float[Array, dim], c: Curvature, atol: float = 1e-05
) -> Array
Check if point x lies in Poincaré ball.
Source code in hyperbolix/manifolds/poincare.py
| def is_in_manifold(self, x: Float[Array, "dim"], c: Curvature, atol: float = 1e-5) -> Array:
"""Check if point x lies in Poincaré ball."""
return _is_in_manifold(self._cast(x), c, atol)
|
is_in_tangent_space
is_in_tangent_space(
v: Float[Array, dim], x: Float[Array, dim], c: Curvature
) -> Array
Check if vector v lies in tangent space at point x.
Source code in hyperbolix/manifolds/poincare.py
| def is_in_tangent_space(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature) -> Array:
"""Check if vector v lies in tangent space at point x."""
return _is_in_tangent_space(self._cast(v), self._cast(x), c)
|
conformal_factor(
x: Float[Array, "... dim"], c: Curvature
) -> Float[Array, "... 1"]
Numerically stable conformal factor lambda(x) = 2 / (1 - c||x||^2).
Batch-compatible version that handles arbitrary leading dimensions.
Source code in hyperbolix/manifolds/poincare.py
| def conformal_factor(self, x: Float[Array, "... dim"], c: Curvature) -> Float[Array, "... 1"]:
"""Numerically stable conformal factor lambda(x) = 2 / (1 - c||x||^2).
Batch-compatible version that handles arbitrary leading dimensions.
"""
return _conformal_factor_batch(self._cast(x), c)
|
embed_spatial_0
embed_spatial_0(
v_spatial: Float[Array, "... n"],
) -> Float[Array, "... n"]
Identity embedding (the ball has no time coordinate). Kept for API parity.
Source code in hyperbolix/manifolds/poincare.py
| def embed_spatial_0(self, v_spatial: Float[Array, "... n"]) -> Float[Array, "... n"]:
"""Identity embedding (the ball has no time coordinate). Kept for API parity."""
return _embed_spatial_0(self._cast(v_spatial))
|
compute_mlr_pp
compute_mlr_pp(
x: Float[Array, "batch in_dim"],
z: Float[Array, "out_dim in_dim"],
r: Float[Array, "out_dim 1"],
c: Curvature,
clamping_factor: float,
smoothing_factor: float,
min_enorm: float = 1e-15,
) -> Float[Array, "batch out_dim"]
Compute HNN++ multinomial linear regression on the Poincare ball.
Source code in hyperbolix/manifolds/poincare.py
| def compute_mlr_pp(
self,
x: Float[Array, "batch in_dim"],
z: Float[Array, "out_dim in_dim"],
r: Float[Array, "out_dim 1"],
c: Curvature,
clamping_factor: float,
smoothing_factor: float,
min_enorm: float = 1e-15,
) -> Float[Array, "batch out_dim"]:
"""Compute HNN++ multinomial linear regression on the Poincare ball."""
return _compute_mlr_pp(self._cast(x), self._cast(z), self._cast(r), c, clamping_factor, smoothing_factor, min_enorm)
|
beta_concat
beta_concat(
points: Float[Array, "M n_i"], c: Curvature
) -> Float[Array, n]
Beta-concatenation of M equal-dimensional Poincaré ball points.
Source code in hyperbolix/manifolds/poincare.py
| def beta_concat(self, points: Float[Array, "M n_i"], c: Curvature) -> Float[Array, "n"]:
"""Beta-concatenation of M equal-dimensional Poincaré ball points."""
return _beta_concat(self._cast(points), c)
|
busemann
busemann(
x: Float[Array, dim], v: Float[Array, dim], c: Curvature
) -> Float[Array, ""]
Closed-form Poincaré Busemann function B^v(x) = (1/√c)·log(‖v - √c·x‖²/(1 - c‖x‖²)).
Point-to-horosphere coordinate (Chen et al. 2026, Eq. 3). v must be a unit
direction — it is not normalized internally. Single point (d,) → scalar; use
:func:jax.vmap for batching and over a direction set. B^v(origin) = 0, and it
matches :meth:Hyperboloid.busemann under the Poincaré↔Hyperboloid isometry.
See Also
For an asymmetric quasimetric energy, compose this Busemann coordinate with an
external Euclidean quasimetric (e.g. IQE/MRN). Do not reach for
:meth:apollonian_dist expecting asymmetry — it is a coboundary (symmetrizes to
√c·dist) and cannot deliver circulation.
Source code in hyperbolix/manifolds/poincare.py
| def busemann(self, x: Float[Array, "dim"], v: Float[Array, "dim"], c: Curvature) -> Float[Array, ""]:
"""Closed-form Poincaré Busemann function ``B^v(x) = (1/√c)·log(‖v - √c·x‖²/(1 - c‖x‖²))``.
Point-to-horosphere coordinate (Chen et al. 2026, Eq. 3). ``v`` must be a *unit*
direction — it is **not** normalized internally. Single point ``(d,)`` → scalar; use
:func:`jax.vmap` for batching and over a direction set. ``B^v(origin) = 0``, and it
matches :meth:`Hyperboloid.busemann` under the Poincaré↔Hyperboloid isometry.
See Also
--------
For an *asymmetric* quasimetric energy, compose this Busemann coordinate with an
external Euclidean quasimetric (e.g. IQE/MRN). Do **not** reach for
:meth:`apollonian_dist` expecting asymmetry — it is a coboundary (symmetrizes to
``√c·dist``) and cannot deliver circulation.
"""
return _busemann(self._cast(x), self._cast(v), c)
|
Hyperboloid
The hyperboloid (Lorentz) model with Minkowski geometry.
Lorentz Operations
The Hyperboloid class includes specialized operations for convolutional layers:
lorentz_boost: Lorentz boost transformationdistance_rescale: Distance-based rescalinghcat: Lorentz direct concatenation for convolutionslog_radius_concat: log-radius–preserving concatenation (digamma-scaled hcat; Shi et al. 2026, Sec. 4.3)
hyperbolix.manifolds.hyperboloid.Hyperboloid
Hyperboloid(dtype: dtype = jnp.float32, *, c: float = 1.0)
Bases: ManifoldBase
Hyperboloid manifold with automatic dtype casting.
Provides all manifold operations with automatic casting of array inputs
to the specified dtype. This eliminates the need for manual casting and
provides better numerical stability control.
Args:
dtype: Target JAX dtype for computations (default: jnp.float32)
c: Curvature value (default: 1.0). Must be positive.
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds.hyperboloid import Hyperboloid, VERSION_DEFAULT
>>>
>>> # Create manifold with float64 for better precision
>>> manifold = Hyperboloid(dtype=jnp.float64)
>>>
>>> # Custom curvature
>>> manifold = Hyperboloid(c=0.5)
>>> c = manifold.c # returns 0.5
Source code in hyperbolix/manifolds/_base.py
| def __init__(
self,
dtype: jnp.dtype = jnp.float32,
*,
c: float = 1.0,
) -> None:
self.dtype = dtype
self._c_val = c
|
create_origin
create_origin(
c: Curvature, dim: int
) -> Float[Array, dim_plus_1]
Create hyperboloid origin [1/√c, 0, ..., 0].
Source code in hyperbolix/manifolds/hyperboloid.py
| def create_origin(self, c: Curvature, dim: int) -> Float[Array, "dim_plus_1"]:
"""Create hyperboloid origin [1/√c, 0, ..., 0]."""
return _create_origin(c, dim, self.dtype)
|
minkowski_inner
minkowski_inner(
x: Float[Array, dim_plus_1], y: Float[Array, dim_plus_1]
) -> Float[Array, ""]
Compute Minkowski inner product ⟨x, y⟩_L = -x₀y₀ + ⟨x_rest, y_rest⟩.
Source code in hyperbolix/manifolds/hyperboloid.py
| def minkowski_inner(self, x: Float[Array, "dim_plus_1"], y: Float[Array, "dim_plus_1"]) -> Float[Array, ""]:
"""Compute Minkowski inner product ⟨x, y⟩_L = -x₀y₀ + ⟨x_rest, y_rest⟩."""
return _minkowski_inner(self._cast(x), self._cast(y))
|
proj
proj(
x: Float[Array, dim_plus_1], c: Curvature
) -> Float[Array, dim_plus_1]
Project point onto hyperboloid.
Source code in hyperbolix/manifolds/hyperboloid.py
| def proj(self, x: Float[Array, "dim_plus_1"], c: Curvature) -> Float[Array, "dim_plus_1"]:
"""Project point onto hyperboloid."""
return _proj(self._cast(x), c)
|
proj_batch
proj_batch(
x: Float[Array, "... dim_plus_1"], c: Curvature
) -> Float[Array, "... dim_plus_1"]
Project batched points onto hyperboloid (handles arbitrary leading dimensions).
Source code in hyperbolix/manifolds/hyperboloid.py
| def proj_batch(self, x: Float[Array, "... dim_plus_1"], c: Curvature) -> Float[Array, "... dim_plus_1"]:
"""Project batched points onto hyperboloid (handles arbitrary leading dimensions)."""
return _proj_batch(self._cast(x), c)
|
addition
addition(
x: Float[Array, dim_plus_1],
y: Float[Array, dim_plus_1],
c: Curvature,
) -> Float[Array, dim_plus_1]
Lorentz gyrovector addition x ⊕ y = Exp_x(PT_{0→x}(Log_0(y))).
The intrinsic Lorentz addition of Chen et al. (2025b) / Shi et al. (2026, Eq. 1).
Forms a gyrocommutative gyrogroup with identity = origin and inverse
⊖x = (-1) ⊙ x = [x₀, -x_s]; matches Poincaré Möbius addition under the
stereographic isometry. scalar_mul provides the companion gyro scaling (Eq. 2).
Source code in hyperbolix/manifolds/hyperboloid.py
| def addition(
self, x: Float[Array, "dim_plus_1"], y: Float[Array, "dim_plus_1"], c: Curvature
) -> Float[Array, "dim_plus_1"]:
"""Lorentz gyrovector addition ``x ⊕ y = Exp_x(PT_{0→x}(Log_0(y)))``.
The intrinsic Lorentz addition of Chen et al. (2025b) / Shi et al. (2026, Eq. 1).
Forms a gyrocommutative gyrogroup with identity = origin and inverse
``⊖x = (-1) ⊙ x = [x₀, -x_s]``; matches Poincaré Möbius addition under the
stereographic isometry. ``scalar_mul`` provides the companion gyro scaling (Eq. 2).
"""
return _addition(self._cast(x), self._cast(y), c)
|
scalar_mul
scalar_mul(
r: float, x: Float[Array, dim_plus_1], c: Curvature
) -> Float[Array, dim_plus_1]
Scalar multiplication on hyperboloid.
Source code in hyperbolix/manifolds/hyperboloid.py
| def scalar_mul(self, r: float, x: Float[Array, "dim_plus_1"], c: Curvature) -> Float[Array, "dim_plus_1"]:
"""Scalar multiplication on hyperboloid."""
x = self._cast(x)
r_cast = jnp.asarray(r, dtype=x.dtype)
return _scalar_mul(r_cast, x, c) # type: ignore[arg-type]
|
dist
dist(
x: Float[Array, dim_plus_1],
y: Float[Array, dim_plus_1],
c: Curvature,
version_idx: int = VERSION_DEFAULT,
) -> Float[Array, ""]
Compute geodesic distance between hyperboloid points.
Source code in hyperbolix/manifolds/hyperboloid.py
| def dist(
self,
x: Float[Array, "dim_plus_1"],
y: Float[Array, "dim_plus_1"],
c: Curvature,
version_idx: int = VERSION_DEFAULT,
) -> Float[Array, ""]:
"""Compute geodesic distance between hyperboloid points."""
return _dist(self._cast(x), self._cast(y), c, version_idx)
|
dist_0
dist_0(
x: Float[Array, dim_plus_1],
c: Curvature,
version_idx: int = VERSION_DEFAULT,
) -> Float[Array, ""]
Compute geodesic distance from hyperboloid origin.
Source code in hyperbolix/manifolds/hyperboloid.py
| def dist_0(self, x: Float[Array, "dim_plus_1"], c: Curvature, version_idx: int = VERSION_DEFAULT) -> Float[Array, ""]:
"""Compute geodesic distance from hyperboloid origin."""
return _dist_0(self._cast(x), c, version_idx)
|
sqdist
sqdist(
x: Float[Array, dim_plus_1],
y: Float[Array, dim_plus_1],
c: Curvature,
) -> Float[Array, ""]
Squared Lorentzian distance d_L²(x, y) = -2/c - 2⟨x,y⟩_L (Law et al. 2019).
A fast, acosh-free dissimilarity that is monotone in the geodesic :meth:dist
(d_L² = (2/c)(cosh(√c·d) - 1)) but is not the squared geodesic distance and not
a true metric. Prefer it where a smooth monotone distance proxy suffices (contrastive /
commitment / attention scores); use :meth:dist when the geodesic length is required.
Source code in hyperbolix/manifolds/hyperboloid.py
| def sqdist(self, x: Float[Array, "dim_plus_1"], y: Float[Array, "dim_plus_1"], c: Curvature) -> Float[Array, ""]:
"""Squared Lorentzian distance ``d_L²(x, y) = -2/c - 2⟨x,y⟩_L`` (Law et al. 2019).
A fast, ``acosh``-free dissimilarity that is monotone in the geodesic :meth:`dist`
(``d_L² = (2/c)(cosh(√c·d) - 1)``) but is **not** the squared geodesic distance and **not**
a true metric. Prefer it where a smooth monotone distance proxy suffices (contrastive /
commitment / attention scores); use :meth:`dist` when the geodesic length is required.
"""
return _sqdist(self._cast(x), self._cast(y), c)
|
expmap
expmap(
v: Float[Array, dim_plus_1],
x: Float[Array, dim_plus_1],
c: Curvature,
) -> Float[Array, dim_plus_1]
Exponential map: map tangent vector v at point x to manifold.
Source code in hyperbolix/manifolds/hyperboloid.py
| def expmap(self, v: Float[Array, "dim_plus_1"], x: Float[Array, "dim_plus_1"], c: Curvature) -> Float[Array, "dim_plus_1"]:
"""Exponential map: map tangent vector v at point x to manifold."""
return _expmap(self._cast(v), self._cast(x), c)
|
expmap_0
expmap_0(
v: Float[Array, dim_plus_1], c: Curvature
) -> Float[Array, dim_plus_1]
Exponential map from origin.
Source code in hyperbolix/manifolds/hyperboloid.py
| def expmap_0(self, v: Float[Array, "dim_plus_1"], c: Curvature) -> Float[Array, "dim_plus_1"]:
"""Exponential map from origin."""
return _expmap_0(self._cast(v), c)
|
retraction
retraction(
v: Float[Array, dim_plus_1],
x: Float[Array, dim_plus_1],
c: Curvature,
) -> Float[Array, dim_plus_1]
Retraction: first-order approximation of exponential map.
Source code in hyperbolix/manifolds/hyperboloid.py
| def retraction(
self, v: Float[Array, "dim_plus_1"], x: Float[Array, "dim_plus_1"], c: Curvature
) -> Float[Array, "dim_plus_1"]:
"""Retraction: first-order approximation of exponential map."""
return _retraction(self._cast(v), self._cast(x), c)
|
logmap
logmap(
y: Float[Array, dim_plus_1],
x: Float[Array, dim_plus_1],
c: Curvature,
) -> Float[Array, dim_plus_1]
Logarithmic map: map point y to tangent space at point x.
Source code in hyperbolix/manifolds/hyperboloid.py
| def logmap(self, y: Float[Array, "dim_plus_1"], x: Float[Array, "dim_plus_1"], c: Curvature) -> Float[Array, "dim_plus_1"]:
"""Logarithmic map: map point y to tangent space at point x."""
return _logmap(self._cast(y), self._cast(x), c)
|
logmap_0
logmap_0(
y: Float[Array, dim_plus_1], c: Curvature
) -> Float[Array, dim_plus_1]
Logarithmic map from origin.
Source code in hyperbolix/manifolds/hyperboloid.py
| def logmap_0(self, y: Float[Array, "dim_plus_1"], c: Curvature) -> Float[Array, "dim_plus_1"]:
"""Logarithmic map from origin."""
return _logmap_0(self._cast(y), c)
|
ptransp
ptransp(
v: Float[Array, dim_plus_1],
x: Float[Array, dim_plus_1],
y: Float[Array, dim_plus_1],
c: Curvature,
) -> Float[Array, dim_plus_1]
Parallel transport tangent vector v from point x to point y.
Source code in hyperbolix/manifolds/hyperboloid.py
| def ptransp(
self, v: Float[Array, "dim_plus_1"], x: Float[Array, "dim_plus_1"], y: Float[Array, "dim_plus_1"], c: Curvature
) -> Float[Array, "dim_plus_1"]:
"""Parallel transport tangent vector v from point x to point y."""
return _ptransp(self._cast(v), self._cast(x), self._cast(y), c)
|
ptransp_0
ptransp_0(
v: Float[Array, dim_plus_1],
y: Float[Array, dim_plus_1],
c: Curvature,
) -> Float[Array, dim_plus_1]
Parallel transport tangent vector v from origin to point y.
Source code in hyperbolix/manifolds/hyperboloid.py
| def ptransp_0(
self, v: Float[Array, "dim_plus_1"], y: Float[Array, "dim_plus_1"], c: Curvature
) -> Float[Array, "dim_plus_1"]:
"""Parallel transport tangent vector v from origin to point y."""
return _ptransp_0(self._cast(v), self._cast(y), c)
|
tangent_inner
tangent_inner(
u: Float[Array, dim_plus_1],
v: Float[Array, dim_plus_1],
x: Float[Array, dim_plus_1],
c: Curvature,
) -> Float[Array, ""]
Compute inner product of tangent vectors u and v at point x.
Source code in hyperbolix/manifolds/hyperboloid.py
| def tangent_inner(
self, u: Float[Array, "dim_plus_1"], v: Float[Array, "dim_plus_1"], x: Float[Array, "dim_plus_1"], c: Curvature
) -> Float[Array, ""]:
"""Compute inner product of tangent vectors u and v at point x."""
return _tangent_inner(self._cast(u), self._cast(v), self._cast(x), c)
|
tangent_norm
tangent_norm(
v: Float[Array, dim_plus_1],
x: Float[Array, dim_plus_1],
c: Curvature,
) -> Float[Array, ""]
Compute norm of tangent vector v at point x.
Source code in hyperbolix/manifolds/hyperboloid.py
| def tangent_norm(self, v: Float[Array, "dim_plus_1"], x: Float[Array, "dim_plus_1"], c: Curvature) -> Float[Array, ""]:
"""Compute norm of tangent vector v at point x."""
return _tangent_norm(self._cast(v), self._cast(x), c)
|
egrad2rgrad
egrad2rgrad(
grad: Float[Array, dim_plus_1],
x: Float[Array, dim_plus_1],
c: Curvature,
) -> Float[Array, dim_plus_1]
Convert Euclidean gradient to Riemannian gradient.
Source code in hyperbolix/manifolds/hyperboloid.py
| def egrad2rgrad(
self, grad: Float[Array, "dim_plus_1"], x: Float[Array, "dim_plus_1"], c: Curvature
) -> Float[Array, "dim_plus_1"]:
"""Convert Euclidean gradient to Riemannian gradient."""
return _egrad2rgrad(self._cast(grad), self._cast(x), c)
|
tangent_proj
tangent_proj(
v: Float[Array, dim_plus_1],
x: Float[Array, dim_plus_1],
c: Curvature,
) -> Float[Array, dim_plus_1]
Project vector v onto tangent space at point x.
Source code in hyperbolix/manifolds/hyperboloid.py
| def tangent_proj(
self, v: Float[Array, "dim_plus_1"], x: Float[Array, "dim_plus_1"], c: Curvature
) -> Float[Array, "dim_plus_1"]:
"""Project vector v onto tangent space at point x."""
return _tangent_proj(self._cast(v), self._cast(x), c)
|
is_in_manifold
is_in_manifold(
x: Float[Array, dim_plus_1],
c: Curvature,
atol: float = 0.0001,
) -> Array
Check if point x lies on hyperboloid.
Source code in hyperbolix/manifolds/hyperboloid.py
| def is_in_manifold(self, x: Float[Array, "dim_plus_1"], c: Curvature, atol: float = 1e-4) -> Array:
"""Check if point x lies on hyperboloid."""
return _is_in_manifold(self._cast(x), c, atol)
|
is_in_tangent_space
is_in_tangent_space(
v: Float[Array, dim_plus_1],
x: Float[Array, dim_plus_1],
c: Curvature,
) -> Array
Check if vector v lies in tangent space at point x.
Source code in hyperbolix/manifolds/hyperboloid.py
| def is_in_tangent_space(self, v: Float[Array, "dim_plus_1"], x: Float[Array, "dim_plus_1"], c: Curvature) -> Array:
"""Check if vector v lies in tangent space at point x."""
return _is_in_tangent_space(self._cast(v), self._cast(x), c)
|
hcat
hcat(
points: Float[Array, "N n"], c: Curvature = 1.0
) -> Float[Array, dN_plus_1]
Hyperbolic concatenation of N points into one point.
Source code in hyperbolix/manifolds/hyperboloid.py
| def hcat(
self,
points: Float[Array, "N n"],
c: Curvature = 1.0,
) -> Float[Array, "dN_plus_1"]:
"""Hyperbolic concatenation of N points into one point."""
return _hcat(self._cast(points), c)
|
log_radius_concat
log_radius_concat(
points: Float[Array, "N n"], c: Curvature = 1.0
) -> Float[Array, dN_plus_1]
Log-radius-preserving concatenation of N points (Shi et al. 2026, Sec. 4.3).
Like :meth:hcat, but rescales each block's spatial part by a digamma factor so the
expected log spatial radius is invariant to the number of concatenated blocks. Reduces
to :meth:hcat when N == 1.
Source code in hyperbolix/manifolds/hyperboloid.py
| def log_radius_concat(
self,
points: Float[Array, "N n"],
c: Curvature = 1.0,
) -> Float[Array, "dN_plus_1"]:
"""Log-radius-preserving concatenation of N points (Shi et al. 2026, Sec. 4.3).
Like :meth:`hcat`, but rescales each block's spatial part by a digamma factor so the
expected log spatial radius is invariant to the number of concatenated blocks. Reduces
to :meth:`hcat` when ``N == 1``.
"""
return _log_radius_concat(self._cast(points), c)
|
embed_spatial_0
embed_spatial_0(
v_spatial: Float[Array, "... n"],
) -> Float[Array, "... n_plus_1"]
Embed spatial vector as tangent vector at origin.
Source code in hyperbolix/manifolds/hyperboloid.py
| def embed_spatial_0(self, v_spatial: Float[Array, "... n"]) -> Float[Array, "... n_plus_1"]:
"""Embed spatial vector as tangent vector at origin."""
return _embed_spatial_0(self._cast(v_spatial))
|
compute_mlr
compute_mlr(
x: Float[Array, "batch in_dim"],
z: Float[Array, "out_dim in_dim_minus_1"],
r: Float[Array, "out_dim 1"],
c: Curvature,
clamping_factor: float,
smoothing_factor: float,
min_enorm: float = 1e-15,
) -> Float[Array, "batch out_dim"]
Compute multinomial linear regression on hyperboloid.
Source code in hyperbolix/manifolds/hyperboloid.py
| def compute_mlr(
self,
x: Float[Array, "batch in_dim"],
z: Float[Array, "out_dim in_dim_minus_1"],
r: Float[Array, "out_dim 1"],
c: Curvature,
clamping_factor: float,
smoothing_factor: float,
min_enorm: float = 1e-15,
) -> Float[Array, "batch out_dim"]:
"""Compute multinomial linear regression on hyperboloid."""
return _compute_mlr(self._cast(x), self._cast(z), self._cast(r), c, clamping_factor, smoothing_factor, min_enorm)
|
busemann
busemann(
x: Float[Array, dim_plus_1],
v: Float[Array, dim],
c: Curvature,
) -> Float[Array, ""]
Closed-form Lorentz Busemann function B^v(x) = (1/√c)·log(√c·(x_t - ⟨x_s, v⟩)).
Point-to-horosphere coordinate (Chen et al. 2026, Eq. 4). v must be a unit
spatial direction (dim d) — it is not normalized internally. Single point
(d+1,) → scalar; use :func:jax.vmap for batching and over a direction set.
B^v(origin) = 0.
See Also
For an asymmetric quasimetric energy, compose this Busemann coordinate with an
external Euclidean quasimetric (e.g. IQE/MRN). Do not reach for
:meth:Poincare.apollonian_dist expecting asymmetry — it is a coboundary
(symmetrizes to √c·dist) and cannot deliver circulation.
Source code in hyperbolix/manifolds/hyperboloid.py
| def busemann(self, x: Float[Array, "dim_plus_1"], v: Float[Array, "dim"], c: Curvature) -> Float[Array, ""]:
"""Closed-form Lorentz Busemann function ``B^v(x) = (1/√c)·log(√c·(x_t - ⟨x_s, v⟩))``.
Point-to-horosphere coordinate (Chen et al. 2026, Eq. 4). ``v`` must be a *unit*
spatial direction (dim ``d``) — it is **not** normalized internally. Single point
``(d+1,)`` → scalar; use :func:`jax.vmap` for batching and over a direction set.
``B^v(origin) = 0``.
See Also
--------
For an *asymmetric* quasimetric energy, compose this Busemann coordinate with an
external Euclidean quasimetric (e.g. IQE/MRN). Do **not** reach for
:meth:`Poincare.apollonian_dist` expecting asymmetry — it is a coboundary
(symmetrizes to ``√c·dist``) and cannot deliver circulation.
"""
return _busemann(self._cast(x), self._cast(v), c)
|
Proper Velocity
The Proper Velocity (PV) model — an unconstrained \(\mathbb{R}^n\) representation of hyperbolic geometry rooted in special relativity's proper velocity (Ungar 2022, Ch. 10). PV is algebraically a gyrovector space isomorphic to the Poincaré ball via \(\pi(x) = (\beta_x / (1 + \beta_x)) \cdot x\), and carries a Riemannian metric making that isomorphism an isometry.
Why Proper Velocity?
Unlike the bounded Poincaré ball (points must stay in the open unit ball) and the constrained hyperboloid (points must satisfy \(\langle x, x \rangle_L = -1/c\)), PV points live in all of \(\mathbb{R}^n\) with no constraint. This gives:
- No projection step after updates — the manifold is \(\mathbb{R}^n\)
- Better numerical stability for large radii (no boundary collapse, no Lorentz constraint drift)
- Drop-in with Euclidean optimizers — the retraction reduces to \(x + v\)
The metric \(g_x(u,v) = \langle u, v\rangle - c\beta_x^2\langle x,u\rangle\langle x,v\rangle\) still gives sectional curvature \(-c\); only the coordinates change.
Convention
The paper uses curvature \(K < 0\) with \(\beta_x = 1/\sqrt{1 - K\|x\|^2}\). Hyperbolix keeps the \(c > 0\) convention (sectional curvature \(-c\)), substituting \(K = -c\), so \(\beta_x = 1/\sqrt{1 + c\|x\|^2}\).
hyperbolix.manifolds.proper_velocity.ProperVelocity
ProperVelocity(
dtype: dtype = jnp.float32, *, c: float = 1.0
)
Bases: ManifoldBase
Proper Velocity (PV) manifold with automatic dtype casting.
PV is an unconstrained representation of hyperbolic geometry rooted in
special relativity's proper velocity (Ungar 2022, Ch. 10). Points live
in R^n without any manifold constraint, which gives better numerical
stability for large radii than the bounded Poincaré ball or the
constrained hyperboloid (Chen et al. 2026, Tables 1-3).
Args:
dtype: Target JAX dtype for computations (default: jnp.float32).
c: Curvature value (default: 1.0). Must be positive.
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds.proper_velocity import ProperVelocity
>>>
>>> manifold = ProperVelocity(dtype=jnp.float64)
>>> x = jnp.array([0.1, 0.2], dtype=jnp.float32)
>>> y = jnp.array([0.3, 0.4], dtype=jnp.float32)
>>> d = manifold.dist(x, y, c=1.0)
>>> d.dtype # float64
Source code in hyperbolix/manifolds/_base.py
| def __init__(
self,
dtype: jnp.dtype = jnp.float32,
*,
c: float = 1.0,
) -> None:
self.dtype = dtype
self._c_val = c
|
create_origin
create_origin(c: Curvature, dim: int) -> Float[Array, dim]
Create the PV origin (zero vector in R^n).
Source code in hyperbolix/manifolds/proper_velocity.py
| def create_origin(self, c: Curvature, dim: int) -> Float[Array, "dim"]:
"""Create the PV origin (zero vector in R^n)."""
return _create_origin(c, dim, self.dtype)
|
beta
beta(
x: Float[Array, dim], c: Curvature
) -> Float[Array, ""]
PV beta factor β_x = 1/√(1 + c·||x||²).
Source code in hyperbolix/manifolds/proper_velocity.py
| def beta(self, x: Float[Array, "dim"], c: Curvature) -> Float[Array, ""]:
"""PV beta factor β_x = 1/√(1 + c·||x||²)."""
return _beta(self._cast(x), c)
|
proj
proj(
x: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Projection onto PV (replaces non-finite values; PV is unconstrained).
Source code in hyperbolix/manifolds/proper_velocity.py
| def proj(self, x: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Projection onto PV (replaces non-finite values; PV is unconstrained)."""
return _proj(self._cast(x), c)
|
addition
addition(
x: Float[Array, dim], y: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
PV gyroaddition x ⊕_U y.
Source code in hyperbolix/manifolds/proper_velocity.py
| def addition(self, x: Float[Array, "dim"], y: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""PV gyroaddition x ⊕_U y."""
return _addition(self._cast(x), self._cast(y), c)
|
scalar_mul
scalar_mul(
r: float, x: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
PV scalar multiplication r ⊗_U x.
Source code in hyperbolix/manifolds/proper_velocity.py
| def scalar_mul(self, r: float, x: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""PV scalar multiplication r ⊗_U x."""
x = self._cast(x)
r_cast = jnp.asarray(r, dtype=x.dtype)
return _scalar_mul(r_cast, x, c) # type: ignore[arg-type]
|
dist
dist(
x: Float[Array, dim],
y: Float[Array, dim],
c: Curvature,
version_idx: int = VERSION_DEFAULT,
) -> Float[Array, ""]
Geodesic distance between PV points.
Source code in hyperbolix/manifolds/proper_velocity.py
| def dist(
self,
x: Float[Array, "dim"],
y: Float[Array, "dim"],
c: Curvature,
version_idx: int = VERSION_DEFAULT,
) -> Float[Array, ""]:
"""Geodesic distance between PV points."""
del version_idx # only one implementation currently
return _dist(self._cast(x), self._cast(y), c)
|
dist_0
dist_0(
x: Float[Array, dim],
c: Curvature,
version_idx: int = VERSION_DEFAULT,
) -> Float[Array, ""]
Geodesic distance from the PV origin.
Source code in hyperbolix/manifolds/proper_velocity.py
| def dist_0(self, x: Float[Array, "dim"], c: Curvature, version_idx: int = VERSION_DEFAULT) -> Float[Array, ""]:
"""Geodesic distance from the PV origin."""
del version_idx
return _dist_0(self._cast(x), c)
|
expmap
expmap(
v: Float[Array, dim], x: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Exponential map at x.
Source code in hyperbolix/manifolds/proper_velocity.py
| def expmap(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Exponential map at x."""
return _expmap(self._cast(v), self._cast(x), c)
|
expmap_0
expmap_0(
v: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Exponential map from the origin.
Source code in hyperbolix/manifolds/proper_velocity.py
| def expmap_0(self, v: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Exponential map from the origin."""
return _expmap_0(self._cast(v), c)
|
logmap
logmap(
y: Float[Array, dim], x: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Logarithmic map at x.
Source code in hyperbolix/manifolds/proper_velocity.py
| def logmap(self, y: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Logarithmic map at x."""
return _logmap(self._cast(y), self._cast(x), c)
|
logmap_0
logmap_0(
y: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Logarithmic map to the origin.
Source code in hyperbolix/manifolds/proper_velocity.py
| def logmap_0(self, y: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Logarithmic map to the origin."""
return _logmap_0(self._cast(y), c)
|
retraction
retraction(
v: Float[Array, dim], x: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Euclidean retraction (exact for PV).
Source code in hyperbolix/manifolds/proper_velocity.py
| def retraction(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Euclidean retraction (exact for PV)."""
return _retraction(self._cast(v), self._cast(x), c)
|
ptransp
ptransp(
v: Float[Array, dim],
x: Float[Array, dim],
y: Float[Array, dim],
c: Curvature,
) -> Float[Array, dim]
Parallel transport v from T_x PV to T_y PV.
Source code in hyperbolix/manifolds/proper_velocity.py
| def ptransp(
self,
v: Float[Array, "dim"],
x: Float[Array, "dim"],
y: Float[Array, "dim"],
c: Curvature,
) -> Float[Array, "dim"]:
"""Parallel transport v from T_x PV to T_y PV."""
return _ptransp(self._cast(v), self._cast(x), self._cast(y), c)
|
ptransp_0
ptransp_0(
v: Float[Array, dim], y: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Parallel transport v from T_0 PV to T_y PV.
Source code in hyperbolix/manifolds/proper_velocity.py
| def ptransp_0(self, v: Float[Array, "dim"], y: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Parallel transport v from T_0 PV to T_y PV."""
return _ptransp_0(self._cast(v), self._cast(y), c)
|
tangent_inner
tangent_inner(
u: Float[Array, dim],
v: Float[Array, dim],
x: Float[Array, dim],
c: Curvature,
) -> Float[Array, ""]
Riemannian inner product ⟨u, v⟩_x.
Source code in hyperbolix/manifolds/proper_velocity.py
| def tangent_inner(
self,
u: Float[Array, "dim"],
v: Float[Array, "dim"],
x: Float[Array, "dim"],
c: Curvature,
) -> Float[Array, ""]:
"""Riemannian inner product ⟨u, v⟩_x."""
return _tangent_inner(self._cast(u), self._cast(v), self._cast(x), c)
|
tangent_norm
tangent_norm(
v: Float[Array, dim], x: Float[Array, dim], c: Curvature
) -> Float[Array, ""]
Riemannian norm ||v||_x.
Source code in hyperbolix/manifolds/proper_velocity.py
| def tangent_norm(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature) -> Float[Array, ""]:
"""Riemannian norm ||v||_x."""
return _tangent_norm(self._cast(v), self._cast(x), c)
|
tangent_proj
tangent_proj(
v: Float[Array, dim], x: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Tangent-space projection (identity for PV).
Source code in hyperbolix/manifolds/proper_velocity.py
| def tangent_proj(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Tangent-space projection (identity for PV)."""
return _tangent_proj(self._cast(v), self._cast(x), c)
|
egrad2rgrad
egrad2rgrad(
grad: Float[Array, dim],
x: Float[Array, dim],
c: Curvature,
) -> Float[Array, dim]
Convert Euclidean gradient to Riemannian gradient.
Source code in hyperbolix/manifolds/proper_velocity.py
| def egrad2rgrad(self, grad: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature) -> Float[Array, "dim"]:
"""Convert Euclidean gradient to Riemannian gradient."""
return _egrad2rgrad(self._cast(grad), self._cast(x), c)
|
is_in_manifold
is_in_manifold(
x: Float[Array, dim], c: Curvature, atol: float = 1e-05
) -> Array
Check that all entries are finite (PV has no constraint).
Source code in hyperbolix/manifolds/proper_velocity.py
| def is_in_manifold(self, x: Float[Array, "dim"], c: Curvature, atol: float = 1e-5) -> Array:
"""Check that all entries are finite (PV has no constraint)."""
return _is_in_manifold(self._cast(x), c, atol)
|
is_in_tangent_space
is_in_tangent_space(
v: Float[Array, dim], x: Float[Array, dim], c: Curvature
) -> Array
Check that v has finite entries (T_x PV = R^n).
Source code in hyperbolix/manifolds/proper_velocity.py
| def is_in_tangent_space(self, v: Float[Array, "dim"], x: Float[Array, "dim"], c: Curvature) -> Array:
"""Check that v has finite entries (T_x PV = R^n)."""
return _is_in_tangent_space(self._cast(v), self._cast(x), c)
|
embed_spatial_0
embed_spatial_0(
v_spatial: Float[Array, "... n"],
) -> Float[Array, "... n"]
Identity embedding (no time coordinate). Kept for API parity.
Source code in hyperbolix/manifolds/proper_velocity.py
| def embed_spatial_0(self, v_spatial: Float[Array, "... n"]) -> Float[Array, "... n"]:
"""Identity embedding (no time coordinate). Kept for API parity."""
return _embed_spatial_0(self._cast(v_spatial))
|
compute_mlr
compute_mlr(
x: Float[Array, "batch in_dim"],
z: Float[Array, "out_dim in_dim"],
r: Float[Array, "out_dim 1"],
c: Curvature,
clamping_factor: float,
smoothing_factor: float,
min_enorm: float = 1e-15,
) -> Float[Array, "batch out_dim"]
PV multinomial logistic regression (paper Thm 5.2).
Source code in hyperbolix/manifolds/proper_velocity.py
| def compute_mlr(
self,
x: Float[Array, "batch in_dim"],
z: Float[Array, "out_dim in_dim"],
r: Float[Array, "out_dim 1"],
c: Curvature,
clamping_factor: float,
smoothing_factor: float,
min_enorm: float = 1e-15,
) -> Float[Array, "batch out_dim"]:
"""PV multinomial logistic regression (paper Thm 5.2)."""
return _compute_mlr(
self._cast(x),
self._cast(z),
self._cast(r),
c,
clamping_factor,
smoothing_factor,
min_enorm,
)
|
Product Manifold
Heterogeneous-curvature product space \(P = M_1 \times M_2 \times \dots \times M_n\) where each factor \(M_i\) can be any base manifold (Poincaré, Hyperboloid, Euclidean, Proper Velocity) with its own curvature \(c_i\). Points are represented as flat concatenated arrays of shape (total_dim,).
The geodesic distance on a product Riemannian manifold is Pythagorean over component distances:
\[d_P(x, y) \;=\; \sqrt{\sum_{i=1}^{n} d_{M_i}(x_i, y_i)^2}\]
where \(x_i\), \(y_i\) are the per-factor slices of the flat points.
Per-factor c argument
Every geometry method (dist, expmap, logmap, proj, origin, …) takes a positional c argument that must be a sequence of length n_factors — one curvature per factor. There is no scalar fallback and no broadcast: pass product.curvatures for static curvatures, or a tuple built from LearnableCurvature calls for trainable ones. ProductManifold satisfies the Manifold protocol — the protocol-level Curvature type unions scalar and sequence-of-scalars, so isinstance(product, Manifold) is True and generic code typed against Manifold accepts product instances. The product itself has no c attribute — read factor-stored values via product.curvatures.
Static vs learnable curvature
Factor instances may carry an initial c (Hyperboloid(c=1.0)), but ProductManifold never reads it in its geometry methods — it is exposed via product.curvatures as a convenience default. For learnable curvature, instantiate one LearnableCurvature per factor on your nnx.Module and pass c=(self.curv_a(), self.curv_b(), ...) to the product. See the Manifolds User Guide — Curvature in ProductManifold for the full pattern.
hyperbolix.manifolds.product.ProductManifold
ProductManifold(
*factors: tuple[ManifoldBase, int],
dtype: dtype = jnp.float32,
)
Product manifold P = M1 x M2 x ... x Mn.
Each factor is specified as (manifold_instance, dim) where dim is
the point dimension (array slice width): ambient for Hyperboloid (d+1),
spatial for Poincaré/Euclidean/ProperVelocity (d).
Curvature is supplied at call time as a per-factor sequence:
product.dist(x, y, (c0, c1, ...)) # explicit tuple
product.dist(x, y, product.curvatures) # use factor-stored values
Factor instances may carry an initial curvature value (Hyperboloid(c=1.0)),
but that value is not consulted by ProductManifold's methods — it is
only useful as a static default to read off via product.curvatures.
ProductManifold satisfies the Manifold protocol with c typed as
Curvature (a union of scalar and sequence-of-scalars), so generic code
written against Manifold accepts product instances too. The product
itself has no c attribute — there is no single scalar curvature.
Args:
*factors: Tuples of (manifold_instance, dim). At least one required.
dtype: Target JAX dtype for computations (default: jnp.float32).
Examples:
>>> product = ProductManifold((Hyperboloid(c=1.0), 5), (Poincare(c=0.1), 3))
>>> c = product.curvatures # (1.0, 0.1)
>>> x = product.origin(c)
>>> parts = product.split(x)
>>> d = product.dist(x, y, c)
From a signature (repeated factors):
>>> product = ProductManifold.from_signature(
... (Hyperboloid, 3, 8, 1.0), # 8 copies of 3D-ambient Hyperboloid
... (Poincare, 2, 4, 0.1), # 4 copies of 2D Poincaré
... )
Source code in hyperbolix/manifolds/product.py
| def __init__(
self,
*factors: tuple[ManifoldBase, int],
dtype: jnp.dtype = jnp.float32,
) -> None:
if len(factors) < 1:
raise ValueError("ProductManifold requires at least one factor")
self.dtype = dtype
manifolds = []
dims = []
slices = []
pos = 0
for i, (manifold, dim) in enumerate(factors):
if not isinstance(manifold, ManifoldBase):
raise TypeError(f"Factor {i}: expected ManifoldBase instance, got {type(manifold).__name__}")
if dim < 1:
raise ValueError(f"Factor {i}: dim must be >= 1, got {dim}")
manifolds.append(manifold)
dims.append(dim)
slices.append(slice(pos, pos + dim))
pos += dim
self._factors = tuple(manifolds)
self._dims = tuple(dims)
self._slices = tuple(slices)
self._total_dim = pos
|
curvatures
property
curvatures: tuple[ScalarCurvature, ...]
Per-factor curvatures stored on the factor instances.
Useful as the c argument when curvature is static, e.g.
product.dist(x, y, product.curvatures). For learnable curvature,
build the tuple from LearnableCurvature calls instead.
split
split(x: Float[Array, total_dim]) -> tuple[Array, ...]
Split a product point into per-factor components.
Source code in hyperbolix/manifolds/product.py
| def split(self, x: Float[Array, "total_dim"]) -> tuple[Array, ...]:
"""Split a product point into per-factor components."""
return tuple(x[s] for s in self._slices)
|
combine
combine(*parts: Array) -> Float[Array, total_dim]
Combine per-factor components into a product point.
Source code in hyperbolix/manifolds/product.py
| def combine(self, *parts: Array) -> Float[Array, "total_dim"]:
"""Combine per-factor components into a product point."""
if len(parts) != len(self._factors):
raise ValueError(f"Expected {len(self._factors)} parts, got {len(parts)}")
return jnp.concatenate(list(parts))
|
origin
origin(c: Curvature) -> Float[Array, total_dim]
Construct the product origin point under per-factor curvatures.
Uses expmap_0(zeros, c[i]) per factor, which returns the manifold
origin for all manifold types without isinstance checks.
Source code in hyperbolix/manifolds/product.py
| def origin(self, c: Curvature) -> Float[Array, "total_dim"]:
"""Construct the product origin point under per-factor curvatures.
Uses ``expmap_0(zeros, c[i])`` per factor, which returns the manifold
origin for all manifold types without isinstance checks.
"""
cs = self._validate_c(c)
parts = []
for m, dim, c_i in zip(self._factors, self._dims, cs, strict=True):
zero_tangent = jnp.zeros(dim, dtype=self.dtype)
parts.append(m.expmap_0(zero_tangent, c_i))
return jnp.concatenate(parts)
|
proj
proj(
x: Float[Array, total_dim], c: Curvature
) -> Float[Array, total_dim]
Project point onto product manifold (per-factor projection).
Source code in hyperbolix/manifolds/product.py
| def proj(
self,
x: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, "total_dim"]:
"""Project point onto product manifold (per-factor projection)."""
cs = self._validate_c(c)
x = self._cast(x)
parts = self.split(x)
return jnp.concatenate([m.proj(p, c_i) for m, p, c_i in zip(self._factors, parts, cs, strict=True)])
|
addition
addition(
x: Float[Array, total_dim],
y: Float[Array, total_dim],
c: Curvature,
) -> Float[Array, total_dim]
Manifold addition (per-factor).
Source code in hyperbolix/manifolds/product.py
| def addition(
self,
x: Float[Array, "total_dim"],
y: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, "total_dim"]:
"""Manifold addition (per-factor)."""
cs = self._validate_c(c)
x, y = self._cast(x), self._cast(y)
x_parts, y_parts = self.split(x), self.split(y)
return jnp.concatenate(
[m.addition(xp, yp, c_i) for m, xp, yp, c_i in zip(self._factors, x_parts, y_parts, cs, strict=True)]
)
|
scalar_mul
scalar_mul(
r: float, x: Float[Array, total_dim], c: Curvature
) -> Float[Array, total_dim]
Scalar multiplication (same scalar r applied to all factors).
Source code in hyperbolix/manifolds/product.py
| def scalar_mul(
self,
r: float,
x: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, "total_dim"]:
"""Scalar multiplication (same scalar r applied to all factors)."""
cs = self._validate_c(c)
x = self._cast(x)
parts = self.split(x)
return jnp.concatenate([m.scalar_mul(r, p, c_i) for m, p, c_i in zip(self._factors, parts, cs, strict=True)])
|
expmap
expmap(
v: Float[Array, total_dim],
x: Float[Array, total_dim],
c: Curvature,
) -> Float[Array, total_dim]
Exponential map (per-factor).
Source code in hyperbolix/manifolds/product.py
| def expmap(
self,
v: Float[Array, "total_dim"],
x: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, "total_dim"]:
"""Exponential map (per-factor)."""
cs = self._validate_c(c)
v, x = self._cast(v), self._cast(x)
v_parts, x_parts = self.split(v), self.split(x)
return jnp.concatenate(
[m.expmap(vp, xp, c_i) for m, vp, xp, c_i in zip(self._factors, v_parts, x_parts, cs, strict=True)]
)
|
expmap_0
expmap_0(
v: Float[Array, total_dim], c: Curvature
) -> Float[Array, total_dim]
Exponential map from origin (per-factor).
Source code in hyperbolix/manifolds/product.py
| def expmap_0(
self,
v: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, "total_dim"]:
"""Exponential map from origin (per-factor)."""
cs = self._validate_c(c)
v = self._cast(v)
parts = self.split(v)
return jnp.concatenate([m.expmap_0(p, c_i) for m, p, c_i in zip(self._factors, parts, cs, strict=True)])
|
logmap
logmap(
y: Float[Array, total_dim],
x: Float[Array, total_dim],
c: Curvature,
) -> Float[Array, total_dim]
Logarithmic map (per-factor).
Source code in hyperbolix/manifolds/product.py
| def logmap(
self,
y: Float[Array, "total_dim"],
x: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, "total_dim"]:
"""Logarithmic map (per-factor)."""
cs = self._validate_c(c)
y, x = self._cast(y), self._cast(x)
y_parts, x_parts = self.split(y), self.split(x)
return jnp.concatenate(
[m.logmap(yp, xp, c_i) for m, yp, xp, c_i in zip(self._factors, y_parts, x_parts, cs, strict=True)]
)
|
logmap_0
logmap_0(
y: Float[Array, total_dim], c: Curvature
) -> Float[Array, total_dim]
Logarithmic map to origin (per-factor).
Source code in hyperbolix/manifolds/product.py
| def logmap_0(
self,
y: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, "total_dim"]:
"""Logarithmic map to origin (per-factor)."""
cs = self._validate_c(c)
y = self._cast(y)
parts = self.split(y)
return jnp.concatenate([m.logmap_0(p, c_i) for m, p, c_i in zip(self._factors, parts, cs, strict=True)])
|
retraction
retraction(
v: Float[Array, total_dim],
x: Float[Array, total_dim],
c: Curvature,
) -> Float[Array, total_dim]
Retraction (per-factor).
Source code in hyperbolix/manifolds/product.py
| def retraction(
self,
v: Float[Array, "total_dim"],
x: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, "total_dim"]:
"""Retraction (per-factor)."""
cs = self._validate_c(c)
v, x = self._cast(v), self._cast(x)
v_parts, x_parts = self.split(v), self.split(x)
return jnp.concatenate(
[m.retraction(vp, xp, c_i) for m, vp, xp, c_i in zip(self._factors, v_parts, x_parts, cs, strict=True)]
)
|
ptransp
ptransp(
v: Float[Array, total_dim],
x: Float[Array, total_dim],
y: Float[Array, total_dim],
c: Curvature,
) -> Float[Array, total_dim]
Parallel transport of v from x to y (per-factor).
Source code in hyperbolix/manifolds/product.py
| def ptransp(
self,
v: Float[Array, "total_dim"],
x: Float[Array, "total_dim"],
y: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, "total_dim"]:
"""Parallel transport of v from x to y (per-factor)."""
cs = self._validate_c(c)
v, x, y = self._cast(v), self._cast(x), self._cast(y)
v_parts, x_parts, y_parts = self.split(v), self.split(x), self.split(y)
return jnp.concatenate(
[
m.ptransp(vp, xp, yp, c_i)
for m, vp, xp, yp, c_i in zip(self._factors, v_parts, x_parts, y_parts, cs, strict=True)
]
)
|
ptransp_0
ptransp_0(
v: Float[Array, total_dim],
y: Float[Array, total_dim],
c: Curvature,
) -> Float[Array, total_dim]
Parallel transport from origin to y (per-factor).
Source code in hyperbolix/manifolds/product.py
| def ptransp_0(
self,
v: Float[Array, "total_dim"],
y: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, "total_dim"]:
"""Parallel transport from origin to y (per-factor)."""
cs = self._validate_c(c)
v, y = self._cast(v), self._cast(y)
v_parts, y_parts = self.split(v), self.split(y)
return jnp.concatenate(
[m.ptransp_0(vp, yp, c_i) for m, vp, yp, c_i in zip(self._factors, v_parts, y_parts, cs, strict=True)]
)
|
tangent_inner
tangent_inner(
u: Float[Array, total_dim],
v: Float[Array, total_dim],
x: Float[Array, total_dim],
c: Curvature,
) -> Float[Array, ""]
Riemannian inner product (sum of per-factor inner products).
Source code in hyperbolix/manifolds/product.py
| def tangent_inner(
self,
u: Float[Array, "total_dim"],
v: Float[Array, "total_dim"],
x: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, ""]:
"""Riemannian inner product (sum of per-factor inner products)."""
cs = self._validate_c(c)
u, v, x = self._cast(u), self._cast(v), self._cast(x)
u_parts, v_parts, x_parts = self.split(u), self.split(v), self.split(x)
inners = jnp.stack(
[
m.tangent_inner(up, vp, xp, c_i)
for m, up, vp, xp, c_i in zip(self._factors, u_parts, v_parts, x_parts, cs, strict=True)
]
)
return jnp.sum(inners)
|
tangent_norm
tangent_norm(
v: Float[Array, total_dim],
x: Float[Array, total_dim],
c: Curvature,
) -> Float[Array, ""]
Riemannian norm (sqrt of tangent inner product with itself).
Source code in hyperbolix/manifolds/product.py
| def tangent_norm(
self,
v: Float[Array, "total_dim"],
x: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, ""]:
"""Riemannian norm (sqrt of tangent inner product with itself)."""
return jnp.sqrt(jnp.maximum(self.tangent_inner(v, v, x, c), 0.0))
|
egrad2rgrad
egrad2rgrad(
grad: Float[Array, total_dim],
x: Float[Array, total_dim],
c: Curvature,
) -> Float[Array, total_dim]
Convert Euclidean gradient to Riemannian gradient (per-factor).
Source code in hyperbolix/manifolds/product.py
| def egrad2rgrad(
self,
grad: Float[Array, "total_dim"],
x: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, "total_dim"]:
"""Convert Euclidean gradient to Riemannian gradient (per-factor)."""
cs = self._validate_c(c)
grad, x = self._cast(grad), self._cast(x)
g_parts, x_parts = self.split(grad), self.split(x)
return jnp.concatenate(
[m.egrad2rgrad(gp, xp, c_i) for m, gp, xp, c_i in zip(self._factors, g_parts, x_parts, cs, strict=True)]
)
|
tangent_proj
tangent_proj(
v: Float[Array, total_dim],
x: Float[Array, total_dim],
c: Curvature,
) -> Float[Array, total_dim]
Project vector onto tangent space (per-factor).
Source code in hyperbolix/manifolds/product.py
| def tangent_proj(
self,
v: Float[Array, "total_dim"],
x: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, "total_dim"]:
"""Project vector onto tangent space (per-factor)."""
cs = self._validate_c(c)
v, x = self._cast(v), self._cast(x)
v_parts, x_parts = self.split(v), self.split(x)
return jnp.concatenate(
[m.tangent_proj(vp, xp, c_i) for m, vp, xp, c_i in zip(self._factors, v_parts, x_parts, cs, strict=True)]
)
|
component_dist
component_dist(
x: Float[Array, total_dim],
y: Float[Array, total_dim],
c: Curvature,
) -> Float[Array, n_factors]
Per-factor geodesic distances as a vector.
Source code in hyperbolix/manifolds/product.py
| def component_dist(
self,
x: Float[Array, "total_dim"],
y: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, "n_factors"]:
"""Per-factor geodesic distances as a vector."""
cs = self._validate_c(c)
x, y = self._cast(x), self._cast(y)
x_parts, y_parts = self.split(x), self.split(y)
return jnp.stack([m.dist(xp, yp, c_i) for m, xp, yp, c_i in zip(self._factors, x_parts, y_parts, cs, strict=True)])
|
dist
dist(
x: Float[Array, total_dim],
y: Float[Array, total_dim],
c: Curvature,
) -> Float[Array, ""]
Product distance (L2/Riemannian): d = sqrt(sum d_i^2).
Source code in hyperbolix/manifolds/product.py
| def dist(
self,
x: Float[Array, "total_dim"],
y: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, ""]:
"""Product distance (L2/Riemannian): d = sqrt(sum d_i^2)."""
d_per_factor = self.component_dist(x, y, c)
return jnp.sqrt(jnp.sum(d_per_factor**2))
|
dist_0
dist_0(
x: Float[Array, total_dim], c: Curvature
) -> Float[Array, ""]
Distance from origin (L2/Riemannian).
Source code in hyperbolix/manifolds/product.py
| def dist_0(
self,
x: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, ""]:
"""Distance from origin (L2/Riemannian)."""
cs = self._validate_c(c)
x = self._cast(x)
parts = self.split(x)
d_sq = jnp.stack([m.dist_0(p, c_i) ** 2 for m, p, c_i in zip(self._factors, parts, cs, strict=True)])
return jnp.sqrt(jnp.sum(d_sq))
|
dist_l1
dist_l1(
x: Float[Array, total_dim],
y: Float[Array, total_dim],
c: Curvature,
) -> Float[Array, ""]
L1 product distance: d = sum d_i.
Source code in hyperbolix/manifolds/product.py
| def dist_l1(
self,
x: Float[Array, "total_dim"],
y: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, ""]:
"""L1 product distance: d = sum d_i."""
return jnp.sum(self.component_dist(x, y, c))
|
dist_min
dist_min(
x: Float[Array, total_dim],
y: Float[Array, total_dim],
c: Curvature,
) -> Float[Array, ""]
Min product distance: d = min d_i.
Warning:
This is NOT a true metric: positive-definiteness fails. Two
distinct points x != y that agree on any single factor
(x_i == y_i) yield d = 0. Use only when the
"closest-factor" interpretation is semantically intended;
for an actual distance use dist (L2) or dist_l1.
Source code in hyperbolix/manifolds/product.py
| def dist_min(
self,
x: Float[Array, "total_dim"],
y: Float[Array, "total_dim"],
c: Curvature,
) -> Float[Array, ""]:
"""Min product distance: d = min d_i.
Warning:
This is NOT a true metric: positive-definiteness fails. Two
distinct points ``x != y`` that agree on any single factor
(``x_i == y_i``) yield ``d = 0``. Use only when the
"closest-factor" interpretation is semantically intended;
for an actual distance use ``dist`` (L2) or ``dist_l1``.
"""
return jnp.min(self.component_dist(x, y, c))
|
is_in_manifold
is_in_manifold(
x: Float[Array, total_dim], c: Curvature
) -> Array
Check if point is on product manifold (all factors valid).
Source code in hyperbolix/manifolds/product.py
| def is_in_manifold(
self,
x: Float[Array, "total_dim"],
c: Curvature,
) -> Array:
"""Check if point is on product manifold (all factors valid)."""
cs = self._validate_c(c)
x = self._cast(x)
parts = self.split(x)
checks = [m.is_in_manifold(p, c_i) for m, p, c_i in zip(self._factors, parts, cs, strict=True)]
return jnp.all(jnp.stack(checks))
|
is_in_tangent_space
is_in_tangent_space(
v: Float[Array, total_dim],
x: Float[Array, total_dim],
c: Curvature,
) -> Array
Check if vector is in tangent space at x (all factors valid).
Source code in hyperbolix/manifolds/product.py
| def is_in_tangent_space(
self,
v: Float[Array, "total_dim"],
x: Float[Array, "total_dim"],
c: Curvature,
) -> Array:
"""Check if vector is in tangent space at x (all factors valid)."""
cs = self._validate_c(c)
v, x = self._cast(v), self._cast(x)
v_parts, x_parts = self.split(v), self.split(x)
checks = [
m.is_in_tangent_space(vp, xp, c_i) for m, vp, xp, c_i in zip(self._factors, v_parts, x_parts, cs, strict=True)
]
return jnp.all(jnp.stack(checks))
|
from_signature
classmethod
from_signature(
*specs: tuple, dtype: dtype = jnp.float32
) -> ProductManifold
Create product manifold from a signature specification.
Each spec is one of:
- (ManifoldClass, dim, count) — uses c=1.0 as factor init.
- (ManifoldClass, dim, count, curvature) — sets factor init c.
The curvature here is an initial / default value stored on the
factor; it is only consulted via product.curvatures. Geometry
methods take c at call time.
Euclidean factors silently ignore the curvature field
(Euclidean has fixed c=0; see Euclidean.__init__).
Args:
*specs: Signature tuples (3- or 4-tuple as described above).
dtype: Target JAX dtype.
Returns:
ProductManifold instance.
Examples:
>>> pm = ProductManifold.from_signature(
... (Hyperboloid, 3, 8, 1.0),
... (Poincare, 2, 4, 0.1),
... (Euclidean, 8, 1),
... )
Source code in hyperbolix/manifolds/product.py
| @classmethod
def from_signature(
cls,
*specs: tuple,
dtype: jnp.dtype = jnp.float32,
) -> "ProductManifold":
"""Create product manifold from a signature specification.
Each spec is one of:
- ``(ManifoldClass, dim, count)`` — uses ``c=1.0`` as factor init.
- ``(ManifoldClass, dim, count, curvature)`` — sets factor init ``c``.
The ``curvature`` here is an **initial / default** value stored on the
factor; it is only consulted via ``product.curvatures``. Geometry
methods take ``c`` at call time.
Euclidean factors silently ignore the ``curvature`` field
(Euclidean has fixed ``c=0``; see ``Euclidean.__init__``).
Args:
*specs: Signature tuples (3- or 4-tuple as described above).
dtype: Target JAX dtype.
Returns:
ProductManifold instance.
Examples:
>>> pm = ProductManifold.from_signature(
... (Hyperboloid, 3, 8, 1.0),
... (Poincare, 2, 4, 0.1),
... (Euclidean, 8, 1),
... )
"""
from .euclidean import Euclidean
factors: list[tuple[ManifoldBase, int]] = []
for spec in specs:
if len(spec) == 3:
manifold_cls, dim, count = spec
curvature = 1.0
elif len(spec) == 4:
manifold_cls, dim, count, curvature = spec
else:
raise ValueError(
"Expected 3-tuple (ManifoldClass, dim, count) or 4-tuple "
f"(ManifoldClass, dim, count, curvature); got {len(spec)}-tuple"
)
for _ in range(count):
if issubclass(manifold_cls, Euclidean):
instance = manifold_cls(dtype=dtype)
else:
instance = manifold_cls(dtype=dtype, c=curvature)
factors.append((instance, dim))
return cls(*factors, dtype=dtype)
|
Isometry Mappings
Distance-preserving maps between the Poincaré ball, hyperboloid, and Proper
Velocity (PV) models — all coordinate models of the same hyperbolic space.
Provides Poincaré ↔ Hyperboloid, Poincaré ↔ PV (PVNN Eq. 4), and the direct
Hyperboloid ↔ PV map (PV coordinates are the space-like part of the 4-velocity).
hyperbolix.manifolds.isometry_mappings
Isometry mappings between hyperbolic manifold models.
This module implements distance-preserving transformations (isometries) between
different models of hyperbolic geometry. All functions operate on single points
and use JAX's vmap for batch operations.
Supported Models (curvature c > 0, sectional curvature -c):
- Hyperboloid model (Lorentz model): Points in R^(d+1) satisfying ⟨x,x⟩_L = -1/c
- Poincaré ball model: Points in R^d with ||y||² < 1/c
- Proper Velocity (PV) model: Unconstrained points in R^d (Chen et al. 2026)
Provided maps (all exact, distance-preserving, mutually consistent):
- Poincaré ↔ Hyperboloid: poincare_to_hyperboloid / hyperboloid_to_poincare
(curvature-aware stereographic projection through [-1/√c, 0, ..., 0]).
- Poincaré ↔ PV: poincare_to_pv / pv_to_poincare
(the gyro-isomorphism of PVNN Eq. 4, an isometry by their Thm 4.2).
- Hyperboloid ↔ PV: hyperboloid_to_pv / pv_to_hyperboloid
(direct map — PV coordinates are the space-like part of the 4-velocity).
JIT Compilation & Batching
All functions work with single points and return single points.
Use jax.vmap for batch operations:
>>> import jax
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds import isometry_mappings
>>>
>>> # Single point conversion
>>> x_hyp = jnp.array([1.0, 0.1, 0.2]) # Hyperboloid point
>>> y_poinc = isometry_mappings.hyperboloid_to_poincare(x_hyp, c=1.0)
>>>
>>> # Batch conversion with vmap
>>> x_batch = jnp.array([[1.0, 0.1, 0.2], [1.1, 0.15, 0.25]])
>>> convert_batch = jax.vmap(isometry_mappings.hyperboloid_to_poincare, in_axes=(0, None))
>>> y_batch = convert_batch(x_batch, 1.0)
References:
Wikipedia: Hyperboloid model
https://en.wikipedia.org/wiki/Hyperboloid_model#Relation_to_other_models
Chen et al. "Proper Velocity Neural Networks." ICLR 2026 (PV ↔ Poincaré, Eq. 4).
hyperboloid_to_poincare
hyperboloid_to_poincare(
x: Float[Array, dim_plus_1], c: Curvature
) -> Float[Array, dim]
Convert hyperboloid point to Poincaré ball via stereographic projection.
Projects the hyperboloid point onto the hyperplane t = 0 by intersecting
with a line through [-1/√c, 0, ..., 0]. This implements the canonical
isometry between the two models (radius-1/√c Poincaré ball).
Formula:
y_i = x_i / (√c·t + 1)
where x = [t, x_1, ..., x_n] on hyperboloid (t = x₀ ≥ 1/√c)
Args:
x: Point on hyperboloid, shape (dim+1,). Should satisfy ⟨x,x⟩_L = -1/c.
c: Curvature (positive)
Returns:
Point in Poincaré ball, shape (dim,). Satisfies ||y||² < 1/c.
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds import isometry_mappings
>>>
>>> # Convert hyperboloid origin to Poincaré origin
>>> x_origin = jnp.array([1.0, 0.0, 0.0]) # c=1.0 origin
>>> y = isometry_mappings.hyperboloid_to_poincare(x_origin, c=1.0)
>>> bool(jnp.allclose(y, jnp.zeros(2)))
True
References:
Wikipedia: Hyperboloid model - Relation to other models
Source code in hyperbolix/manifolds/isometry_mappings.py
| def hyperboloid_to_poincare(
x: Float[Array, "dim_plus_1"],
c: Curvature,
) -> Float[Array, "dim"]:
"""Convert hyperboloid point to Poincaré ball via stereographic projection.
Projects the hyperboloid point onto the hyperplane t = 0 by intersecting
with a line through [-1/√c, 0, ..., 0]. This implements the canonical
isometry between the two models (radius-1/√c Poincaré ball).
Formula:
y_i = x_i / (√c·t + 1)
where x = [t, x_1, ..., x_n] on hyperboloid (t = x₀ ≥ 1/√c)
Args:
x: Point on hyperboloid, shape (dim+1,). Should satisfy ⟨x,x⟩_L = -1/c.
c: Curvature (positive)
Returns:
Point in Poincaré ball, shape (dim,). Satisfies ||y||² < 1/c.
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds import isometry_mappings
>>>
>>> # Convert hyperboloid origin to Poincaré origin
>>> x_origin = jnp.array([1.0, 0.0, 0.0]) # c=1.0 origin
>>> y = isometry_mappings.hyperboloid_to_poincare(x_origin, c=1.0)
>>> bool(jnp.allclose(y, jnp.zeros(2)))
True
References:
Wikipedia: Hyperboloid model - Relation to other models
"""
sqrt_c = jnp.sqrt(c)
t = x[0] # Temporal component
x_spatial = x[1:] # Spatial components (x_1, ..., x_n)
# Curvature-aware stereographic projection: y_i = x_i / (√c·t + 1).
# Since t ≥ 1/√c on the hyperboloid, √c·t ≥ 1 and the denominator ≥ 2 — stable.
denominator = jnp.maximum(sqrt_c * t + 1.0, MIN_DENOM)
return x_spatial / denominator
|
poincare_to_hyperboloid
poincare_to_hyperboloid(
y: Float[Array, dim], c: Curvature
) -> Float[Array, dim_plus_1]
Convert Poincaré ball point to hyperboloid via inverse stereographic projection.
Inverts the stereographic projection to map points from the Poincaré ball
back to the hyperboloid. This implements the canonical isometry between
the two models (radius-1/√c Poincaré ball).
Formula:
t = (1 + c·||y||²) / ((1 - c·||y||²)·√c)
x_i = 2·y_i / (1 - c·||y||²)
where y = [y_1, ..., y_n] in Poincaré ball (||y||² < 1/c)
Args:
y: Point in Poincaré ball, shape (dim,). Should satisfy ||y||² < 1/c.
c: Curvature (positive)
Returns:
Point on hyperboloid, shape (dim+1,). Satisfies ⟨x,x⟩_L = -1/c.
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds import isometry_mappings
>>>
>>> # Convert Poincaré origin to hyperboloid origin
>>> y_origin = jnp.array([0.0, 0.0])
>>> x = isometry_mappings.poincare_to_hyperboloid(y_origin, c=1.0)
>>> bool(jnp.allclose(x, jnp.array([1.0, 0.0, 0.0])))
True
References:
Wikipedia: Hyperboloid model - Relation to other models
Source code in hyperbolix/manifolds/isometry_mappings.py
| def poincare_to_hyperboloid(
y: Float[Array, "dim"],
c: Curvature,
) -> Float[Array, "dim_plus_1"]:
"""Convert Poincaré ball point to hyperboloid via inverse stereographic projection.
Inverts the stereographic projection to map points from the Poincaré ball
back to the hyperboloid. This implements the canonical isometry between
the two models (radius-1/√c Poincaré ball).
Formula:
t = (1 + c·||y||²) / ((1 - c·||y||²)·√c)
x_i = 2·y_i / (1 - c·||y||²)
where y = [y_1, ..., y_n] in Poincaré ball (||y||² < 1/c)
Args:
y: Point in Poincaré ball, shape (dim,). Should satisfy ||y||² < 1/c.
c: Curvature (positive)
Returns:
Point on hyperboloid, shape (dim+1,). Satisfies ⟨x,x⟩_L = -1/c.
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds import isometry_mappings
>>>
>>> # Convert Poincaré origin to hyperboloid origin
>>> y_origin = jnp.array([0.0, 0.0])
>>> x = isometry_mappings.poincare_to_hyperboloid(y_origin, c=1.0)
>>> bool(jnp.allclose(x, jnp.array([1.0, 0.0, 0.0])))
True
References:
Wikipedia: Hyperboloid model - Relation to other models
"""
y_sqnorm = jnp.dot(y, y)
sqrt_c = jnp.sqrt(c)
# Curvature-aware inverse stereographic projection. The spatial part scales
# by the Poincaré conformal factor 1/(1 - c·||y||²); only the time component
# carries the extra 1/√c, so the two denominators differ.
one_minus = jnp.maximum(1.0 - c * y_sqnorm, MIN_DENOM)
t = (1.0 + c * y_sqnorm) / (one_minus * sqrt_c)
x_spatial = 2.0 * y / one_minus
# Concatenate temporal and spatial components: [t, x_1, ..., x_n]
return jnp.concatenate([t[None], x_spatial])
|
pv_to_poincare
pv_to_poincare(
x: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Convert a Proper Velocity point to the Poincaré ball.
Implements the gyro-isomorphism π_{PV→P} of the PVNN paper (Chen et al.
2026, Eq. 4), proven to be a Riemannian isometry (their Thm 4.2). With
hyperbolix's c > 0 convention (K = -c):
Formula:
y = x / (1 + √(1 + c·||x||²))
This is the numerically stable form of y = β_x/(1 + β_x)·x with the PV
beta factor β_x = 1/√(1 + c·||x||²): dividing numerator and denominator
by β_x turns it into x / (1 + 1/β_x). The denominator is ≥ 2, so the map
never blows up — every finite PV point lands strictly inside the radius-1/√c
ball.
Args:
x: Point in PV space (unconstrained R^n), shape (dim,).
c: Curvature (positive).
Returns:
Point in the Poincaré ball, shape (dim,). Satisfies ||y||² < 1/c.
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds import isometry_mappings
>>>
>>> # PV origin (0) maps to Poincaré origin (0)
>>> y = isometry_mappings.pv_to_poincare(jnp.zeros(2), c=1.0)
>>> bool(jnp.allclose(y, jnp.zeros(2)))
True
References:
Chen et al. "Proper Velocity Neural Networks." ICLR 2026, Eq. 4.
Source code in hyperbolix/manifolds/isometry_mappings.py
| def pv_to_poincare(
x: Float[Array, "dim"],
c: Curvature,
) -> Float[Array, "dim"]:
"""Convert a Proper Velocity point to the Poincaré ball.
Implements the gyro-isomorphism π_{PV→P} of the PVNN paper (Chen et al.
2026, Eq. 4), proven to be a Riemannian isometry (their Thm 4.2). With
hyperbolix's c > 0 convention (K = -c):
Formula:
y = x / (1 + √(1 + c·||x||²))
This is the numerically stable form of ``y = β_x/(1 + β_x)·x`` with the PV
beta factor ``β_x = 1/√(1 + c·||x||²)``: dividing numerator and denominator
by β_x turns it into ``x / (1 + 1/β_x)``. The denominator is ≥ 2, so the map
never blows up — every finite PV point lands strictly inside the radius-1/√c
ball.
Args:
x: Point in PV space (unconstrained R^n), shape (dim,).
c: Curvature (positive).
Returns:
Point in the Poincaré ball, shape (dim,). Satisfies ||y||² < 1/c.
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds import isometry_mappings
>>>
>>> # PV origin (0) maps to Poincaré origin (0)
>>> y = isometry_mappings.pv_to_poincare(jnp.zeros(2), c=1.0)
>>> bool(jnp.allclose(y, jnp.zeros(2)))
True
References:
Chen et al. "Proper Velocity Neural Networks." ICLR 2026, Eq. 4.
"""
beta_inv = jnp.sqrt(1.0 + c * jnp.dot(x, x)) # 1/β_x = √(1 + c·||x||²)
return x / (1.0 + beta_inv)
|
poincare_to_pv
poincare_to_pv(
y: Float[Array, dim], c: Curvature
) -> Float[Array, dim]
Convert a Poincaré ball point to Proper Velocity space.
Inverse of :func:pv_to_poincare — the map π_{P→PV} of the PVNN paper
(Chen et al. 2026, Eq. 4). With hyperbolix's c > 0 convention (K = -c):
Formula:
x = 2·y / (1 - c·||y||²) = λ(y)·y
where λ(y) = 2/(1 - c·||y||²) is the Poincaré conformal factor. As y
approaches the ball boundary (||y||² → 1/c) the image grows without bound —
expected, since PV is the unconstrained R^n model. The denominator is
guarded by MIN_DENOM to avoid division by zero at the boundary.
Args:
y: Point in the Poincaré ball, shape (dim,). Should satisfy ||y||² < 1/c.
c: Curvature (positive).
Returns:
Point in PV space (unconstrained R^n), shape (dim,).
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds import isometry_mappings
>>>
>>> # Poincaré origin (0) maps to PV origin (0)
>>> x = isometry_mappings.poincare_to_pv(jnp.zeros(2), c=1.0)
>>> bool(jnp.allclose(x, jnp.zeros(2)))
True
References:
Chen et al. "Proper Velocity Neural Networks." ICLR 2026, Eq. 4.
Source code in hyperbolix/manifolds/isometry_mappings.py
| def poincare_to_pv(
y: Float[Array, "dim"],
c: Curvature,
) -> Float[Array, "dim"]:
"""Convert a Poincaré ball point to Proper Velocity space.
Inverse of :func:`pv_to_poincare` — the map π_{P→PV} of the PVNN paper
(Chen et al. 2026, Eq. 4). With hyperbolix's c > 0 convention (K = -c):
Formula:
x = 2·y / (1 - c·||y||²) = λ(y)·y
where λ(y) = 2/(1 - c·||y||²) is the Poincaré conformal factor. As y
approaches the ball boundary (||y||² → 1/c) the image grows without bound —
expected, since PV is the *unconstrained* R^n model. The denominator is
guarded by ``MIN_DENOM`` to avoid division by zero at the boundary.
Args:
y: Point in the Poincaré ball, shape (dim,). Should satisfy ||y||² < 1/c.
c: Curvature (positive).
Returns:
Point in PV space (unconstrained R^n), shape (dim,).
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds import isometry_mappings
>>>
>>> # Poincaré origin (0) maps to PV origin (0)
>>> x = isometry_mappings.poincare_to_pv(jnp.zeros(2), c=1.0)
>>> bool(jnp.allclose(x, jnp.zeros(2)))
True
References:
Chen et al. "Proper Velocity Neural Networks." ICLR 2026, Eq. 4.
"""
denominator = jnp.maximum(1.0 - c * jnp.dot(y, y), MIN_DENOM)
return 2.0 * y / denominator
|
pv_to_hyperboloid
pv_to_hyperboloid(
x: Float[Array, dim], c: Curvature
) -> Float[Array, dim_plus_1]
Convert a Proper Velocity point to the hyperboloid model.
Uses the direct relation "proper velocity is the spatial part of the
(dimensionless) 4-velocity": the PV coordinates are exactly the space-like
hyperboloid components, and the time component is reconstructed from the
Lorentz constraint ⟨z,z⟩_L = -1/c.
Formula:
z = [√(1/c + ||x||²), x_1, ..., x_n]
This is an exact isometry (it equals
poincare_to_hyperboloid(pv_to_poincare(x, c), c)) but avoids the
near-boundary 1/(1 - c·||y||²) blow-up of the composed route — it is just a
concatenation. The time component is ≥ 1/√c > 0, so the result is always a
valid, numerically stable hyperboloid point.
Args:
x: Point in PV space (unconstrained R^n), shape (dim,).
c: Curvature (positive).
Returns:
Point on the hyperboloid, shape (dim+1,). Satisfies ⟨z,z⟩_L = -1/c, z₀ > 0.
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds import isometry_mappings
>>>
>>> # PV origin maps to the hyperboloid origin [1/√c, 0, ...]
>>> z = isometry_mappings.pv_to_hyperboloid(jnp.zeros(2), c=1.0)
>>> bool(jnp.allclose(z, jnp.array([1.0, 0.0, 0.0])))
True
References:
Chen et al. "Proper Velocity Neural Networks." ICLR 2026.
Source code in hyperbolix/manifolds/isometry_mappings.py
| def pv_to_hyperboloid(
x: Float[Array, "dim"],
c: Curvature,
) -> Float[Array, "dim_plus_1"]:
"""Convert a Proper Velocity point to the hyperboloid model.
Uses the direct relation "proper velocity is the spatial part of the
(dimensionless) 4-velocity": the PV coordinates are exactly the space-like
hyperboloid components, and the time component is reconstructed from the
Lorentz constraint ⟨z,z⟩_L = -1/c.
Formula:
z = [√(1/c + ||x||²), x_1, ..., x_n]
This is an exact isometry (it equals
``poincare_to_hyperboloid(pv_to_poincare(x, c), c)``) but avoids the
near-boundary 1/(1 - c·||y||²) blow-up of the composed route — it is just a
concatenation. The time component is ≥ 1/√c > 0, so the result is always a
valid, numerically stable hyperboloid point.
Args:
x: Point in PV space (unconstrained R^n), shape (dim,).
c: Curvature (positive).
Returns:
Point on the hyperboloid, shape (dim+1,). Satisfies ⟨z,z⟩_L = -1/c, z₀ > 0.
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds import isometry_mappings
>>>
>>> # PV origin maps to the hyperboloid origin [1/√c, 0, ...]
>>> z = isometry_mappings.pv_to_hyperboloid(jnp.zeros(2), c=1.0)
>>> bool(jnp.allclose(z, jnp.array([1.0, 0.0, 0.0])))
True
References:
Chen et al. "Proper Velocity Neural Networks." ICLR 2026.
"""
time = jnp.sqrt(1.0 / c + jnp.dot(x, x)) # z₀ = √(1/c + ||x||²) ≥ 1/√c
return jnp.concatenate([time[None], x])
|
hyperboloid_to_pv
hyperboloid_to_pv(
x: Float[Array, dim_plus_1], c: Curvature
) -> Float[Array, dim]
Convert a hyperboloid point to Proper Velocity space.
Inverse of :func:pv_to_hyperboloid: the PV coordinates are exactly the
space-like components of the hyperboloid point, so the map simply drops the
time component.
Formula:
x = [z_1, ..., z_n] (drop the time component z₀)
The curvature c is unused (the relation is curvature-independent) but is
kept in the signature for API symmetry with the other mappings.
Args:
x: Point on the hyperboloid, shape (dim+1,). Should satisfy ⟨x,x⟩_L = -1/c.
c: Curvature (positive). Unused.
Returns:
Point in PV space (unconstrained R^n), shape (dim,).
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds import isometry_mappings
>>>
>>> # Hyperboloid origin [1/√c, 0, ...] maps to the PV origin (0)
>>> x = isometry_mappings.hyperboloid_to_pv(jnp.array([1.0, 0.0, 0.0]), c=1.0)
>>> bool(jnp.allclose(x, jnp.zeros(2)))
True
References:
Chen et al. "Proper Velocity Neural Networks." ICLR 2026.
Source code in hyperbolix/manifolds/isometry_mappings.py
| def hyperboloid_to_pv(
x: Float[Array, "dim_plus_1"],
c: Curvature,
) -> Float[Array, "dim"]:
"""Convert a hyperboloid point to Proper Velocity space.
Inverse of :func:`pv_to_hyperboloid`: the PV coordinates are exactly the
space-like components of the hyperboloid point, so the map simply drops the
time component.
Formula:
x = [z_1, ..., z_n] (drop the time component z₀)
The curvature ``c`` is unused (the relation is curvature-independent) but is
kept in the signature for API symmetry with the other mappings.
Args:
x: Point on the hyperboloid, shape (dim+1,). Should satisfy ⟨x,x⟩_L = -1/c.
c: Curvature (positive). Unused.
Returns:
Point in PV space (unconstrained R^n), shape (dim,).
Examples:
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds import isometry_mappings
>>>
>>> # Hyperboloid origin [1/√c, 0, ...] maps to the PV origin (0)
>>> x = isometry_mappings.hyperboloid_to_pv(jnp.array([1.0, 0.0, 0.0]), c=1.0)
>>> bool(jnp.allclose(x, jnp.zeros(2)))
True
References:
Chen et al. "Proper Velocity Neural Networks." ICLR 2026.
"""
del c # curvature-independent: PV coords are the hyperboloid spatial part
return x[1:]
|
Usage Examples
Basic Distance Computation
import jax.numpy as jnp
from hyperbolix.manifolds import Poincare
poincare = Poincare()
x = jnp.array([0.1, 0.2])
y = jnp.array([0.3, -0.1])
c = 1.0
# Compute distance (default: VERSION_MOBIUS_DIRECT)
distance = poincare.dist(x, y, c)
Float64 Precision
from hyperbolix.manifolds import Poincare
import jax.numpy as jnp
# High-precision manifold
poincare_f64 = Poincare(dtype=jnp.float64)
x = jnp.array([0.1, 0.2]) # float32 input
distance = poincare_f64.dist(x, y, c=1.0) # automatically cast to float64
print(distance.dtype) # float64
Batched Operations with vmap
import jax
from hyperbolix.manifolds import Hyperboloid
hyperboloid = Hyperboloid()
c = 1.0
# Batch of ambient points (d+1 dimensions)
x_batch = jax.random.normal(jax.random.PRNGKey(0), (100, 4))
y_batch = jax.random.normal(jax.random.PRNGKey(1), (100, 4))
# Project to hyperboloid
x_proj = jax.vmap(hyperboloid.proj, in_axes=(0, None))(x_batch, c)
y_proj = jax.vmap(hyperboloid.proj, in_axes=(0, None))(y_batch, c)
# Compute distances
distances = jax.vmap(hyperboloid.dist, in_axes=(0, 0, None))(x_proj, y_proj, c)
Exponential and Logarithmic Maps
from hyperbolix.manifolds import Poincare
import jax.numpy as jnp
poincare = Poincare()
# Point on manifold
x = poincare.proj(jnp.array([0.2, 0.3]), c=1.0)
# Tangent vector
v = jnp.array([0.1, -0.05])
# Exponential map (move along geodesic)
y = poincare.expmap(v, x, c=1.0)
# Logarithmic map (inverse operation)
v_recovered = poincare.logmap(y, x, c=1.0)
Proper Velocity Operations
import jax
import jax.numpy as jnp
from hyperbolix.manifolds import ProperVelocity
pv = ProperVelocity()
c = 1.0
# Points live in unconstrained R^n — no projection step needed
x = jnp.array([0.3, 0.5])
y = jnp.array([-0.2, 0.4])
# Geodesic distance (asinh form — stable over all of R^n)
d = pv.dist(x, y, c)
# Exp/log maps at the origin
v = jnp.array([0.1, -0.2])
y_moved = pv.expmap_0(v, c)
v_recovered = pv.logmap_0(y_moved, c)
# Euclidean gradient -> Riemannian gradient under the PV metric
grad_euc = jnp.array([1.0, 0.0])
grad_riem = pv.egrad2rgrad(grad_euc, x, c)
# Retraction is exact Euclidean addition (PV is unconstrained)
x_next = pv.retraction(v, x, c)
assert jnp.allclose(x_next, x + v)
Product Manifolds (Mixed Curvature)
import jax
import jax.numpy as jnp
from hyperbolix.manifolds import (
ProductManifold, Hyperboloid, Poincare, Euclidean,
)
# Build P = H^5(c=1.0) x P^3(c=0.1) x E^4
product = ProductManifold(
(Hyperboloid(c=1.0), 5), # 5 = ambient dim (d+1) for hyperboloid
(Poincare(c=0.1), 3), # 3 = spatial dim for poincaré
(Euclidean(), 4), # 4 = standard euclidean dim
)
# Per-factor curvatures must be passed at call time as a sequence.
c = product.curvatures # (1.0, 0.1, 0.0) — static default
o = product.origin(c) # shape (12,)
# Pythagorean product distance: d_P = sqrt(sum d_i^2)
x = product.origin(c)
y = product.origin(c) # generated elsewhere; here we use o for illustration
d_l2 = product.dist(x, y, c) # scalar
d_per_factor = product.component_dist(x, y, c) # shape (3,) per-factor distances
# Batch with vmap: broadcast c with None across the batch.
dist_batch = jax.vmap(product.dist, in_axes=(0, 0, None))
# distances = dist_batch(xs, ys, c)
# Repeated-factor construction via from_signature
mixed = ProductManifold.from_signature(
(Hyperboloid, 5, 4, 1.0), # 4 copies of H^4(c=1.0)
(Poincare, 3, 2, 0.1), # 2 copies of P^3(c=0.1)
(Euclidean, 4, 1), # 1 copy of E^4
)
# For learnable curvature, build the sequence from LearnableCurvature calls on
# your nnx.Module:
# c = (self.curv_h(), self.curv_p(), 0.0)
# d = product.dist(x, y, c)
Isometry Mappings
from hyperbolix.manifolds import isometry_mappings
import jax.numpy as jnp
# Hyperboloid point (ambient coordinates, d+1 dims)
x_hyperboloid = jnp.array([1.5, 0.5, 0.3]) # Must satisfy Lorentz constraint
# Map to Poincaré ball (intrinsic coordinates, d dims)
x_poincare = isometry_mappings.hyperboloid_to_poincare(x_hyperboloid, c=1.0)
# Map back (round-trip)
x_hyperboloid_recovered = isometry_mappings.poincare_to_hyperboloid(x_poincare, c=1.0)
Numerical Considerations
Float32 Precision
Float32 can cause numerical issues, especially in the Poincaré ball near the boundary. Use Poincare(dtype=jnp.float64) for:
- High curvature values (
c > 1.0)
- Points near manifold boundaries
- Deep neural networks with many layers
See the Numerical Stability guide for details.