greeners_ols/
wls.rs

1use crate::ols::OlsResult;
2use crate::ols::OLS;
3use greeners_core::error::GreenersError;
4use greeners_core::{CovarianceType, DataFrame, Formula};
5use ndarray::{Array1, Array2};
6
7/// Weighted Least Squares estimator.
8///
9/// Transforms the model by sqrt(w): y* = sqrt(w)*y, X* = sqrt(w)*X,
10/// then runs OLS on the transformed data.
11pub struct WLS;
12
13impl WLS {
14    /// Fit WLS from a formula, DataFrame, and weight column name.
15    pub fn from_formula(
16        formula: &Formula,
17        data: &DataFrame,
18        weights: &Array1<f64>,
19        cov_type: CovarianceType,
20    ) -> Result<OlsResult, GreenersError> {
21        let (y, x) = data.to_design_matrix(formula)?;
22        let var_names = data.formula_var_names(formula)?;
23        Self::fit_with_names(&y, &x, weights, cov_type, Some(var_names))
24    }
25
26    /// Fit WLS from arrays.
27    pub fn fit(
28        y: &Array1<f64>,
29        x: &Array2<f64>,
30        weights: &Array1<f64>,
31        cov_type: CovarianceType,
32    ) -> Result<OlsResult, GreenersError> {
33        Self::fit_with_names(y, x, weights, cov_type, None)
34    }
35
36    /// Fit WLS with variable names.
37    pub fn fit_with_names(
38        y: &Array1<f64>,
39        x: &Array2<f64>,
40        weights: &Array1<f64>,
41        cov_type: CovarianceType,
42        variable_names: Option<Vec<String>>,
43    ) -> Result<OlsResult, GreenersError> {
44        let n = y.len();
45        if weights.len() != n {
46            return Err(GreenersError::ShapeMismatch(format!(
47                "weights length ({}) must match observations ({})",
48                weights.len(),
49                n
50            )));
51        }
52        if x.nrows() != n {
53            return Err(GreenersError::ShapeMismatch(format!(
54                "X rows ({}) must match y length ({})",
55                x.nrows(),
56                n
57            )));
58        }
59
60        // Validate weights: must be positive
61        if weights.iter().any(|&w| w <= 0.0 || !w.is_finite()) {
62            return Err(GreenersError::InvalidOperation(
63                "Weights must be positive and finite".into(),
64            ));
65        }
66
67        // Transform: multiply by sqrt(w)
68        let sqrt_w = weights.mapv(f64::sqrt);
69
70        let y_star = &sqrt_w * y;
71        let mut x_star = x.clone();
72        for i in 0..n {
73            x_star.row_mut(i).mapv_inplace(|val| val * sqrt_w[i]);
74        }
75
76        let has_intercept =
77            (0..x.ncols()).any(|j| x.column(j).iter().all(|&val| (val - 1.0).abs() < 1e-12));
78
79        OLS::fit_internal(
80            &y_star,
81            &x_star,
82            cov_type,
83            variable_names,
84            Some(has_intercept),
85        )
86    }
87}
greeners_ols/wls.rs

greeners_ols/
wls.rs
