numeris/stats/
chi_squared.rs

1use crate::FloatScalar;
2use crate::special::{lgamma, gamma_inc};
3use super::{ContinuousDistribution, StatsError, quantile_newton, normal_quantile_standard};
4
5/// Chi-squared distribution with k degrees of freedom.
6///
7/// Special case of Gamma(k/2, 1/2).
8///
9/// # Example
10///
11/// ```
12/// use numeris::stats::{ChiSquared, ContinuousDistribution};
13///
14/// let chi2 = ChiSquared::new(3.0_f64).unwrap();
15/// assert!((chi2.mean() - 3.0).abs() < 1e-14);
16/// assert!((chi2.variance() - 6.0).abs() < 1e-14);
17/// ```
18#[derive(Debug, Clone, Copy)]
19pub struct ChiSquared<T> {
20    k: T, // degrees of freedom
21}
22
23impl<T: FloatScalar> ChiSquared<T> {
24    /// Create a chi-squared distribution with `k` degrees of freedom. Requires `k > 0`.
25    pub fn new(k: T) -> Result<Self, StatsError> {
26        if k <= T::zero() {
27            return Err(StatsError::InvalidParameter);
28        }
29        Ok(Self { k })
30    }
31}
32
33impl<T: FloatScalar> ChiSquared<T> {
34    /// Draw a random sample from this distribution.
35    ///
36    /// Chi-squared(k) = Gamma(k/2, 1/2), so we sample Gamma(k/2, 1) * 2.
37    pub fn sample(&self, rng: &mut super::Rng) -> T {
38        let two = T::one() + T::one();
39        rng.next_gamma(self.k / two) * two
40    }
41
42    /// Fill a fixed-size array with independent samples.
43    pub fn sample_array<const K: usize>(&self, rng: &mut super::Rng) -> [T; K] {
44        let mut out = [T::zero(); K];
45        for v in out.iter_mut() {
46            *v = self.sample(rng);
47        }
48        out
49    }
50}
51
52impl<T: FloatScalar> ContinuousDistribution<T> for ChiSquared<T> {
53    fn pdf(&self, x: T) -> T {
54        if x <= T::zero() {
55            return T::zero();
56        }
57        self.ln_pdf(x).exp()
58    }
59
60    fn ln_pdf(&self, x: T) -> T {
61        if x <= T::zero() {
62            return T::neg_infinity();
63        }
64        let one = T::one();
65        let two = one + one;
66        let half_k = self.k / two;
67        (half_k - one) * x.ln() - x / two - half_k * two.ln() - lgamma(half_k)
68    }
69
70    fn cdf(&self, x: T) -> T {
71        if x <= T::zero() {
72            return T::zero();
73        }
74        let two = T::one() + T::one();
75        gamma_inc(self.k / two, x / two).unwrap_or(T::nan())
76    }
77
78    fn quantile(&self, p: T) -> T {
79        let two = T::one() + T::one();
80        let nine = T::from(9.0).unwrap();
81        // Wilson-Hilferty approximation
82        let z = normal_quantile_standard(p);
83        let v = T::one() - two / (nine * self.k) + z * (two / (nine * self.k)).sqrt();
84        let x0 = if v > T::zero() { self.k * v * v * v } else { self.mean() };
85        let hi = self.mean() + T::from(40.0).unwrap() * self.variance().sqrt();
86        quantile_newton(|x| self.cdf(x), |x| self.pdf(x), p, x0, T::zero(), hi)
87    }
88
89    fn mean(&self) -> T {
90        self.k
91    }
92
93    fn variance(&self) -> T {
94        let two = T::one() + T::one();
95        two * self.k
96    }
97}
numeris/stats/chi_squared.rs

numeris/stats/
chi_squared.rs
