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

dml_core/
data.rs
