romberg.rs

use core::ops::RangeInclusive;

pub fn integral<F: Fn(f64) -> f64>(range: &RangeInclusive<f64>, f: F, err: f64) -> f64 {
    let memo = [trapezoid_area(range, &f)];

    romberg(range, &f, 2, err, &memo)
}

fn romberg<F: Fn(f64) -> f64>(
    range: &RangeInclusive<f64>,
    f: &F,
    k: usize,
    err: f64,
    memo: &[f64],
) -> f64 {
    let mut new_memo = vec![0.0; k];

    new_memo[0] = trapezoid(range, &f, (2usize).pow(k as u32 - 1), memo[0]);

    (1..k).for_each(|i| {
        let numer = (4.0f64).powi(i as i32) * new_memo[i - 1] - memo[i - 1];
        let denom = (4.0f64).powi(i as i32) - 1.0;

        new_memo[i] = numer / denom;
    });

    let current = new_memo[k - 1];
    let past = memo[k - 2];

    if (current - past).abs() < err {
        current
    } else {
        romberg(range, f, k + 1, err, &new_memo)
    }
}

#[inline]
fn trapezoid<F: Fn(f64) -> f64>(range: &RangeInclusive<f64>, f: F, n: usize, memo: f64) -> f64 {
    let start = *range.start();
    let end = *range.end();

    let h = (end - start) / (n as f64);

    let y_internal = h
        * (1..n)
            .filter(|j| j & 1 != 0)
            .map(|j| start + h * (j as f64))
            .map(f)
            .sum::<f64>();

    0.5 * memo + y_internal
}

#[inline]
fn trapezoid_area<F: Fn(f64) -> f64>(range: &RangeInclusive<f64>, f: F) -> f64 {
    let x1 = *range.start();
    let x2 = *range.end();

    let y1 = f(x1);
    let y2 = f(x2);

    0.5 * (y1 + y2) * (x2 - x1)
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_integral() {
        let f = |x: f64| x * x;
        let range = 3.0..=9.0;

        let area = integral(&range, &f, std::f64::EPSILON);
        assert_eq!(area as i64, (729 - 27) / 3);
    }

    #[test]
    fn test_trapezoid() {
        let f = |x: f64| x;
        let range = 1.0..=3.0;

        let area = trapezoid_area(&range, &f);
        assert_eq!(area as i64, 4);

        let area = trapezoid(&range, &f, 2, area);
        assert_eq!(area as i64, 4);
    }
}