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

dml_core/learner/
mod.rs
