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

//! Two-dimensional interpolation.
//!
//! # Example
//! ```
//! use interp::interp2d::Interp2D;
//!
//! let grid = ndarray::array![
//!     [0.0f64, 0.0, 0.0],
//!     [0.0, 1.0, 0.0],
//!     [0.0, 0.0, 0.0]
//! ];
//! let xs = vec![0.0, 1.0, 2.0];
//! let ys = vec![3.0, 4.0, 5.0];
//!
//! // Create interpolator struct.
//! let interp = Interp2D::new(xs, ys, grid);
//!
//! // Evaluate the interpolated data at some point.
//! assert!((interp.eval_no_extrapolation((1.0, 4.0)).unwrap() - 1.0).abs() < 1e-6);
//! ```

use crate::{find_closest_neighbours_indices, is_monotonic};
use ndarray::{ArrayBase, Axis, Data, Ix2, OwnedRepr};
use std::ops::{Add, Div, Mul, Sub};

// /// Define the extrapolation behaviour when a point outside of the sample grid should
// /// be evaluated.
// #[derive(Copy, Clone, Debug)]
// pub enum ExtrapolationMode {
//     /// Don't extrapolate and panic.
//     Panic,
//     /// Stay at the closest sample value.
//     Constant,
//     /// Extrapolate with closest sample derivative.
//     Linear,
// }

/// Two dimensional bilinear interpolator for data on an irregular grid.
///
/// * `C1`: Coordinate type of first axis.
/// * `C2`: Coordinate type of second axis.
/// * `Z`: Value type.
/// * `S`: Array representation of `Z` values.
#[derive(Debug)]
pub struct Interp2D<C1, C2, Z, S>
where
    S: Data<Elem = Z>,
{
    x: Vec<C1>,
    y: Vec<C2>,
    z: ArrayBase<S, Ix2>,
}

impl<C1, C2, Z, S> Clone for Interp2D<C1, C2, Z, S>
where
    C1: Clone,
    C2: Clone,
    S: Data<Elem = Z>,
    ArrayBase<S, Ix2>: Clone,
{
    fn clone(&self) -> Self {
        Self {
            x: self.x.clone(),
            y: self.y.clone(),
            z: self.z.clone(),
        }
    }
}

/// Interpolate over a two dimensional *normalized* grid cell.
/// The sample values (`xij`) are located on the corners of the normalized quadratic grid cell:
///
/// ```txt
/// [[x10, x11]
///  [x00, x01]]
/// ```
///
/// The normalized grid cell starts at `(0, 0)` and ends at `(1, 1)`
///
/// `alpha` and `beta` are the normalized coordinates.
/// The parameters `alpha` and `beta` give the position of the interpolation point.
///
/// `alpha` and `beta` should range from `0.0` to `1.0` for interpolation, otherwise
/// the value is *extrapolated*.
fn interpolate2d_bilinear<C, Z>(x00: Z, x10: Z, x01: Z, x11: Z, alpha: C, beta: C) -> Z
where
    C: Copy + Add<Output = C> + Sub<Output = C>,
    Z: Copy + Mul<C, Output = Z> + Add<Output = Z> + Sub<Output = Z>,
{
    // debug_assert!(alpha >= C::zero() && alpha <= C::one(), "Cannot extrapolate.");
    // debug_assert!(beta >= C::zero() && beta <= C::one(), "Cannot extrapolate.");
    x00 + (x10 - x00) * alpha + (x01 - x00) * beta + (x00 + x11 - x10 - x01) * alpha * beta
}

#[test]
fn test_interpolate2d() {
    // Tolerance for test.
    let tol = 1e-6f64;

    assert!((interpolate2d_bilinear(1.0f64, 2., 3., 4., 0., 0.) - 1.).abs() < tol);
    assert!((interpolate2d_bilinear(1.0f64, 2., 3., 4., 1., 0.) - 2.).abs() < tol);
    assert!((interpolate2d_bilinear(1.0f64, 2., 3., 4., 0., 1.) - 3.).abs() < tol);
    assert!((interpolate2d_bilinear(1.0f64, 2., 3., 4., 1., 1.) - 4.).abs() < tol);
    assert!((interpolate2d_bilinear(0.0f64, 1., 1., 0., 0.5, 0.5) - 0.5).abs() < tol);
}

