stats_claw/distributions/sampling/
chi_squared.rs

1//! Chi-squared distribution numerics, for the
2//! [`ChiSquaredDistribution`].
3//!
4//! Equivalent to `scipy.stats.chi2(degrees_of_freedom)`: the distribution of a
5//! sum of `k` squared standard normals. The density is closed form, the CDF is
6//! the regularized lower incomplete gamma `gamma_p`, the quantile inverts it by
7//! bisection, and sampling sums squared `standard_normal` draws.
8//!
9//! # Examples
10//!
11//! ```
12//! use stats_claw::distributions::{Cdf, Moments};
13//! use stats_claw::distributions::ChiSquaredDistribution;
14//!
15//! let d = ChiSquaredDistribution { degrees_of_freedom: 5, ..Default::default() };
16//! // Mean equals the degrees of freedom.
17//! assert_eq!(d.mean(), Some(5.0));
18//! // CDF at 0 is 0; CDF at a very large value is ≈ 1.
19//! assert!(d.cdf(0.0).abs() < 1e-12);
20//! assert!(d.cdf(1_000.0) > 1.0 - 1e-9);
21//! ```
22
23use super::super::{Cdf, LogCdf, Moments, Pdf, Quantile, Sample, bisection_quantile, count_to_f64};
24use crate::distributions::ChiSquaredDistribution;
25use crate::rng::SplitMix64;
26use crate::special::{gamma_p, ln_gamma, ln_gamma_p, ln_gamma_q};
27
28impl ChiSquaredDistribution {
29    /// Degrees of freedom as a float (`k`).
30    fn k(&self) -> f64 {
31        count_to_f64(self.degrees_of_freedom)
32    }
33}
34
35impl Pdf for ChiSquaredDistribution {
36    fn pdf(&self, x: f64) -> f64 {
37        if x < 0.0 {
38            return 0.0;
39        }
40        if x == 0.0 {
41            // Density is 0 for k>2, 0.5 for k=2, and diverges for k<2; scipy's
42            // grid starts at 0 so report the k≥2 closed-form limit.
43            return if (self.k() - 2.0).abs() < 1e-12 {
44                0.5
45            } else {
46                0.0
47            };
48        }
49        let half_k = 0.5 * self.k();
50        let log_norm = half_k.mul_add(2.0_f64.ln(), ln_gamma(half_k));
51        let tail = (-0.5f64).mul_add(x, -log_norm);
52        let ln_density = (half_k - 1.0).mul_add(x.ln(), tail);
53        ln_density.exp()
54    }
55}
56
57impl Cdf for ChiSquaredDistribution {
58    fn cdf(&self, x: f64) -> f64 {
59        if x <= 0.0 {
60            0.0
61        } else {
62            gamma_p(0.5 * self.k(), 0.5 * x)
63        }
64    }
65}
66
67impl LogCdf for ChiSquaredDistribution {
68    fn logsf(&self, x: f64) -> f64 {
69        // sf(x) = Q(k/2, x/2); its log stays finite where Q underflows.
70        if x <= 0.0 {
71            0.0
72        } else {
73            ln_gamma_q(0.5 * self.k(), 0.5 * x)
74        }
75    }
76    fn logcdf(&self, x: f64) -> f64 {
77        // cdf(x) = P(k/2, x/2).
78        if x <= 0.0 {
79            f64::NEG_INFINITY
80        } else {
81            ln_gamma_p(0.5 * self.k(), 0.5 * x)
82        }
83    }
84}
85
86impl Quantile for ChiSquaredDistribution {
87    fn quantile(&self, p: f64) -> f64 {
88        let k = self.k();
89        let hi = 20.0f64.mul_add((2.0 * k).sqrt(), k).max(1.0);
90        bisection_quantile(p, 0.0, hi, |x| self.cdf(x))
91    }
92}
93
94impl Moments for ChiSquaredDistribution {
95    fn mean(&self) -> Option<f64> {
96        Some(self.k())
97    }
98    fn variance(&self) -> Option<f64> {
99        Some(2.0 * self.k())
100    }
101}
102
103impl Sample for ChiSquaredDistribution {
104    fn sample(&self, rng: &mut SplitMix64) -> f64 {
105        let mut sum = 0.0;
106        for _ in 0..self.degrees_of_freedom.max(0) {
107            let z = rng.standard_normal();
108            sum = z.mul_add(z, sum);
109        }
110        sum
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    /// The mean equals the degrees of freedom and variance twice that.
119    #[test]
120    fn moments_are_df_and_twice_df() {
121        let d = ChiSquaredDistribution {
122            degrees_of_freedom: 5,
123            ..Default::default()
124        };
125        assert_eq!(d.mean(), Some(5.0));
126        assert_eq!(d.variance(), Some(10.0));
127    }
128
129    /// The CDF is monotone and saturates toward 1 in the upper tail.
130    #[test]
131    fn cdf_saturates() {
132        let d = ChiSquaredDistribution {
133            degrees_of_freedom: 3,
134            ..Default::default()
135        };
136        assert!(d.cdf(0.0).abs() < 1e-12);
137        assert!(d.cdf(1e6) > 1.0 - 1e-9);
138    }
139
140    /// `logsf` agrees with the log of the linear survival in the body and matches
141    /// scipy `chi2.logsf` deep in the tail (df = 3, x = 200).
142    #[test]
143    fn logsf_matches_scipy_and_stays_finite() {
144        let d = ChiSquaredDistribution {
145            degrees_of_freedom: 3,
146            ..Default::default()
147        };
148        let body = d.logsf(2.0);
149        assert!(
150            ((body - (1.0 - d.cdf(2.0)).ln()) / body.abs()).abs() < 1e-9,
151            "logsf body was {body}"
152        );
153        let tail = d.logsf(200.0);
154        let want = -97.571_669_639_663_45; // scipy.stats.chi2.logsf(200, 3)
155        assert!(tail.is_finite(), "logsf(200) was {tail}");
156        assert!(
157            ((tail - want) / want).abs() < 1e-9,
158            "logsf(200) = {tail}, want {want}"
159        );
160    }
161}
stats_claw/distributions/sampling/chi_squared.rs

stats_claw/distributions/sampling/
chi_squared.rs
