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)
}
}