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 std::{
    fmt::{Debug, Formatter},
    sync::Arc,
};

use nalgebra::DMatrix;
use strafe_trait::{Assumption, Manipulator, StatisticalEstimate, StatisticalEstimator};
use strafe_type::{Alpha64, ModelMatrix, Rational64};

use crate::funcs::{confidence_interval, predict, prediction_interval};

#[derive(Copy, Clone, Debug)]
pub struct LinearEstimator {
    alpha: Alpha64,
    scale: Option<f64>,
    df: Option<Rational64>,
}

impl StatisticalEstimator for LinearEstimator {
    type Input = (ModelMatrix, DMatrix<f64>, DMatrix<f64>, DMatrix<f64>);
    type Output = Result<LinearEstimate, Box<dyn std::error::Error>>;

    fn assumptions() -> Vec<Box<dyn Assumption>> {
        todo!()
    }

    fn estimate(&self, (x, y, b, w): &Self::Input) -> Self::Output {
        let b_temp = b.clone();
        let x = x.matrix();

        Ok(LinearEstimate {
            predictor_estimate: Manipulator::new(Arc::new(move |x: &ModelMatrix| {
                ModelMatrix::from(predict(&x.matrix().insert_column(0, 1.0), &b_temp))
            })),
            predictor_confidence_interval: {
                let (lower, upper) =
                    prediction_interval(&x, &y, &b, &w, self.alpha, self.scale, self.df);
                let mut pi: ModelMatrix = DMatrix::<f64>::from_iterator(
                    x.nrows(),
                    2,
                    lower.into_iter().chain(upper.into_iter()).cloned(),
                )
                .into();
                pi.set_name_index(0, "Lower Prediction Interval");
                pi.set_name_index(1, "Upper Prediction Interval");

                pi
            },
            predictions_estimate: ModelMatrix::from(predict(&x, b)),
            predictions_confidence_interval: {
                let (lower, upper) =
                    confidence_interval(&x, &y, &b, &w, self.alpha, self.scale, self.df);
                let mut pi: ModelMatrix = DMatrix::<f64>::from_iterator(
                    x.nrows(),
                    2,
                    lower.into_iter().chain(upper.into_iter()).cloned(),
                )
                .into();
                pi.set_name_index(0, "Lower Confidence Interval");
                pi.set_name_index(1, "Upper Confidence Interval");

                pi
            },
            alpha: self.alpha,
        })
    }
}

impl Default for LinearEstimator {
    fn default() -> Self {
        Self {
            alpha: 0.05.into(),
            scale: None,
            df: None,
        }
    }
}

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

    pub fn with_alpha<A: Into<Alpha64>>(self, alpha: A) -> Self {
        Self {
            alpha: alpha.into(),
            ..self
        }
    }

    pub fn with_scale(self, scale: f64) -> Self {
        Self {
            scale: Some(scale),
            ..self
        }
    }

    pub fn with_df<R: Into<Rational64>>(self, df: R) -> Self {
        Self {
            df: Some(df.into()),
            ..self
        }
    }
}

#[derive(Clone)]
pub struct LinearEstimate {
    pub predictor_estimate: Manipulator<ModelMatrix, ModelMatrix>,
    pub predictor_confidence_interval: ModelMatrix,
    pub predictions_estimate: ModelMatrix,
    pub predictions_confidence_interval: ModelMatrix,
    alpha: Alpha64,
}

impl Debug for LinearEstimate {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RankLinearEstimate")
            // .field("predictor_estimate", &self.predictor_estimate)
            .field(
                "predictor_confidence_interval",
                &self.predictor_confidence_interval,
            )
            .field("predictions_estimate", &self.predictions_estimate)
            .field(
                "predictions_confidence_interval",
                &self.predictions_confidence_interval,
            )
            .field("alpha", &self.alpha)
            .finish()
    }
}

impl StatisticalEstimate<Manipulator<ModelMatrix, ModelMatrix>, ModelMatrix> for LinearEstimate {
    fn alpha(&self) -> Alpha64 {
        self.alpha
    }

    fn estimate(&self) -> Manipulator<ModelMatrix, ModelMatrix> {
        self.predictor_estimate.clone()
    }

    fn confidence_interval(&self) -> ModelMatrix {
        self.predictor_confidence_interval.clone()
    }
}

impl StatisticalEstimate<ModelMatrix, ModelMatrix> for LinearEstimate {
    fn alpha(&self) -> Alpha64 {
        self.alpha
    }

    fn estimate(&self) -> ModelMatrix {
        self.predictions_estimate.clone()
    }

    fn confidence_interval(&self) -> ModelMatrix {
        self.predictions_confidence_interval.clone()
    }
}