geographdb-core 0.5.4

//! Natural gradient descent via the Fisher-Rao information metric.
//!
//! Standard gradient descent moves in Euclidean parameter space, which is
//! sensitive to reparameterisation. The natural gradient instead moves along
//! the steepest ascent direction measured in *KL-divergence* (Fisher-Rao
//! distance), making steps that are invariant to how parameters are scaled.
//!
//! For a categorical distribution with logit parameters θ the diagonal Fisher
//! information matrix is F_ii = p_i (1 − p_i).  Preconditioning the gradient
//! by F⁻¹ gives the natural gradient: g_nat_i = g_i / (p_i (1 − p_i)).
//!
//! This module provides:
//! - `softmax`           — stable softmax (no overflow)
//! - `diagonal_fisher`   — diagonal of the Fisher information matrix
//! - `natural_gradient`  — Fisher-preconditioned gradient
//! - `kl_divergence`     — KL(p ‖ q) for step-size measurement
//! - `fisher_rao_dist`   — Bhattacharyya angle: 2 arccos(Σ √(p_i q_i))
//! - `compare_steps`     — one-shot comparison of vanilla vs natural update

// ── Distributions ─────────────────────────────────────────────────────────────

/// Numerically stable softmax.
pub fn softmax(logits: &[f32]) -> Vec<f32> {
    let max = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let exps: Vec<f32> = logits.iter().map(|&l| (l - max).exp()).collect();
    let sum: f32 = exps.iter().sum();
    exps.iter().map(|&e| e / sum).collect()
}

/// Diagonal of the Fisher information matrix for a categorical distribution.
///
/// For logit parameterisation, F_ii = p_i (1 − p_i).
pub fn diagonal_fisher(probs: &[f32]) -> Vec<f32> {
    probs.iter().map(|&p| p * (1.0 - p)).collect()
}

// ── Information Geometry ──────────────────────────────────────────────────────

/// KL divergence KL(p ‖ q) = Σ p_i ln(p_i / q_i).
///
/// Terms where p_i = 0 contribute zero; q_i = 0 with p_i > 0 returns ∞.
pub fn kl_divergence(p: &[f32], q: &[f32]) -> f32 {
    assert_eq!(p.len(), q.len());
    p.iter()
        .zip(q.iter())
        .map(|(&pi, &qi)| {
            if pi <= 0.0 {
                0.0
            } else {
                pi * (pi / qi.max(1e-30)).ln()
            }
        })
        .sum()
}

/// Fisher-Rao geodesic distance (Bhattacharyya angle).
///
/// d(p, q) = 2 arccos(Σ √(p_i · q_i))
///
/// This is the geodesic distance on the statistical manifold of categorical
/// distributions under the Fisher information metric.
pub fn fisher_rao_dist(p: &[f32], q: &[f32]) -> f32 {
    assert_eq!(p.len(), q.len());
    let bc: f32 = p
        .iter()
        .zip(q.iter())
        .map(|(&pi, &qi)| (pi * qi).sqrt())
        .sum();
    2.0 * bc.clamp(-1.0, 1.0).acos()
}

// ── Natural Gradient ──────────────────────────────────────────────────────────

/// Fisher-preconditioned (natural) gradient.
///
/// g_nat_i = g_i / (F_ii + ε) = g_i / (p_i (1 − p_i) + ε)
///
/// `eps` regularises directions where a component probability is near 0 or 1
/// (Fisher information collapses at the boundary of the simplex).
pub fn natural_gradient(grad: &[f32], probs: &[f32], eps: f32) -> Vec<f32> {
    assert_eq!(grad.len(), probs.len());
    let fisher = diagonal_fisher(probs);
    grad.iter()
        .zip(fisher.iter())
        .map(|(&g, &f)| g / (f + eps))
        .collect()
}

// ── Step Comparison ───────────────────────────────────────────────────────────

