use rug::Float;
use crate::ball::{ArbBall, IntervalEval};
use crate::kernel::{ExprId, ExprPool};
pub(crate) const VERIFY_PREC: u32 = 192;
const T_MAX: f64 = 4.4;
const LEVELS: u32 = 10;
const CONVERGED: f64 = 1e-20;
const DECAY_FLOOR: f64 = 1e-8;
pub(crate) fn fl(prec: u32, v: f64) -> Float {
Float::with_val(prec, v)
}
pub(crate) fn pi_ball(prec: u32) -> ArbBall {
let mid = Float::with_val(prec, rug::float::Constant::Pi);
let rad = Float::with_val(prec, 16u32) >> prec;
ArbBall { mid, rad, prec }
}
pub(crate) fn numeric_ball(expr: ExprId, pool: &ExprPool) -> Option<ArbBall> {
eval_ball(expr, &[], pool, VERIFY_PREC)
}
pub(crate) fn eval_ball(
expr: ExprId,
bindings: &[(ExprId, Float)],
pool: &ExprPool,
prec: u32,
) -> Option<ArbBall> {
let mut ev = IntervalEval::new(prec);
ev.bind(super::pi(pool), pi_ball(prec));
for (sym, val) in bindings {
ev.bind(
*sym,
ArbBall {
mid: val.clone(),
rad: fl(prec, 0.0),
prec,
},
);
}
let ball = ev.eval(expr, pool)?;
if !ball.mid.is_finite() || !ball.rad.is_finite() {
return None;
}
Some(ball)
}
pub(crate) enum Sample {
Value(Float),
Unevaluable,
}
#[derive(Clone)]
pub(crate) enum Region {
Real,
Positive,
Interval(Float, Float),
}
pub(crate) enum QuadOutcome {
Value { value: Float, est_err: Float },
Divergent,
Inconclusive,
}
#[derive(Clone, Debug)]
pub(crate) struct Window {
pub centre: Float,
pub scale: Float,
}
impl Window {
pub fn union(&self, other: &Window, prec: u32) -> Window {
let lo_a = Float::with_val(prec, &self.centre - &self.scale);
let hi_a = Float::with_val(prec, &self.centre + &self.scale);
let lo_b = Float::with_val(prec, &other.centre - &other.scale);
let hi_b = Float::with_val(prec, &other.centre + &other.scale);
let lo = if lo_a < lo_b { lo_a } else { lo_b };
let hi = if hi_a > hi_b { hi_a } else { hi_b };
Window {
centre: mul(&Float::with_val(prec, &lo + &hi), &fl(prec, 0.5), prec),
scale: mul(&Float::with_val(prec, &hi - &lo), &fl(prec, 0.5), prec),
}
}
}
pub(crate) enum Located {
At(Window),
Vanishes,
NotFound,
}
pub(crate) fn quadrature(
g: &mut impl FnMut(&Float) -> Sample,
region: Region,
prec: u32,
) -> QuadOutcome {
let window = match region {
Region::Interval(_, _) => None,
Region::Real | Region::Positive => {
let log_scale = matches!(region, Region::Positive);
match locate_where(g, prec, log_scale) {
Located::At(w) => Some(w),
Located::Vanishes => {
return QuadOutcome::Value {
value: fl(prec, 0.0),
est_err: fl(prec, 0.0),
}
}
Located::NotFound => return QuadOutcome::Inconclusive,
}
}
};
quadrature_at(g, region, window.as_ref(), prec)
}
pub(crate) fn quadrature_at(
g: &mut impl FnMut(&Float) -> Sample,
region: Region,
window: Option<&Window>,
prec: u32,
) -> QuadOutcome {
match region {
Region::Interval(lo, hi) => tanh_sinh(g, &lo, &hi, prec),
Region::Real => match window {
Some(w) => sinh_sinh(g, &w.centre, &w.scale, prec, false),
None => QuadOutcome::Inconclusive,
},
Region::Positive => match window {
Some(w) => sinh_sinh(g, &w.centre, &w.scale, prec, true),
None => QuadOutcome::Inconclusive,
},
}
}
pub(crate) fn locate(
g: &mut impl FnMut(&Float) -> Sample,
prec: u32,
log_scale: bool,
) -> Option<Window> {
match locate_where(g, prec, log_scale) {
Located::At(w) => Some(w),
_ => None,
}
}
pub(crate) fn locate_where(
g: &mut impl FnMut(&Float) -> Sample,
prec: u32,
log_scale: bool,
) -> Located {
let mut probe = |s: f64| -> Option<Float> {
let arg = if log_scale {
Float::with_val(prec, fl(prec, s).exp())
} else {
fl(prec, s)
};
match g(&arg) {
Sample::Value(v) if v.is_finite() => {
let mut a = v;
a.abs_mut();
if log_scale {
a *= &arg;
}
Some(a)
}
_ => None,
}
};
let mut grid: Vec<f64> = vec![0.0];
for k in -24i32..=48 {
let v = (2f64).powf(f64::from(k) / 2.0);
grid.push(v);
grid.push(-v);
}
let mut best: Option<(f64, Float)> = None;
let mut evaluated = 0usize;
for &s in &grid {
if let Some(v) = probe(s) {
evaluated += 1;
let better = match &best {
None => v > 0,
Some((_, bv)) => &v > bv,
};
if better {
best = Some((s, v));
}
}
}
let Some((centre, peak)) = best else {
return if evaluated >= grid.len() / 2 {
Located::Vanishes
} else {
Located::NotFound
};
};
if peak == 0 {
return Located::Vanishes;
}
let floor = peak * fl(prec, DECAY_FLOOR);
let start = centre.abs().max(1.0) * 1e-9;
let mut edges = [0.0f64; 2];
for (i, sign) in [-1.0f64, 1.0].iter().enumerate() {
let mut r = start;
let mut settled = false;
for _ in 0..140 {
let v = probe(centre + sign * r).unwrap_or_else(|| fl(prec, 0.0));
if v < floor {
settled = true;
break;
}
r *= 2.0;
}
if !settled {
return Located::NotFound;
}
edges[i] = centre + sign * r;
}
Located::At(Window {
centre: fl(prec, (edges[0] + edges[1]) / 2.0),
scale: fl(prec, (edges[1] - edges[0]) / 2.0),
})
}
fn tanh_sinh(
g: &mut impl FnMut(&Float) -> Sample,
lo: &Float,
hi: &Float,
prec: u32,
) -> QuadOutcome {
let half_width = mul(&Float::with_val(prec, hi - lo), &fl(prec, 0.5), prec);
let pi_half = mul(&pi_f(prec), &fl(prec, 0.5), prec);
let (lo, hi) = (lo.clone(), hi.clone());
let node = move |t: &Float| -> (Float, Float) {
let u = mul(&pi_half, &Float::with_val(prec, t.clone().sinh()), prec);
let two_u = mul(&u, &fl(prec, 2.0), prec);
let x = if u <= 0 {
let d = Float::with_val(
prec,
mul(&half_width, &fl(prec, 2.0), prec)
/ Float::with_val(prec, 1.0 + Float::with_val(prec, (-two_u.clone()).exp())),
);
Float::with_val(prec, &lo + &d)
} else {
let d = Float::with_val(
prec,
mul(&half_width, &fl(prec, 2.0), prec)
/ Float::with_val(prec, 1.0 + Float::with_val(prec, two_u.clone().exp())),
);
Float::with_val(prec, &hi - &d)
};
let ch = Float::with_val(prec, u.clone().cosh());
let num = mul(
&mul(&half_width, &pi_half, prec),
&Float::with_val(prec, t.clone().cosh()),
prec,
);
let w = Float::with_val(prec, &num / mul(&ch, &ch, prec));
(x, w)
};
de_levels(g, &node, prec)
}
fn mul(a: &Float, b: &Float, prec: u32) -> Float {
Float::with_val(prec, a * b)
}
fn pi_f(prec: u32) -> Float {
Float::with_val(prec, rug::float::Constant::Pi)
}
fn sinh_sinh(
g: &mut impl FnMut(&Float) -> Sample,
centre: &Float,
scale: &Float,
prec: u32,
log_scale: bool,
) -> QuadOutcome {
let pi_half = mul(&pi_f(prec), &fl(prec, 0.5), prec);
let node = |t: &Float| -> (Float, Float) {
let u = mul(&pi_half, &Float::with_val(prec, t.clone().sinh()), prec);
let s = Float::with_val(
prec,
centre + mul(scale, &Float::with_val(prec, u.clone().sinh()), prec),
);
let w = mul(
&mul(scale, &pi_half, prec),
&mul(
&Float::with_val(prec, t.clone().cosh()),
&Float::with_val(prec, u.clone().cosh()),
prec,
),
prec,
);
if log_scale {
let x = Float::with_val(prec, s.clone().exp());
let w = mul(&w, &x, prec);
(x, w)
} else {
(s, w)
}
};
de_levels(g, &node, prec)
}
fn within_magnitude(x: &Float) -> bool {
match x.get_exp() {
None => true, Some(e) => e.abs() <= 4096,
}
}
fn de_levels(
g: &mut impl FnMut(&Float) -> Sample,
node: &impl Fn(&Float) -> (Float, Float),
prec: u32,
) -> QuadOutcome {
let mut acc = fl(prec, 0.0);
let mut max_term = fl(prec, 0.0);
let mut edge_term = fl(prec, 0.0);
let mut outermost = 0.0f64;
let mut any_value = false;
let mut prev: Option<Float> = None;
for level in 0..LEVELS {
let h = 0.5f64 / f64::from(1u32 << level);
let n = (T_MAX / h) as i64;
let step = if level == 0 { 1 } else { 2 };
let start = if level == 0 { -n } else { -n + (n + 1) % 2 };
let mut k = start;
while k <= n {
let t = fl(prec, k as f64 * h);
k += step;
let (x, w) = node(&t);
if !x.is_finite() || !w.is_finite() || !within_magnitude(&x) {
continue;
}
match g(&x) {
Sample::Value(v) => {
any_value = true;
let term = Float::with_val(prec, &w * &v);
if !term.is_finite() {
return QuadOutcome::Divergent;
}
let mut mag = term.clone();
mag.abs_mut();
if mag > max_term {
max_term = mag.clone();
}
let abs_t = t.to_f64().abs();
if abs_t > outermost {
outermost = abs_t;
edge_term = mag.clone();
} else if abs_t >= outermost && mag > edge_term {
edge_term = mag.clone();
}
acc += term;
}
Sample::Unevaluable => {}
}
}
if !any_value {
return QuadOutcome::Inconclusive;
}
let value = Float::with_val(prec, &acc * fl(prec, h));
if max_term > 0 && edge_term > Float::with_val(prec, &max_term * fl(prec, 1e-20)) {
return QuadOutcome::Divergent;
}
if let Some(p) = &prev {
let diff = Float::with_val(prec, &value - p).abs();
let denom = {
let mut a = value.clone();
a.abs_mut();
if a < 1 {
a = fl(prec, 1.0);
}
a
};
if diff <= Float::with_val(prec, &denom * fl(prec, CONVERGED)) {
return QuadOutcome::Value {
value,
est_err: diff,
};
}
}
prev = Some(value);
}
QuadOutcome::Inconclusive
}
#[cfg(test)]
mod tests {
use super::*;
fn run(f: impl Fn(&Float) -> Float, region: Region) -> Option<Float> {
let mut g = |v: &Float| Sample::Value(f(v));
match quadrature(&mut g, region, VERIFY_PREC) {
QuadOutcome::Value { value, .. } => Some(value),
QuadOutcome::Divergent => {
println!("DIVERGENT");
None
}
QuadOutcome::Inconclusive => {
println!("INCONCLUSIVE");
None
}
}
}
fn agrees(v: &Float, expected: &Float) -> bool {
let d = Float::with_val(VERIFY_PREC, v - expected).abs();
d < Float::with_val(VERIFY_PREC, 1e-30)
}
fn gauss(x: &Float) -> Float {
let e = mul(&mul(x, x, VERIFY_PREC), &fl(VERIFY_PREC, -0.5), VERIFY_PREC);
Float::with_val(VERIFY_PREC, e.exp())
}
#[test]
fn gaussian_over_the_line_is_root_two_pi() {
let v = run(gauss, Region::Real).expect("should converge");
let expected = Float::with_val(
VERIFY_PREC,
mul(&pi_f(VERIFY_PREC), &fl(VERIFY_PREC, 2.0), VERIFY_PREC).sqrt(),
);
assert!(agrees(&v, &expected), "{} vs {}", v, expected);
}
#[test]
fn a_shifted_gaussian_is_found_and_integrated() {
let v = run(
|x| {
let d = Float::with_val(VERIFY_PREC, x - fl(VERIFY_PREC, 40.0));
let e = mul(
&mul(&d, &d, VERIFY_PREC),
&fl(VERIFY_PREC, -0.125),
VERIFY_PREC,
);
Float::with_val(VERIFY_PREC, e.exp())
},
Region::Real,
)
.expect("should converge");
let expected = Float::with_val(
VERIFY_PREC,
mul(&pi_f(VERIFY_PREC), &fl(VERIFY_PREC, 8.0), VERIFY_PREC).sqrt(),
);
assert!(agrees(&v, &expected), "{} vs {}", v, expected);
}
#[test]
fn a_divergent_integrand_is_reported_divergent_not_valued() {
let mut g = |v: &Float| {
let e = mul(&mul(v, v, VERIFY_PREC), &fl(VERIFY_PREC, 0.5), VERIFY_PREC);
Sample::Value(Float::with_val(VERIFY_PREC, e.exp()))
};
assert!(matches!(
quadrature(&mut g, Region::Real, VERIFY_PREC),
QuadOutcome::Divergent
));
}
}