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

dml_core/
data.rs
