///|
/// Average potential outcomes for an arbitrary treatment level.
pub struct DoubleMLAPO {
data : DoubleMLData
treatment_level : Double
n_folds : Int
n_rep : Int
seed : Int
propensity_clip : Double
// v0.60.0+: injected nuisance learners. Defaults to
// `LearnerDispatch::linear_regression()` so v0.59.0 callers
// see byte-identical results. v0.61.0+ will plumb these
// through `cross_fit_predict` for `g_hat` and `m_hat`; for
// v0.60.0 they're stored on the struct and returned via the
// accessors but not yet consumed internally.
ml_g : LearnerDispatch
ml_m : LearnerDispatch
g_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. APO policy score:
// psi_a[i] = -1 (constant)
// psi_b[i] = g[i] + treated[i] * (y[i] - g[i]) / m[i]
// where `treated[i] = 1{d[i] == treatment_level}`,
// `g = E[Y | X]`, `m = P(D = treatment_level | X)` (clipped).
// Length `n_obs`. Populated by `fit(...)`.
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 DoubleMLAPO with @moonbitlang/core/debug.Debug::{to_repr}
///|
pub fn DoubleMLAPO::new(
data : DoubleMLData,
treatment_level? : Double = 1.0,
n_folds? : Int = 2,
n_rep? : Int = 1,
seed? : Int = 3141,
propensity_clip? : Double = 1.0e-6,
ml_g? : LearnerDispatch = LearnerDispatch::linear_regression(),
ml_m? : LearnerDispatch = LearnerDispatch::linear_regression(),
) -> DoubleMLAPO {
try {
require(n_folds >= 2)
require(n_folds <= data.n_obs())
require(n_rep >= 1)
require(propensity_clip > 0.0)
require(propensity_clip < 0.5)
{
data,
treatment_level,
n_folds,
n_rep,
seed,
propensity_clip,
ml_g,
ml_m,
g_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(), -1.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 DoubleMLAPO::n_obs(self : DoubleMLAPO) -> Int {
self.data.n_obs()
}
///|
/// Number of features (covariate columns).
pub fn DoubleMLAPO::n_features(self : DoubleMLAPO) -> Int {
self.data.n_features()
}
///|
/// Accessor for the outcome-nuisance learner used by the most
/// recent `fit(...)` call. v0.60.0+.
pub fn DoubleMLAPO::learner_g(self : DoubleMLAPO) -> LearnerDispatch {
self.ml_g
}
///|
/// Accessor for the propensity-score learner used by the most
/// recent `fit(...)` call. v0.60.0+.
pub fn DoubleMLAPO::learner_m(self : DoubleMLAPO) -> LearnerDispatch {
self.ml_m
}
///|
pub fn DoubleMLAPO::coef(self : DoubleMLAPO) -> Double {
try {
require(self.fitted)
self.coef
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
pub fn DoubleMLAPO::se(self : DoubleMLAPO) -> Double {
try {
require(self.fitted)
self.se
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
pub fn DoubleMLAPO::confint(self : DoubleMLAPO) -> (Double, Double) {
try {
require(self.fitted)
(self.coef - 1.96 * self.se, self.coef + 1.96 * self.se)
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
pub fn DoubleMLAPO::predictions_g(self : DoubleMLAPO) -> Array[Double] {
try {
require(self.fitted)
self.g_hat
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
pub fn DoubleMLAPO::predictions_m(self : DoubleMLAPO) -> Array[Double] {
try {
require(self.fitted)
self.m_hat
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
fn indicator_level(d : Array[Double], level : Double) -> Array[Double] {
let out = Array::make(d.length(), 0.0)
for i = 0; i < d.length(); i = i + 1 {
out[i] = if d[i] == level { 1.0 } else { 0.0 }
}
out
}
///|
/// Cross-fit the APO nuisances. The order is asymmetric: the
/// conditional outcome `g` is fit only on the treated subset (the
/// same as the upstream `DoubleMLAPO`), while the propensity `m`
/// is fit on the *full* training fold (including controls). The
/// effect is that rows with `treated = 0` receive their `g` value
/// from the previous fold's treated-only fit (or stay at the
/// default 0.0 if no fold has yet seen them); rows with `treated = 1`
/// receive both a `g` (on the treated subset) and a propensity
/// update. Matches the upstream `DoubleMLAPO` convention.
///
/// v0.63.0+: routes `ml_g` (for the conditional outcome on the
/// treated subset) and `ml_m` (for the propensity on the full
/// fold) through `LearnerDispatch`. Defaults preserve v0.62.2
/// byte-equality (default `LearnerDispatch::linear_regression()`
/// gives the v0.62.2 OLS path).
fn cross_fit_apo(
ml_g : LearnerDispatch,
ml_m : LearnerDispatch,
x : Matrix,
y : Array[Double],
treated : Array[Double],
folds : Array[Fold],
clip : Double,
) -> (Array[Double], Array[Double]) {
let n = x.rows()
let g = Array::make(n, 0.0)
let m = Array::make(n, 0.0)
for fold in folds {
let tr = fold.train_indices()
let te = fold.test_indices()
let tg = filter_indices(tr, treated)
if tg.length() > 0 {
let p = fit_predict_one_dispatch(
ml_g,
slice_matrix_rows(x, tg),
slice_vector(y, tg),
slice_matrix_rows(x, te),
)
for k = 0; k < te.length(); k = k + 1 {
g[te[k]] = p[k]
}
}
let p = fit_predict_one_dispatch(
ml_m,
slice_matrix_rows(x, tr),
slice_vector(treated, tr),
slice_matrix_rows(x, te),
)
for k = 0; k < te.length(); k = k + 1 {
m[te[k]] = p[k]
}
}
(g, clip_vec(m, clip, 1.0 - clip))
}
///|
pub fn DoubleMLAPO::fit(
self : DoubleMLAPO,
ml_g? : LearnerDispatch = self.ml_g,
ml_m? : LearnerDispatch = self.ml_m,
) -> DoubleMLAPO {
// v0.65.0+: cluster-data dispatch — when `cluster_vars` is
// non-empty, route through `fit_cluster` (cluster-aware
// folds, unit-level cluster-robust SE).
if self.data.is_cluster_data() {
return self.fit_cluster(ml_g~, ml_m~)
}
let n = self.n_obs()
let treated = indicator_level(self.data.d, self.treatment_level)
// Accumulate `g` and `m` directly in the storage arrays; divide
// by `n_rep` after the loop.
let g = Array::make(n, 0.0)
let m = Array::make(n, 0.0)
for r = 0; r < self.n_rep; r = r + 1 {
let (gr, mr) = cross_fit_apo(
ml_g,
ml_m,
self.data.x,
self.data.y,
treated,
kfold(n, self.n_folds, self.seed + r),
self.propensity_clip,
)
for i = 0; i < n; i = i + 1 {
g[i] = g[i] + gr[i]
m[i] = m[i] + mr[i]
}
}
let inv = 1.0 / self.n_rep.to_double()
for i = 0; i < n; i = i + 1 {
g[i] = g[i] * inv
m[i] = m[i] * inv
}
// `pa` is the APO score's `psi_a` component, which is structurally
// `-1` for every observation (the potential-outcome score has no
// treatment-side variation). `pb` is the `psi_b` component with
// the IPW-style centring `g + treated * (y - g) / m`.
let pa = Array::make(n, -1.0)
let pb = Array::make(n, 0.0)
for i = 0; i < n; i = i + 1 {
pb[i] = g[i] + treated[i] * (self.data.y[i] - g[i]) / m[i]
}
// Score + variance via the shared `var_est` helper.
let (theta, se) = var_est(pa, pb)
{
data: self.data,
treatment_level: self.treatment_level,
n_folds: self.n_folds,
n_rep: self.n_rep,
seed: self.seed,
propensity_clip: self.propensity_clip,
ml_g,
ml_m,
g_hat: g,
m_hat: m,
coef: theta,
se,
fitted: true,
// v0.61.0: persist `pa` / `pb` as `psi_a` / `psi_b` for
// the multiplier bootstrap (`pa` is constant -1, `pb` is
// the IPW-centred potential-outcome score).
psi_a: pa,
psi_b: pb,
boot_t_stat: [],
boot_method: "",
n_rep_boot: 0,
boot_seed: 0,
}
}
///|
/// v0.65.0+: clustered-DML path for `DoubleMLAPO`. Folds
/// partition whole units (`kfold` on unique cluster ids,
/// expanded to row folds); nuisances are cross-fitted under
/// cluster folds; `psi_a = -1` (constant), `psi_b = g +
/// treated * (y - g) / m` is computed from the cluster
/// nuisances; coefficient is the fold-weighted ratio of
/// cluster score sums and SE is unit-level cluster-robust
/// (`cluster_causal_param_and_se`). The cluster path
/// differs from the row-level path only in the fold partition
/// and the two aggregation steps; the per-row score elements
/// are identical, so a single nuisances cross-fit (with
/// cluster-respecting folds) feeds both paths.
fn DoubleMLAPO::fit_cluster(
self : DoubleMLAPO,
ml_g~ : LearnerDispatch,
ml_m~ : LearnerDispatch,
max_attempts? : Int = 1,
) -> DoubleMLAPO {
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)
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 treated = indicator_level(self.data.d, self.treatment_level)
let coefs : Array[Double] = Array::make(nrep, 0.0)
let ses : Array[Double] = Array::make(nrep, 0.0)
let mut g : Array[Double] = Array::make(n, 0.0)
let mut m : Array[Double] = Array::make(n, 0.0)
for r = 0; r < nrep; r = r + 1 {
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,
)
let (g_r, m_r) = cross_fit_apo(
ml_g,
ml_m,
self.data.x,
self.data.y,
treated,
folds_row,
self.propensity_clip,
)
g = g_r
m = m_r
let pa : Array[Double] = Array::make(n, 0.0)
let pb : Array[Double] = Array::make(n, 0.0)
for i = 0; i < n; i = i + 1 {
pa[i] = -1.0
pb[i] = g[i] + treated[i] * (self.data.y[i] - g[i]) / m[i]
}
let (t, s) = cluster_causal_param_and_se(
pa,
pb,
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.65.0+: per-observation IF (constant -1 + IPW score)
// persisted for the multiplier bootstrap; recomputed from
// the last rep's cluster-aware nuisances so the stored
// arrays align with `g_hat` / `m_hat` and `coef`.
let pa : Array[Double] = Array::make(n, 0.0)
let pb : Array[Double] = Array::make(n, 0.0)
for i = 0; i < n; i = i + 1 {
pa[i] = -1.0
pb[i] = g[i] + treated[i] * (self.data.y[i] - g[i]) / m[i]
}
{
data: self.data,
treatment_level: self.treatment_level,
n_folds: self.n_folds,
n_rep: self.n_rep,
seed: self.seed,
propensity_clip: self.propensity_clip,
ml_g,
ml_m,
g_hat: g,
m_hat: m,
coef,
se,
fitted: true,
psi_a: pa,
psi_b: pb,
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 `DoubleMLAPO`. The
/// per-observation influence function is
///
/// psi[i] = psi_a[i] + theta * psi_b[i]
/// = -1 + theta * pb[i]
///
/// where `pb[i] = g[i] + treated[i] * (y[i] - g[i]) / m[i]`
/// is the IPW-centred potential-outcome score, computed at
/// the fitted `coef` from the stored `g_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`); APO is the
/// `n_thetas=1` case.
pub fn DoubleMLAPO::bootstrap(
self : DoubleMLAPO,
method_name? : String = "normal",
n_rep_boot? : Int = 500,
seed? : Int = 2024,
) -> DoubleMLAPO {
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 DoubleMLAPO::bootstrap): " +
m,
)
}
// Compute psi[i] = psi_a[i] + coef * psi_b[i] and
// ss_psi = sum(psi[i]^2) once. `psi_a[i] = -1` and
// `psi_b` is the IPW-centred potential-outcome score,
// both populated by `fit(...)`.
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())
}
}
///|
/// v0.65.0+: tune the (ml_g, ml_m) nuisance-learner pair via
/// MSE-on-g_hat cross-fitting. The chosen `(learner_g,
/// learner_m)` is then re-fit on the FINAL-FIT fold partition
/// (`self.n_folds`) under `DoubleMLAPO::fit`. `treatment_level`
/// is held fixed at `self.treatment_level`.
///
/// `param_set` is an `Array[TuneParam]`; each entry is a
/// `(learner_g, learner_m)` pair. `scoring_method` is
/// `"MSE"` (default), `"RMSE"`, or `"NegMSE"`. Returns a
/// re-fitted `DoubleMLAPO` with the chosen pair applied.
pub fn DoubleMLAPO::tune(
self : DoubleMLAPO,
param_set~ : Array[TuneParam],
scoring_method? : String = "MSE",
n_folds_tune? : Int = 5,
seed? : Int = 3141,
) -> DoubleMLAPO {
try {
require(param_set.length() > 0)
require(n_folds_tune >= 2)
let scoring = TuneScoring::parse(scoring_method)
let folds_tune = kfold(self.n_obs(), n_folds_tune, seed)
let n = self.n_obs()
let scores : Array[Double] = Array::make(param_set.length(), 0.0)
for i = 0; i < param_set.length(); i = i + 1 {
let c = param_set[i]
let g_hat_c = cross_fit_predict_dispatch(
c.learner_l,
self.data.x,
self.data.y,
folds_tune,
)
scores[i] = if g_hat_c.length() == n {
tune_score_outcome(self.data.y, g_hat_c, scoring)
} else {
TUNE_SCORE_FAIL_SENTINEL
}
}
let best_idx = if scoring is NegMSE {
let mut bi = 0
let mut bv = scores[0]
for i = 1; i < scores.length(); i = i + 1 {
if scores[i] > bv {
bv = scores[i]
bi = i
}
}
bi
} else {
let mut bi = 0
let mut bv = scores[0]
for i = 1; i < scores.length(); i = i + 1 {
if scores[i] < bv {
bv = scores[i]
bi = i
}
}
bi
}
let best_param = param_set[best_idx]
self.fit(ml_g=best_param.learner_l, ml_m=best_param.learner_m)
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// v0.66.0+: Cinelli & Hazlett (2020) omitted-variable bias
/// analysis. Outcome residual is `y - g_hat` (the APO is
/// conditioned on `treated=1`); the Riesz-representer
/// variance is `mean(psi_a^2) = 1` for the APO score
/// (`psi_a = -1` constant). Routes through the shared
/// `irm_style_sensitivity` helper.
pub fn DoubleMLAPO::sensitivity_analysis(
self : DoubleMLAPO,
cf_y? : Double = 0.05,
cf_d? : Double = 0.05,
) -> SensitivityResult raise {
require(self.fitted)
let g_hat = self.predictions_g()
let n = g_hat.length()
let residuals : Array[Double] = Array::make(n, 0.0)
for i = 0; i < n; i = i + 1 {
residuals[i] = self.data.y[i] - g_hat[i]
}
irm_style_sensitivity(self.coef, residuals, self.psi_a, cf_y, cf_d)
}
///|
/// Average potential outcomes *symmetric* across multiple treatment
/// levels. v0.49.0: full upstream parity — `treatment_levels`
/// validation in `new` (rejects duplicates and levels not in
/// `data.d`), the `causal_contrast` method (level-by-level delta
/// and SE vs a reference level), and the
/// `treatment_levels()` / `n_treatment_levels()` / `fitted()`
/// accessors.
///
/// Each treatment level is fit with the same fold partition
/// (the parent does not currently route a shared partition to the
/// child `DoubleMLAPO`; each child draws its own folds via
/// `kfold`. v0.50.0+ plans to wire the parent through a
/// `fit_with_splits` helper to share one stratified partition,
/// but the v0.49.0 implementation is the v0.50.0 PR target) and
/// the same closed-form `LinearRegression` learner. The
/// `causal_contrast(reference_levels)` method then returns the
/// level-by-level difference `coefs[i] - coefs[ref]` for one or
/// more reference levels, matching the upstream
/// `DoubleMLAPOS.causal_contrast` semantics.
pub struct DoubleMLAPOS {
data : DoubleMLData
treatment_levels : Array[Double]
n_folds : Int
n_rep : Int
seed : Int
propensity_clip : Double
// v0.60.0+: injected nuisance learners (forwarded to
// each child `DoubleMLAPO` on fit). Same forward-compat
// story as the rest of the v0.60.0 estimator set.
ml_g : LearnerDispatch
ml_m : LearnerDispatch
coefs : Array[Double]
ses : Array[Double]
fitted : Bool
// v0.63.0+: multiplier-bootstrap state. `boot_t_stat[j]` is
// the length-`n_rep_boot` t-stat array for `coefs[j]`
// (the APO coefficient at treatment level `j`).
// Populated by `bootstrap(...)`; empty until then.
boot_t_stat : Array[Array[Double]]
boot_method : String
n_rep_boot : Int
boot_seed : Int
} derive(Debug)
///|
pub extend DoubleMLAPOS with @moonbitlang/core/debug.Debug::{to_repr}
///|
pub fn DoubleMLAPOS::new(
data : DoubleMLData,
treatment_levels : Array[Double],
n_folds? : Int = 2,
n_rep? : Int = 1,
seed? : Int = 3141,
propensity_clip? : Double = 1.0e-6,
ml_g? : LearnerDispatch = LearnerDispatch::linear_regression(),
ml_m? : LearnerDispatch = LearnerDispatch::linear_regression(),
) -> DoubleMLAPOS {
try {
require(treatment_levels.length() >= 1)
require(n_folds >= 2)
require(n_folds <= data.n_obs())
require(n_rep >= 1)
require(propensity_clip > 0.0)
require(propensity_clip < 0.5)
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
// v0.49.0: every requested treatment level must be present in
// the data's treatment assignment (`np.unique(data.d)` in
// upstream). Catch duplicates within the request itself too.
for i = 0; i < treatment_levels.length(); i = i + 1 {
let lvl = treatment_levels[i]
let mut dup = false
for j = 0; j < i; j = j + 1 {
if treatment_levels[j] == lvl {
dup = true
break
}
}
if dup {
abort(
"DoubleMLAPOS: treatment_levels contains a duplicate entry: " +
lvl.to_string(),
)
}
let mut in_data = false
for k = 0; k < data.d.length(); k = k + 1 {
if data.d[k] == lvl {
in_data = true
break
}
}
if !in_data {
abort(
"DoubleMLAPOS: treatment_level " +
lvl.to_string() +
" is not present in data.d",
)
}
}
{
data,
treatment_levels,
n_folds,
n_rep,
seed,
propensity_clip,
ml_g,
ml_m,
coefs: Array::make(treatment_levels.length(), 0.0),
ses: Array::make(treatment_levels.length(), 0.0),
fitted: false,
boot_t_stat: [],
boot_method: "",
n_rep_boot: 0,
boot_seed: 0,
}
}
///|
pub fn DoubleMLAPOS::fit(
self : DoubleMLAPOS,
ml_g? : LearnerDispatch = self.ml_g,
ml_m? : LearnerDispatch = self.ml_m,
) -> DoubleMLAPOS {
try {
ignore(ml_g)
ignore(ml_m)
require(self.treatment_levels.length() >= 1)
let c = Array::make(self.treatment_levels.length(), 0.0)
let s = Array::make(self.treatment_levels.length(), 0.0)
for j = 0; j < self.treatment_levels.length(); j = j + 1 {
// v0.49.0: the parent APOS fits each child `DoubleMLAPO` with
// `n_rep = self.n_rep` so the child draws its own fold
// partition (`n_rep` total fold draws per treatment level).
// The parent does not currently route a shared stratified
// partition to the child — that's a v0.50.0+ target. The
// resulting fold-draw count is `n_rep * n_treatment_levels`,
// matching the upstream `DoubleMLAPOS.fit` total.
let z = DoubleMLAPO::new(
self.data,
treatment_level=self.treatment_levels[j],
n_folds=self.n_folds,
n_rep=self.n_rep,
seed=self.seed,
propensity_clip=self.propensity_clip,
).fit()
c[j] = z.coef()
s[j] = z.se()
}
{
data: self.data,
treatment_levels: self.treatment_levels,
n_folds: self.n_folds,
n_rep: self.n_rep,
seed: self.seed,
propensity_clip: self.propensity_clip,
ml_g,
ml_m,
coefs: c,
ses: s,
fitted: true,
boot_t_stat: [],
boot_method: "",
n_rep_boot: 0,
boot_seed: 0,
}
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// Accessor for the outcome-nuisance learner used by the most
/// recent `fit(...)` call. v0.60.0+.
pub fn DoubleMLAPOS::learner_g(self : DoubleMLAPOS) -> LearnerDispatch {
self.ml_g
}
///|
/// Accessor for the propensity-score learner used by the most
/// recent `fit(...)` call. v0.60.0+.
pub fn DoubleMLAPOS::learner_m(self : DoubleMLAPOS) -> LearnerDispatch {
self.ml_m
}
///|
pub fn DoubleMLAPOS::coefs(self : DoubleMLAPOS) -> Array[Double] {
try {
require(self.fitted)
self.coefs
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
pub fn DoubleMLAPOS::ses(self : DoubleMLAPOS) -> Array[Double] {
try {
require(self.fitted)
self.ses
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// v0.49.0: the requested treatment levels, in user-supplied order.
pub fn DoubleMLAPOS::treatment_levels(self : DoubleMLAPOS) -> Array[Double] {
self.treatment_levels
}
///|
/// v0.49.0: number of requested treatment levels.
pub fn DoubleMLAPOS::n_treatment_levels(self : DoubleMLAPOS) -> Int {
self.treatment_levels.length()
}
///|
/// Per-level multiplier bootstrap for `DoubleMLAPOS`. Each
/// treatment level's child `DoubleMLAPO` is re-fit (deterministic
/// given `seed`) and its `bootstrap(...)` invoked; the per-level
/// `boot_t_stat` arrays are concatenated into a length-`n_levels`
/// `Array[Array[Double]]` and also written to `self.boot_t_stat`.
///
/// `method_name` selects the multiplier distribution: `"normal"`
/// (default), `"Bayes"`, `"wild"` — see `DoubleMLAPO::bootstrap`.
/// `n_rep_boot` defaults to 500; `seed` defaults to 2024.
pub fn DoubleMLAPOS::bootstrap(
self : DoubleMLAPOS,
method_name? : String = "normal",
n_rep_boot? : Int = 500,
seed? : Int = 2024,
) -> DoubleMLAPOS {
try {
require(self.fitted)
require(
method_name == "normal" || method_name == "Bayes" || method_name == "wild",
)
require(n_rep_boot >= 2)
let n_levels = self.treatment_levels.length()
let per_level : Array[Array[Double]] = Array::make(n_levels, [])
for j = 0; j < n_levels; j = j + 1 {
let child = DoubleMLAPO::new(
self.data,
treatment_level=self.treatment_levels[j],
n_folds=self.n_folds,
n_rep=self.n_rep,
seed=self.seed,
propensity_clip=self.propensity_clip,
ml_g=self.ml_g,
ml_m=self.ml_m,
)
let fitted_child = child.fit()
let booted = fitted_child.bootstrap(method_name~, n_rep_boot~, seed~)
per_level[j] = booted.boot_t_stat
}
{
..self,
boot_t_stat: per_level,
boot_method: method_name,
n_rep_boot,
boot_seed: seed,
}
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// Per-level Wald confidence intervals `(coef - z * se, coef + z * se)`
/// for `DoubleMLAPOS`. Returns one `(lo, hi)` tuple per treatment
/// level, in user-supplied order. `level` defaults to 0.95 (z =
/// 1.959963984540054 for the standard normal).
pub fn DoubleMLAPOS::confint(
self : DoubleMLAPOS,
level? : Double = 0.95,
joint? : Bool = false,
) -> Array[(Double, Double)] {
try {
require(self.fitted)
require(level > 0.0 && level < 1.0)
if joint {
require(self.boot_t_stat.length() > 0)
}
let alpha = 1.0 - level
let mut z = norm_ppf(1.0 - alpha / 2.0)
let n_levels = self.treatment_levels.length()
let out : Array[(Double, Double)] = Array::make(n_levels, (0.0, 0.0))
if joint {
// Joint critical value: (1 - alpha) quantile of max |t|
// over the n_levels t-statistics per bootstrap rep.
// `boot_t_stat : Array[Array[Double]]` of length
// `n_levels`; each per-level array is length
// `n_rep_boot`.
let n_boot = self.n_rep_boot
let max_t_arr : Array[Double] = Array::make(n_boot, 0.0)
for b = 0; b < n_boot; b = b + 1 {
let mut mx = 0.0
for j = 0; j < n_levels; j = j + 1 {
let per_level = self.boot_t_stat[j]
let t : Double = per_level[b]
let abs_t : Double = if t < 0.0 { -t } else { t }
if abs_t > mx {
mx = abs_t
}
}
max_t_arr[b] = mx
}
// Inline empirical_quantile (insertion sort, no helper).
max_t_arr.sort()
let idx = ((n_boot - 1).to_double() * (1.0 - alpha)).to_int()
z = max_t_arr[idx]
}
for j = 0; j < n_levels; j = j + 1 {
let lo = self.coefs[j] - z * self.ses[j]
let hi = self.coefs[j] + z * self.ses[j]
out[j] = (lo, hi)
}
out
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// Per-level `boot_t_stat` accessors (v0.63.0+).
pub fn DoubleMLAPOS::boot_t_stats(self : DoubleMLAPOS) -> Array[Array[Double]] {
try {
require(self.fitted)
self.boot_t_stat
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
pub fn DoubleMLAPOS::boot_method(self : DoubleMLAPOS) -> String {
self.boot_method
}
///|
pub fn DoubleMLAPOS::n_rep_boot(self : DoubleMLAPOS) -> Int {
self.n_rep_boot
}
///|
pub fn DoubleMLAPOS::boot_seed(self : DoubleMLAPOS) -> Int {
self.boot_seed
}
///|
/// v0.49.0: whether `fit` has been called.
pub fn DoubleMLAPOS::fitted(self : DoubleMLAPOS) -> Bool {
self.fitted
}
///|
/// v0.49.0: causal contrasts between the requested treatment levels
/// and the supplied reference level(s). Returns one
/// `Array[Double]` of length `2 * treatment_levels.length() - 1`
/// per reference level: the ref-level slot is a single `0.0` (its
/// contrast is trivially zero), and every other slot is a
/// `(delta, se)` pair where `delta = coefs[i] - coefs[ref_idx]`
/// and `se = sqrt(se[i]^2 + se[ref_idx]^2)`. `ref_idx` is the
/// position of the reference level in `treatment_levels`. The
/// layout is interleaved (`[0.0, delta_0, se_0, delta_1, se_1, ...]`
/// for a 2-level input) rather than a `[(level, coef, se), ...]`
/// table — see `apo_test.mbt::apos_causal_contrast_with_reference`
/// for the actual indices.
///
/// The SE of each contrast is computed via the standard
/// `var(psi_a) + var(psi_b) - 2*cov(psi_a, psi_b)` style
/// approximation; for v0.49.0 we take the conservative
/// `sqrt(se_a^2 + se_b^2)` route (same as the upstream
/// `causal_contrast` summary table), which is exact when the
/// per-level psi_a and psi_b are independent across levels
/// (true under stratified kfold with disjoint train indices).
pub fn DoubleMLAPOS::causal_contrast(
self : DoubleMLAPOS,
reference_levels : Array[Double],
) -> Array[Array[Double]] {
// v0.49.0: pre-condition check is local; ref_indices is built
// up in the same scope that consumes it. `abort` (not raise) is
// used so the function signature stays clean.
if !self.fitted {
abort("precondition failed at DoubleMLAPOS::causal_contrast: not fitted")
}
if reference_levels.length() < 1 {
abort(
"precondition failed at DoubleMLAPOS::causal_contrast: reference_levels is empty",
)
}
let ref_indices : Array[Int] = []
for r = 0; r < reference_levels.length(); r = r + 1 {
let ref_lvl = reference_levels[r]
let mut found = false
for i = 0; i < self.treatment_levels.length(); i = i + 1 {
if self.treatment_levels[i] == ref_lvl {
ignore(ref_indices.push(i))
found = true
break
}
}
if !found {
abort(
"DoubleMLAPOS::causal_contrast: reference_level " +
ref_lvl.to_string() +
" is not in treatment_levels",
)
}
}
let results : Array[Array[Double]] = []
for r = 0; r < ref_indices.length(); r = r + 1 {
let ref_idx = ref_indices[r]
let row : Array[Double] = []
for i = 0; i < self.treatment_levels.length(); i = i + 1 {
if i == ref_idx {
ignore(row.push(0.0))
} else {
// delta = coef_i - coef_ref, se = sqrt(se_i^2 + se_ref^2)
let delta = self.coefs[i] - self.coefs[ref_idx]
let se = (self.ses[i] * self.ses[i] +
self.ses[ref_idx] * self.ses[ref_idx]).sqrt()
ignore(row.push(delta))
ignore(row.push(se))
}
}
results.push(row)
}
results
}