///|
/// Data container for `DoubleMLPLIV`. In addition to the covariates
/// `x`, the outcome `y` and the treatment `d`, PLIV needs an
/// instrumental variable `z`. The port supports a *single*
/// instrument (1-D array) — the multi-instrument case is not
/// implemented.
pub struct DoubleMLPLIVData {
x : Matrix
y : Array[Double]
d : Array[Double]
z : Array[Double]
cluster_vars : Array[Int]
} derive(Debug)
///|
pub extend DoubleMLPLIVData with @moonbitlang/core/debug.Debug::{to_repr}
///|
/// Build a `DoubleMLPLIVData` from an `n x p` feature matrix, an
/// outcome vector of length `n`, a treatment vector of length `n`
/// and an instrument vector of length `n`.
pub fn DoubleMLPLIVData::new(
x : Matrix,
y : Array[Double],
d : Array[Double],
z : Array[Double],
cluster_vars? : Array[Int] = [],
) -> DoubleMLPLIVData {
try {
require(x.nrows == y.length())
require(x.nrows == d.length())
require(x.nrows == z.length())
if cluster_vars.length() > 0 {
require(cluster_vars.length() == x.nrows)
}
{ x, y, d, z, cluster_vars, }
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// Number of observations.
pub fn DoubleMLPLIVData::n_obs(self : DoubleMLPLIVData) -> Int {
self.x.rows()
}
///|
/// Number of features.
pub fn DoubleMLPLIVData::n_features(self : DoubleMLPLIVData) -> Int {
self.x.cols()
}
///|
/// True iff the data is set up for clustered inference (a
/// non-empty cluster_vars vector was passed to DoubleMLPLIVData::new).
pub fn DoubleMLPLIVData::is_cluster_data(self : DoubleMLPLIVData) -> Bool {
self.cluster_vars.length() > 0
}
///|
/// Length of the cluster_vars vector (0 when not clustered).
pub fn DoubleMLPLIVData::n_cluster_vars(self : DoubleMLPLIVData) -> Int {
self.cluster_vars.length()
}
///|
/// Double / debiased machine learning estimator for the *partially
/// linear IV regression model* (PLIV) of Chernozhukov et al. (2018) with
/// the *partialling out* score:
///
/// Y = D * theta_0 + g_0(X) + zeta, E[zeta | D, X] = 0
/// D = m_0(X) + V, E[V | X] = 0
/// Z = ell_0(X) + xi, E[xi | X] = 0,
/// Cov(Z, V) != 0 (relevance)
///
/// with the *partialling out* (single-instrument) score
///
/// l_hat = E_hat[Y | X]
/// r_hat = E_hat[D | X]
/// m_hat = E_hat[Z | X]
/// u_hat = Y - l_hat
/// w_hat = D - r_hat
/// v_hat = Z - m_hat
/// psi_a = -w_hat * v_hat
/// psi_b = v_hat * u_hat
/// psi(theta) = theta * psi_a + psi_b
///
/// with point estimate and variance
///
/// theta_hat = -mean(psi_b) / mean(psi_a)
/// = mean(v_hat * u_hat) / mean(w_hat * v_hat)
/// J = mean(psi_a)
/// gamma = mean(psi(theta_hat)^2)
/// sigma2 = gamma / (J^2 * n)
/// se = sqrt(sigma2).
///
/// The port supports only a *single* instrument and the *partialling
/// out* score (the `IV-type` score, which would need an additional
/// `ml_g` learner, is not implemented). All three nuisance functions
/// are estimated with the same closed-form `LinearRegression` learner.
pub struct DoubleMLPLIV {
data : DoubleMLPLIVData
n_folds : Int
n_rep : Int
seed : Int
// v0.59.0+: injected nuisance learner (replaces the v0.57.0
// hardcoded `LinearRegression`). Defaults to OLS so v0.57.0
// callers see byte-identical results.
learner : LearnerDispatch
l_hat : Array[Double]
r_hat : Array[Double]
m_hat : Array[Double]
coef : Double
se : Double
fitted : Bool
// v0.61.0+: per-observation influence function components
// for the multiplier bootstrap. PLIV partialling-out score:
// psi_a[i] = -w_hat[i] * v_hat[i]
// psi_b[i] = v_hat[i] * u_hat[i]
// where `w_hat = d - r_hat`, `v_hat = z - m_hat`,
// `u_hat = y - l_hat`. Length `n_obs`. Populated by
// `fit(...)` from the last rep's nuisances.
psi_a : Array[Double]
psi_b : Array[Double]
// v0.61.0+: multiplier bootstrap state. `boot_t_stat` is a
// length-`n_rep_boot` array of t-statistics for `coef`.
// Populated by `bootstrap(...)`; empty until then.
boot_t_stat : Array[Double]
boot_method : String
n_rep_boot : Int
boot_seed : Int
} derive(Debug)
///|
pub extend DoubleMLPLIV with @moonbitlang/core/debug.Debug::{to_repr}
///|
pub fn DoubleMLPLIV::new(
data : DoubleMLPLIVData,
n_folds? : Int = 2,
n_rep? : Int = 1,
seed? : Int = 3141,
learner? : LearnerDispatch = LearnerDispatch::linear_regression(),
) -> DoubleMLPLIV {
try {
require(n_folds >= 2)
require(n_folds <= data.n_obs())
require(n_rep >= 1)
{
data,
n_folds,
n_rep,
seed,
learner,
l_hat: Array::make(data.n_obs(), 0.0),
r_hat: Array::make(data.n_obs(), 0.0),
m_hat: Array::make(data.n_obs(), 0.0),
coef: 0.0,
se: 0.0,
fitted: false,
psi_a: Array::make(data.n_obs(), 0.0),
psi_b: Array::make(data.n_obs(), 0.0),
boot_t_stat: [],
boot_method: "",
n_rep_boot: 0,
boot_seed: 0,
}
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
pub fn DoubleMLPLIV::n_obs(self : DoubleMLPLIV) -> Int {
self.data.n_obs()
}
///|
pub fn DoubleMLPLIV::coef(self : DoubleMLPLIV) -> Double {
try {
require(self.fitted)
self.coef
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
pub fn DoubleMLPLIV::se(self : DoubleMLPLIV) -> Double {
try {
require(self.fitted)
self.se
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
pub fn DoubleMLPLIV::confint(self : DoubleMLPLIV) -> (Double, Double) {
try {
require(self.fitted)
let lo = self.coef - 1.96 * self.se
let hi = self.coef + 1.96 * self.se
(lo, hi)
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
pub fn DoubleMLPLIV::predictions_l(self : DoubleMLPLIV) -> Array[Double] {
self.l_hat
}
///|
pub fn DoubleMLPLIV::predictions_r(self : DoubleMLPLIV) -> Array[Double] {
self.r_hat
}
///|
pub fn DoubleMLPLIV::predictions_m(self : DoubleMLPLIV) -> Array[Double] {
self.m_hat
}
///|
/// Run the PLIV estimation. The default learner is a closed-form
/// `LinearRegression`; a different `Learner` can be supplied for
/// experiments. The result is stored on the object and the object is
/// returned for chaining.
///
/// Per-repetition behaviour: each repetition `r` cross-fits the
/// `l / r / m` nuisances from its own folds (seed `self.seed + r`),
/// computes its own `(theta_r, se_r)` from the partialling-out
/// single-instrument score, and the two arrays are then aggregated
/// by `aggregate_coef_se` (median of thetas, then SE from the median
/// of `(theta_r + 1.96 * se_r)`). For `n_rep == 1` the aggregator
/// returns the single `(theta_1, se_1)` exactly, so the byte-equality
/// with the previous "average then estimate" implementation is
/// preserved. The `predictions_l / r / m` accessors return the
/// nuisances from the *last* repetition (the conventional choice in
/// upstream `doubleml`), not a cross-rep average.
pub fn DoubleMLPLIV::fit(
self : DoubleMLPLIV,
learner? : LearnerDispatch = LearnerDispatch::linear_regression(),
max_attempts? : Int = 1,
) -> DoubleMLPLIV {
try {
require(max_attempts >= 1)
if self.data.is_cluster_data() {
return self.fit_cluster(learner~, max_attempts~)
}
let n = self.n_obs()
let nrep = self.n_rep
let coefs : Array[Double] = Array::make(nrep, 0.0)
let ses : Array[Double] = Array::make(nrep, 0.0)
// hold the last rep's predictions; final values land in l_hat / r_hat / m_hat
let mut l_pred : Array[Double] = Array::make(n, 0.0)
let mut r_pred : Array[Double] = Array::make(n, 0.0)
let mut m_pred : Array[Double] = Array::make(n, 0.0)
for r = 0; r < nrep; r = r + 1 {
let folds = kfold(n, self.n_folds, self.seed + r)
l_pred = cross_fit_predict_dispatch(
learner,
self.data.x,
self.data.y,
folds,
)
r_pred = cross_fit_predict_dispatch(
learner,
self.data.x,
self.data.d,
folds,
)
m_pred = cross_fit_predict_dispatch(
learner,
self.data.x,
self.data.z,
folds,
)
// Score (partialling out, single instrument) for THIS rep's nuisances only
let y = self.data.y
let d = self.data.d
let z = self.data.z
let psi_a : Array[Double] = Array::make(n, 0.0)
let psi_b : Array[Double] = Array::make(n, 0.0)
for i = 0; i < n; i = i + 1 {
let u = y[i] - l_pred[i]
let w = d[i] - r_pred[i]
let v = z[i] - m_pred[i]
psi_a[i] = -w * v
psi_b[i] = v * u
}
let (coef_r, se_r) = var_est(psi_a, psi_b)
coefs[r] = coef_r
ses[r] = se_r
}
// last iteration's predictions are now in l_pred / r_pred / m_pred
let (coef, se) = aggregate_coef_se(coefs, ses)
// v0.61.0: per-observation influence function for the
// multiplier bootstrap. Recompute `psi_a / psi_b` from the
// last rep's nuisances so the stored arrays align with
// `l_hat` / `r_hat` / `m_hat` and `coef` (matches the
// v0.20.0+ DID convention).
let psi_a : Array[Double] = Array::make(n, 0.0)
let psi_b : Array[Double] = Array::make(n, 0.0)
for i = 0; i < n; i = i + 1 {
let u = self.data.y[i] - l_pred[i]
let w = self.data.d[i] - r_pred[i]
let v = self.data.z[i] - m_pred[i]
psi_a[i] = -w * v
psi_b[i] = v * u
}
{
data: self.data,
n_folds: self.n_folds,
n_rep: self.n_rep,
seed: self.seed,
// v0.59.0+: persist per-fit learner override on the struct
learner,
l_hat: l_pred,
r_hat: r_pred,
m_hat: m_pred,
coef,
se,
fitted: true,
psi_a,
psi_b,
boot_t_stat: [],
boot_method: "",
n_rep_boot: 0,
boot_seed: 0,
}
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// Clustered-DML path for `DoubleMLPLIV`. Same shape as
/// `DoubleMLPLR::fit_cluster`: folds are drawn over the
/// unique unit ids, expanded to row folds via
/// `expand_unit_folds_to_rows`; coefficient is the
/// fold-weighted ratio of cluster score sums
/// (`est_coef_cluster`); variance is unit-level
/// cluster-robust (`var_est_cluster`). All three nuisances
/// (`l = E[Y|X]`, `r = E[D|X]`, `m = E[Z|X]`) are
/// cross-fitted with cluster-respecting folds; the per-row
/// score elements are the same as the row-level path
/// (`psi_a = -w_hat * v_hat`, `psi_b = v_hat * u_hat`).
fn DoubleMLPLIV::fit_cluster(
self : DoubleMLPLIV,
learner~ : LearnerDispatch,
max_attempts? : Int = 1,
) -> DoubleMLPLIV {
try {
require(max_attempts >= 1)
let cluster = self.data.cluster_vars
let n = self.n_obs()
let nrep = self.n_rep
let uniq = unique_units(cluster)
let n_units = uniq.length()
require(self.n_folds <= n_units)
// v0.36.0: build_row_unit_map raises ClusterDataError on
// malformed cluster vector; catch and re-abort to preserve
// pre-v0.36.0 behavior.
let row_unit = build_row_unit_map(cluster, uniq) catch {
ClusterDataError::MissingUnit(g) =>
abort(
"expand_unit_folds_to_rows: row without a unit id (unit_id=" +
g.to_string() +
")",
)
}
let unit_rows : Array[Array[Int]] = Array::makei(n_units, fn(_) {
let rows : Array[Int] = []
rows
})
for i = 0; i < n; i = i + 1 {
unit_rows[row_unit[i]].push(i)
}
let coefs : Array[Double] = Array::make(nrep, 0.0)
let ses : Array[Double] = Array::make(nrep, 0.0)
let mut l_pred : Array[Double] = Array::make(n, 0.0)
let mut r_pred : Array[Double] = Array::make(n, 0.0)
let mut m_pred : Array[Double] = Array::make(n, 0.0)
for r = 0; r < nrep; r = r + 1 {
// v0.40.0: retry loop on J-floor (see plr.mbt::fit_cluster).
let mut theta_r = 0.0
let mut se_r = 0.0
let mut attempt = 0
let mut succeeded = false
while attempt < max_attempts && !succeeded {
let rep_seed = self.seed + r + attempt * nrep
let folds_u = kfold(n_units, self.n_folds, rep_seed)
let (folds_row, unit_fold, fold_n_units) = expand_unit_folds_to_rows(
cluster, folds_u, row_unit,
)
l_pred = cross_fit_predict_dispatch(
learner,
self.data.x,
self.data.y,
folds_row,
)
r_pred = cross_fit_predict_dispatch(
learner,
self.data.x,
self.data.d,
folds_row,
)
m_pred = cross_fit_predict_dispatch(
learner,
self.data.x,
self.data.z,
folds_row,
)
let y = self.data.y
let d = self.data.d
let z = self.data.z
let psi_a : Array[Double] = Array::make(n, 0.0)
let psi_b : Array[Double] = Array::make(n, 0.0)
for i = 0; i < n; i = i + 1 {
let u = y[i] - l_pred[i]
let w = d[i] - r_pred[i]
let v = z[i] - m_pred[i]
psi_a[i] = -w * v
psi_b[i] = v * u
}
let (t, s) = cluster_causal_param_and_se(
psi_a,
psi_b,
folds_row,
fold_n_units,
unit_rows,
unit_fold,
folds_u.length(),
self.n_folds,
) catch {
_ => {
attempt = attempt + 1
(0.0, 0.0)
}
}
theta_r = t
se_r = s
succeeded = true
}
if !succeeded {
abort(
"var_est_cluster: J-floor fired " +
max_attempts.to_string() +
" times for rep=" +
r.to_string() +
" (cluster SE numerically unstable across multiple fold splits, try a different seed or larger n_units)",
)
}
coefs[r] = theta_r
ses[r] = se_r
}
let (coef, se) = aggregate_coef_se(coefs, ses)
// v0.61.0: per-observation influence function for the
// multiplier bootstrap. Same convention as `fit()`:
// recompute from the last rep's cross-fitted nuisances so
// the stored arrays align with `l_hat` / `r_hat` /
// `m_hat` and `coef`.
let psi_a : Array[Double] = Array::make(n, 0.0)
let psi_b : Array[Double] = Array::make(n, 0.0)
for i = 0; i < n; i = i + 1 {
let u = self.data.y[i] - l_pred[i]
let w = self.data.d[i] - r_pred[i]
let v = self.data.z[i] - m_pred[i]
psi_a[i] = -w * v
psi_b[i] = v * u
}
{
data: self.data,
n_folds: self.n_folds,
n_rep: self.n_rep,
seed: self.seed,
// v0.59.0+: persist per-fit learner override on the
// cluster-path return struct.
learner,
l_hat: l_pred,
r_hat: r_pred,
m_hat: m_pred,
coef,
se,
fitted: true,
psi_a,
psi_b,
boot_t_stat: [],
boot_method: "",
n_rep_boot: 0,
boot_seed: 0,
}
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// v0.61.0+: multiplier bootstrap for `DoubleMLPLIV`. The
/// per-observation influence function is
///
/// psi[i] = psi_a[i] + theta * psi_b[i]
/// = -w*v + theta * v*u
///
/// where `w = d - r_hat`, `v = z - m_hat`, `u = y - l_hat`
/// are computed at the fitted `coef` from the last rep's
/// cross-fitted nuisances `l_hat` / `r_hat` / `m_hat`.
/// Draws `n_rep_boot` weight vectors of length `n_obs` from
/// the chosen multiplier distribution, and returns a fitted
/// model with `boot_t_stat[b] = sum_i w[b, i] * psi[i] /
/// (sqrt(n) * se_psi)` populated where
/// `se_psi = sqrt(mean(psi^2))`.
///
/// `method_name` selects the multiplier distribution:
/// - `"normal"` (default): `w[i] ~ N(0, 1)`.
/// - `"Bayes"`: `w[i] = exp(1) - 1` (mean 0, var 1).
/// - `"wild"`: `w[i] = x[i] / sqrt(2) + (y[i]^2 - 1) / 2`
/// with `x, y ~ N(0, 1)`.
///
/// Calling `bootstrap` requires the model to be fitted; calling
/// on an un-fit model aborts with `PreconditionError`. The
/// helper is `did_bootstrap_t_stat` (v0.55.0 extracted from
/// `DoubleMLDIDCrossSection::bootstrap`); PLIV is the
/// `n_thetas=1` case.
pub fn DoubleMLPLIV::bootstrap(
self : DoubleMLPLIV,
method_name? : String = "normal",
n_rep_boot? : Int = 500,
seed? : Int = 2024,
) -> DoubleMLPLIV {
try {
require(self.fitted)
require(
method_name == "normal" || method_name == "Bayes" || method_name == "wild",
)
require(n_rep_boot >= 2)
let n = self.n_obs()
// Draw weights. Shape: (n_rep_boot, n_obs).
let weights = draw_bootstrap_weights(method_name, n_rep_boot, n, seed) catch {
BootstrapMethodError::UnknownMethod(m) =>
abort(
"draw_bootstrap_weights: unknown method (set in DoubleMLPLIV::bootstrap): " +
m,
)
}
// Compute psi[i] = psi_a[i] + coef * psi_b[i] and
// ss_psi = sum(psi[i]^2) once. Both `psi_a` and `psi_b`
// were populated by `fit(...)` from the last rep's
// nuisances.
let psi : Array[Double] = Array::make(n, 0.0)
let mut ss_psi = 0.0
for i = 0; i < n; i = i + 1 {
let psi_i = self.psi_a[i] + self.coef * self.psi_b[i]
psi[i] = psi_i
ss_psi = ss_psi + psi_i * psi_i
}
let n_d = n.to_double()
let se_psi = (ss_psi / n_d).sqrt()
if se_psi <= 0.0 {
// Degenerate: psi sums to 0. Cannot divide.
let boot_t_stat_zero : Array[Double] = Array::make(n_rep_boot, 0.0)
return {
..self,
boot_t_stat: boot_t_stat_zero,
boot_method: method_name,
n_rep_boot,
boot_seed: seed,
}
}
// n_thetas=1 case.
let se_flat : Array[Double] = [se_psi]
let boot_t_stat = did_bootstrap_t_stat(
weights, psi, se_flat, n_rep_boot, n, 1,
)
{
..self,
boot_t_stat,
boot_method: method_name,
n_rep_boot,
boot_seed: seed,
}
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}