use nalgebra::DMatrix;
use strafe_trait::ModelBuilder;
use strafe_type::{Alpha64, ModelMatrix};
use crate::regression::lm::{linear_regression_standardized, LeastSquaresRegression};
pub struct LeastSquaresRegressionBuilder {
pub(crate) x: Option<ModelMatrix>,
pub(crate) y: Option<ModelMatrix>,
pub(crate) alpha: Alpha64,
pub(crate) w: Option<ModelMatrix>,
}
impl LeastSquaresRegressionBuilder {
pub fn new() -> Self {
Self {
x: None,
y: None,
alpha: 0.05.into(),
w: None,
}
}
}
impl ModelBuilder for LeastSquaresRegressionBuilder {
type Model = LeastSquaresRegression;
fn with_x(self, x: &ModelMatrix) -> Self {
Self {
x: Some(x.clone()),
..self
}
}
fn with_y(self, y: &ModelMatrix) -> Self {
Self {
y: Some(y.clone()),
..self
}
}
fn with_weights(self, weights: &ModelMatrix) -> Self {
Self {
w: Some(weights.clone()),
..self
}
}
fn with_alpha<A: Into<Alpha64>>(self, alpha: A) -> Self {
Self {
alpha: alpha.into(),
..self
}
}
fn build(self) -> Self::Model {
let b = linear_regression_standardized(
&self.x.as_ref().unwrap().matrix(),
&self.y.as_ref().unwrap().matrix(),
);
let x1 = self.x.as_ref().unwrap().matrix().insert_column(0, 1.0);
let w = ModelMatrix::from(DMatrix::from_iterator(
x1.shape().0,
1,
(0..x1.shape().0).map(|_| 1.0),
));
let model_data = (
ModelMatrix::from(x1.clone()),
self.y.as_ref().unwrap().matrix(),
b.clone(),
w.matrix(),
);
LeastSquaresRegression {
x: self.x.as_ref().unwrap().clone(),
y: self.y.as_ref().unwrap().clone(),
alpha: self.alpha,
x1,
w,
b,
model_data,
}
}
}