Skip to content

Commit eba02cf

Browse files
authored
feat: DML infrastructure - cross-fitting, shared DR scores, ridge, learner protocol (PR-B0) (#789)
1 parent 2d41cbe commit eba02cf

19 files changed

Lines changed: 3276 additions & 43 deletions

CHANGELOG.md

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
88
## [3.10.0] - 2026-08-22
99

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

DEFERRED.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ provenance and AI-review deviation-documentation: a row here (or in
2020

2121
| Issue | Location | PR | Priority |
2222
|-------|----------|----|----------|
23+
| 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 |
2324
| 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 |
2425
| 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 |
2526
| 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.
5253

5354
| Issue | Location | PR | Priority |
5455
|-------|----------|----|----------|
56+
| 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 |
5557
| `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 |
5658
| 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 |
5759
| 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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ Related tracking surfaces:
2121

2222
| Issue | Location | Origin | Effort | Priority |
2323
|-------|----------|--------|--------|----------|
24+
| 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 |
2425
| 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 |
2526
| 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 |
2627
| 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 |
Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
"""DoubleML parity spike for the Chang (2020) Case 1 (repeated outcomes) score.
2+
3+
Hand-rolled DML2 cross-fitted Chang estimator vs DoubleMLDID under identical
4+
folds, learners, and clipping. This is the reproducible anchor cited by the
5+
REGISTRY "Cross-fitting, DR-score, and ridge infrastructure (DML)" section's
6+
global-p-hat Note.
7+
8+
Environment (side venv; doubleml/sklearn are NEVER diff-diff dependencies):
9+
10+
python -m venv .venv-doubleml
11+
.venv-doubleml/bin/pip install "doubleml==0.11.4" scikit-learn
12+
.venv-doubleml/bin/python benchmarks/doubleml/chang_case1_parity.py
13+
14+
Observed transcript (2026-08-22, doubleml 0.11.4, sklearn 1.9.0, macOS arm64):
15+
16+
DoubleML ATT = 3.260530717619 SE = 0.173029660048
17+
hand Chang (global p) ATT = 3.260530717619 SE = 0.173029660048
18+
diff vs DoubleML: ATT -4.441e-16 SE +5.551e-17
19+
hand Chang (fold-mean p) ATT = 3.262788670382 (finite-sample gap +2.258e-03)
20+
PARITY OK (global-p ATT and SE within 1e-10 of DoubleML)
21+
22+
The global-p-hat convention (p-hat = full-sample treated share) matches
23+
DoubleML exactly; the fold-mean convention from the paper's proofs differs
24+
only in finite samples. The SE parity uses the augmented score
25+
psi_bar_i = summand_i - D_i * theta / p_hat.
26+
"""
27+
28+
import numpy as np
29+
from doubleml import DoubleMLDID
30+
from doubleml.data import DoubleMLDIDData
31+
from sklearn.linear_model import LinearRegression, LogisticRegression
32+
33+
rng = np.random.default_rng(42)
34+
N, d = 500, 5
35+
X = rng.standard_normal((N, d))
36+
g0 = 1 / (1 + np.exp(-(X[:, 0] - 0.5 * X[:, 1])))
37+
D = (rng.uniform(size=N) < g0).astype(float)
38+
# Delta Y = ell(X) + theta*D + noise, theta = 3
39+
ell0 = X[:, 0] + 0.5 * X[:, 2] ** 2
40+
dY = ell0 + 3.0 * D + rng.standard_normal(N)
41+
42+
K = 5
43+
TRIM = 1e-2 # match DoubleML clipping default behavior
44+
45+
# Fixed fold assignment shared by both implementations.
46+
perm = rng.permutation(N)
47+
folds = [np.sort(perm[i::K]) for i in range(K)]
48+
smpls = [(np.setdiff1d(np.arange(N), te), te) for te in folds] # (train, test)
49+
50+
51+
def _nuisances(tr, te):
52+
lg = LogisticRegression(penalty=None, solver="lbfgs", max_iter=1000).fit(X[tr], D[tr])
53+
g_hat = np.clip(lg.predict_proba(X[te])[:, 1], TRIM, 1 - TRIM)
54+
ctrl = tr[D[tr] == 0] # outcome nuisance fit on untreated complement only
55+
lr = LinearRegression().fit(X[ctrl], dY[ctrl])
56+
return g_hat, lr.predict(X[te])
57+
58+
59+
# --- DoubleML reference ----------------------------------------------------
60+
data = DoubleMLDIDData.from_arrays(X, dY, D)
61+
m = DoubleMLDID(
62+
data,
63+
ml_g=LinearRegression(),
64+
ml_m=LogisticRegression(penalty=None, solver="lbfgs", max_iter=1000),
65+
n_folds=K,
66+
n_rep=1,
67+
score="observational",
68+
in_sample_normalization=False,
69+
clipping_threshold=TRIM,
70+
draw_sample_splitting=False,
71+
)
72+
m.set_sample_splitting([smpls])
73+
m.fit()
74+
att_dml, se_dml = float(m.coef[0]), float(m.se[0])
75+
76+
# --- hand-rolled Chang Case 1, DML2 cross-fitting --------------------------
77+
p_glob = D.mean() # global treated share (the library convention)
78+
theta_k_glob = np.empty(K)
79+
theta_k_fold = np.empty(K)
80+
for k, (tr, te) in enumerate(smpls):
81+
g_hat, ell_hat = _nuisances(tr, te)
82+
w_num = (D[te] - g_hat) / (1 - g_hat)
83+
resid = dY[te] - ell_hat
84+
theta_k_glob[k] = np.mean(w_num * resid / p_glob)
85+
theta_k_fold[k] = np.mean(w_num * resid / D[te].mean())
86+
theta_glob = theta_k_glob.mean()
87+
theta_fold = theta_k_fold.mean()
88+
89+
# Variance from the augmented score psi_bar = summand - D*theta/p (global p).
90+
sig2_k = np.empty(K)
91+
for k, (tr, te) in enumerate(smpls):
92+
g_hat, ell_hat = _nuisances(tr, te)
93+
summand = (D[te] - g_hat) / (p_glob * (1 - g_hat)) * (dY[te] - ell_hat)
94+
psi_bar = summand - D[te] * theta_glob / p_glob
95+
sig2_k[k] = np.mean(psi_bar**2)
96+
se_glob = np.sqrt(sig2_k.mean() / N)
97+
98+
print(f"DoubleML ATT = {att_dml:.12f} SE = {se_dml:.12f}")
99+
print(f"hand Chang (global p) ATT = {theta_glob:.12f} SE = {se_glob:.12f}")
100+
print(f" diff vs DoubleML: ATT {theta_glob - att_dml:+.3e} SE {se_glob - se_dml:+.3e}")
101+
print(
102+
f"hand Chang (fold-mean p) ATT = {theta_fold:.12f} "
103+
f"(finite-sample gap {theta_fold - theta_glob:+.3e})"
104+
)
105+
106+
assert abs(theta_glob - att_dml) < 1e-10, "global-p ATT parity broken"
107+
assert abs(se_glob - se_dml) < 1e-10, "global-p SE parity broken"
108+
print("PARITY OK (global-p ATT and SE within 1e-10 of DoubleML)")

0 commit comments

Comments
 (0)