ndarray_polyfit/
lib.rs

1use ndarray::Array1;
2use ndarray::Array2;
3use ndarray_linalg::Solve;
4
5/// @param x_values The x-values
6/// @param y_values The y-values
7/// @param polynomial_degree The degree of the polynomial. I. e. 2 for a parabola.
8/// @return Degree of monomials increases with the vector index
9pub fn polyfit(
10    x_values: Array1<f64>,
11    y_values: Array1<f64>,
12    polynomial_degree: usize,
13) -> Result<Array1<f64>, &'static str> {
14    let number_of_columns = polynomial_degree + 1;
15    let number_of_rows = x_values.len();
16
17    // Construct the Vandermonde matrix
18    let mut a = Array2::zeros((number_of_rows, number_of_columns));
19
20    for row in 0..number_of_rows {
21        // First column is always 1
22        a[[row, 0]] = 1.0;
23
24        for col in 1..number_of_columns {
25            a[[row, col]] = x_values[row].powf(col as f64);
26        }
27    }
28
29    // Solve using least squares method
30    match a.solve(&y_values) {
31        Ok(mat) => Ok(mat),
32        Err(_error) => Err("Linear system could not be solved"),
33    }
34}
35
36#[cfg(test)]
37mod tests {
38    use super::*;
39    use approx::assert_abs_diff_eq;
40    use ndarray::array;
41
42    #[test]
43    fn test_polyfit() {
44        let x_values = array![0.0, 1.0, 2.0];
45        let y_values = array![0.0, 1.0, 4.0];
46
47        let result = polyfit(x_values, y_values, 2).unwrap();
48
49        assert_abs_diff_eq!(result[0], 0.0, epsilon = 1e-6);
50        assert_abs_diff_eq!(result[1], 0.0, epsilon = 1e-6);
51        assert_abs_diff_eq!(result[2], 1.0, epsilon = 1e-6);
52    }
53}