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

dml_core/
data.rs
