r2rs-stats 0.1.1

// "Whatever you do, work at it with all your heart, as working for the Lord,
// not for human masters, since you know that you will receive an inheritance
// from the Lord as a reward. It is the Lord Christ you are serving."
// (Col 3:23-24)

use std::error::Error;

use nalgebra::DMatrix;
use num_traits::Float;
use strafe_trait::{
    Manipulator, Model, Statistic, StatisticalEstimate, StatisticalEstimator, StatisticalTest,
};
use strafe_type::ModelMatrix;

use crate::{
    funcs::{residual_sum_of_squares, residuals},
    regression::lm::LeastSquaresRegression,
    tests::{LenientAdjustedRSquaredTest, LinearEstimator, TCoefficientBuilder},
};

impl Model for LeastSquaresRegression {
    fn get_x(&self) -> ModelMatrix {
        self.x.clone()
    }

    fn get_x1(&self) -> ModelMatrix {
        ModelMatrix::from(&self.x1)
    }

    fn get_y(&self) -> ModelMatrix {
        self.y.clone()
    }

    fn get_weights(&self) -> ModelMatrix {
        self.w.clone()
    }

    fn get_intercept(&self) -> bool {
        true
    }

    fn manipulator(
        &self,
    ) -> Result<
        Box<dyn StatisticalEstimate<Manipulator<ModelMatrix, ModelMatrix>, ModelMatrix>>,
        Box<dyn Error>,
    > {
        Ok(Box::new(
            LinearEstimator::new()
                .with_alpha(self.alpha)
                .estimate(&self.model_data)?,
        ))
    }

    fn determination(&self) -> Result<Box<dyn Statistic>, Box<dyn Error>> {
        Ok(Box::new(
            LenientAdjustedRSquaredTest::new()
                .with_alpha(self.alpha)
                .test(&self.model_data)?,
        ))
    }
    fn parameters(
        &self,
    ) -> Result<Vec<Box<dyn StatisticalEstimate<f64, (f64, f64)>>>, Box<dyn Error>> {
        Ok(TCoefficientBuilder::new()
            .with_alpha(self.alpha)
            .estimate(&self.model_data)?
            .into_iter()
            .map(|c| {
                let x: Box<dyn StatisticalEstimate<f64, (f64, f64)>> = Box::new(c);
                x
            })
            .collect())
    }

    fn predictions(
        &self,
    ) -> Result<Box<dyn StatisticalEstimate<ModelMatrix, ModelMatrix>>, Box<dyn Error>> {
        Ok(Box::new(
            LinearEstimator::new()
                .with_alpha(self.alpha)
                .estimate(&self.model_data)?,
        ))
    }

    fn residuals(&self) -> Result<ModelMatrix, Box<dyn Error>> {
        let mut weighted_x = self.x1.clone();
        weighted_x
            .column_iter_mut()
            .for_each(|mut row| row.component_mul_assign(&self.w.matrix()));
        let weighted_y = self.y.matrix().component_mul(&self.w.matrix());

        Ok(ModelMatrix::from(residuals(
            &weighted_y,
            &weighted_x,
            &self.b,
        )))
    }

    fn standardized_residuals(&self) -> Result<ModelMatrix, Box<dyn Error>> {
        let x = self.x1.clone();
        let y = self.y.matrix();
        let b = self.b.clone();
        let hat = self.leverage();

        let mut ret = self.residuals()?.matrix();
        let sd =
            (residual_sum_of_squares(&x, &y, &b) / (x.nrows() as f64 - x.ncols() as f64)).sqrt();

        for (r_i, h_ii) in ret.iter_mut().zip(hat.into_iter()) {
            *r_i = *r_i / (sd * (1.0 - h_ii).sqrt());
        }

        Ok(ModelMatrix::from(ret))
    }

    fn studentized_residuals(&self) -> Result<ModelMatrix, Box<dyn Error>> {
        let x = self.x1.clone();
        let y = self.y.matrix();
        let b = self.b.clone();
        let hat = self.leverage();
        let sum = residual_sum_of_squares(&x, &y, &b);
        let denom = x.nrows() as f64 - x.ncols() as f64;
        let sigma = self
            .residuals()?
            .matrix()
            .into_iter()
            .zip(hat.iter())
            .map(|(&res, &hat)| {
                if hat < 1.0 {
                    ((sum - res.powi(2) / (1.0 - hat)) / denom).sqrt()
                } else {
                    (sum / denom).sqrt()
                }
            })
            .collect::<Vec<_>>();

        let mut ret = self.residuals()?.matrix();
        for ((r_i, sigma_i), h_ii) in ret.iter_mut().zip(sigma.into_iter()).zip(hat.into_iter()) {
            *r_i = *r_i / (sigma_i * (1.0 - h_ii).sqrt());
        }

        Ok(ModelMatrix::from(ret))
    }

    fn variance(&self) -> Result<DMatrix<f64>, Box<dyn Error>> {
        let cov_unscaled = (self.x1.transpose() * &self.x1)
            .pseudo_inverse(f64::epsilon())
            .unwrap();
        let rss = self
            .residuals()?
            .matrix()
            .into_iter()
            .map(|r| r.powi(2))
            .sum::<f64>();
        let rdf = self.x1.nrows() as f64 - self.parameters()?.len() as f64;
        let sigma = (rss / rdf).sqrt();
        Ok(sigma.powi(2) * cov_unscaled)
    }
}