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    fn predict_proba(&self, x: ArrayView2<'_, f64>) -> Result<Array1<f64>>;
54    fn fresh(&self) -> Box<dyn Classifier>;
55}
56
57/// Centred (and, with weights, `√w`-scaled) data for the linear learners, so
58/// the intercept drops out: `b = ȳ − x̄ᵀβ`.
59pub(crate) struct Centered {
60    pub x_mean: Array1<f64>,
61    pub y_mean: f64,
62    pub xc: Array2<f64>,
63    pub yc: Array1<f64>,
64}
65
66/// With `w`, rows are scaled by `√w` after weighted centring. If
67/// `normalize_weights`, weights are first rescaled to sum to `n` (as
68/// scikit-learn's `Lasso` does; `Ridge` uses them as given).
69pub(crate) fn center(
70    x: ArrayView2<'_, f64>,
71    y: ArrayView1<'_, f64>,
72    w: Option<ArrayView1<'_, f64>>,
73    normalize_weights: bool,
74) -> Result<Centered> {
75    let n = x.nrows();
76    if y.len() != n {
77        return Err(DmlError::DimensionMismatch("fit x/y rows".into()));
78    }
79    if n == 0 {
80        return Err(DmlError::InvalidArgument("empty training set".into()));
81    }
82    let Some(w) = w else {
83        let x_mean = x.mean_axis(Axis(0)).expect("n > 0");
84        let y_mean = y.mean().expect("n > 0");
85        return Ok(Centered {
86            xc: &x - &x_mean,
87            yc: &y - y_mean,
88            x_mean,
89            y_mean,
90        });
91    };
92    if w.len() != n {
93        return Err(DmlError::DimensionMismatch("sample weight length".into()));
94    }
95    let total = w.sum();
96    if !w.iter().all(|&v| v >= 0.0 && v.is_finite()) || total <= 0.0 {
97        return Err(DmlError::InvalidArgument(
98            "sample weights must be non-negative, finite and not all zero".into(),
99        ));
100    }
101    let w: Array1<f64> = if normalize_weights {
102        &w * (n as f64 / total)
103    } else {
104        w.to_owned()
105    };
106    let total = w.sum();
107    let x_mean = x.t().dot(&w) / total;
108    let y_mean = y.dot(&w) / total;
109    let sqrt_w = w.mapv(f64::sqrt);
110    let xc = (&x - &x_mean) * sqrt_w.view().insert_axis(Axis(1));
111    let yc = (&y - y_mean) * &sqrt_w;
112    Ok(Centered {
113        x_mean,
114        y_mean,
115        xc,
116        yc,
117    })
118}
dml_core/learner/mod.rs

dml_core/learner/
mod.rs
