Fast Newton-Raphson inverse_cdf for Chi and InverseGamma

src/distribution/chi.rs

@@ -148,6 +148,32 @@ impl ContinuousCDF<f64, f64> for Chi {
             gamma::gamma_ur(self.freedom() as f64 / 2.0, x * x / 2.0)
         }
     }
+
+    /// Calculates the inverse cumulative distribution function for the chi
+    /// distribution at `p`, i.e. the `p`-quantile.
+    ///
+    /// # Panics
+    ///
+    /// If `p` is not in `[0, 1]`.
+    fn inverse_cdf(&self, p: f64) -> f64 {
+        if !(0.0..=1.0).contains(&p) {
+            panic!("p must be in [0, 1]")
+        }
+        if p == 0.0 {
+            return self.min();
+        }
+        if p == 1.0 {
+            return self.max();
+        }
+        // The chi cdf has no closed-form inverse; solve it with the shared
+        // safeguarded Newton search instead of the generic bisection default.
+        super::internal::newton_raphson_quantile(
+            p,
+            |x| self.cdf(x),
+            |x| self.sf(x),
+            |x| self.pdf(x),
+        )
+    }
 }
 
 impl Min<f64> for Chi {
@@ -553,4 +579,23 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn test_inverse_cdf_p0_p1() {
+        let d = create_ok(3);
+        assert_eq!(d.inverse_cdf(0.0), d.min());
+        assert_eq!(d.inverse_cdf(1.0), d.max());
+    }
+
+    #[test]
+    #[should_panic(expected = "p must be in [0, 1]")]
+    fn test_inverse_cdf_p_above_one() {
+        create_ok(3).inverse_cdf(1.0 + f64::EPSILON);
+    }
+
+    #[test]
+    #[should_panic(expected = "p must be in [0, 1]")]
+    fn test_inverse_cdf_p_below_zero() {
+        create_ok(3).inverse_cdf(-1e-300);
+    }
 }

src/distribution/internal.rs

@@ -37,6 +37,93 @@ pub fn integral_bisection_search<K: Num + Clone, T: Num + PartialOrd>(
     }
 }
 
