dml_core/learner/
mod.rs

1//! Nuisance learners: the [`Regressor`] / [`Classifier`] traits and the
2//! built-in implementations.
3
4use ndarray::{Array1, Array2, ArrayView1, ArrayView2, Axis};
5
6use crate::error::{DmlError, Result};
7
8mod lasso;
9mod logistic;
10mod ols;
11
12pub use lasso::Lasso;
13pub use logistic::LogisticRegression;
14pub use ols::Ols;
15
16/// A regression learner for nuisance functions such as `E[Y | X]` or `E[D | X]`.
17///
18/// `fresh` returns an unfitted copy with the same hyperparameters, so
19/// cross-fitting can train an independent model on each fold.
20///
21/// `Send + Sync` lets folds be fitted on parallel threads.
22pub trait Regressor: Send + Sync {
23    fn fit(&mut self, x: ArrayView2<'_, f64>, y: ArrayView1<'_, f64>) -> Result<()>;
24
25    /// Fit minimising a `w`-weighted loss (needed by the [`RLearner`](crate::RLearner)).
26    /// The default reports that weights are unsupported.
27    fn fit_weighted(
28        &mut self,
29        _x: ArrayView2<'_, f64>,
30        _y: ArrayView1<'_, f64>,
31        _w: ArrayView1<'_, f64>,
32    ) -> Result<()> {
33        Err(DmlError::InvalidArgument(
34            "this learner does not support sample weights".into(),
35        ))
36    }
37
38    fn predict(&self, x: ArrayView2<'_, f64>) -> Result<Array1<f64>>;
39    fn fresh(&self) -> Box<dyn Regressor>;
40}
41
42/// A classification learner returning `P(label = 1 | X)`, used for
43/// propensity scores in models with a binary treatment such as [`Irm`](crate::Irm).
44pub trait Classifier: Send + Sync {
45    fn fit(&mut self, x: ArrayView2<'_, f64>, y: ArrayView1<'_, f64>) -> Result<()>;
46    fn predict_proba(&self, x: ArrayView2<'_, f64>) -> Result<Array1<f64>>;
47    fn fresh(&self) -> Box<dyn Classifier>;
48}
49
50/// Centred (and, with weights, `√w`-scaled) data for the linear learners, so
51/// the intercept drops out: `b = ȳ − x̄ᵀβ`.
52pub(crate) struct Centered {
53    pub x_mean: Array1<f64>,
54    pub y_mean: f64,
55    pub xc: Array2<f64>,
56    pub yc: Array1<f64>,
57}
58
59/// With `w`, rows are scaled by `√w` after weighted centring. If
60/// `normalize_weights`, weights are first rescaled to sum to `n` (as
61/// scikit-learn's `Lasso` does; `Ridge` uses them as given).
62pub(crate) fn center(
63    x: ArrayView2<'_, f64>,
64    y: ArrayView1<'_, f64>,
65    w: Option<ArrayView1<'_, f64>>,
66    normalize_weights: bool,
67) -> Result<Centered> {
68    let n = x.nrows();
69    if y.len() != n {
70        return Err(DmlError::DimensionMismatch("fit x/y rows".into()));
71    }
72    if n == 0 {
73        return Err(DmlError::InvalidArgument("empty training set".into()));
74    }
75    let Some(w) = w else {
76        let x_mean = x.mean_axis(Axis(0)).expect("n > 0");
77        let y_mean = y.mean().expect("n > 0");
78        return Ok(Centered {
79            xc: &x - &x_mean,
80            yc: &y - y_mean,
81            x_mean,
82            y_mean,
83        });
84    };
85    if w.len() != n {
86        return Err(DmlError::DimensionMismatch("sample weight length".into()));
87    }
88    let total = w.sum();
89    if !w.iter().all(|&v| v >= 0.0 && v.is_finite()) || total <= 0.0 {
90        return Err(DmlError::InvalidArgument(
91            "sample weights must be non-negative, finite and not all zero".into(),
92        ));
93    }
94    let w: Array1<f64> = if normalize_weights {
95        &w * (n as f64 / total)
96    } else {
97        w.to_owned()
98    };
99    let total = w.sum();
100    let x_mean = x.t().dot(&w) / total;
101    let y_mean = y.dot(&w) / total;
102    let sqrt_w = w.mapv(f64::sqrt);
103    let xc = (&x - &x_mean) * sqrt_w.view().insert_axis(Axis(1));
104    let yc = (&y - y_mean) * &sqrt_w;
105    Ok(Centered {
106        x_mean,
107        y_mean,
108        xc,
109        yc,
110    })
111}
dml_core/learner/mod.rs

dml_core/learner/
mod.rs
