dml-core 0.14.0

Double / debiased machine learning for causal inference in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

//! Shared machinery for models whose score is non-linear in θ (quantile
//! models, the logistic partially linear model): one repetition's estimate
//! and score, assembled into a [`DmlFit`] with DoubleML's variance estimate.

use ndarray::{Array1, Array2, Array3};

use super::{Name, NuisanceTarget};
use crate::compat::np_mean;
use crate::data::DmlData;
use crate::fit::{aggregate, DmlFit, NuisanceLoss};

/// Result of one repetition of a model with a non-linear score.
pub(super) struct RepFit {
    /// Mean preliminary estimate, the next repetition's start value.
    pub next_start: f64,
    pub theta: f64,
    pub psi: Array1<f64>,
    /// Mean score derivative `J`.
    pub jacobian: f64,
    pub predictions: Vec<(Name, Array1<f64>)>,
    pub targets: Vec<NuisanceTarget>,
}

/// Assemble a [`DmlFit`] from the repetitions of a non-linear model. The score
/// is stored linearised at the estimate (`ψ_a = J`, `ψ_b = ψ − J·θ̂`), so the
/// bootstrap and linear combinations work as for linear scores.
pub(super) fn assemble(name: String, data: &DmlData, n_folds: usize, reps: Vec<RepFit>) -> DmlFit {
    let n = data.n_obs();
    let n_rep = reps.len();
    let mut all_coef = Array2::zeros((1, n_rep));
    let mut all_se = Array2::zeros((1, n_rep));
    let mut psi = Array3::zeros((n, 1, n_rep));
    let mut psi_a = Array3::zeros((n, 1, n_rep));
    let mut psi_b = Array3::zeros((n, 1, n_rep));
    let mut psi_deriv = Array2::zeros((1, n_rep));
    let mut predictions: Vec<(String, Array3<f64>)> = Vec::new();
    let mut losses: Vec<NuisanceLoss> = Vec::new();
    for (r, rep) in reps.into_iter().enumerate() {
        // DoubleML: σ² = γ / (n·J²) with γ the mean squared score.
        let jac = rep.jacobian;
        let squares: Vec<f64> = rep.psi.iter().map(|v| v * v).collect();
        let gamma = np_mean(&squares);
        let sigma2 = 1.0 / (n as f64 * (jac * jac)) * gamma;
        all_coef[[0, r]] = rep.theta;
        all_se[[0, r]] = sigma2.sqrt();
        psi_deriv[[0, r]] = jac;
        psi.slice_mut(ndarray::s![.., 0, r]).assign(&rep.psi);
        psi_a.slice_mut(ndarray::s![.., 0, r]).fill(jac);
        psi_b
            .slice_mut(ndarray::s![.., 0, r])
            .assign(&(&rep.psi - jac * rep.theta));
        for target in &rep.targets {
            let Some((_, pred)) = rep.predictions.iter().find(|(nm, _)| *nm == target.name) else {
                continue;
            };
            let slot = match losses.iter().position(|l| l.name == target.name) {
                Some(i) => i,
                None => {
                    losses.push(NuisanceLoss {
                        name: target.name.to_string(),
                        kind: target.kind,
                        loss: Array2::from_elem((1, n_rep), f64::NAN),
                        targets: Array3::from_elem((n, n_rep, 1), f64::NAN),
                    });
                    losses.len() - 1
                }
            };
            losses[slot].loss[[0, r]] = target.loss(pred);
            losses[slot]
                .targets
                .slice_mut(ndarray::s![.., r, 0])
                .assign(&target.target);
        }
        for (name, pred) in rep.predictions {
            let slot = match predictions.iter().position(|(nm, _)| *nm == name) {
                Some(i) => i,
                None => {
                    predictions
                        .push((name.to_string(), Array3::from_elem((n, n_rep, 1), f64::NAN)));
                    predictions.len() - 1
                }
            };
            predictions[slot]
                .1
                .slice_mut(ndarray::s![.., r, 0])
                .assign(&pred);
        }
    }
    let (coef, se) = aggregate(&all_coef, &all_se);
    DmlFit {
        model: name,
        treatment_names: data.treatment_names().to_vec(),
        n_obs: n,
        n_folds,
        coef,
        se,
        all_coef,
        all_se,
        psi,
        psi_a,
        psi_b,
        psi_deriv,
        predictions,
        sensitivity: None,
        cate: None,
        losses,
        cluster: None,
    }
}