/// One step comparison of vanilla gradient descent vs natural gradient descent.
pub struct StepComparison {
    /// Logits before the step (reference distribution p₀).
    pub logits_before: Vec<f32>,
    /// Vanilla update: θ ← θ − lr · g.
    pub logits_vanilla: Vec<f32>,
    /// Natural update: θ ← θ − lr · F⁻¹ g.
    pub logits_natural: Vec<f32>,
    /// Euclidean distance traveled in logit space (vanilla).
    pub euclid_vanilla: f32,
    /// Euclidean distance traveled in logit space (natural).
    pub euclid_natural: f32,
    /// KL(p₀ ‖ p_vanilla) — information divergence of vanilla step.
    pub kl_vanilla: f32,
    /// KL(p₀ ‖ p_natural) — information divergence of natural step.
    pub kl_natural: f32,
    /// Fisher-Rao geodesic distance (vanilla).
    pub fisher_rao_vanilla: f32,
    /// Fisher-Rao geodesic distance (natural).
    pub fisher_rao_natural: f32,
}

/// Compare one vanilla gradient step vs one natural gradient step.
///
/// Both start from `logits` and move against `grad` with step size `lr`.
/// The natural gradient is regularised by `eps` to avoid division by near-zero
/// Fisher diagonals.
pub fn compare_steps(logits: &[f32], grad: &[f32], lr: f32, eps: f32) -> StepComparison {
    let probs_before = softmax(logits);
    let nat_grad = natural_gradient(grad, &probs_before, eps);

    let logits_vanilla: Vec<f32> = logits.iter().zip(grad).map(|(&t, &g)| t - lr * g).collect();
    let logits_natural: Vec<f32> = logits
        .iter()
        .zip(&nat_grad)
        .map(|(&t, &ng)| t - lr * ng)
        .collect();

    let probs_vanilla = softmax(&logits_vanilla);
    let probs_natural = softmax(&logits_natural);

    let euclid_vanilla = logits
        .iter()
        .zip(&logits_vanilla)
        .map(|(&a, &b)| (a - b).powi(2))
        .sum::<f32>()
        .sqrt();
    let euclid_natural = logits
        .iter()
        .zip(&logits_natural)
        .map(|(&a, &b)| (a - b).powi(2))
        .sum::<f32>()
        .sqrt();

    StepComparison {
        logits_before: logits.to_vec(),
        logits_vanilla,
        logits_natural,
        euclid_vanilla,
        euclid_natural,
        kl_vanilla: kl_divergence(&probs_before, &probs_vanilla),
        kl_natural: kl_divergence(&probs_before, &probs_natural),
        fisher_rao_vanilla: fisher_rao_dist(&probs_before, &probs_vanilla),
        fisher_rao_natural: fisher_rao_dist(&probs_before, &probs_natural),
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    const EPS: f32 = 1e-5;

    #[test]
    fn test_softmax_sums_to_one() {
        let logits = vec![1.0f32, 2.0, 0.5, -1.0, 3.0];
        let p = softmax(&logits);
        let sum: f32 = p.iter().sum();
        assert!((sum - 1.0).abs() < EPS, "softmax must sum to 1, got {sum}");
    }

    #[test]
    fn test_softmax_argmax_preserved() {
        let logits = vec![0.1f32, 5.0, 0.3, -2.0];
        let p = softmax(&logits);
        let argmax = p
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .map(|(i, _)| i)
            .unwrap();
        assert_eq!(argmax, 1, "largest logit must map to largest probability");
    }

    #[test]
    fn test_softmax_uniform_logits() {
        let logits = vec![0.0f32; 4];
        let p = softmax(&logits);
        for &pi in &p {
            assert!((pi - 0.25).abs() < EPS, "uniform logits → uniform probs");
        }
    }

    #[test]
    fn test_diagonal_fisher_uniform() {
        // Uniform distribution over 4 classes: p_i = 0.25, F_ii = 0.25*0.75 = 0.1875
        let probs = vec![0.25f32; 4];
        let f = diagonal_fisher(&probs);
        for &fi in &f {
            assert!(
                (fi - 0.1875).abs() < EPS,
                "F_ii for uniform = 0.1875, got {fi}"
            );
        }
    }

    #[test]
    fn test_diagonal_fisher_concentrated() {
        // If p_0 ≈ 1.0, F_00 ≈ 0 (flat landscape at boundary).
        let probs = vec![0.999f32, 0.0005, 0.0005];
        let f = diagonal_fisher(&probs);
        assert!(
            f[0] < 0.01,
            "Fisher near 1.0 should be near 0, got {}",
            f[0]
        );
    }

    #[test]
    fn test_natural_gradient_scales_by_fisher_inverse() {
        // For uniform probs, F_ii = p*(1-p).  g_nat_i = g_i / F_ii.
        let probs = vec![0.5f32, 0.5];
        let grad = vec![1.0f32, -1.0];
        let nat = natural_gradient(&grad, &probs, 0.0);
        // F_ii = 0.5 * 0.5 = 0.25, so g_nat = g / 0.25 = 4*g
        assert!(
            (nat[0] - 4.0).abs() < EPS,
            "natural grad should be 4x vanilla at p=0.5"
        );
        assert!(
            (nat[1] + 4.0).abs() < EPS,
            "natural grad should be 4x vanilla at p=0.5"
        );
    }

    #[test]
    fn test_natural_gradient_eps_regularization() {
        // p_i near 0 → F_ii near 0; eps prevents division explosion.
        let probs = vec![0.0001f32, 0.9999];
        let grad = vec![1.0f32, 0.0];
        let nat_eps = natural_gradient(&grad, &probs, 1e-3);
        let nat_no_eps = natural_gradient(&grad, &probs, 1e-30);
        // With eps, the output is bounded.
        assert!(nat_eps[0].is_finite(), "eps should prevent NaN/inf");
        // Without eps it's large but finite here (since F_ii = 0.0001*0.9999 > 0).
        assert!(
            nat_no_eps[0] > nat_eps[0],
            "eps reduces magnitude near boundary"
        );
    }

    #[test]
    fn test_kl_divergence_self_is_zero() {
        let p = vec![0.2f32, 0.5, 0.3];
        let kl = kl_divergence(&p, &p);
        assert!(kl.abs() < EPS, "KL(p||p) must be zero, got {kl}");
    }

    #[test]
    fn test_kl_divergence_positive() {
        let p = vec![0.7f32, 0.3];
        let q = vec![0.3f32, 0.7];
        assert!(
            kl_divergence(&p, &q) > 0.0,
            "KL divergence between different distributions must be positive"
        );
    }

    #[test]
    fn test_fisher_rao_dist_self_is_zero() {
        let p = vec![0.4f32, 0.4, 0.2];
        let d = fisher_rao_dist(&p, &p);
        assert!(
            d.abs() < EPS,
            "Fisher-Rao distance to self must be zero, got {d}"
        );
    }

    #[test]
    fn test_fisher_rao_dist_nonnegative() {
        let p = vec![0.6f32, 0.4];
        let q = vec![0.2f32, 0.8];
        assert!(fisher_rao_dist(&p, &q) >= 0.0);
    }

    #[test]
    fn test_natural_grad_smaller_fisher_rao_step() {
        // Flat gradient on logits with one dominant class.  The natural gradient
        // corrects for the curvature of the simplex so both updates move a more
        // principled KL distance.  With a high-curvature direction (p close to 1),
        // the vanilla step overshoots in Fisher-Rao space; the natural step dampens
        // it.
        //
        // Setup: logits tilted so class 0 dominates; gradient pushes class 0 up more.
        let logits = vec![3.0f32, 0.5, 0.5, 0.5];
        let grad = vec![1.0f32, 0.1, 0.1, 0.1]; // pushes already-dominant class
        let lr = 0.3;

        let cmp = compare_steps(&logits, &grad, lr, 1e-4);

        // The vanilla step drives p_0 even higher, a large move in simplex space.
        // The natural step rescales the dominant component down (small F_ii at
        // boundary → large g_nat → but softmax re-normalises logits).
        // Both should be finite and positive.
        assert!(cmp.kl_vanilla.is_finite() && cmp.kl_vanilla >= 0.0);
        assert!(cmp.kl_natural.is_finite() && cmp.kl_natural >= 0.0);
        assert!(cmp.fisher_rao_vanilla.is_finite());
        assert!(cmp.fisher_rao_natural.is_finite());
    }
}