#![forbid(unsafe_code)]
#![warn(missing_docs)]
#[cfg(feature = "manifold")]
mod manifold;
#[cfg(feature = "manifold")]
pub use manifold::FisherRaoSimplex;
mod projection;
pub use projection::{fisher_orthogonal_project, fisher_orthogonal_project_full};
use thiserror::Error;
pub use logp::hellinger;
#[derive(Debug, Error)]
pub enum Error {
#[error(transparent)]
Logp(#[from] logp::Error),
#[error("interpolation parameter t={t} outside [0, 1]")]
InvalidT {
t: f64,
},
#[error("length mismatch: {a_name} has {a_len} elements, {b_name} has {b_len}")]
LengthMismatch {
a_name: &'static str,
a_len: usize,
b_name: &'static str,
b_len: usize,
},
#[error(
"zero entry at index {idx}: log-space operations require strictly positive distributions"
)]
ZeroEntry {
idx: usize,
},
}
pub type Result<T> = core::result::Result<T, Error>;
fn validate_t(t: f64) -> Result<()> {
if !(0.0..=1.0).contains(&t) {
return Err(Error::InvalidT { t });
}
Ok(())
}
fn validate_same_len(
a: &[f64],
b: &[f64],
a_name: &'static str,
b_name: &'static str,
) -> Result<()> {
if a.len() != b.len() {
return Err(Error::LengthMismatch {
a_name,
a_len: a.len(),
b_name,
b_len: b.len(),
});
}
Ok(())
}
fn validate_strictly_positive(p: &[f64]) -> Result<()> {
for (i, &v) in p.iter().enumerate() {
if v <= 0.0 {
return Err(Error::ZeroEntry { idx: i });
}
}
Ok(())
}
pub fn rao_distance_categorical(p: &[f64], q: &[f64], tol: f64) -> Result<f64> {
let mut bc = logp::bhattacharyya_coeff(p, q, tol)?;
bc = bc.clamp(0.0, 1.0);
if (1.0 - bc).abs() <= 10.0 * tol {
bc = 1.0;
}
Ok(2.0 * bc.acos())
}
pub fn fisher_rao_geodesic(p: &[f64], q: &[f64], t: f64, tol: f64) -> Result<Vec<f64>> {
validate_t(t)?;
logp::validate_simplex(p, tol)?;
logp::validate_simplex(q, tol)?;
validate_same_len(p, q, "p", "q")?;
let bc = logp::bhattacharyya_coeff(p, q, tol)?.clamp(0.0, 1.0);
let omega = bc.acos();
if omega < tol {
return Ok(p.to_vec());
}
let sin_omega = omega.sin();
let s0 = ((1.0 - t) * omega).sin();
let s1 = (t * omega).sin();
let result: Vec<f64> = p
.iter()
.zip(q.iter())
.map(|(&pi, &qi)| {
let theta = (s0 * pi.sqrt() + s1 * qi.sqrt()) / sin_omega;
theta * theta
})
.collect();
Ok(result)
}
pub fn m_geodesic(p: &[f64], q: &[f64], t: f64, tol: f64) -> Result<Vec<f64>> {
validate_t(t)?;
logp::validate_simplex(p, tol)?;
logp::validate_simplex(q, tol)?;
validate_same_len(p, q, "p", "q")?;
let result: Vec<f64> = p
.iter()
.zip(q.iter())
.map(|(&pi, &qi)| (1.0 - t) * pi + t * qi)
.collect();
Ok(result)
}
pub fn e_geodesic(p: &[f64], q: &[f64], t: f64, tol: f64) -> Result<Vec<f64>> {
validate_t(t)?;
logp::validate_simplex(p, tol)?;
logp::validate_simplex(q, tol)?;
validate_same_len(p, q, "p", "q")?;
validate_strictly_positive(p)?;
validate_strictly_positive(q)?;
let log_unnorm: Vec<f64> = p
.iter()
.zip(q.iter())
.map(|(&pi, &qi)| (1.0 - t) * pi.ln() + t * qi.ln())
.collect();
let lse = logp::log_sum_exp(&log_unnorm);
let result: Vec<f64> = log_unnorm.iter().map(|&v| (v - lse).exp()).collect();
Ok(result)
}
pub fn alpha_geodesic(p: &[f64], q: &[f64], t: f64, alpha: f64, tol: f64) -> Result<Vec<f64>> {
validate_t(t)?;
logp::validate_simplex(p, tol)?;
logp::validate_simplex(q, tol)?;
validate_same_len(p, q, "p", "q")?;
if (alpha - 1.0).abs() < tol {
return e_geodesic(p, q, t, tol);
}
let m = (1.0 - alpha) / 2.0; if m < 0.0 {
validate_strictly_positive(p)?;
validate_strictly_positive(q)?;
}
let unnorm: Vec<f64> = p
.iter()
.zip(q.iter())
.map(|(&pi, &qi)| {
let mixed = (1.0 - t) * pi.powf(m) + t * qi.powf(m);
mixed.powf(1.0 / m)
})
.collect();
let sum: f64 = unnorm.iter().sum();
Ok(unnorm.iter().map(|&v| v / sum).collect())
}
pub fn fisher_information_diagonal(p: &[f64], tol: f64) -> Result<Vec<f64>> {
logp::validate_simplex(p, tol)?;
validate_strictly_positive(p)?;
Ok(p.iter().map(|&pi| 1.0 / pi).collect())
}
pub fn natural_gradient(p: &[f64], euclidean_grad: &[f64]) -> Result<Vec<f64>> {
validate_same_len(p, euclidean_grad, "p", "euclidean_grad")?;
Ok(p.iter()
.zip(euclidean_grad.iter())
.map(|(&pi, &gi)| pi * gi)
.collect())
}
#[cfg(test)]
mod tests {
use super::*;
use proptest::prelude::*;
fn simplex_vec(len: usize) -> impl Strategy<Value = Vec<f64>> {
prop::collection::vec(0.0f64..10.0, len).prop_map(|mut v| {
let s: f64 = v.iter().sum();
if s == 0.0 {
v[0] = 1.0;
return v;
}
for x in v.iter_mut() {
*x /= s;
}
v
})
}
fn positive_simplex_vec(len: usize) -> impl Strategy<Value = Vec<f64>> {
prop::collection::vec(0.01f64..10.0, len).prop_map(|mut v| {
let s: f64 = v.iter().sum();
for x in v.iter_mut() {
*x /= s;
}
v
})
}
const TOL: f64 = 1e-9;
proptest! {
#[test]
fn rao_is_symmetric(p in simplex_vec(8), q in simplex_vec(8)) {
let d1 = rao_distance_categorical(&p, &q, TOL).unwrap();
let d2 = rao_distance_categorical(&q, &p, TOL).unwrap();
prop_assert!((d1 - d2).abs() < 1e-12);
}
#[test]
fn rao_self_is_zero(p in simplex_vec(12)) {
let d = rao_distance_categorical(&p, &p, TOL).unwrap();
prop_assert!(d.abs() < 1e-12);
}
#[test]
fn rao_is_bounded(p in simplex_vec(10), q in simplex_vec(10)) {
let d = rao_distance_categorical(&p, &q, TOL).unwrap();
prop_assert!(d >= -1e-12);
prop_assert!(d <= core::f64::consts::PI + 1e-12);
}
#[test]
fn hellinger_is_symmetric(p in simplex_vec(8), q in simplex_vec(8)) {
let h1 = hellinger(&p, &q, TOL).unwrap();
let h2 = hellinger(&q, &p, TOL).unwrap();
prop_assert!((h1 - h2).abs() < 1e-12);
}
#[test]
fn hellinger_is_bounded(p in simplex_vec(8), q in simplex_vec(8)) {
let h = hellinger(&p, &q, TOL).unwrap();
prop_assert!(h >= -1e-12);
prop_assert!(h <= 1.0 + 1e-12);
}
#[test]
fn rao_matches_bhattacharyya_formula(p in simplex_vec(10), q in simplex_vec(10)) {
let rao = rao_distance_categorical(&p, &q, TOL).unwrap();
let mut bc = logp::bhattacharyya_coeff(&p, &q, TOL).unwrap();
bc = bc.clamp(0.0, 1.0);
if (1.0 - bc).abs() <= 10.0 * TOL {
bc = 1.0;
}
let expected = 2.0 * bc.acos();
prop_assert!((rao - expected).abs() < 1e-12, "rao={rao} expected={expected} bc={bc}");
}
}
proptest! {
#[test]
fn rao_triangle_inequality(
p in simplex_vec(5),
q in simplex_vec(5),
r in simplex_vec(5),
) {
let dpq = rao_distance_categorical(&p, &q, TOL).unwrap();
let dqr = rao_distance_categorical(&q, &r, TOL).unwrap();
let dpr = rao_distance_categorical(&p, &r, TOL).unwrap();
prop_assert!(dpr <= dpq + dqr + 1e-10,
"triangle inequality violated: d(p,r)={dpr} > d(p,q)+d(q,r)={}", dpq + dqr);
}
}
#[test]
fn rao_disjoint_equals_pi() {
let p = vec![1.0, 0.0, 0.0, 0.0];
let q = vec![0.0, 1.0, 0.0, 0.0];
let d = rao_distance_categorical(&p, &q, 1e-12).unwrap();
assert!(
(d - core::f64::consts::PI).abs() < 1e-10,
"expected pi, got {d}"
);
}
proptest! {
#[test]
fn hellinger_self_distance_zero(p in simplex_vec(8)) {
let h = hellinger(&p, &p, TOL).unwrap();
prop_assert!(h.abs() < 1e-7, "hellinger(p,p) = {h}, expected 0");
}
}
proptest! {
#[test]
fn fisher_rao_geodesic_endpoints(
p in positive_simplex_vec(6),
q in positive_simplex_vec(6),
) {
let g0 = fisher_rao_geodesic(&p, &q, 0.0, TOL).unwrap();
let g1 = fisher_rao_geodesic(&p, &q, 1.0, TOL).unwrap();
for i in 0..p.len() {
prop_assert!((g0[i] - p[i]).abs() < 1e-10,
"endpoint t=0: g0[{i}]={} != p[{i}]={}", g0[i], p[i]);
prop_assert!((g1[i] - q[i]).abs() < 1e-10,
"endpoint t=1: g1[{i}]={} != q[{i}]={}", g1[i], q[i]);
}
}
}
proptest! {
#[test]
fn fisher_rao_geodesic_midpoint_minimizes_sum_distances(
p in positive_simplex_vec(5),
q in positive_simplex_vec(5),
) {
let mid = fisher_rao_geodesic(&p, &q, 0.5, TOL).unwrap();
let d_full = rao_distance_categorical(&p, &q, TOL).unwrap();
let d_pm = rao_distance_categorical(&p, &mid, TOL).unwrap();
let d_mq = rao_distance_categorical(&mid, &q, TOL).unwrap();
prop_assert!((d_pm + d_mq - d_full).abs() < 1e-8,
"d(p,mid)+d(mid,q)={} != d(p,q)={}", d_pm + d_mq, d_full);
}
}
proptest! {
#[test]
fn m_geodesic_endpoints(
p in simplex_vec(6),
q in simplex_vec(6),
) {
let g0 = m_geodesic(&p, &q, 0.0, TOL).unwrap();
let g1 = m_geodesic(&p, &q, 1.0, TOL).unwrap();
for i in 0..p.len() {
prop_assert!((g0[i] - p[i]).abs() < 1e-14,
"m-geodesic t=0: g0[{i}]={} != p[{i}]={}", g0[i], p[i]);
prop_assert!((g1[i] - q[i]).abs() < 1e-14,
"m-geodesic t=1: g1[{i}]={} != q[{i}]={}", g1[i], q[i]);
}
}
}
proptest! {
#[test]
fn e_geodesic_endpoints(
p in positive_simplex_vec(6),
q in positive_simplex_vec(6),
) {
let g0 = e_geodesic(&p, &q, 0.0, TOL).unwrap();
let g1 = e_geodesic(&p, &q, 1.0, TOL).unwrap();
for i in 0..p.len() {
prop_assert!((g0[i] - p[i]).abs() < 1e-10,
"e-geodesic t=0: g0[{i}]={} != p[{i}]={}", g0[i], p[i]);
prop_assert!((g1[i] - q[i]).abs() < 1e-10,
"e-geodesic t=1: g1[{i}]={} != q[{i}]={}", g1[i], q[i]);
}
}
}
proptest! {
#[test]
fn e_geodesic_kl_monotone(
p in positive_simplex_vec(5),
q in positive_simplex_vec(5),
) {
let mut prev_kl = 0.0_f64; for step in 1..=10 {
let t = step as f64 / 10.0;
let gamma_t = e_geodesic(&p, &q, t, TOL).unwrap();
let kl_t = logp::kl_divergence(&p, &gamma_t, TOL).unwrap();
prop_assert!(kl_t >= prev_kl - 1e-10,
"KL not monotone: KL(p, gamma({t}))={kl_t} < prev={prev_kl}");
prev_kl = kl_t;
}
}
}
proptest! {
#[test]
fn natural_gradient_on_uniform_equals_euclidean_scaled(
g in prop::collection::vec(-10.0f64..10.0, 5),
) {
let n = g.len();
let p: Vec<f64> = vec![1.0 / n as f64; n];
let ng = natural_gradient(&p, &g).unwrap();
for i in 0..n {
let expected = g[i] / n as f64;
prop_assert!((ng[i] - expected).abs() < 1e-14,
"natural_gradient[{i}]={} != g[{i}]/n={}", ng[i], expected);
}
}
}
#[test]
fn natural_gradient_length_mismatch_errors() {
let p = vec![0.5, 0.5];
let g = vec![1.0, 2.0, 3.0];
let result = natural_gradient(&p, &g);
assert!(result.is_err());
let err_msg = result.unwrap_err().to_string();
assert!(
err_msg.contains("length mismatch"),
"expected length mismatch error, got: {err_msg}"
);
}
#[test]
fn fisher_rao_geodesic_invalid_t_errors() {
let p = vec![0.5, 0.3, 0.2];
let q = vec![0.2, 0.3, 0.5];
let neg = fisher_rao_geodesic(&p, &q, -0.1, 1e-12);
assert!(neg.is_err());
assert!(neg.unwrap_err().to_string().contains("outside [0, 1]"));
let over = fisher_rao_geodesic(&p, &q, 1.5, 1e-12);
assert!(over.is_err());
assert!(over.unwrap_err().to_string().contains("outside [0, 1]"));
}
#[test]
fn alpha_geodesic_subsumes_m_and_e() {
let p = [0.5, 0.3, 0.2];
let q = [0.2, 0.5, 0.3];
for &t in &[0.0, 0.25, 0.5, 0.75, 1.0] {
let a = alpha_geodesic(&p, &q, t, -1.0, 1e-12).unwrap();
let m = m_geodesic(&p, &q, t, 1e-12).unwrap();
for (x, y) in a.iter().zip(m.iter()) {
assert!((x - y).abs() < 1e-9, "alpha=-1 != m_geodesic at t={t}");
}
let a1 = alpha_geodesic(&p, &q, t, 1.0, 1e-12).unwrap();
let e = e_geodesic(&p, &q, t, 1e-12).unwrap();
for (x, y) in a1.iter().zip(e.iter()) {
assert!((x - y).abs() < 1e-9, "alpha=+1 != e_geodesic at t={t}");
}
}
}
#[test]
fn alpha_geodesic_endpoints_and_simplex() {
let p = [0.5, 0.3, 0.2];
let q = [0.2, 0.5, 0.3];
for &alpha in &[-1.0, -0.5, 0.0, 0.5, 1.0] {
let g0 = alpha_geodesic(&p, &q, 0.0, alpha, 1e-12).unwrap();
let g1 = alpha_geodesic(&p, &q, 1.0, alpha, 1e-12).unwrap();
for (x, y) in g0.iter().zip(p.iter()) {
assert!((x - y).abs() < 1e-9);
}
for (x, y) in g1.iter().zip(q.iter()) {
assert!((x - y).abs() < 1e-9);
}
for &t in &[0.3, 0.5] {
let g = alpha_geodesic(&p, &q, t, alpha, 1e-12).unwrap();
let s: f64 = g.iter().sum();
assert!((s - 1.0).abs() < 1e-9, "alpha={alpha} t={t} sum={s}");
}
}
}
#[test]
fn alpha_geodesic_negative_exponent_needs_positive() {
let p = [0.5, 0.0, 0.5];
let q = [0.2, 0.5, 0.3];
assert!(alpha_geodesic(&p, &q, 0.5, 2.0, 1e-12).is_err());
assert!(alpha_geodesic(&p, &q, 0.5, 0.0, 1e-12).is_ok());
}
}