Skip to main content

robust_rs/regression/
m_estimator.rs

1//! Regression M-estimation by iteratively reweighted least squares.
2
3use ndarray::{Array1, Array2};
4use robust_rs_core::error::RobustError;
5use robust_rs_core::rho::RhoFunction;
6use robust_rs_core::scale::ScaleEstimator;
7use robust_rs_core::solver::Control;
8
9use crate::estimator::RegressionFit;
10
11/// A configured regression M-estimator: a loss, a scale estimator and solver
12/// controls. Call [`MEstimator::fit`] to fit it to data.
13pub struct MEstimator<R, S> {
14    /// The robust loss.
15    pub rho: R,
16    /// The scale estimator used to standardize residuals.
17    pub scale: S,
18    /// Convergence controls.
19    pub control: Control,
20}
21
22impl<R, S> MEstimator<R, S>
23where
24    R: RhoFunction + Clone + 'static,
25    S: ScaleEstimator,
26{
27    /// Create an M-estimator with default solver controls.
28    pub fn new(rho: R, scale: S) -> Self {
29        Self {
30            rho,
31            scale,
32            control: Control::default(),
33        }
34    }
35
36    /// Fit by IRLS: standardize residuals by the scale, form the weights
37    /// `wᵢ = ρ.weight(rᵢ/s)`, solve the weighted least squares, repeat.
38    pub fn fit(&self, x: &Array2<f64>, y: &Array1<f64>) -> Result<RegressionFit, RobustError> {
39        use crate::wls::weighted_least_squares;
40
41        let n = x.nrows();
42
43        // Initial fit: OLS (unit weights). Dim / rank errors propagate from WLS.
44        let ones = Array1::from_elem(n, 1.0);
45        let mut beta = weighted_least_squares(x, y, &ones)?;
46        let mut resid = y - &x.dot(&beta);
47
48        // Initial scale from the OLS residuals, held fixed through the loop.
49        let scale_est = self.scale.scale(resid.as_slice().expect("contiguous"))?;
50        let s: f64 = scale_est.get();
51
52        let mut converged = false;
53        for _ in 0..self.control.max_iter {
54            let w = resid.mapv(|r| self.rho.weight(r / s));
55            let next = weighted_least_squares(x, y, &w)?;
56
57            // relative L2 change in coefficients (matches Control::tol's semantics)
58            let diff = &next - &beta;
59            let rel = diff.dot(&diff).sqrt() / (next.dot(&next).sqrt() + 1e-12);
60
61            beta = next;
62            resid = y - &x.dot(&beta);
63            if rel <= self.control.tol {
64                converged = true;
65                break;
66            }
67        }
68        if !converged {
69            return Err(RobustError::NonConvergence {
70                iters: self.control.max_iter,
71            });
72        }
73
74        // Weights evaluated at the converged residuals (≈ the last solve's weights).
75        let weights = resid.mapv(|r| self.rho.weight(r / s));
76
77        Ok(RegressionFit {
78            coefficients: beta,
79            scale: scale_est,
80            residuals: resid,
81            weights,
82            rho: Box::new(self.rho.clone()), // R: Clone + 'static → Box<dyn RhoFunction>
83            breakdown_point: 0.0, // plain regression M-estimation: leverage ⇒ 0 breakdown
84        })
85    }
86}