use crate::search::bisect;
use crate::{C64, cis};
use std::ops::{Add, AddAssign, Div, Mul, MulAssign};
const LENTZ_TINY: f64 = 1e-300;
const LENTZ_TOLERANCE: f64 = 1e-16;
const SERIES_TOLERANCE: f64 = 1e-17;
trait Scalar: Copy + Add<Output = Self> + Mul<Output = Self> + MulAssign + AddAssign + Div<f64, Output = Self> {
fn tiny() -> Self;
fn magnitude(self) -> f64;
fn scale(a: f64, x: Self) -> Self;
fn over(a: f64, x: Self) -> Self;
fn recip(self) -> Self;
fn distance_to_one(self) -> f64;
}
impl Scalar for f64 {
fn tiny() -> f64 {
LENTZ_TINY
}
fn magnitude(self) -> f64 {
self.abs()
}
fn scale(a: f64, x: f64) -> f64 {
a * x
}
fn over(a: f64, x: f64) -> f64 {
a / x
}
fn recip(self) -> f64 {
1.0 / self
}
fn distance_to_one(self) -> f64 {
(self - 1.0).abs()
}
}
impl Scalar for C64 {
fn tiny() -> C64 {
C64::new(LENTZ_TINY, 0.0)
}
fn magnitude(self) -> f64 {
self.norm()
}
fn scale(a: f64, x: C64) -> C64 {
a * x
}
fn over(a: f64, x: C64) -> C64 {
a / x
}
fn recip(self) -> C64 {
self.inv()
}
fn distance_to_one(self) -> f64 {
(self - 1.0).norm()
}
}
fn series<T: Scalar>(
mut term: T,
mut sum: T,
max_terms: usize,
mut ratio: impl FnMut(usize) -> T,
divisor: impl Fn(usize) -> f64,
) -> T {
for k in 1..max_terms {
term *= ratio(k);
let t = term / divisor(k);
sum += t;
if t.magnitude() < SERIES_TOLERANCE * sum.magnitude() {
break;
}
}
sum
}
fn lentz<T: Scalar>(f0: T, c0: T, d0: T, max_terms: usize, mut term: impl FnMut(usize) -> (f64, T)) -> T {
let (mut f, mut c, mut d) = (f0, c0, d0);
for k in 1..max_terms {
let (a, b) = term(k);
d = b + T::scale(a, d);
if d.magnitude() < LENTZ_TINY {
d = T::tiny();
}
c = b + T::over(a, c);
if c.magnitude() < LENTZ_TINY {
c = T::tiny();
}
d = d.recip();
let delta = c * d;
f *= delta;
if delta.distance_to_one() < LENTZ_TOLERANCE {
break;
}
}
f
}
fn midpoint_mean(f: impl Fn(f64) -> f64, n: usize) -> f64 {
let h = std::f64::consts::PI / n as f64;
(0..n).map(|k| f(h * (k as f64 + 0.5))).sum::<f64>() / n as f64
}
#[allow(
clippy::unreadable_literal,
clippy::excessive_precision,
reason = "the coefficients of Cody's CALERF, digit for digit"
)]
pub fn erfc(x: f64) -> f64 {
const A: [f64; 5] = [
3.16112374387056560e00,
1.13864154151050156e02,
3.77485237685302021e02,
3.20937758913846947e03,
1.85777706184603153e-1,
];
const B: [f64; 4] =
[2.36012909523441209e01, 2.44024637934444173e02, 1.28261652607737228e03, 2.84423683343917062e03];
const C: [f64; 9] = [
5.64188496988670089e-1,
8.88314979438837594e00,
6.61191906371416295e01,
2.98635138197400131e02,
8.81952221241769090e02,
1.71204761263407058e03,
2.05107837782607147e03,
1.23033935479799725e03,
2.15311535474403846e-8,
];
const D: [f64; 8] = [
1.57449261107098347e01,
1.17693950891312499e02,
5.37181101862009858e02,
1.62138957456669019e03,
3.29079923573345963e03,
4.36261909014324716e03,
3.43936767414372164e03,
1.23033935480374942e03,
];
const P: [f64; 6] = [
3.05326634961232344e-1,
3.60344899949804439e-1,
1.25781726111229246e-1,
1.60837851487422766e-2,
6.58749161529837803e-4,
1.63153871373020978e-2,
];
const Q: [f64; 5] = [
2.56852019228982242e00,
1.87295284992346725e00,
5.27905102951428412e-1,
6.05183413124413191e-2,
2.33520497626869185e-3,
];
const FRAC_1_SQRT_PI: f64 = 5.6418958354775628695e-1;
let y = x.abs();
if y <= 0.46875 {
let ysq = if y > 1.11e-16 { y * y } else { 0.0 };
let (mut num, mut den) = (A[4] * ysq, ysq);
for i in 0..3 {
num = (num + A[i]) * ysq;
den = (den + B[i]) * ysq;
}
return 1.0 - x * (num + A[3]) / (den + B[3]);
}
let r = if y > 27.3 {
0.0
} else {
let r = if y <= 4.0 {
let (mut num, mut den) = (C[8] * y, y);
for i in 0..7 {
num = (num + C[i]) * y;
den = (den + D[i]) * y;
}
(num + C[7]) / (den + D[7])
} else {
let ysq = 1.0 / (y * y);
let (mut num, mut den) = (P[5] * ysq, ysq);
for i in 0..4 {
num = (num + P[i]) * ysq;
den = (den + Q[i]) * ysq;
}
(FRAC_1_SQRT_PI - ysq * (num + P[4]) / (den + Q[4])) / y
};
let u = (y * 16.0).trunc() / 16.0;
(-u * u).exp() * (-(y - u) * (y + u)).exp() * r
};
if x < 0.0 { 2.0 - r } else { r }
}
pub fn erfcinv(y: f64) -> f64 {
if y <= 0.0 {
return f64::INFINITY;
}
if y >= 2.0 {
return f64::NEG_INFINITY;
}
bisect(|x| erfc(x) > y, -27.0, 27.0, 200)
}
#[inline]
pub fn qfunc(x: f64) -> f64 {
0.5 * erfc(x / std::f64::consts::SQRT_2)
}
#[inline]
pub fn qfuncinv(p: f64) -> f64 {
std::f64::consts::SQRT_2 * erfcinv(2.0 * p)
}
pub fn normal_pdf(x: f64) -> f64 {
(-0.5 * x * x).exp() / std::f64::consts::TAU.sqrt()
}
pub fn fresnel(x: f64) -> (f64, f64) {
let ax = x.abs();
let v = if ax < 1.5 {
let a = C64::new(0.0, std::f64::consts::FRAC_PI_2 * ax * ax);
series(C64::new(ax, 0.0), C64::new(ax, 0.0), 60, |k| a / k as f64, |k| (2 * k + 1) as f64)
} else {
let z = C64::new(1.0, -1.0) * (std::f64::consts::PI.sqrt() / 2.0 * ax);
let f = lentz(z, z, C64::new(0.0, 0.0), 100_000, |k| (k as f64 / 2.0, z));
let erfc = cis(std::f64::consts::FRAC_PI_2 * ax * ax) / (std::f64::consts::PI.sqrt() * f);
C64::new(0.5, 0.5) * (1.0 - erfc)
};
(v.re.copysign(x), v.im.copysign(x))
}
pub fn bessel_j0(x: f64) -> f64 {
let n = 64 + 2 * x.abs().ceil() as usize;
midpoint_mean(|theta| (x * theta.sin()).cos(), n)
}
pub fn ln_bessel_k(nu: f64, x: f64) -> f64 {
if x.is_nan() || nu.is_nan() || x <= 0.0 {
return f64::NAN;
}
let nu = nu.abs();
if x == f64::INFINITY {
return f64::NEG_INFINITY;
}
if nu == f64::INFINITY {
return f64::INFINITY;
}
if nu >= 0.05 && !(nu / x).is_finite() {
return -std::f64::consts::LN_2 + ln_gamma(nu) + nu * (std::f64::consts::LN_2 - x.ln());
}
let f = |t: f64| {
let a = nu * t;
a + (-2.0 * a).exp().ln_1p() - std::f64::consts::LN_2 - x * t.cosh()
};
let t_max = if nu > 0.0 { (nu / x).asinh() } else { 0.0 };
let peak = f(t_max);
let width = 1.0 / (x * t_max.cosh()).sqrt().max(1.0);
let h = 0.1 * width.min(1.0);
let mut acc = 0.5 * (f(0.0) - peak).exp();
let mut k = 1usize;
loop {
let t = k as f64 * h;
let v = f(t) - peak;
if v.is_nan() {
return f64::NAN;
}
acc += v.exp();
if t > t_max && v < -50.0 {
break;
}
k += 1;
}
peak + (acc * h).ln()
}
pub fn bessel_k(nu: f64, x: f64) -> f64 {
ln_bessel_k(nu, x).exp()
}
pub fn bessel_i0e(x: f64) -> f64 {
let ax = x.abs();
if ax <= 60.0 {
let n = 32 + 8 * ax.sqrt().ceil() as usize + ax.ceil() as usize;
midpoint_mean(|theta| (ax * (theta.cos() - 1.0)).exp(), n)
} else {
let ratio = |k: usize| {
let odd = (2 * k - 1) as f64;
odd * odd / (k as f64 * 8.0 * ax)
};
series(1.0, 1.0, 30, ratio, |_| 1.0) / (std::f64::consts::TAU * ax).sqrt()
}
}
pub fn erf(x: f64) -> f64 {
if x.abs() >= 2.0 {
return (1.0 - erfc(x.abs())).copysign(x);
}
let x2 = x * x;
series(x, x, 60, |k| -x2 / k as f64, |k| (2 * k + 1) as f64) * std::f64::consts::FRAC_2_SQRT_PI
}
#[allow(clippy::unreadable_literal, reason = "the Lanczos coefficients, digit for digit")]
pub fn ln_gamma(x: f64) -> f64 {
const G: f64 = 7.0;
const C: [f64; 9] = [
0.999_999_999_999_809_9,
676.5203681218851,
-1259.1392167224028,
771.323_428_777_653_1,
-176.615_029_162_140_6,
12.507343278686905,
-0.13857109526572012,
9.984_369_578_019_572e-6,
1.5056327351493116e-7,
];
if x < 0.5 {
let pi = std::f64::consts::PI;
return (pi / (pi * x).sin()).ln() - ln_gamma(1.0 - x);
}
let x = x - 1.0;
let t = x + G + 0.5;
let a = C[0] + C.iter().enumerate().skip(1).map(|(i, c)| c / (x + i as f64)).sum::<f64>();
0.5 * std::f64::consts::TAU.ln() + (x + 0.5) * t.ln() - t + a.ln()
}
pub fn gamma_p(a: f64, x: f64) -> f64 {
if x <= 0.0 {
return 0.0;
}
let ln_front = a * x.ln() - x - ln_gamma(a);
if x < a + 1.0 {
let mut ap = a;
let ratio = |_| {
ap += 1.0;
x / ap
};
series(1.0 / a, 1.0 / a, 10_001, ratio, |_| 1.0) * ln_front.exp()
} else {
let mut b = x + 1.0 - a;
let d = 1.0 / b;
let term = |i: usize| {
b += 2.0;
(-(i as f64) * (i as f64 - a), b)
};
let h = lentz(d, 1.0 / LENTZ_TINY, d, 10_000, term);
1.0 - ln_front.exp() * h
}
}
pub fn ln_binomial_pmf(n: usize, p: f64) -> Vec<f64> {
if p <= 0.0 || p >= 1.0 {
let certain = if p <= 0.0 { 0 } else { n };
return (0..=n).map(|k| if k == certain { 0.0 } else { f64::NEG_INFINITY }).collect();
}
let (lp, lq) = (p.ln(), (1.0 - p).ln());
let mut terms = Vec::with_capacity(n + 1);
let mut log_c = 0.0;
for k in 0..=n {
if k > 0 {
log_c += ((n + 1 - k) as f64).ln() - (k as f64).ln();
}
terms.push(log_c + k as f64 * lp + (n - k) as f64 * lq);
}
terms
}
pub fn beta_inc(a: f64, b: f64, x: f64) -> f64 {
if x <= 0.0 {
return 0.0;
}
if x >= 1.0 {
return 1.0;
}
if x > (a + 1.0) / (a + b + 2.0) {
return 1.0 - beta_inc(b, a, 1.0 - x);
}
let ln_front = a * x.ln() + b * (1.0 - x).ln() + ln_gamma(a + b) - ln_gamma(a) - ln_gamma(b);
let term = |k: usize| {
let m = (k / 2) as f64;
let d = if k % 2 == 1 {
-(a + m) * (a + b + m) * x / ((a + 2.0 * m) * (a + 2.0 * m + 1.0))
} else {
m * (b - m) * x / ((a + 2.0 * m - 1.0) * (a + 2.0 * m))
};
(d, 1.0)
};
ln_front.exp() / (a * lentz(1.0, 1.0, 0.0, 10_000, term))
}
pub fn beta_inv(q: f64, a: f64, b: f64) -> f64 {
if q <= 0.0 {
return 0.0;
}
if q >= 1.0 {
return 1.0;
}
bisect(|x| beta_inc(a, b, x) < q, 0.0, 1.0, 200)
}
pub fn expint_e1(x: f64) -> f64 {
if x <= 0.0 {
return f64::INFINITY;
}
if x <= 1.0 {
const EULER: f64 = 0.577_215_664_901_532_9;
-EULER - x.ln() - series(1.0, 0.0, 60, |k| -x / k as f64, |k| k as f64)
} else {
let mut b = x + 1.0;
let d = 1.0 / b;
let term = |i: usize| {
b += 2.0;
(-((i * i) as f64), b)
};
lentz(d, 1.0 / LENTZ_TINY, d, 10_000, term) * (-x).exp()
}
}