feat: DML infrastructure - cross-fitting, shared DR scores, ridge, learner protocol (PR-B0)

CHANGELOG.md

@@ -8,6 +8,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+- **Internal DML infrastructure for the upcoming `DMLDiD` estimator (zero
+  behavior change).** Private modules `diff_diff/_crossfit.py` (replayable
+  unit/cluster-level K-fold assignment + out-of-fold nuisance prediction,
+  `BootstrapReplaySpec`-style by-value state capture), `diff_diff/_learners.py`
+  (duck-typed `RegressorLearner`/`ClassifierLearner` protocol + native
+  linear/ridge/logit learners), and `diff_diff/_dr_scores.py` (shared DR panel
+  scores). `ContinuousDiD`'s DRDID-parity influence function relocated verbatim
+  to `_dr_scores.drdid_panel_inf_func`, pinned by committed two-tier oracles
+  (`tests/test_dr_scores.py`); the Chang (2020) Neyman-orthogonal Case 1 score
+  pair ships alongside, anchored to DoubleML at machine precision by the
+  committed parity spike (`benchmarks/doubleml/chang_case1_parity.py`). New
+  `linalg.solve_ridge` (unpenalized-intercept ridge, glmnet-style standardized
+  penalty, closed-form frozen-preprocessing LOOCV, `p >= n` support). See the
+  new REGISTRY section "Cross-fitting, DR-score, and ridge infrastructure
+  (DML)".
 - **LPDiD non-absorbing SEs now anchored against the authors' Stata `lpdid` package
   end-to-end (no library behavior change).** New committed-golden arm
   (`benchmarks/stata/generate_lpdid_nonabsorbing_golden.do` →

DEFERRED.md

@@ -20,6 +20,7 @@ provenance and AI-review deviation-documentation: a row here (or in
 
 | Issue | Location | PR | Priority |
 |-------|----------|----|----------|
+| Chang (2020) Case 3 (multilevel treatment intensity) DML estimator: the paper prints no Case 3 regularity assumptions or variance theorem (Theorems 2/4 cover Cases 1-2 only; "same arguments" p. 13), the printed multilevel algorithm carries typos (review Gaps items 5-6), and no oracle exists (`DoubleMLDIDMulti` is staggered TIMING, not intensity) — needs the variance derivation + overlap assumptions before implementation | `docs/methodology/papers/chang-2020-review.md` | dml-b0 | Low |
 | LWDiD PSM matching variance: implement the Abadie-Imbens (2006) matching variance (matched-control reuse + first-stage matching uncertainty) so `estimation_method='psm'` can report valid inference instead of the current fail-closed NaN tuple (point retained; naive var(diffs)/n was invalid under with-replacement reuse) | `diff_diff/lwdid.py` | #588 | Low |
 | LWDiD cohort-relative sensitivity exclusions: `robustness_pre_periods` / `sensitivity_no_anticipation` currently reject multi-cohort staggered inputs because their exclusion windows are defined relative to the EARLIEST adoption (later cohorts' own pre periods fall inside the global post window and survive every restriction, mislabeling the specification). Supporting staggered inputs needs per-cohort window semantics (exclude the last k periods of each cohort's own `t < g` window before its transformation), which the current row-subset design cannot express — a per-cohort masking derivation + its aggregation contract | `diff_diff/lwdid_sensitivity.py` | #588 | Low |
 | LWDiD estimation-time anticipation offset (CS-parity gap): no `anticipation=` parameter exists — no-anticipation is hard-assumed and the only surface is the single-cohort `sensitivity_no_anticipation` refit diagnostic. The LW device for anticipation is the anchor-window shift — drop the last k pre-treatment periods from the transformation/anchor window (`Ybar_{i,S0}` with `S0 < S-1`, per the REGISTRY LWDiD Anticipation edge case; NOT a cohort recode g → g − k) — and exposing an `anticipation=` parameter needs that derivation checked against LW 2025 NAC/CNAS before implementation | `diff_diff/lwdid.py` | LWDiD release audit | Low |
@@ -52,6 +53,7 @@ exists but parity can't be verified without a local toolchain.
 
 | Issue | Location | PR | Priority |
 |-------|----------|----|----------|
+| Chang (2020) published-version cross-check of the arXiv `p_hat_k` typo: the library adopted the global treated-share convention on DoubleML parity evidence (REGISTRY DML-infrastructure Note; spike `benchmarks/doubleml/chang_case1_parity.py`), but the published *Econometrics Journal* text (paywalled) has not been checked against the arXiv printing to see whether it resolves the algorithm-vs-proof contradiction | `docs/methodology/papers/chang-2020-review.md` | dml-b0 | Low |
 | `StaggeredTripleDifference` R cross-validation (the engine is shared with `TripleDifference`'s staggered mode since 3(b), so this covers both surfaces): CSV fixtures not committed (gitignored); tests skip without local R + `triplediff`. Commit fixtures or generate deterministically. | `tests/test_methodology_staggered_triple_diff.py` | #245 | Medium |
 | Staggered DDD R parity (both surfaces - one shared engine since 3(b)): benchmark only tests the no-covariate path (`xformla=~1`). Add covariate-adjusted scenarios + aggregation-SE parity assertions. | `benchmarks/R/benchmark_staggered_triplediff.R` | #245 | Medium |
 | Staggered DDD per-cohort group-effect SEs include WIF (both surfaces - one shared engine since 3(b)) (conservative vs R's `wif=NULL`); documented in REGISTRY. Could override the mixin for an exact R match (verification needs R `triplediff`). | `_staggered_triple_diff_engine.py` | #245 | Low |

TODO.md

@@ -21,6 +21,7 @@ Related tracking surfaces:
 
 | Issue | Location | Origin | Effort | Priority |
 |-------|----------|--------|--------|----------|
+| Consolidate the remaining estimator-entangled DR/logit score variants (`staggered.py::_doubly_robust` + RC twins, `triple_diff.py`, `lwdid.py`, `wooldridge.py`) onto the shared `_dr_scores.py` module, each migration with its own committed oracle capture (the ContinuousDiD lift's two-tier pattern in `tests/test_dr_scores.py`); and add a ridge vcov path to `solve_ridge` if an estimator ever needs analytical ridge inference | `diff_diff/_dr_scores.py` | dml-b0 | Mid | Low |
 | hc2/hc2_bm floor `1 - h_ii` at 1e-10 in the shared leverage meat, fabricating finite (if inflated) variances for leverage-one observations - hc3 now fails closed there (LWDiD fix wave) but the pre-existing hc2 family behavior is released surface; decide fail-closed vs keep-floor for hc2/hc2_bm | `diff_diff/linalg.py` | #588 | Quick | Low |
 | Numeric between-period cohorts (e.g. `first_treat=4.5` with integer times) are rejected by LWDiD while CallawaySantAnna estimates them and LWDiD's own datetime/Period cohorts map to the next observed period — close the dtype asymmetry by adopting the next-observed-period mapping for numeric cohorts too (contract documented in REGISTRY cohort-encodings Note + `docs/api/lwdid.rst` Input Contract). Lands only after PR #588 merges | `diff_diff/lwdid.py` | #588 | Quick | Low |
 | Implement the LW 2026 eq. 7.9/7.10 unit-average cohort estimand (regress per-unit post-average transformed outcomes on `[1, D_g]` vs never-treated) as an alternative to the documented cell-mass `cohort_effects` convention (REGISTRY within-cohort aggregation Note; the two differ on unbalanced panels, where cell-mass weights units by observed post periods). Needs the 7.10 regression + its covariance on the NT path. Lands only after PR #588 merges | `diff_diff/lwdid_staggered.py` | #588 | Quick | Low |

benchmarks/doubleml/chang_case1_parity.py

@@ -0,0 +1,108 @@
+"""DoubleML parity spike for the Chang (2020) Case 1 (repeated outcomes) score.
+
+Hand-rolled DML2 cross-fitted Chang estimator vs DoubleMLDID under identical
+folds, learners, and clipping. This is the reproducible anchor cited by the
+REGISTRY "Cross-fitting, DR-score, and ridge infrastructure (DML)" section's
+global-p-hat Note.
+
+Environment (side venv; doubleml/sklearn are NEVER diff-diff dependencies):
+
+    python -m venv .venv-doubleml
+    .venv-doubleml/bin/pip install "doubleml==0.11.4" scikit-learn
+    .venv-doubleml/bin/python benchmarks/doubleml/chang_case1_parity.py
+
+Observed transcript (2026-08-22, doubleml 0.11.4, sklearn 1.9.0, macOS arm64):
+
+    DoubleML              ATT = 3.260530717619   SE = 0.173029660048
+    hand Chang (global p) ATT = 3.260530717619   SE = 0.173029660048
+      diff vs DoubleML:   ATT -4.441e-16   SE +5.551e-17
+    hand Chang (fold-mean p) ATT = 3.262788670382   (finite-sample gap +2.258e-03)
+    PARITY OK (global-p ATT and SE within 1e-10 of DoubleML)
+
+The global-p-hat convention (p-hat = full-sample treated share) matches
+DoubleML exactly; the fold-mean convention from the paper's proofs differs
+only in finite samples. The SE parity uses the augmented score
+psi_bar_i = summand_i - D_i * theta / p_hat.
+"""
+
+import numpy as np
+from doubleml import DoubleMLDID
+from doubleml.data import DoubleMLDIDData
+from sklearn.linear_model import LinearRegression, LogisticRegression
+
+rng = np.random.default_rng(42)
+N, d = 500, 5
+X = rng.standard_normal((N, d))
+g0 = 1 / (1 + np.exp(-(X[:, 0] - 0.5 * X[:, 1])))
+D = (rng.uniform(size=N) < g0).astype(float)
+# Delta Y = ell(X) + theta*D + noise, theta = 3
+ell0 = X[:, 0] + 0.5 * X[:, 2] ** 2
+dY = ell0 + 3.0 * D + rng.standard_normal(N)
+
+K = 5
+TRIM = 1e-2  # match DoubleML clipping default behavior
+
+# Fixed fold assignment shared by both implementations.
+perm = rng.permutation(N)
+folds = [np.sort(perm[i::K]) for i in range(K)]
+smpls = [(np.setdiff1d(np.arange(N), te), te) for te in folds]  # (train, test)
+
+
+def _nuisances(tr, te):
+    lg = LogisticRegression(penalty=None, solver="lbfgs", max_iter=1000).fit(X[tr], D[tr])
+    g_hat = np.clip(lg.predict_proba(X[te])[:, 1], TRIM, 1 - TRIM)
+    ctrl = tr[D[tr] == 0]  # outcome nuisance fit on untreated complement only
+    lr = LinearRegression().fit(X[ctrl], dY[ctrl])
+    return g_hat, lr.predict(X[te])
+
+
+# --- DoubleML reference ----------------------------------------------------
+data = DoubleMLDIDData.from_arrays(X, dY, D)
+m = DoubleMLDID(
+    data,
+    ml_g=LinearRegression(),
+    ml_m=LogisticRegression(penalty=None, solver="lbfgs", max_iter=1000),
+    n_folds=K,
+    n_rep=1,
+    score="observational",
+    in_sample_normalization=False,
+    clipping_threshold=TRIM,
+    draw_sample_splitting=False,
+)
+m.set_sample_splitting([smpls])
+m.fit()
+att_dml, se_dml = float(m.coef[0]), float(m.se[0])
+
+# --- hand-rolled Chang Case 1, DML2 cross-fitting --------------------------
+p_glob = D.mean()  # global treated share (the library convention)
+theta_k_glob = np.empty(K)
+theta_k_fold = np.empty(K)
+for k, (tr, te) in enumerate(smpls):
+    g_hat, ell_hat = _nuisances(tr, te)
+    w_num = (D[te] - g_hat) / (1 - g_hat)
+    resid = dY[te] - ell_hat
+    theta_k_glob[k] = np.mean(w_num * resid / p_glob)
+    theta_k_fold[k] = np.mean(w_num * resid / D[te].mean())
+theta_glob = theta_k_glob.mean()
+theta_fold = theta_k_fold.mean()
+
+# Variance from the augmented score psi_bar = summand - D*theta/p (global p).
+sig2_k = np.empty(K)
+for k, (tr, te) in enumerate(smpls):
+    g_hat, ell_hat = _nuisances(tr, te)
+    summand = (D[te] - g_hat) / (p_glob * (1 - g_hat)) * (dY[te] - ell_hat)
+    psi_bar = summand - D[te] * theta_glob / p_glob
+    sig2_k[k] = np.mean(psi_bar**2)
+se_glob = np.sqrt(sig2_k.mean() / N)
+
+print(f"DoubleML              ATT = {att_dml:.12f}   SE = {se_dml:.12f}")
+print(f"hand Chang (global p) ATT = {theta_glob:.12f}   SE = {se_glob:.12f}")
+print(f"  diff vs DoubleML:   ATT {theta_glob - att_dml:+.3e}   SE {se_glob - se_dml:+.3e}")
+print(
+    f"hand Chang (fold-mean p) ATT = {theta_fold:.12f}   "
+    f"(finite-sample gap {theta_fold - theta_glob:+.3e})"
+)
+
+assert abs(theta_glob - att_dml) < 1e-10, "global-p ATT parity broken"
+assert abs(se_glob - se_dml) < 1e-10, "global-p SE parity broken"
+print("PARITY OK (global-p ATT and SE within 1e-10 of DoubleML)")