dml_core/model/mod.rs
1//! DML models with linear scores: [`Plr`], [`Irm`], [`Pliv`], and the
2//! [`LinearScoreModel`] trait for adding your own.
3
4use ndarray::{Array1, ArrayView1, ArrayView2};
5
6use crate::crossfit::Fold;
7use crate::data::DmlData;
8use crate::error::Result;
9use crate::score::LinearScore;
10use crate::sensitivity::SensitivityParts;
11
12mod apo;
13pub(crate) mod did;
14mod iivm;
15pub(crate) mod irm;
16mod lplr;
17mod nonlinear;
18mod pliv;
19mod plpr;
20mod plr;
21mod quantile;
22mod rdflex;
23mod ssm;
24
25pub use apo::{Apo, Apos, AposFit};
26pub use did::{Did, DidCs, DidScore};
27pub use iivm::Iivm;
28pub use irm::{Irm, IrmScore};
29pub use lplr::{Lplr, LplrScore, TreatmentLearner};
30pub use pliv::Pliv;
31pub use plpr::{PanelApproach, Plpr, StaticPanel};
32pub use plr::{Plr, PlrScore};
33pub use quantile::{Cvar, Lpq, Pq, Qte, QteFit};
34pub use rdflex::{FirstStage, RdFlex, RdFlexFit, RdKernel};
35pub use ssm::{Ssm, SsmScore};
36
37/// The data a model sees when estimating the effect of one treatment.
38///
39/// `#[non_exhaustive]`: fields may be added in future versions, so code outside
40/// this crate reads the fields but cannot construct the struct.
41#[derive(Debug, Clone, Copy)]
42#[non_exhaustive]
43pub struct ScoreInput<'a> {
44 pub y: ArrayView1<'a, f64>,
45 /// The treatment being estimated.
46 pub d: ArrayView1<'a, f64>,
47 /// Covariates, including any other treatments.
48 pub x: ArrayView2<'a, f64>,
49 pub z: Option<ArrayView1<'a, f64>>,
50 /// 0/1 time-period indicator (repeated cross-sections).
51 pub t: Option<ArrayView1<'a, f64>>,
52 /// 0/1 selection indicator (sample selection models).
53 pub s: Option<ArrayView1<'a, f64>>,
54 /// Index of the sample-splitting repetition being estimated.
55 pub rep: usize,
56}
57
58/// Ingredients for best-linear-predictor inference on effect heterogeneity
59/// ([`DmlFit::cate`](crate::DmlFit::cate), [`DmlFit::gate`](crate::DmlFit::gate)):
60/// the BLP regresses `signal − θ·theta_factor` on the basis, with every row of
61/// the basis multiplied by `basis_scale` if given (partially linear models,
62/// where the signal is the outcome residual and the scale the treatment
63/// residual).
64#[derive(Debug, Clone)]
65pub struct CateSignal {
66 pub signal: Array1<f64>,
67 pub theta_factor: Option<Array1<f64>>,
68 pub basis_scale: Option<Array1<f64>>,
69 /// The signal is a (doubly robust) signal of each unit's treatment effect,
70 /// so a policy tree may be fitted on it ([`DmlFit::policy_tree`](crate::DmlFit::policy_tree)).
71 pub is_effect: bool,
72}
73
74/// Score elements plus the cross-fitted nuisance predictions behind them.
75///
76/// Build it with [`ScoreEstimate::new`] (the struct is `#[non_exhaustive]`, so
77/// fields can be added without breaking custom models).
78#[derive(Debug, Clone)]
79#[non_exhaustive]
80pub struct ScoreEstimate {
81 pub score: LinearScore,
82 /// `(name, out-of-fold predictions)` — e.g. `("ml_l", ℓ̂)`.
83 pub predictions: Vec<(&'static str, Array1<f64>)>,
84 /// Ingredients for [sensitivity analysis](crate::sensitivity), if the
85 /// model supports it.
86 pub sensitivity: Option<SensitivityParts>,
87 /// What each nuisance learner predicts, for out-of-sample diagnostics.
88 pub targets: Vec<NuisanceTarget>,
89 /// Signal for CATE / GATE inference, if the model supports it.
90 pub cate: Option<CateSignal>,
91}
92
93/// Loss used to evaluate a nuisance learner out of sample.
94#[derive(Debug, Clone, Copy, PartialEq, Eq)]
95pub enum LossKind {
96 /// Root mean squared error (regression learners).
97 Rmse,
98 /// Log loss with predictions clipped to `[1e-15, 1 − 1e-15]` (classifiers).
99 LogLoss,
100}
101
102/// The target a nuisance learner predicts; `NaN` marks rows where the
103/// nuisance does not apply (e.g. `g₀` on treated rows).
104#[derive(Debug, Clone)]
105pub struct NuisanceTarget {
106 pub name: &'static str,
107 pub target: Array1<f64>,
108 pub kind: LossKind,
109}
110
111impl NuisanceTarget {
112 pub fn regression(name: &'static str, target: Array1<f64>) -> Self {
113 Self {
114 name,
115 target,
116 kind: LossKind::Rmse,
117 }
118 }
119
120 pub fn classification(name: &'static str, target: Array1<f64>) -> Self {
121 Self {
122 name,
123 target,
124 kind: LossKind::LogLoss,
125 }
126 }
127
128 /// Keep the target only where `keep` is true (`NaN` elsewhere).
129 pub fn only_where(mut self, keep: impl Fn(usize) -> bool) -> Self {
130 for (i, v) in self.target.iter_mut().enumerate() {
131 if !keep(i) {
132 *v = f64::NAN;
133 }
134 }
135 self
136 }
137
138 /// Out-of-sample loss of `prediction` over the rows with a target.
139 pub fn loss(&self, prediction: &Array1<f64>) -> f64 {
140 let mut total = 0.0;
141 let mut count = 0usize;
142 for (&t, &p) in self.target.iter().zip(prediction) {
143 if t.is_nan() {
144 continue;
145 }
146 total += match self.kind {
147 LossKind::Rmse => (t - p).powi(2),
148 LossKind::LogLoss => {
149 let p = p.clamp(1e-15, 1.0 - 1e-15);
150 -(t * p.ln() + (1.0 - t) * (1.0 - p).ln())
151 }
152 };
153 count += 1;
154 }
155 if count == 0 {
156 return f64::NAN;
157 }
158 let mean = total / count as f64;
159 match self.kind {
160 LossKind::Rmse => mean.sqrt(),
161 LossKind::LogLoss => mean,
162 }
163 }
164}
165
166impl ScoreEstimate {
167 pub fn new(score: LinearScore, predictions: Vec<(&'static str, Array1<f64>)>) -> Self {
168 Self {
169 score,
170 predictions,
171 sensitivity: None,
172 targets: Vec::new(),
173 cate: None,
174 }
175 }
176
177 /// Attach the signal for CATE / GATE inference.
178 pub fn with_cate_signal(mut self, cate: CateSignal) -> Self {
179 self.cate = Some(cate);
180 self
181 }
182
183 /// Attach the nuisance targets used for learner diagnostics.
184 pub fn with_targets(mut self, targets: Vec<NuisanceTarget>) -> Self {
185 self.targets = targets;
186 self
187 }
188
189 /// Attach the model's sensitivity ingredients.
190 pub fn with_sensitivity(mut self, parts: SensitivityParts) -> Self {
191 self.sensitivity = Some(parts);
192 self
193 }
194}
195
196/// A DML model whose orthogonal score is linear in θ.
197///
198/// Implement this to plug a new score into the shared machinery (repeated
199/// cross-fitting, aggregation, inference, bootstrap) via [`fit_model`](crate::fit_model).
200pub trait LinearScoreModel {
201 /// Label for summaries, e.g. `"PLR (partialling out)"`.
202 fn name(&self) -> String;
203
204 /// Reject data the model cannot handle (called once before fitting).
205 fn validate(&self, _data: &DmlData) -> Result<()> {
206 Ok(())
207 }
208
209 /// Cross-fit the nuisances on `folds` and return the score elements.
210 fn estimate_score(&self, input: &ScoreInput<'_>, folds: &[Fold]) -> Result<ScoreEstimate>;
211
212 /// Strata for the default sample splits (e.g. treatment levels), so each
213 /// fold holds about the same share of every stratum. `None` (the default)
214 /// draws plain K-fold splits.
215 fn strata(&self, _data: &DmlData) -> Option<Vec<usize>> {
216 None
217 }
218}
219
220/// Builder methods shared by every model: sample-splitting settings and `fit`.
221///
222/// (`macro_rules!` generates code at compile time; here it pastes the same
223/// methods into each model's `impl` block instead of repeating them.)
224macro_rules! dml_model_api {
225 () => {
226 /// Number of cross-fitting folds (default 5).
227 pub fn n_folds(mut self, n_folds: usize) -> Self {
228 self.splitting.n_folds = n_folds;
229 self
230 }
231
232 /// Number of independent sample splits (default 1). Estimates are
233 /// aggregated across repetitions by the median, as in DoubleML.
234 pub fn n_rep(mut self, n_rep: usize) -> Self {
235 self.splitting.n_rep = n_rep;
236 self
237 }
238
239 /// Seed for the sample splits (default 0).
240 pub fn seed(mut self, seed: u64) -> Self {
241 self.splitting.seed = seed;
242 self
243 }
244
245 /// Draw seeded sample splits and fit. With clustered data
246 /// ([`DmlData::with_clusters`](crate::DmlData::with_clusters)) whole clusters are assigned to folds
247 /// and standard errors are cluster-robust.
248 pub fn fit(&self, data: &$crate::DmlData) -> $crate::Result<$crate::DmlFit> {
249 match data.clusters() {
250 Some(clusters) => {
251 let splits = self.splitting.draw_clustered(clusters)?;
252 $crate::fit_model_clustered(self, data, &splits)
253 }
254 None => {
255 let splits = match $crate::LinearScoreModel::strata(self, data) {
256 Some(strata) => self.splitting.draw_stratified(&strata)?,
257 None => self.splitting.draw(data.n_obs())?,
258 };
259 $crate::fit_model(self, data, &splits)
260 }
261 }
262 }
263
264 /// Fit clustered data with explicit cluster folds, one
265 /// `Vec<ClusterFold>` per repetition.
266 pub fn fit_with_cluster_folds(
267 &self,
268 data: &$crate::DmlData,
269 splits: &[Vec<$crate::ClusterFold>],
270 ) -> $crate::Result<$crate::DmlFit> {
271 $crate::fit_model_clustered(self, data, splits)
272 }
273
274 /// Sensitivity benchmark for the covariate columns `benchmarking_set`:
275 /// refits this model without them (with the same sample splitting, so
276 /// `long` must come from [`fit`](Self::fit) with the current settings)
277 /// and compares it with `long`.
278 pub fn sensitivity_benchmark(
279 &self,
280 data: &$crate::DmlData,
281 long: &$crate::DmlFit,
282 benchmarking_set: &[usize],
283 ) -> $crate::Result<$crate::SensitivityBenchmark> {
284 let short = self.fit(&data.without_covariates(benchmarking_set)?)?;
285 long.sensitivity_benchmark(&short)
286 }
287
288 /// Fit with caller-supplied splits, one `Vec<Fold>` per repetition
289 /// (e.g. to reproduce another implementation exactly).
290 pub fn fit_with_folds(
291 &self,
292 data: &$crate::DmlData,
293 splits: &[Vec<$crate::Fold>],
294 ) -> $crate::Result<$crate::DmlFit> {
295 $crate::fit_model(self, data, splits)
296 }
297 };
298}
299pub(crate) use dml_model_api;