feat: TWFE weight diagnostics (twfeweights R port, re-scoped from #753) - #812
Open
wenddymacro wants to merge 5 commits into
Open
feat: TWFE weight diagnostics (twfeweights R port, re-scoped from #753)#812wenddymacro wants to merge 5 commits into
wenddymacro wants to merge 5 commits into
Conversation
Generator + committed goldens for the upcoming `attgt_weights` / `decompose_twfe_weights` surface, ported from Brantly Callaway's `twfeweights` R package (MIT). Three fixtures: `mpdta` (real; non-1..T time labels), `sim_staggered` (equal cohorts, real pre-trend so pretrend_bias != 0), and `unbalanced_cohorts` (120/70/60 — breaks the p_g == 1/3 degeneracy that would let a cohort-share bug pass silently on the equal-cohort fixture). Pins `twfe_weights`/`attO_weights`/`att_simple_weights`, `implicit_twfe_weights` (no-cov, covariate, gmin1), `implicit_aipw_weights`, `twfe_cov_bal`/`aipw_cov_bal` + the summary roll-up, and the two two-period kernels that will have no public Python surface. The no-covariate decomposition is generated with a TIME-INVARIANT covariate rather than `xformula = ~1`: upstream builds an nT x 0 model matrix on that branch and `fixest::demean` segfaults on a zero-column matrix (reproduced in isolation, fixest 0.14.2 / R 4.6.1). Double-demeaning annihilates a time-invariant regressor exactly, so the call is numerically the `~1` branch — and the Python test will assert both `covariates=None` and `covariates=[<col>]` against this one golden, proving the equivalence rather than assuming it. R is only needed to regenerate the JSON, never to run the tests. Co-Authored-By: Claude <noreply@anthropic.com>
…iners Result containers for the incoming TWFE implicit-weight diagnostics, landed ahead of the compute module so they pin the output schema and the Diagnostic contract before any math depends on them. Both subclass Diagnostic: they assess what a regression implicitly weights rather than estimating an effect, so neither carries the estimator quintet. The headline scalars are deliberately named `implied_att` and `estimate` rather than `att` so they do not read as inference-bearing. Output columns are diff-diff's (`group`, `time`, `post`, `weight`, `att`), not R's (`time.period`, `attgt`). Covariate balance is a result-object method rather than a mutate-in-place second pass as in R: `covariate_balance(level="summary"|"cell")` reads a table computed at construction time, so the result never retains the raw panel. It raises with the fix inlined when balance was not requested. Names are clearly separated from the existing dCDH surface (`twowayfeweights` / `TWFEWeightsResult`), which weights (unit, time) cells; these weight ATT(g,t) parameters. Roster (M-091) and the shared construction fixture updated; the roster test auto-enrolls both classes. Co-Authored-By: Claude <noreply@anthropic.com>
One entry point folding R's three separate weight functions behind
`aggregation=`: "twfe" (twfe_weights), "overall" (attO_weights, ATT^O), and
"simple" (att_simple_weights, ATT^simple). Reports what each estimand
implicitly puts on every group-time effect, plus the negative-weight share
that makes the staggered-TWFE pathology legible.
Takes a fitted CallawaySantAnnaResults as the primary input, reading cohort
masses off the aggregation bookkeeping so no raw panel is needed; a
(gt_frame, data=, unit=, time=, first_treat=) fallback consumes
`result.to_dataframe("group_time")` verbatim.
Design restrictions are hard errors, not warnings, each naming its fix:
aggregation="twfe" needs base_period="universal" and
control_group="never_treated" (matching R's own stop()s), and no aggregation
accepts a repeated-cross-section or unbalanced-fallback fit, whose cohort
shares are not comparable across periods.
Deviation from R: cohorts and periods are mapped to positional time before
the (maxT - g + 1)/T arithmetic. R evaluates that on raw labels, which is
only correct on consecutive integers; positional time is bit-identical there
(mpdta 2003..2007 -> 1..5 both give 4/5 at g=2004) and correct on gapped
grids. Pinned by a test that remaps periods to 10,20,30,40,50.
R's keep_untreated= is not exposed: it synthesizes G=0 rows that are
excluded from every normalization and contribute exactly zero.
Parity: machine precision (max |dw| = 4.7e-16) against R twfeweights 0.9.0
on 3 fixtures x 3 aggregations. Primary assertions feed R's own ATT(g,t)
back in, isolating this module from CallawaySantAnna-vs-`did` parity; a
separate, deliberately looser class covers the composed end-to-end path.
Also fixes a generator bug: R stores `post` as a FACTOR, so as.integer()
emitted level codes 1/2 rather than 0/1.
Known-red until the docs commit: test_doc_deps_integrity.py wants a
docs/doc-deps.yaml entry, which lands with the API page.
Co-Authored-By: Claude <noreply@anthropic.com>
…ance
Re-derives a TWFE estimate from its ATT(g,t) building blocks, reporting the
implicit weight on each cell, the contribution of PRE-treatment cells
(`pretrend_bias` - parallel-trends violations rather than treatment), and
implicit-weight covariate balance via `result.covariate_balance()`.
Takes the raw panel rather than a fitted CS result because it re-estimates:
it double-demeans treatment and covariates and forms its own group-time
contrasts, so there is no ATT(g,t) table it could consume, and a CS result
carries no panel by design. The two surfaces are tied by an identity that
the suite pins:
attgt_weights(cs, aggregation="twfe").implied_att
== decompose_twfe_weights(panel, ...).estimate
Two numerical points, both found by disagreeing with the goldens and then
proving which side was right:
1. Covariates that double-demeaning ANNIHILATES are now dropped before the
projection, judged against each column's own PRE-demeaning norm. A
time-invariant regressor leaves a column of pure rounding noise (~1e-16
against a raw scale of ~1); regressing on it amplifies that by ~1e16 and
silently corrupted the per-cell weights. A rank test on the demeaned
matrix alone cannot see this - there, 1e-16 is simply the largest pivot.
With the fix, covariates=None and covariates=[<time-invariant col>] agree
to 1e-15 on every fixture, which is exactly the equivalence the
no-covariate golden relies on.
2. Cells whose comparison-group implicit weights are constant AND average to
zero make `resid / mean(resid)` a 0/0. On sim_staggered (equal cohorts at
g in {0,3,4}, T=5) this happens exactly at t=3, where
-E_3[D] + mean_t E_t[D] = -1/3 + 1/3. We take the limit (a constant over
its own mean is one); R divides the rounding errors and lands ~3e-4 away.
Verified against a hand-computed contrast that needs none of this module:
ours is exact to 4.4e-16. The weights on such cells cancel exactly in the
aggregate, so `estimate` is unaffected - the suite gates `estimate`
tightly on every fixture and relaxes only the per-cell and
decomposition/remainder-split assertions, on cells DETECTED as degenerate
rather than on a hard-coded fixture.
Parity vs R twfeweights 0.9.0: estimate and per-cell weights at machine
precision on all 3 fixtures x 4 configurations; all 11 balance statistics at
machine precision, including `frac_treated_extreme`, which required
reproducing BMisc's weighted-ECDF plus `stats:::quantile.ecdf`'s
pseudo-sample reconstruction rather than a plain quantile.
method="aipw" is not implemented yet and raises listing the accepted values.
Co-Authored-By: Claude <noreply@anthropic.com>
Plotting (replacing upstream's ggtwfeweights S3 methods) and every documentation surface the new API owes. plot_twfe_weights(result, kind="auto"|"weights"|"balance") lives beside plot_bacon in visualization/_diagnostic.py and dispatches on either result type. The weights view puts weight on x and ATT(g,t) on y with zero lines, so negative-weight cells sit visibly left of the axis; the balance view plots unweighted against implicitly-weighted covariate differences with a no-improvement diagonal. "auto" picks balance when a balance table is present. Docs: a REGISTRY.md section carrying the weight equations, the cross-surface identity, the tolerance table with per-gate rationale, and eleven explicit Note/Deviation-from-R entries - including the fixest zero-column segfault and its root cause, the annihilated-covariate drop, and the 0/0-cell limit, so the two places we deliberately differ from R are recorded rather than discovered later by a reviewer. Two paragraphs separate this surface from `twowayfeweights` (dCDH, weights (unit, time) cells) and from BaconDecomposition (decomposes into 2x2 comparisons), since all three are "TWFE weight" diagnostics and the distinction is the thing a reader most needs. Also: docs/api/twfe_weights.rst with runnable examples, four api/index.rst registrations (2 result classes, the plot, 2 functions, toctree), doc-deps.yaml group + sources entries, a README one-liner in Diagnostics & Sensitivity, llms.txt catalog entry, llms-full.txt API + result blocks, a references.rst sub-entry naming the upstream package and its MIT copyright, and a changelog.d fragment. This closes the doc-deps gate the attgt_weights commit left red. Verified: 14378 tests collect clean; docs IA, doc-deps integrity, diagnostic roster, guides, changelog-fragment, serialization and all visualization suites green (903 passed, 43 skipped). Co-Authored-By: Claude <noreply@anthropic.com>
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
1 participant
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.
Re-scoped follow-up to #753, branched off current
main. Onlytwfeweights(MIT, © 2023 Brantly Callaway);
ptetoolsandbadcontrolsare droppedentirely per your license read.
Against your six requirements
1. Post-lasso dropped.
did_post_lasso/did_post_lasso_ra/PostLassoResultare not ported —R/did_post_lasso.R:69has the leftoverbrowser()and references undefined variables, so there is no runnablereference. No sklearn dependency anywhere. No
quality_reports/.2. API consolidated to 5 public symbols (from 21 upstream exports):
attgt_weights(results, aggregation="twfe"|"overall"|"simple")twfe_weights,attO_weights,att_simple_weightsdecompose_twfe_weights(data, ..., method="fwl")implicit_twfe_weightsATTGTWeightsResult,TWFEDecompositionResultplot_twfe_weights()ggtwfeweightsS3 methodsCovariate balance is a result-object method —
result.covariate_balance(level="summary"|"cell")— rather than R's mutate-in-place second pass, so the result never retains the
raw panel. Everything else is private; the two-period kernels and the four
balance statistics are pinned directly by the parity suite since they have no
public surface.
3. House conventions.
attgt_weightstakes a fittedCallawaySantAnnaResultsas the primary input, reading cohort masses off itsaggregation bookkeeping so no raw panel is needed; the fallback consumes
result.to_dataframe("group_time")verbatim. Params areoutcome/unit/time/first_treat. Both results subclassDiagnosticwithsummary()/to_dict()/to_dataframe()and carry no inference quintet — theheadline scalars are named
implied_attandestimatedeliberately. Outputcolumns are ours (
group,time,post,weight,att), nevertime.period/attgt. Plotting isplot_twfe_weights()indiff_diff/visualization/.4. Output parity, and the fixest waiver is gone. I found the root cause:
fixest::demean()segfaults on a zero-column matrix, andxformula = ~1is the only branch that builds one (
model.matrix(~-1, data)→nT × 0).Reproduced in isolation on R 4.6.1 / fixest 0.14.2. It is not a property of any
fixture.
The no-covariate golden is therefore generated with a time-invariant
covariate — double-demeaning annihilates it exactly, so the call is
numerically the
~1branch — and the parity test asserts bothcovariates=Noneandcovariates=[<that column>]against that one golden, sothe equivalence is proven rather than assumed. On
mpdta,twfe_weights(att_gt(...))andimplicit_twfe_weights(xformula = ~lpop)$estboth give
-0.03654894.Goldens:
benchmarks/data/twfeweights_golden.json+ three sibling panel CSVs,regenerated by
benchmarks/R/generate_twfeweights_golden.R. R is never neededto run the tests; they
pytest.skipif a fixture file is absent. Tolerancesare module constants with per-gate rationale, tabulated in REGISTRY.
Three fixtures:
mpdta(real, non-1..Tlabels),sim_staggered(equalcohorts, a real pre-trend so
pretrend_bias != 0), andunbalanced_cohorts(120/70/60 — breaks the
p_g == 1/3degeneracy that would let a cohort-sharebug pass silently on the equal-cohort fixture).
Results: ATT(g,t) weights match at machine precision (max 4.7e-16) on all
3 fixtures × 3 aggregations; FWL
estimateand cell weights likewise; all 11balance statistics at machine precision.
5. Docs. REGISTRY section (equations, cross-surface identity, tolerance
table, 11 Note/Deviation entries),
docs/api/twfe_weights.rst, 4api/index.rstregistrations,doc-deps.yaml,references.rst, one READMEline,
llms.txt+llms-full.txt, changelog fragment. MIT attribution: theupstream copyright notice is reproduced verbatim in the module docstring.
6. Branched off current
main.Two places I deliberately differ from R (both in REGISTRY)
Both were found by disagreeing with the goldens and then working out which
side was right.
Annihilated covariates are dropped before the projection. A time-invariant
regressor leaves a column of pure rounding noise after double-demeaning
(~1e-16 against a raw scale of ~1); regressing on it amplifies that by ~1e16
and silently corrupts the per-cell weights. The test is scale-relative
(demeaned norm vs the column's own raw norm) because a rank test on the
demeaned matrix alone cannot see it — there, 1e-16 is simply the largest
pivot. With this,
covariates=Noneandcovariates=[<time-invariant col>]agree to 1e-15.
Cells with a 0/0 normalizer report the limit, not the noise. For the
never-treated group the double-demeaned treatment is constant within a period
(
-E_t[D] + mean_t E_t[D]), and onsim_staggeredthat constant isanalytically zero at
t=3(-1/3 + 1/3). We take the limit (a constantover its own mean is one); R divides the two rounding errors and lands ~3e-4
away. Checked against a hand-computed contrast that uses none of this module:
ours is exact to 4.4e-16. The aggregate is unaffected either way — the
weights on those cells cancel exactly (
w(3,3) + w(4,3) = 0), which is whyestimatestill matches R to 1e-15. The parity suite gatesestimatetightlyeverywhere and relaxes only the per-cell assertions, on cells detected as
degenerate rather than hard-coded to a fixture.
Naming — your call
I need names distinct from the existing
twowayfeweights/TWFEWeightsResult(dCDH) surface, which weights (unit, time) cells where these weight
ATT(g,t) parameters. Implemented as set A; B and C are mechanical renames
if you prefer one:
attgt_weights/ATTGTWeightsResult/decompose_twfe_weights/TWFEDecompositionResult. Puts the object beingweighted in the head-noun position; shares no token with
twowayfeweightsand diverges at character 1 for autocomplete.
estimand_weights/EstimandWeightsResult/twfe_weight_decomposition/TWFEWeightDecompositionResult. Generalizesbetter if an
aggregation="event_study"is added later; vaguer about whatthe rows are.
callaway_twfe_weights/ … — author-prefixed. I did not use it sincethe library has
BaconDecomposition, notgoodman_bacon_decompose, and itreads confusingly next to
CallawaySantAnna.plot_twfe_weightsis your name, kept verbatim.Scope note
method=accepts"fwl"only. Upstream'simplicit_aipw_weightsis left fora follow-up rather than growing this PR — the old port's AIPW was numerically
wrong against R (score form vs R's
gamma0/gamma0_tildeAIPW-weights form),so it needs a from-scratch implementation and its own parity work. The goldens
already carry the AIPW references, so that PR will not need R re-run. Recorded
in REGISTRY as a documented deviation, not a silent gap.
Verification
ruff/blackclean,mypy diff_diffat zero errors, 14378 tests collectclean. Green locally: the new suites (
test_twfe_weights.py,test_twfe_weights_parity.py), plus docs IA, doc-deps integrity, diagnosticroster (M-091), guides, changelog fragments, serialization and all
visualization suites — 903 passed / 43 skipped. I did not complete a full
local
pytestrun (pure-Python mode, no Rust backend built locally); happy toadd
ready-for-ciwhenever you want CI on it.🤖 Generated with Claude Code