1use ndarray::Array1;
2use ndarray::Array2;
3use ndarray_linalg::Solve;
4
5pub 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 let mut a = Array2::zeros((number_of_rows, number_of_columns));
19
20 for row in 0..number_of_rows {
21 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 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}