puruspe/
beta.rs

1//! This module provides implementations of the beta function and related functions.
2//!
3//! It includes the following main functions:
4//! - `beta`: Calculates the beta function.
5//! - `betai`: Calculates the regularized incomplete beta function.
6//! - `invbetai`: Calculates the inverse of the regularized incomplete beta function.
7
8use crate::{ln_gamma, EPS, FPMIN, W, Y};
9const SWITCH: usize = 3000;
10
11/// Calculates the beta function.
12///
13/// The beta function is defined as:
14///
15/// $$ B(z,w) = \int_0^1 t^{z-1} (1-t)^{w-1} dt $$
16///
17/// It can also be expressed in terms of gamma functions:
18///
19/// $$ B(z,w) = \frac{\Gamma(z)\Gamma(w)}{\Gamma(z+w)} $$
20///
21/// # Arguments
22///
23/// * `z` - First parameter
24/// * `w` - Second parameter
25///
26/// # Returns
27///
28/// The value of the beta function $B(z,w)$
29pub fn beta(z: f64, w: f64) -> f64 {
30    (ln_gamma(z) + ln_gamma(w) - ln_gamma(z + w)).exp()
31}
32
33// =============================================================================
34// Incomplete Beta function
35// =============================================================================
36/// Calculates the regularized incomplete beta function.
37///
38/// The regularized incomplete beta function is defined as:
39///
40/// $$ I_x(a,b) = \frac{B(x;a,b)}{B(a,b)} = \frac{1}{B(a,b)} \int_0^x t^{a-1} (1-t)^{b-1} dt $$
41///
42/// where $B(a,b)$ is the beta function and $B(x;a,b)$ is the incomplete beta function.
43///
44/// # Arguments
45///
46/// * `a` - First shape parameter
47/// * `b` - Second shape parameter
48/// * `x` - Upper limit of integration (between 0 and 1)
49///
50/// # Returns
51///
52/// The value of the regularized incomplete beta function $I_x(a,b)$
53///
54/// # Panics
55///
56/// Panics if `a` ≤ 0, if `b` ≤ 0 or if x is not in the range `0..=1`.
57pub fn betai(a: f64, b: f64, x: f64) -> f64 {
58    assert!(a > 0f64 && b > 0f64, "Bad a or b in routine betai");
59    assert!((0f64..=1f64).contains(&x), "Bad x in routine betai");
60    if x == 0f64 || x == 1f64 {
61        return x;
62    }
63    let switch = SWITCH as f64;
64    if a > switch && b > switch {
65        return betaiapprox(a, b, x);
66    }
67    let bt = (ln_gamma(a + b) - ln_gamma(a) - ln_gamma(b) + a * x.ln() + b * (1f64 - x).ln()).exp();
68    if x < (a + 1f64) / (a + b * 2f64) {
69        bt * betacf(a, b, x) / a
70    } else {
71        1f64 - bt * betacf(b, a, 1f64 - x) / b
72    }
73}
74
75/// Continued fraction beta
76fn betacf(a: f64, b: f64, x: f64) -> f64 {
77    let qab = a + b;
78    let qap = a + 1f64;
79    let qam = a - 1f64;
80    let mut c = 1f64;
81    let mut d = 1f64 - qab * x / qap;
82    if d.abs() < FPMIN {
83        d = FPMIN;
84    }
85    d = 1f64 / d;
86    let mut h = d;
87    for m in 1..10000 {
88        let m = m as f64;
89        let m2 = 2f64 * m;
90        let mut aa = m * (b - m) * x / ((qam + m2) * (a + m2));
91        d = 1f64 + aa * d;
92        if d.abs() < FPMIN {
93            d = FPMIN;
94        }
95        c = 1f64 + aa / c;
96        if c.abs() < FPMIN {
97            c = FPMIN;
98        }
99        d = 1f64 / d;
100        h *= d * c;
101        aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2));
102        d = 1f64 + aa * d;
103        if d.abs() < FPMIN {
104            d = FPMIN;
105        }
106        c = 1f64 + aa / c;
107        if c.abs() < FPMIN {
108            c = FPMIN;
109        }
110        d = 1f64 / d;
111        let del = d * c;
112        h *= del;
113        if (del - 1f64).abs() <= EPS {
114            break;
115        }
116    }
117    h
118}
119
120/// Incomplete beta by Gauss Legendre quadrature
121fn betaiapprox(a: f64, b: f64, x: f64) -> f64 {
122    let a1 = a - 1f64;
123    let b1 = b - 1f64;
124    let mu = a / (a + b);
125    let lnmu = mu.ln();
126    let lnmuc = (1f64 - mu).ln();
127    let mut t = (a * b / ((a + b).powi(2) * (a + b + 1f64))).sqrt();
128    let xu = if x > a / (a + b) {
129        if x >= 1f64 {
130            return 1f64;
131        }
132        1f64.min((mu + 10f64 * t).max(x + 5f64 * t))
133    } else {
134        if x <= 0f64 {
135            return 0f64;
136        }
137        0f64.max((mu - 10f64 * t).min(x - 5f64 * t))
138    };
139    let mut sum = 0f64;
140    for j in 0..18 {
141        t = x + (xu - x) * Y[j];
142        sum += W[j] * (a1 * (t.ln() - lnmu) + b1 * (1f64 - t).ln() - lnmuc).exp();
143    }
144    let ans = sum
145        * (xu - x)
146        * (a1 * lnmu - ln_gamma(a) + b1 * lnmuc - ln_gamma(b) + ln_gamma(a + b)).exp();
147    if ans > 0f64 {
148        1f64 - ans
149    } else {
150        -ans
151    }
152}
153
154// =============================================================================
155// Inverse of Incomplete Beta function
156// =============================================================================
157/// Calculates the inverse of the regularized incomplete beta function.
158///
159/// This function finds $x$ such that:
160///
161/// $$ p = I_x(a,b) = \frac{1}{B(a,b)} \int_0^x t^{a-1} (1-t)^{b-1} dt $$
162///
163/// # Arguments
164///
165/// * `p` - The probability value (between 0 and 1)
166/// * `a` - First shape parameter
167/// * `b` - Second shape parameter
168///
169/// # Returns
170///
171/// The value of $x$ for which $I_x(a,b) = p$
172pub fn invbetai(p: f64, a: f64, b: f64) -> f64 {
173    let a1 = a - 1f64;
174    let b1 = b - 1f64;
175    let mut t: f64;
176    let mut x: f64;
177    let mut u: f64;
178    if p <= 0f64 {
179        return 0f64;
180    } else if p >= 1f64 {
181        return 1f64;
182    } else if a >= 1f64 && b >= 1f64 {
183        let pp = if p < 0.5 { p } else { 1f64 - p };
184        t = (-2f64 * pp.ln()).sqrt();
185        x = (2.30753 + t * 0.27061) / (1f64 + t * (0.99229 + t * 0.04481)) - t;
186        if p < 0.5 {
187            x = -x;
188        }
189        let al = (x.powi(2) - 3f64) / 6f64;
190        let h = 2f64 / (1f64 / (2f64 * a - 1f64) + 1f64 / (2f64 * b - 1f64));
191        let w = (x * (al + h).sqrt() / h)
192            - (1f64 / (2f64 * b - 1f64) - 1f64 / (2f64 * a - 1f64))
193                * (al + 5f64 / 6f64 - 2f64 / (3f64 * h));
194        x = a / (a + b * (2f64 * w).exp());
195    } else {
196        let lna = (a / (a + b)).ln();
197        let lnb = (b / (a + b)).ln();
198        t = (a * lna).exp() / a;
199        u = (b * lnb).exp() / b;
200        let w = t + u;
201        x = if p < t / w {
202            (a * w * p).powf(1f64 / a)
203        } else {
204            1f64 - (b * w * (1f64 - p)).powf(1f64 / b)
205        };
206    }
207    let afac = -ln_gamma(a) - ln_gamma(b) + ln_gamma(a + b);
208    for j in 0..10 {
209        if x == 0f64 || x == 1f64 {
210            return x;
211        }
212        let err = betai(a, b, x) - p;
213        t = (a1 * x.ln() + b1 * (1f64 - x).ln() + afac).exp();
214        u = err / t;
215        t = u / (1f64 - 0.5 * 1f64.min(u * (a1 / x - b1 / (1f64 - x))));
216        x -= t;
217        if x <= 0f64 {
218            x = 0.5 * (x + t);
219        }
220        if x >= 1f64 {
221            x = 0.5 * (x + t + 1f64);
222        }
223        if t.abs() < EPS * x && j > 0 {
224            break;
225        }
226    }
227    x
228}
puruspe/beta.rs

puruspe/
beta.rs
