dml_core/
cate.rs

1//! Meta-learners for the conditional average treatment effect
2//! `τ(x) = E[Y(1) − Y(0) | X = x]` of a binary treatment.
3//!
4//! | learner | idea | reference |
5//! |---|---|---|
6//! | [`SLearner`] | one model `μ(x, d)`; `τ̂ = μ(x,1) − μ(x,0)` | — |
7//! | [`TLearner`] | separate `μ₀`, `μ₁` per arm; `τ̂ = μ₁ − μ₀` | — |
8//! | [`XLearner`] | impute individual effects, model them per arm, blend by propensity | Künzel et al. (2019) |
9//! | [`RLearner`] | residual-on-residual weighted regression (Robinson decomposition) | Nie & Wager (2021) |
10//! | [`DrLearner`] | regress the cross-fitted AIPW pseudo-outcome on `x` | Kennedy (2023) |
11//!
12//! The R- and DR-learners cross-fit their nuisances. The DR-learner's
13//! pseudo-outcome is a valid orthogonal signal, so it also supports
14//! inference on the CATE through [`CateFit::blp`] / [`CateFit::gate`].
15
16use ndarray::{concatenate, Array1, Array2, ArrayView1, ArrayView2, Axis};
17
18use crate::blp::{as_column, Blp};
19use crate::crossfit::{cross_fit_predict, cross_fit_predict_proba, Fold, KFold};
20use crate::data::{gather, DmlData};
21use crate::error::{DmlError, Result};
22use crate::learner::{Classifier, Regressor};
23use crate::model::irm::{ipw_residual, irm_nuisances};
24
25/// How a fitted CATE model turns covariates into effects.
26enum Predictor {
27    /// μ([x, 1]) − μ([x, 0])
28    Single(Box<dyn Regressor>),
29    /// μ₁(x) − μ₀(x)
30    TwoModels {
31        mu0: Box<dyn Regressor>,
32        mu1: Box<dyn Regressor>,
33    },
34    /// e(x)·τ₀(x) + (1 − e(x))·τ₁(x)
35    Blend {
36        tau0: Box<dyn Regressor>,
37        tau1: Box<dyn Regressor>,
38        propensity: Box<dyn Classifier>,
39    },
40    /// τ̂(x) = f(x) from a final-stage regression
41    Direct(Box<dyn Regressor>),
42}
43
44/// A fitted CATE model.
45pub struct CateFit {
46    method: &'static str,
47    n_features: usize,
48    predictor: Predictor,
49    pseudo_outcomes: Option<Array1<f64>>,
50}
51
52impl std::fmt::Debug for CateFit {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        f.debug_struct("CateFit")
55            .field("method", &self.method)
56            .field("n_features", &self.n_features)
57            .finish_non_exhaustive()
58    }
59}
60
61impl CateFit {
62    /// `"S-learner"`, `"T-learner"`, …
63    pub fn method(&self) -> &str {
64        self.method
65    }
66
67    /// Estimated effects `τ̂(x)` for each row of `x`.
68    pub fn effect(&self, x: ArrayView2<'_, f64>) -> Result<Array1<f64>> {
69        if x.ncols() != self.n_features {
70            return Err(DmlError::DimensionMismatch(format!(
71                "x has {} columns, the model was fitted with {}",
72                x.ncols(),
73                self.n_features
74            )));
75        }
76        match &self.predictor {
77            Predictor::Single(model) => {
78                let n = x.nrows();
79                let x1 = concatenate![Axis(1), x, Array2::ones((n, 1))];
80                let x0 = concatenate![Axis(1), x, Array2::zeros((n, 1))];
81                Ok(model.predict(x1.view())? - model.predict(x0.view())?)
82            }
83            Predictor::TwoModels { mu0, mu1 } => Ok(mu1.predict(x)? - mu0.predict(x)?),
84            Predictor::Blend {
85                tau0,
86                tau1,
87                propensity,
88            } => {
89                let e = propensity.predict_proba(x)?;
90                let (t0, t1) = (tau0.predict(x)?, tau1.predict(x)?);
91                Ok(&e * &t0 + &(1.0 - &e) * &t1)
92            }
93            Predictor::Direct(model) => model.predict(x),
94        }
95    }
96
97    /// Average of `τ̂(x)` over the rows of `x`.
98    pub fn average_effect(&self, x: ArrayView2<'_, f64>) -> Result<f64> {
99        Ok(self.effect(x)?.mean().unwrap_or(f64::NAN))
100    }
101
102    /// DR-learner only: the cross-fitted AIPW pseudo-outcomes `φᵢ`.
103    pub fn pseudo_outcomes(&self) -> Option<ArrayView1<'_, f64>> {
104        self.pseudo_outcomes.as_ref().map(|p| p.view())
105    }
106
107    /// DR-learner only: best linear predictor of the CATE on `basis` with
108    /// HC0 inference. A basis of ones gives the ATE with its DML standard error.
109    pub fn blp(&self, basis: ArrayView2<'_, f64>) -> Result<Blp> {
110        Blp::fit(as_column(self.signal()?).view(), basis)
111    }
112
113    /// DR-learner only: group average treatment effects for 0/1 group dummies.
114    pub fn gate(&self, groups: ArrayView2<'_, f64>) -> Result<Blp> {
115        Blp::fit_gate(as_column(self.signal()?).view(), groups)
116    }
117
118    fn signal(&self) -> Result<ArrayView1<'_, f64>> {
119        self.pseudo_outcomes().ok_or_else(|| {
120            DmlError::InvalidArgument("BLP / GATE inference is available for the DR-learner".into())
121        })
122    }
123}
124
125/// Validate a single binary treatment and split the data by arm.
126struct Arms<'a> {
127    x: ArrayView2<'a, f64>,
128    y: ArrayView1<'a, f64>,
129    d: ArrayView1<'a, f64>,
130    treated: Vec<usize>,
131    control: Vec<usize>,
132}
133
134impl<'a> Arms<'a> {
135    fn new(data: &'a DmlData) -> Result<Self> {
136        if data.n_treat() != 1 || !data.is_binary_treatment(0) {
137            return Err(DmlError::InvalidArgument(
138                "CATE learners need exactly one binary (0/1) treatment".into(),
139            ));
140        }
141        let d = data.treatment(0);
142        let treated: Vec<usize> = (0..d.len()).filter(|&i| d[i] == 1.0).collect();
143        let control: Vec<usize> = (0..d.len()).filter(|&i| d[i] == 0.0).collect();
144        if treated.is_empty() || control.is_empty() {
145            return Err(DmlError::InvalidArgument(
146                "need both treated and control observations".into(),
147            ));
148        }
149        Ok(Self {
150            x: data.x(),
151            y: data.y(),
152            d,
153            treated,
154            control,
155        })
156    }
157
158    /// Fresh copy of `template` fitted on one arm.
159    fn fit_arm(
160        &self,
161        template: &dyn Regressor,
162        rows: &[usize],
163        target: ArrayView1<'_, f64>,
164    ) -> Result<Box<dyn Regressor>> {
165        let mut model = template.fresh();
166        model.fit(
167            self.x.select(Axis(0), rows).view(),
168            gather(target, rows).view(),
169        )?;
170        Ok(model)
171    }
172}
173
174/// Folds for the cross-fitted learners: explicit, or seeded K-fold.
175fn resolve_folds(n: usize, n_folds: usize, seed: u64, folds: Option<&[Fold]>) -> Result<Vec<Fold>> {
176    match folds {
177        Some(f) => Ok(f.to_vec()),
178        None => KFold::new(n_folds, seed)?.split(n),
179    }
180}
181
182/// S-learner: one outcome model on `[x, d]`.
183pub struct SLearner {
184    model: Box<dyn Regressor>,
185}
186
187impl SLearner {
188    pub fn new(model: impl Regressor + 'static) -> Self {
189        Self::from_boxed(Box::new(model))
190    }
191
192    pub fn from_boxed(model: Box<dyn Regressor>) -> Self {
193        Self { model }
194    }
195
196    pub fn fit(&self, data: &DmlData) -> Result<CateFit> {
197        let arms = Arms::new(data)?;
198        let xd = concatenate![Axis(1), arms.x, arms.d.insert_axis(Axis(1))];
199        let mut model = self.model.fresh();
200        model.fit(xd.view(), arms.y)?;
201        Ok(CateFit {
202            method: "S-learner",
203            n_features: data.n_features(),
204            predictor: Predictor::Single(model),
205            pseudo_outcomes: None,
206        })
207    }
208}
209
210/// T-learner: separate outcome models for treated and controls.
211pub struct TLearner {
212    model: Box<dyn Regressor>,
213}
214
215impl TLearner {
216    pub fn new(model: impl Regressor + 'static) -> Self {
217        Self::from_boxed(Box::new(model))
218    }
219
220    pub fn from_boxed(model: Box<dyn Regressor>) -> Self {
221        Self { model }
222    }
223
224    pub fn fit(&self, data: &DmlData) -> Result<CateFit> {
225        let arms = Arms::new(data)?;
226        let mu0 = arms.fit_arm(self.model.as_ref(), &arms.control, arms.y)?;
227        let mu1 = arms.fit_arm(self.model.as_ref(), &arms.treated, arms.y)?;
228        Ok(CateFit {
229            method: "T-learner",
230            n_features: data.n_features(),
231            predictor: Predictor::TwoModels { mu0, mu1 },
232            pseudo_outcomes: None,
233        })
234    }
235}
236
237/// X-learner (Künzel et al., 2019).
238///
239/// 1. Fit `μ₀`, `μ₁` per arm (`outcome` learner).
240/// 2. Impute effects: `Y − μ₀(X)` for treated, `μ₁(X) − Y` for controls.
241/// 3. Model the imputed effects per arm (`effect` learner): `τ₁`, `τ₀`.
242/// 4. Blend: `τ̂(x) = e(x)·τ₀(x) + (1 − e(x))·τ₁(x)` with propensity `e`.
243pub struct XLearner {
244    outcome: Box<dyn Regressor>,
245    effect: Box<dyn Regressor>,
246    propensity: Box<dyn Classifier>,
247}
248
249impl XLearner {
250    pub fn new(
251        outcome: impl Regressor + 'static,
252        effect: impl Regressor + 'static,
253        propensity: impl Classifier + 'static,
254    ) -> Self {
255        Self::from_boxed(Box::new(outcome), Box::new(effect), Box::new(propensity))
256    }
257
258    pub fn from_boxed(
259        outcome: Box<dyn Regressor>,
260        effect: Box<dyn Regressor>,
261        propensity: Box<dyn Classifier>,
262    ) -> Self {
263        Self {
264            outcome,
265            effect,
266            propensity,
267        }
268    }
269
270    pub fn fit(&self, data: &DmlData) -> Result<CateFit> {
271        let arms = Arms::new(data)?;
272        let mu0 = arms.fit_arm(self.outcome.as_ref(), &arms.control, arms.y)?;
273        let mu1 = arms.fit_arm(self.outcome.as_ref(), &arms.treated, arms.y)?;
274        // Imputed individual effects, laid out over all rows; each arm's
275        // effect model only reads its own rows.
276        let mut imputed = Array1::<f64>::zeros(arms.y.len());
277        let x_t = arms.x.select(Axis(0), &arms.treated);
278        let x_c = arms.x.select(Axis(0), &arms.control);
279        for (&i, m0) in arms.treated.iter().zip(mu0.predict(x_t.view())?) {
280            imputed[i] = arms.y[i] - m0;
281        }
282        for (&i, m1) in arms.control.iter().zip(mu1.predict(x_c.view())?) {
283            imputed[i] = m1 - arms.y[i];
284        }
285        let tau1 = arms.fit_arm(self.effect.as_ref(), &arms.treated, imputed.view())?;
286        let tau0 = arms.fit_arm(self.effect.as_ref(), &arms.control, imputed.view())?;
287        let mut propensity = self.propensity.fresh();
288        propensity.fit(arms.x, arms.d)?;
289        Ok(CateFit {
290            method: "X-learner",
291            n_features: data.n_features(),
292            predictor: Predictor::Blend {
293                tau0,
294                tau1,
295                propensity,
296            },
297            pseudo_outcomes: None,
298        })
299    }
300}
301
302/// R-learner (Nie & Wager, 2021): with cross-fitted `m̂(x) = E[Y | X]` and
303/// propensity `ê(x)`, minimise `Σ (Ỹ − τ(X)·D̃)²`, i.e. a weighted regression
304/// of `Ỹ/D̃` on `X` with weights `D̃²`. The `effect` learner must support
305/// [`Regressor::fit_weighted`].
306pub struct RLearner {
307    outcome: Box<dyn Regressor>,
308    propensity: Box<dyn Classifier>,
309    effect: Box<dyn Regressor>,
310    n_folds: usize,
311    seed: u64,
312    clipping_threshold: f64,
313}
314
315impl RLearner {
316    pub fn new(
317        outcome: impl Regressor + 'static,
318        propensity: impl Classifier + 'static,
319        effect: impl Regressor + 'static,
320    ) -> Self {
321        Self::from_boxed(Box::new(outcome), Box::new(propensity), Box::new(effect))
322    }
323
324    pub fn from_boxed(
325        outcome: Box<dyn Regressor>,
326        propensity: Box<dyn Classifier>,
327        effect: Box<dyn Regressor>,
328    ) -> Self {
329        Self {
330            outcome,
331            propensity,
332            effect,
333            n_folds: 5,
334            seed: 0,
335            clipping_threshold: 1e-2,
336        }
337    }
338
339    pub fn n_folds(mut self, n_folds: usize) -> Self {
340        self.n_folds = n_folds;
341        self
342    }
343
344    pub fn seed(mut self, seed: u64) -> Self {
345        self.seed = seed;
346        self
347    }
348
349    /// Propensities are clipped to `[t, 1 − t]` (default 0.01).
350    pub fn clipping_threshold(mut self, t: f64) -> Self {
351        self.clipping_threshold = t;
352        self
353    }
354
355    pub fn fit(&self, data: &DmlData) -> Result<CateFit> {
356        self.fit_impl(data, None)
357    }
358
359    pub fn fit_with_folds(&self, data: &DmlData, folds: &[Fold]) -> Result<CateFit> {
360        self.fit_impl(data, Some(folds))
361    }
362
363    fn fit_impl(&self, data: &DmlData, folds: Option<&[Fold]>) -> Result<CateFit> {
364        let arms = Arms::new(data)?;
365        check_clipping(self.clipping_threshold)?;
366        let folds = resolve_folds(data.n_obs(), self.n_folds, self.seed, folds)?;
367        let m_hat = cross_fit_predict(self.outcome.as_ref(), arms.x, arms.y, &folds)?;
368        let t = self.clipping_threshold;
369        let e_hat = cross_fit_predict_proba(self.propensity.as_ref(), arms.x, arms.d, &folds)?
370            .mapv(|p| p.clamp(t, 1.0 - t));
371        let y_res: Array1<f64> = &arms.y - &m_hat;
372        let d_res: Array1<f64> = &arms.d - &e_hat;
373        let target = &y_res / &d_res;
374        let weights = &d_res * &d_res;
375        let mut effect = self.effect.fresh();
376        effect.fit_weighted(arms.x, target.view(), weights.view())?;
377        Ok(CateFit {
378            method: "R-learner",
379            n_features: data.n_features(),
380            predictor: Predictor::Direct(effect),
381            pseudo_outcomes: None,
382        })
383    }
384}
385
386/// DR-learner (Kennedy, 2023): cross-fit `ĝ₀`, `ĝ₁` and the propensity `m̂`
387/// exactly as [`Irm`](crate::Irm) does, form the AIPW pseudo-outcome
388/// `φ = ĝ₁ − ĝ₀ + D(Y − ĝ₁)/m̂ − (1 − D)(Y − ĝ₀)/(1 − m̂)`, and regress it
389/// on `X` with the `effect` learner.
390pub struct DrLearner {
391    outcome: Box<dyn Regressor>,
392    propensity: Box<dyn Classifier>,
393    effect: Box<dyn Regressor>,
394    n_folds: usize,
395    seed: u64,
396    clipping_threshold: f64,
397}
398
399impl DrLearner {
400    pub fn new(
401        outcome: impl Regressor + 'static,
402        propensity: impl Classifier + 'static,
403        effect: impl Regressor + 'static,
404    ) -> Self {
405        Self::from_boxed(Box::new(outcome), Box::new(propensity), Box::new(effect))
406    }
407
408    pub fn from_boxed(
409        outcome: Box<dyn Regressor>,
410        propensity: Box<dyn Classifier>,
411        effect: Box<dyn Regressor>,
412    ) -> Self {
413        Self {
414            outcome,
415            propensity,
416            effect,
417            n_folds: 5,
418            seed: 0,
419            clipping_threshold: 1e-2,
420        }
421    }
422
423    pub fn n_folds(mut self, n_folds: usize) -> Self {
424        self.n_folds = n_folds;
425        self
426    }
427
428    pub fn seed(mut self, seed: u64) -> Self {
429        self.seed = seed;
430        self
431    }
432
433    /// Propensities are clipped to `[t, 1 − t]` (default 0.01).
434    pub fn clipping_threshold(mut self, t: f64) -> Self {
435        self.clipping_threshold = t;
436        self
437    }
438
439    pub fn fit(&self, data: &DmlData) -> Result<CateFit> {
440        self.fit_impl(data, None)
441    }
442
443    pub fn fit_with_folds(&self, data: &DmlData, folds: &[Fold]) -> Result<CateFit> {
444        self.fit_impl(data, Some(folds))
445    }
446
447    fn fit_impl(&self, data: &DmlData, folds: Option<&[Fold]>) -> Result<CateFit> {
448        let arms = Arms::new(data)?;
449        check_clipping(self.clipping_threshold)?;
450        let folds = resolve_folds(data.n_obs(), self.n_folds, self.seed, folds)?;
451        let (g0, g1, m) = irm_nuisances(
452            self.outcome.as_ref(),
453            self.propensity.as_ref(),
454            arms.x,
455            arms.y,
456            arms.d,
457            &folds,
458            self.clipping_threshold,
459            crate::propensity::Calibration::None,
460        )?;
461        let d = arms.d.to_owned();
462        let phi = &g1 - &g0 + &ipw_residual(arms.y, &d, &g0, &g1, &m);
463        let mut effect = self.effect.fresh();
464        effect.fit(arms.x, phi.view())?;
465        Ok(CateFit {
466            method: "DR-learner",
467            n_features: data.n_features(),
468            predictor: Predictor::Direct(effect),
469            pseudo_outcomes: Some(phi),
470        })
471    }
472}
473
474fn check_clipping(t: f64) -> Result<()> {
475    if t > 0.0 && t < 0.5 {
476        Ok(())
477    } else {
478        Err(DmlError::InvalidArgument(
479            "clipping_threshold must be in (0, 0.5)".into(),
480        ))
481    }
482}
dml_core/cate.rs

dml_core/
cate.rs
