te1d/

func.rs

//! Describes real-valued functions defined on the real line.

use std::rc::Rc;

use ndarray::{Array, Array1, Array2, array};

use crate::calc::{ColVec};
use crate::linalg::{LinalgError, solve_gauss};


/// Temperature-dependent material properties must implement [`RealValuedFunction`] trait.
pub trait RealValuedFunction {
    /// Evaluate the function values at the given *increasing* points `incr_xs`.
    fn ats_incr(&self, incr_xs: &ColVec) -> ColVec;

    /// Evaluate the function value at a given point `x`.
    #[inline]
    fn at(&self, x: f64) -> f64 {
        let fs = self.ats_incr(&array![x]);
        fs[0]
    }

    /// Evaluate the function values at the given points `xs`.
    #[inline]
    fn ats(&self, xs: &ColVec) -> ColVec {
        xs.mapv(|e| { self.at(e) })
    }
}

/// Piecewise functions must implement [`RealValuedFunction`] trait and [`PiecewiseFunction`] trait.
pub trait PiecewiseFunction: RealValuedFunction {
    /// Evaluate the function value in a specific segment.
    fn at_segment(&self, segment_idx: usize, x: f64) -> f64;

    /// Return the function value at a beyond-left-end point.
    fn at_left(&self) -> f64;

    /// Return the function value at a beyond-right-end point.
    fn at_right(&self) -> f64;

    /// Evaluate function values at the given *increasing* points `input_incr_xs`.
    /// * `input_incr_xs` must be increasing.
    /// * `data_strict_incr_xs` must be *strictly increasing*.
    fn piecewise_ats_incr(&self, input_incr_xs: &ColVec, data_strict_incr_xs: &ColVec, data_fs: &ColVec) -> ColVec {
        let data_xs = data_strict_incr_xs;
        let input_xs = input_incr_xs;
        let n = data_xs.len();
        let (xmin, xmax): (f64, f64) = (data_xs[0], data_xs[n-1]);
        let mut result: ColVec = input_xs.to_owned();
        let mut x_prev: f64 = input_xs[0] - 1.0;
        let mut x_prev_segment_idx: usize = 0;
        let last_segment_idx: usize = data_xs.len()-2;

        for (i, &x) in input_xs.iter().enumerate() {
            if x < x_prev {
                panic!("`input_incr_xs` must be increasing!");
            }
            if x < xmin {
                result[i] = self.at_left();
                x_prev = x;
                continue;
            }
            if x > xmax {
                result[i] = self.at_right();
                x_prev = x;
                continue;
            }
            let x_segment_idx: usize = self.search_segment(data_xs, x, x_prev_segment_idx, last_segment_idx);
            let xl: f64 = data_xs[x_segment_idx];
            let xr: f64 = data_xs[x_segment_idx+1];
            if (x < xl) || (x > xr) {
                panic!("Impossible: check if `data_strict_incr_xs` is strictly increasing!");
            }

            if x == xl {
                result[i] = data_fs[x_segment_idx];
            }
            if x == xr {
                result[i] = data_fs[x_segment_idx+1];
            } else {
                result[i] = self.at_segment(x_segment_idx, x);
            }
            x_prev = x;
            x_prev_segment_idx = x_segment_idx;
        }

        result
    }

    /// Finds the segment in which the given number is contained.
    /// * Binary search is used.
    /// * `start_segment_idx` and `end_segment_idx` are the lowest and largest index of segment to be returned.
    #[inline]
    fn search_segment(&self, data_strict_incr_xs: &ColVec, x: f64, start_segment_idx: usize, end_segment_idx: usize) -> usize {
        let mut left_idx: usize = start_segment_idx;
        let mut right_idx: usize = end_segment_idx;
        let mut mid_idx: usize = (left_idx + right_idx)/2;

        for _ in 0..(end_segment_idx-start_segment_idx)+1 {
            mid_idx = (left_idx + right_idx)/2;
            if left_idx >= right_idx {
                break;
            }
            if x < data_strict_incr_xs[mid_idx] {
                right_idx = mid_idx-1;
            } else if x > data_strict_incr_xs[mid_idx+1] {
                left_idx = mid_idx+1;
            } else {
                break;
            }
        }

        mid_idx
    }

}


