diff --git a/.github/workflows/notebooks.yml b/.github/workflows/notebooks.yml index c28fc9c9..369fead1 100644 --- a/.github/workflows/notebooks.yml +++ b/.github/workflows/notebooks.yml @@ -8,6 +8,8 @@ on: - 'diff_diff/**' - 'pyproject.toml' - '.github/workflows/notebooks.yml' + # the interop drift guard runs only in this workflow (balance installed) + - 'tests/test_t26_composition_drift_calibration_drift.py' pull_request: branches: [main] types: [opened, synchronize, reopened, labeled, unlabeled] @@ -16,6 +18,8 @@ on: - 'diff_diff/**' - 'pyproject.toml' - '.github/workflows/notebooks.yml' + # the interop drift guard runs only in this workflow (balance installed) + - 'tests/test_t26_composition_drift_calibration_drift.py' schedule: # Weekly Sunday 6am UTC — smoke test that notebooks still execute cleanly - cron: '0 6 * * 0' @@ -58,11 +62,16 @@ jobs: --nbmake-timeout=600 \ --ignore=docs/tutorials/06_power_analysis.ipynb \ --ignore=docs/tutorials/10_trop.ipynb \ + --ignore=docs/tutorials/26_composition_drift_calibration.ipynb \ -v \ --tb=short # Excluded notebooks (too slow for pure-Python CI without Rust backend): # 06_power_analysis — SyntheticDiD simulate_power Monte Carlo (>600s) # 10_trop — LOOCV grid search (>600s) + # Excluded notebooks (external interop dependency): + # 26_composition_drift_calibration — requires the balance package; + # runs in the isolated interop-notebooks job below so this job's + # minimal env keeps enforcing that tutorials add no dependencies - name: Upload failed notebook outputs if: failure() @@ -71,3 +80,59 @@ jobs: name: failed-notebook-outputs path: docs/tutorials/*.ipynb retention-days: 7 + + interop-notebooks: + name: Execute balance-interop notebook + # Same ready-for-ci label gate as execute-notebooks (keep in sync). + if: >- + github.event_name != 'pull_request' + || (contains(github.event.pull_request.labels.*.name, 'ready-for-ci') + && (github.event.action != 'labeled' && github.event.action != 'unlabeled' + || github.event.label.name == 'ready-for-ci')) + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + + - name: Set up Python + uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6 + with: + # 3.12+: balance's legacy numpy<2 / scipy<1.14 / scikit-learn<1.4 + # pins apply only to python < 3.12 + python-version: '3.12' + + - name: Install dependencies + # balance is a tutorial-only dependency: it is installed ONLY in this + # isolated job, never in package requirements or the main notebooks + # job. The weekly cron on this workflow doubles as a cross-package + # integration smoke of diff-diff HEAD against latest PyPI balance. + run: | + pip install numpy pandas scipy matplotlib nbmake pytest ipykernel "balance>=0.21" + # Add repo root to Python path so Jupyter kernels can import diff_diff + # (pip install -e . requires the Rust/maturin toolchain; .pth avoids that) + python -c "import site; print(site.getsitepackages()[0])" | xargs -I{} sh -c 'echo "$PWD" > {}/diff_diff_dev.pth' + + - name: Execute interop notebook + env: + DIFF_DIFF_BACKEND: python + run: | + pytest --nbmake docs/tutorials/26_composition_drift_calibration.ipynb \ + --nbmake-timeout=600 \ + -v \ + --tb=short + + - name: Run interop drift guard + # balance is present only in this job, so this is the drift test's + # CI home (it importorskips balance everywhere else). + env: + DIFF_DIFF_BACKEND: python + run: | + pytest tests/test_t26_composition_drift_calibration_drift.py -v --tb=short + + - name: Upload failed notebook outputs + if: failure() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7 + with: + name: failed-interop-notebook-outputs + path: docs/tutorials/26_composition_drift_calibration.ipynb + retention-days: 7 diff --git a/CHANGELOG.md b/CHANGELOG.md index b3bba70b..c41a7bb1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,37 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [3.6.2] - 2026-07-03 ### Added +- **balance interop launch: composition-drift tutorial + `interop-notebooks` CI job.** Meta's + `balance` package (>= 0.21) ships a one-way adapter `balance.interop.diff_diff` + (facebookresearch/balance PR #465) whose `balance[did]` extra pins `diff-diff>=3.3,<4`. + `docs/tutorials/26_composition_drift_calibration.ipynb` is the diff-diff-side companion to + balance's `balance_diff_diff_brfss` tutorial, telling the failure-mode half of the story: a + BRFSS-style smoking-ban DGP with no systematic arm-specific trends (parallel trends hold in + expectation; planted ATT -3.0pp, realized -2.98pp under a rarely-binding probability floor) where + treatment-correlated non-response drift biases the design-weight Callaway-Sant'Anna ATT to + ~-4.1pp with *clean pre-trends*; a per-wave national rake fails (~-4.4pp - margins satisfied in + aggregate while arm-level composition is untouched); per-state raking with balance (BRFSS's own + granularity, population-count totals) recovers ~-3.2pp. Also covers the seam both ways (native + `SurveyDesign` + `aggregate_survey` vs `bd.to_panel_for_did`/`bd.fit_did`, exact-parity assert), + a 3-estimator x 2-weighting sweep, and `as_balance_diagnostic` cross-package diagnostics. + `tests/test_t26_composition_drift_calibration_drift.py` re-derives every quoted number + (auto-skips without balance). balance stays out of package requirements: the tutorial runs in a + new isolated `interop-notebooks` job in `notebooks.yml` (python 3.12, installs + `balance>=0.21`, also the drift guard's CI home; the workflow's weekly cron doubles as a + cross-package integration smoke against latest PyPI balance), and the main notebooks job env is + unchanged. +- **`balance.interop.diff_diff` contract tests.** `tests/test_balance_interop_contract.py` pins + the diff-diff surface Meta's balance adapter consumes, importing no balance code: + `aggregate_survey` forwarded-params superset + `(panel, SurveyDesign)` return schema, the + `SurveyDesign` 15-field dataclass contract (plus TSL / replicate constructions), estimator and + short-alias resolution (`CS`/`DiD`/`BJS`/`HAD`) with `survey_design=` accepted by all 17 + promised `fit()` signatures, the `_balance_adjustment` setattr provenance side-channel + (guards against future `__slots__`), the CallawaySantAnna pweight-only guard, and the + `SurveyMetadata.design_effect`/`effective_n`/`sum_weights` attribute names read by + `as_balance_diagnostic`. Docs handoff closing the survey-roadmap Phase 8g gap: "Weight + calibration with balance" section in `docs/api/prep.rst`, calibration pointers in + `llms.txt`/`llms-full.txt`/`llms-practitioner.txt` and `README.md` Survey Support, and + Deville & Särndal (1992) + Sarig, Galili & Eilat (2023) in `docs/references.rst`. - **`SyntheticControl` ADH-2015 §4 tail diagnostics** (two opt-in `SyntheticControlResults` methods, closing the last two ADH-2015 §4 checklist items). `regression_weights()` reports the implied donor weights `W^reg = X0a'(X0a X0a')^{-1} X1a` of the regression counterfactual diff --git a/README.md b/README.md index 651f0edb..c5e46f4a 100644 --- a/README.md +++ b/README.md @@ -137,6 +137,7 @@ Most estimators accept an optional `survey_design` parameter (or `survey=` / `we - **Variance methods**: Taylor Series Linearization (TSL via Binder 1983), replicate weights (BRR / Fay / JK1 / JKn / SDR), survey-aware bootstrap - **Diagnostics**: DEFF per coefficient, effective n, subpopulation analysis, weight trimming, CV on estimates - **Repeated cross-sections**: `CallawaySantAnna(panel=False)` for BRFSS, ACS, CPS +- **Weight calibration / raking**: upstream by design - pair with Meta's [balance](https://import-balance.org/) package, whose `balance.interop.diff_diff` adapter hands raked samples straight to diff-diff; see the [composition-drift tutorial](https://diff-diff.readthedocs.io/en/stable/tutorials/26_composition_drift_calibration.html) No other Python or R DiD package offers design-based variance estimation for modern heterogeneity-robust estimators. diff --git a/diff_diff/guides/llms-full.txt b/diff_diff/guides/llms-full.txt index 8da9890c..c96f5dab 100644 --- a/diff_diff/guides/llms-full.txt +++ b/diff_diff/guides/llms-full.txt @@ -1188,6 +1188,22 @@ read-throughs for compatibility with external adapters that the canonical names; assume the flat aliases are present on every staggered class unless explicitly noted otherwise. +**balance interop.** Meta's `balance` package (>= 0.21) ships the +one-way adapter `balance.interop.diff_diff` (`pip install "balance[did]"`, +pins `diff-diff>=3.3,<4`): `to_survey_design(sample)` builds a +`SurveyDesign` from a balance `Sample`'s active weight column plus the +convention columns `stratum`/`psu`/`fpc`; `to_panel_for_did(sample, by=, +outcomes=)` wraps `diff_diff.aggregate_survey` to collapse respondent +microdata into a unit-period panel plus second-stage design; +`fit_did(sample, estimator=, ...)` resolves any exported estimator by +name or short alias (`CS`/`DiD`/`BJS`/`HAD`) and forwards +`survey_design=`, attaching the source Sample to the result as +`_balance_adjustment` for provenance; `as_balance_diagnostic(sample, +res)` joins balance's ASMD/Kish-ESS with `res.survey_metadata`'s +DEFF/effective-n into one flat dict. The diff-diff surface it consumes +is pinned by `tests/test_balance_interop_contract.py`; the workflow is +demonstrated in Tutorial 26 (composition drift & calibration). + ### DiDResults Returned by `DifferenceInDifferences.fit()` and `TwoWayFixedEffects.fit()`. diff --git a/diff_diff/guides/llms-practitioner.txt b/diff_diff/guides/llms-practitioner.txt index 2088f6c4..7b0b39f9 100644 --- a/diff_diff/guides/llms-practitioner.txt +++ b/diff_diff/guides/llms-practitioner.txt @@ -48,6 +48,11 @@ Key questions to answer: units. - Is there treatment effect heterogeneity you should preserve rather than average over? +- If the data are a survey: are the weights calibrated (raked) at the + granularity of your comparison units? Non-response drift that correlates + with treatment timing does NOT difference out of a DiD; calibrate upstream + with Meta's balance package first — see Tutorial 26: + docs/tutorials/26_composition_drift_calibration.ipynb. ```python # After estimation, the target parameter is available as: diff --git a/diff_diff/guides/llms.txt b/diff_diff/guides/llms.txt index 5268ce18..1c9ce3fd 100644 --- a/diff_diff/guides/llms.txt +++ b/diff_diff/guides/llms.txt @@ -101,6 +101,7 @@ Full practitioner guide: call `diff_diff.get_llm_guide("practitioner")` - [16 Survey DiD](https://diff-diff.readthedocs.io/en/stable/tutorials/16_survey_did.html): Survey-weighted DiD — SurveyDesign, strata/PSU/FPC, replicate weights, subpopulation analysis, DEFF diagnostics - [16 Wooldridge ETWFE](https://diff-diff.readthedocs.io/en/stable/tutorials/16_wooldridge_etwfe.html): Wooldridge (2023, 2025) ETWFE — saturated OLS, logit/Poisson (ASF-based ATT), aggregation types - [22 HAD Survey-Weighted Workflow](https://diff-diff.readthedocs.io/en/stable/tutorials/22_had_survey_design.html): HeterogeneousAdoptionDiD + did_had_pretest_workflow under SurveyDesign(strata, psu, weights, fpc) — BRFSS-shape panel, modest SE inflation explanation, Phase 4.5 C0 QUG-deferred verdict +- [26 Composition Drift & Calibration](https://diff-diff.readthedocs.io/en/stable/tutorials/26_composition_drift_calibration.html): When differential non-response biases the DiD itself — per-state raking with Meta's balance package, `balance.interop.diff_diff` adapter, raking-granularity lesson (requires `pip install balance`) ## Survey Support @@ -110,6 +111,7 @@ Most estimators accept an optional `survey_design` parameter (`SyntheticControl` - **Variance methods**: Taylor Series Linearization (TSL), replicate weights (BRR/Fay/JK1/JKn/SDR), survey-aware bootstrap - **Diagnostics**: DEFF per coefficient, effective n, subpopulation analysis, weight trimming, CV on estimates - **Repeated cross-sections**: `CallawaySantAnna(panel=False)` for BRFSS, ACS, CPS +- **Weight calibration / raking**: upstream by design — pair with Meta's [balance](https://import-balance.org/) package (>= 0.21), whose `balance.interop.diff_diff` adapter (`to_survey_design` / `to_panel_for_did` / `fit_did` / `as_balance_diagnostic`, `pip install "balance[did]"`) hands raked samples straight to diff-diff estimators; see [Tutorial 26](https://diff-diff.readthedocs.io/en/stable/tutorials/26_composition_drift_calibration.html) for when calibration is essential for the causal estimand itself - **Compatibility matrix**: [Survey Design Support](https://diff-diff.readthedocs.io/en/stable/choosing_estimator.html#survey-design-support) No R or Python package offers design-based variance estimation for modern heterogeneity-robust DiD estimators. R's `did`, `fixest`, `synthdid`, and `didimputation` accept flat weight vectors only. diff --git a/docs/api/prep.rst b/docs/api/prep.rst index f6012e87..8660e26c 100644 --- a/docs/api/prep.rst +++ b/docs/api/prep.rst @@ -366,6 +366,41 @@ Example # treatment="treated", time="post", survey_design=stage2, # ) +Weight calibration with balance +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``SurveyDesign`` expects **pre-calibrated** weights: post-stratification, +raking, and calibration are deliberately out of scope for diff-diff and +remain upstream. Meta's `balance `_ package +(>= 0.21) is the recommended companion - it rakes survey samples to +population margins and ships a dedicated adapter, +``balance.interop.diff_diff``, that hands the calibrated sample straight +to diff-diff (``to_survey_design`` / ``to_panel_for_did`` / ``fit_did`` / +``as_balance_diagnostic``; installable via ``pip install "balance[did]"``). + +The handoff needs no adapter if you prefer the native seam - calibrated +weights are just a column:: + + design = SurveyDesign(weights="raked_wt", strata="strat", psu="psu") + panel, stage2 = aggregate_survey( + microdata, by=["state", "year"], outcomes="smoking_rate", + survey_design=design, + ) + result = CallawaySantAnna().fit( + panel, outcome="smoking_rate_mean", unit="state", time="year", + first_treat="g", survey_design=stage2, + ) + +**When calibration matters for the causal estimand** (not just +descriptives): non-response drift that is differential by treatment arm +and time does *not* difference out of a DiD. See the +:doc:`composition-drift tutorial <../tutorials/26_composition_drift_calibration>` +for a worked BRFSS-style failure mode - including why raking granularity +must match the comparison units (state-level raking, as BRFSS itself +does, not a pooled national rake) - and the companion +`balance tutorial `_ +for the robust case (common drift) and descriptive-estimand repair. + Data Validation --------------- diff --git a/docs/doc-deps.yaml b/docs/doc-deps.yaml index 972cb88f..b9fcfbcb 100644 --- a/docs/doc-deps.yaml +++ b/docs/doc-deps.yaml @@ -846,6 +846,9 @@ sources: type: roadmap - path: docs/tutorials/16_survey_did.ipynb type: tutorial + - path: docs/tutorials/26_composition_drift_calibration.ipynb + type: tutorial + note: "balance interop: calibration handoff + composition-drift failure mode" - path: README.md section: "Survey Support" type: user_guide @@ -953,6 +956,9 @@ sources: docs: - path: docs/api/prep.rst type: api_reference + - path: docs/tutorials/26_composition_drift_calibration.ipynb + type: tutorial + note: "aggregate_survey is the seam balance.interop.diff_diff wraps" - path: docs/practitioner_getting_started.rst type: user_guide - path: docs/practitioner_decision_tree.rst diff --git a/docs/index.rst b/docs/index.rst index 38ce2725..59dadd96 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -84,6 +84,7 @@ Quick Links tutorials/21_had_pretest_workflow tutorials/22_had_survey_design tutorials/23_spillover_tva + tutorials/26_composition_drift_calibration .. toctree:: :maxdepth: 1 diff --git a/docs/references.rst b/docs/references.rst index e52c6d50..3d2cbcbb 100644 --- a/docs/references.rst +++ b/docs/references.rst @@ -106,6 +106,14 @@ Survey-Design Inference (Taylor-Series Linearization) The "when to weight" framework distinguishing precision, endogenous-sampling, and population-effect motivations for survey weights; cited in REGISTRY.md ``## Survey Data Support`` -> "Weighted Estimation". +- **Deville, J.-C. & Särndal, C.-E. (1992).** "Calibration Estimators in Survey Sampling." *Journal of the American Statistical Association*, 87(418), 376-382. https://doi.org/10.1080/01621459.1992.10475217 + + The calibration/raking framework underlying post-stratified survey weights. diff-diff deliberately keeps calibration upstream (``SurveyDesign`` expects pre-calibrated weights); the ``docs/api/prep.rst`` "Weight calibration with balance" section and Tutorial 26 document the handoff. + +- **Sarig, T., Galili, T. & Eilat, R. (2023).** "balance - a Python package for balancing biased data samples." *arXiv:2307.06024* (stat.CO). https://arxiv.org/abs/2307.06024 + + Meta's ``balance`` package, the recommended upstream calibration companion. Its ``balance.interop.diff_diff`` adapter (balance >= 0.21) hands raked samples to diff-diff's survey-aware estimators; Tutorial 26 (``docs/tutorials/26_composition_drift_calibration.ipynb``) demonstrates the workflow, and ``tests/test_balance_interop_contract.py`` pins the consumed surface. + Placebo Tests and DiD Diagnostics --------------------------------- diff --git a/docs/survey-roadmap.md b/docs/survey-roadmap.md index 58ea5447..06d4f1d4 100644 --- a/docs/survey-roadmap.md +++ b/docs/survey-roadmap.md @@ -109,13 +109,17 @@ Files: `benchmarks/R/benchmark_realdata_*.R`, `tests/test_survey_real_data.py`, - **Multi-stage design**: not yet documented. Single-stage (strata + PSU) is sufficient per Lumley (2004) Section 2.2. -- **Post-stratification / calibration**: not yet documented. `SurveyDesign` - expects pre-calibrated weights. `samplics` is the most complete Python - option (post-stratification, raking, GREG) but is in read-only mode — - active development has moved to `svy`, which is not yet publicly - released. `weightipy` is actively maintained for raking. Weight - calibration is out of scope for diff-diff today, though building this - capability is a future possibility. +- **Post-stratification / calibration**: DOCUMENTED (2026-07). `SurveyDesign` + expects pre-calibrated weights; calibration stays upstream by design. The + recommended companion is Meta's `balance` package (>= 0.21), which ships a + dedicated `balance.interop.diff_diff` adapter (`pip install "balance[did]"`). + The handoff is documented in `docs/api/prep.rst` ("Weight calibration with + balance"), demonstrated end-to-end in + `docs/tutorials/26_composition_drift_calibration.ipynb` (including when + calibration is essential for the causal estimand, not just descriptives), + and the consumed diff-diff surface is pinned by + `tests/test_balance_interop_contract.py`. `samplics` (read-only; successor + `svy` not yet released) and `weightipy` remain alternatives. ### Phase 10: Survey Completeness (v2.9.0–v3.0) diff --git a/docs/tutorials/26_composition_drift_calibration.ipynb b/docs/tutorials/26_composition_drift_calibration.ipynb new file mode 100644 index 00000000..caa3f6f1 --- /dev/null +++ b/docs/tutorials/26_composition_drift_calibration.ipynb @@ -0,0 +1,242 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "cell-00", + "metadata": {}, + "source": "# When Who Answers Changes: Survey Calibration for Causal DiD\n\nSurvey samples change composition over time: response rates decline, and they decline unevenly across demographic groups. This tutorial shows when that drift breaks a difference-in-differences estimate itself - not just descriptive statistics - and how to fix it by pairing diff-diff with Meta's [balance](https://import-balance.org/) calibration package.\n\n**A matched pair.** This notebook is the diff-diff-side companion to balance's [`balance_diff_diff_brfss` tutorial](https://github.com/facebookresearch/balance/blob/main/tutorials/balance_diff_diff_brfss.ipynb). Their lesson is reassurance: when non-response drift is *common* to treated and control units, it badly biases descriptive trends (which raking repairs) but largely **differences out** of the DiD. This notebook is the warning that completes the pair: when the drift is *differential* - correlated with treatment timing - the DiD is biased too, pre-trend tests won't catch it, and calibration becomes essential for the **causal** estimand. Along the way we hit a subtlety their setting never triggers: raking *granularity*.\n\n**This tutorial covers:**\n\n1. A BRFSS-style staggered smoking-ban dataset with no arm-specific trends by construction - population parallel trends hold in expectation (planted ATT -3.0pp, realized -2.98pp after a rarely-binding probability floor) - but differentially drifting non-response\n2. Why treatment-correlated composition drift does **not** difference out of a DiD - with clean pre-trends the whole way\n3. A per-wave *national* rake as a false fix, and raking at the granularity of your comparison units (state-year, as BRFSS itself does) as the real one\n4. The diff-diff/balance seam both ways: the native `SurveyDesign` + `aggregate_survey` three-liner, and the `balance.interop.diff_diff` adapter (`to_panel_for_did` / `fit_did`) - with an exact-parity check\n5. An estimator sweep (CallawaySantAnna / SunAbraham / ImputationDiD): composition bias is a data problem, not an estimator problem\n6. Cross-package diagnostics: `survey_metadata` design effects and `as_balance_diagnostic`\n7. A practical checklist: when calibration matters for causal vs descriptive estimands\n\n**Prerequisites:** [Tutorial 01: Basic DiD](01_basic_did.ipynb), [Tutorial 02: Staggered DiD](02_staggered_did.ipynb), [Tutorial 16: Survey DiD](16_survey_did.ipynb); ideally the [balance quickstart](https://import-balance.org/docs/tutorials/quickstart/).\n\n**Requirements:** `pip install diff-diff \"balance>=0.21\" matplotlib`. balance is needed **only for this tutorial** - it is not a diff-diff dependency. (`pip install \"balance[did]\"` installs both packages at once.)\n" + }, + { + "cell_type": "markdown", + "id": "cell-01", + "metadata": {}, + "source": "## 1. Two regimes of non-response drift\n\n\"Do the survey weights matter?\" is really two questions - one per estimand:\n\n| Non-response drift | Descriptive estimand (a prevalence trend) | Causal DiD estimand (an ATT) |\n|---|---|---|\n| **Common** - same in treated and control units | **Biased** - raking to population margins repairs it | **Robust** - composition changes difference out ([balance's tutorial](https://github.com/facebookresearch/balance/blob/main/tutorials/balance_diff_diff_brfss.ipynb)) |\n| **Differential** - correlated with treatment x time | Biased | **Biased** - the drift *is* treatment x post shaped, so the DiD reads it as treatment effect (**this tutorial**) |\n\nThe dangerous cell is the bottom-right one, because nothing in the standard DiD toolkit flags it: the fit converges, standard errors look fine, and - as we will see - the pre-treatment event-study coefficients stay clean.\n" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-02", + "metadata": {}, + "outputs": [], + "source": "import logging\nimport warnings\n\nimport numpy as np\nimport pandas as pd\n\nimport balance # >= 0.21 for balance.interop.diff_diff\nfrom balance import Sample\nfrom balance.interop import diff_diff as bd\n\n# Quiet balance's per-call INFO logging (must run AFTER `import balance` -\n# its __init__ re-arms the logger) and pandas FutureWarnings.\nlogging.getLogger(\"balance\").setLevel(logging.ERROR)\nwarnings.filterwarnings(\"ignore\", category=FutureWarning)\n# diff-diff normalizes pweights to mean 1 and warns to say so; expected here.\nwarnings.filterwarnings(\n \"ignore\", message=\".*weights normalized.*\", category=UserWarning\n)\n\nimport diff_diff\nfrom diff_diff import CallawaySantAnna, SurveyDesign, aggregate_survey\n\ntry:\n import matplotlib.pyplot as plt\n\n plt.style.use(\"seaborn-v0_8-whitegrid\")\n HAS_MATPLOTLIB = True\nexcept ImportError:\n HAS_MATPLOTLIB = False\n\nprint(f\"diff-diff {diff_diff.__version__} | balance {balance.__version__}\")\n" + }, + { + "cell_type": "markdown", + "id": "cell-03", + "metadata": {}, + "source": "## 2. The data: a BRFSS-shaped survey around staggered smoking bans\n\nWe simulate respondent-level microdata mirroring the public-use [BRFSS](https://www.cdc.gov/brfss/annual_data/annual_2024.html) file: 50 states x 7 years (2018-2024), with 1,200 adults invited per state-year, of whom roughly 580-860 respond (~265,000 respondent rows in total). Comprehensive indoor-smoking bans take effect in 10 states in 2020 and 10 more in 2022; 30 states never adopt. The outcome is **current smoking** (binary `smoker`). Each respondent carries raking demographics (`age_band`, `educ_cat`), a complex design (`stratum`, `psu`, `fpc`), and a demographic-blind `design_weight` standing in for BRFSS `_LLCPWT`. A separate 20,000-row \"ACS\" frame supplies the population margins (and per-state adult population counts) we will rake to.\n\nThe generator is built so every headline claim is checkable against ground truth:\n\n- **No arm-specific trends by construction.** Smoking prevalence is additive and linear in education, age, and time; demographic margins are identical in every state and year. Trends differ *by demographic* (smoking falls fastest among the college-educated) but identically across treatment arms, so population parallel trends hold **in expectation**; mean-zero PSU-year shocks (kept in deliberately - they create the design effects the diagnostics section reads) add seed-specific noise but no systematic drift.\n- **The planted treatment effect is -3.0pp** in every treated state-year. One honest wrinkle: probabilities are floored at 1%, and that floor binds for ~2% of treated-post person-years (elderly college-educated respondents in low-smoking states), so the **realized population ATT is -2.98pp**. The generator computes it from the clipped potential outcomes and returns it as `truth[\"realized_att_pp\"]` - that is the truth line we hold every estimate against below.\n- **Non-response depends only on the raking observables** (`age_band`, `educ_cat`) - never on smoking itself. Response rates decline over time everywhere, concentrated among younger and less-educated adults (the real BRFSS pattern). That is the *common* drift.\n- **The differential twist** (`differential=True`): after a state's ban takes effect, response among `hs_or_less` adults falls an extra 7pp per event year - enforcement publicity raises refusals in the low-SES communities where smoking is concentrated, and the implementation consumes the follow-up-call budget. Scenario A (`differential=False`) and scenario B share the identical invited population and outcomes; they differ **only in who answers the phone**.\n\nBecause non-response is driven entirely by variables inside the raking margins (missing-at-random given the margins), calibration *can* fix everything - the interesting question is what kind of calibration. To run this on real data, replace this one cell with `pyreadstat.read_xport(...)` calls per BRFSS year.\n" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-04", + "metadata": {}, + "outputs": [], + "source": "N_STATES = 50\nYEARS = np.arange(2018, 2025)\nN_INVITED = 1200\nN_STRATA = 5\nPSUS_PER_STATE = 8\nFPC_PSUS_PER_STRATUM = 200.0\n\nAGE_BANDS = [\"18-34\", \"35-49\", \"50-64\", \"65+\"]\nAGE_SHARES = np.array([0.30, 0.25, 0.25, 0.20])\nEDUC_CATS = [\"hs_or_less\", \"some_college\", \"college_plus\"]\nEDUC_SHARES = np.array([0.35, 0.30, 0.35])\n\nBASE_EDUC_PP = np.array([22.0, 15.0, 9.0])\nAGE_ADJ_PP = np.array([2.0, 3.0, 1.0, -2.0])\nTREND_COMMON_PP = 0.25\nTREND_EDUC_PP = np.array([-0.15, 0.0, 0.10])\nTRUE_ATT_PP = -3.0\nSTATE_RE_SD_PP = 1.5\nPSU_SHOCK_SD_PP = 0.8\nP_CLIP_PP = (1.0, 60.0)\n\nR_BASE = 0.70\nR_AGE_SHIFT = np.array([-0.10, -0.02, 0.03, 0.08])\nR_EDUC_SHIFT = np.array([-0.09, 0.00, 0.06])\nR_COMMON_DRIFT_EDUC = np.array([0.015, 0.0075, 0.0])\nR_COMMON_DRIFT_YOUNG = 0.010\nR_DIFF_DRIFT_PER_EVENT_YEAR = 0.07\nR_CLIP = (0.10, 0.95)\n\nTARGET_N = 20_000\nSEED = 20260704\n\n\ndef simulate_brfss_smoking(differential, seed=SEED, drift_start_offset=0):\n \"\"\"BRFSS-style microdata around staggered smoking bans.\n\n No arm-specific trends by construction: population parallel trends\n hold in expectation (mean-zero PSU-year shocks add noise, not drift),\n and the planted effect is -3.0pp (realized population ATT ~-2.98pp\n after the probability floor). All SYSTEMATIC estimator bias in this\n notebook comes from sample composition. Scenarios A/B share invited\n respondents and outcomes for the same seed - they differ only in who\n responds.\n \"\"\"\n rng = np.random.default_rng(seed)\n\n perm = rng.permutation(N_STATES)\n g_of_state = np.zeros(N_STATES, dtype=int)\n g_of_state[perm[:10]] = 2020\n g_of_state[perm[10:20]] = 2022\n stratum_of_state = rng.integers(0, N_STRATA, size=N_STATES)\n state_pop = rng.lognormal(mean=np.log(4e6), sigma=0.6, size=N_STATES)\n state_re = np.clip(rng.normal(0.0, STATE_RE_SD_PP, size=N_STATES), -3.0, 3.0)\n psu_shock = rng.normal(\n 0.0, PSU_SHOCK_SD_PP, size=(N_STATES, PSUS_PER_STATE, len(YEARS))\n )\n\n n_inv = N_STATES * len(YEARS) * N_INVITED\n state = np.repeat(np.arange(N_STATES), len(YEARS) * N_INVITED)\n year = np.tile(np.repeat(YEARS, N_INVITED), N_STATES)\n age_idx = rng.choice(len(AGE_BANDS), size=n_inv, p=AGE_SHARES)\n educ_idx = rng.choice(len(EDUC_CATS), size=n_inv, p=EDUC_SHARES)\n psu_idx = rng.integers(0, PSUS_PER_STATE, size=n_inv)\n u_respond = rng.uniform(size=n_inv)\n u_smoker = rng.uniform(size=n_inv)\n weight_jitter = rng.uniform(0.85, 1.15, size=n_inv)\n\n k = year - YEARS[0]\n year_idx = year - YEARS[0]\n g = g_of_state[state]\n treated_post = (g > 0) & (year >= g)\n\n base_pp = (\n BASE_EDUC_PP[educ_idx]\n + AGE_ADJ_PP[age_idx]\n + state_re[state]\n + psu_shock[state, psu_idx, year_idx]\n - (TREND_COMMON_PP + TREND_EDUC_PP[educ_idx]) * k\n )\n p_pp = np.clip(base_pp + TRUE_ATT_PP * treated_post, *P_CLIP_PP)\n smoker = (u_smoker < p_pp / 100.0).astype(int)\n\n r = (\n R_BASE\n + R_AGE_SHIFT[age_idx]\n + R_EDUC_SHIFT[educ_idx]\n - R_COMMON_DRIFT_EDUC[educ_idx] * k\n - R_COMMON_DRIFT_YOUNG * k * (age_idx == 0)\n )\n if differential:\n event_time = year - g - drift_start_offset\n hit = (g > 0) & (event_time >= 0) & (educ_idx == 0)\n r = r - R_DIFF_DRIFT_PER_EVENT_YEAR * (event_time + 1) * hit\n r = np.clip(r, *R_CLIP)\n responded = u_respond < r\n\n micro = pd.DataFrame(\n {\n \"id\": np.arange(n_inv)[responded],\n \"state\": state[responded],\n \"year\": year[responded],\n \"g\": g[responded],\n \"smoker\": smoker[responded],\n \"age_band\": np.array(AGE_BANDS)[age_idx[responded]],\n \"educ_cat\": np.array(EDUC_CATS)[educ_idx[responded]],\n \"stratum\": stratum_of_state[state[responded]],\n \"psu\": state[responded] * 100 + psu_idx[responded],\n \"fpc\": FPC_PSUS_PER_STRATUM,\n \"design_weight\": (state_pop[state] / N_INVITED * weight_jitter)[\n responded\n ],\n }\n )\n\n rng_t = np.random.default_rng(seed + 1)\n target_df = pd.DataFrame(\n {\n \"id\": np.arange(TARGET_N),\n \"age_band\": np.array(AGE_BANDS)[\n rng_t.choice(len(AGE_BANDS), size=TARGET_N, p=AGE_SHARES)\n ],\n \"educ_cat\": np.array(EDUC_CATS)[\n rng_t.choice(len(EDUC_CATS), size=TARGET_N, p=EDUC_SHARES)\n ],\n }\n )\n\n kk = YEARS - YEARS[0]\n cell = (\n BASE_EDUC_PP[None, :, None]\n + AGE_ADJ_PP[None, None, :]\n - (TREND_COMMON_PP + TREND_EDUC_PP[None, :, None]) * kk[:, None, None]\n )\n tp = (g_of_state[None, :] > 0) & (YEARS[:, None] >= g_of_state[None, :])\n base_prev = np.einsum(\"tea,e,a->t\", cell, EDUC_SHARES, AGE_SHARES)\n w_s = state_pop / state_pop.sum()\n pop_prev = base_prev + TRUE_ATT_PP * (tp * w_s[None, :]).sum(axis=1)\n # Realized population ATT: the probability floor P_CLIP_PP[0] binds for\n # ~2% of treated-post person-years, attenuating the planted -3.0pp.\n y1 = np.clip(base_pp + TRUE_ATT_PP, *P_CLIP_PP)\n y0 = np.clip(base_pp, *P_CLIP_PP)\n w_pop = state_pop[state]\n realized_att_pp = ((y1 - y0) * w_pop)[treated_post].sum() / w_pop[\n treated_post\n ].sum()\n truth = {\n \"true_att_pp\": TRUE_ATT_PP,\n \"realized_att_pp\": float(realized_att_pp),\n \"floor_bind_share\": float(\n ((y1 - y0) > TRUE_ATT_PP + 1e-12)[treated_post].mean()\n ),\n \"pop_prevalence_by_year\": dict(zip(YEARS.tolist(), pop_prev / 100.0)),\n \"g_of_state\": g_of_state,\n \"state_pop\": state_pop,\n }\n return micro, target_df, truth\n" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-05", + "metadata": {}, + "outputs": [], + "source": "micro_common, target, truth = simulate_brfss_smoking(differential=False)\nmicro_diff, _, _ = simulate_brfss_smoking(differential=True)\n\nby_cell = micro_diff.groupby([\"state\", \"year\"]).size()\nprint(f\"scenario A respondents: {len(micro_common):,}\")\nprint(f\"scenario B respondents: {len(micro_diff):,}\")\nprint(f\"respondents per state-year (B): {by_cell.min()}-{by_cell.max()}\")\nTRUE_ATT = truth[\"realized_att_pp\"]\nprint(f\"planted ATT: {truth['true_att_pp']}pp | realized population ATT: \"\n f\"{TRUE_ATT:+.2f}pp (probability floor binds for \"\n f\"{truth['floor_bind_share']:.1%} of treated-post person-years)\")\ntarget.head(3)\n" + }, + { + "cell_type": "markdown", + "id": "cell-06", + "metadata": {}, + "source": "## 3. The native seam: `SurveyDesign` + `aggregate_survey` + `CallawaySantAnna`\n\ndiff-diff's survey path consumes calibrated weights as *just a column*. The handoff from microdata to a modern staggered estimator is three calls:\n\n1. **`SurveyDesign(weights=..., strata=..., psu=..., fpc=...)`** - declare the complex design on the microdata (column *names*, no copies).\n2. **`aggregate_survey(micro, by=[\"state\", \"year\"], outcomes=\"smoker\", ...)`** - collapse respondents to a state-year panel of design-weighted prevalences with full Taylor-linearized precision tracking, returning the panel *and* a pre-configured second-stage `SurveyDesign` (population weights + state-level clustering, taken from the first `by` column).\n3. **`CallawaySantAnna(...).fit(panel, ..., survey_design=second_stage)`** - the staggered DiD with design-based variance.\n\nOne convention worth noticing now: we named the design columns `stratum` / `psu` / `fpc`. Those are exactly the defaults in `balance.interop.conventions.DEFAULT_DESIGN_COLUMNS`, so when we reach the balance adapter in section 8 it will wire the same design automatically.\n\nWe wrap the three calls in a helper because we are about to fit the same model under four different weight columns - that swap being a one-argument change is the point of the design.\n" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-07", + "metadata": {}, + "outputs": [], + "source": "def fit_survey_cs(micro, weights_col):\n \"\"\"Native seam: microdata + a weight column -> survey-weighted CS fit.\"\"\"\n design = SurveyDesign(\n weights=weights_col, strata=\"stratum\", psu=\"psu\", fpc=\"fpc\"\n )\n panel, second_stage = aggregate_survey(\n micro, by=[\"state\", \"year\"], outcomes=\"smoker\", survey_design=design\n )\n panel = panel.merge(\n micro[[\"state\", \"g\"]].drop_duplicates(), on=\"state\", how=\"left\"\n )\n cs = CallawaySantAnna(\n estimation_method=\"reg\",\n control_group=\"not_yet_treated\",\n base_period=\"universal\",\n )\n return cs.fit(\n panel,\n outcome=\"smoker_mean\",\n unit=\"state\",\n time=\"year\",\n first_treat=\"g\",\n survey_design=second_stage,\n aggregate=\"all\",\n )\n\n\ndef show_att(res, label):\n lo, hi = res.overall_conf_int\n print(\n f\"{label}: ATT = {res.overall_att * 100:+.2f}pp \"\n f\"(SE {res.overall_se * 100:.2f}, 95% CI [{lo * 100:+.2f}, {hi * 100:+.2f}])\"\n )\n" + }, + { + "cell_type": "markdown", + "id": "cell-08", + "metadata": {}, + "source": "## 4. Scenario A - common drift: the descriptive trend breaks, the DiD does not\n\nFirst the regime the balance tutorial covers. Response rates decline over time, concentrated among younger and less-educated adults, but *identically* in treated and control states. The design-weighted national prevalence drifts away from the truth (the sample is increasingly educated, and educated adults smoke less - so the naive trend overstates the decline). Yet the DiD contrast is untouched, because both arms mis-measure levels the same way.\n" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-09", + "metadata": {}, + "outputs": [], + "source": "res_a = fit_survey_cs(micro_common, \"design_weight\")\nshow_att(res_a, \"Scenario A, design weights\")\n\npop_series = pd.Series(truth[\"pop_prevalence_by_year\"])\ndesign_series = micro_common.groupby(\"year\")[[\"smoker\", \"design_weight\"]].apply(\n lambda d: (d.smoker * d.design_weight).sum() / d.design_weight.sum()\n)\ngap_pp = (design_series - pop_series) * 100\nprint(f\"descriptive gap, design-weighted vs population: \"\n f\"{gap_pp.loc[2018]:+.2f}pp (2018) -> {gap_pp.loc[2024]:+.2f}pp (2024)\")\n\nif HAS_MATPLOTLIB:\n fig, ax = plt.subplots(figsize=(8, 4.5))\n ax.plot(pop_series.index, pop_series * 100, \"k-\", lw=2, label=\"population (truth)\")\n ax.plot(design_series.index, design_series * 100, \"C1--\", lw=2, marker=\"o\",\n label=\"design-weighted sample\")\n ax.set_xlabel(\"year\")\n ax.set_ylabel(\"smoking prevalence (%)\")\n ax.set_title(\"Scenario A: common drift biases the descriptive trend\")\n ax.legend()\n plt.tight_layout()\n plt.show()\n" + }, + { + "cell_type": "markdown", + "id": "cell-10", + "metadata": {}, + "source": "The ATT comes back at about **-3.1pp against the -2.98pp realized truth** - the causal estimate is robust even though the descriptive series visibly drifts below the truth. This is exactly the reassurance half of the story; for the full descriptive-repair workflow (per-wave raking, Love plots, effective sample sizes) see the [balance tutorial](https://github.com/facebookresearch/balance/blob/main/tutorials/balance_diff_diff_brfss.ipynb) - we will not repeat it here.\n\nNow we break it.\n" + }, + { + "cell_type": "markdown", + "id": "cell-11", + "metadata": {}, + "source": "## 5. Scenario B - differential drift: the causal estimand breaks\n\nSame states, same bans, same outcomes, same common drift - plus one mechanism: once a ban takes effect, response among `hs_or_less` adults in that state falls an extra 7pp per event year. The composition damage is easy to see if you know to look at it:\n" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-12", + "metadata": {}, + "outputs": [], + "source": "hs_share = (\n micro_diff.assign(hs=lambda d: (d.educ_cat == \"hs_or_less\").astype(float))\n .groupby([\"year\", \"g\"])[[\"hs\", \"design_weight\"]]\n .apply(lambda d: (d.hs * d.design_weight).sum() / d.design_weight.sum())\n .unstack()\n)\nprint(\"design-weighted hs_or_less share by cohort (population share: 0.350):\")\nprint(hs_share.round(3).to_string())\n\nif HAS_MATPLOTLIB:\n fig, ax = plt.subplots(figsize=(8, 4.5))\n styles = {0: (\"C0\", \"never treated (30 states)\"),\n 2020: (\"C3\", \"ban in 2020 (10 states)\"),\n 2022: (\"C1\", \"ban in 2022 (10 states)\")}\n for gval, (color, label) in styles.items():\n ax.plot(hs_share.index, hs_share[gval], color=color, marker=\"o\", label=label)\n ax.axhline(0.35, color=\"gray\", ls=\":\", label=\"population share (0.35)\")\n for gval, color in [(2020, \"C3\"), (2022, \"C1\")]:\n ax.axvline(gval, color=color, ls=\"--\", alpha=0.4)\n ax.set_xlabel(\"year\")\n ax.set_ylabel(\"hs_or_less share of design-weighted sample\")\n ax.set_title(\"Scenario B: who answers changes exactly when treatment starts\")\n ax.legend()\n plt.tight_layout()\n plt.show()\n" + }, + { + "cell_type": "markdown", + "id": "cell-13", + "metadata": {}, + "source": "Each cohort's low-education share peels away from ~0.30 **exactly at its adoption year**, sliding to ~0.11 (2020 cohort) and ~0.19 (2022 cohort) by 2024 while never-treated states stay flat.\n\nWhy this bias cannot difference out: the design weights are demographic-blind, so a state-year's measured prevalence is $m(s,t) = \\sum_e \\sigma_e(s,t)\\, p_e(t)$, where $\\sigma_e$ is the *realized sample* share of education group $e$ and $p_e$ its true prevalence. A DiD passes any additive term that is common across arms - that killed the composition term in scenario A. But here $\\sigma_{hs}$ falls **only in treated states, only after adoption**: the composition artifact $\\Delta\\sigma_{hs} \\times (p_{hs} - \\bar{p}_{rest})$ sits precisely in the treatment x post cell of the design. To the estimator, it *is* an ATT. With a ~12pp gap in smoking prevalence between `hs_or_less` and the (share-weighted) rest, a 17pp share collapse manufactures roughly an extra -2pp of \"effect\" at long horizons, ramping up from ~-0.3pp at adoption - which masquerades as a treatment effect that *grows over time*.\n" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-14", + "metadata": {}, + "outputs": [], + "source": "res_design = fit_survey_cs(micro_diff, \"design_weight\")\nshow_att(res_design, \"Scenario B, design weights\")\nprint(f\"true ATT (realized): {TRUE_ATT:+.2f}pp\")\n\nes_design = res_design.event_study_effects\nif HAS_MATPLOTLIB:\n fig, ax = plt.subplots(figsize=(8, 4.5))\n ev = sorted(es_design)\n eff = np.array([es_design[e][\"effect\"] for e in ev]) * 100\n se = np.array([es_design[e][\"se\"] for e in ev]) * 100\n ax.errorbar(ev, eff, yerr=1.96 * np.nan_to_num(se), fmt=\"o-\", color=\"C1\",\n capsize=3, label=\"design weights\")\n ax.axhline(0, color=\"gray\", lw=0.8)\n ax.axhline(TRUE_ATT, color=\"k\", ls=\":\", label=f\"true ATT ({TRUE_ATT:.2f}pp)\")\n ax.axvline(-0.5, color=\"gray\", ls=\"--\", alpha=0.5)\n ax.set_xlabel(\"event time (years since ban)\")\n ax.set_ylabel(\"effect on smoking prevalence (pp)\")\n ax.set_title(\"Scenario B, design weights: clean pre-trends, biased ATT\")\n ax.legend()\n plt.tight_layout()\n plt.show()\n" + }, + { + "cell_type": "markdown", + "id": "cell-15", + "metadata": {}, + "source": "Three things just happened, and together they are the core lesson of this tutorial:\n\n1. **The ATT is badly overstated**: about **-4.1pp against a realized truth of -2.98pp** (roughly +38%), and the 95% CI (~[-4.9, -3.4]) *excludes* the truth. You would publish \"bans cut smoking by 4 points\" with confidence.\n2. **The dynamics are fake**: the event study \"builds\" to about -5pp at event year 4. That growth is the response-rate ramp, not policy.\n3. **The pre-trends are clean** (every pre-treatment coefficient is well inside its confidence band), because the drift starts *at* adoption. A pre-trend test cannot catch at-adoption composition drift - it certifies parallel trends of *who you measured*, not of the population.\n\nThe standard toolkit passes this regression with flying colors. The composition plot above is the only thing that flagged it.\n" + }, + { + "cell_type": "markdown", + "id": "cell-16", + "metadata": {}, + "source": "## 6. Fix attempt 1: per-wave national raking (the natural first move)\n\nThe textbook response - and the balance quickstart recipe - is to rake each survey wave to population margins: reweight so the sample's `age_band` x `educ_cat` distribution matches the ACS frame, year by year. We wrap balance's `Sample -> set_target -> adjust(method=\"rake\")` in a helper that (a) rakes within each cell of a chosen `granularity`, and (b) rescales each cell's raked weights to its population **count** - in real practice raking targets are population counts, so weight totals carry the state scale rather than balance's internal per-call normalization.\n\nOne mechanical detail: balance casts ids to `str`, so we align the returned weights back to the original rows by string id and assert nothing went missing.\n" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-17", + "metadata": {}, + "outputs": [], + "source": "RAKE_VARS = [\"age_band\", \"educ_cat\"]\n\n\ndef rake_to_population(micro, target_df, granularity, weight_name, cell_totals):\n \"\"\"Rake design weights to population margins within each granularity cell.\n\n cell_totals maps each groupby key to the population COUNT the cell's\n raked weights must sum to.\n Returns (micro + weight_name column, {cell_key: adjusted balance Sample}).\n \"\"\"\n target_sample = Sample.from_frame(target_df, id_column=\"id\")\n cols = [\"id\", *RAKE_VARS, \"smoker\", \"design_weight\"]\n w_new = pd.Series(np.nan, index=micro.index)\n adjusted = {}\n for key, cell in micro.groupby(granularity):\n if isinstance(key, tuple) and len(key) == 1:\n key = key[0]\n s = Sample.from_frame(\n cell[cols].copy(),\n id_column=\"id\",\n weight_column=\"design_weight\",\n outcome_columns=[\"smoker\"],\n )\n adj = s.set_target(target_sample).adjust(method=\"rake\", variables=RAKE_VARS)\n w = adj.df.set_index(\"id\")[adj.weight_column]\n aligned = w.reindex(cell[\"id\"].astype(str).values).to_numpy()\n assert not np.isnan(aligned).any(), f\"NaN raked weights in cell {key}\"\n aligned = aligned * (cell_totals[key] / aligned.sum())\n w_new.loc[cell.index] = aligned\n adjusted[key] = adj\n out = micro.copy()\n out[weight_name] = w_new\n return out, adjusted\n\n\nstate_pop = truth[\"state_pop\"]\n\nmicro_diff, adj_national = rake_to_population(\n micro_diff, target, [\"year\"], \"w_national\",\n cell_totals={int(y): state_pop.sum() for y in YEARS},\n)\nres_national = fit_survey_cs(micro_diff, \"w_national\")\nshow_att(res_national, \"Scenario B, per-wave NATIONAL rake\")\n\nm24 = micro_diff[micro_diff.year == 2024].assign(\n hs=lambda d: (d.educ_cat == \"hs_or_less\").astype(float)\n)\nfor gval, label in [(0, \"never treated\"), (2020, \"2020 cohort\")]:\n sub = m24[m24.g == gval]\n share = (sub.hs * sub.w_national).sum() / sub.w_national.sum()\n print(f\"2024 hs_or_less share after national rake, {label}: \"\n f\"{share:.3f} (population: 0.350)\")\n" + }, + { + "cell_type": "markdown", + "id": "cell-18", + "metadata": {}, + "source": "**It got *worse*** - about **-4.4pp**. And the composition printout says why: after the national rake, never-treated states sit at ~0.42 low-education share while the 2020 cohort sits at ~0.18. The rake satisfied the national margin *in aggregate* - by pushing extra weight onto low-education respondents everywhere, over-correcting the states that still had them and leaving the treated states (which barely have any left) still far below the margin.\n\nRaking equalizes composition **at the level you rake at**. A DiD compares *states*; a national rake constrains only the national mixture, so the treated-vs-control composition gap - the thing biasing the DiD - survives, and the redistributed weight can even amplify it.\n\n**The rule: rake at (or below) the granularity of the units your design compares.** This is not exotic - it is what BRFSS itself does: `_LLCPWT` is raked *within each state* to state-level demographic control totals. (The balance tutorial's per-wave national rake was the right call *there*: under common drift, every state needs the same correction, so the national rake fixes each state too. Under differential drift, that shortcut collapses.)\n" + }, + { + "cell_type": "markdown", + "id": "cell-19", + "metadata": {}, + "source": "## 7. Fix attempt 2: rake each state-year to the margins\n\nSame helper, `granularity=[\"state\", \"year\"]` - 350 small rakes (~750 respondents each, a couple of seconds total), each rescaled to its state's adult-population count from the ACS frame. State totals are then constant across years by construction, so the second-stage population weights recover the true state scale, uncontaminated by response-rate levels.\n" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-20", + "metadata": {}, + "outputs": [], + "source": "micro_diff, adj_cell = rake_to_population(\n micro_diff, target, [\"state\", \"year\"], \"w_raked\",\n cell_totals={\n (st, int(y)): state_pop[st] for st in range(N_STATES) for y in YEARS\n },\n)\nres_raked = fit_survey_cs(micro_diff, \"w_raked\")\n\nshow_att(res_design, \"design weights \")\nshow_att(res_national, \"national rake \")\nshow_att(res_raked, \"state-year rake \")\nprint(f\"true ATT (realized): {TRUE_ATT:+.2f}pp\")\n\nm24 = micro_diff[micro_diff.year == 2024].assign(\n hs=lambda d: (d.educ_cat == \"hs_or_less\").astype(float)\n)\nfor gval, label in [(0, \"never treated\"), (2020, \"2020 cohort\"), (2022, \"2022 cohort\")]:\n sub = m24[m24.g == gval]\n share = (sub.hs * sub.w_raked).sum() / sub.w_raked.sum()\n print(f\"2024 hs_or_less share after state-year rake, {label}: {share:.3f}\")\n\nes_raked = res_raked.event_study_effects\nif HAS_MATPLOTLIB:\n fig, ax = plt.subplots(figsize=(8.5, 5))\n for es, color, label in [(es_design, \"C1\", \"design weights\"),\n (es_raked, \"C0\", \"state-year raked\")]:\n ev = sorted(es)\n eff = np.array([es[e][\"effect\"] for e in ev]) * 100\n se = np.array([es[e][\"se\"] for e in ev]) * 100\n ax.errorbar(ev, eff, yerr=1.96 * np.nan_to_num(se), fmt=\"o-\", color=color,\n capsize=3, label=label, alpha=0.9)\n ax.axhline(0, color=\"gray\", lw=0.8)\n ax.axhline(TRUE_ATT, color=\"k\", ls=\":\", lw=1.5,\n label=f\"true ATT ({TRUE_ATT:.2f}pp)\")\n ax.axvline(-0.5, color=\"gray\", ls=\"--\", alpha=0.5)\n ax.set_xlabel(\"event time (years since ban)\")\n ax.set_ylabel(\"effect on smoking prevalence (pp)\")\n ax.set_title(\"Composition drift manufactures dynamics; state-level raking removes them\")\n ax.legend()\n plt.tight_layout()\n plt.show()\n" + }, + { + "cell_type": "markdown", + "id": "cell-21", + "metadata": {}, + "source": "Recovered: the state-year rake lands at about **-3.2pp (SE 0.5)** - within half a standard error of the realized -2.98pp truth - the fake build-up flattens onto the truth line at every horizon, and all three cohorts' low-education shares sit at 0.354 against the 0.350 population margin. The scoreboard:\n\n| weights | ATT (pp) | vs realized truth (-2.98) |\n|---|---|---|\n| design weights | ~ -4.1 | ~+38% overstated, CI excludes truth |\n| per-wave national rake | ~ -4.4 | worse - margins met in aggregate only |\n| **state-year rake, population-count totals** | **~ -3.2** | recovered (within 0.5 SE) |\n" + }, + { + "cell_type": "markdown", + "id": "cell-22", + "metadata": {}, + "source": "## 8. The adapter: `balance.interop.diff_diff`\n\nEverything above used the native seam - calibrated weights as a column. balance 0.21 ships an adapter that packages the same handoff when your object in hand is a balance `Sample`:\n\n- **`bd.to_panel_for_did(sample, by=, outcomes=)`** builds the first-stage `SurveyDesign` from the Sample's *active* weight column (auto-wiring our `stratum`/`psu`/`fpc` convention columns), strips balance's bookkeeping columns, and calls `diff_diff.aggregate_survey` - returning the same `(panel, second_stage_design)` pair.\n- **`bd.fit_did(sample, estimator=, ...)`** resolves any diff-diff estimator by name (or short alias: `\"CS\"`, `\"DiD\"`, `\"BJS\"`, `\"HAD\"`), splits your kwargs between `__init__` and `fit()` by signature, forwards `survey_design=`, and attaches the source Sample to the result as `_balance_adjustment` for provenance.\n\nBoth paths must agree exactly - and we assert it:\n" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-23", + "metadata": {}, + "outputs": [], + "source": "keep = [\"id\", \"state\", \"year\", \"smoker\", \"age_band\", \"educ_cat\",\n \"stratum\", \"psu\", \"fpc\", \"w_raked\"]\nsample = Sample.from_frame(\n micro_diff[keep].copy(),\n id_column=\"id\",\n weight_column=\"w_raked\",\n outcome_columns=[\"smoker\"],\n)\n\npanel_df, second_stage = bd.to_panel_for_did(\n sample, by=[\"state\", \"year\"], outcomes=\"smoker\"\n)\nprint(f\"panel: {panel_df.shape[0]} state-years, \"\n f\"second-stage weights={second_stage.weights!r}, psu={second_stage.psu!r}\")\n\n# g is unit-constant metadata: re-join it at the panel level\npanel_df = panel_df.merge(\n micro_diff[[\"state\", \"g\"]].drop_duplicates(), on=\"state\", how=\"left\"\n)\npanel_df[\"panel_id\"] = np.arange(len(panel_df))\npanel_sample = Sample.from_frame(\n panel_df,\n id_column=\"panel_id\",\n weight_column=second_stage.weights,\n outcome_columns=[\"smoker_mean\"],\n)\n\nres_adapter = bd.fit_did(\n panel_sample,\n estimator=\"CallawaySantAnna\",\n outcome=\"smoker_mean\",\n time=\"year\",\n unit=\"state\",\n treatment_first=\"g\",\n design_columns={\"psu\": \"state\"}, # match the native second-stage design\n estimation_method=\"reg\",\n control_group=\"not_yet_treated\",\n base_period=\"universal\",\n aggregate=\"all\",\n)\n\nassert np.isclose(res_adapter.overall_att, res_raked.overall_att, rtol=1e-12)\nprint(f\"native ATT: {res_raked.overall_att * 100:+.4f}pp\")\nprint(f\"adapter ATT: {res_adapter.overall_att * 100:+.4f}pp (exact match)\")\nprint(f\"provenance attached: {hasattr(res_adapter, '_balance_adjustment')}\")\n" + }, + { + "cell_type": "markdown", + "id": "cell-24", + "metadata": {}, + "source": "Use whichever side of the seam matches your pipeline: if you live in balance (Samples, adjustment lineage, per-wave diagnostics), the adapter keeps that lineage attached to the diff-diff result; if you live in DataFrames, the native three-liner needs no extra dependency. The contract between the packages - `aggregate_survey`'s signature, `SurveyDesign`'s fields, estimator names, the provenance side-channel - is pinned on the diff-diff side by `tests/test_balance_interop_contract.py`, so both routes stay stable.\n" + }, + { + "cell_type": "markdown", + "id": "cell-25", + "metadata": {}, + "source": "## 9. Does the fix depend on the estimator?\n\nA tempting escape: \"maybe a more robust estimator handles it.\" It cannot - the composition artifact lives in the *data*, in the treatment x post cell, upstream of any identification strategy. `fit_did`'s one-string estimator swap makes the sweep trivial:\n" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-26", + "metadata": {}, + "outputs": [], + "source": "ESTIMATOR_KWARGS = {\n \"CallawaySantAnna\": dict(estimation_method=\"reg\",\n control_group=\"not_yet_treated\",\n base_period=\"universal\", aggregate=\"all\"),\n \"SunAbraham\": dict(control_group=\"never_treated\"),\n \"ImputationDiD\": {},\n}\n\nrows = []\nfor wcol, wlabel in [(\"design_weight\", \"design\"), (\"w_raked\", \"state-raked\")]:\n s = Sample.from_frame(\n micro_diff[keep[:-1] + [wcol]].copy(),\n id_column=\"id\", weight_column=wcol, outcome_columns=[\"smoker\"],\n )\n pdf, ss = bd.to_panel_for_did(s, by=[\"state\", \"year\"], outcomes=\"smoker\")\n pdf = pdf.merge(micro_diff[[\"state\", \"g\"]].drop_duplicates(),\n on=\"state\", how=\"left\")\n pdf[\"panel_id\"] = np.arange(len(pdf))\n ps = Sample.from_frame(pdf, id_column=\"panel_id\", weight_column=ss.weights,\n outcome_columns=[\"smoker_mean\"])\n for name, extra in ESTIMATOR_KWARGS.items():\n r = bd.fit_did(\n ps, estimator=name, outcome=\"smoker_mean\", time=\"year\",\n unit=\"state\", treatment_first=\"g\",\n design_columns={\"psu\": \"state\"}, **extra,\n )\n rows.append({\"weights\": wlabel, \"estimator\": name,\n \"att_pp\": round(r.overall_att * 100, 2),\n \"se_pp\": round(r.overall_se * 100, 2)})\n\nsweep = pd.DataFrame(rows).pivot(index=\"estimator\", columns=\"weights\",\n values=\"att_pp\")\nprint(f\"overall ATT (pp) by estimator and weighting (realized truth: {TRUE_ATT:.2f}):\")\nprint(sweep.round(2).to_string())\n" + }, + { + "cell_type": "markdown", + "id": "cell-27", + "metadata": {}, + "source": "All three estimators tell the same story: roughly -3.8 to -4.2pp under design weights, roughly -2.9 to -3.2pp once the composition is fixed. Callaway-Sant'Anna's doubly-robust machinery, Sun-Abraham's interaction weighting, and Borusyak-Jaravel-Spiess's imputation efficiency all faithfully estimate the effect *in the data they are given*. Composition bias is a data problem; fix it in the data.\n" + }, + { + "cell_type": "markdown", + "id": "cell-28", + "metadata": {}, + "source": "## 10. Diagnostics across the seam\n\nEach package owns half the diagnostic picture. diff-diff's `survey_metadata` describes the *second stage* - how much precision the population weighting costs across states. balance's diagnostics describe each *raking cell* - how hard the calibration had to work, which is itself an early-warning signal: the design effect of the rake blows up exactly where composition was most distorted. `bd.as_balance_diagnostic` joins the two into one flat dict.\n" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-29", + "metadata": {}, + "outputs": [], + "source": "sm = res_raked.survey_metadata\nprint(f\"second stage: design_effect={sm.design_effect:.2f}, \"\n f\"effective_n={sm.effective_n:.1f} of {sm.n_psu} states, \"\n f\"df_survey={sm.df_survey}\")\n\n# per-cell raking cost: never-treated vs treated state, 2024\nnever_state = int(micro_diff.loc[micro_diff.g == 0, \"state\"].iloc[0])\ntreated_state = int(micro_diff.loc[micro_diff.g == 2020, \"state\"].iloc[0])\nfor label, st in [(\"never-treated\", never_state), (\"2020-cohort\", treated_state)]:\n diag = bd.as_balance_diagnostic(adj_cell[(st, 2024)], res_adapter)\n print(f\"\\n{label} state {st}, 2024 raking cell:\")\n print(f\" balance_design_effect: {diag['balance_design_effect']:.2f}\")\n print(f\" balance_kish_ess: {diag['balance_kish_ess']:.0f}\")\n print(f\" balance_asmd_max_post: {diag['balance_asmd_max_post']:.4f}\")\n print(f\" att (full panel fit): {diag['att'] * 100:+.2f}pp\")\n" + }, + { + "cell_type": "markdown", + "id": "cell-30", + "metadata": {}, + "source": "Read the contrast: in a never-treated state the 2024 rake barely works (design effect near 1), while in a 2020-cohort state it pays a large design effect to rebuild the missing low-education mass - and post-adjustment imbalance (`asmd_max_post`) is ~0 in both, confirming the margins were hit. **Rising per-cell raking design effects concentrated in treated-post cells are the operational fingerprint of differential drift.**\n\nThe practical monitoring rule this tutorial suggests: alongside your pre-trend plot, always plot **design-weighted demographic shares by cohort over time** (the section 5 plot) and the **per-cell raking design effects**. Both are one groupby away, and they catch what the pre-trend test structurally cannot.\n" + }, + { + "cell_type": "markdown", + "id": "cell-31", + "metadata": {}, + "source": "## 11. When do you need calibration? A checklist\n\n| Your estimand | Drift type | What you need |\n|---|---|---|\n| Descriptive (levels, trends) | any | Calibrated weights, per wave ([balance tutorial](https://github.com/facebookresearch/balance/blob/main/tutorials/balance_diff_diff_brfss.ipynb)) |\n| Causal DiD | common across arms | Design weights suffice for the point estimate; calibration still improves descriptives and SEs |\n| Causal DiD | **correlated with treatment timing** | **Calibration at the granularity of your comparison units** (this tutorial) |\n\nThree caveats that decide whether raking is *enough*:\n\n- **Margins must come from a policy-unaffected source.** We raked to ACS-style demographics. If the policy itself changed the margins you rake to, calibration bakes the effect away.\n- **Raking fixes drift on observables inside your margins.** Here non-response depended only on `age_band` x `educ_cat` - missing-at-random given the margins - so raking was exact. If response depends on the *outcome itself within* demographic cells (heavier smokers refusing regardless of education), no reweighting on demographics can fix it. That regime is the subject of Sant'Anna & Xu (2023), \"Difference-in-Differences with Compositional Changes\" ([arXiv:2304.13925](https://arxiv.org/abs/2304.13925)) - a model-based route diff-diff tracks on its roadmap.\n- **Drift can pre-date adoption.** Our mechanism switched on at the effective date, which is why the pre-trends stayed clean. As an exercise, rerun with `simulate_brfss_smoking(differential=True, drift_start_offset=-2)` - drift beginning with the legislative campaign two years early - and watch the design-weight pre-treatment coefficients light up. When your event study *does* flag pre-trends in survey data, composition drift belongs on the suspect list right next to genuine trend violations.\n" + }, + { + "cell_type": "markdown", + "id": "cell-32", + "metadata": {}, + "source": "## Summary\n\n**Key takeaways:**\n\n1. Survey weights matter for causal estimands, not just descriptive ones - but only under the right failure mode. Common non-response drift differences out of a DiD; drift **correlated with treatment timing does not** (it is treatment x post shaped, so the estimator reads it as effect).\n2. In our BRFSS-style smoking-ban setting, differential drift inflated a -2.98pp (realized) true ATT to ~-4.1pp with a fake growing dynamic profile - **while every pre-trend test stayed clean**. Pre-trend tests certify who you measured, not the population.\n3. **Raking granularity must match your comparison units.** A per-wave national rake made things *worse* (margins met in aggregate, arm-level composition untouched); raking each state-year to population-count totals - BRFSS's own practice - recovered ~-3.2pp.\n4. The diff-diff/balance seam works both ways and agrees exactly: native (`SurveyDesign` + `aggregate_survey` + `fit(survey_design=...)`) or the `balance.interop.diff_diff` adapter (`to_panel_for_did` / `fit_did`, with `_balance_adjustment` provenance).\n5. No estimator choice rescues distorted data: CS, Sun-Abraham, and ImputationDiD were all equally biased before raking and all recovered after.\n6. Monitor **demographic shares by cohort** and **per-cell raking design effects** alongside pre-trends; they are the early-warning system for composition drift.\n\n**Quick reference** - the whole workflow:\n\n```python\n# 1. rake each comparison cell to population margins (balance)\nadj = (Sample.from_frame(cell_df, id_column=\"id\", weight_column=\"design_weight\")\n .set_target(acs_sample).adjust(method=\"rake\", variables=[\"age_band\", \"educ_cat\"]))\n\n# 2. hand the calibrated weights to diff-diff (native seam)\ndesign = SurveyDesign(weights=\"w_raked\", strata=\"stratum\", psu=\"psu\", fpc=\"fpc\")\npanel, stage2 = aggregate_survey(micro, by=[\"state\", \"year\"], outcomes=\"smoker\",\n survey_design=design)\nresult = CallawaySantAnna().fit(panel, outcome=\"smoker_mean\", unit=\"state\",\n time=\"year\", first_treat=\"g\", survey_design=stage2)\n```\n\n**Related tutorials:** [Tutorial 16: Survey DiD](16_survey_did.ipynb) (SurveyDesign in depth: replicate weights, subpopulations, DEFF) - [Tutorial 17: Brand Awareness Survey](17_brand_awareness_survey.ipynb) - [Tutorial 22: HAD Survey-Weighted Workflow](22_had_survey_design.ipynb) - [balance quickstart](https://import-balance.org/docs/tutorials/quickstart/) and the [balance x diff-diff BRFSS tutorial](https://github.com/facebookresearch/balance/blob/main/tutorials/balance_diff_diff_brfss.ipynb).\n\n**References:**\n\n- Callaway, B. & Sant'Anna, P. H. C. (2021). \"Difference-in-Differences with Multiple Time Periods.\" *Journal of Econometrics*, 225(2), 200-230.\n- Sant'Anna, P. H. C. & Xu, Q. (2023). \"Difference-in-Differences with Compositional Changes.\" [arXiv:2304.13925](https://arxiv.org/abs/2304.13925).\n- Deville, J.-C. & Särndal, C.-E. (1992). \"Calibration Estimators in Survey Sampling.\" *JASA*, 87(418), 376-382.\n- Deming, W. E. & Stephan, F. F. (1940). \"On a Least Squares Adjustment of a Sampled Frequency Table When the Expected Marginal Totals are Known.\" *Annals of Mathematical Statistics*, 11(4), 427-444.\n- Groves, R. M. & Peytcheva, E. (2008). \"The Impact of Nonresponse Rates on Nonresponse Bias: A Meta-Analysis.\" *Public Opinion Quarterly*, 72(2), 167-189.\n- Solon, G., Haider, S. J. & Wooldridge, J. M. (2015). \"What Are We Weighting For?\" *Journal of Human Resources*, 50(2), 301-316.\n- Lumley, T. (2004). \"Analysis of Complex Survey Samples.\" *Journal of Statistical Software*, 9(8).\n- Sarig, T., Galili, T. & Eilat, R. (2023). \"balance - a Python package for balancing biased data samples.\" [arXiv:2307.06024](https://arxiv.org/abs/2307.06024).\n" + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/tutorials/README.md b/docs/tutorials/README.md index ebd543a0..6b7154b1 100644 --- a/docs/tutorials/README.md +++ b/docs/tutorials/README.md @@ -135,6 +135,15 @@ Power-analysis decision guide for geo experiments (framed on a 50-state staggere - When a clean-tail 2×2 is unbiased, the small-holdout and few-clusters caveats, and a CS-vs-2×2 decision guide - Fully self-contained: runs live (no committed data files) +### 26. Composition Drift & Survey Calibration with balance (`26_composition_drift_calibration.ipynb`) +The failure-mode companion to Meta balance's `balance_diff_diff_brfss` tutorial: when non-response drift correlates with treatment timing, the design-weight DiD itself is biased and calibration becomes essential for the *causal* estimand: +- BRFSS-style smoking-ban DGP with no systematic arm-specific trends (parallel trends hold in expectation; planted ATT -3.0pp, realized -2.98pp) and treatment-correlated non-response drift +- Design-weight Callaway-Sant'Anna overstates the ATT ~35% with *clean pre-trends* (pre-trend tests are not a safety net) +- A per-wave national rake fails (margins satisfied in aggregate, arm-level composition untouched); per-state raking - BRFSS's own granularity - recovers the truth +- The seam both ways: native `SurveyDesign` + `aggregate_survey` 3-liner vs the `balance.interop.diff_diff` adapter (`to_panel_for_did` / `fit_did`), with an exact-parity assert +- Estimator sweep (CS / SunAbraham / ImputationDiD), `survey_metadata` DEFF diagnostics, and `as_balance_diagnostic` cross-package diagnostics +- Requires `pip install "balance>=0.21"` (this tutorial only); fully self-contained data + ## Running the Notebooks 1. Install diff-diff with dependencies: diff --git a/tests/test_balance_interop_contract.py b/tests/test_balance_interop_contract.py new file mode 100644 index 00000000..bc872b76 --- /dev/null +++ b/tests/test_balance_interop_contract.py @@ -0,0 +1,321 @@ +"""Contract tests for the surface consumed by ``balance.interop.diff_diff``. + +Meta's `balance` package (>=0.21) ships a one-way adapter, +``balance.interop.diff_diff`` (facebookresearch/balance PR #465), whose +``pip install balance[did]`` extra pins ``diff-diff>=3.3.0,<4``. The adapter +depends on specific diff-diff behaviors beyond the flat result aliases +already pinned in ``tests/test_result_aliases.py``: + +- ``diff_diff.aggregate_survey`` (top-level export): forwarded params and + the ``(panel_df, second_stage_design)`` return pair + (``to_panel_for_did`` wraps it verbatim); +- ``SurveyDesign`` field names (``balance.interop.conventions`` + hard-codes ``DEFAULT_DESIGN_COLUMNS`` and the adapter validates + ``design_columns=`` overrides against the dataclass fields); +- estimator resolution by class name and short alias via + ``getattr(diff_diff, name)`` (``fit_did``'s dispatch), with ``fit()`` + accepting ``survey_design=``; +- results accepting attribute attachment (``fit_did`` sets a + ``_balance_adjustment`` provenance side-channel via ``setattr``); +- the pweight-only guard on staggered estimators (the adapter defaults + ``weight_type="pweight"`` because of it); +- ``SurveyMetadata`` attribute names read by ``as_balance_diagnostic`` + (``design_effect`` / ``effective_n`` / ``sum_weights``) via defensive + ``getattr(..., None)`` - a rename would silently NULL balance's + diagnostics rather than raise. + +These tests intentionally import NO balance code: they pin OUR public +surface so a diff-diff refactor cannot silently break ``balance[did]``. +""" + +import dataclasses +import inspect + +import numpy as np +import pandas as pd +import pytest + +import diff_diff +from diff_diff import CallawaySantAnna, SurveyDesign, aggregate_survey + +# --------------------------------------------------------------------------- +# Shared tiny survey micro-frame + one fitted survey CS result +# --------------------------------------------------------------------------- + +N_UNITS = 12 +YEARS = [2019, 2020, 2021, 2022] + + +def _make_micro(seed=0): + """Respondent-level microdata: 12 units x 4 years x 40 respondents.""" + rng = np.random.default_rng(seed) + n_per = 40 + unit = np.repeat(np.arange(N_UNITS), len(YEARS) * n_per) + year = np.tile(np.repeat(YEARS, n_per), N_UNITS) + g = np.where(unit < 4, 2021, 0)[np.arange(len(unit))] + y = ( + 1.0 + + 0.1 * (year - YEARS[0]) + + 0.5 * unit / N_UNITS + - 1.0 * ((g > 0) & (year >= g)) + + rng.normal(0, 0.5, len(unit)) + ) + return pd.DataFrame( + { + "unit": unit, + "year": year, + "g": g, + "y": y, + "w": rng.uniform(0.5, 2.0, len(unit)), + "stratum": unit % 3, + "psu": unit * 10 + rng.integers(0, 4, len(unit)), + } + ) + + +@pytest.fixture(scope="module") +def micro(): + return _make_micro() + + +@pytest.fixture(scope="module") +def panel_and_design(micro): + design = SurveyDesign(weights="w", strata="stratum", psu="psu") + return aggregate_survey(micro, by=["unit", "year"], outcomes="y", survey_design=design) + + +@pytest.fixture(scope="module") +def fitted_cs(panel_and_design, micro): + panel, second_stage = panel_and_design + panel = panel.merge(micro[["unit", "g"]].drop_duplicates(), on="unit", how="left") + cs = CallawaySantAnna(estimation_method="reg", base_period="universal") + return cs.fit( + panel, + outcome="y_mean", + unit="unit", + time="year", + first_treat="g", + survey_design=second_stage, + ) + + +# --------------------------------------------------------------------------- +# 1-2. aggregate_survey: signature superset + return contract +# --------------------------------------------------------------------------- + +# Params to_panel_for_did forwards verbatim (balance/interop/diff_diff.py, +# aggregate_survey call). The adapter routes lonely_psu into the +# first-stage SurveyDesign, NOT into aggregate_survey. +ADAPTER_FORWARDED_PARAMS = { + "data", + "by", + "outcomes", + "survey_design", + "covariates", + "min_n", + "second_stage_weights", +} + + +def test_aggregate_survey_signature_superset(): + # Superset (not exact-set) pinning: an ADDITIVE optional param cannot + # break balance's keyword-arg adapter, so it must not break this test. + params = set(inspect.signature(aggregate_survey).parameters) + missing = ADAPTER_FORWARDED_PARAMS - params + assert not missing, ( + f"aggregate_survey lost parameter(s) {sorted(missing)} that " + "balance.interop.diff_diff.to_panel_for_did forwards verbatim." + ) + + +def test_aggregate_survey_return_pair_and_panel_schema(panel_and_design): + panel, second_stage = panel_and_design + assert isinstance(panel, pd.DataFrame) + assert isinstance(second_stage, SurveyDesign) + # to_panel_for_did documents the {outcome}_mean/_se/_n/_precision cells + # and wires the second-stage design's weights column into the panel. + for col in ["y_mean", "y_se", "y_n", "y_precision"]: + assert col in panel.columns, f"panel lost column {col}" + assert second_stage.weights in panel.columns + # The first `by` element becomes the second-stage clustering variable. + assert second_stage.psu == "unit" + + +# --------------------------------------------------------------------------- +# 3. SurveyDesign field-name contract +# --------------------------------------------------------------------------- + +# Normative list: balance/interop/conventions.py DEFAULT_DESIGN_COLUMNS plus +# the adapter's _ALLOWED_DESIGN_FIELDS validate against these exact names. +SURVEY_DESIGN_FIELDS = { + "weights", + "strata", + "psu", + "fpc", + "weight_type", + "nest", + "lonely_psu", + "replicate_weights", + "replicate_method", + "replicate_strata", + "fay_rho", + "combined_weights", + "replicate_scale", + "replicate_rscales", + "mse", +} + + +def test_survey_design_field_names_exact(): + fields = {f.name for f in dataclasses.fields(SurveyDesign)} + assert fields == SURVEY_DESIGN_FIELDS, ( + "SurveyDesign dataclass fields changed " + f"(removed: {sorted(SURVEY_DESIGN_FIELDS - fields)}, " + f"added: {sorted(fields - SURVEY_DESIGN_FIELDS)}). The exact pin is " + "intentional friction: balance's adapter enumerates these names in " + "conventions.py/_ALLOWED_DESIGN_FIELDS, so ANY change (adds included) " + "should be consciously synced with the balance maintainers before " + "updating this list. Removed/renamed fields break balance[did] " + "outright." + ) + + +def test_survey_design_tsl_construction(): + # The TSL combo to_survey_design builds by default (auto-wired + # stratum/psu/fpc convention columns + the adapter's lonely_psu default). + design = SurveyDesign( + weights="w", + strata="stratum", + psu="psu", + fpc="fpc", + nest=True, + lonely_psu="adjust", + ) + assert design.weight_type == "pweight" # adapter's documented default + + +def test_survey_design_replicate_construction(): + # Replicate combo (mutually exclusive with strata/psu/fpc). + design = SurveyDesign( + weights="w", + replicate_weights=["rep_1", "rep_2"], + replicate_method="JK1", + ) + assert design.replicate_method == "JK1" + + +# --------------------------------------------------------------------------- +# 4. Estimator resolution by name and alias +# --------------------------------------------------------------------------- + +# The 17 estimator names promised in balance/interop/diff_diff.py's module +# docstring ("weight_type='pweight' is ... compatible with ..."). +ADAPTER_DOCSTRING_ESTIMATORS = [ + "CallawaySantAnna", + "StackedDiD", + "ImputationDiD", + "HeterogeneousAdoptionDiD", + "TwoStageDiD", + "WooldridgeDiD", + "TROP", + "StaggeredTripleDifference", + "ChaisemartinDHaultfoeuille", + "TripleDifference", + "SyntheticDiD", + "EfficientDiD", + "DifferenceInDifferences", + "TwoWayFixedEffects", + "MultiPeriodDiD", + "SunAbraham", + "ContinuousDiD", +] + +# Short aliases promised in fit_did's docstring; removing one breaks +# balance's documented examples (verified exports in diff_diff/__init__.py). +ADAPTER_DOCSTRING_ALIASES = { + "CS": "CallawaySantAnna", + "DiD": "DifferenceInDifferences", + "BJS": "ImputationDiD", + "HAD": "HeterogeneousAdoptionDiD", +} + + +@pytest.mark.parametrize("name", ADAPTER_DOCSTRING_ESTIMATORS) +def test_estimator_resolves_with_survey_design_fit(name): + cls = getattr(diff_diff, name, None) + assert cls is not None and inspect.isclass(cls), ( + f"diff_diff.{name} is no longer an exported class; " + "balance.interop.diff_diff.fit_did resolves estimators via " + "getattr(diff_diff, name)." + ) + fit_params = set(inspect.signature(cls.fit).parameters) + assert "survey_design" in fit_params, ( + f"{name}.fit() no longer accepts survey_design=; fit_did would " + "warn and silently run the fit without the balance-built design." + ) + + +@pytest.mark.parametrize("alias,target", sorted(ADAPTER_DOCSTRING_ALIASES.items())) +def test_estimator_alias_resolves(alias, target): + cls = getattr(diff_diff, alias, None) + assert cls is not None, f"short alias diff_diff.{alias} was removed" + assert cls.__name__ == target + + +# --------------------------------------------------------------------------- +# 5. Provenance side-channel: setattr on results must keep working +# --------------------------------------------------------------------------- + + +def test_result_accepts_provenance_attribute(fitted_cs): + # fit_did(preserve_adjustment=True) attaches the balance Sample via + # setattr(results, "_balance_adjustment", sample). A future __slots__ + # on result dataclasses would break this silently for balance users. + sentinel = object() + fitted_cs._balance_adjustment = sentinel + assert fitted_cs._balance_adjustment is sentinel + del fitted_cs._balance_adjustment + # ...while the flat aliases stay read-only properties (full alias + # coverage lives in tests/test_result_aliases.py - not duplicated here). + with pytest.raises(AttributeError): + fitted_cs.att = 0.0 + + +# --------------------------------------------------------------------------- +# 6. pweight-only guard (why the adapter defaults weight_type="pweight") +# --------------------------------------------------------------------------- + + +def test_cs_rejects_fweight_design(micro): + data = micro.copy() + data["w_int"] = 2 # fweights must be non-negative integers + design = SurveyDesign(weights="w_int", weight_type="fweight") + cs = CallawaySantAnna(estimation_method="reg") + with pytest.raises(ValueError, match="pweight"): + cs.fit( + data, + outcome="y", + unit="unit", + time="year", + first_treat="g", + survey_design=design, + ) + + +# --------------------------------------------------------------------------- +# 7. SurveyMetadata attribute names read by as_balance_diagnostic +# --------------------------------------------------------------------------- + + +def test_survey_metadata_attribute_contract(fitted_cs): + sm = fitted_cs.survey_metadata + assert sm is not None, ( + "survey-fitted CallawaySantAnnaResults.survey_metadata is None; " + "as_balance_diagnostic would silently report None diagnostics." + ) + for attr in ("design_effect", "effective_n", "sum_weights"): + value = getattr(sm, attr, None) + assert value is not None and np.isfinite(value), ( + f"SurveyMetadata.{attr} missing or non-finite; balance reads it " + "via getattr(sm, ..., None) and would silently emit None." + ) diff --git a/tests/test_openai_review.py b/tests/test_openai_review.py index ba4b8f96..17cb9512 100644 --- a/tests/test_openai_review.py +++ b/tests/test_openai_review.py @@ -2127,7 +2127,7 @@ class TestCiWorkflowLabelEventGuard: # fails the test, even in multi-job files. EXPECTED_JOBS = { "rust-test.yml": ("rust-tests", "python-tests", "python-fallback"), - "notebooks.yml": ("execute-notebooks",), + "notebooks.yml": ("execute-notebooks", "interop-notebooks"), "docs-tests.yml": ("doc-snippets", "sphinx-build", "docs-deps-py39-smoke"), # release-build-check.yml is a single reusable-workflow caller job gated on # ready-for-ci (it build-tests the PyPI release path on PRs); lock its guard too. diff --git a/tests/test_t26_composition_drift_calibration_drift.py b/tests/test_t26_composition_drift_calibration_drift.py new file mode 100644 index 00000000..e35fcf3d --- /dev/null +++ b/tests/test_t26_composition_drift_calibration_drift.py @@ -0,0 +1,452 @@ +"""Drift detection for Tutorial 26 +(``docs/tutorials/26_composition_drift_calibration.ipynb``). + +The tutorial narrative quotes seed-specific numbers (planted ATT -3.0pp +with realized population ATT ~-2.98pp; design-weight CS ~-4.1pp with +clean pre-trends; national per-wave rake ~-4.4pp as the "false fix"; +state-year rake ~-3.2pp as the recovery; 2024 composition shares). If library numerics drift, the prose can go +stale silently while ``pytest --nbmake`` still passes - it only checks +that cells execute. These asserts re-derive the headline numbers using +the locked T26 DGP duplicated below (verbatim from the notebook SS2 code +cell) and check them against tolerance bands around the quoted values. + +Requires the ``balance`` package (>=0.21) for the raking acts - the +whole module skips when balance is absent (main-suite CI legs). It DOES +run in the ``interop-notebooks`` CI job, where balance is installed; +that job is this guard's CI home. + +The DGP-builder constants below MUST stay in sync with the notebook SS2 +code cell; ``test_notebook_dgp_constants_match`` catches silent drift on +those values. +""" + +from __future__ import annotations + +import logging +import warnings + +import numpy as np +import pandas as pd +import pytest + +balance = pytest.importorskip( + "balance", + minversion="0.21", + reason="balance>=0.21 required (interop-notebooks CI job / local)", +) +from balance import Sample # noqa: E402 +from balance.interop import diff_diff as bd # noqa: E402 + +from diff_diff import CallawaySantAnna, SurveyDesign, aggregate_survey # noqa: E402 + +logging.getLogger("balance").setLevel(logging.ERROR) + +TRUE_ATT_QUOTED = -3.0 # pp, planted coefficient (realized ATT ~-2.98pp after floor) + +# --------------------------------------------------------------------------- +# Locked DGP - duplicated verbatim from the notebook SS2 cell. Keep in sync. +# --------------------------------------------------------------------------- + +N_STATES = 50 +YEARS = np.arange(2018, 2025) +N_INVITED = 1200 +N_STRATA = 5 +PSUS_PER_STATE = 8 +FPC_PSUS_PER_STRATUM = 200.0 + +AGE_BANDS = ["18-34", "35-49", "50-64", "65+"] +AGE_SHARES = np.array([0.30, 0.25, 0.25, 0.20]) +EDUC_CATS = ["hs_or_less", "some_college", "college_plus"] +EDUC_SHARES = np.array([0.35, 0.30, 0.35]) + +BASE_EDUC_PP = np.array([22.0, 15.0, 9.0]) +AGE_ADJ_PP = np.array([2.0, 3.0, 1.0, -2.0]) +TREND_COMMON_PP = 0.25 +TREND_EDUC_PP = np.array([-0.15, 0.0, 0.10]) +TRUE_ATT_PP = -3.0 +STATE_RE_SD_PP = 1.5 +PSU_SHOCK_SD_PP = 0.8 +P_CLIP_PP = (1.0, 60.0) + +R_BASE = 0.70 +R_AGE_SHIFT = np.array([-0.10, -0.02, 0.03, 0.08]) +R_EDUC_SHIFT = np.array([-0.09, 0.00, 0.06]) +R_COMMON_DRIFT_EDUC = np.array([0.015, 0.0075, 0.0]) +R_COMMON_DRIFT_YOUNG = 0.010 +R_DIFF_DRIFT_PER_EVENT_YEAR = 0.07 +R_CLIP = (0.10, 0.95) + +TARGET_N = 20_000 +SEED = 20260704 + +RAKE_VARS = ["age_band", "educ_cat"] + + +def simulate_brfss_smoking(differential, seed=SEED, drift_start_offset=0): + """Duplicated from the notebook SS2 code cell. Keep in sync. + + No arm-specific trends by construction: population parallel trends hold + in expectation (mean-zero PSU-year shocks add noise, not drift); the + planted effect is -3.0pp (realized population ATT ~-2.98pp after the + probability floor). All SYSTEMATIC estimator bias comes from sample + composition. + """ + rng = np.random.default_rng(seed) + + perm = rng.permutation(N_STATES) + g_of_state = np.zeros(N_STATES, dtype=int) + g_of_state[perm[:10]] = 2020 + g_of_state[perm[10:20]] = 2022 + stratum_of_state = rng.integers(0, N_STRATA, size=N_STATES) + state_pop = rng.lognormal(mean=np.log(4e6), sigma=0.6, size=N_STATES) + state_re = np.clip(rng.normal(0.0, STATE_RE_SD_PP, size=N_STATES), -3.0, 3.0) + psu_shock = rng.normal(0.0, PSU_SHOCK_SD_PP, size=(N_STATES, PSUS_PER_STATE, len(YEARS))) + + n_inv = N_STATES * len(YEARS) * N_INVITED + state = np.repeat(np.arange(N_STATES), len(YEARS) * N_INVITED) + year = np.tile(np.repeat(YEARS, N_INVITED), N_STATES) + age_idx = rng.choice(len(AGE_BANDS), size=n_inv, p=AGE_SHARES) + educ_idx = rng.choice(len(EDUC_CATS), size=n_inv, p=EDUC_SHARES) + psu_idx = rng.integers(0, PSUS_PER_STATE, size=n_inv) + u_respond = rng.uniform(size=n_inv) + u_smoker = rng.uniform(size=n_inv) + weight_jitter = rng.uniform(0.85, 1.15, size=n_inv) + + k = year - YEARS[0] + year_idx = year - YEARS[0] + g = g_of_state[state] + treated_post = (g > 0) & (year >= g) + + base_pp = ( + BASE_EDUC_PP[educ_idx] + + AGE_ADJ_PP[age_idx] + + state_re[state] + + psu_shock[state, psu_idx, year_idx] + - (TREND_COMMON_PP + TREND_EDUC_PP[educ_idx]) * k + ) + p_pp = np.clip(base_pp + TRUE_ATT_PP * treated_post, *P_CLIP_PP) + smoker = (u_smoker < p_pp / 100.0).astype(int) + + r = ( + R_BASE + + R_AGE_SHIFT[age_idx] + + R_EDUC_SHIFT[educ_idx] + - R_COMMON_DRIFT_EDUC[educ_idx] * k + - R_COMMON_DRIFT_YOUNG * k * (age_idx == 0) + ) + if differential: + event_time = year - g - drift_start_offset + hit = (g > 0) & (event_time >= 0) & (educ_idx == 0) + r = r - R_DIFF_DRIFT_PER_EVENT_YEAR * (event_time + 1) * hit + r = np.clip(r, *R_CLIP) + responded = u_respond < r + + micro = pd.DataFrame( + { + "id": np.arange(n_inv)[responded], + "state": state[responded], + "year": year[responded], + "g": g[responded], + "smoker": smoker[responded], + "age_band": np.array(AGE_BANDS)[age_idx[responded]], + "educ_cat": np.array(EDUC_CATS)[educ_idx[responded]], + "stratum": stratum_of_state[state[responded]], + "psu": state[responded] * 100 + psu_idx[responded], + "fpc": FPC_PSUS_PER_STRATUM, + "design_weight": (state_pop[state] / N_INVITED * weight_jitter)[responded], + } + ) + + rng_t = np.random.default_rng(seed + 1) + target_df = pd.DataFrame( + { + "id": np.arange(TARGET_N), + "age_band": np.array(AGE_BANDS)[ + rng_t.choice(len(AGE_BANDS), size=TARGET_N, p=AGE_SHARES) + ], + "educ_cat": np.array(EDUC_CATS)[ + rng_t.choice(len(EDUC_CATS), size=TARGET_N, p=EDUC_SHARES) + ], + } + ) + + kk = YEARS - YEARS[0] + cell = ( + BASE_EDUC_PP[None, :, None] + + AGE_ADJ_PP[None, None, :] + - (TREND_COMMON_PP + TREND_EDUC_PP[None, :, None]) * kk[:, None, None] + ) + tp = (g_of_state[None, :] > 0) & (YEARS[:, None] >= g_of_state[None, :]) + base_prev = np.einsum("tea,e,a->t", cell, EDUC_SHARES, AGE_SHARES) + w_s = state_pop / state_pop.sum() + pop_prev = base_prev + TRUE_ATT_PP * (tp * w_s[None, :]).sum(axis=1) + # Realized population ATT: the probability floor P_CLIP_PP[0] binds for + # ~2% of treated-post person-years, attenuating the planted -3.0pp. + y1 = np.clip(base_pp + TRUE_ATT_PP, *P_CLIP_PP) + y0 = np.clip(base_pp, *P_CLIP_PP) + w_pop = state_pop[state] + realized_att_pp = ((y1 - y0) * w_pop)[treated_post].sum() / w_pop[treated_post].sum() + truth = { + "true_att_pp": TRUE_ATT_PP, + "realized_att_pp": float(realized_att_pp), + "floor_bind_share": float(((y1 - y0) > TRUE_ATT_PP + 1e-12)[treated_post].mean()), + "pop_prevalence_by_year": dict(zip(YEARS.tolist(), pop_prev / 100.0)), + "g_of_state": g_of_state, + "state_pop": state_pop, + } + return micro, target_df, truth + + +def rake_to_population(micro, target_df, granularity, weight_name, cell_totals): + """Duplicated verbatim from the notebook SS5 code cell. Keep in sync.""" + target_sample = Sample.from_frame(target_df, id_column="id") + cols = ["id", *RAKE_VARS, "smoker", "design_weight"] + w_new = pd.Series(np.nan, index=micro.index) + adjusted = {} + for key, cell in micro.groupby(granularity): + if isinstance(key, tuple) and len(key) == 1: + key = key[0] + s = Sample.from_frame( + cell[cols].copy(), + id_column="id", + weight_column="design_weight", + outcome_columns=["smoker"], + ) + adj = s.set_target(target_sample).adjust(method="rake", variables=RAKE_VARS) + w = adj.df.set_index("id")[adj.weight_column] + aligned = w.reindex(cell["id"].astype(str).values).to_numpy() + assert not np.isnan(aligned).any(), f"NaN raked weights in cell {key}" + aligned = aligned * (cell_totals[key] / aligned.sum()) + w_new.loc[cell.index] = aligned + adjusted[key] = adj + out = micro.copy() + out[weight_name] = w_new + return out, adjusted + + +def fit_survey_cs(micro, weights_col): + """Duplicated from the notebook SS3 code cell (native seam). Keep in sync.""" + design = SurveyDesign(weights=weights_col, strata="stratum", psu="psu", fpc="fpc") + panel, second_stage = aggregate_survey( + micro, by=["state", "year"], outcomes="smoker", survey_design=design + ) + panel = panel.merge(micro[["state", "g"]].drop_duplicates(), on="state", how="left") + cs = CallawaySantAnna( + estimation_method="reg", + control_group="not_yet_treated", + base_period="universal", + ) + return cs.fit( + panel, + outcome="smoker_mean", + unit="state", + time="year", + first_treat="g", + survey_design=second_stage, + aggregate="all", + ) + + +# --------------------------------------------------------------------------- +# Shared pipeline run (module-scoped: the rakes dominate the ~10s runtime) +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def pipeline(): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + micro_a, _, _ = simulate_brfss_smoking(differential=False) + micro, target, truth = simulate_brfss_smoking(differential=True) + state_pop = truth["state_pop"] + res_a = fit_survey_cs(micro_a, "design_weight") + res_design = fit_survey_cs(micro, "design_weight") + micro, _ = rake_to_population( + micro, + target, + ["year"], + "w_national", + cell_totals={int(y): state_pop.sum() for y in YEARS}, + ) + res_national = fit_survey_cs(micro, "w_national") + micro, _ = rake_to_population( + micro, + target, + ["state", "year"], + "w_raked", + cell_totals={(st, int(y)): state_pop[st] for st in range(N_STATES) for y in YEARS}, + ) + res_raked = fit_survey_cs(micro, "w_raked") + return { + "micro": micro, + "truth": truth, + "A_design": res_a, + "B_design": res_design, + "B_national": res_national, + "B_raked": res_raked, + } + + +def test_planted_and_realized_truth(pipeline): + # Planted coefficient is exactly -3.0pp; the 1pp probability floor binds + # for ~2% of treated-post person-years, so the realized population ATT + # (quoted throughout the tutorial as the truth line) is ~-2.98pp. + truth = pipeline["truth"] + assert truth["true_att_pp"] == TRUE_ATT_QUOTED + assert ( + -3.0 <= truth["realized_att_pp"] <= -2.95 + ), f"realized ATT drifted: {truth['realized_att_pp']:.4f}pp" + assert 0.005 <= truth["floor_bind_share"] <= 0.05 + + +def test_scenario_a_design_weights_robust(pipeline): + att = pipeline["A_design"].overall_att * 100 + assert -3.5 <= att <= -2.7, f"scenario A design ATT drifted: {att:.2f}pp" + + +def test_scenario_b_design_weights_overstate(pipeline): + res = pipeline["B_design"] + att = res.overall_att * 100 + assert -4.6 <= att <= -3.7, f"scenario B design ATT drifted: {att:.2f}pp" + # The quoted story: the design-weight CI EXCLUDES the realized truth + # (~-2.98pp; -2.95 is the upper edge of the realized-truth band pinned + # in test_planted_and_realized_truth). + hi = res.overall_conf_int[1] + assert hi * 100 < -2.95, ( + f"design-weight CI upper bound {hi*100:.2f}pp no longer excludes " + "the realized truth (~-2.98pp)" + ) + # Pre-trends stay clean (drift starts at adoption): |pre| below 1.5pp. + es = res.event_study_effects + max_pre = max(abs(es[e]["effect"]) * 100 for e in es if e < -1) + assert max_pre < 1.5, f"pre-trend coefficient drifted: {max_pre:.2f}pp" + + +def test_national_rake_is_not_a_fix(pipeline): + att_nat = pipeline["B_national"].overall_att * 100 + att_des = pipeline["B_design"].overall_att * 100 + realized = pipeline["truth"]["realized_att_pp"] + assert -4.9 <= att_nat <= -4.0, f"national-rake ATT drifted: {att_nat:.2f}pp" + # Quoted lesson: national raking does NOT move the estimate toward truth. + assert abs(att_nat - realized) >= abs(att_des - realized) - 0.1 + + +def test_state_rake_recovers_truth(pipeline): + res = pipeline["B_raked"] + att = res.overall_att * 100 + se = res.overall_se * 100 + realized = pipeline["truth"]["realized_att_pp"] + assert -3.6 <= att <= -2.8, f"state-rake ATT drifted: {att:.2f}pp" + assert ( + abs(att - realized) <= 2 * se + ), f"state-rake ATT {att:.2f}pp not within 2 SE ({se:.2f}) of realized truth" + + +def test_composition_shares_2024(pipeline): + micro = pipeline["micro"] + m24 = micro[micro.year == 2024].assign(hs=lambda d: (d.educ_cat == "hs_or_less").astype(float)) + + def share(sub, wcol): + return (sub.hs * sub[wcol]).sum() / sub[wcol].sum() + + treated = m24[m24.g == 2020] + never = m24[m24.g == 0] + # Quoted: treated design-weighted hs share collapses (~0.11)... + assert share(treated, "design_weight") < 0.15 + # ...and state-year raking restores BOTH arms to the 0.35 margin. + assert abs(share(treated, "w_raked") - 0.35) < 0.02 + assert abs(share(never, "w_raked") - 0.35) < 0.02 + + +def test_native_adapter_parity(pipeline): + """The notebook asserts the native seam and bd.fit_did agree exactly.""" + micro = pipeline["micro"] + keep = [ + "id", + "state", + "year", + "smoker", + "age_band", + "educ_cat", + "stratum", + "psu", + "fpc", + "w_raked", + ] + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + sample = Sample.from_frame( + micro[keep].copy(), + id_column="id", + weight_column="w_raked", + outcome_columns=["smoker"], + ) + panel_df, second_stage = bd.to_panel_for_did( + sample, by=["state", "year"], outcomes="smoker" + ) + panel_df = panel_df.merge(micro[["state", "g"]].drop_duplicates(), on="state", how="left") + panel_df["panel_id"] = np.arange(len(panel_df)) + panel_sample = Sample.from_frame( + panel_df, + id_column="panel_id", + weight_column=second_stage.weights, + outcome_columns=["smoker_mean"], + ) + res_adapter = bd.fit_did( + panel_sample, + estimator="CallawaySantAnna", + outcome="smoker_mean", + time="year", + unit="state", + treatment_first="g", + design_columns={"psu": "state"}, + estimation_method="reg", + control_group="not_yet_treated", + base_period="universal", + aggregate="all", + ) + np.testing.assert_allclose(res_adapter.overall_att, pipeline["B_raked"].overall_att, rtol=1e-12) + assert hasattr(res_adapter, "_balance_adjustment") + + +def test_notebook_dgp_constants_match(): + """Sync guard: the notebook's DGP constants must match this module's + locked copy, so a notebook-only edit can't silently invalidate the + quoted numbers (t25 precedent). + + CI isolation note: CI legs that copy ``tests/`` without ``docs/`` skip + gracefully here (nbmake separately verifies execution).""" + import json + from pathlib import Path + + nb_path = ( + Path(__file__).resolve().parents[1] + / "docs" + / "tutorials" + / "26_composition_drift_calibration.ipynb" + ) + if not nb_path.exists(): + pytest.skip(f"Notebook not found at {nb_path}; sync guard is local-dev only.") + with nb_path.open() as f: + nb = json.load(f) + src = "\n".join("".join(c["source"]) for c in nb["cells"] if c["cell_type"] == "code") + for needle in ( + # locked constants + "SEED = 20260704", + "N_INVITED = 1200", + "TRUE_ATT_PP = -3.0", + "BASE_EDUC_PP = np.array([22.0, 15.0, 9.0])", + "TREND_EDUC_PP = np.array([-0.15, 0.0, 0.10])", + "R_DIFF_DRIFT_PER_EVENT_YEAR = 0.07", + "R_COMMON_DRIFT_EDUC = np.array([0.015, 0.0075, 0.0])", + 'RAKE_VARS = ["age_band", "educ_cat"]', + "TARGET_N = 20_000", + # load-bearing logic lines (mechanism, rescale, realized truth, fit) + "r = r - R_DIFF_DRIFT_PER_EVENT_YEAR * (event_time + 1) * hit", + "aligned = aligned * (cell_totals[key] / aligned.sum())", + "y1 = np.clip(base_pp + TRUE_ATT_PP, *P_CLIP_PP)", + 'estimation_method="reg"', + 'control_group="not_yet_treated"', + 'base_period="universal"', + ): + assert needle in src, f"notebook SS2 missing locked constant: {needle!r}"