dml_core/learner/
ols.rs

1use ndarray::{Array1, ArrayView1, ArrayView2};
2
3use super::{center, Centered, Regressor};
4use crate::error::{DmlError, Result};
5use crate::linalg::cholesky_solve;
6
7/// Ordinary least squares with an intercept and optional ridge penalty.
8///
9/// `ridge` (default 0) is added to the diagonal for slopes only, which keeps
10/// the fit well defined when covariates are collinear.
11#[derive(Debug, Clone, Default)]
12pub struct Ols {
13    ridge: f64,
14    coef: Option<Array1<f64>>,
15    intercept: f64,
16}
17
18impl Ols {
19    pub fn new() -> Self {
20        Self::default()
21    }
22
23    pub fn with_ridge(ridge: f64) -> Self {
24        Self {
25            ridge: ridge.max(0.0),
26            ..Self::default()
27        }
28    }
29
30    pub fn coef(&self) -> Option<&Array1<f64>> {
31        self.coef.as_ref()
32    }
33
34    pub fn intercept(&self) -> f64 {
35        self.intercept
36    }
37}
38
39impl Ols {
40    fn fit_impl(
41        &mut self,
42        x: ArrayView2<'_, f64>,
43        y: ArrayView1<'_, f64>,
44        w: Option<ArrayView1<'_, f64>>,
45    ) -> Result<()> {
46        let p = x.ncols();
47        // Centre to absorb the intercept, then solve (Xc'Xc + λI) b = Xc'yc.
48        let Centered {
49            x_mean,
50            y_mean,
51            xc,
52            yc,
53        } = center(x, y, w, false)?;
54        let mut xtx = xc.t().dot(&xc);
55        for j in 0..p {
56            xtx[[j, j]] += self.ridge;
57        }
58        let xty = xc.t().dot(&yc);
59        let coef = if p == 0 {
60            Array1::zeros(0)
61        } else {
62            cholesky_solve(&xtx, &xty)?
63        };
64        self.intercept = y_mean - x_mean.dot(&coef);
65        self.coef = Some(coef);
66        Ok(())
67    }
68}
69
70impl Regressor for Ols {
71    fn fit(&mut self, x: ArrayView2<'_, f64>, y: ArrayView1<'_, f64>) -> Result<()> {
72        self.fit_impl(x, y, None)
73    }
74
75    /// Weighted least squares: minimises `Σ wᵢ (yᵢ − b − xᵢᵀβ)² + ridge·‖β‖²`.
76    fn fit_weighted(
77        &mut self,
78        x: ArrayView2<'_, f64>,
79        y: ArrayView1<'_, f64>,
80        w: ArrayView1<'_, f64>,
81    ) -> Result<()> {
82        self.fit_impl(x, y, Some(w))
83    }
84
85    fn predict(&self, x: ArrayView2<'_, f64>) -> Result<Array1<f64>> {
86        let coef = self.coef.as_ref().ok_or(DmlError::NotFitted)?;
87        if x.ncols() != coef.len() {
88            return Err(DmlError::DimensionMismatch(
89                "Ols::predict feature count".into(),
90            ));
91        }
92        Ok(x.dot(coef) + self.intercept)
93    }
94
95    fn fresh(&self) -> Box<dyn Regressor> {
96        Box::new(Ols::with_ridge(self.ridge))
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use ndarray::array;
104
105    #[test]
106    fn recovers_exact_linear_fit() {
107        let x = array![[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [2.0, 1.0], [0.5, 3.0]];
108        let y: Array1<f64> = x
109            .rows()
110            .into_iter()
111            .map(|r| 1.5 + 2.0 * r[0] - 0.5 * r[1])
112            .collect();
113        let mut m = Ols::new();
114        m.fit(x.view(), y.view()).unwrap();
115        let c = m.coef().unwrap();
116        assert!((c[0] - 2.0).abs() < 1e-10 && (c[1] + 0.5).abs() < 1e-10);
117        assert!((m.intercept() - 1.5).abs() < 1e-10);
118    }
119
120    #[test]
121    fn predict_before_fit_errors() {
122        assert_eq!(
123            Ols::new().predict(array![[1.0]].view()).unwrap_err(),
124            DmlError::NotFitted
125        );
126    }
127
128    #[test]
129    fn integer_weights_equal_row_duplication() {
130        let x = array![[0.0], [1.0], [2.0], [3.0]];
131        let y = array![1.0, 0.5, 2.5, 2.0];
132        let mut weighted = Ols::new();
133        weighted
134            .fit_weighted(x.view(), y.view(), array![1.0, 2.0, 1.0, 3.0].view())
135            .unwrap();
136        let xd = array![[0.0], [1.0], [1.0], [2.0], [3.0], [3.0], [3.0]];
137        let yd = array![1.0, 0.5, 0.5, 2.5, 2.0, 2.0, 2.0];
138        let mut dup = Ols::new();
139        dup.fit(xd.view(), yd.view()).unwrap();
140        assert!((weighted.coef().unwrap()[0] - dup.coef().unwrap()[0]).abs() < 1e-12);
141        assert!((weighted.intercept() - dup.intercept()).abs() < 1e-12);
142    }
143
144    #[test]
145    fn ridge_handles_collinearity() {
146        let x = array![[1.0, 2.0], [2.0, 4.0], [3.0, 6.0]];
147        let y = array![1.0, 2.0, 3.0];
148        assert!(Ols::new().fit(x.view(), y.view()).is_err());
149        assert!(Ols::with_ridge(1e-6).fit(x.view(), y.view()).is_ok());
150    }
151}
dml_core/learner/ols.rs

dml_core/learner/
ols.rs
