use astro_float::{BigFloat, Consts, RoundingMode};
use crate::base::errors::SymplexError;
fn unevaluable(reason: impl Into<String>) -> SymplexError {
SymplexError::Unevaluable {
reason: reason.into(),
}
}
fn pos(x: &BigFloat) -> bool {
x.is_positive() && !x.is_zero()
}
fn is_npint(x: &BigFloat) -> bool {
x.is_zero() || (x.is_int() && x.is_negative())
}
fn is_even(x: &BigFloat) -> bool {
if x.is_zero() {
return true;
}
let mut half = x.clone();
match x.exponent() {
Some(e) => {
half.set_exponent(e - 1);
half.is_int()
}
None => true,
}
}
fn gamma_sign(x: &BigFloat) -> i32 {
if pos(x) || is_even(&x.floor()) { 1 } else { -1 }
}
pub(super) fn exact_sum(a: &BigFloat, b: &BigFloat, prec: usize, rm: RoundingMode) -> BigFloat {
let span = |x: &BigFloat| -> Option<(i64, i64)> {
let e = i64::from(x.exponent()?);
let m = x.mantissa_max_bit_len().unwrap_or(64) as i64;
(!x.is_zero()).then_some((e, e - m))
};
let p = match (span(a), span(b)) {
(Some((ea, la)), Some((eb, lb))) => {
let bits = ea.max(eb) + 2 - la.min(lb);
usize::try_from(bits).unwrap_or(prec).min(prec + (1 << 18))
}
(Some((ea, la)), None) | (None, Some((ea, la))) => {
usize::try_from(ea - la + 2).unwrap_or(prec)
}
(None, None) => prec,
}
.max(64);
a.add(b, p, rm)
}
fn int(n: i32) -> BigFloat {
BigFloat::from_i32(n, 64)
}
pub(super) fn gamma_ratio(
num: &[BigFloat],
den: &[BigFloat],
prec: usize,
rm: RoundingMode,
cc: &mut Consts,
) -> Result<BigFloat, SymplexError> {
if num.iter().chain(den).any(|x| x.is_nan() || x.is_inf()) {
return Err(unevaluable("Gamma quotient of a non-finite argument"));
}
let (poles_num, mut regular_num): (Vec<&BigFloat>, Vec<BigFloat>) = split_poles(num);
let (poles_den, mut regular_den): (Vec<&BigFloat>, Vec<BigFloat>) = split_poles(den);
if poles_num.len() < poles_den.len() {
return Ok(BigFloat::new(prec));
}
if poles_num.len() > poles_den.len() {
return Err(unevaluable(
"the Gamma quotient has a pole here (more poles of Gamma above than below): the value is infinite",
));
}
let mut sign = 1;
for (i, j) in poles_num.into_iter().zip(poles_den) {
if !is_even(&exact_sum(i, j, prec, rm)) {
sign = -sign;
}
regular_num.push(exact_sum(&int(1), &j.neg(), prec, rm));
regular_den.push(exact_sum(&int(1), &i.neg(), prec, rm));
}
let unit = |x: &BigFloat| x.cmp(&int(1)) == Some(0) || x.cmp(&int(2)) == Some(0);
regular_num.retain(|x| !unit(x));
regular_den.retain(|x| !unit(x));
for x in regular_num.iter().chain(®ular_den) {
sign *= gamma_sign(x);
}
let mut wp = prec + 64;
loop {
let mut total = BigFloat::new(wp);
let mut size = 0.0f64;
for (x, plus) in regular_num
.iter()
.map(|x| (x, true))
.chain(regular_den.iter().map(|x| (x, false)))
{
let l = super::arb_log_gamma(x, wp, rm, cc)?.0;
size += super::bigfloat_to_f64_rounded(&l, rm)
.unwrap_or(f64::INFINITY)
.abs();
total = if plus {
total.add(&l, wp, rm)
} else {
total.sub(&l, wp, rm)
};
}
let needed = prec + 32 + (size + 1.0).log2().ceil().max(0.0) as usize;
if needed <= wp || !size.is_finite() {
let v = total.exp(wp, rm, cc);
if v.is_inf() || v.is_nan() || (v.is_zero() && !total.is_zero()) {
return Err(unevaluable(
"the Gamma quotient is beyond the floating-point exponent range",
));
}
let v = super::round_to(v, prec, rm);
return Ok(if sign < 0 { v.neg() } else { v });
}
wp = needed;
}
}
fn split_poles(xs: &[BigFloat]) -> (Vec<&BigFloat>, Vec<BigFloat>) {
let mut poles = Vec::new();
let mut regular = Vec::new();
for x in xs {
if is_npint(x) {
poles.push(x);
} else {
regular.push(x.clone());
}
}
(poles, regular)
}
pub(super) fn rising_factorial(
x: &BigFloat,
n: &BigFloat,
prec: usize,
rm: RoundingMode,
cc: &mut Consts,
) -> Result<BigFloat, SymplexError> {
if n.is_zero() {
return Ok(BigFloat::from_i32(1, prec));
}
let xn = exact_sum(x, n, prec, rm);
gamma_ratio(&[xn], std::slice::from_ref(x), prec, rm, cc)
}
pub(super) fn falling_factorial(
x: &BigFloat,
n: &BigFloat,
prec: usize,
rm: RoundingMode,
cc: &mut Consts,
) -> Result<BigFloat, SymplexError> {
if n.is_zero() {
return Ok(BigFloat::from_i32(1, prec));
}
let x1 = exact_sum(x, &int(1), prec, rm);
let xn1 = exact_sum(&x1, &n.neg(), prec, rm);
gamma_ratio(&[x1], &[xn1], prec, rm, cc)
}
pub(super) fn binomial(
n: &BigFloat,
k: &BigFloat,
prec: usize,
rm: RoundingMode,
cc: &mut Consts,
) -> Result<BigFloat, SymplexError> {
let n1 = exact_sum(n, &int(1), prec, rm);
let k1 = exact_sum(k, &int(1), prec, rm);
let nk1 = exact_sum(&n1, &k.neg(), prec, rm);
gamma_ratio(&[n1], &[k1, nk1], prec, rm, cc)
}
pub(super) fn binomial_meets_pole(
n: &BigFloat,
k: &BigFloat,
prec: usize,
rm: RoundingMode,
) -> bool {
let n1 = exact_sum(n, &int(1), prec, rm);
let k1 = exact_sum(k, &int(1), prec, rm);
let nk1 = exact_sum(&n1, &k.neg(), prec, rm);
is_npint(&n1) || is_npint(&k1) || is_npint(&nk1)
}
pub(super) fn catalan(
x: &BigFloat,
prec: usize,
rm: RoundingMode,
cc: &mut Consts,
) -> Result<BigFloat, SymplexError> {
if x.is_nan() || x.is_inf() {
return Err(unevaluable("catalan of a non-finite argument"));
}
if is_npint(x) && !x.is_zero() {
return Ok(if x.cmp(&int(-1)) == Some(0) {
int(-1).div(&int(2), prec, rm)
} else {
BigFloat::new(prec)
});
}
let numerator = exact_sum(&exact_sum(x, x, prec, rm), &int(1), prec, rm);
if is_npint(&numerator) {
return Err(unevaluable(
"catalan has a pole at a negative half-integer (the value is infinite)",
));
}
let x1 = exact_sum(x, &int(1), prec, rm);
let x2 = exact_sum(x, &int(2), prec, rm);
gamma_ratio(&[numerator], &[x1, x2], prec, rm, cc)
}
fn cos_pi(x: &BigFloat, wp: usize, rm: RoundingMode, cc: &mut Consts) -> BigFloat {
let half = BigFloat::from_f64(0.5, 64);
let mut k = x.floor();
let mut f = exact_sum(x, &k.neg(), wp, rm);
if f.cmp(&half).is_some_and(|c| c > 0) {
f = exact_sum(&f, &int(-1), wp, rm);
k = exact_sum(&k, &int(1), wp, rm);
}
let g = exact_sum(&half, &f.abs().neg(), wp, rm);
let s = if g.is_zero() {
BigFloat::new(wp)
} else {
cc.pi(wp, rm).clone().mul(&g, wp, rm).sin(wp, rm, cc)
};
if is_even(&k) { s } else { s.neg() }
}
pub(super) fn fibonacci(
x: &BigFloat,
lucas: bool,
prec: usize,
rm: RoundingMode,
cc: &mut Consts,
) -> Result<BigFloat, SymplexError> {
if x.is_nan() || x.is_inf() {
return Err(unevaluable("fibonacci of a non-finite argument"));
}
let int_bits = x
.exponent()
.map_or(0, |e| usize::try_from(e).unwrap_or(0))
.min(1 << 20);
let cap = super::cancellation_cap(prec);
let mut extra = 32 + int_bits;
loop {
let wp = prec + extra;
let sqrt5 = int(5).sqrt(wp, rm);
let phi = int(1).add(&sqrt5, wp, rm).div(&int(2), wp, rm);
let t = x.mul(&phi.ln(wp, rm, cc), wp, rm);
let up = t.exp(wp, rm, cc);
let down = t.neg().exp(wp, rm, cc);
if up.is_inf() || up.is_nan() || down.is_inf() || down.is_nan() {
return Err(unevaluable(
"fibonacci of this argument is beyond the floating-point exponent range",
));
}
let cosine = cos_pi(x, wp, rm, cc);
let tail = cosine.mul(&down, wp, rm);
let sum = if lucas {
up.add(&tail, wp, rm)
} else {
up.sub(&tail, wp, rm)
};
let top = [&up, &tail]
.into_iter()
.filter(|v| !v.is_zero())
.filter_map(|v| v.exponent())
.max()
.unwrap_or(0);
let lost = match sum.exponent() {
Some(e) if !sum.is_zero() => {
usize::try_from(i64::from(top) - i64::from(e)).unwrap_or(0)
}
_ => wp,
};
if lost + 16 <= extra {
let value = if lucas { sum } else { sum.div(&sqrt5, wp, rm) };
return Ok(super::round_to(value, prec, rm));
}
if extra >= cap {
return Err(super::special_exhausted(prec));
}
extra = (lost + 32).max(2 * extra).min(cap);
}
}
pub(super) fn harmonic(
z: &BigFloat,
prec: usize,
rm: RoundingMode,
cc: &mut Consts,
) -> Result<BigFloat, SymplexError> {
if z.is_nan() || z.is_inf() {
return Err(unevaluable("harmonic of a non-finite argument"));
}
if z.is_zero() {
return Ok(BigFloat::new(prec));
}
if z.is_int() && z.is_negative() {
return Err(unevaluable(
"harmonic has a pole at a negative integer (the value is infinite)",
));
}
let z1 = exact_sum(z, &int(1), prec, rm);
let cap = super::cancellation_cap(prec);
let mut extra = 32;
loop {
let wp = prec + extra;
let psi = super::arb_digamma(&z1, wp, rm, cc)?;
let gamma = super::arb_euler_gamma(wp, rm, cc)?;
let h = psi.add(&gamma, wp, rm);
let top = psi
.exponent()
.unwrap_or(0)
.max(gamma.exponent().unwrap_or(0));
let lost = match h.exponent() {
Some(e) if !h.is_zero() => usize::try_from(i64::from(top) - i64::from(e)).unwrap_or(0),
_ => wp,
};
if lost + 16 <= extra {
return Ok(super::round_to(h, prec, rm));
}
if extra >= cap {
return Err(super::special_exhausted(prec));
}
extra = (lost + 32).max(2 * extra).min(cap);
}
}
const MAX_ORDER: f64 = 9.0e15;
pub(super) fn polygamma_large(
n: &BigFloat,
x: &BigFloat,
prec: usize,
rm: RoundingMode,
cc: &mut Consts,
) -> Result<BigFloat, SymplexError> {
if x.is_nan() || x.is_inf() {
return Err(unevaluable("polygamma of a non-finite argument"));
}
if is_npint(x) {
return Err(unevaluable("polygamma at non-positive integer pole"));
}
let s = exact_sum(n, &int(1), prec, rm);
let sf = super::bigfloat_to_f64_rounded(&s, rm)?;
if !(2.0..=MAX_ORDER).contains(&sf) {
return Err(SymplexError::NotImplemented(format!(
"polygamma of order {} (at most {MAX_ORDER:e})",
sf - 1.0
)));
}
let odd = !is_even(&s);
let cap = super::cancellation_cap(prec);
let mut extra = 32;
loop {
let wp = prec + extra;
let (scale, z) = scaled_hurwitz(&s, sf, odd, x, wp, rm, cc)?;
let lost = match z.exponent() {
Some(e) if !z.is_zero() => usize::try_from(-i64::from(e)).unwrap_or(0),
_ => wp,
};
if lost + 16 > extra {
if extra >= cap {
return Err(super::special_exhausted(prec));
}
extra = (lost + 32).max(2 * extra).min(cap);
continue;
}
let lg = super::arb_log_gamma(&s, wp, rm, cc)?.0;
let ln_d = scale.ln(wp, rm, cc);
let size = super::bigfloat_to_f64_rounded(&lg, rm)?.abs()
+ sf * super::bigfloat_to_f64_rounded(&ln_d, rm)?.abs();
let lp = prec + 32 + (size + 1.0).log2().ceil().max(0.0) as usize;
let lg = super::arb_log_gamma(&s, lp, rm, cc)?.0;
let s_ln_d = s.mul(&scale.ln(lp, rm, cc), lp, rm);
let ln_z = z.abs().ln(lp, rm, cc);
let total = lg.sub(&s_ln_d, lp, rm).add(&ln_z, lp, rm);
let v = total.exp(lp, rm, cc);
if v.is_inf() || v.is_nan() || v.is_zero() {
return Err(unevaluable(
"polygamma of this order is beyond the floating-point exponent range",
));
}
let negative = odd != z.is_negative();
let v = super::round_to(v, prec, rm);
return Ok(if negative { v.neg() } else { v });
}
}
fn scaled_hurwitz(
s: &BigFloat,
sf: f64,
odd: bool,
x: &BigFloat,
wp: usize,
rm: RoundingMode,
cc: &mut Consts,
) -> Result<(BigFloat, BigFloat), SymplexError> {
if pos(x) {
return Ok((x.clone(), positive_part(s, sf, x, x, wp, rm, cc)?));
}
let m = x.floor().neg();
let f = exact_sum(x, &m, wp, rm);
let one_minus_f = exact_sum(&int(1), &f.neg(), wp, rm);
let mf = super::bigfloat_to_f64_rounded(&m, rm)?;
if odd && f == one_minus_f {
let start = exact_sum(&f, &m, wp, rm);
return Ok((
start.clone(),
positive_part(s, sf, &start, &start, wp, rm, cc)?,
));
}
let g = if f < one_minus_f {
f.clone()
} else {
one_minus_f.clone()
};
let negligible = |t: &BigFloat, sum: &BigFloat| -> bool {
match (t.exponent(), sum.exponent()) {
(_, _) if t.is_zero() => true,
(Some(te), Some(se)) if !sum.is_zero() => {
i64::from(se) - i64::from(te) > wp as i64 + 32
}
_ => false,
}
};
let mut z = BigFloat::new(wp);
let mut j = 0.0f64;
let mut u = f.clone(); let mut v = one_minus_f.clone(); while j < mf {
let a = scaled_term(s, &u, &g, wp, rm, cc);
let b = scaled_term(s, &v, &g, wp, rm, cc);
let pair = if odd {
a.sub(&b, wp, rm)
} else {
a.add(&b, wp, rm)
};
z = z.add(&pair, wp, rm);
if negligible(&a, &z) && negligible(&b, &z) {
return Ok((g, z));
}
j += 1.0;
u = exact_sum(&u, &int(1), wp, rm);
v = exact_sum(&v, &int(1), wp, rm);
if j > 1.0e6 {
return Err(SymplexError::NotImplemented(
"polygamma of large order far left of 0 (more than a million terms)".into(),
));
}
}
let rest = positive_part(s, sf, &u, &g, wp, rm, cc)?;
Ok((g, z.add(&rest, wp, rm)))
}
fn scaled_term(
s: &BigFloat,
u: &BigFloat,
g: &BigFloat,
wp: usize,
rm: RoundingMode,
cc: &mut Consts,
) -> BigFloat {
let s_bits = s.exponent().map_or(0, |e| usize::try_from(e).unwrap_or(0));
let p = wp + s_bits + 2 * (usize::BITS - wp.leading_zeros()) as usize + 16;
let l = u.div(g, p, rm).ln(p, rm, cc);
s.mul(&l, p, rm).neg().exp(wp, rm, cc)
}
fn positive_part(
s: &BigFloat,
sf: f64,
f: &BigFloat,
g: &BigFloat,
wp: usize,
rm: RoundingMode,
cc: &mut Consts,
) -> Result<BigFloat, SymplexError> {
let switch = (sf / std::f64::consts::PI).max(wp as f64) + 1.0;
let mut sum = BigFloat::new(wp);
let mut u = f.clone();
for _ in 0..1_000_000 {
let uf = super::bigfloat_to_f64_rounded(&u, rm)?;
if uf >= switch {
let tail = euler_maclaurin_tail(s, &u, g, wp, rm, cc);
return Ok(sum.add(&tail, wp, rm));
}
let t = scaled_term(s, &u, g, wp, rm, cc);
sum = sum.add(&t, wp, rm);
let u1 = exact_sum(&u, &int(1), wp, rm);
let next = scaled_term(s, &u1, g, wp, rm, cc);
let factor = 1.0 + (uf + 1.0) / (sf - 1.0);
let bound_exp =
next.exponent().map_or(i64::MIN / 2, i64::from) + factor.log2().ceil() as i64 + 1;
let sum_exp = sum.exponent().map_or(i64::MIN / 2, i64::from);
if next.is_zero() || sum_exp - bound_exp > wp as i64 + 8 {
return Ok(sum);
}
u = u1;
}
Err(SymplexError::NotImplemented(
"polygamma of large order: the Hurwitz zeta sum did not settle in a million terms".into(),
))
}
fn euler_maclaurin_tail(
s: &BigFloat,
a: &BigFloat,
g: &BigFloat,
wp: usize,
rm: RoundingMode,
cc: &mut Consts,
) -> BigFloat {
let one = BigFloat::from_i32(1, wp);
let s_minus_1 = s.sub(&one, wp, rm);
let mut bracket = a
.div(&s_minus_1, wp, rm)
.add(&BigFloat::from_f64(0.5, wp), wp, rm);
let a2 = a.mul(a, wp, rm);
let mut coef = s.div(&a.mul(&BigFloat::from_i32(2, wp), wp, rm), wp, rm);
let mut prev: Option<i64> = None;
for j in 1usize..=(wp / 2 + 64) {
if j > 1 {
let jj = j as i64;
let p1 = s.add(&BigFloat::from_i64(2 * jj - 3, 64), wp, rm);
let p2 = s.add(&BigFloat::from_i64(2 * jj - 2, 64), wp, rm);
let q = BigFloat::from_i64((2 * jj - 1) * (2 * jj), 64);
coef = coef
.mul(&p1, wp, rm)
.mul(&p2, wp, rm)
.div(&q, wp, rm)
.div(&a2, wp, rm);
}
let b = super::ratio_to_bigfloat(&super::bernoulli::even(j), wp, rm);
let term = b.mul(&coef, wp, rm);
let te = term.exponent().map_or(i64::MIN / 2, i64::from);
if prev.is_some_and(|p| te > p) {
break; }
prev = Some(te);
bracket = bracket.add(&term, wp, rm);
let be = bracket.exponent().map_or(0, i64::from);
if term.is_zero() || be - te > wp as i64 + 8 {
break;
}
}
scaled_term(s, a, g, wp, rm, cc).mul(&bracket, wp, rm)
}