/// Find the value of `f` at `(x, y)`
/// given four corner values `vij = f(xi, yj)` for all `(i, j) in [0, 1] x [0, 1]`.
fn interp2d<C1, C2, Z>(
    (x, y): (C1, C2),
    (x0, x1): (C1, C1),
    (y0, y1): (C2, C2),
    (v00, v10, v01, v11): (Z, Z, Z, Z),
) -> Z
where
    C1: Copy + Sub<Output = C1> + Div,
    C2: Copy + Sub<Output = C2> + Div<Output = <C1 as Div>::Output>,
    Z: Copy + Mul<<C1 as Div>::Output, Output = Z> + Add<Output = Z> + Sub<Output = Z>,
    <C1 as Div>::Output:
        Copy + Add<Output = <C1 as Div>::Output> + Sub<Output = <C1 as Div>::Output>,
{
    let dx = x1 - x0;
    let dy = y1 - y0;

    let alpha = (x - x0) / dx;
    let beta = (y - y0) / dy;

    interpolate2d_bilinear(v00, v10, v01, v11, alpha, beta)
}

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

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

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

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

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

    /// Get the y-coordinate values.
    pub fn ys(&self) -> &Vec<C2> {
        &self.y
    }

    /// Get the raw z values.
    pub fn z(&self) -> &ArrayBase<S, Ix2> {
        &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) -> Interp2D<C1, C2, Z2, OwnedRepr<Z2>>
    where
        C1: PartialOrd + Clone,
        C2: PartialOrd + Clone,
    {
        Interp2D {
            x: self.x.clone(),
            y: self.y.clone(),
            z: self.z.map(f),
        }
    }

    /// Apply the function `f` to each x coordinate. Panics if the ordering of the coordinates
    /// is not monotonic after applyint `f`.
    pub fn map_x_axis<Xnew>(self, f: impl Fn(C1) -> Xnew) -> Interp2D<Xnew, C2, Z, S>
    where
        Xnew: PartialOrd,
    {
        let xnew = self.x.into_iter().map(f).collect();
        assert!(is_monotonic(&xnew));
        Interp2D {
            x: xnew,
            y: self.y,
            z: self.z,
        }
    }
    /// Apply the function `f` to each y coordinate. Panics if the ordering of the coordinates
    /// is not monotonic after applyint `f`.
    pub fn map_y_axis<Ynew>(self, f: impl Fn(C2) -> Ynew) -> Interp2D<C1, Ynew, Z, S>
    where
        Ynew: PartialOrd,
    {
        let ynew = self.y.into_iter().map(f).collect();
        assert!(is_monotonic(&ynew));
        Interp2D {
            x: self.x,
            y: ynew,
            z: self.z,
        }
    }
}

impl<C1, C2, Z, S> Interp2D<C1, C2, Z, S>
where
    C1: Copy + Sub<Output = C1> + Div + PartialOrd,
    C2: Copy + Sub<Output = C2> + Div<Output = <C1 as Div>::Output> + PartialOrd,
    Z: Copy + Mul<<C1 as Div>::Output, Output = Z> + Add<Output = Z> + Sub<Output = Z>,
    <C1 as Div>::Output:
        Copy + Add<Output = <C1 as Div>::Output> + Sub<Output = <C1 as Div>::Output>,
    S: Data<Elem = Z> + Clone,
{
    /// Evaluate the sampled function by interpolation at `(x, y)`.
    ///
    /// If `(x, y)` 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, y): (C1, C2)) -> Z {
        // Find closest grid points.
        let (x0, x1) = find_closest_neighbours_indices(&self.x, x);
        let (y0, y1) = find_closest_neighbours_indices(&self.y, y);

        interp2d(
            (x, y),
            (self.x[x0], self.x[x1]),
            (self.y[y0], self.y[y1]),
            (
                self.z[[x0, y0]],
                self.z[[x1, y0]],
                self.z[[x0, y1]],
                self.z[[x1, y1]],
            ),
        )
    }

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

