1use 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 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 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 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 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 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 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 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 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 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}