rs_ml/regression/
linear.rs

1//! Linear regression models
2
3use ndarray::{Array1, Array2, Axis};
4use ndarray_linalg::Inverse;
5
6use crate::Estimator;
7
8use super::Regressor;
9
10/// Estimator for a linear ordinary least squares model
11#[derive(Debug, Clone, Copy)]
12pub struct OrdinaryLeastSquaresEstimator;
13
14/// Fitted regression model for ordinary least squares
15#[derive(Debug, Clone)]
16pub struct OrdinaryLeastSquaresRegressor {
17    beta: Array2<f64>,
18}
19
20impl Estimator<(&Array2<f64>, &Array1<f64>)> for OrdinaryLeastSquaresEstimator {
21    type Estimator = OrdinaryLeastSquaresRegressor;
22
23    fn fit(&self, input: &(&Array2<f64>, &Array1<f64>)) -> Option<Self::Estimator> {
24        let (x, y) = input;
25
26        let nrows = x.nrows();
27        let mut x_added_one = x.to_owned().clone();
28        x_added_one.push_column(Array1::ones(nrows).view()).ok()?;
29
30        let binding = y.view().insert_axis(Axis(0));
31        let transformed_y = binding.t();
32        let inv_gram_matrix: Array2<f64> = x_added_one.t().dot(&x_added_one).inv().ok()?;
33
34        let beta = inv_gram_matrix.dot(&x_added_one.t().dot(&transformed_y));
35
36        Some(OrdinaryLeastSquaresRegressor { beta })
37    }
38}
39
40impl Regressor<Array2<f64>, Array1<f64>> for OrdinaryLeastSquaresRegressor {
41    fn predict(&self, input: &Array2<f64>) -> Option<Array1<f64>> {
42        let nrows = input.nrows();
43        let mut x_added_one = input.to_owned().clone();
44        x_added_one.push_column(Array1::ones(nrows).view()).ok()?;
45
46        let y = x_added_one.dot(&self.beta);
47
48        let binding = y.t().remove_axis(Axis(0));
49        Some(binding.to_owned())
50    }
51}