+/// Quantile `F^{-1}(p)` for a continuous distribution supported on `(0, ∞)`
+/// whose cdf has no closed-form inverse, found with a safeguarded Newton–Raphson
+/// step (Numerical Recipes' `rtsafe`); `cdf`, `sf` and `pdf` are the
+/// distribution's own functions and `p` must lie strictly inside `(0, 1)`.
+///
+/// The bracket `[low, high]` is kept as an invariant and only ever tightened, so
+/// a Newton step that is non-finite, would leave the bracket, or is not shrinking
+/// at least as fast as a bisection falls back to bisection; the iterate can never
+/// escape the support — the guard that keeps a quantile sitting against a
+/// boundary from diverging to NaN (cf. Gamma in #382) — and the bracket is halved
+/// often enough to reach the root within the iteration budget. In the upper half
+/// the survival function is inverted rather than the cdf: as `cdf` saturates to
+/// one it can no longer resolve a deep upper-tail quantile, whereas `sf` stays
+/// well conditioned there.
+pub fn newton_raphson_quantile(
+    p: f64,
+    cdf: impl Fn(f64) -> f64,
+    sf: impl Fn(f64) -> f64,
+    pdf: impl Fn(f64) -> f64,
+) -> f64 {
+    // Bracket the quantile within a factor of two, `cdf(low) <= p <= cdf(high)`,
+    // by walking a unit interval out toward it. Moving *both* ends keeps the
+    // bracket tight when the quantile is far from 1 (deep in either tail), so the
+    // Newton phase starts close and converges in a handful of steps rather than
+    // bisecting the whole span.
+    let mut low = 1.0;
+    let mut high = 2.0;
+    while cdf(low) > p {
+        high = low;
+        low /= 2.0;
+    }
+    while cdf(high) < p {
+        low = high;
+        high *= 2.0;
+    }
+
+    // Solve `sf(x) = 1 - p` in the upper half and `cdf(x) = p` otherwise; either
+    // way the residual is increasing in `x` with derivative `pdf(x)`.
+    let upper = p > 0.5;
+    let target = if upper { 1.0 - p } else { p };
+
+    // A *relative* accuracy target: quantiles span many orders of magnitude, so
+    // an absolute tolerance (as in `prec::convergence`) is meaningless deep in a
+    // tail where the quantile itself is far smaller than that tolerance.
+    let accuracy = crate::prec::DEFAULT_RELATIVE_ACC;
+    const MAX_ITERATIONS: usize = 100;
+    let mut x = (low + high) / 2.0;
+    let mut last_step = high - low;
+    for _ in 0..MAX_ITERATIONS {
+        let residual = if upper {
+            target - sf(x)
+        } else {
+            cdf(x) - target
+        };
+        let newton = x - residual / pdf(x);
+        // A full Newton step below the relative tolerance means we have
+        // converged; accept it before the bracket bookkeeping below, which would
+        // otherwise reject a converged step that has rounded onto the very
+        // endpoint we are about to move to `x`.
+        if (newton - x).abs() <= accuracy * x.abs() {
+            return newton;
+        }
+        // Tighten the bracket by the sign of the (increasing) residual.
+        if residual >= 0.0 {
+            high = x;
+        } else {
+            low = x;
+        }
+        // Take the Newton step only while it stays strictly inside the bracket
+        // *and* shrinks at least as fast as a bisection would, else bisect. The
+        // second half of that test is what bounds the iteration count: a tail
+        // where the step is merely linear — the inverse gamma's `exp(-b/x)`
+        // advances `b/x` by one per step regardless of how far the root is —
+        // otherwise creeps toward the root and runs out of iterations short of
+        // it, e.g. `InverseGamma(1, 1).inverse_cdf(1e-200)` off by 4%.
+        let step = (newton - x).abs();
+        if newton.is_finite() && newton > low && newton < high && 2.0 * step <= last_step {
+            x = newton;
+            last_step = step;
+        } else {
+            last_step = (high - low) / 2.0;
+            x = low + last_step;
+        }
+    }
+    x
+}
+
 #[cfg(test)]
 macro_rules! testing_boiler {
     ($($arg_name:ident: $arg_ty:ty),+; $dist:ty; $dist_err:ty) => {
@@ -460,6 +547,38 @@ mod test {
         }
     }
 
+    #[cfg(feature = "std")]
+    #[test]
+    fn test_newton_raphson_quantile() {
+        use super::newton_raphson_quantile;
+        // Exponential(λ) has the closed-form quantile -ln(1 - p)/λ, so it is an
+        // exact oracle for the shared solver across both tails.
+        for lambda in [0.5f64, 1.0, 4.0] {
+            let cdf = |x: f64| -(-lambda * x).exp_m1();
+            let sf = |x: f64| (-lambda * x).exp();
+            let pdf = |x: f64| lambda * (-lambda * x).exp();
+            for p in [
+                1e-12,
+                1e-8,
+                1e-4,
+                1e-2,
+                0.25,
+                0.5,
+                0.75,
+                0.99,
+                1.0 - 1e-6,
+                1.0 - 1e-12,
+            ] {
+                let got = newton_raphson_quantile(p, cdf, sf, pdf);
+                let want = -(-p).ln_1p() / lambda; // -ln(1 - p)/λ, accurate in both tails
+                assert!(
+                    (got - want).abs() <= 1e-12 * want,
+                    "Exp({lambda}).inverse_cdf({p}) = {got}, want {want}"
+                );
+            }
+        }
+    }
+
     pub mod boiler_tests {
         use crate::distribution::{Beta, BetaError};
         use crate::statistics::*;

src/distribution/inverse_gamma.rs

@@ -173,6 +173,33 @@ impl ContinuousCDF<f64, f64> for InverseGamma {
             gamma::gamma_lr(self.shape, self.rate / x)
         }
     }
+
+    /// Calculates the inverse cumulative distribution function for the inverse
+    /// gamma distribution at `p`, i.e. the `p`-quantile.
+    ///
+    /// # Panics
+    ///
+    /// If `p` is not in `[0, 1]`.
+    fn inverse_cdf(&self, p: f64) -> f64 {
+        if !(0.0..=1.0).contains(&p) {
+            panic!("p must be in [0, 1]")
+        }
+        if p == 0.0 {
+            return self.min();
+        }
+        if p == 1.0 {
+            return self.max();
+        }
+        // The inverse gamma cdf has no closed-form inverse; solve it with the
+        // shared safeguarded Newton search instead of the generic bisection
+        // default, which loses precision far into the (very heavy) upper tail.
+        super::internal::newton_raphson_quantile(
+            p,
+            |x| self.cdf(x),
+            |x| self.sf(x),
+            |x| self.pdf(x),
+        )
+    }
 }
 
 impl Min<f64> for InverseGamma {
@@ -513,4 +540,46 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn test_inverse_cdf_deep_lower_tail() {
+        // A Newton step on `exp(-rate/x)` only advances `rate/x` by one however
+        // far the root is, so without the bisection-progress guard the search
+        // runs out of iterations short of it. References from mpmath (dps=60).
+        let cases: &[(f64, f64, f64, f64)] = &[
+            (1.0, 1.0, 1e-100, 0.0043429448190325185),
+            (1.0, 1.0, 1e-200, 0.0021714724095162593),
+            (1.0, 1.0, 1e-300, 0.0014476482730108394),
+            (0.01, 1.0, 1e-200, 0.00222287682578071),
+            (0.01, 1.0, 1e-300, 0.0014711980613782147),
+            (100.0, 1.0, 1e-100, 0.0020694527796494867),
+            (100.0, 1.0, 1e-200, 0.0013193408747368574),
+            (1.0, 400.0, 1e-300, 0.5790593092043358),
+            (0.5, 1e6, 1e-200, 2188.7520668986626),
+        ];
+        for &(a, b, p, expected) in cases {
+            let q = InverseGamma::new(a, b).unwrap().inverse_cdf(p);
+            let relerr = ((q - expected) / expected).abs();
+            assert!(relerr <= 1e-14, "InverseGamma({a}, {b}).inverse_cdf({p}) = {q}, want {expected} (relerr {relerr:e})");
+        }
+    }
+
+    #[test]
+    fn test_inverse_cdf_p0_p1() {
+        let d = create_ok(1.0, 1.0);
+        assert_eq!(d.inverse_cdf(0.0), d.min());
+        assert_eq!(d.inverse_cdf(1.0), d.max());
+    }
+
+    #[test]
+    #[should_panic(expected = "p must be in [0, 1]")]
+    fn test_inverse_cdf_p_above_one() {
+        create_ok(1.0, 1.0).inverse_cdf(1.0 + f64::EPSILON);
+    }
+
+    #[test]
+    #[should_panic(expected = "p must be in [0, 1]")]
+    fn test_inverse_cdf_p_below_zero() {
+        create_ok(1.0, 1.0).inverse_cdf(-1e-300);
+    }
 }