Skip to content

feat: port twfeweights/ptetools/badcontrols R packages (R-compat) - #753

Open
wenddymacro wants to merge 53 commits into
igerber:mainfrom
wenddymacro:feat/twfeweights-r-compat
Open

feat: port twfeweights/ptetools/badcontrols R packages (R-compat)#753
wenddymacro wants to merge 53 commits into
igerber:mainfrom
wenddymacro:feat/twfeweights-r-compat

Conversation

@wenddymacro

Copy link
Copy Markdown
Contributor

Summary

Ports the full public API of the R packages twfeweights, ptetools, and badcontrols into diff_diff/, with R/Python numeric parity on the implemented subset.

Scope

  • twfeweights: twfe_weights (incl. AIPW + balance diagnostics) + post-lasso block
  • ptetools: pte() main loop, did_attgt, full QTT/QoTT block, process_dose_gt (splines2-compatible bspline_basis), mboot2/mboot_se_and_crit, RCS/dose/mboot, covid_attgt, attgt_noif, aggregations, and matplotlib/Plotly plotting wrappers (ggpte, ggpte_cont, plot_qtt, autoplot/plot counterparts)
  • badcontrols: parametric + ML cross-fit DR bad-control score with shared fold ingress

R parity verified

  • twfe_weights, ptetools did_attgt, badcontrols continuous + binary imputation
  • qtt_crit_val to ~1e-13; overall QTT on single-cohort R panel to 0.0
  • bspline_basis matches splines2::bSpline/dbs exactly (golden values pinned in tests)
  • Dropped-parity surfaces documented in REGISTRY.md (implicit_twfe_weights — R fixest segfaults on fixture; did_post_lasso — R source contains a browser() debug path)

Tests

New-feature suites (tests/test_{ptetools,twfeweights,badcontrols}_*.py + test_r_parity_new_features.py): 82 passed under DIFF_DIFF_BACKEND=python. Ruff/black/mypy clean.

yiyi added 30 commits August 6, 2026 09:23
yiyi added 21 commits August 6, 2026 14:47
- process_dose_gt consumes an R-style gt_results dict + ptep options and
  returns a complete DoseResult: ATT(d)/ACRT(d) curves, per-dose multiplier-
  bootstrap SEs, pointwise/simultaneous critical values, and overall ATT/ACRT
  with SEs and influence functions.
- bspline_basis reproduces splines2::bSpline / dbs exactly (clamped boundary
  knots, intercept=False drops first basis column, derivative via the knot/
  coefficient transform); golden parity pinned against live R output.
- mboot_se_and_crit turns mboot2 draws into R-style IQR bootstrap SEs and a
  sup-t critical value using R quantile(type=1).
- DoseResult extended to the full dose_obj surface while keeping
  pte_dose_results backward-compatible; new exports + docs + CHANGELOG entry.
- tests/test_ptetools_process_dose_gt.py: splines2 golden parity, knot
  validation, end-to-end point estimates, seed reproducibility, order and
  missing-field rejection.
Port the quantile-treatment-effects machinery and extend the
twfeweights/badcontrols/ptetools R-compat layer:

- pte_qtt / PTEQTTResult, compute_pte (g,t) loop, qtt/qott aggregation,
  qtt_empirical_bootstrap super-t bands, block_boot_sample, _qtt_crit_val
- ggpte / ggpte_cont event-study and dose plotting wrappers plus
  autoplot/plot methods on PTE, QTT, emp-boot, and dose result objects
- attgt_noif container, covid_attgt DRDID levels/changes score
- PTEResults.aggregate() influence-function SEs, CIs, to_dataframe levels,
  dynamic multiplier-bootstrap bands
- dr_ml_attgt bad-controls cell wrapper, mp_weights_obj twfeweights support
- _NotSupplied copy/deepcopy for result containers
R's badcontrols::dr_ml_attgt always cross-fits the parametric nuisances
(OLS m/omega, logit p), so a full-sample Python fit was not fold-mirror
parity: different set.seed gave different R ATs but the Python call gave
one. This makes dr_parametric_bad_control cross-fit like R:

- split folds per treatment arm (treated/control each see every fold)
- m0/nu0/omega0 OLS + p2 logit fit on the training folds, evaluated on
  the held-out fold; in-sample fitted outcomes feed the nu/omega targets.
- ingress fold_ids (validated 0..n_folds-1) for exact shared-fold parity.
- keep the max(propensity)>0.99 -> imputation fallback guard from R.

New tests pin fold-dependence, fold_ids reproducibility/validation, the
imputation fallback, and gt_data from two_by_two_subset. REGISTRY gains a
Bad Controls section; grf-vs-sklearn ML nuisance documented as not
parity-able.

mypy note: env fails on numpy 2.5.1 .pyi under 3.12 target 3.10 (pre-
existing, not from this change).
…licit weights

The R parity harness pins twfe_weights/attO_weights/att_simple_weights to
1e-8. Two other twfeweights functions cannot be byte-parity-tested and are
now documented as REGISTRY deviations rather than claimed silently:

- did_post_lasso: the R reference source is incomplete (carries a
  browser() debug path); Python's LassoCV AIPW is a faithful reading of the
  paper, verified for internal consistency only.
- implicit_twfe_weights: fixest segfaults on the parity fixture, so no
  stable R reference number exists; the closed-form FWL decomposition is
  verified for self-consistency (alpha_weight*attgt reconstructs TWFE).

Adds a twfe_weights methodology section + TOC entry and a CHANGELOG note.
…1) like R

The high-level pte() wrapper built its influence surface with NaN for
off-support units and no (n/n1) sample-size correction, diverging from R's
compute.pte (ptetools/R/pte.R:137-141), which zero-pads with
rep(0, n); this.inf_func[disidx] <- (n/n1)*attgt. The lower-level
compute_pte already matched R; the wrapper did not.

pte() now:
- zero-fills off-support unit entries (not NaN)
- scales each cell influence function by (n / n1) for overall-vs-cell sizes
- keeps base-period-skip cells as a full-NA column (as both R and compute_pte do)

New test pins the (units, cells) surface: no NaN anywhere, off-support
entries zero, and placed entries equal (n/n1)*did_attgt(...).inf_func per
estimable cell. REGISTRY gains a ptetools influence-surface note.
@wenddymacro
wenddymacro requested a review from igerber as a code owner August 7, 2026 07:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant