dml_core/model/
ssm.rs

1//! Sample selection model (SSM): the average treatment effect of a binary
2//! treatment when the outcome is observed only for a selected subsample
3//! (attrition, non-response), as DoubleML's `DoubleMLSSM` (Bia, Huber and
4//! Lafférs, 2024).
5//!
6//! ```text
7//! Y = g(D, X) + U,   observed only where S = 1
8//! m(X) = P(D = 1 | X),   π(D, X) = P(S = 1 | D, X)
9//! ```
10//!
11//! - [`SsmScore::MissingAtRandom`]: selection depends on the treatment and the
12//!   covariates only. `π` is learned from `(X, D)`, `m` from `X`, and the
13//!   outcome models on the selected treated / controls.
14//! - [`SsmScore::Nonignorable`]: selection may depend on unobservables, with
15//!   an instrument `Z` for selection. Each training fold is split in halves
16//!   (stratified by `D + 2S`): `π` is learned from `(X, D, Z)` on the first
17//!   half, and `m` and the outcome models on the second, with the predicted
18//!   selection probability as an extra covariate.
19
20use ndarray::{concatenate, Array1, ArrayView1, ArrayView2, Axis};
21
22use super::quantile::half_split;
23use super::{dml_model_api, LinearScoreModel, NuisanceTarget, ScoreEstimate, ScoreInput};
24use crate::crossfit::{cross_fit_predict_proba, cross_fit_predict_subset, Fold, SampleSplitting};
25use crate::data::{gather, is_binary, DmlData};
26use crate::error::{DmlError, Result};
27use crate::learner::{Classifier, Regressor};
28use crate::propensity::{adjust_propensity, Calibration};
29use crate::score::LinearScore;
30
31/// Selection mechanism of [`Ssm`].
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
33pub enum SsmScore {
34    /// Selection depends on `D` and `X` only.
35    #[default]
36    MissingAtRandom,
37    /// Selection may depend on unobservables; needs an instrument for selection.
38    Nonignorable,
39}
40
41/// Average treatment effect of a binary treatment when the outcome is
42/// observed only where the selection indicator `S` is 1 (attrition,
43/// non-response), as DoubleML's `DoubleMLSSM`. Requires a selection indicator
44/// ([`DmlData::with_selection_indicator`]) and, for
45/// [`SsmScore::Nonignorable`], an instrument for selection
46/// ([`DmlData::with_instrument`]).
47///
48/// The score is doubly robust in the outcome models of the selected treated
49/// and controls, the propensity `P(D = 1 | X)` and the selection probability
50/// `P(S = 1 | D, X)`. Under nonignorable selection each training fold is
51/// split in halves: the selection probability is learned on the first, and
52/// the other nuisances on the second with the predicted selection
53/// probability as a covariate.
54pub struct Ssm {
55    ml_g: Box<dyn Regressor>,
56    ml_pi: Box<dyn Classifier>,
57    ml_m: Box<dyn Classifier>,
58    score: SsmScore,
59    clipping_threshold: f64,
60    calibration: Calibration,
61    normalize_ipw: bool,
62    splitting: SampleSplitting,
63}
64
65impl Ssm {
66    /// `ml_g` learns `E[Y | D, X]` on the selected, `ml_pi` the selection
67    /// probability `P(S = 1 | D, X)`, `ml_m` the propensity `P(D = 1 | X)`.
68    pub fn new(
69        ml_g: impl Regressor + 'static,
70        ml_pi: impl Classifier + 'static,
71        ml_m: impl Classifier + 'static,
72    ) -> Self {
73        Self::from_boxed(Box::new(ml_g), Box::new(ml_pi), Box::new(ml_m))
74    }
75
76    pub fn from_boxed(
77        ml_g: Box<dyn Regressor>,
78        ml_pi: Box<dyn Classifier>,
79        ml_m: Box<dyn Classifier>,
80    ) -> Self {
81        Self {
82            ml_g,
83            ml_pi,
84            ml_m,
85            score: SsmScore::default(),
86            clipping_threshold: 1e-2,
87            calibration: Calibration::None,
88            normalize_ipw: false,
89            splitting: SampleSplitting::default(),
90        }
91    }
92
93    pub fn score(mut self, score: SsmScore) -> Self {
94        self.score = score;
95        self
96    }
97
98    /// Treatment propensities are clipped to `[t, 1 − t]` (default `t = 0.01`).
99    /// Selection probabilities are not clipped, as in DoubleML.
100    pub fn clipping_threshold(mut self, t: f64) -> Self {
101        self.clipping_threshold = t;
102        self
103    }
104    /// Calibrate the propensity scores before clipping (default: none), as
105    /// DoubleML's `ps_processor_config`.
106    pub fn calibration(mut self, calibration: Calibration) -> Self {
107        self.calibration = calibration;
108        self
109    }
110
111    /// Normalise the inverse-probability weights within treated and controls.
112    pub fn normalize_ipw(mut self, normalize: bool) -> Self {
113        self.normalize_ipw = normalize;
114        self
115    }
116
117    dml_model_api!();
118
119    /// Nonignorable selection: nested cross-fitting within each fold.
120    /// Returns `(g0, g1, π, m)` out of fold.
121    fn nonignorable_nuisances(
122        &self,
123        input: &ScoreInput<'_>,
124        s: ArrayView1<'_, f64>,
125        z: ArrayView1<'_, f64>,
126        folds: &[Fold],
127    ) -> Result<[Array1<f64>; 4]> {
128        let ScoreInput { y, d, x, .. } = *input;
129        let n = y.len();
130        let strata: Vec<f64> = (0..n).map(|i| d[i] + 2.0 * s[i]).collect();
131        let dxz = concatenate![Axis(1), x, d.insert_axis(Axis(1)), z.insert_axis(Axis(1))];
132        let fit_fold = |fold: &Fold| -> Result<[Array1<f64>; 4]> {
133            let (train1, train2) = half_split(&fold.train, &strata)?;
134            // Selection probability from the first half, predicted for every
135            // observation and added to the covariates.
136            let pi_all = fit_proba(self.ml_pi.as_ref(), dxz.view(), &train1, s, dxz.view())?;
137            let xpi = concatenate![Axis(1), x, pi_all.view().insert_axis(Axis(1))];
138            let xpi_test = xpi.select(Axis(0), &fold.test);
139            let m = fit_proba(self.ml_m.as_ref(), xpi.view(), &train2, d, xpi_test.view())?;
140            // Outcome models on the selected of the second half, in index order
141            // (DoubleML's `np.intersect1d`).
142            let mut sorted2 = train2.clone();
143            sorted2.sort_unstable();
144            let outcome = |level: f64| -> Result<Array1<f64>> {
145                let rows: Vec<usize> = sorted2
146                    .iter()
147                    .copied()
148                    .filter(|&i| d[i] == level && s[i] == 1.0)
149                    .collect();
150                if rows.is_empty() {
151                    return Err(DmlError::InvalidArgument(
152                        "a training fold has no selected observations in a treatment arm".into(),
153                    ));
154                }
155                let mut model = self.ml_g.fresh();
156                model.fit(xpi.select(Axis(0), &rows).view(), gather(y, &rows).view())?;
157                model.predict(xpi_test.view())
158            };
159            let pi_test = gather(pi_all.view(), &fold.test);
160            Ok([outcome(0.0)?, outcome(1.0)?, pi_test, m])
161        };
162        #[cfg(feature = "parallel")]
163        let per_fold: Vec<[Array1<f64>; 4]> = {
164            use rayon::prelude::*;
165            folds.par_iter().map(fit_fold).collect::<Result<_>>()?
166        };
167        #[cfg(not(feature = "parallel"))]
168        let per_fold: Vec<[Array1<f64>; 4]> = folds.iter().map(fit_fold).collect::<Result<_>>()?;
169
170        let mut out: [Array1<f64>; 4] = std::array::from_fn(|_| Array1::from_elem(n, f64::NAN));
171        for (fold, preds) in folds.iter().zip(per_fold) {
172            for (dst, src) in out.iter_mut().zip(preds) {
173                for (k, &i) in fold.test.iter().enumerate() {
174                    dst[i] = src[k];
175                }
176            }
177        }
178        Ok(out)
179    }
180}
181
182/// Fit a fresh copy of `learner` on `rows` (in that order); predict `x_pred`.
183fn fit_proba(
184    learner: &dyn Classifier,
185    x: ArrayView2<'_, f64>,
186    rows: &[usize],
187    labels: ArrayView1<'_, f64>,
188    x_pred: ArrayView2<'_, f64>,
189) -> Result<Array1<f64>> {
190    let mut model = learner.fresh();
191    model.fit(x.select(Axis(0), rows).view(), gather(labels, rows).view())?;
192    model.predict_proba(x_pred)
193}
194
195impl LinearScoreModel for Ssm {
196    fn name(&self) -> String {
197        match self.score {
198            SsmScore::MissingAtRandom => "SSM (missing at random)".into(),
199            SsmScore::Nonignorable => "SSM (nonignorable)".into(),
200        }
201    }
202
203    fn validate(&self, data: &DmlData) -> Result<()> {
204        if !(0..data.n_treat()).all(|j| data.is_binary_treatment(j)) {
205            return Err(DmlError::InvalidArgument(
206                "SSM needs binary (0/1) treatments".into(),
207            ));
208        }
209        if data.selection().is_none() {
210            return Err(DmlError::InvalidArgument(
211                "SSM needs a selection indicator (DmlData::with_selection_indicator)".into(),
212            ));
213        }
214        if data.n_instruments() > 1 {
215            return Err(DmlError::InvalidArgument(
216                "SSM needs exactly one instrument".into(),
217            ));
218        }
219        if self.score == SsmScore::Nonignorable && data.z().is_none() {
220            return Err(DmlError::InvalidArgument(
221                "the nonignorable score needs an instrument for selection (DmlData::with_instrument)"
222                    .into(),
223            ));
224        }
225        if !(self.clipping_threshold > 0.0 && self.clipping_threshold < 0.5) {
226            return Err(DmlError::InvalidArgument(
227                "clipping_threshold must be in (0, 0.5)".into(),
228            ));
229        }
230        Ok(())
231    }
232
233    /// Stratify by treatment and selection (`D + 2S`), as DoubleML.
234    fn strata(&self, data: &DmlData) -> Option<Vec<usize>> {
235        let (d, s) = (data.treatment(0), data.selection()?);
236        Some(
237            d.iter()
238                .zip(s)
239                .map(|(&d, &s)| (d + 2.0 * s) as usize)
240                .collect(),
241        )
242    }
243
244    fn estimate_score(&self, input: &ScoreInput<'_>, folds: &[Fold]) -> Result<ScoreEstimate> {
245        let ScoreInput { y, d, x, z, s, .. } = *input;
246        let s =
247            s.ok_or_else(|| DmlError::InvalidArgument("SSM needs a selection indicator".into()))?;
248        if !is_binary(d) {
249            return Err(DmlError::InvalidArgument(
250                "SSM needs a binary treatment".into(),
251            ));
252        }
253        let n = y.len();
254        let [g0, g1, pi, m] = match self.score {
255            SsmScore::MissingAtRandom => {
256                let dx = concatenate![Axis(1), x, d.insert_axis(Axis(1))];
257                let pi = cross_fit_predict_proba(self.ml_pi.as_ref(), dx.view(), s, folds)?;
258                let m = cross_fit_predict_proba(self.ml_m.as_ref(), x, d, folds)?;
259                let selected = |level: f64| -> Vec<bool> {
260                    (0..n).map(|i| d[i] == level && s[i] == 1.0).collect()
261                };
262                let g1 = cross_fit_predict_subset(self.ml_g.as_ref(), x, y, folds, &selected(1.0))?;
263                let g0 = cross_fit_predict_subset(self.ml_g.as_ref(), x, y, folds, &selected(0.0))?;
264                [g0, g1, pi, m]
265            }
266            SsmScore::Nonignorable => {
267                let z = z.ok_or_else(|| {
268                    DmlError::InvalidArgument("the nonignorable score needs an instrument".into())
269                })?;
270                self.nonignorable_nuisances(input, s, z, folds)?
271            }
272        };
273        let t = self.clipping_threshold;
274        let m = adjust_propensity(m, d, folds, t, self.calibration)?;
275
276        // ψ_b = ψ_b1 − ψ_b0, each an IPW-corrected outcome prediction; written
277        // in DoubleML's operation order (and its sequential sums) so results
278        // agree to the last bit.
279        let treat: Array1<f64> = d.mapv(|v| f64::from(v == 1.0));
280        let control: Array1<f64> = d.mapv(|v| f64::from(v == 0.0));
281        let (w1, w0) = if self.normalize_ipw {
282            let sum = |a: Array1<f64>| a.iter().fold(0.0, |acc, &v| acc + v);
283            let w1 = treat.sum() / sum(&treat * &s / &(&pi * &m));
284            let w0 = control.sum() / sum(&control * &s / &(&pi * &(1.0 - &m)));
285            (Some(w1), Some(w0))
286        } else {
287            (None, None)
288        };
289        let psi_b: Array1<f64> = (0..n)
290            .map(|i| {
291                let c1 = treat[i] * s[i] * (y[i] - g1[i]) / (m[i] * pi[i]);
292                let c0 = control[i] * s[i] * (y[i] - g0[i]) / ((1.0 - m[i]) * pi[i]);
293                let b1 = w1.map_or(c1, |w| w * c1) + g1[i];
294                let b0 = w0.map_or(c0, |w| w * c0) + g0[i];
295                b1 - b0
296            })
297            .collect();
298        let psi_a = Array1::from_elem(n, -1.0);
299        Ok(ScoreEstimate::new(
300            LinearScore::new(psi_a, psi_b)?,
301            vec![("ml_g_d0", g0), ("ml_g_d1", g1), ("ml_pi", pi), ("ml_m", m)],
302        )
303        // As DoubleML, the outcome models are evaluated on every observation.
304        .with_targets(vec![
305            NuisanceTarget::regression("ml_g_d0", y.to_owned()),
306            NuisanceTarget::regression("ml_g_d1", y.to_owned()),
307            NuisanceTarget::classification("ml_pi", s.to_owned()),
308            NuisanceTarget::classification("ml_m", d.to_owned()),
309        ]))
310    }
311}
dml_core/model/ssm.rs

dml_core/model/
ssm.rs
