Skip to main content

zenith_float_num/
dist.rs

1//! Real distribution PDF / CDF / PMF kernels on [`ExactNum`].
2
3use crate::defs::WORD_BIT_SIZE;
4use crate::Consts;
5use crate::Error;
6use crate::ExactNum;
7use crate::RoundingMode;
8use crate::NAN;
9
10fn dist_nan() -> ExactNum {
11    ExactNum::nan(Some(Error::InvalidArgument))
12}
13
14fn work_p(p: usize) -> usize {
15    p.saturating_add(WORD_BIT_SIZE)
16}
17
18fn finite_pos(x: &ExactNum) -> bool {
19    x.is_positive() && !x.is_inf()
20}
21
22fn nn_int(x: &ExactNum) -> bool {
23    x.is_int() && !x.is_negative()
24}
25
26impl ExactNum {
27    /// Standard-form normal density
28    /// \(\varphi(x;\mu,\sigma)=\exp(-(x-\mu)^2/(2\sigma^2))/(\sigma\sqrt{2\pi})\).
29    ///
30    /// `sigma ≤ 0` or a non-finite parameter is `NaN`.
31    ///
32    /// # Precision
33    ///
34    /// - Algorithm: `exp` / `sqrt` / `π` at `p + WORD_BIT_SIZE`, then one rounding to `p`.
35    /// - Bound: working precision `p + WORD_BIT_SIZE`.
36    /// - MPFR oracle: no (composite of existing leaves).
37    pub fn normal_pdf(
38        &self,
39        mu: &Self,
40        sigma: &Self,
41        p: usize,
42        rm: RoundingMode,
43        cc: &mut Consts,
44    ) -> Self {
45        if self.is_nan() || mu.is_nan() || sigma.is_nan() {
46            return NAN;
47        }
48        if !finite_pos(sigma) {
49            return dist_nan();
50        }
51        if self.is_inf() {
52            return ExactNum::from_u8(0, p);
53        }
54        let pw = work_p(p);
55        let none = RoundingMode::None;
56        let z = self.sub(mu, pw, none);
57        let two = ExactNum::from_u8(2, pw);
58        let sig2 = sigma.mul(sigma, pw, none);
59        let expo = z
60            .mul(&z, pw, none)
61            .div(&two.mul(&sig2, pw, none), pw, none)
62            .neg();
63        let num = expo.exp(pw, none, cc);
64        let two_pi = two.mul(&cc.pi(pw, none), pw, none);
65        let den = sigma.mul(&two_pi.sqrt(pw, none), pw, none);
66        num.div(&den, p, rm)
67    }
68
69    /// Normal CDF \(\Phi(x;\mu,\sigma)=(1+\mathrm{erf}((x-\mu)/(\sigma\sqrt{2})))/2\).
70    ///
71    /// `sigma ≤ 0` is `NaN`.
72    ///
73    /// # Precision
74    ///
75    /// - Algorithm: existing `erf` at `p + WORD_BIT_SIZE`.
76    /// - Bound: working precision `p + WORD_BIT_SIZE`.
77    /// - MPFR oracle: no (composite of `erf`).
78    pub fn normal_cdf(
79        &self,
80        mu: &Self,
81        sigma: &Self,
82        p: usize,
83        rm: RoundingMode,
84        cc: &mut Consts,
85    ) -> Self {
86        if self.is_nan() || mu.is_nan() || sigma.is_nan() {
87            return NAN;
88        }
89        if !finite_pos(sigma) {
90            return dist_nan();
91        }
92        let pw = work_p(p);
93        let none = RoundingMode::None;
94        let two = ExactNum::from_u8(2, pw);
95        let z = self
96            .sub(mu, pw, none)
97            .div(&sigma.mul(&two.sqrt(pw, none), pw, none), pw, none);
98        let one = ExactNum::from_u8(1, pw);
99        one.add(&z.erf(pw, none, cc), pw, none).div(&two, p, rm)
100    }
101
102    /// Gamma density on the scale parameterization
103    /// \(x^{\alpha-1}e^{-x/\beta}/(\beta^\alpha\Gamma(\alpha))\).
104    ///
105    /// Requires `x ≥ 0`, `alpha > 0`, `beta > 0`.
106    ///
107    /// # Precision
108    ///
109    /// - Algorithm: `pow` / `exp` / `gamma` at `p + WORD_BIT_SIZE`.
110    /// - Bound: working precision `p + WORD_BIT_SIZE`.
111    /// - MPFR oracle: no.
112    pub fn gamma_pdf(
113        &self,
114        alpha: &Self,
115        beta: &Self,
116        p: usize,
117        rm: RoundingMode,
118        cc: &mut Consts,
119    ) -> Self {
120        if self.is_nan() || alpha.is_nan() || beta.is_nan() {
121            return NAN;
122        }
123        if self.is_negative() || !finite_pos(alpha) || !finite_pos(beta) {
124            return dist_nan();
125        }
126        let pw = work_p(p);
127        let none = RoundingMode::None;
128        let one = ExactNum::from_u8(1, pw);
129        let am1 = alpha.sub(&one, pw, none);
130        let xb = self.div(beta, pw, none);
131        let num = self
132            .pow(&am1, pw, none, cc)
133            .mul(&xb.neg().exp(pw, none, cc), pw, none);
134        let den = beta
135            .pow(alpha, pw, none, cc)
136            .mul(&alpha.gamma(pw, none, cc), pw, none);
137        num.div(&den, p, rm)
138    }
139
140    /// Beta density \(x^{\alpha-1}(1-x)^{\beta-1}/B(\alpha,\beta)\) with
141    /// \(B(\alpha,\beta)=\Gamma(\alpha)\Gamma(\beta)/\Gamma(\alpha+\beta)\).
142    ///
143    /// Requires `x ∈ [0, 1]`, `alpha > 0`, `beta > 0`.
144    ///
145    /// # Precision
146    ///
147    /// - Algorithm: `pow` / `gamma` at `p + WORD_BIT_SIZE`.
148    /// - Bound: working precision `p + WORD_BIT_SIZE`.
149    /// - MPFR oracle: no.
150    pub fn beta_pdf(
151        &self,
152        alpha: &Self,
153        beta: &Self,
154        p: usize,
155        rm: RoundingMode,
156        cc: &mut Consts,
157    ) -> Self {
158        if self.is_nan() || alpha.is_nan() || beta.is_nan() {
159            return NAN;
160        }
161        if self.is_negative()
162            || matches!(self.cmp(&ExactNum::from_u8(1, p)), Some(c) if c > 0)
163            || !finite_pos(alpha)
164            || !finite_pos(beta)
165        {
166            return dist_nan();
167        }
168        let pw = work_p(p);
169        let none = RoundingMode::None;
170        let one = ExactNum::from_u8(1, pw);
171        let am1 = alpha.sub(&one, pw, none);
172        let bm1 = beta.sub(&one, pw, none);
173        let num = self.pow(&am1, pw, none, cc).mul(
174            &one.sub(self, pw, none).pow(&bm1, pw, none, cc),
175            pw,
176            none,
177        );
178        let bfn = alpha
179            .gamma(pw, none, cc)
180            .mul(&beta.gamma(pw, none, cc), pw, none)
181            .div(&alpha.add(beta, pw, none).gamma(pw, none, cc), pw, none);
182        num.div(&bfn, p, rm)
183    }
184
185    /// Poisson PMF \(\lambda^k e^{-\lambda}/k!\) for a non-negative integer `self` \(= k\).
186    ///
187    /// # Precision
188    ///
189    /// - Algorithm: `pow` / `exp` / \(\Gamma(k+1)\) at `p + WORD_BIT_SIZE`.
190    /// - Bound: working precision `p + WORD_BIT_SIZE`.
191    /// - MPFR oracle: no.
192    pub fn poisson_pmf(&self, lambda: &Self, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
193        if self.is_nan() || lambda.is_nan() {
194            return NAN;
195        }
196        if !nn_int(self) || lambda.is_negative() {
197            return dist_nan();
198        }
199        let pw = work_p(p);
200        let none = RoundingMode::None;
201        let one = ExactNum::from_u8(1, pw);
202        let kf = self.add(&one, pw, none).gamma(pw, none, cc);
203        lambda
204            .pow(self, pw, none, cc)
205            .mul(&lambda.neg().exp(pw, none, cc), pw, none)
206            .div(&kf, p, rm)
207    }
208
209    /// Binomial PMF \(\binom{n}{k} \mathrm{prob}^k (1-\mathrm{prob})^{n-k}\).
210    ///
211    /// `self` is \(k\). Requires non-negative integers `k ≤ n` and `prob ∈ [0, 1]`.
212    ///
213    /// # Precision
214    ///
215    /// - Algorithm: multiplicative binomial coefficient, then `pow`.
216    /// - Bound: working precision `p + WORD_BIT_SIZE`.
217    /// - MPFR oracle: no.
218    pub fn binomial_pmf(
219        &self,
220        n: &Self,
221        prob: &Self,
222        p: usize,
223        rm: RoundingMode,
224        cc: &mut Consts,
225    ) -> Self {
226        if self.is_nan() || n.is_nan() || prob.is_nan() {
227            return NAN;
228        }
229        if !nn_int(self)
230            || !nn_int(n)
231            || matches!(self.cmp(n), Some(c) if c > 0)
232            || prob.is_negative()
233            || matches!(prob.cmp(&ExactNum::from_u8(1, p)), Some(c) if c > 0)
234        {
235            return dist_nan();
236        }
237        let pw = work_p(p);
238        let none = RoundingMode::None;
239        let c = binom_mul(n, self, pw);
240        let q = ExactNum::from_u8(1, pw).sub(prob, pw, none);
241        let nmk = n.sub(self, pw, none);
242        c.mul(&prob.pow(self, pw, none, cc), pw, none)
243            .mul(&q.pow(&nmk, pw, none, cc), pw, none)
244            .set_prec_val(p, rm)
245    }
246
247    /// Chi-squared CDF \(P(k/2, x/2)=\gamma(k/2, x/2)/\Gamma(k/2)\).
248    ///
249    /// `self` is \(x\). Requires `x ≥ 0` and `k > 0`.
250    ///
251    /// # Precision
252    ///
253    /// - Algorithm: lower `gammainc` over `gamma` at `p + WORD_BIT_SIZE`.
254    /// - Bound: working precision `p + WORD_BIT_SIZE`.
255    /// - MPFR oracle: no.
256    pub fn chi_squared_cdf(&self, k: &Self, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
257        if self.is_nan() || k.is_nan() {
258            return NAN;
259        }
260        if self.is_negative() || !finite_pos(k) {
261            return dist_nan();
262        }
263        let pw = work_p(p);
264        let none = RoundingMode::None;
265        let two = ExactNum::from_u8(2, pw);
266        let s = k.div(&two, pw, none);
267        let xh = self.div(&two, pw, none);
268        s.gammainc(&xh, pw, none, cc)
269            .div(&s.gamma(pw, none, cc), p, rm)
270    }
271
272    /// Student-\(t\) density via \(\Gamma\):
273    /// \(\Gamma((\nu+1)/2)/(\sqrt{\nu\pi}\,\Gamma(\nu/2))\,(1+x^2/\nu)^{-(\nu+1)/2}\).
274    ///
275    /// `self` is \(x\). Requires `nu > 0`.
276    ///
277    /// # Precision
278    ///
279    /// - Algorithm: `gamma` / `pow` / `sqrt` at `p + WORD_BIT_SIZE`.
280    /// - Bound: working precision `p + WORD_BIT_SIZE`.
281    /// - MPFR oracle: no.
282    pub fn student_t_pdf(&self, nu: &Self, p: usize, rm: RoundingMode, cc: &mut Consts) -> Self {
283        if self.is_nan() || nu.is_nan() {
284            return NAN;
285        }
286        if !finite_pos(nu) {
287            return dist_nan();
288        }
289        let pw = work_p(p);
290        let none = RoundingMode::None;
291        let one = ExactNum::from_u8(1, pw);
292        let two = ExactNum::from_u8(2, pw);
293        let np1 = nu.add(&one, pw, none);
294        let half_np1 = np1.div(&two, pw, none);
295        let half_n = nu.div(&two, pw, none);
296        let pref = half_np1.gamma(pw, none, cc).div(
297            &nu.mul(&cc.pi(pw, none), pw, none).sqrt(pw, none).mul(
298                &half_n.gamma(pw, none, cc),
299                pw,
300                none,
301            ),
302            pw,
303            none,
304        );
305        let body = one
306            .add(&self.mul(self, pw, none).div(nu, pw, none), pw, none)
307            .pow(&half_np1.neg(), pw, none, cc);
308        pref.mul(&body, p, rm)
309    }
310}
311
312fn binom_mul(n: &ExactNum, k: &ExactNum, pw: usize) -> ExactNum {
313    let none = RoundingMode::None;
314    let one = ExactNum::from_u8(1, pw);
315    if k.is_zero() {
316        return one;
317    }
318    let mut i = one.clone();
319    let mut c = one.clone();
320    let nmk = n.sub(k, pw, none);
321    loop {
322        let term = nmk.add(&i, pw, none);
323        c = c.mul(&term, pw, none).div(&i, pw, none);
324        if i.cmp(k) == Some(0) {
325            return c;
326        }
327        i = i.add(&one, pw, none);
328        if i.cmp(k) == Some(1) {
329            return c;
330        }
331    }
332}
333
334impl ExactNum {
335    fn set_prec_val(mut self, p: usize, rm: RoundingMode) -> Self {
336        let _ = self.set_precision(p, rm);
337        self
338    }
339}
340
341#[cfg(test)]
342mod tests {
343    use super::*;
344
345    #[test]
346    fn dist_plan_golds() {
347        let p = 256;
348        let rm = RoundingMode::ToEven;
349        let mut cc = Consts::new().unwrap();
350        let zero = ExactNum::from_u8(0, p);
351        let one = ExactNum::from_u8(1, p);
352        let two = ExactNum::from_u8(2, p);
353
354        let np = zero.normal_pdf(&zero, &one, p, rm, &mut cc);
355        let two_pi = two.mul(&cc.pi(p, rm), p, rm);
356        let want_np = two_pi.sqrt(p, rm).reciprocal(p, rm);
357        assert_eq!(np.cmp(&want_np), Some(0));
358
359        let nc = zero.normal_cdf(&zero, &one, p, rm, &mut cc);
360        let half = one.div(&two, p, rm);
361        assert_eq!(nc.cmp(&half), Some(0));
362
363        let gp = one.gamma_pdf(&one, &one, p, rm, &mut cc);
364        let em1 = one.neg().exp(p, rm, &mut cc);
365        assert_eq!(gp.cmp(&em1), Some(0));
366
367        let po = zero.poisson_pmf(&one, p, rm, &mut cc);
368        assert_eq!(po.cmp(&em1), Some(0));
369
370        // χ²(2) CDF at the 95% table quantile −2 ln(0.05) = 2 ln 20 is exactly 19/20.
371        let twenty = ExactNum::from_u8(20, p);
372        let x95 = two.mul(&twenty.ln(p, rm, &mut cc), p, rm);
373        let chi = x95.chi_squared_cdf(&two, p, rm, &mut cc);
374        let table = ExactNum::from_u8(19, p).div(&ExactNum::from_u8(20, p), p, rm);
375        assert_eq!(chi.cmp(&table), Some(0));
376
377        assert!(one.normal_pdf(&zero, &zero, p, rm, &mut cc).is_nan());
378        assert!(one.neg().poisson_pmf(&one, p, rm, &mut cc).is_nan());
379    }
380}