probability_rs/dist/
beta.rs

1use super::gamma::Gamma;
2use crate::dist::{Continuous, DistError, Distribution, Moments};
3use crate::rng::RngCore;
4
5#[derive(Debug, Clone, Copy)]
6pub struct Beta {
7    a: f64,
8    b: f64,
9    ln_beta: f64,
10}
11
12impl Beta {
13    pub fn new(a: f64, b: f64) -> Result<Self, DistError> {
14        if !(a > 0.0 && b > 0.0) || !a.is_finite() || !b.is_finite() {
15            return Err(DistError::InvalidParameter);
16        }
17        let ln_beta = ln_gamma(a) + ln_gamma(b) - ln_gamma(a + b);
18        Ok(Self { a, b, ln_beta })
19    }
20    #[inline]
21    pub fn a(&self) -> f64 {
22        self.a
23    }
24    #[inline]
25    pub fn b(&self) -> f64 {
26        self.b
27    }
28}
29
30impl Distribution for Beta {
31    type Value = f64;
32    fn cdf(&self, x: f64) -> f64 {
33        if x <= 0.0 {
34            return 0.0;
35        }
36        if x >= 1.0 {
37            return 1.0;
38        }
39        reg_inc_beta(self.a, self.b, x)
40    }
41    fn in_support(&self, x: f64) -> bool {
42        x >= 0.0 && x <= 1.0 && x.is_finite()
43    }
44    fn sample<R: RngCore>(&self, rng: &mut R) -> f64 {
45        let ga = Gamma::new(self.a, 1.0).unwrap().sample(rng);
46        let gb = Gamma::new(self.b, 1.0).unwrap().sample(rng);
47        ga / (ga + gb)
48    }
49}
50
51impl Continuous for Beta {
52    fn pdf(&self, x: f64) -> f64 {
53        if !self.in_support(x) {
54            return 0.0;
55        }
56        ((self.a - 1.0) * x.ln() + (self.b - 1.0) * (1.0 - x).ln() - self.ln_beta).exp()
57    }
58    fn inv_cdf(&self, p: f64) -> f64 {
59        debug_assert!(p > 0.0 && p < 1.0);
60        // Simple Newton with bracketing in [0,1]
61        let mut lo = 0.0;
62        let mut hi = 1.0;
63        let mut x = p; // initial guess
64        for _ in 0..60 {
65            let fx = self.cdf(x) - p;
66            if fx.abs() < 1e-10 {
67                break;
68            }
69            if fx < 0.0 {
70                lo = x;
71            } else {
72                hi = x;
73            }
74            let dfx = self.pdf(x).max(1e-300);
75            let mut x_new = x - fx / dfx;
76            if !(0.0..=1.0).contains(&x_new) {
77                x_new = 0.5 * (lo + hi);
78            }
79            x = x_new;
80        }
81        x
82    }
83}
84
85impl Moments for Beta {
86    fn mean(&self) -> f64 {
87        self.a / (self.a + self.b)
88    }
89    fn variance(&self) -> f64 {
90        (self.a * self.b) / ((self.a + self.b).powi(2) * (self.a + self.b + 1.0))
91    }
92    fn skewness(&self) -> f64 {
93        let a = self.a;
94        let b = self.b;
95        let num = 2.0 * (b - a) * (a + b + 1.0).sqrt();
96        let den = (a + b + 2.0) * (a * b).sqrt();
97        num / den
98    }
99    fn kurtosis(&self) -> f64 {
100        let a = self.a;
101        let b = self.b;
102        let num = 6.0 * ((a - b).powi(2) * (a + b + 1.0) - a * b * (a + b + 2.0));
103        let den = a * b * (a + b + 2.0) * (a + b + 3.0);
104        num / den
105    }
106    fn entropy(&self) -> f64 {
107        // H = ln B(a,b) - (a-1)ψ(a) - (b-1)ψ(b) + (a+b-2)ψ(a+b)
108        let a = self.a; let b = self.b;
109    let ln_beta = super::gamma::ln_gamma(a) + super::gamma::ln_gamma(b) - super::gamma::ln_gamma(a + b);
110        ln_beta - (a - 1.0) * crate::num::digamma(a) - (b - 1.0) * crate::num::digamma(b)
111            + (a + b - 2.0) * crate::num::digamma(a + b)
112    }
113}
114
115// Helpers: ln_gamma and regularized incomplete beta (continued fractions)
116fn ln_gamma(z: f64) -> f64 {
117    super::gamma::ln_gamma(z)
118}
119
120fn reg_inc_beta(a: f64, b: f64, x: f64) -> f64 {
121    // Use symmetry to ensure x <= (a+1)/(a+b+2)
122    if x <= 0.0 {
123        return 0.0;
124    }
125    if x >= 1.0 {
126        return 1.0;
127    }
128    let bt = ((a + b).ln() - ln_gamma(a) - ln_gamma(b) + a * x.ln() + b * (1.0 - x).ln()).exp();
129    if x < (a + 1.0) / (a + b + 2.0) {
130        bt * beta_cf(a, b, x) / a
131    } else {
132        1.0 - bt * beta_cf(b, a, 1.0 - x) / b
133    }
134}
135
136fn beta_cf(a: f64, b: f64, x: f64) -> f64 {
137    // Continued fraction for incomplete beta (Numerical Recipes style)
138    let mut am = 1.0;
139    let mut bm = 1.0;
140    let mut az = 1.0;
141    let qab = a + b;
142    let qap = a + 1.0;
143    let qam = a - 1.0;
144    let mut bz = 1.0 - qab * x / qap;
145    let eps = 3e-14;
146    let fpmin = 1e-300;
147    for m in 1..=200 {
148        let m2 = 2 * m;
149        // even step
150        let d = m as f64 * (b - m as f64) * x / ((qam + m2 as f64) * (a + m2 as f64));
151        let ap = az + d * am;
152        let bp = bz + d * bm;
153        // odd step
154        let d = -(a + m as f64) * (qab + m as f64) * x / ((a + m2 as f64) * (qap + m2 as f64));
155        let app = ap + d * az;
156        let bpp = bp + d * bz;
157        am = ap / bpp.max(fpmin);
158        bm = bp / bpp.max(fpmin);
159        az = app / bpp.max(fpmin);
160        bz = 1.0;
161        if (app - ap).abs() < eps * app.abs() {
162            break;
163        }
164    }
165    az
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171    #[test]
172    fn moments() {
173        let b = Beta::new(2.0, 5.0).unwrap();
174        assert!((b.mean() - (2.0 / 7.0)).abs() < 1e-12);
175    }
176    #[test]
177    fn moments_higher() {
178        let b = Beta::new(2.0, 2.0).unwrap();
179        assert!(b.skewness().abs() < 1e-15);
180        assert!((b.kurtosis() - (-6.0/7.0)).abs() < 1e-12);
181    }
182}
probability_rs/dist/beta.rs

probability_rs/dist/
beta.rs