/// Use [`Polynomial`] to implement [`RealValuedFunction`] when the *coefficients* of a polynomial is given.
/// * There is no performance difference between [`Polynomial::ats_incr()`] and [`Polynomial::ats()`].
/// * The `coefs` is the coefficients of a polynomial. For example, `coefs=[a2, a1, a0]` means the quadratic polynomial `a2 x^2 + a1 x + a0`.
/// * [`Polynomial::deriv()`] gives the derivative of the polynomial.
/// * [`Polynomial::from_interp()`] interpolates data and return a polynomial. **Do not use** it for 7-or-more-degree polynomials. It is numerically unstable. Use [`BarycentricPolynomial`](crate::interp::BarycentricPolynomial) whenever possible.
/// 
/// # Examples
/// ```
/// use ndarray::prelude::*;
/// use te1d::prelude::*;
/// use te1d::func::Polynomial;
/// 
/// // create a new polynomial
/// let f = Polynomial::new(&array![1.5, -1.2, 0.5, 0.7]);
/// let result = f.ats(&array![-1.0, -0.5, 0.0, 0.3, 0.9]);
/// let sol = array![-2.5, -0.0375, 0.7, 0.7825, 1.2715];
/// assert!(all_close(&result, &sol, 1e-05, 1e-08));
/// assert!(is_close(f.at(0.75), 1.0328125, 1e-05, 1e-08));
/// assert_eq!(&result, &f.ats_incr(&array![-1.0, -0.5, 0.0, 0.3, 0.9]));
/// 
/// // find the derivative
/// let deriv = f.deriv();
/// assert_eq!(&deriv, &Polynomial::new(&array![4.5, -2.4, 0.5]));
/// 
/// // polynomial interpolation
/// let xs: ColVec = array![0.0, 0.5, 2.5, 3.0];
/// let fs: ColVec = array![1.0, 0.0, 2.0, 3.0];
/// let f = Polynomial::from_interp(&xs, &fs).unwrap();
/// assert!(all_close(&f.ats(&xs), &fs, 1e-05, 1e-08));
/// assert_eq!(f.deg, 3);
/// ```
#[derive(Debug, PartialEq, Clone)]
pub struct Polynomial {
    /// The coefficients of the polynomial.
    pub coefs: ColVec,
    /// The degree of the polynomial.
    pub deg: usize,
}

impl RealValuedFunction for Polynomial {
    fn at(&self, x: f64) -> f64 {
        let mut result: f64 = 0.0;
        for coef in self.coefs.iter() {
            result *= x;
            result += coef;
        }

        result
    }

    #[inline]
    fn ats_incr(&self, incr_xs: &ColVec) -> ColVec {
        self.ats(incr_xs)
    }
}

impl Polynomial {
    /// Construct a new polynomial by giving its coefficients.
    pub fn new(coefs: &ColVec) -> Self {
        Polynomial {
            coefs: coefs.to_owned(),
            deg: coefs.len()-1,
        }
    }

    /// Compute the derivative of the polynomial.
    pub fn deriv(&self) -> Self {
        let deg = self.deg;
        let mut deriv_coefs: ColVec = Array1::default(deg);
        for i in 0..deg {
            deriv_coefs[i] = (deg - i) as f64 * self.coefs[i];
        }

        Polynomial::new(&deriv_coefs)
    }

    /// Perform the polynomial interpolation.
    pub fn from_interp(xs: &ColVec, fs: &ColVec) -> Result<Self, LinalgError> {
        let num_data = xs.len();
        if num_data != fs.len() {
            panic!("The lengths of `xs` and `fs` should be equal!");
        }
        if num_data < 2 {
            panic!("At least two points are required!");
        }
        
        // construct Vandermonde-like matrix.
        let mut mat: Array2<f64> = Array::default((num_data, num_data));
        for row in 0..num_data {
            for col in 0..num_data {
                mat[[row, col]] = xs[row].powi((num_data-1-col) as i32);
            }
        }
        // solve the matrix
        match solve_gauss(&mat, &fs) {
            Ok(b) => return Ok(Polynomial::new(&b)),
            Err(err) => return Err(err),
        }
    }
}


