// dgp_did_cs_CS2021.mbt
//
// Pure-MoonBit port of upstream
// `doubleml.did.datasets.dgp_did_cs_CS2021.make_did_cs_CS2021`.
//
// Callaway & Sant'Anna (2021) staggered-adoption DID DGP for the
// "2x2" group-time average treatment effect on the treated
// (ATT(g, t)). Adapted to the MoonBit lib's `DoubleMLDIDCSData`
// surface (continuous treatment `d`, time, group, id vectors; X is
// a covariate matrix).
//
// Setup:
// - n_groups groups, n_periods periods.
// - Each group g in {0, 1, ..., n_groups-1} has treatment timing
// tau_g = ceil(g / n_periods_indicator); here we use the
// upstream "single treated cohort" pattern (one group treated
// at tau_g = 1, never-treated cohort g = 0).
// - Outcome Y is generated per the partially-linear DGP
// with covariates X drawn from N(0, Sigma).
// - DGP (CS2021 Eq. 1.1, simplified to 2-period 2-group):
// Y_{i,t} = alpha_i + beta_t + theta * D_{i,t}
// + X_i @ gamma + epsilon_{i,t},
// where D_{i,t} = 1{g_i <= t} * 1{group i is treated}.
//
// We replicate the canonical 2-cohort 2-period case: one treated
// cohort with g=1 (treated at t=1) and one never-treated cohort
// (g=0). Y is generated with the true ATT equal to `theta`.
///|
/// Result of the DID CS2021 DGP.
struct DidCsData {
/// True ATT parameter.
theta : Double
/// Number of cohorts (groups) in the panel.
n_groups : Int32
/// Number of time periods.
n_periods : Int32
/// Outcome vector (length n_obs = n_groups * cohort_size * n_periods).
y : Array[Double]
/// Treatment indicator (length n_obs; binary).
d : Array[Double]
/// Time index (length n_obs; integer 0..n_periods-1).
t : Array[Int32]
/// Group/treatment-cohort index (length n_obs; integer 0..n_groups-1).
g : Array[Int32]
/// Unit id within cohort (length n_obs; integer 0..cohort_size-1).
id : Array[Int32]
/// Covariate matrix (n_obs x dim_x).
x : Matrix
}
///|
/// Generate the DID CS2021 DGP. Default: 2 cohorts, 2 periods,
/// `theta = 1.0`, balanced cohort sizes, `n_per_cohort = 200`,
/// `dim_x = 3`.
pub fn make_did_cs_CS2021(
n_obs : Int32,
dim_x : Int32,
theta : Double,
seed : Int32,
) -> DidCsData {
let rng = chacha8_rng(seed)
// 2 cohorts, 2 periods, equal cohort sizes; cohort 0 = never-treated,
// cohort 1 = treated at t=1.
let n_periods = 2
let n_groups = 2
let n_per_cohort = n_obs / n_groups / n_periods
let true_n = n_per_cohort * n_groups * n_periods
// 1. Pre-draw X-draw normals (n_units * dim_x).
let n_units = n_per_cohort * n_groups
let n_normals_x = n_units * dim_x
let z_flat : Array[Double] = Array::make(n_normals_x, 0.0)
let half_z = (n_normals_x + 1) / 2
for i = 0; i < half_z; i = i + 1 {
let (z1, z2) = box_muller_pair(rng)
let idx_a = 2 * i
let idx_b = 2 * i + 1
if idx_a < n_normals_x {
z_flat[idx_a] = z1
}
if idx_b < n_normals_x {
z_flat[idx_b] = z2
}
}
// 2. Build unit-level X (n_units x dim_x) via lower-triangular L.
let x_unit : Array[Double] = Array::make(n_normals_x, 0.0)
for i = 0; i < n_units; i = i + 1 {
for k = 0; k < dim_x; k = k + 1 {
let mut s = 0.0
let mut acc = 1.0
let mut j = k
while j >= 0 {
s = s + acc * z_flat[i * dim_x + j]
acc = acc * 0.7
if j == 0 {
break
}
j = j - 1
}
x_unit[i * dim_x + k] = s
}
}
// 3. Pre-draw time-varying noise (n_units * n_periods).
let eps : Array[Double] = Array::make(true_n, 0.0)
let half_n = (true_n + 1) / 2
for i = 0; i < half_n; i = i + 1 {
let (z1, z2) = box_muller_pair(rng)
let idx_a = 2 * i
let idx_b = 2 * i + 1
if idx_a < true_n {
eps[idx_a] = z1
}
if idx_b < true_n {
eps[idx_b] = z2
}
}
// 3b. Per-unit intercept `alpha_i`, a standard normal draw. v0.148.0.
//
// This REPLACED `alpha_i = 0.5 * unit_in_cohort`, a deterministic ramp over
// the within-cohort index reaching 99.5 with SD ~ 28.8. The ramp was not a
// small modelling choice, it dominated everything else and silently made the
// whole fixture unusable. Measured on make_did_cs_CS2021(800, 3, 1.0, 3141):
//
// var(y) = 828.42 (the ramp's share)
// var(X . gamma) ~ 4.3
// corr(X_k, y) <= 0.027 for every k -> r^2 <= 7e-4
// mse(g_hat) = 871.78
// mse(cell mean) = 829.09 -> g_hat 5.2% WORSE than a constant
//
// So E[Y | D, T, X] carried no covariate signal at all: every conditional
// nuisance was effectively a constant, `sigma2 = mean((y - g_hat)^2)`
// collapsed onto `var(y)`, and the sensitivity machinery had nothing to
// condition on. `dgp_recovery_test`'s `max_abs_err < 5.0` never bit either,
// because it was reading a zero from a cell that was never estimated.
//
// A draw is kept deliberately rather than dropping the term outright:
// unobserved unit heterogeneity IS the hidden confounding that `nu2` exists
// to price, so removing it entirely would leave the fixture unable to
// represent the thing it is for. SD 1 puts it on the same scale as the true
// treatment effect (`theta = 1.0`) and makes it ~19% of the covariate
// signal, i.e. a minority that must be partialled out rather than a
// majority that drowns everything.
//
// Upstream has no unit intercept at all -- its outcome is
// `210 + (t/T) * (27.4*W1 + 13.7*W2 + 13.7*W3 + 13.7*W4)`, driven by
// covariates. This term is a port addition, kept on the measurement above.
let alpha : Array[Double] = Array::make(n_units, 0.0)
let half_a = (n_units + 1) / 2
for i = 0; i < half_a; i = i + 1 {
let (z1, z2) = box_muller_pair(rng)
let idx_a = 2 * i
let idx_b = 2 * i + 1
if idx_a < n_units {
alpha[idx_a] = z1
}
if idx_b < n_units {
alpha[idx_b] = z2
}
}
// 4. Build panel y, d, t, g, id. Layout: outer = unit, inner = period.
let y : Array[Double] = Array::make(true_n, 0.0)
let d : Array[Double] = Array::make(true_n, 0.0)
let t : Array[Int32] = Array::make(true_n, 0)
let g : Array[Int32] = Array::make(true_n, 0)
let id : Array[Int32] = Array::make(true_n, 0)
for u = 0; u < n_units; u = u + 1 {
let cohort = u / n_per_cohort
let unit_in_cohort = u % n_per_cohort
let is_treated = cohort > 0
let tau_g = if is_treated { 1 } else { -1 } // never-treated sentinel
for period = 0; period < n_periods; period = period + 1 {
let idx = u * n_periods + period
// Treatment indicator: post-treatment for treated cohort.
d[idx] = if is_treated && period >= tau_g { 1.0 } else { 0.0 }
// Outcome: alpha_i + beta_t + theta * D + X_i @ gamma + eps
let mut x_dot_gamma = 0.0
for k = 0; k < dim_x; k = k + 1 {
x_dot_gamma = x_dot_gamma + x_unit[u * dim_x + k]
}
let alpha_i = alpha[u]
let beta_t = 0.3 * period.to_double()
y[idx] = alpha_i + beta_t + theta * d[idx] + x_dot_gamma + eps[idx]
t[idx] = period
g[idx] = cohort
id[idx] = unit_in_cohort
}
}
// Expand X to panel (each unit's X is repeated n_periods times).
let x_panel_flat : Array[Double] = Array::make(true_n * dim_x, 0.0)
for u = 0; u < n_units; u = u + 1 {
for p = 0; p < n_periods; p = p + 1 {
let panel_row = u * n_periods + p
for k = 0; k < dim_x; k = k + 1 {
x_panel_flat[panel_row * dim_x + k] = x_unit[u * dim_x + k]
}
}
}
let x_panel = Matrix::from_array(x_panel_flat, true_n, dim_x)
{ theta, n_groups, n_periods, y, d, t, g, id, x: x_panel, }
}
///|
/// Get the true ATT parameter.
pub fn DidCsData::theta_get(self : DidCsData) -> Double {
self.theta
}
///|
/// Get the outcome vector.
pub fn DidCsData::y_get(self : DidCsData) -> Array[Double] {
self.y
}
///|
/// Get the treatment vector.
pub fn DidCsData::d_get(self : DidCsData) -> Array[Double] {
self.d
}
///|
/// Get the time-period index.
pub fn DidCsData::t_get(self : DidCsData) -> Array[Int32] {
self.t
}
///|
/// Get the group (cohort) index.
pub fn DidCsData::g_get(self : DidCsData) -> Array[Int32] {
self.g
}
///|
/// Get the unit id (within cohort).
pub fn DidCsData::id_get(self : DidCsData) -> Array[Int32] {
self.id
}
///|
/// Get the covariate matrix (panel-level, n_obs x dim_x).
pub fn DidCsData::x_get(self : DidCsData) -> Matrix {
self.x
}
///|
/// Get the number of cohorts.
pub fn DidCsData::n_groups_get(self : DidCsData) -> Int32 {
self.n_groups
}
///|
/// Get the number of periods.
pub fn DidCsData::n_periods_get(self : DidCsData) -> Int32 {
self.n_periods
}