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

dml_core/learner/
mod.rs
