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

dml_core/learner/
mod.rs