/// Use [`CustomPiecewiseFunctionConstExt`] to construct a piecewise function with *constant extrapolation*, implementing [`RealValuedFunction`].
/// * [`CustomPiecewiseFunctionConstExt::ats_incr()`] is faster than [`CustomPiecewiseFunctionConstExt::ats()`] if the input argument is increasing.
/// 
/// # Examples
/// ```
/// use std::rc::Rc;
/// use ndarray::prelude::*;
/// use te1d::prelude::*;
/// use te1d::func::{CustomPiecewiseFunctionConstExt, ExplicitFunction};
/// 
/// let incr_xs: ColVec = array![0.0, 1.0, 2.0];
/// let fs: ColVec = array![0.0, 1.0, 2.0];
/// let func1 = ExplicitFunction::new(Rc::new(|x: f64| { x }));
/// let func2 = ExplicitFunction::new(Rc::new(|x: f64| { x.powi(2) }));
/// let f = CustomPiecewiseFunctionConstExt::new(&incr_xs, &fs, &vec![Rc::new(func1), Rc::new(func2)], 0.0, 3.0);
/// let result = f.ats_incr(&array![-1.0, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5]);
/// let sol: ColVec = array![0.0, 0.0, 0.5, 1.0, 2.25, 2.0, 3.0];
/// println!("{}, {}", result, sol);
/// assert!(all_close(&result, &sol, 1e-05, 1e-08));
/// ```
#[derive(Clone)]
pub struct CustomPiecewiseFunctionConstExt {
    /// `x`-coordinates of the break points, must be strictly increasing.
    pub xs: ColVec,
    /// `y`-coordinates of the break points, same length as `xs`.
    pub fs: ColVec,
    /// The functions in each piece.
    funcs: Vec<Rc<dyn RealValuedFunction>>,
    /// Value to return for `x < xs[0]`.
    pub left: f64,
    /// Value to return for `x > xs[xs.len()-1]`.
    pub right: f64,
    /// The number of raw data points.
    pub num_data: usize,
}

impl PiecewiseFunction for CustomPiecewiseFunctionConstExt {
    #[inline]
    fn at_segment(&self, segment_idx: usize, x: f64) -> f64 {
        self.funcs[segment_idx].at(x)
    }

    #[inline]
    fn at_left(&self) -> f64 {
        self.left
    }

    #[inline]
    fn at_right(&self) -> f64 {
        self.right
    }
}

impl RealValuedFunction for CustomPiecewiseFunctionConstExt {
    /// `incr_xs` must be increasing.
    #[inline]
    fn ats_incr(&self, incr_xs: &ColVec) -> ColVec {
        self.piecewise_ats_incr(incr_xs, &self.xs, &self.fs)
    }
}

impl CustomPiecewiseFunctionConstExt {
    /// Construct a new piecewise function.
    /// 
    /// # Arguments
    /// - `strict_incr_xs`: `x`-coordinates of the break points, must be strictly increasing
    /// - `fs`: `y`-coordinates of the break points, same length as `xs`
    /// - `left`: Value to return for `x < strict_incr_xs[0]`
    /// - `right`: Value to return for `x > strict_incr_xs[xs.len()-1]`
    /// 
    /// # Panics
    /// * when the lengths of `incr_xs` and `fs` are different.
    /// * when the length of `incr_xs` is less than two.
    /// * when the lengths of `incr_xs` is not the length of `funcs` plus one.
    pub fn new(strict_incr_xs: &ColVec, fs: &ColVec, funcs: &Vec<Rc<dyn RealValuedFunction>>, left: f64, right: f64) -> Self {
        let num_data: usize = strict_incr_xs.len();
        if num_data != fs.len() {
            panic!("The lengths of `strict_incr_xs` and `fs` should be equal!");
        }
        if num_data < 2 {
            panic!("At least two points are required!");
        }
        if num_data != funcs.len()+1 {
            panic!("The lengths of `strict_incr_xs` must be the length of `funcs` plus one!");
        }

        CustomPiecewiseFunctionConstExt {
            xs: strict_incr_xs.to_owned(), fs: fs.to_owned(), funcs: funcs.clone(), left, right, num_data
        }
    }
}


