dml_core/
bootstrap.rs

1//! Multiplier bootstrap for joint confidence intervals and Romano–Wolf
2//! multiple-testing adjustment.
3
4use ndarray::{Array1, Array2, Array3, ArrayView3, Axis};
5use rand::{Rng, SeedableRng};
6use rand_chacha::ChaCha8Rng;
7use rand_distr::{Exp1, StandardNormal};
8
9use crate::error::{DmlError, Result};
10use crate::fit::{check_level, confint_from, DmlFit};
11use crate::stats::{median, quantile_mut};
12
13/// Distribution of the multiplier weights `ξ_i` (all mean 0, variance 1).
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
15pub enum BootstrapMethod {
16    /// `ξ ~ N(0, 1)`
17    #[default]
18    Normal,
19    /// `ξ ~ Exp(1) − 1`
20    Bayes,
21    /// Mammen-type: `ξ = a/√2 + (b² − 1)/2`, with `a, b ~ N(0, 1)`
22    Wild,
23}
24
25/// Draw an `(n_boot, n)` matrix of multiplier weights.
26pub fn draw_weights(
27    method: BootstrapMethod,
28    n_boot: usize,
29    n: usize,
30    rng: &mut impl Rng,
31) -> Array2<f64> {
32    let shape = (n_boot, n);
33    match method {
34        BootstrapMethod::Normal => {
35            Array2::from_shape_simple_fn(shape, || rng.sample(StandardNormal))
36        }
37        BootstrapMethod::Bayes => {
38            Array2::from_shape_simple_fn(shape, || rng.sample::<f64, _>(Exp1) - 1.0)
39        }
40        BootstrapMethod::Wild => Array2::from_shape_simple_fn(shape, || {
41            let a: f64 = rng.sample(StandardNormal);
42            let b: f64 = rng.sample(StandardNormal);
43            a / std::f64::consts::SQRT_2 + (b * b - 1.0) / 2.0
44        }),
45    }
46}
47
48/// Bootstrap distribution of the t-statistics of a [`DmlFit`].
49///
50/// Keeps its own copy of the per-repetition estimates (small `(k, n_rep)`
51/// arrays), so it is independent of the fit it came from.
52#[derive(Debug, Clone)]
53pub struct Bootstrap {
54    all_coef: Array2<f64>,
55    all_se: Array2<f64>,
56    /// Shape `(n_boot, k, n_rep)`.
57    t_stats: Array3<f64>,
58}
59
60impl DmlFit {
61    /// Multiplier bootstrap with `n_boot` draws per repetition:
62    /// `t*_b = Σ_i ξ_bi · ψ_i / (J · n · se)`.
63    pub fn bootstrap(
64        &self,
65        method: BootstrapMethod,
66        n_boot: usize,
67        seed: u64,
68    ) -> Result<Bootstrap> {
69        if n_boot == 0 {
70            return Err(DmlError::InvalidArgument("n_boot must be positive".into()));
71        }
72        let mut rng = ChaCha8Rng::seed_from_u64(seed);
73        let weights: Vec<Array2<f64>> = (0..self.n_rep())
74            .map(|_| draw_weights(method, n_boot, self.n_obs, &mut rng))
75            .collect();
76        self.bootstrap_with_weights(&weights)
77    }
78
79    /// Bootstrap with caller-supplied weights, one `(n_boot, n)` matrix per
80    /// repetition (e.g. to reproduce DoubleML's draws exactly).
81    pub fn bootstrap_with_weights(&self, weights: &[Array2<f64>]) -> Result<Bootstrap> {
82        if self.cluster.is_some() {
83            return Err(DmlError::InvalidArgument(
84                "the multiplier bootstrap is not available for clustered data (as in DoubleML)"
85                    .into(),
86            ));
87        }
88        if weights.len() != self.n_rep() {
89            return Err(DmlError::DimensionMismatch(format!(
90                "{} weight matrices for {} repetitions",
91                weights.len(),
92                self.n_rep()
93            )));
94        }
95        let n_boot = weights[0].nrows();
96        let (n, k) = (self.n_obs, self.n_treat());
97        let mut t_stats = Array3::<f64>::zeros((n_boot, k, self.n_rep()));
98        for (r, w) in weights.iter().enumerate() {
99            if w.dim() != (n_boot, n) {
100                return Err(DmlError::DimensionMismatch(
101                    "every weight matrix must be (n_boot, n_obs)".into(),
102                ));
103            }
104            // scaled_psi[i, j] = ψ_ij / (J_j · n · se_j)
105            let mut scaled = self.psi.index_axis(Axis(2), r).to_owned();
106            for j in 0..k {
107                let denom = self.psi_deriv[[j, r]] * n as f64 * self.all_se[[j, r]];
108                scaled.column_mut(j).mapv_inplace(|v| v / denom);
109            }
110            t_stats.index_axis_mut(Axis(2), r).assign(&w.dot(&scaled));
111        }
112        Ok(Bootstrap {
113            all_coef: self.all_coef.clone(),
114            all_se: self.all_se.clone(),
115            t_stats,
116        })
117    }
118}
119
120impl Bootstrap {
121    /// Bootstrapped t-statistics, shape `(n_boot, k, n_rep)`.
122    pub fn t_stats(&self) -> ArrayView3<'_, f64> {
123        self.t_stats.view()
124    }
125
126    /// Per-repetition critical values: the `level` quantile of `max_j |t*_j|`.
127    pub fn critical_values(&self, level: f64) -> Result<Array1<f64>> {
128        check_level(level)?;
129        Ok(self
130            .t_stats
131            .axis_iter(Axis(2))
132            .map(|rep| {
133                let mut max_abs: Vec<f64> = rep
134                    .rows()
135                    .into_iter()
136                    .map(|row| row.iter().fold(0.0_f64, |m, v| m.max(v.abs())))
137                    .collect();
138                quantile_mut(&mut max_abs, level)
139            })
140            .collect())
141    }
142
143    /// Simultaneous confidence intervals covering all treatments jointly at
144    /// `level`, shape `(k, 2)`.
145    pub fn joint_confint(&self, level: f64) -> Result<Array2<f64>> {
146        Ok(confint_from(
147            &self.all_coef,
148            &self.all_se,
149            &self.critical_values(level)?,
150        ))
151    }
152
153    /// Romano–Wolf step-down adjusted p-values (median over repetitions).
154    pub fn romano_wolf(&self) -> Array1<f64> {
155        let (k, n_rep) = self.all_coef.dim();
156        let mut all = Array2::<f64>::zeros((k, n_rep));
157        for r in 0..n_rep {
158            let boot = self.t_stats.index_axis(Axis(2), r);
159            let abs_t: Vec<f64> = (0..k)
160                .map(|j| (self.all_coef[[j, r]] / self.all_se[[j, r]]).abs())
161                .collect();
162            // Treatments in decreasing order of |t|.
163            let mut order: Vec<usize> = (0..k).collect();
164            order.sort_by(|&a, &b| abs_t[b].total_cmp(&abs_t[a]));
165            let mut running = 0.0_f64;
166            for (step, &j) in order.iter().enumerate() {
167                let remaining = &order[step..];
168                let exceed = boot
169                    .rows()
170                    .into_iter()
171                    .filter(|row| {
172                        remaining.iter().fold(0.0_f64, |m, &i| m.max(row[i].abs())) >= abs_t[j]
173                    })
174                    .count();
175                let p = (exceed as f64 / boot.nrows() as f64).min(1.0);
176                running = running.max(p);
177                all[[j, r]] = running;
178            }
179        }
180        all.rows()
181            .into_iter()
182            .map(|row| median(&row.to_vec()))
183            .collect()
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    #[test]
192    fn weights_have_mean_zero_unit_variance() {
193        let mut rng = ChaCha8Rng::seed_from_u64(1);
194        for method in [
195            BootstrapMethod::Normal,
196            BootstrapMethod::Bayes,
197            BootstrapMethod::Wild,
198        ] {
199            let w = draw_weights(method, 200, 500, &mut rng);
200            let mean = w.mean().unwrap();
201            let var = w.mapv(|v| (v - mean).powi(2)).mean().unwrap();
202            assert!(mean.abs() < 0.02, "{method:?} mean {mean}");
203            assert!((var - 1.0).abs() < 0.03, "{method:?} var {var}");
204        }
205    }
206}
dml_core/bootstrap.rs

dml_core/
bootstrap.rs
