Metadata-Version: 2.4
Name: liudist
Version: 0.3.0
Summary: Fast exact Fisher-Rao (Rao) distance between Liouville distributions
Author-email: Georgy Meshcheryakov <iam@georgy.top>
License: MIT
Keywords: Fisher-Rao,Rao distance,Liouville distribution,Dirichlet,information geometry,Riemannian geometry,JAX,GPU
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax>=0.4
Requires-Dist: numpy
Provides-Extra: reference
Requires-Dist: scipy; extra == "reference"
Provides-Extra: test
Requires-Dist: scipy; extra == "test"
Requires-Dist: pytest; extra == "test"
Dynamic: license-file

# liudist

**Fast, exact Fisher–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

```bash
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](#precision-float32)).

---

## Quick start

```python
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)` (drops the coupling) |
| `lorentzian_embedding(X, r, betas=…)` | **coupling-aware** embedding → `(m+1, N)` + signature; keeps the total/`n` term the great-arc drops |
| `total_arclength(X, r)` | the total/depth coordinate `σ` (1-D Fisher arc length of `A=Σ(x+r)`) → `(N,)` |
| `pairwise_distance_matrix(X, r, method=…, betas=…)` | **full N×N** Fisher–Rao matrix via one signed Gram (GEMM); coupling retained |
| `refine_neighbors(X, r, k=…)` | exact-refine each law's `k` nearest neighbours (`distance_x_batched`) → exact k-NN graph |
| `to_euclidean(D)` | classical MDS: a distance matrix → true Euclidean coordinates |
| `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:

```python
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** (great-arc)
  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**. That error
  vanishes for *shape* differences as `m` grows, **but not for depth** (total-changing)
  differences — there the coupling is an O(1) fraction of the metric at every `m`, so the
  great-arc loses the total/`n` (depth / library size / cell size) signal. Use
  `lorentzian_embedding` when the total can differ across laws.
- **`lorentzian_embedding(X, r, betas=…)`** → `(m+1, N)` (+ signature): the **coupling-aware**
  embedding. Appends the single **total coordinate** `σ = ∫₀^A √ψ'(s) ds`, `A = Σ(x+r)` (a
  *timelike* axis, signature `(+…+, −)`) so that `d_x² ≈ ‖Δφ‖² − Δσ²` — an isometric
  embedding into Minkowski `ℝ^{m,1}`. This is a **provably tighter upper bound** than the
  great-arc (`d_x² ≤ ‖Δφ‖² − Δσ² ≤ ‖Δφ‖²`), tighter by the coupling term `Δσ²`;
  it restores the depth signal the great-arc drops. Measured vs. an independent
  geodesic-shooting ground truth (`Δ ≈ 0.15` regime): great-arc error → Lorentzian error is
  ≈1 % → 0.07 % for a shape displacement and, for a depth (total-changing) displacement,
  **tens of percent → ~0.01 %** (e.g. 52 % → 0.014 %). See [NOTES.md §5](NOTES.md).

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

### Pairwise distance matrices at scale (coupling retained)
Because the Lorentzian estimate `d_x² = ‖Δφ‖² − Δσ²` is a signed inner product, the full
`N × N` distance matrix is a **single signed Gram / GEMM** on the `(m+1)`-dim vectors — no
coupling dropped:

```python
D = liudist.pairwise_distance_matrix(X, r)                 # (N, N), Lorentzian (default)
D = liudist.pairwise_distance_matrix(X, r, betas=betas)    # total Liouville distance
graph = liudist.refine_neighbors(X, r, k=32)               # exact-refine each law's k-NN
```

Build is `O(N m Q)`; the matrix is one `O(N² m)` multiply. For `N ~ 5000`, `m ~ 1e4` this
is **seconds** (coupling retained), versus the ~`N²/2 ≈ 1.3e7` geodesic solves — days — an
exact all-pairs matrix would need. For accuracy-critical pairs, `refine_neighbors` replaces
each law's `k` nearest approximate neighbours with the exact `distance_x` (`N·k` solves, not
`N²/2`); since the Lorentzian is an upper bound, over-retrieving never hides a truly-close
pair. `to_euclidean(D)` runs classical MDS if a downstream method needs point coordinates
rather than a distance matrix.

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

Full derivations — the metric and its Sherman–Morrison inverse, the Christoffel collapse
and geodesic ODE, the curvature-corrected mean, and the Gamma–Dirichlet decomposition
behind the coupling-aware embedding — are in **[NOTES.md](NOTES.md)**.

---

## Validation

```bash
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`.
- **Coupling-aware embedding is an upper bound.** `lorentzian_embedding` /
  `pairwise_distance_matrix` give `d_x² ≈ ‖Δφ‖² − Δσ²`, a provably tight **upper** bound
  (`≥ 0` always); the small residual is the σ-graph curvature (shrinks with `m`). It keeps
  the coupling the great-arc drops — but for machine-precision distances on specific pairs,
  refine with `distance_x` / `refine_neighbors`. The matrix is *pseudo*-Euclidean (one
  timelike axis): ideal for a distance matrix or any kernel/Gram method; use `to_euclidean`
  if a method needs literal Euclidean points.
