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

dml_core/learner/
mod.rs
