probl_engine/
quadrature.rs

1//! Adaptive numerical integration on bounded intervals, for `@integrate`
2//! (docs/semantics.md, section 13): the 21-point Gauss–Kronrod rule of
3//! QUADPACK, bisecting the part with the largest error estimate until the
4//! estimate is within the tolerance. An estimate isn't a bound: a feature
5//! narrower than the rule's nodes can go unseen, which is why callers give
6//! the places where the integrand has features as breakpoints.
7
8use crate::dist::Budget;
9use crate::error::{OpError, OpResult};
10
11// The rule's constants, as QUADPACK publishes them.
12
13/// The Kronrod nodes on [-1, 1], from the outside in; the odd ones are the
14/// 10-point Gauss nodes.
15#[allow(clippy::excessive_precision)]
16const NODES: [f64; 11] = [
17    0.995657163025808080735527280689003,
18    0.973906528517171720077964012084452,
19    0.930157491355708226001207180059508,
20    0.865063366688984510732096688423493,
21    0.780817726586416897063717578345042,
22    0.679409568299024406234327365114874,
23    0.562757134668604683339000099272694,
24    0.433395394129247190799265943165784,
25    0.294392862701460198131126603103866,
26    0.148874338981631210884826001129720,
27    0.0,
28];
29
30#[allow(clippy::excessive_precision)]
31const KRONROD: [f64; 11] = [
32    0.011694638867371874278064396062192,
33    0.032558162307964727478818972459390,
34    0.054755896574351996031381300244580,
35    0.075039674810919952767043140916190,
36    0.093125454583697605535065465083366,
37    0.109387158802297641899210590325805,
38    0.123491976262065851077958109831074,
39    0.134709217311473325928054001771707,
40    0.142775938577060080797094273138717,
41    0.147739104901338491374841515972068,
42    0.149445554002916905664936468389821,
43];
44
45#[allow(clippy::excessive_precision)]
46const GAUSS: [f64; 5] = [
47    0.066671344308688137593568809893332,
48    0.149451349150580593145776339657697,
49    0.219086362515982043995534934228163,
50    0.269266719309996355091226921569469,
51    0.295524224714752870173892994651338,
52];
53
54/// The most parts one integral is bisected into.
55pub const MAX_PARTS: usize = 4_000;
56
57/// An integral's value and its error estimate.
58#[derive(Clone, Copy, Debug, PartialEq)]
59pub struct Estimate {
60    pub value: f64,
61    pub error: f64,
62}
63
64/// The rule on `[a, b]`: the Kronrod estimate, and its difference from the
65/// Gauss one as the error. Each evaluation costs a unit of work.
66pub fn rule(f: &mut impl FnMut(f64) -> f64, a: f64, b: f64, budget: &mut Budget) -> OpResult<Estimate> {
67    budget.work(21)?;
68    let center = 0.5 * (a + b);
69    let half = 0.5 * (b - a);
70    let fc = f(center);
71    let mut kronrod = KRONROD[10] * fc;
72    let mut gauss = 0.0;
73    for (i, &x) in NODES[..10].iter().enumerate() {
74        let pair = f(center - half * x) + f(center + half * x);
75        kronrod += KRONROD[i] * pair;
76        if i % 2 == 1 {
77            gauss += GAUSS[i / 2] * pair;
78        }
79    }
80    let value = kronrod * half;
81    let error = ((kronrod - gauss) * half).abs();
82    Ok(Estimate { value, error })
83}
84
85/// One part of an integral: its interval and its estimate.
86#[derive(Clone, Copy, Debug)]
87pub struct Part {
88    pub a: f64,
89    pub b: f64,
90    pub estimate: Estimate,
91}
92
93/// `f` integrated over the intervals between consecutive `breaks`, by parts
94/// bisected where the error is largest until the total error is at most
95/// `tolerance` times the total, and `absolute` (the scale below which an
96/// error doesn't matter). The parts, in order, or an error if the
97/// tolerance can't be met, or a value isn't finite.
98pub fn adaptive(
99    f: &mut impl FnMut(f64) -> f64,
100    breaks: &[f64],
101    tolerance: f64,
102    absolute: f64,
103    budget: &mut Budget,
104) -> OpResult<Vec<Part>> {
105    let mut parts: Vec<Part> = Vec::with_capacity(breaks.len() + 16);
106    for w in breaks.windows(2) {
107        if w[1] > w[0] {
108            parts.push(Part {
109                a: w[0],
110                b: w[1],
111                estimate: rule(f, w[0], w[1], budget)?,
112            });
113        }
114    }
115    loop {
116        let total: f64 = parts.iter().map(|p| p.estimate.value).sum();
117        let error: f64 = parts.iter().map(|p| p.estimate.error).sum();
118        if !total.is_finite() || !error.is_finite() {
119            return Err(OpError::unsupported(
120                "numerical integration didn't converge: it met a value that isn't finite",
121            )
122            .help("the integral may not exist; or use `@mode sample(runs: 10_000)`"));
123        }
124        if error <= (tolerance * total.abs()).max(absolute) {
125            break;
126        }
127        if parts.len() >= MAX_PARTS {
128            return Err(not_converged(tolerance));
129        }
130        let (i, _) = parts
131            .iter()
132            .enumerate()
133            .max_by(|x, y| x.1.estimate.error.total_cmp(&y.1.estimate.error))
134            .expect("parts");
135        let Part { a, b, .. } = parts[i];
136        let mid = 0.5 * (a + b);
137        if mid <= a || mid >= b {
138            return Err(not_converged(tolerance));
139        }
140        let left = Part {
141            a,
142            b: mid,
143            estimate: rule(f, a, mid, budget)?,
144        };
145        let right = Part {
146            a: mid,
147            b,
148            estimate: rule(f, mid, b, budget)?,
149        };
150        parts[i] = left;
151        parts.insert(i + 1, right);
152    }
153    Ok(parts)
154}
155
156pub fn not_converged(tolerance: f64) -> OpError {
157    OpError::unsupported(format!(
158        "numerical integration didn't converge to its relative tolerance of {tolerance:e}"
159    ))
160    .help("loosen it, like `@integrate 1e-6`, or use `@mode sample(runs: 10_000)`")
161}
162
163/// `f` over `[a, b]`, to `tolerance`.
164pub fn integral(
165    f: &mut impl FnMut(f64) -> f64,
166    a: f64,
167    b: f64,
168    tolerance: f64,
169    absolute: f64,
170    budget: &mut Budget,
171) -> OpResult<Estimate> {
172    let parts = adaptive(f, &[a, b], tolerance, absolute, budget)?;
173    Ok(Estimate {
174        value: parts.iter().map(|p| p.estimate.value).sum(),
175        error: parts.iter().map(|p| p.estimate.error).sum(),
176    })
177}
178
179/// `f` over `[a, b]`, to `tolerance` relative to the integral of `|f|`, so
180/// that an integral near zero, like a centred moment's, converges: `None` if
181/// it doesn't.
182pub fn moment(f: impl Fn(f64) -> f64, a: f64, b: f64, tolerance: f64) -> Option<f64> {
183    let mut budget = Budget::unlimited();
184    let scale = integral(&mut |u| f(u).abs(), a, b, tolerance, 0.0, &mut budget).ok()?;
185    if scale.value == 0.0 {
186        return Some(0.0);
187    }
188    integral(&mut |u| f(u), a, b, tolerance, tolerance * scale.value, &mut budget)
189        .ok()
190        .map(|e| e.value)
191}
192
193/// An integral over [0, 1], by parts with their running totals, to answer
194/// for any interval and to invert.
195#[derive(Clone, Debug)]
196pub struct Table {
197    parts: Vec<Part>,
198    /// The total before each part.
199    before: Vec<f64>,
200}
201
202impl Table {
203    pub fn new(parts: Vec<Part>) -> Table {
204        let mut before = Vec::with_capacity(parts.len());
205        let mut sum = 0.0;
206        for p in &parts {
207            before.push(sum);
208            sum += p.estimate.value;
209        }
210        Table { parts, before }
211    }
212
213    pub fn total(&self) -> f64 {
214        self.parts
215            .last()
216            .map_or(0.0, |p| self.before[self.parts.len() - 1] + p.estimate.value)
217    }
218
219    /// The part that holds `u`.
220    fn find(&self, u: f64) -> usize {
221        self.parts.partition_point(|p| p.b <= u).min(self.parts.len() - 1)
222    }
223
224    /// The integral from 0 to `u`: whole parts from the totals, and the
225    /// rule on the rest of the part `u` is in.
226    pub fn up_to(&self, f: &mut impl FnMut(f64) -> f64, u: f64, budget: &mut Budget) -> OpResult<f64> {
227        if u <= 0.0 {
228            return Ok(0.0);
229        }
230        if u >= 1.0 {
231            return Ok(self.total());
232        }
233        let i = self.find(u);
234        let p = self.parts[i];
235        let partial = if u <= p.a {
236            0.0
237        } else if u >= p.b {
238            p.estimate.value
239        } else {
240            rule(f, p.a, u, budget)?.value
241        };
242        Ok(self.before[i] + partial)
243    }
244
245    /// The integral from `a` to `b`.
246    pub fn between(&self, f: &mut impl FnMut(f64) -> f64, a: f64, b: f64, budget: &mut Budget) -> OpResult<f64> {
247        if b <= a {
248            return Ok(0.0);
249        }
250        Ok((self.up_to(f, b, budget)? - self.up_to(f, a, budget)?).max(0.0))
251    }
252
253    /// The `u` with the integral from 0 to it `target`: in its part, by
254    /// Newton's steps, since `f` is the integral's derivative, kept within
255    /// a bracket that bisection narrows when a step leaves it.
256    pub fn inverse(&self, f: &mut impl FnMut(f64) -> f64, target: f64, budget: &mut Budget) -> OpResult<f64> {
257        if target <= 0.0 {
258            return Ok(self.parts.first().map_or(0.0, |p| p.a));
259        }
260        if target >= self.total() {
261            return Ok(self.parts.last().map_or(1.0, |p| p.b));
262        }
263        let i = self.before.partition_point(|&b| b <= target).saturating_sub(1);
264        let p = self.parts[i];
265        let need = target - self.before[i];
266        let mass = p.estimate.value;
267        let (mut lo, mut hi) = (p.a, p.b);
268        let mut u = p.a + (p.b - p.a) * (need / mass).clamp(0.0, 1.0);
269        for _ in 0..60 {
270            let diff = rule(f, p.a, u, budget)?.value - need;
271            if diff.abs() <= 1e-12 * mass {
272                break;
273            }
274            if diff < 0.0 {
275                lo = u;
276            } else {
277                hi = u;
278            }
279            let slope = f(u);
280            let step = if slope > 0.0 { u - diff / slope } else { f64::NAN };
281            u = if step > lo && step < hi { step } else { 0.5 * (lo + hi) };
282            if hi - lo <= f64::EPSILON * hi.abs().max(1e-300) {
283                break;
284            }
285        }
286        Ok(u)
287    }
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293
294    fn budget() -> Budget {
295        Budget::unlimited()
296    }
297
298    #[test]
299    fn smooth_integrals_are_exact_to_rounding() {
300        let e = integral(&mut |x: f64| x.exp(), 0.0, 1.0, 1e-12, 0.0, &mut budget()).unwrap();
301        assert!((e.value - (1f64.exp() - 1.0)).abs() < 1e-14, "{e:?}");
302        // A narrow peak, found from breakpoints graded away from it.
303        let mut peak = |x: f64| (-((x - 0.3) / 1e-4).powi(2)).exp();
304        let mut breaks = vec![0.0, 0.3, 1.0];
305        for k in 0..12 {
306            let d = 1e-4 * 2f64.powi(k);
307            breaks.extend([0.3 - d, 0.3 + d]);
308        }
309        breaks.retain(|b| (0.0..=1.0).contains(b));
310        breaks.sort_by(f64::total_cmp);
311        let parts = adaptive(&mut peak, &breaks, 1e-10, 0.0, &mut budget()).unwrap();
312        let total: f64 = parts.iter().map(|p| p.estimate.value).sum();
313        assert!((total - 1e-4 * std::f64::consts::PI.sqrt()).abs() < 1e-13, "{total}");
314    }
315
316    #[test]
317    fn tables_answer_for_intervals_and_invert() {
318        let mut f = |x: f64| 2.0 * x;
319        let parts = adaptive(&mut f, &[0.0, 1.0], 1e-12, 0.0, &mut budget()).unwrap();
320        let t = Table::new(parts);
321        let mut b = budget();
322        assert!((t.total() - 1.0).abs() < 1e-14);
323        assert!((t.between(&mut f, 0.25, 0.5, &mut b).unwrap() - 0.1875).abs() < 1e-14);
324        assert!((t.inverse(&mut f, 0.25, &mut b).unwrap() - 0.5).abs() < 1e-12);
325    }
326
327    #[test]
328    fn what_doesnt_converge_says_so() {
329        let e = integral(&mut |x: f64| 1.0 / x, 0.0, 1.0, 1e-9, 0.0, &mut budget()).unwrap_err();
330        assert!(e.message.contains("numerical integration didn't converge"), "{e:?}");
331    }
332}
probl_engine/quadrature.rs

probl_engine/
quadrature.rs
