r2rs_stats/regression/glm/

model.rs

// "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::real::Real;
use strafe_trait::{
    Manipulator, Model, Statistic, StatisticalEstimate, StatisticalEstimator, StatisticalTest,
};
use strafe_type::ModelMatrix;

use crate::{
    funcs::residual_sum_of_squares,
    regression::glm::{family::core::Family, link::core::Link, GeneralizedLinearRegression},
    tests::{LenientAdjustedRSquaredTest, LinearEstimator, ZCoefficientBuilder},
};

impl<L: Link, F: Family<L>> Model for GeneralizedLinearRegression<L, F> {
    fn get_x(&self) -> ModelMatrix {
        self.x.clone()
    }

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

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

    fn get_weights(&self) -> ModelMatrix {
        ModelMatrix::from(&self.final_weights)
    }

    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(ZCoefficientBuilder::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>> {
        let mut linear_prediction = LinearEstimator::new()
            .with_alpha(self.alpha)
            .estimate(&self.model_data)?;

        let mut linear_prediction_ret = linear_prediction.clone();
        let linear_prediction_estimate: ModelMatrix = linear_prediction.estimate();

        linear_prediction_ret.predictions_estimate = ModelMatrix::from(
            self.family
                .link()
                .link_inverse(&linear_prediction_estimate.matrix()),
        );
        let lin_pred = linear_prediction_ret
            .predictions_confidence_interval
            .matrix();
        let lin_pred_lower = self
            .family
            .link()
            .link_inverse(&ModelMatrix::from(lin_pred.column(0).as_slice()).matrix());
        let lin_pred_upper = self
            .family
            .link()
            .link_inverse(&ModelMatrix::from(lin_pred.column(1).as_slice()).matrix());
        linear_prediction_ret.predictions_confidence_interval =
            ModelMatrix::from(DMatrix::<f64>::from_iterator(
                lin_pred.nrows(),
                2,
                lin_pred_lower
                    .into_iter()
                    .chain(lin_pred_upper.into_iter())
                    .cloned(),
            ));

        Ok(Box::new(linear_prediction_ret))
    }

    fn residuals(&self) -> Result<ModelMatrix, Box<dyn Error>> {
        let pred = self.predictions()?.estimate();

        Ok(ModelMatrix::from(
            (self.y.matrix() - pred.matrix()).component_div(&self.family.link().mu_eta(&self.eta)),
        ))
    }

    fn studentized_residuals(&self) -> Result<ModelMatrix, Box<dyn Error>> {
        let x = self.x1.matrix();
        let y = self.y.matrix();
        let b = self.b.clone();
        let r = self.residuals_deviance()?.matrix();
        let pear_res = self.residuals_pearson()?.matrix();
        let sum = residual_sum_of_squares(&x, &y, &b);
        let denom = x.nrows() as f64 - x.ncols() as f64;
        let hat = self.leverage();
        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<_>>();

        Ok(ModelMatrix::from(
            r.into_iter()
                .zip(pear_res.into_iter())
                .zip(hat.into_iter())
                .zip(sigma.into_iter())
                .map(|(((r, pear_res), hat), sigma)| {
                    let mut ret =
                        r.signum() * (r.powi(2) + (hat * pear_res.powi(2)) / (1.0 - hat)).sqrt();
                    if self.family.set_dispersion().is_none() {
                        ret /= sigma;
                    }
                    ret
                })
                .collect::<Vec<_>>(),
        ))
    }

    fn standardized_residuals(&self) -> Result<ModelMatrix, Box<dyn Error>> {
        let res = self.residuals_deviance()?.matrix();
        let disp = self.dispersion();
        let hat = self.leverage();
        Ok(ModelMatrix::from(
            res.into_iter()
                .zip(hat.into_iter())
                .map(|(res, hat)| res / (disp * (1.0 - hat)).sqrt())
                .collect::<Vec<_>>(),
        ))
    }

    fn variance(&self) -> Result<DMatrix<f64>, Box<dyn Error>> {
        let mut w = self.final_weights.clone();
        w.iter_mut().for_each(|w_i| *w_i = w_i.sqrt());

        let mut weighted_x = self.x1.matrix();
        weighted_x
            .column_iter_mut()
            .for_each(|mut row| row.component_mul_assign(&w));

        let cov_unscaled = (weighted_x.clone().transpose() * weighted_x.clone())
            .pseudo_inverse(f64::epsilon())
            .unwrap();

        Ok(self.dispersion() * cov_unscaled)
    }
}

impl<L: Link, F: Family<L>> GeneralizedLinearRegression<L, F> {
    pub fn dispersion(&self) -> f64 {
        if let Some(disp) = self.family.set_dispersion() {
            disp
        } else {
            let df_residual = self.x1.matrix().nrows() - self.b.len();
            let resid = self.residuals().unwrap().matrix();
            self.final_weights
                .iter()
                .zip(resid.into_iter())
                .map(|(w_i, r_i)| w_i * r_i.powi(2))
                .sum::<f64>()
                / df_residual as f64
        }
    }

    pub fn residuals_pearson(&self) -> Result<ModelMatrix, Box<dyn Error>> {
        let r = self
            .y
            .matrix()
            .iter()
            .zip(self.mu.iter())
            .map(|(y, mu)| y - mu)
            .collect::<Vec<_>>();
        let wts = self.w.matrix();
        let var = self.family.variance(&self.mu);
        Ok(ModelMatrix::from(
            r.into_iter()
                .zip(wts.into_iter())
                .zip(var.into_iter())
                .map(|((r, w), v)| r * w.sqrt() / v.sqrt())
                .collect::<Vec<_>>(),
        ))
    }

    pub fn residuals_deviance(&self) -> Result<ModelMatrix, Box<dyn Error>> {
        let y = self.y.matrix();
        let wts = self.w.matrix();
        let mu = self.mu.clone();
        let dev_res = self.family.residual_deviance(&y, &mu, &wts);
        let ret = dev_res
            .into_iter()
            .zip(y.into_iter())
            .zip(mu.into_iter())
            .map(|((dev_res, y), mu)| {
                let d_res = dev_res.max(0.0).sqrt();
                if y > mu {
                    d_res
                } else {
                    -d_res
                }
            })
            .collect::<Vec<_>>();

        Ok(ModelMatrix::from(ret))
    }
}