dml_core/
score.rs

1//! Linear Neyman-orthogonal scores and the DML2 estimator.
2
3use ndarray::Array1;
4
5use crate::error::{DmlError, Result};
6
7/// Score elements of a model whose orthogonal score is linear in θ:
8///
9/// ```text
10/// ψ(W; θ, η) = ψ_a(W; η) · θ + ψ_b(W; η)
11/// ```
12///
13/// PLR, PLIV and IRM all have this form, so one solver serves every model.
14#[derive(Debug, Clone)]
15pub struct LinearScore {
16    pub psi_a: Array1<f64>,
17    pub psi_b: Array1<f64>,
18}
19
20/// θ̂ and its standard error for one treatment and one repetition.
21#[derive(Debug, Clone)]
22pub(crate) struct ScoreSolution {
23    pub theta: f64,
24    pub se: f64,
25    /// ψ evaluated at θ̂.
26    pub psi: Array1<f64>,
27    /// `J = mean(ψ_a)`, the derivative of the mean score in θ.
28    pub j: f64,
29}
30
31impl LinearScore {
32    pub fn new(psi_a: Array1<f64>, psi_b: Array1<f64>) -> Result<Self> {
33        if psi_a.len() != psi_b.len() || psi_a.is_empty() {
34            return Err(DmlError::DimensionMismatch(
35                "psi_a and psi_b must have the same non-zero length".into(),
36            ));
37        }
38        Ok(Self { psi_a, psi_b })
39    }
40
41    /// DML2: solve `mean(ψ_a)·θ + mean(ψ_b) = 0` over all folds pooled, then
42    /// `σ² = mean(ψ²) / J²` and `se = √(σ² / n)`.
43    pub(crate) fn solve(&self) -> Result<ScoreSolution> {
44        let n = self.psi_a.len() as f64;
45        let j = self.psi_a.sum() / n;
46        if j == 0.0 || !j.is_finite() {
47            return Err(DmlError::InvalidArgument(
48                "mean of psi_a is zero or not finite: theta is not identified".into(),
49            ));
50        }
51        let theta = -(self.psi_b.sum() / n) / j;
52        let psi = &self.psi_a * theta + &self.psi_b;
53        let sigma2 = psi.dot(&psi) / n / (j * j);
54        Ok(ScoreSolution {
55            theta,
56            se: (sigma2 / n).sqrt(),
57            psi,
58            j,
59        })
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66    use ndarray::array;
67
68    #[test]
69    fn solves_mean_score_equation() {
70        let s = LinearScore::new(array![-1.0, -2.0, -3.0], array![2.0, 1.0, 3.0]).unwrap();
71        let sol = s.solve().unwrap();
72        assert!((sol.theta - 1.0).abs() < 1e-15);
73        assert!(sol.psi.sum().abs() < 1e-12);
74        assert_eq!(sol.j, -2.0);
75        // ψ = [1, -1, 0]: σ² = (2/3) / 4, se = √(σ²/3)
76        assert!((sol.se - (2.0_f64 / 3.0 / 4.0 / 3.0).sqrt()).abs() < 1e-15);
77    }
78
79    #[test]
80    fn unidentified_score_errors() {
81        let s = LinearScore::new(array![0.0, 0.0], array![1.0, 2.0]).unwrap();
82        assert!(s.solve().is_err());
83        assert!(LinearScore::new(array![1.0], array![1.0, 2.0]).is_err());
84    }
85}
dml_core/score.rs

dml_core/
score.rs
