use crate::error::GreenersError;
use crate::ols::OlsResult;
use crate::{CovarianceType, DataFrame, Formula, OLS};
use ndarray::{Array1, Array2};
pub struct WLS;
impl WLS {
pub fn from_formula(
formula: &Formula,
data: &DataFrame,
weights: &Array1<f64>,
cov_type: CovarianceType,
) -> Result<OlsResult, GreenersError> {
let (y, x) = data.to_design_matrix(formula)?;
let var_names = data.formula_var_names(formula)?;
Self::fit_with_names(&y, &x, weights, cov_type, Some(var_names))
}
pub fn fit(
y: &Array1<f64>,
x: &Array2<f64>,
weights: &Array1<f64>,
cov_type: CovarianceType,
) -> Result<OlsResult, GreenersError> {
Self::fit_with_names(y, x, weights, cov_type, None)
}
pub fn fit_with_names(
y: &Array1<f64>,
x: &Array2<f64>,
weights: &Array1<f64>,
cov_type: CovarianceType,
variable_names: Option<Vec<String>>,
) -> Result<OlsResult, GreenersError> {
let n = y.len();
if weights.len() != n {
return Err(GreenersError::ShapeMismatch(format!(
"weights length ({}) must match observations ({})",
weights.len(),
n
)));
}
if x.nrows() != n {
return Err(GreenersError::ShapeMismatch(format!(
"X rows ({}) must match y length ({})",
x.nrows(),
n
)));
}
if weights.iter().any(|&w| w <= 0.0 || !w.is_finite()) {
return Err(GreenersError::InvalidOperation(
"Weights must be positive and finite".into(),
));
}
let sqrt_w = weights.mapv(f64::sqrt);
let y_star = &sqrt_w * y;
let mut x_star = x.clone();
for i in 0..n {
x_star.row_mut(i).mapv_inplace(|val| val * sqrt_w[i]);
}
let has_intercept = (0..x.ncols()).any(|j| {
x.column(j).iter().all(|&val| (val - 1.0).abs() < 1e-12)
});
OLS::fit_internal(&y_star, &x_star, cov_type, variable_names, Some(has_intercept))
}
}