stats_claw/distributions/sampling/
chi_squared.rs

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

stats_claw/distributions/sampling/
chi_squared.rs