/// Use [`ExplicitFunction`] to implement [`RealValuedFunction`] when an explicit formula of a real-valued function is known.
/// * There is no performance difference between [`ExplicitFunction::ats_incr()`] and [`ExplicitFunction::ats()`].
/// 
/// # Examples
/// ```
/// use std::rc::Rc;
/// use ndarray::prelude::*;
/// use te1d::prelude::*;
/// use te1d::func::ExplicitFunction;
/// 
/// let f = ExplicitFunction::new(Rc::new(|x: f64| {x.powi(2)}));
/// 
/// let result = f.ats(&array![-1.0, -0.5, 0.0, 0.3, 0.9]);
/// let sol = array![1.0, 0.25, 0.0, 0.09, 0.81];
/// assert_eq!(result, sol);
/// assert_eq!(f.at(0.75), 0.5625);
/// assert_eq!(&result, &f.ats_incr(&array![-1.0, -0.5, 0.0, 0.3, 0.9]));
/// ```
#[derive(Clone)]
pub struct ExplicitFunction {
    f: Rc<dyn Fn(f64) -> f64>,
}

impl RealValuedFunction for ExplicitFunction {
    #[inline]
    fn at(&self, x: f64) -> f64 {
        (self.f)(x)
    }

    #[inline]
    fn ats_incr(&self, incr_xs: &ColVec) -> ColVec {
        self.ats(incr_xs)
    }
}

impl ExplicitFunction {
    /// # Arguments
    /// * `f`: an explicit function
    pub fn new(f: Rc<dyn Fn(f64) -> f64>) -> Self {
        ExplicitFunction { f: f.clone() }
    }
}


/// Use [`ConstantFunction`] to implement [`RealValuedFunction`] when a constant function is known.
/// * There is no performance difference between [`ConstantFunction::ats_incr()`] and [`ConstantFunction::ats()`].
/// 
/// # Examples
/// ```
/// use std::rc::Rc;
/// use ndarray::prelude::*;
/// use te1d::prelude::*;
/// use te1d::func::ConstantFunction;
/// 
/// let f = ConstantFunction::new(2.0);
/// 
/// let result = f.ats(&array![-1.0, -0.5, 0.0, 0.3, 0.9]);
/// let sol = array![2.0, 2.0, 2.0, 2.0, 2.0];
/// assert_eq!(result, sol);
/// assert_eq!(f.at(0.75), 2.0);
/// assert_eq!(&result, &f.ats_incr(&array![-1.0, -0.5, 0.0, 0.3, 0.9]));
/// ```
#[derive(Debug, PartialEq, Clone)]
pub struct ConstantFunction {
    const_value: f64,
}

impl RealValuedFunction for ConstantFunction {
    #[inline]
    fn at(&self, _x: f64) -> f64 {
        self.const_value
    }

    #[inline]
    fn ats_incr(&self, incr_xs: &ColVec) -> ColVec {
        self.ats(incr_xs)
    }
}

impl ConstantFunction {
    /// # Arguments
    /// * `f`: an explicit function
    pub fn new(const_value: f64) -> Self {
        ConstantFunction { const_value }
    }
}


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

    #[test]
    #[should_panic]
    fn test_panic_piecewise_function_const_ext_new() {
        CustomPiecewiseFunctionConstExt::new(&array![0.0, 0.5, 1.0], &array![0.0, 1.0], &vec![], 0.0, 1.0);  // must panic; different length
        CustomPiecewiseFunctionConstExt::new(&array![0.0], &array![1.0], &vec![], 0.0, 1.0);  // must panic; too few points
        let f = ExplicitFunction::new(Rc::new(|x: f64| {x.powi(2)}));
        CustomPiecewiseFunctionConstExt::new(&array![0.0, 1.0, 2.0], &array![0.0, 1.0, 2.0], &vec![Rc::new(f)], 0.0, 1.0);  // must panic; number of break points and the number of functions do not match
    }
}