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

dml_core/learner/
mod.rs
