liudist 0.1.0

Fast exact Fisher-Rao (Rao) distance between Liouville distributions

liudist

Fast, exact Fisher–Rao (Rao) distance between Liouville distributions.

Two Liouville laws that share the shape vector r but differ in the component vector x and the scalar beta sit on a Riemannian product manifold, so the squared Fisher–Rao distance splits cleanly:

d² = d_x²  +  d_beta²

• d_beta is closed form: sqrt(R)·|log(beta_y/beta_x)|, R = Σ r.
• d_x is the Fisher–Rao distance on the Dirichlet manifold Dir(x+r), whose metric is g(x) = diag(ψ'(x+r)) − ψ'(Σ(x+r))·11ᵀ (a diagonal + rank-1 coupling). There is no closed form; liudist solves it exactly with a GPU-native, preconditioned discrete-geodesic optimiser.

At m = 3×10⁴ components the exact d_x takes ~0.16 s on an RTX A5000 — about 250× faster than a straightforward scipy reference — and the method scales to m ~ 1e5. The rank-1 coupling (the geometrically important term) is retained throughout; nothing is dropped.

---

Install

pip install ./liudist            # from the directory containing pyproject.toml
# or, no install — just put this folder on PYTHONPATH:
#   cd liudist && python -c "import liudist"

Requirements: jax (with a CUDA build for GPU; CPU works too) and numpy. scipy is optional — only liudist.reference (the independent ground-truth solver) and the tests need it.

float64 is enabled on import (the method needs it for full accuracy). JAX uses the GPU automatically if it sees one. For a ~5× float32 run, set import jax; jax.config.update("jax_enable_x64", False) before importing liudist (see Precision).

---

Quick start

import numpy as np, liudist

m = 30000
rng = np.random.default_rng(0)
r = rng.uniform(0.01, 2.0, m)                                   # shape vector  (> 0)
x = rng.exponential(2.0, m)                                     # law 1 components (>= 0)
y = np.clip(x + (2/np.sqrt(m)) * rng.standard_normal(m), 0, None)   # law 2

# total Fisher-Rao distance (x-part + closed-form beta-part)
d = liudist.distance(x, y, r, beta_x=1.0, beta_y=1.2)

# exact x-part alone, with solver diagnostics
dx, info = liudist.distance_x(x, y, r, return_info=True)
print(dx, info["nit"])          # e.g. (17, 15, 0)  iterations per stage

# cheap O(m) approximation, no geodesic solve
da = liudist.approx_distance_x(x, y, r, kind="midpoint")

Run the full demo: python example.py.

---

API

function	description
distance(x, y, r, beta_x=1, beta_y=1, **kw)	total exact distance √(d_x²+d_beta²); forwards **kw to distance_x
distance_x(x, y, r, **kw)	exact x-part (the GPU solver); returns float, or (float, info) with return_info=True
distance_beta(beta_x, beta_y, r)	closed-form beta factor
approx_distance_x(x, y, r, kind=...)	fast O(m) approximation; kind ∈ {midpoint, endpoint, avg, product}
distance_x_batched(X, y, r, **kw)	experimental — B laws (rows of X) vs one reference y; returns d (B,)
logmap(mu, X, r, kind=...)	Riemannian log map(s) at mu: v_k = log_mu(p_k); kind ∈ {approx, picard, newton}
exp_map(mu, V, r)	geodesic exp map (inverse of logmap); V is (m,) or (N, m)
frechet_mean(X, r, method=...)	geodesic center of mass of the x-part; method ∈ {curvature, karcher}
barycenter(X, r, betas=...)	full center of mass → x* (and beta* = geometric mean if betas given)
curvature_corrected_mean, beta_barycenter	the closed-form pieces of the center
tangent_embedding(X, r, ...)	Euclidean chart at a common reference → (m, N) (or (m+1, N) with betas)
arclength_embedding(X, r)	reference-free product embedding → (m, N) (best at scale)
whiten(mu, r, V)	metric whitening g(mu)^{1/2} v (so ‖w‖₂ = ‖v‖_g)
metric_components(x, r), metric_quad(x, r, v), apply_Ginv(x, r, v)	the metric a=ψ'(x+r), c=ψ'(Σ), ‖v‖²_g, and g⁻¹v
liudist.reference.distance_x_reference(x, y, r)	independent NumPy/SciPy ODE ground truth (small m)

Inputs are 1-D arrays: x, y ≥ 0, r > 0, all length m; beta > 0. X is (N, m).

Key distance_x options (all keep the converged distance unchanged; they trade speed/accuracy/robustness)

option	default	meaning
gtol	1e-6	gradient tolerance. The distance reaches machine precision long before the gradient does, so 1e-6 gives full accuracy in the realistic regime while saving iterations. Tighten for very large separations.
precond	'full'	'full' whitens the whole metric (removes the rank-1 coupling mode); 'diag' whitens only the diagonal.
path_precond	True	seed L-BFGS with the inverse path-Laplacian (kills the ~K² conditioning; ~5× fewer iterations).
coarse_resolve	True	re-solve the coarse grid from the (robust) fine geodesic — prevents a rare silent ~1% error at very large separations.
fast_special	True	branchless trigamma instead of polygamma (~3× faster, distance rel err ~5e-10). Set False for exact polygamma (~1e-14) at ~3× the cost.
K	16	base discretisation; Richardson uses K and 2K.

---

Performance & accuracy

Exact d_x, fixed_total regime (per-coordinate step ~1/√m), RTX A5000, warm:

m	liudist.distance_x	speedup vs scipy¹	rel. error²
1 000	0.04 s	~99×	5e-10
10 000	0.08 s	~178×	4e-10
30 000	0.16 s	~246×	4e-10

¹ vs a scipy L-BFGS-B discrete-geodesic solver (the natural CPU baseline; ~39 s at m=3e4). ² vs the exact-polygamma solver / NumPy ODE ground truth. With fast_special=False the error is ~2e-14 at ~3× the time.

Repeated many-laws-vs-one-reference at m=3e4 reuses the compile: ~0.08 s/law.

Cheap approximations (approx_distance_x) are O(m) and ~10³–10⁵× faster than the exact solve; 'midpoint' is the recommended one (rel err ~|Δ|²). 'product' drops the coupling and is only accurate when the coupling is negligible (large m, small per-coord step).

Precision: float32

The metric's coupling term is not a catastrophic cancellation, so float32 is safe and ~5× faster end-to-end (the trigamma is the bottleneck and fp64 is ~1/64 rate on the A5000), at ~1e-6 accuracy. To use it, disable x64 before importing:

import jax; jax.config.update("jax_enable_x64", False)
import liudist           # whole pipeline now float32

fast_special (the trigamma approximation) and float32 are not additive — both remove the same fp64 cost — so in float32 they coincide. Use fast_special for float64-class accuracy at 3×; drop to float32 only if ~1e-6 suffices.

---

Tangent space: log maps, center of mass, embeddings

Because the manifold is a Riemannian product M_x × M_β, everything below factorises: the β-part is closed form and the work is on the Dirichlet x-part. These tools put N laws into one Euclidean chart at m ~ 1e4, N ~ 1e3 in seconds — the building block for clustering / PCA / nearest-neighbour search under the Fisher–Rao metric.

Riemannian log map — logmap(mu, X, r, kind=...)

log_mu(p) is the initial geodesic velocity reaching p at t=1, with ‖log_mu(p)‖_g(mu) = d_x(mu, p). X is one law (m,) or a stack (N, m); three back-ends trade cost vs accuracy (cheap-scalable → exact-small):

kind	cost	accuracy	use
'approx' (default)	O(Nm), no integration	2nd-order (few %; coupling/step-limited)	the scalable workhorse at N~1e3, m~1e4
'picard'	O(N·m·n_steps), batched	near-exact for small/moderate dispersion	a better-than-2nd-order batch map
'newton' / 'exact'	per-point Newton–GMRES	machine accurate (~5e-10)	validation / small–medium m

exp_map(mu, V, r) is the inverse (RK4 geodesic integrator). Round-trip exp_mu(log_mu(p)) = p to ~1e-14 (newton).

Geodesic center of mass — frechet_mean(X, r) / barycenter

argmin_q ½ Σ_k w_k d_x²(q, x_k). The default method='curvature' is a closed-form, O(Nm) estimate (xbar + ¼ g(xbar)⁻¹(b⊙M₂ − e·s₂·1), one Sherman–Morrison solve); its excess objective is ~disp⁴ and its Riemannian gradient scales ~disp² (≈4e-4 at disp 0.1). method='karcher' polishes it with batched Picard log maps (a few steps; degrades gracefully — info['diverged'] — if dispersion is too large for Picard). barycenter(..., betas=...) adds the closed-form β* (weighted geometric mean).

This is the recommended center at scale — exact Karcher (N exact log maps/iteration) is neither feasible nor needed there.

Euclidean embeddings

Map each law to a vector so ordinary L2 geometry approximates Fisher–Rao:

• tangent_embedding(X, r, reference='center', logmap_kind=...) → (m, N): log-map every law at a common reference (default the Fréchet-mean center, which minimises distortion), then whiten by the metric so ‖wⱼ−wₖ‖₂ ≈ d_x(pⱼ,pₖ). With betas, appends one extra (already-whitened) √R(log βₖ − log β_u) row → (m+1, N).
• arclength_embedding(X, r) → (m, N): reference-free product embedding (each coordinate → its exact 1-D Fisher arc length). Pairwise L2 equals the product-metric distance exactly; the only error is the rank-1 coupling (small, → 0 as m grows). The most accurate scalable embedding (4–13× better than the approx tangent map, stable near the boundary).

whiten(mu, r, V) exposes the O(m) metric square root alone (‖whiten(v)‖₂ = ‖v‖_g).

Timing (RTX A5000, warm, m = 30 000, N = 1000)

op	time
frechet_mean (curvature)	~1.5 s
logmap (approx)	~0.8 s
tangent_embedding	~0.9 s
arclength_embedding	~0.2 s

All are closed-form / O(Nm) (or O(NmQ)); anything that integrates geodesics per point (newton log maps, exact Karcher) is for validation / small–medium m, not this scale.

How it works

d_x² is the minimum of a discrete-geodesic energy E = K·Σ‖ΔP_k‖²_g over interior path nodes, with Richardson(K, 2K) extrapolation to remove the discretisation bias. The whole L-BFGS minimisation runs in a single jax.lax.while_loop on the device (no scipy, no host transfers). Three things make it converge in ~25 iterations instead of ~250: full-metric whitening (Hessian ≈ I in the coordinate directions, coupling mode removed), a path-Laplacian L-BFGS seed (H₀ = (L⁻¹⊗I)/2K, the Brownian-bridge Green's function — kills the K² path conditioning), and multigrid warm-start of the fine solve from the coarse one. The dominant cost is the trigamma ψ', replaced by a branchless fixed-shift approximation (fast_special) that is ~12× faster than polygamma at ~1e-9 accuracy and whose derivative is the matching tetragamma — so autodiff gets the fast backward pass for free.

---

Validation

pip install ./liudist[test]
python -m pytest -q                 # or:  python tests/test_liudist.py

The suite checks distance_x against an independent NumPy/SciPy shooting solver at small m, that fast_special matches exact polygamma to ~1e-9, the product-manifold identity d² = d_x² + d_beta², and that the batched solver matches the per-lane one. The tangent tools are checked too: the exact (newton) log map reproduces the reference distance and round-trips through exp_map; whitening is an isometry (‖whiten(v)‖₂ = ‖v‖_g); the arc-length embedding's pairwise L2 equals the product distance; the curvature mean has a small Riemannian gradient that Karcher drives to ~0; and beta_barycenter is the geometric mean.

---

Notes & limitations

• JAX version. The single-lane solver reuses JAX's Wolfe line search from jax._src.scipy.optimize.line_search (private API; tested on jax 0.4–0.10). If a future JAX moves it, the import raises a clear error.
• distance_x_batched is experimental. It is numerically correct but the batch runs to the slowest lane's iteration count, and very large separations may not reach machine precision. For accuracy-critical large separations, loop distance_x instead (it reuses the compile).
• Domain. Requires x ≥ 0, r > 0 so that α = x+r > 0; near the simplex boundary (α → 0) the metric is stiff but the solver and the fast trigamma remain accurate (the small-α part of the trigamma is computed exactly).
• Shared r. The product split and this solver assume both laws share the shape vector r.