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