libreda-interp 0.0.3

// Copyright (c) 2021-2021 Thomas Kramer.
// SPDX-FileCopyrightText: 2022 Thomas Kramer <code@tkramer.ch>
//
// SPDX-License-Identifier: GPL-3.0-or-later

//! One-dimensional linear interpolation.
//!
//! # Examples
//! ```
//! use interp::interp1d::Interp1D;
//!
//! let xs = vec![0.0, 1.0, 2.0];
//! let zs = vec![0.0f64, 1.0, 0.0];
//!
//! let interp = Interp1D::new(xs, zs);
//!
//! assert!((interp.eval_no_extrapolation(1.0).unwrap() - 1.0).abs() < 1e-6);
//! ```

use crate::{find_closest_neighbours_indices, is_monotonic};
use std::ops::{Add, Div, Mul, Sub};

///
/// * `C`: Coordinate type.
/// * `Z`: Value type.
#[derive(Debug, Clone)]
pub struct Interp1D<C, Z> {
    /// Index.
    x: Vec<C>,
    /// Samples.
    z: Vec<Z>,
}

/// Interpolate between two values `x0` and `x1`.
/// `alpha` should range from `0.0` to `1.0` for interpolation, otherwise
/// the value is *extrapolated*.
fn interpolate1d<C, Z>(z0: Z, z1: Z, alpha: C) -> Z
where
    C: Copy + Sub<Output = C>,
    Z: Copy + Mul<C, Output = Z> + Add<Output = Z> + Sub<Output = Z>,
{
    z0 - z0 * alpha + z1 * alpha
}

#[test]
fn test_interpolate1d() {
    assert!((interpolate1d(1.0f64, 2., 0.) - 1.).abs() < 1e-6);
    assert!((interpolate1d(1.0f64, 2., 1.) - 2.).abs() < 1e-6);
    assert!((interpolate1d(1.0f64, 2., 0.5) - 1.5).abs() < 1e-6);
}

/// Find the value of `f(x)`
/// given two sample values `vi = f(xi)` for all `i in [0, 1]`.
fn interp1d<C, Z>(x: C, (x0, x1): (C, C), (v0, v1): (Z, Z)) -> Z
where
    C: Copy + Sub<Output = C> + Div,
    Z: Copy + Mul<<C as Div>::Output, Output = Z> + Add<Output = Z> + Sub<Output = Z>,
    <C as Div>::Output: Copy + Add<Output = <C as Div>::Output> + Sub<Output = <C as Div>::Output>,
{
    let dx = x1 - x0;

    let alpha = (x - x0) / dx;

    interpolate1d(v0, v1, alpha)
}

impl<C, Z> Interp1D<C, Z> {
    /// Create a new interpolation engine.
    ///
    /// Interpolates values which are sampled on a rectangular grid.
    ///
    /// # Parameters
    /// * `x`: The x-coordinates. Must be monotonic.
    /// * `z`: The values `z(x)` for each grid point defined by the `x` coordinates.
    ///
    /// # Panics
    /// Panics when
    /// * dimensions of the indices and z-values don't match.
    /// * one axis is empty.
    /// * x values are not monotonic.
    pub fn new(x: Vec<C>, z: Vec<Z>) -> Self
    where
        C: PartialOrd,
    {
        assert_eq!(z.len(), x.len(), "x-axis length mismatch.");
        assert!(!x.is_empty());

        assert!(is_monotonic(&x), "x values must be monotonic.");

        Self { x, z }
    }

    /// Get the boundaries of the sample range
    /// as `(x0, x1)` tuple.
    pub fn bounds(&self) -> (C, C)
    where
        C: Copy,
    {
        (self.x[0], self.x[self.x.len() - 1])
    }

    /// Check if the `x` coordinate lies within the defined sample range.
    pub fn is_within_bounds(&self, x: C) -> bool
    where
        C: PartialOrd + Copy,
    {
        let (x0, x1) = self.bounds();
        x0 <= x && x <= x1
    }

    /// Get the x-coordinate values.
    pub fn xs(&self) -> &Vec<C> {
        &self.x
    }

    /// Get the raw z values.
    pub fn z(&self) -> &Vec<Z> {
        &self.z
    }

    /// Create a new interpolated table by applying a function elementwise to the values.
    pub fn map_values<Z2>(&self, f: impl Fn(&Z) -> Z2) -> Interp1D<C, Z2>
    where
        C: PartialOrd + Clone,
    {
        Interp1D {
            x: self.x.clone(),
            z: self.z.iter().map(f).collect(),
        }
    }

    /// Apply a function to the coordinates values of the axis.
    /// Panics if the new coordinates are not monotonically increasing.
    pub fn map_axis<Xnew>(self, f: impl Fn(C) -> Xnew) -> Interp1D<Xnew, Z>
    where
        Xnew: PartialOrd,
    {
        let xnew = self.x.into_iter().map(f).collect();
        assert!(is_monotonic(&xnew));
        Interp1D { x: xnew, z: self.z }
    }
}

impl<C, Z> Interp1D<C, Z>
where
    C: Copy + Sub<Output = C> + Div + PartialOrd,
    Z: Copy + Mul<<C as Div>::Output, Output = Z> + Add<Output = Z> + Sub<Output = Z>,
    <C as Div>::Output: Copy + Add<Output = <C as Div>::Output> + Sub<Output = <C as Div>::Output>,
{
    /// Evaluate the sampled function by interpolation at `x`.
    ///
    /// If `x` lies out of the sampled range then the function is silently *extrapolated*.
    /// The boundaries of the sample range can be queried with `bounds()`.
    pub fn eval(&self, x: C) -> Z {
        // Find closest grid points.

        let (x0, x1) = find_closest_neighbours_indices(&self.x, x);

        interp1d(x, (self.x[x0], self.x[x1]), (self.z[x0], self.z[x1]))
    }

    /// Returns the same value as `eval()` as long as `x` is within the
    /// range of the samples. Otherwise `None` is returned instead of an extrapolation.
    pub fn eval_no_extrapolation(&self, x: C) -> Option<Z> {
        if self.is_within_bounds(x) {
            Some(self.eval(x))
        } else {
            None
        }
    }
}

#[test]
fn test_interp1d() {
    let xs = vec![0.0f64, 1.0, 2.0];
    let zs = vec![0.0, 1.0, 0.0];

    let interp = Interp1D::new(xs, zs);

    assert!((interp.eval_no_extrapolation(1.0).unwrap() - 1.0).abs() < 1e-6);
    assert!((interp.eval_no_extrapolation(2.0).unwrap() - 0.0).abs() < 1e-6);
    assert!((interp.eval_no_extrapolation(1.5).unwrap() - 0.5).abs() < 1e-6);
}

#[test]
fn test_interp1d_map_values() {
    let xs = vec![0.0f64, 1.0, 2.0];
    let zs = vec![0.0, 1.0, 0.0];

    let interp = Interp1D::new(xs, zs);
    let interp = interp.map_values(|x| 2. * x);

    assert!((interp.eval_no_extrapolation(1.0).unwrap() - 2.0).abs() < 1e-6);
}

#[test]
fn test_interp1d_map_axis() {
    let xs = vec![0.0f64, 1.0, 2.0];
    let zs = vec![0.0, 1.0, 0.0];

    let interp = Interp1D::new(xs, zs);
    let interp = interp.map_axis(|x| 2. * x);

    assert!((interp.eval_no_extrapolation(2.0).unwrap() - 1.0).abs() < 1e-6);
}