Implementation of online first and second order ordinary least squares regression

main.rs

/// A simple online linear regression algorithm.
/// See https://en.wikipedia.org/wiki/Simple_linear_regression#Fitting_the_regression_line
/// for details.
pub struct FirstOrderOLS {
    x: f64,
    y: f64,
    xx: f64,
    xy: f64,
    n: usize,
}

impl FirstOrderOLS {
    /// Create a new instance of the algorithm.
    pub fn new() -> Self {
        Self {
            x: 0.0,
            y: 0.0,
            xx: 0.0,
            xy: 0.0,
            n: 0,
        }
    }

    /// Add a new point to the algorithm in O(1).
    pub fn add_point(&mut self, x: f64, y: f64) {
        self.x += x;
        self.y += y;
        self.xx += x * x;
        self.xy += x * y;
        self.n += 1;
    }

    /// Remove a point from the algorithm in O(1).
    pub fn remove_point(&mut self, x: f64, y: f64) {
        self.x -= x;
        self.y -= y;
        self.xx -= x * x;
        self.xy -= x * y;
        self.n -= 1;
    }

    /// Reset the algorithm to its initial state.
    pub fn reset(&mut self) {
        self.x = 0.0;
        self.y = 0.0;
        self.xx = 0.0;
        self.xy = 0.0;
        self.n = 0;
    }

    /// Return the coefficients of the regression line.
    /// The first value is the slope, the second value is the intercept.
    pub fn coeffs(&self) -> (f64, f64) {
        let slope = (self.n as f64 * self.xy - self.x * self.y)
            / (self.n as f64 * self.xx - self.x * self.x);
        let intercept = (self.y - slope * self.x) / self.n as f64;
        (slope, intercept)
    }
}

////////////////////////////////////////////////////////////////////////////////

/// A simple online quadratic regression algorithm.
/// See https://en.wikipedia.org/wiki/Polynomial_regression#Matrix_form_and_calculation_of_estimates
/// for details.
pub struct SecondOrderOLS {
    x: f64,
    y: f64,
    xx: f64,
    xy: f64,
    xxx: f64,
    xxy: f64,
    xxxx: f64,
    n: usize,
}

impl SecondOrderOLS {
    /// Create a new instance of the algorithm.
    pub fn new() -> Self {
        Self {
            x: 0.0,
            y: 0.0,
            xx: 0.0,
            xy: 0.0,
            xxx: 0.0,
            xxy: 0.0,
            xxxx: 0.0,
            n: 0,
        }
    }

    /// Add a new point to the algorithm in O(1).
    pub fn add_point(&mut self, x: f64, y: f64) {
        self.x += x;
        self.y += y;
        self.xx += x * x;
        self.xy += x * y;
        self.xxx += x * x * x;
        self.xxy += x * x * y;
        self.xxxx += x * x * x * x;
        self.n += 1;
    }

    /// Remove a point from the algorithm in O(1).
    pub fn remove_point(&mut self, x: f64, y: f64) {
        self.x -= x;
        self.y -= y;
        self.xx -= x * x;
        self.xy -= x * y;
        self.xxx -= x * x * x;
        self.xxy -= x * x * y;
        self.xxxx -= x * x * x * x;
        self.n -= 1;
    }

    /// Reset the algorithm to its initial state.
    pub fn reset(&mut self) {
        self.x = 0.0;
        self.y = 0.0;
        self.xx = 0.0;
        self.xy = 0.0;
        self.xxx = 0.0;
        self.xxy = 0.0;
        self.xxxx = 0.0;
        self.n = 0;
    }

    /// Return the coefficients of the regression parabola.
    /// The first value is the quadratic coefficient, the second value is the linear coefficient,
    /// and the third value is the constant coefficient.
    pub fn coeffs(&self) -> (f64, f64, f64) {
        let n = self.n as f64;
        let a = (-n * self.xx * self.xxy + n * self.xxx * self.xy + self.x * self.x * self.xxy
            - self.x * self.xx * self.xy
            - self.x * self.xxx * self.y
            + self.xx * self.xx * self.y)
            / (-n * self.xx * self.xxxx + n * self.xxx * self.xxx + self.x * self.x * self.xxxx
                - 2.0 * self.x * self.xx * self.xxx
                + self.xx * self.xx * self.xx);

        let b = (n * self.xxx * self.xxy - n * self.xxxx * self.xy - self.x * self.xx * self.xxy
            + self.x * self.xxxx * self.y
            + self.xx * self.xx * self.xy
            - self.xx * self.xxx * self.y)
            / (-n * self.xx * self.xxxx + n * self.xxx * self.xxx + self.x * self.x * self.xxxx
                - 2.0 * self.x * self.xx * self.xxx
                + self.xx * self.xx * self.xx);

        let c = (-self.x * self.xxx * self.xxy
            + self.x * self.xxxx * self.xy
            + self.xx * self.xx * self.xxy
            - self.xx * self.xxx * self.xy
            - self.xx * self.xxxx * self.y
            + self.xxx * self.xxx * self.y)
            / (-n * self.xx * self.xxxx + n * self.xxx * self.xxx + self.x * self.x * self.xxxx
                - 2.0 * self.x * self.xx * self.xxx
                + self.xx * self.xx * self.xx);

        (a, b, c)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    #[test]
    fn test_first_order() {
        let mut model = FirstOrderOLS::new();
        for i in 0..1000 {
            let i = i as f64;
            model.add_point(i, 2.0 * i + 1.0);
        }
        let (slope, intercept) = model.coeffs();
        assert!(slope - 2.0 < 1e-6);
        assert!(intercept - 1.0 < 1e-6);
    }
    #[test]
    fn test_second_order() {
        let mut model = SecondOrderOLS::new();
        for i in 0..1000 {
            let i = i as f64;
            model.add_point(i, 3.0 * i * i + 2.0 * i + 1.0);
        }
        let (a, b, c) = model.coeffs();
        assert!(a - 3.0 < 1e-6);
        assert!(b - 2.0 < 1e-6);
        assert!(c - 1.0 < 1e-6);
    }
}