r2rs-stats 0.1.1

Statistics programming for Rust based on R's stats package
Documentation
// "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 nalgebra::DMatrix;
use r2rs_nmath::{distribution::NormalBuilder, traits::Distribution};
use strafe_type::{Alpha64, FloatConstraint, ModelMatrix};

use crate::tests::z_coefficient::ZCoefficient;

#[derive(Copy, Clone, Debug)]
pub struct ZCoefficientBuilder {
    alpha: Alpha64,
}

impl Default for ZCoefficientBuilder {
    fn default() -> Self {
        Self { alpha: 0.05.into() }
    }
}

impl ZCoefficientBuilder {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_alpha(self, alpha: Alpha64) -> Self {
        Self { alpha }
    }

    pub(crate) fn build(
        &self,
        (x, _y, b, w): &(ModelMatrix, DMatrix<f64>, DMatrix<f64>, DMatrix<f64>),
    ) -> Result<Vec<ZCoefficient>, Box<dyn std::error::Error>> {
        let names = x.column_names();
        let x = x.matrix();

        let mut weights = w.clone();
        weights.iter_mut().for_each(|w_i| *w_i = w_i.sqrt());

        let mut weighted_x = x.clone();
        weighted_x
            .column_iter_mut()
            .for_each(|mut row| row.component_mul_assign(&weights));

        let n = x.shape().0;
        let k = x.shape().1 - 1;

        let num_coef = x.shape().1;
        let weighted_x_r = weighted_x.clone().qr().r();
        let covariance = (weighted_x_r.transpose() * &weighted_x_r)
            .pseudo_inverse(f64::EPSILON)
            .expect("Inverse Error");

        let mut ret = Vec::new();
        for j in 0..num_coef {
            let estimate = b[(j, 0)];
            let standard_error = (covariance.diagonal()[j]).sqrt();
            let z = estimate / standard_error;

            let mut normal_distr_builder = NormalBuilder::new();
            let normal_distr = normal_distr_builder.build();

            let p = normal_distr.probability(-z.abs(), true).unwrap() * 2.0;
            let confidence_value = normal_distr
                .quantile(self.alpha.unwrap() / (2.0 * (n - k - 1) as f64), false)
                .unwrap();

            let lower_bound = estimate - (confidence_value * standard_error);
            let upper_bound = estimate + (confidence_value * standard_error);

            ret.push(ZCoefficient {
                name: names[j].clone(),
                alpha: (self.alpha.unwrap() / (k + 1) as f64).into(),
                estimate,
                confidence_interval: (lower_bound, upper_bound),
                z,
                p,
            });
        }

        Ok(ret)
    }
}