impl<C1, C2, Z, S> Interp2D<C1, C2, Z, S>
where
    Z: Clone,
    S: Data<Elem = Z> + Clone,
    OwnedRepr<Z>: Data<Elem = Z>,
{
    /// Swap the input variables.
    pub fn swap_variables(self) -> Interp2D<C2, C1, Z, OwnedRepr<Z>> {
        let Self { x, y, z } = self;
        Interp2D {
            x: y,
            y: x,
            z: z.t().to_owned(),
        }
    }
}

#[test]
fn test_interp2d_on_view() {
    let grid = ndarray::array![[0.0f64, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]];
    let xs = vec![0.0, 1.0, 2.0];
    let ys = vec![3.0, 4.0, 5.0];

    let interp = Interp2D::new(xs, ys, grid.view());

    // Tolerance for test.
    let tol = 1e-6f64;

    assert!((interp.eval_no_extrapolation((1.0, 4.0)).unwrap() - 1.0).abs() < tol);
}

#[test]
fn test_interp2d() {
    let grid = ndarray::array![[0.0f64, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]];
    let xs = vec![0.0, 1.0, 2.0];
    let ys = vec![3.0, 4.0, 5.0];

    let interp = Interp2D::new(xs, ys, grid);

    // Tolerance for test.
    let tol = 1e-6f64;

    assert!((interp.eval_no_extrapolation((1.0, 4.0)).unwrap() - 1.0).abs() < tol);
    assert!((interp.eval_no_extrapolation((0.0, 3.0)).unwrap() - 0.0).abs() < tol);
    assert!((interp.eval_no_extrapolation((0.5, 3.5)).unwrap() - 0.25).abs() < tol);

    // // Swap input variables.
    // let interp = interp.swap_variables();
    //
    // assert!((interp.eval_no_extrapolation((4.0, 1.0)).unwrap() - 1.0).abs() < tol);
    // assert!((interp.eval_no_extrapolation((3.0, 0.0)).unwrap() - 0.0).abs() < tol);
    // assert!((interp.eval_no_extrapolation((3.5, 0.5)).unwrap() - 0.25).abs() < tol);
}

#[test]
fn test_map_values() {
    let grid = ndarray::array![[0.0f64, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]];
    let xs = vec![0.0, 1.0, 2.0];
    let ys = vec![3.0, 4.0, 5.0];

    let interp = Interp2D::new(xs, ys, grid);

    let interp_doubled = interp.map_values(|&v| 2. * v);

    // Tolerance for test.
    let tol = 1e-6;

    assert!((interp_doubled.eval_no_extrapolation((1.0, 4.0)).unwrap() - 2.0).abs() < tol);
}

#[test]
fn test_map_x_axis() {
    let grid = ndarray::array![[0.0f64, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]];
    let xs = vec![0.0, 1.0, 2.0];
    let ys = vec![3.0, 4.0, 5.0];

    let interp = Interp2D::new(xs, ys, grid);

    let shift = 1.0;
    let interp_shifted = interp.map_x_axis(|x| x + shift);

    // Tolerance for test.
    let tol = 1e-6;

    assert!(
        (interp_shifted
            .eval_no_extrapolation((1.0 + shift, 4.0))
            .unwrap()
            - 1.0)
            .abs()
            < tol
    );
}

#[test]
fn test_map_y_axis() {
    let grid = ndarray::array![[0.0f64, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]];
    let xs = vec![0.0, 1.0, 2.0];
    let ys = vec![3.0, 4.0, 5.0];

    let interp = Interp2D::new(xs, ys, grid);

    let shift = 1.0;
    let interp_shifted = interp.map_y_axis(|y| y + shift);

    // Tolerance for test.
    let tol = 1e-6;

    assert!(
        (interp_shifted
            .eval_no_extrapolation((1.0, 4.0 + shift))
            .unwrap()
            - 1.0)
            .abs()
            < tol
    );
}