1use crate::{
3 Arg, Atan, Erfc, Erfcinv, Float, Ln, Log1p, Unary, betainc, betaincinv, cos, erfcinv, exp, expm1, gammainc, gammaincc, gammaincinv, lgamma, ln, log1p, pow,
4 random::Items, sin, slice::{CHUNK, kept, written}
5};
6use rand::{Rng, distr::{Distribution as _, Open01, StandardUniform}};
7use std::{f64::consts::{FRAC_1_SQRT_2, PI, SQRT_2, TAU}, mem::MaybeUninit};
8
9const SQRT_TAU: f64 = 2.5066282746310002;
11
12#[derive(Clone, Copy, Debug, PartialEq)]
14pub struct Distribution(Law);
15
16#[derive(Clone, Copy, Debug, PartialEq, Eq)]
18pub struct Invalid(pub &'static str);
19
20impl std::fmt::Display for Invalid { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.write_str(self.0) } }
21
22impl std::error::Error for Invalid {}
23
24#[derive(Clone, Copy, Debug, PartialEq)]
25enum Law {
26 Normal(Normal), Uniform(Uniform), Beta(Beta), Cauchy(Cauchy), Exponential(Exponential), Fisher(Fisher), Gamma(Gamma), InverseGamma(InverseGamma),
27 Laplace(Laplace), Logistic(Logistic), LogNormal(LogNormal), StudentT(StudentT), Weibull(Weibull), Binomial(Binomial), Poisson(Poisson)
28}
29
30macro_rules! each_law {
32 ($law:expr, $d:ident => $body:expr) => {
33 match $law {
34 Law::Normal($d) => $body,
35 Law::Uniform($d) => $body,
36 Law::Beta($d) => $body,
37 Law::Cauchy($d) => $body,
38 Law::Exponential($d) => $body,
39 Law::Fisher($d) => $body,
40 Law::Gamma($d) => $body,
41 Law::InverseGamma($d) => $body,
42 Law::Laplace($d) => $body,
43 Law::Logistic($d) => $body,
44 Law::LogNormal($d) => $body,
45 Law::StudentT($d) => $body,
46 Law::Weibull($d) => $body,
47 Law::Binomial($d) => $body,
48 Law::Poisson($d) => $body,
49 }
50 };
51}
52
53trait Formulas {
55 fn pdf(&self, x: f64) -> f64;
56 fn cdf(&self, x: f64) -> f64;
57 fn quantile(&self, p: f64) -> f64;
58 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 { self.quantile(open(rng)) }
59 fn words(&self) -> usize { 1 }
61 fn samples(&self, items: &mut Items, out: &mut [f64]) { for o in out { *o = self.sample(&mut items.item()); } }
62 #[inline(always)]
63 fn pdf_map<F: Float>(&self) -> impl Fn(F) -> F + Copy { move |x| narrowed(x, |x| self.pdf(x)) }
64 #[inline(always)]
65 fn cdf_map<F: Float>(&self) -> impl Fn(F) -> F + Copy { move |x| narrowed(x, |x| self.cdf(x)) }
66 #[inline(always)]
67 fn quantile_map<F: Float>(&self) -> impl Fn(F) -> F + Copy { move |p| narrowed(p, |p| self.quantile(p)) }
68 #[inline(always)]
69 fn pdf_many<F: Float>(&self, x: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool where Ln: Unary<F, Out = F> {
70 let pdf = self.pdf_map();
71 each!([x], n, out, |[x]| pdf(x))
72 }
73 #[inline(always)]
74 fn cdf_many<F: Float>(&self, x: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool where Erfc: Unary<F, Out = F>, Atan: Unary<F, Out = F> {
75 let cdf = self.cdf_map();
76 each!([x], n, out, |[x]| cdf(x))
77 }
78 #[inline(always)]
79 fn quantile_many<F: Float>(&self, p: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool where Erfcinv: Unary<F, Out = F>, Log1p: Unary<F, Out = F>,
80 Ln: Unary<F, Out = F> { let q = self.quantile_map(); each!([p], n, out, |[p]| q(p)) }
81}
82
83macro_rules! through_slices {
85 () => {
86 fn pdf(&self, x: f64) -> f64 {
87 one_item(|out| self.pdf_many(Arg::One(x), 1, out))
88 }
89 fn cdf(&self, x: f64) -> f64 {
90 one_item(|out| self.cdf_many(Arg::One(x), 1, out))
91 }
92 fn quantile(&self, p: f64) -> f64 {
93 one_item(|out| self.quantile_many(Arg::One(p), 1, out))
94 }
95 };
96}
97
98#[inline(always)]
100fn one_item<F: Float>(many: impl FnOnce(&mut [MaybeUninit<F>]) -> bool) -> F {
101 let mut out = [MaybeUninit::uninit()];
102 many(&mut out);
103 (unsafe { written(&mut out) })[0]
104}
105
106#[inline(always)]
108fn narrowed<F: Float>(x: F, f: impl Fn(f64) -> f64) -> F { F::narrow(f(x.wide())) }
109
110#[derive(Clone, Copy, Debug, PartialEq)]
112pub struct Pdf(pub Distribution);
113forms!(Pdf at real: |k, [x]| one_item(|out| k.slice([Arg::One(x)], 1, out)), many |k, [x], n, out| each_law!(&k.0 .0, d => d.pdf_many(x, n, out)));
114
115#[derive(Clone, Copy, Debug, PartialEq)]
117pub struct Cdf(pub Distribution);
118forms!(Cdf at real: |k, [x]| one_item(|out| k.slice([Arg::One(x)], 1, out)), many |k, [x], n, out| each_law!(&k.0 .0, d => d.cdf_many(x, n, out)));
119
120#[derive(Clone, Copy, Debug, PartialEq)]
122pub struct Quantile(pub Distribution);
123forms!(Quantile at real: |k, [p]| one_item(|out| k.slice([Arg::One(p)], 1, out)), many |k, [p], n, out| quantiles(p, n, out, |p, n, out| each_law!(&k.0 .0, d => d.quantile_many(p, n, out))));
124
125macro_rules! native_law {
126 ($( $object:ident $(valid $valid:ident)? ),+) => {$(
127 impl crate::facts::Formula<1> for $object {
128 fn supports(&self, dtype: crate::buffer::DType) -> bool { matches!(dtype, crate::buffer::DType::Float(_)) }
129 const FAILS_NAN: bool = true;
130 $(fn emit_valid<F: Float, O: crate::recipe::Ops<F> + ?Sized>(&self, _: crate::domain::Variant, o: &O, [p]: [O::V; 1], _: O::V) -> Option<(crate::error::Kind, O::B)> { Some((crate::error::Kind::Domain, $valid(o, p))) })?
131 fn real<F: crate::float::Native>(&self, _: &crate::recipe::NativeOps<{ crate::FUSED }>, [x]: [F; 1]) -> Option<F> {
132 Some(F::narrow(if F::WIDE { self.one(x.wide()) } else { f64::from(self.one(x.wide() as f32)) }))
133 }
134 fn native_slices(&self) -> Option<crate::facts::NativeSlices<Self, 1>> { Some(crate::facts::NativeSlices::unary()) }
135 }
136 )+};
137}
138native_law!(Pdf, Cdf, Quantile valid probability);
139
140fn probability<F: Float, O: crate::recipe::Ops<F> + ?Sized>(o: &O, p: O::V) -> O::B {
141 o.and(o.less_equal(o.number(F::zero()), p), o.less_equal(p, o.number(F::one())))
142}
143
144fn quantiles<F: Float>(p: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>], law: impl Fn(Arg<'_, F>, usize, &mut [MaybeUninit<F>]) -> bool) -> bool {
146 let inside = |p: F| probability(&crate::recipe::Floats, p);
147 let all = match p { Arg::Many(p) => p[..n].iter().fold(true, |all, &p| all & inside(p)), Arg::One(p) => inside(p) };
148 if all { return law(p, n, out); }
149 let within: Vec<F> = (0..n).map(|i| if inside(p.nth(i)) { p.nth(i) } else { F::narrow(0.5) }).collect();
150 let ok = law(Arg::Many(&within), n, out);
151 (0..n).filter(|&i| !inside(p.nth(i))).fold(ok, |_, i| { out[i].write(F::nan()); false })
152}
153
154fn composed<F: Float>(
156 x: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>], inner: impl Fn(F) -> F, values: impl Fn(&mut [F], &mut [MaybeUninit<F>]), finish: impl Fn(F, F) -> F
157) -> bool {
158 let one;
159 let (x, copies) = match x { Arg::Many(x) => (&x[..n], 0), Arg::One(v) => { one = [v]; (&one[..n.min(1)], n.saturating_sub(1)) } };
160 let mut args = Vec::with_capacity(x.len().min(CHUNK));
161 let mut ok = true;
162 for (x, out) in x.chunks(CHUNK).zip(out.chunks_mut(CHUNK)) {
163 args.clear();
164 args.extend(x.iter().map(|&v| inner(v)));
165 values(&mut args, &mut out[..x.len()]);
166 let results = unsafe { written(&mut out[..x.len()]) };
167 let mut clean = true;
168 for (r, &v) in results.iter_mut().zip(x) {
170 *r = finish(v, *r);
171 clean &= !r.is_nan();
172 }
173 ok &= clean || results.iter().zip(x).all(|(&r, &v)| kept(r, [v]));
174 }
175 if copies > 0 {
176 let r = unsafe { written(&mut out[..1]) }[0];
177 out[1..n].fill(MaybeUninit::new(r));
178 }
179 ok
180}
181
182fn through<F: Float, G: Unary<F, Out = F>>(g: G) -> impl Fn(&mut [F], &mut [MaybeUninit<F>]) { move |a, v| { g.slice([Arg::Many(a)], a.len(), v); } }
184
185fn raised<F: Float>(k: F) -> impl Fn(&mut [F], &mut [MaybeUninit<F>]) {
187 let e = crate::pow::Exponent::of(k.wide(), crate::FUSED);
188 move |a, v| crate::pow::powers(e, k, a, v)
189}
190
191impl Distribution {
192 pub fn pdf(&self, x: f64) -> f64 { Pdf(*self).one(x) }
194 pub fn cdf(&self, x: f64) -> f64 { Cdf(*self).one(x) }
196 pub fn quantile(&self, p: f64) -> f64 { Quantile(*self).one(p) }
198 pub fn words(&self) -> usize { each_law!(&self.0, d => d.words()) }
200 pub fn samples(&self, items: &mut Items, out: &mut [f64]) { each_law!(&self.0, d => d.samples(items, out)) }
202 pub fn is_discrete(&self) -> bool { matches!(self.0, Law::Binomial(_) | Law::Poisson(_)) }
204}
205
206fn checked(parameters: &[f64], ok: bool, message: &'static str) -> Result<(), Invalid> {
208 if !parameters.iter().all(|p| p.is_finite()) { Err(Invalid("distribution parameters must be finite")) } else if ok { Ok(()) } else { Err(Invalid(message)) }
209}
210
211const WHOLE: u64 = 1 << 53;
213
214impl Distribution {
215 pub fn normal(mean: f64, sd: f64) -> Result<Self, Invalid> {
217 checked(&[mean, sd], sd > 0.0, "the standard deviation must be positive")?;
218 Ok(Self(Law::Normal(Normal { mean, sd })))
219 }
220 pub fn uniform(min: f64, max: f64) -> Result<Self, Invalid> {
222 checked(&[min, max, max - min], min < max, "the minimum must be below the maximum")?;
223 Ok(Self(Law::Uniform(Uniform { min, max })))
224 }
225 pub fn beta(a: f64, b: f64) -> Result<Self, Invalid> {
227 checked(&[a, b], a > 0.0 && b > 0.0, "the shapes must be positive")?;
228 Ok(Self(Law::Beta(Beta::new(a, b))))
229 }
230 pub fn cauchy(location: f64, scale: f64) -> Result<Self, Invalid> {
232 checked(&[location, scale], scale > 0.0, "the scale must be positive")?;
233 Ok(Self(Law::Cauchy(Cauchy { location, scale })))
234 }
235 pub fn chi_squared(df: f64) -> Result<Self, Invalid> {
237 checked(&[df], df > 0.0, "the degrees of freedom must be positive")?;
238 Ok(Self(Law::Gamma(Gamma::new(df / 2.0, 2.0))))
239 }
240 pub fn exponential(rate: f64) -> Result<Self, Invalid> {
242 checked(&[rate], rate > 0.0, "the rate must be positive")?;
243 Ok(Self(Law::Exponential(Exponential { rate })))
244 }
245 pub fn fisher(df1: f64, df2: f64) -> Result<Self, Invalid> {
247 checked(&[df1, df2], df1 > 0.0 && df2 > 0.0, "the degrees of freedom must be positive")?;
248 Ok(Self(Law::Fisher(Fisher { df1, df2, beta: Beta::new(df1 / 2.0, df2 / 2.0) })))
249 }
250 pub fn gamma(shape: f64, scale: f64) -> Result<Self, Invalid> {
252 checked(&[shape, scale, 1.0 / scale], shape > 0.0 && scale > 0.0, "the shape and scale must be positive")?;
253 Ok(Self(Law::Gamma(Gamma::new(shape, scale))))
254 }
255 pub fn inverse_gamma(shape: f64, scale: f64) -> Result<Self, Invalid> {
257 checked(&[shape, scale], shape > 0.0 && scale > 0.0, "the shape and scale must be positive")?;
258 Ok(Self(Law::InverseGamma(InverseGamma { shape, scale, ln_gamma: lgamma(shape) })))
259 }
260 pub fn laplace(location: f64, scale: f64) -> Result<Self, Invalid> {
262 checked(&[location, scale], scale > 0.0, "the scale must be positive")?;
263 Ok(Self(Law::Laplace(Laplace { location, scale })))
264 }
265 pub fn logistic(location: f64, scale: f64) -> Result<Self, Invalid> {
267 checked(&[location, scale], scale > 0.0, "the scale must be positive")?;
268 Ok(Self(Law::Logistic(Logistic { location, scale })))
269 }
270 pub fn log_normal(location: f64, scale: f64) -> Result<Self, Invalid> {
272 checked(&[location, scale], scale > 0.0, "the scale must be positive")?;
273 Ok(Self(Law::LogNormal(LogNormal { location, scale })))
274 }
275 pub fn student_t(df: f64) -> Result<Self, Invalid> {
277 checked(&[df], df > 0.0, "the degrees of freedom must be positive")?;
278 let factor = exp(lgamma((df + 1.0) / 2.0) - lgamma(df / 2.0)) / (df * PI).sqrt();
279 Ok(Self(Law::StudentT(StudentT { df, factor })))
280 }
281 pub fn weibull(shape: f64, scale: f64) -> Result<Self, Invalid> {
283 checked(&[shape, scale], shape > 0.0 && scale > 0.0, "the shape and scale must be positive")?;
284 Ok(Self(Law::Weibull(Weibull { shape, scale })))
285 }
286 pub fn binomial(n: u64, p: f64) -> Result<Self, Invalid> {
288 checked(&[p], (0.0..=1.0).contains(&p), "the probability must be from 0 to 1")?;
289 if n > WHOLE { return Err(Invalid("binomial trials exceed the sampler's whole-number range")); }
290 Ok(Self(Law::Binomial(Binomial { n: n as f64, p })))
291 }
292 pub fn poisson(rate: f64) -> Result<Self, Invalid> {
294 checked(&[rate], rate >= 0.0, "the rate must not be negative")?;
295 if rate > WHOLE as f64 { return Err(Invalid("the Poisson rate exceeds the sampler's whole-number range")); }
296 Ok(Self(Law::Poisson(Poisson { rate })))
297 }
298}
299
300pub mod standard {
302 pub const NORMAL: [f64; 2] = [0.0, 1.0];
303 pub const UNIFORM: [f64; 2] = [0.0, 1.0];
304 pub const CAUCHY: [f64; 2] = [0.0, 1.0];
305 pub const EXPONENTIAL: [f64; 1] = [1.0];
306 pub const LAPLACE: [f64; 2] = [0.0, 1.0];
307 pub const LOGISTIC: [f64; 2] = [0.0, 1.0];
308 pub const LOG_NORMAL: [f64; 2] = [0.0, 1.0];
309}
310
311#[inline(always)]
313fn times<F: Float>(c: F, l: F) -> F { if c == F::zero() { F::zero() } else { c * l } }
314
315fn open<R: Rng + ?Sized>(rng: &mut R) -> f64 { Open01.sample(rng) }
317
318fn standard_normal<R: Rng + ?Sized>(rng: &mut R) -> f64 {
320 let (u, v): (f64, f64) = (open(rng), StandardUniform.sample(rng));
321 (-2.0 * ln(u)).sqrt() * cos(TAU * v)
322}
323
324fn gamma_words(shape: f64) -> usize { if shape < 1.0 { 4 } else { 3 } }
326
327fn standard_gamma<R: Rng + ?Sized>(shape: f64, rng: &mut R) -> f64 {
329 if shape < 1.0 { return standard_gamma(shape + 1.0, rng) * exp(ln(open(rng)) / shape); }
330 let d = shape - 1.0 / 3.0;
331 let c = 1.0 / (9.0 * d).sqrt();
332 loop {
333 let x = standard_normal(rng);
334 let v = 1.0 + c * x;
335 if v <= 0.0 { continue; }
336 let v = v * v * v;
337 let u = open(rng);
338 if u < 1.0 - 0.0331 * x * x * x * x || ln(u) < 0.5 * x * x + d * (1.0 - v + ln(v)) { return d * v; }
339 }
340}
341
342#[derive(Clone, Copy, Debug, PartialEq)]
343struct Normal { mean: f64, sd: f64 }
344impl Normal {
345 #[inline(always)]
347 fn erfc_line<F: Float>(&self) -> impl Fn(F) -> F + Copy + use<F> {
348 let (mean, scale) = (F::narrow(self.mean), F::narrow(FRAC_1_SQRT_2 / self.sd));
349 move |x| (mean - x) * scale
350 }
351 #[inline(always)]
353 fn from_erfcinv<F: Float>(&self) -> impl Fn(F) -> F + Copy + use<F> {
354 let (mean, step) = (F::narrow(self.mean), F::narrow(self.sd * SQRT_2));
355 move |e| mean - step * e
356 }
357 #[inline(always)]
359 fn density<F: Float>(&self) -> impl Fn(F) -> F + Copy + use<F> {
360 let (mean, rsd, half, height) = (F::narrow(self.mean), F::narrow(self.sd.recip()), F::narrow(-0.5), F::narrow((self.sd * SQRT_TAU).recip()));
361 move |x| { let z = (x - mean) * rsd; exp(half * z * z) * height }
362 }
363}
364
365impl Formulas for Normal {
366 through_slices!();
367 #[inline(always)]
368 fn pdf_map<F: Float>(&self) -> impl Fn(F) -> F + Copy { self.density() }
369 #[inline(always)]
370 fn cdf_many<F: Float>(&self, x: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool where Erfc: Unary<F, Out = F> {
371 let half = F::narrow(0.5);
372 composed(x, n, out, self.erfc_line(), through(Erfc), #[inline(always)] move |_, e| half * e)
373 }
374 #[inline(always)]
375 fn quantile_many<F: Float>(&self, p: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool where Erfcinv: Unary<F, Out = F> {
376 let (two, from) = (F::narrow(2.0), self.from_erfcinv());
377 composed(p, n, out, #[inline(always)] move |p| two * p, through(Erfcinv), #[inline(always)] move |_, e| from(e))
378 }
379 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 { self.mean + self.sd * standard_normal(rng) }
380 fn words(&self) -> usize { 2 }
382 fn samples(&self, items: &mut Items, out: &mut [f64]) {
383 for pair in out.chunks_mut(2) {
384 let mut item = items.item();
385 let (u, v): (f64, f64) = (open(&mut item), StandardUniform.sample(&mut item));
386 let r = self.sd * (-2.0 * ln(u)).sqrt();
387 pair[0] = self.mean + r * cos(TAU * v);
388 if let Some(second) = pair.get_mut(1) { *second = self.mean + r * sin(TAU * v); }
389 }
390 }
391}
392
393#[derive(Clone, Copy, Debug, PartialEq)]
394struct Uniform { min: f64, max: f64 }
395
396impl Uniform {
397 #[inline(always)]
399 fn bounds<F: Float>(&self) -> (F, F, F) { let (min, max) = (F::narrow(self.min), F::narrow(self.max)); (min, max, max - min) }
400}
401
402impl Formulas for Uniform {
403 through_slices!();
404 #[inline(always)]
405 fn pdf_map<F: Float>(&self) -> impl Fn(F) -> F + Copy {
406 let (min, max, width) = self.bounds::<F>();
407 let height = width.recip();
408 move |x| if x >= min && x <= max { height } else { F::zero() }
409 }
410 #[inline(always)]
411 fn cdf_map<F: Float>(&self) -> impl Fn(F) -> F + Copy {
412 let (min, _, width) = self.bounds::<F>();
413 move |x| {
414 let c = (x - min) / width;
415 if c < F::zero() { F::zero() } else if c > F::one() { F::one() } else { c }
416 }
417 }
418 #[inline(always)]
419 fn quantile_map<F: Float>(&self) -> impl Fn(F) -> F + Copy { let (min, _, width) = self.bounds::<F>(); move |p| min + p * width }
420 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 { let u: f64 = StandardUniform.sample(rng); self.min + u * (self.max - self.min) }
421}
422
423#[derive(Clone, Copy, Debug, PartialEq)]
425struct Beta { a: f64, b: f64, ln_beta: f64 }
426
427impl Beta { fn new(a: f64, b: f64) -> Self { Self { a, b, ln_beta: lgamma(a) + lgamma(b) - lgamma(a + b) } } }
428
429impl Formulas for Beta {
430 fn pdf(&self, x: f64) -> f64 {
431 if !(0.0..=1.0).contains(&x) { return 0.0; }
432 exp(times(self.a - 1.0, ln(x)) + times(self.b - 1.0, log1p(-x)) - self.ln_beta)
433 }
434 fn cdf(&self, x: f64) -> f64 { betainc(self.a, self.b, x.clamp(0.0, 1.0)) }
435 fn quantile(&self, p: f64) -> f64 { betaincinv(self.a, self.b, p) }
436 fn words(&self) -> usize { gamma_words(self.a) + gamma_words(self.b) }
437 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 { let x = standard_gamma(self.a, rng); x / (x + standard_gamma(self.b, rng)) }
438}
439
440#[derive(Clone, Copy, Debug, PartialEq)]
441struct Cauchy { location: f64, scale: f64 }
442
443impl Formulas for Cauchy {
444 through_slices!();
445 #[inline(always)]
446 fn pdf_map<F: Float>(&self) -> impl Fn(F) -> F + Copy {
447 let (location, rscale, height) = (F::narrow(self.location), F::narrow(self.scale.recip()), F::narrow((PI * self.scale).recip()));
448 move |x| { let z = (x - location) * rscale; height / (F::one() + z * z) }
449 }
450 #[inline(always)]
452 fn cdf_many<F: Float>(&self, x: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool where Atan: Unary<F, Out = F> {
453 let (location, rscale, rpi, half) = (F::narrow(self.location), F::narrow(self.scale.recip()), F::narrow(PI.recip()), F::narrow(0.5));
454 composed(x, n, out, move |x| {
455 let z = (x - location) * rscale;
456 let (num, den) = if z.abs() <= F::one() { (z, F::one()) } else { (-F::one(), z) };
457 num / den
458 }, through(Atan), move |x, a| {
459 let (z, a) = ((x - location) * rscale, a * rpi);
460 if z.abs() <= F::one() { half + a } else if z < F::zero() { a } else { F::one() + a }
461 })
462 }
463 #[inline(always)]
465 fn quantile_many<F: Float>(&self, p: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool {
466 let (location, scale, pi, low, high, half) =
467 (F::narrow(self.location), F::narrow(self.scale), F::narrow(PI), F::narrow(0.25), F::narrow(0.75), F::narrow(0.5));
468 composed(p, n, out, move |p| {
469 pi * if p < low { p } else if p > high { F::one() - p } else { p - half }
470 }, |a, v| {
471 each!([Arg::Many(a)], a.len(), v, |[x]| F::from_working(crate::trig::formula(
472 &crate::recipe::NativeOps::<{ crate::FUSED }>::DEFAULT, F::working(x), crate::trig::Kind::Tan
473 )));
474 }, move |p, t| {
475 let (num, den) = if p < low { (-scale, t) } else if p > high { (scale, t) } else { (scale * t, F::one()) };
476 location + num / den
477 })
478 }
479}
480
481#[derive(Clone, Copy, Debug, PartialEq)]
482struct Exponential { rate: f64 }
483
484impl Formulas for Exponential {
485 through_slices!();
486 #[inline(always)]
487 fn pdf_map<F: Float>(&self) -> impl Fn(F) -> F + Copy {
488 let rate = F::narrow(self.rate);
489 move |x| { let d = rate * exp(-rate * x); if x < F::zero() { F::zero() } else { d } }
490 }
491 #[inline(always)]
492 fn cdf_map<F: Float>(&self) -> impl Fn(F) -> F + Copy {
493 let rate = F::narrow(self.rate);
494 move |x| { let c = -expm1(-rate * x); if x <= F::zero() { F::zero() } else { c } }
495 }
496 #[inline(always)]
497 fn quantile_many<F: Float>(&self, p: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool where Log1p: Unary<F, Out = F> {
498 let mean = F::narrow(self.rate.recip());
499 composed(p, n, out, |p| -p, through(Log1p), move |_, l| -l * mean)
500 }
501 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 { -ln(open(rng)) / self.rate }
502}
503
504#[derive(Clone, Copy, Debug, PartialEq)]
506struct Fisher { df1: f64, df2: f64, beta: Beta }
507
508impl Fisher { fn y(&self, x: f64) -> f64 { self.df1 * x / (self.df1 * x + self.df2) } }
509
510impl Formulas for Fisher {
511 fn pdf(&self, x: f64) -> f64 {
512 if x < 0.0 || x.is_infinite() { return 0.0; }
513 let d = self.df1 * x + self.df2;
514 self.beta.pdf(self.y(x)) * self.df1 * self.df2 / (d * d)
515 }
516 fn cdf(&self, x: f64) -> f64 { if x <= 0.0 { 0.0 } else if x.is_infinite() { 1.0 } else { self.beta.cdf(self.y(x)) } }
517 fn quantile(&self, p: f64) -> f64 { let y = self.beta.quantile(p); self.df2 * y / (self.df1 * (1.0 - y)) }
518 fn words(&self) -> usize { gamma_words(self.df1 / 2.0) + gamma_words(self.df2 / 2.0) }
519 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
520 let x = standard_gamma(self.df1 / 2.0, rng) / self.df1;
521 x / (standard_gamma(self.df2 / 2.0, rng) / self.df2)
522 }
523}
524
525#[derive(Clone, Copy, Debug, PartialEq)]
527struct Gamma { shape: f64, scale: f64, ln_gamma: f64 }
528
529impl Gamma { fn new(shape: f64, scale: f64) -> Self { Self { shape, scale, ln_gamma: lgamma(shape) } } }
530
531impl Formulas for Gamma {
532 fn pdf(&self, x: f64) -> f64 {
533 if x < 0.0 || x.is_infinite() { return 0.0; }
534 exp(times(self.shape - 1.0, ln(x)) - x / self.scale - self.ln_gamma - self.shape * ln(self.scale))
535 }
536 fn cdf(&self, x: f64) -> f64 { if x <= 0.0 { 0.0 } else { gammainc(self.shape, x / self.scale) } }
537 fn quantile(&self, p: f64) -> f64 { self.scale * gammaincinv(self.shape, p) }
538 fn words(&self) -> usize { gamma_words(self.shape) }
539 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 { self.scale * standard_gamma(self.shape, rng) }
540}
541
542#[derive(Clone, Copy, Debug, PartialEq)]
544struct InverseGamma { shape: f64, scale: f64, ln_gamma: f64 }
545
546impl Formulas for InverseGamma {
547 fn pdf(&self, x: f64) -> f64 {
548 if x <= 0.0 { return 0.0; }
549 exp(self.shape * ln(self.scale) - self.ln_gamma - (self.shape + 1.0) * ln(x) - self.scale / x)
550 }
551 fn cdf(&self, x: f64) -> f64 { if x <= 0.0 { 0.0 } else { gammaincc(self.shape, self.scale / x) } }
552 fn quantile(&self, p: f64) -> f64 { self.scale / gammaincinv(self.shape, 1.0 - p) }
553 fn words(&self) -> usize { gamma_words(self.shape) }
554 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 { self.scale / standard_gamma(self.shape, rng) }
555}
556
557#[derive(Clone, Copy, Debug, PartialEq)]
558struct Laplace { location: f64, scale: f64 }
559
560impl Formulas for Laplace {
561 through_slices!();
562 #[inline(always)]
563 fn pdf_map<F: Float>(&self) -> impl Fn(F) -> F + Copy {
564 let (location, rscale, height) = (F::narrow(self.location), F::narrow(self.scale.recip()), F::narrow((2.0 * self.scale).recip()));
565 move |x| exp(-((x - location) * rscale).abs()) * height
566 }
567 #[inline(always)]
569 fn cdf_map<F: Float>(&self) -> impl Fn(F) -> F + Copy {
570 let (location, rscale, half) = (F::narrow(self.location), F::narrow(self.scale.recip()), F::narrow(0.5));
571 move |x| {
572 let z = (x - location) * rscale;
573 let e = half * exp(-z.abs());
574 if z < F::zero() { e } else { F::one() - e }
575 }
576 }
577 #[inline(always)]
578 fn quantile_many<F: Float>(&self, p: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool where Ln: Unary<F, Out = F> {
579 let (location, scale, half) = (F::narrow(self.location), F::narrow(self.scale), F::narrow(0.5));
580 composed(p, n, out, move |p| {
581 let r = if p < half { p } else { F::one() - p };
582 r + r
583 }, through(Ln), move |p, l| if p < half { location + scale * l } else { location - scale * l })
584 }
585}
586
587#[derive(Clone, Copy, Debug, PartialEq)]
589struct Logistic { location: f64, scale: f64 }
590
591impl Formulas for Logistic {
592 through_slices!();
593 #[inline(always)]
594 fn pdf_map<F: Float>(&self) -> impl Fn(F) -> F + Copy {
595 let (location, rscale) = (F::narrow(self.location), F::narrow(self.scale.recip()));
596 move |x| {
597 let e = exp(-((x - location) * rscale).abs());
598 let d = F::one() + e;
599 e / (d * d) * rscale
600 }
601 }
602 #[inline(always)]
603 fn cdf_map<F: Float>(&self) -> impl Fn(F) -> F + Copy {
604 let (location, rscale) = (F::narrow(self.location), F::narrow(self.scale.recip()));
605 move |x| {
606 let z = (x - location) * rscale;
607 let e = exp(-z.abs());
608 let r = (F::one() + e).recip();
609 if z < F::zero() { e * r } else { r }
610 }
611 }
612 #[inline(always)]
613 fn quantile_many<F: Float>(&self, p: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool where Ln: Unary<F, Out = F> {
614 let (location, scale) = (F::narrow(self.location), F::narrow(self.scale));
615 composed(p, n, out, |p| p / (F::one() - p), through(Ln), move |_, l| location + scale * l)
616 }
617}
618
619#[derive(Clone, Copy, Debug, PartialEq)]
620struct LogNormal { location: f64, scale: f64 }
621impl LogNormal {
622 #[inline(always)]
624 fn normal(&self) -> Normal { Normal { mean: self.location, sd: self.scale } }
625}
626
627impl Formulas for LogNormal {
628 through_slices!();
629 #[inline(always)]
631 fn pdf_many<F: Float>(&self, x: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool where Ln: Unary<F, Out = F> {
632 let pdf = self.normal().density();
633 composed(x, n, out, #[inline(always)] |x| if x > F::zero() { x } else { F::one() }, through(Ln), #[inline(always)] move |x, l| {
634 let p = pdf(l) / x;
635 if x <= F::zero() { F::zero() } else { p }
636 })
637 }
638 #[inline(always)]
639 fn cdf_many<F: Float>(&self, x: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool where Erfc: Unary<F, Out = F> {
640 let (line, half) = (self.normal().erfc_line(), F::narrow(0.5));
641 composed(
642 x, n, out, #[inline(always)] move |x| line(ln(x)), through(Erfc), #[inline(always)] move |x, e| if x <= F::zero() { F::zero() } else { half * e })
643 }
644 #[inline(always)]
645 fn quantile_many<F: Float>(&self, p: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool where Erfcinv: Unary<F, Out = F> {
646 let (two, from) = (F::narrow(2.0), self.normal().from_erfcinv());
647 composed(p, n, out, #[inline(always)] move |p| two * p, through(Erfcinv), #[inline(always)] move |_, e| exp(from(e)))
648 }
649 fn words(&self) -> usize { 2 }
650 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 { exp(self.normal().sample(rng)) }
651}
652
653#[derive(Clone, Copy, Debug, PartialEq)]
655struct StudentT { df: f64, factor: f64 }
656
657impl Formulas for StudentT {
658 fn pdf(&self, t: f64) -> f64 { self.factor * exp(-(self.df + 1.0) / 2.0 * log1p(t * t / self.df)) }
659 fn cdf(&self, t: f64) -> f64 {
660 let (h, d) = (self.df / 2.0, self.df + t * t);
661 let y = t * t / d;
662 let (inside, outside) = if y < 0.5 { let i = betainc(0.5, h, y); (i, 1.0 - i) } else { let o = betainc(h, 0.5, self.df / d); (1.0 - o, o) };
664 if t < 0.0 { 0.5 * outside } else { 0.5 + 0.5 * inside }
665 }
666 fn quantile(&self, p: f64) -> f64 {
667 let h = self.df / 2.0;
668 let q = 2.0 * p.min(1.0 - p);
670 let t = if q > 0.5 { let y = betaincinv(0.5, h, 1.0 - q); (self.df * y / (1.0 - y)).sqrt() }
671 else { let x = betaincinv(h, 0.5, q); (self.df * (1.0 - x) / x).sqrt() };
672 if p < 0.5 { -t } else { t }
673 }
674 fn words(&self) -> usize { 2 + gamma_words(self.df / 2.0) }
675 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 { standard_normal(rng) / (standard_gamma(self.df / 2.0, rng) * 2.0 / self.df).sqrt() }
676}
677
678#[derive(Clone, Copy, Debug, PartialEq)]
679struct Weibull { shape: f64, scale: f64 }
680
681impl Formulas for Weibull {
682 through_slices!();
683 #[inline(always)]
685 fn pdf_many<F: Float>(&self, x: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool {
686 let (rscale, height) = (F::narrow(self.scale.recip()), F::narrow(self.shape / self.scale));
687 composed(x, n, out, move |x| if x < F::zero() { F::one() } else { x * rscale }, raised(F::narrow(self.shape - 1.0)), move |x, w| {
688 let e = exp(-times(x * rscale, w));
689 let d = if e == F::zero() { e } else { height * w * e };
690 if x < F::zero() || x.is_infinite() { F::zero() } else { d }
691 })
692 }
693 #[inline(always)]
695 fn cdf_many<F: Float>(&self, x: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool {
696 let rscale = F::narrow(self.scale.recip());
697 composed(x, n, out, move |x| if x <= F::zero() { F::one() } else { x * rscale }, raised(F::narrow(self.shape)), |x, w| {
698 let c = -expm1(-w);
699 if x <= F::zero() { F::zero() } else { c }
700 })
701 }
702 #[inline(always)]
704 fn quantile_many<F: Float>(&self, p: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool where Log1p: Unary<F, Out = F> {
705 let (degree, scale) = (F::narrow(self.shape), F::narrow(self.scale));
706 composed(p, n, out, |p| -p, move |a, v| {
707 Log1p.slice([Arg::Many(a)], a.len(), v);
708 for (x, &l) in a.iter_mut().zip(unsafe { written(v) }.iter()) { *x = -l; }
709 crate::pow::roots(degree, a, v);
710 }, move |_, w| scale * w)
711 }
712 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 { self.scale * pow(-ln(open(rng)), 1.0 / self.shape) }
713}
714
715fn whole(x: f64, max: f64) -> bool { (0.0..=max).contains(&x) && x.fract() == 0.0 }
717
718fn least(target: f64, guess: f64, max: f64, cdf: impl Fn(f64) -> f64) -> f64 {
720 let reached = |k: f64| cdf(k) >= target * (1.0 - 1e-14);
721 let start = if guess.is_nan() { 0.0 } else { guess.floor().clamp(0.0, max) };
722 let (mut below, mut above) = (start, start);
723 let mut step = 1.0;
724 if reached(start) {
725 loop {
727 below = (above - step).max(-1.0);
728 if below < 0.0 || !reached(below) { break; }
729 above = below;
730 step *= 2.0;
731 }
732 } else {
733 loop {
734 if below >= max { return max; }
735 above = (below + step).min(max);
736 if reached(above) { break; }
737 below = above;
738 step *= 2.0;
739 }
740 }
741 while above - below > 1.0 { let middle = ((below + above) / 2.0).floor(); if reached(middle) { above = middle } else { below = middle } }
742 above
743}
744
745fn normal_guess(mean: f64, sd: f64, p: f64) -> f64 { mean - sd * SQRT_2 * erfcinv(2.0 * p) }
747
748#[derive(Clone, Copy, Debug, PartialEq)]
749struct Binomial { n: f64, p: f64 }
750
751impl Formulas for Binomial {
752 fn pdf(&self, x: f64) -> f64 {
753 if !whole(x, self.n) { return 0.0; }
754 let (n, p) = (self.n, self.p);
755 exp(lgamma(n + 1.0) - lgamma(x + 1.0) - lgamma(n - x + 1.0) + times(x, ln(p)) + times(n - x, log1p(-p)))
756 }
757 fn cdf(&self, x: f64) -> f64 {
758 if x < 0.0 { return 0.0; }
759 if x >= self.n || self.p == 0.0 { return 1.0; }
760 if self.p == 1.0 { return 0.0; }
761 let k = x.floor();
762 betainc(self.n - k, k + 1.0, 1.0 - self.p)
763 }
764 fn quantile(&self, q: f64) -> f64 {
766 let (n, p) = (self.n, self.p);
767 if p == 0.0 { 0.0 } else if p == 1.0 { n } else { least(q, normal_guess(n * p, (n * p * (1.0 - p)).sqrt(), q), n, |k| self.cdf(k)) }
768 }
769 fn words(&self) -> usize { if self.n * self.p.min(1.0 - self.p) < 10.0 { 1 } else { 2 } }
771 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
772 let flip = self.p > 0.5;
773 let p = if flip { 1.0 - self.p } else { self.p };
774 let k = if self.n * p < 10.0 { binomial_inversion(self.n, p, rng) } else { btrd(self.n, p, rng) };
775 if flip { self.n - k } else { k }
776 }
777}
778
779fn binomial_inversion<R: Rng + ?Sized>(n: f64, p: f64, rng: &mut R) -> f64 {
781 let q = 1.0 - p;
782 let first = exp(n * log1p(-p));
783 let bound = n.min(n * p + 10.0 * (n * p * q + 1.0).sqrt());
784 'draw: loop {
785 let (mut u, mut mass, mut k): (f64, f64, f64) = (StandardUniform.sample(rng), first, 0.0);
786 while u > mass {
787 u -= mass;
788 k += 1.0;
789 if k > bound { continue 'draw; }
790 mass *= (n - k + 1.0) * p / (k * q);
791 }
792 return k;
793 }
794}
795
796fn stirling(k: f64) -> f64 {
798 const TABLE: [f64; 10] = [
799 0.08106146679532726, 0.0413406959554093, 0.02767792568499834, 0.020790672103765093, 0.016644691189821193, 0.013876128823070748, 0.01189670994589177,
800 0.010411265261972096, 0.009255462182712733, 0.00833056343336287,
801 ];
802 if k < 10.0 { return TABLE[k as usize]; }
803 let r = 1.0 / (k + 1.0);
804 let r2 = r * r;
805 (1.0 / 12.0 - (1.0 / 360.0 - r2 / 1260.0) * r2) * r
806}
807
808fn btrd<R: Rng + ?Sized>(n: f64, p: f64, rng: &mut R) -> f64 {
810 let q = 1.0 - p;
811 let (m, r) = (((n + 1.0) * p).floor(), p / q);
812 let (nr, npq) = ((n + 1.0) * r, n * p * q);
813 let sq = npq.sqrt();
814 let b = 1.15 + 2.53 * sq;
815 let a = -0.0873 + 0.0248 * b + 0.01 * p;
816 let c = n * p + 0.5;
817 let alpha = (2.83 + 5.1 / b) * sq;
818 let vr = 0.92 - 4.2 / b;
819 let urvr = 0.86 * vr;
820 loop {
821 let mut v: f64 = StandardUniform.sample(rng);
822 if v <= urvr { let u = v / vr - 0.43; return ((2.0 * a / (0.5 - u.abs()) + b) * u + c).floor(); }
823 let u = if v >= vr { let u: f64 = StandardUniform.sample(rng); u - 0.5 } else {
824 let w = v / vr - 0.93;
825 let w2: f64 = StandardUniform.sample(rng);
826 v = w2 * vr;
827 0.5f64.copysign(w) - w
828 };
829 let us = 0.5 - u.abs();
830 let k = ((2.0 * a / us + b) * u + c).floor();
831 if k < 0.0 || k > n { continue; }
832 v *= alpha / (a / (us * us) + b);
833 let km = (k - m).abs();
834 if km <= 15.0 {
835 let mut f = 1.0;
837 let mut i = m.min(k);
838 while i < m.max(k) { i += 1.0; if m < k { f *= nr / i - r } else { v *= nr / i - r } }
839 if v <= f { return k; }
840 continue;
841 }
842 let v = ln(v);
844 let rho = (km / npq) * (((km / 3.0 + 0.625) * km + 1.0 / 6.0) / npq + 0.5);
845 let t = -km * km / (2.0 * npq);
846 if v < t - rho { return k; }
847 if v > t + rho { continue; }
848 let (nm, nk) = (n - m + 1.0, n - k + 1.0);
849 let h = (m + 0.5) * ln((m + 1.0) / (r * nm)) + stirling(m) + stirling(n - m);
850 if v <= h + (n + 1.0) * ln(nm / nk) + (k + 0.5) * ln(nk * r / (k + 1.0)) - stirling(k) - stirling(n - k) { return k; }
851 }
852}
853
854#[derive(Clone, Copy, Debug, PartialEq)]
855struct Poisson { rate: f64 }
856
857impl Formulas for Poisson {
858 fn pdf(&self, x: f64) -> f64 { if !whole(x, f64::INFINITY) { return 0.0; } exp(times(x, ln(self.rate)) - self.rate - lgamma(x + 1.0)) }
859 fn cdf(&self, x: f64) -> f64 {
860 if x < 0.0 { 0.0 } else if self.rate == 0.0 || x.is_infinite() { 1.0 } else { gammaincc(x.floor() + 1.0, self.rate) }
861 }
862 fn quantile(&self, q: f64) -> f64 {
864 if self.rate == 0.0 { 0.0 }
865 else if q == 1.0 { f64::INFINITY }
866 else { least(q, normal_guess(self.rate, self.rate.sqrt(), q), f64::INFINITY, |k| self.cdf(k)) }
867 }
868 fn words(&self) -> usize { if self.rate < 10.0 { (self.rate + 1.0).ceil() as usize } else { 2 } }
870 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
872 let rate = self.rate;
873 if rate < 10.0 {
874 let limit = exp(-rate);
875 let (mut k, mut product) = (0.0, 1.0);
876 loop {
877 let u: f64 = StandardUniform.sample(rng);
878 product *= u;
879 if product <= limit { return k; }
880 k += 1.0;
881 }
882 }
883 let (root, ln_rate) = (rate.sqrt(), ln(rate));
884 let b = 0.931 + 2.53 * root;
885 let a = -0.059 + 0.02483 * b;
886 let ln_alpha = ln(1.1239 + 1.1328 / (b - 3.4));
887 let vr = 0.9277 - 3.6224 / (b - 2.0);
888 loop {
889 let (u, v): (f64, f64) = (StandardUniform.sample(rng), StandardUniform.sample(rng));
890 let u = u - 0.5;
891 let us = 0.5 - u.abs();
892 let k = ((2.0 * a / us + b) * u + rate + 0.43).floor();
893 if us >= 0.07 && v <= vr { return k; }
894 if k < 0.0 || (us < 0.013 && v > us) { continue; }
895 if ln(v) + ln_alpha - ln(a / (us * us) + b) <= -rate + k * ln_rate - lgamma(k + 1.0) { return k; }
896 }
897 }
898}