use std::marker::PhantomData;
use strafe_trait::ModelBuilder;
use strafe_type::{Alpha64, ModelMatrix};
use crate::regression::glm::{
family::core::Family, generalized_linear_model, link::core::Link, GeneralizedLinearRegression,
};
pub struct GeneralizedLinearRegressionBuilder<L: Link, F: Family<L>> {
pub(crate) x: Option<ModelMatrix>,
pub(crate) y: Option<ModelMatrix>,
pub(crate) w: Option<ModelMatrix>,
pub(crate) alpha: Alpha64,
pub(crate) family: F,
pub(crate) _link: PhantomData<L>,
}
impl<L: Link, F: Family<L>> GeneralizedLinearRegressionBuilder<L, F> {
pub fn new(family: F) -> Self {
Self {
x: None,
y: None,
w: None,
alpha: 0.05.into(),
family,
_link: PhantomData::default(),
}
}
}
impl<L: Link, F: Family<L>> ModelBuilder for GeneralizedLinearRegressionBuilder<L, F> {
type Model = GeneralizedLinearRegression<L, F>;
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) -> GeneralizedLinearRegression<L, F> {
let (b, aic, null_deviance, residual_deviance, final_weights, eta, mu) =
generalized_linear_model(
&self.y.as_ref().unwrap().matrix(),
&self.x.as_ref().unwrap().matrix(),
&self.w.as_ref().unwrap().matrix(),
&self.family,
);
let x1 = ModelMatrix::from(
self.x
.as_ref()
.unwrap()
.matrix()
.to_owned()
.insert_column(0, 1.0),
);
let model_data = (
x1.clone(),
self.y.as_ref().unwrap().matrix().clone(),
b.clone(),
final_weights.clone(),
);
GeneralizedLinearRegression {
x: self.x.as_ref().unwrap().clone(),
y: self.y.as_ref().unwrap().clone(),
w: self.w.as_ref().unwrap().clone(),
alpha: self.alpha,
family: self.family,
_link: self._link,
x1,
b,
model_data,
eta,
mu,
aic,
null_deviance,
residual_deviance,
final_weights,
}
}
}