feat: port twfeweights/ptetools/badcontrols R packages (R-compat) - #753
Open
wenddymacro wants to merge 53 commits into
Open
feat: port twfeweights/ptetools/badcontrols R packages (R-compat)#753wenddymacro wants to merge 53 commits into
wenddymacro wants to merge 53 commits into
Conversation
added 30 commits
August 6, 2026 09:23
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Ports the full public API of the R packages
twfeweights,ptetools, andbadcontrolsintodiff_diff/, with R/Python numeric parity on the implemented subset.Scope
twfe_weights(incl. AIPW + balance diagnostics) + post-lasso blockpte()main loop,did_attgt, full QTT/QoTT block,process_dose_gt(splines2-compatiblebspline_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/plotcounterparts)R parity verified
twfe_weights,ptetools did_attgt,badcontrolscontinuous + binary imputationqtt_crit_valto ~1e-13; overall QTT on single-cohort R panel to 0.0bspline_basismatchessplines2::bSpline/dbsexactly (golden values pinned in tests)implicit_twfe_weights— R fixest segfaults on fixture;did_post_lasso— R source contains abrowser()debug path)Tests
New-feature suites (
tests/test_{ptetools,twfeweights,badcontrols}_*.py+test_r_parity_new_features.py): 82 passed underDIFF_DIFF_BACKEND=python. Ruff/black/mypy clean.