fitting 0.5.1

Pure Rust curve fitting library
Documentation
use crate::gaussian::GaussianError;
use crate::linalg;
use crate::linalg::Float;
use ndarray::{array, Array1};

pub fn value<F: Float>(x: F, mu: F, sigma: F, a: F) -> F {
    a * (-(x - mu).powi(2) / (F::from(2).unwrap() * sigma.powi(2))).exp()
}

pub fn values<F: Float>(x_vec: Array1<F>, mu: F, sigma: F, a: F) -> Array1<F> {
    x_vec.iter().map(|x| value(*x, mu, sigma, a)).collect()
}

#[allow(dead_code)]
pub fn fitting_caruanas<F: Float>(
    x_vec: Array1<F>,
    y_vec: Array1<F>,
) -> Result<(F, F, F), GaussianError> {
    let len_x_vec = F::from(x_vec.len()).ok_or(GaussianError::GivenXVecHasNoElement)?;
    let sum_x = x_vec.sum();
    let sum_x_pow2 = x_vec.iter().map(|x| x.powi(2)).sum();
    let sum_x_pow3 = x_vec.iter().map(|x| x.powi(3)).sum();
    let sum_x_pow4 = x_vec.iter().map(|x| x.powi(4)).sum();
    let a = array![
        [len_x_vec, sum_x, sum_x_pow2],
        [sum_x, sum_x_pow2, sum_x_pow3],
        [sum_x_pow2, sum_x_pow3, sum_x_pow4],
    ];

    let sum_log_y = y_vec.iter().map(|y| y.ln()).sum();
    let sum_x_log_y = y_vec
        .iter()
        .zip(x_vec.iter())
        .map(|(y, x)| y.ln() * *x)
        .sum();
    let sum_x_pow2_log_y = y_vec
        .iter()
        .zip(x_vec.iter())
        .map(|(y, x)| y.ln() * x.powi(2))
        .sum();
    let b = array![sum_log_y, sum_x_log_y, sum_x_pow2_log_y];

    let ans_x = linalg::solve(a, b)?;
    let (a, b, c) = (ans_x[0], ans_x[1], ans_x[2]);

    let mu = -b / (F::from(2).unwrap() * c);
    let sigma = (-F::one() / (F::from(2).unwrap() * c)).sqrt();
    let a = (a - (b.powi(2) / (F::from(4).unwrap() * c))).exp();

    Ok((mu, sigma, a))
}

pub fn fitting_guos<F: Float>(
    x_vec: Array1<F>,
    y_vec: Array1<F>,
) -> Result<(F, F, F), GaussianError> {
    // Guo's algorithm doesn't support negative value in y[]
    if y_vec.iter().any(|y| y.is_sign_negative()) {
        return Err(GaussianError::GivenYVecContainsNegativeValue);
    }
    let sum_y_pow2: F = y_vec.iter().map(|y| y.powi(2)).sum();
    let sum_x_y_pow2 = y_vec
        .iter()
        .zip(x_vec.iter())
        .map(|(y, x)| *x * y.powi(2))
        .sum();
    let sum_x_pow2_y_pow2 = y_vec
        .iter()
        .zip(x_vec.iter())
        .map(|(y, x)| x.powi(2) * y.powi(2))
        .sum();
    let sum_x_pow3_y_pow2 = y_vec
        .iter()
        .zip(x_vec.iter())
        .map(|(y, x)| x.powi(3) * y.powi(2))
        .sum();
    let sum_x_pow4_y_pow2 = y_vec
        .iter()
        .zip(x_vec.iter())
        .map(|(y, x)| x.powi(4) * y.powi(2))
        .sum();

    let a = array![
        [sum_y_pow2, sum_x_y_pow2, sum_x_pow2_y_pow2],
        [sum_x_y_pow2, sum_x_pow2_y_pow2, sum_x_pow3_y_pow2],
        [sum_x_pow2_y_pow2, sum_x_pow3_y_pow2, sum_x_pow4_y_pow2],
    ];

    let sum_y_pow2_log_y: F = y_vec.iter().map(|y| y.powi(2) * y.ln()).sum();
    let sum_x_y_pow2_log_y = y_vec
        .iter()
        .zip(x_vec.iter())
        .map(|(y, x)| *x * y.powi(2) * y.ln())
        .sum();
    let sum_x_pow2_y_pow2_log_y = y_vec
        .iter()
        .zip(x_vec.iter())
        .map(|(y, x)| x.powi(2) * y.powi(2) * y.ln())
        .sum();
    let b = array![
        sum_y_pow2_log_y,
        sum_x_y_pow2_log_y,
        sum_x_pow2_y_pow2_log_y,
    ];

    let ans_x = linalg::solve(a, b)?;
    let (a, b, c) = (ans_x[0], ans_x[1], ans_x[2]);

    let mu = -b / (F::from(2).unwrap() * c);
    let sigma = ((-F::one() / (F::from(2).unwrap() * c)) as F).sqrt();
    let a = ((a - (b.powi(2) / (F::from(4).unwrap() * c))) as F).exp();

    Ok((mu, sigma, a))
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_abs_diff_eq;
    use ndarray::{array, Array};

    #[test]
    fn gaussian_value() {
        let (mu, sigma, a): (f64, f64, f64) = (5., 3., 1.);
        let x = 5.;
        let y = value(x, mu, sigma, a);
        assert_eq!(y, a);
    }

    #[test]
    fn gaussian_values() {
        let (mu, sigma, a): (f64, f64, f64) = (5., 3., 1.);
        let x_vec: Array1<f64> = Array::range(1., 10., 1.);
        let y_vec: Array1<f64> = values(x_vec, mu, sigma, a);
        let expected_ans = array![
            0.41111229050718745,
            0.6065306597126334,
            0.8007374029168081,
            0.9459594689067654,
            1.,
            0.9459594689067654,
            0.8007374029168081,
            0.6065306597126334,
            0.41111229050718745
        ];
        assert_abs_diff_eq!(&y_vec, &expected_ans, epsilon = 1e-9);
    }

    #[test]
    fn gaussian_fit_caruanas() {
        let (mu, sigma, a): (f64, f64, f64) = (5., 3., 1.);
        let x_vec: Array1<f64> = Array::range(1., 10., 1.);
        let y_vec: Array1<f64> = values(x_vec.clone(), mu, sigma, a);
        let estimated = fitting_caruanas(x_vec, y_vec).unwrap();
        assert_abs_diff_eq!(
            &array![estimated.0, estimated.1, estimated.2],
            &array![mu, sigma, a],
            epsilon = 1e-9
        );
    }

    #[test]
    fn gaussian_fit_guos() {
        let (mu, sigma, a): (f64, f64, f64) = (5., 3., 1.);
        let x_vec: Array1<f64> = Array::range(1., 10., 1.);
        let y_vec: Array1<f64> = values(x_vec.clone(), mu, sigma, a);
        let estimated = fitting_guos(x_vec, y_vec).unwrap();
        assert_abs_diff_eq!(
            &array![estimated.0, estimated.1, estimated.2],
            &array![mu, sigma, a],
            epsilon = 1e-9
        );
    }

    #[test]
    fn gaussian_fit_guos_y_vec_contains_negative_value() {
        let (mu, sigma, a): (f64, f64, f64) = (5., 3., 1.);
        let x_vec: Array1<f64> = Array::range(1., 10., 1.);
        let y_vec: Array1<f64> = values(x_vec.clone(), mu, sigma, a).map(|y| y - 0.5);

        let err = fitting_guos(x_vec, y_vec).unwrap_err();
        assert_eq!(err, GaussianError::GivenYVecContainsNegativeValue);
    }
}