dml_core/
data.rs

1use std::borrow::Cow;
2
3use ndarray::{concatenate, Array1, Array2, ArrayView1, ArrayView2, Axis};
4
5use crate::error::{DmlError, Result};
6
7/// Data for a DML model.
8///
9/// - `y`: outcome, length `n`
10/// - `d`: one or more treatments, shape `(n, k)` (continuous or 0/1)
11/// - `x`: covariates / confounders, shape `(n, p)`
12/// - `z`: optional instrument(s), length `n` (for [`Pliv`](crate::Pliv), which also takes several,
13///   and, binary, [`Iivm`](crate::Iivm))
14/// - `t`: optional 0/1 time-period indicator (for repeated cross-sections, [`DidCs`](crate::DidCs))
15/// - `s`: optional 0/1 selection indicator, 1 where the outcome is observed (for [`Ssm`](crate::Ssm))
16///
17/// With several treatments each one is estimated in turn, and the other
18/// treatments are added to the covariates (as in DoubleML).
19#[derive(Debug, Clone)]
20pub struct DmlData {
21    y: Array1<f64>,
22    d: Array2<f64>,
23    x: Array2<f64>,
24    z: Option<Array2<f64>>,
25    t: Option<Array1<f64>>,
26    s: Option<Array1<f64>>,
27    external: Option<crate::model::ExternalPredictions>,
28    clusters: Option<Vec<Array1<i64>>>,
29    treatment_names: Vec<String>,
30}
31
32fn all_finite<'a>(mut values: impl Iterator<Item = &'a f64>) -> bool {
33    values.all(|v| v.is_finite())
34}
35
36impl DmlData {
37    /// Single treatment.
38    pub fn new(y: Array1<f64>, d: Array1<f64>, x: Array2<f64>) -> Result<Self> {
39        let n = d.len();
40        let d = d.into_shape_with_order((n, 1)).expect("1-D to (n, 1)");
41        Self::with_treatments(y, d, x)
42    }
43
44    /// Several treatments, one per column of `d`.
45    pub fn with_treatments(y: Array1<f64>, d: Array2<f64>, x: Array2<f64>) -> Result<Self> {
46        let n = y.len();
47        if n < 2 {
48            return Err(DmlError::InvalidArgument(
49                "need at least 2 observations".into(),
50            ));
51        }
52        if d.nrows() != n || x.nrows() != n {
53            return Err(DmlError::DimensionMismatch(format!(
54                "y has {n} rows, d has {}, x has {}",
55                d.nrows(),
56                x.nrows()
57            )));
58        }
59        if d.ncols() == 0 {
60            return Err(DmlError::InvalidArgument(
61                "need at least one treatment".into(),
62            ));
63        }
64        if !all_finite(y.iter()) || !all_finite(d.iter()) || !all_finite(x.iter()) {
65            return Err(DmlError::InvalidArgument(
66                "data contains NaN or infinite values".into(),
67            ));
68        }
69        let treatment_names = if d.ncols() == 1 {
70            vec!["d".to_string()]
71        } else {
72            (0..d.ncols()).map(|j| format!("d{j}")).collect()
73        };
74        Ok(Self {
75            y,
76            d,
77            x,
78            z: None,
79            t: None,
80            s: None,
81            external: None,
82            clusters: None,
83            treatment_names,
84        })
85    }
86
87    /// Attach an instrument `z` (required by [`Pliv`](crate::Pliv)).
88    pub fn with_instrument(self, z: Array1<f64>) -> Result<Self> {
89        let n = z.len();
90        self.with_instruments(z.into_shape_with_order((n, 1)).expect("1-D to (n, 1)"))
91    }
92
93    /// Attach several instruments, one per column ([`Pliv`](crate::Pliv)
94    /// combines them; the other instrumental models need exactly one).
95    pub fn with_instruments(mut self, z: Array2<f64>) -> Result<Self> {
96        if z.nrows() != self.n_obs() {
97            return Err(DmlError::DimensionMismatch(format!(
98                "z has {} rows, expected {}",
99                z.nrows(),
100                self.n_obs()
101            )));
102        }
103        if z.ncols() == 0 {
104            return Err(DmlError::InvalidArgument(
105                "need at least one instrument".into(),
106            ));
107        }
108        if !all_finite(z.iter()) {
109            return Err(DmlError::InvalidArgument(
110                "instrument contains NaN or infinite values".into(),
111            ));
112        }
113        self.z = Some(z);
114        Ok(self)
115    }
116
117    /// Attach a 0/1 time-period indicator (post-period = 1), required by
118    /// [`DidCs`](crate::DidCs) for repeated cross-sections.
119    pub fn with_time_indicator(mut self, t: Array1<f64>) -> Result<Self> {
120        if t.len() != self.n_obs() {
121            return Err(DmlError::DimensionMismatch(format!(
122                "t has {} rows, expected {}",
123                t.len(),
124                self.n_obs()
125            )));
126        }
127        if !is_binary(t.view()) {
128            return Err(DmlError::InvalidArgument(
129                "time indicator must contain only 0 and 1".into(),
130            ));
131        }
132        self.t = Some(t);
133        Ok(self)
134    }
135
136    /// Attach a 0/1 selection indicator (1 where the outcome is observed),
137    /// required by the sample selection model [`Ssm`](crate::Ssm). The
138    /// outcome of unselected observations is not used but must be finite
139    /// (e.g. 0).
140    pub fn with_selection_indicator(mut self, s: Array1<f64>) -> Result<Self> {
141        if s.len() != self.n_obs() {
142            return Err(DmlError::DimensionMismatch(format!(
143                "s has {} rows, expected {}",
144                s.len(),
145                self.n_obs()
146            )));
147        }
148        if !is_binary(s.view()) {
149            return Err(DmlError::InvalidArgument(
150                "selection indicator must contain only 0 and 1".into(),
151            ));
152        }
153        self.s = Some(s);
154        Ok(self)
155    }
156
157    /// Cluster the observations (e.g. by firm or school): sample splitting
158    /// then keeps clusters together and standard errors are cluster-robust.
159    pub fn with_clusters(self, cluster: Array1<i64>) -> Result<Self> {
160        self.set_clusters(vec![cluster])
161    }
162
163    /// Two-way clustering (e.g. by firm and by year).
164    pub fn with_two_way_clusters(self, first: Array1<i64>, second: Array1<i64>) -> Result<Self> {
165        self.set_clusters(vec![first, second])
166    }
167
168    fn set_clusters(mut self, clusters: Vec<Array1<i64>>) -> Result<Self> {
169        if clusters.iter().any(|c| c.len() != self.n_obs()) {
170            return Err(DmlError::DimensionMismatch(
171                "cluster variables need one entry per observation".into(),
172            ));
173        }
174        self.clusters = Some(clusters);
175        Ok(self)
176    }
177
178    /// The cluster variables (one or two), if any.
179    pub fn clusters(&self) -> Option<&[Array1<i64>]> {
180        self.clusters.as_deref()
181    }
182
183    /// The same data without the covariate columns `drop` (for sensitivity
184    /// benchmarks).
185    pub fn without_covariates(&self, drop: &[usize]) -> Result<Self> {
186        let p = self.n_features();
187        if drop.is_empty() || drop.iter().any(|&c| c >= p) {
188            return Err(DmlError::InvalidArgument(format!(
189                "benchmarking columns must be a non-empty subset of the {p} covariates"
190            )));
191        }
192        let keep: Vec<usize> = (0..p).filter(|c| !drop.contains(c)).collect();
193        // External predictions belong to the full covariate set.
194        Ok(Self {
195            x: self.x.select(Axis(1), &keep),
196            external: None,
197            ..self.clone()
198        })
199    }
200
201    /// Names shown in summaries (default `d`, or `d0`, `d1`, …).
202    pub fn with_treatment_names<S: Into<String>>(
203        mut self,
204        names: impl IntoIterator<Item = S>,
205    ) -> Result<Self> {
206        let names: Vec<String> = names.into_iter().map(Into::into).collect();
207        if names.len() != self.n_treat() {
208            return Err(DmlError::DimensionMismatch(format!(
209                "{} names for {} treatments",
210                names.len(),
211                self.n_treat()
212            )));
213        }
214        self.treatment_names = names;
215        Ok(self)
216    }
217
218    pub fn n_obs(&self) -> usize {
219        self.y.len()
220    }
221
222    pub fn n_features(&self) -> usize {
223        self.x.ncols()
224    }
225
226    pub fn n_treat(&self) -> usize {
227        self.d.ncols()
228    }
229
230    pub fn y(&self) -> ArrayView1<'_, f64> {
231        self.y.view()
232    }
233
234    /// All treatments, shape `(n, k)`.
235    pub fn d(&self) -> ArrayView2<'_, f64> {
236        self.d.view()
237    }
238
239    /// Treatment `j`.
240    pub fn treatment(&self, j: usize) -> ArrayView1<'_, f64> {
241        self.d.column(j)
242    }
243
244    pub fn x(&self) -> ArrayView2<'_, f64> {
245        self.x.view()
246    }
247
248    /// The instrument, if exactly one is attached.
249    pub fn z(&self) -> Option<ArrayView1<'_, f64>> {
250        self.z
251            .as_ref()
252            .filter(|z| z.ncols() == 1)
253            .map(|z| z.column(0))
254    }
255
256    /// All instruments, one per column.
257    pub fn instruments(&self) -> Option<ArrayView2<'_, f64>> {
258        self.z.as_ref().map(|z| z.view())
259    }
260
261    pub fn n_instruments(&self) -> usize {
262        self.z.as_ref().map_or(0, |z| z.ncols())
263    }
264
265    /// The 0/1 time-period indicator, if attached.
266    pub fn time(&self) -> Option<ArrayView1<'_, f64>> {
267        self.t.as_ref().map(|t| t.view())
268    }
269
270    /// Use user-supplied out-of-fold predictions for some nuisance learners
271    /// instead of fitting them (DoubleML's `external_predictions`).
272    pub fn with_external_predictions(
273        mut self,
274        external: crate::model::ExternalPredictions,
275    ) -> Self {
276        self.external = Some(external);
277        self
278    }
279
280    pub fn external_predictions(&self) -> Option<&crate::model::ExternalPredictions> {
281        self.external.as_ref()
282    }
283
284    /// The 0/1 selection indicator, if attached.
285    pub fn selection(&self) -> Option<ArrayView1<'_, f64>> {
286        self.s.as_ref().map(|s| s.view())
287    }
288
289    pub fn treatment_names(&self) -> &[String] {
290        &self.treatment_names
291    }
292
293    /// Covariates used when estimating the effect of treatment `j`: `x` plus
294    /// the other treatment columns. Borrowed when there is only one treatment.
295    ///
296    /// (`Cow` = "clone on write": either a borrowed view or an owned matrix.)
297    pub fn covariates_for(&self, j: usize) -> Cow<'_, Array2<f64>> {
298        if self.n_treat() == 1 {
299            return Cow::Borrowed(&self.x);
300        }
301        let others: Vec<usize> = (0..self.n_treat()).filter(|&k| k != j).collect();
302        let d_other = self.d.select(Axis(1), &others);
303        Cow::Owned(concatenate![Axis(1), self.x, d_other])
304    }
305
306    /// True when every value of treatment `j` is exactly 0 or 1.
307    pub fn is_binary_treatment(&self, j: usize) -> bool {
308        is_binary(self.d.column(j))
309    }
310}
311
312/// True when every value is exactly 0 or 1.
313pub(crate) fn is_binary(v: ArrayView1<'_, f64>) -> bool {
314    v.iter().all(|&x| x == 0.0 || x == 1.0)
315}
316
317/// Gather `v[idx]` into a new array.
318pub(crate) fn gather(v: ArrayView1<'_, f64>, idx: &[usize]) -> Array1<f64> {
319    idx.iter().map(|&i| v[i]).collect()
320}
321
322#[cfg(test)]
323mod tests {
324    use super::*;
325    use ndarray::array;
326
327    #[test]
328    fn rejects_mismatched_lengths() {
329        let err = DmlData::new(array![1.0, 2.0], array![0.0], array![[1.0], [2.0]]).unwrap_err();
330        assert!(matches!(err, DmlError::DimensionMismatch(_)));
331    }
332
333    #[test]
334    fn rejects_nan() {
335        let err = DmlData::new(
336            array![1.0, f64::NAN],
337            array![0.0, 1.0],
338            array![[1.0], [2.0]],
339        )
340        .unwrap_err();
341        assert!(matches!(err, DmlError::InvalidArgument(_)));
342    }
343
344    #[test]
345    fn detects_binary_treatment() {
346        let data = DmlData::new(array![1.0, 2.0], array![0.0, 1.0], array![[1.0], [2.0]]).unwrap();
347        assert!(data.is_binary_treatment(0));
348    }
349
350    #[test]
351    fn other_treatments_become_covariates() {
352        let d = array![[1.0, 10.0, 100.0], [2.0, 20.0, 200.0]];
353        let data = DmlData::with_treatments(array![0.0, 1.0], d, array![[5.0], [6.0]]).unwrap();
354        assert_eq!(data.treatment_names(), ["d0", "d1", "d2"]);
355        assert_eq!(
356            *data.covariates_for(1),
357            array![[5.0, 1.0, 100.0], [6.0, 2.0, 200.0]]
358        );
359        assert!(matches!(data.covariates_for(0), Cow::Owned(_)));
360    }
361
362    #[test]
363    fn several_instruments() {
364        let data = DmlData::new(array![1.0, 2.0], array![0.0, 1.0], array![[1.0], [2.0]]).unwrap();
365        let two = data
366            .clone()
367            .with_instruments(array![[1.0, 0.5], [0.0, 2.0]])
368            .unwrap();
369        assert_eq!(two.n_instruments(), 2);
370        assert!(two.z().is_none());
371        assert_eq!(two.instruments().unwrap().dim(), (2, 2));
372        let one = data.with_instrument(array![1.0, 0.0]).unwrap();
373        assert_eq!(one.z().unwrap().to_vec(), vec![1.0, 0.0]);
374    }
375
376    #[test]
377    fn time_indicator_must_be_binary() {
378        let data = DmlData::new(array![1.0, 2.0], array![0.0, 1.0], array![[1.0], [2.0]]).unwrap();
379        assert!(data.clone().with_time_indicator(array![0.0, 2.0]).is_err());
380        assert!(data.clone().with_time_indicator(array![1.0]).is_err());
381        assert!(data
382            .with_time_indicator(array![0.0, 1.0])
383            .unwrap()
384            .time()
385            .is_some());
386    }
387
388    #[test]
389    fn instrument_length_checked() {
390        let data = DmlData::new(array![1.0, 2.0], array![0.0, 1.0], array![[1.0], [2.0]]).unwrap();
391        assert!(data.clone().with_instrument(array![1.0]).is_err());
392        assert!(data
393            .with_instrument(array![1.0, 0.0])
394            .unwrap()
395            .z()
396            .is_some());
397    }
398}
dml_core/data.rs

dml_core/
data.rs
