Interactive notebook

This tutorial is a Jupyter notebook. You can view it on GitHub or download it to run locally.

When Who Answers Changes: Survey Calibration for Causal DiD#

Survey 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 calibration package.

A matched pair. This notebook is the diff-diff-side companion to balance’s `balance_diff_diff_brfss tutorial <facebookresearch/balance>`__. 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.

This tutorial covers:

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
Why treatment-correlated composition drift does not difference out of a DiD - with clean pre-trends the whole way
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
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
An estimator sweep (CallawaySantAnna / SunAbraham / ImputationDiD): composition bias is a data problem, not an estimator problem
Cross-package diagnostics: survey_metadata design effects and as_balance_diagnostic
A practical checklist: when calibration matters for causal vs descriptive estimands

Prerequisites: Tutorial 01: Basic DiD, Tutorial 02: Staggered DiD, Tutorial 16: Survey DiD; ideally the balance quickstart.

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.)

1. Two regimes of non-response drift#

“Do the survey weights matter?” is really two questions - one per estimand:

Non-response drift	Descriptive estimand (a prevalence trend)	Causal DiD estimand (an ATT)
Common - same in treated and control units	Biased - raking to population margins repairs it	Robust - composition changes difference out (balance’s tutorial)
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)

The 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.

[ ]:

import logging
import warnings

import numpy as np
import pandas as pd

import balance  # >= 0.21 for balance.interop.diff_diff
from balance import Sample
from balance.interop import diff_diff as bd

# Quiet balance's per-call INFO logging (must run AFTER `import balance` -
# its __init__ re-arms the logger) and pandas FutureWarnings.
logging.getLogger("balance").setLevel(logging.ERROR)
warnings.filterwarnings("ignore", category=FutureWarning)
# diff-diff normalizes pweights to mean 1 and warns to say so; expected here.
warnings.filterwarnings(
    "ignore", message=".*weights normalized.*", category=UserWarning
)

import diff_diff
from diff_diff import CallawaySantAnna, SurveyDesign, aggregate_survey

try:
    import matplotlib.pyplot as plt

    plt.style.use("seaborn-v0_8-whitegrid")
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False

print(f"diff-diff {diff_diff.__version__} | balance {balance.__version__}")

2. The data: a BRFSS-shaped survey around staggered smoking bans#

We simulate respondent-level microdata mirroring the public-use BRFSS 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.

The generator is built so every headline claim is checkable against ground truth:

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.
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.
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.
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.

Because 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_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


def simulate_brfss_smoking(differential, seed=SEED, drift_start_offset=0):
    """BRFSS-style microdata around staggered smoking bans.

    No arm-specific trends by construction: population parallel trends
    hold in expectation (mean-zero PSU-year shocks add noise, not drift),
    and the planted effect is -3.0pp (realized population ATT ~-2.98pp
    after the probability floor). All SYSTEMATIC estimator bias in this
    notebook comes from sample composition. Scenarios A/B share invited
    respondents and outcomes for the same seed - they differ only in who
    responds.
    """
    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

[ ]:

micro_common, target, truth = simulate_brfss_smoking(differential=False)
micro_diff, _, _ = simulate_brfss_smoking(differential=True)

by_cell = micro_diff.groupby(["state", "year"]).size()
print(f"scenario A respondents: {len(micro_common):,}")
print(f"scenario B respondents: {len(micro_diff):,}")
print(f"respondents per state-year (B): {by_cell.min()}-{by_cell.max()}")
TRUE_ATT = truth["realized_att_pp"]
print(f"planted ATT: {truth['true_att_pp']}pp | realized population ATT: "
      f"{TRUE_ATT:+.2f}pp (probability floor binds for "
      f"{truth['floor_bind_share']:.1%} of treated-post person-years)")
target.head(3)

3. The native seam: `SurveyDesign` + `aggregate_survey` + `CallawaySantAnna`#

diff-diff’s survey path consumes calibrated weights as just a column. The handoff from microdata to a modern staggered estimator is three calls:

``SurveyDesign(weights=…, strata=…, psu=…, fpc=…)`` - declare the complex design on the microdata (column names, no copies).
``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).
``CallawaySantAnna(…).fit(panel, …, survey_design=second_stage)`` - the staggered DiD with design-based variance.

One 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.

We 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.

[ ]:

def fit_survey_cs(micro, weights_col):
    """Native seam: microdata + a weight column -> survey-weighted CS fit."""
    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",
    )


def show_att(res, label):
    lo, hi = res.overall_conf_int
    print(
        f"{label}: ATT = {res.overall_att * 100:+.2f}pp "
        f"(SE {res.overall_se * 100:.2f}, 95% CI [{lo * 100:+.2f}, {hi * 100:+.2f}])"
    )

4. Scenario A - common drift: the descriptive trend breaks, the DiD does not#

First 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.

[ ]:

res_a = fit_survey_cs(micro_common, "design_weight")
show_att(res_a, "Scenario A, design weights")

pop_series = pd.Series(truth["pop_prevalence_by_year"])
design_series = micro_common.groupby("year")[["smoker", "design_weight"]].apply(
    lambda d: (d.smoker * d.design_weight).sum() / d.design_weight.sum()
)
gap_pp = (design_series - pop_series) * 100
print(f"descriptive gap, design-weighted vs population: "
      f"{gap_pp.loc[2018]:+.2f}pp (2018) -> {gap_pp.loc[2024]:+.2f}pp (2024)")

if HAS_MATPLOTLIB:
    fig, ax = plt.subplots(figsize=(8, 4.5))
    ax.plot(pop_series.index, pop_series * 100, "k-", lw=2, label="population (truth)")
    ax.plot(design_series.index, design_series * 100, "C1--", lw=2, marker="o",
            label="design-weighted sample")
    ax.set_xlabel("year")
    ax.set_ylabel("smoking prevalence (%)")
    ax.set_title("Scenario A: common drift biases the descriptive trend")
    ax.legend()
    plt.tight_layout()
    plt.show()

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 - we will not repeat it here.

Now we break it.

5. Scenario B - differential drift: the causal estimand breaks#

Same 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:

[ ]:

hs_share = (
    micro_diff.assign(hs=lambda d: (d.educ_cat == "hs_or_less").astype(float))
    .groupby(["year", "g"])[["hs", "design_weight"]]
    .apply(lambda d: (d.hs * d.design_weight).sum() / d.design_weight.sum())
    .unstack()
)
print("design-weighted hs_or_less share by cohort (population share: 0.350):")
print(hs_share.round(3).to_string())

if HAS_MATPLOTLIB:
    fig, ax = plt.subplots(figsize=(8, 4.5))
    styles = {0: ("C0", "never treated (30 states)"),
              2020: ("C3", "ban in 2020 (10 states)"),
              2022: ("C1", "ban in 2022 (10 states)")}
    for gval, (color, label) in styles.items():
        ax.plot(hs_share.index, hs_share[gval], color=color, marker="o", label=label)
    ax.axhline(0.35, color="gray", ls=":", label="population share (0.35)")
    for gval, color in [(2020, "C3"), (2022, "C1")]:
        ax.axvline(gval, color=color, ls="--", alpha=0.4)
    ax.set_xlabel("year")
    ax.set_ylabel("hs_or_less share of design-weighted sample")
    ax.set_title("Scenario B: who answers changes exactly when treatment starts")
    ax.legend()
    plt.tight_layout()
    plt.show()

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.

Why 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.

[ ]:

res_design = fit_survey_cs(micro_diff, "design_weight")
show_att(res_design, "Scenario B, design weights")
print(f"true ATT (realized): {TRUE_ATT:+.2f}pp")

es_design = res_design.event_study_effects
if HAS_MATPLOTLIB:
    fig, ax = plt.subplots(figsize=(8, 4.5))
    ev = sorted(es_design)
    eff = np.array([es_design[e]["effect"] for e in ev]) * 100
    se = np.array([es_design[e]["se"] for e in ev]) * 100
    ax.errorbar(ev, eff, yerr=1.96 * np.nan_to_num(se), fmt="o-", color="C1",
                capsize=3, label="design weights")
    ax.axhline(0, color="gray", lw=0.8)
    ax.axhline(TRUE_ATT, color="k", ls=":", label=f"true ATT ({TRUE_ATT:.2f}pp)")
    ax.axvline(-0.5, color="gray", ls="--", alpha=0.5)
    ax.set_xlabel("event time (years since ban)")
    ax.set_ylabel("effect on smoking prevalence (pp)")
    ax.set_title("Scenario B, design weights: clean pre-trends, biased ATT")
    ax.legend()
    plt.tight_layout()
    plt.show()

Three things just happened, and together they are the core lesson of this tutorial:

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.
The dynamics are fake: the event study “builds” to about -5pp at event year 4. That growth is the response-rate ramp, not policy.
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.

The standard toolkit passes this regression with flying colors. The composition plot above is the only thing that flagged it.

6. Fix attempt 1: per-wave national raking (the natural first move)#

The 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.

One 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.

[ ]:

RAKE_VARS = ["age_band", "educ_cat"]


def rake_to_population(micro, target_df, granularity, weight_name, cell_totals):
    """Rake design weights to population margins within each granularity cell.

    cell_totals maps each groupby key to the population COUNT the cell's
    raked weights must sum to.
    Returns (micro + weight_name column, {cell_key: adjusted balance Sample}).
    """
    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


state_pop = truth["state_pop"]

micro_diff, adj_national = rake_to_population(
    micro_diff, target, ["year"], "w_national",
    cell_totals={int(y): state_pop.sum() for y in YEARS},
)
res_national = fit_survey_cs(micro_diff, "w_national")
show_att(res_national, "Scenario B, per-wave NATIONAL rake")

m24 = micro_diff[micro_diff.year == 2024].assign(
    hs=lambda d: (d.educ_cat == "hs_or_less").astype(float)
)
for gval, label in [(0, "never treated"), (2020, "2020 cohort")]:
    sub = m24[m24.g == gval]
    share = (sub.hs * sub.w_national).sum() / sub.w_national.sum()
    print(f"2024 hs_or_less share after national rake, {label}: "
          f"{share:.3f} (population: 0.350)")

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.

Raking 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.

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.)

7. Fix attempt 2: rake each state-year to the margins#

Same 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.

[ ]:

micro_diff, adj_cell = rake_to_population(
    micro_diff, 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_diff, "w_raked")

show_att(res_design, "design weights       ")
show_att(res_national, "national rake        ")
show_att(res_raked, "state-year rake      ")
print(f"true ATT (realized): {TRUE_ATT:+.2f}pp")

m24 = micro_diff[micro_diff.year == 2024].assign(
    hs=lambda d: (d.educ_cat == "hs_or_less").astype(float)
)
for gval, label in [(0, "never treated"), (2020, "2020 cohort"), (2022, "2022 cohort")]:
    sub = m24[m24.g == gval]
    share = (sub.hs * sub.w_raked).sum() / sub.w_raked.sum()
    print(f"2024 hs_or_less share after state-year rake, {label}: {share:.3f}")

es_raked = res_raked.event_study_effects
if HAS_MATPLOTLIB:
    fig, ax = plt.subplots(figsize=(8.5, 5))
    for es, color, label in [(es_design, "C1", "design weights"),
                             (es_raked, "C0", "state-year raked")]:
        ev = sorted(es)
        eff = np.array([es[e]["effect"] for e in ev]) * 100
        se = np.array([es[e]["se"] for e in ev]) * 100
        ax.errorbar(ev, eff, yerr=1.96 * np.nan_to_num(se), fmt="o-", color=color,
                    capsize=3, label=label, alpha=0.9)
    ax.axhline(0, color="gray", lw=0.8)
    ax.axhline(TRUE_ATT, color="k", ls=":", lw=1.5,
               label=f"true ATT ({TRUE_ATT:.2f}pp)")
    ax.axvline(-0.5, color="gray", ls="--", alpha=0.5)
    ax.set_xlabel("event time (years since ban)")
    ax.set_ylabel("effect on smoking prevalence (pp)")
    ax.set_title("Composition drift manufactures dynamics; state-level raking removes them")
    ax.legend()
    plt.tight_layout()
    plt.show()

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:

weights	ATT (pp)	vs realized truth (-2.98)
design weights	~ -4.1	~+38% overstated, CI excludes truth
per-wave national rake	~ -4.4	worse - margins met in aggregate only
state-year rake, population-count totals	~ -3.2	recovered (within 0.5 SE)

8. The adapter: `balance.interop.diff_diff`#

Everything 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:

``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.
``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.

Both paths must agree exactly - and we assert it:

[ ]:

keep = ["id", "state", "year", "smoker", "age_band", "educ_cat",
        "stratum", "psu", "fpc", "w_raked"]
sample = Sample.from_frame(
    micro_diff[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"
)
print(f"panel: {panel_df.shape[0]} state-years, "
      f"second-stage weights={second_stage.weights!r}, psu={second_stage.psu!r}")

# g is unit-constant metadata: re-join it at the panel level
panel_df = panel_df.merge(
    micro_diff[["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"},  # match the native second-stage design
    estimation_method="reg",
    control_group="not_yet_treated",
    base_period="universal",
    aggregate="all",
)

assert np.isclose(res_adapter.overall_att, res_raked.overall_att, rtol=1e-12)
print(f"native  ATT: {res_raked.overall_att * 100:+.4f}pp")
print(f"adapter ATT: {res_adapter.overall_att * 100:+.4f}pp  (exact match)")
print(f"provenance attached: {hasattr(res_adapter, '_balance_adjustment')}")

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.

9. Does the fix depend on the estimator?#

A 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:

[ ]:

ESTIMATOR_KWARGS = {
    "CallawaySantAnna": dict(estimation_method="reg",
                             control_group="not_yet_treated",
                             base_period="universal", aggregate="all"),
    "SunAbraham": dict(control_group="never_treated"),
    "ImputationDiD": {},
}

rows = []
for wcol, wlabel in [("design_weight", "design"), ("w_raked", "state-raked")]:
    s = Sample.from_frame(
        micro_diff[keep[:-1] + [wcol]].copy(),
        id_column="id", weight_column=wcol, outcome_columns=["smoker"],
    )
    pdf, ss = bd.to_panel_for_did(s, by=["state", "year"], outcomes="smoker")
    pdf = pdf.merge(micro_diff[["state", "g"]].drop_duplicates(),
                    on="state", how="left")
    pdf["panel_id"] = np.arange(len(pdf))
    ps = Sample.from_frame(pdf, id_column="panel_id", weight_column=ss.weights,
                           outcome_columns=["smoker_mean"])
    for name, extra in ESTIMATOR_KWARGS.items():
        r = bd.fit_did(
            ps, estimator=name, outcome="smoker_mean", time="year",
            unit="state", treatment_first="g",
            design_columns={"psu": "state"}, **extra,
        )
        rows.append({"weights": wlabel, "estimator": name,
                     "att_pp": round(r.overall_att * 100, 2),
                     "se_pp": round(r.overall_se * 100, 2)})

sweep = pd.DataFrame(rows).pivot(index="estimator", columns="weights",
                                 values="att_pp")
print(f"overall ATT (pp) by estimator and weighting (realized truth: {TRUE_ATT:.2f}):")
print(sweep.round(2).to_string())

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.

10. Diagnostics across the seam#

Each 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.

[ ]:

sm = res_raked.survey_metadata
print(f"second stage: design_effect={sm.design_effect:.2f}, "
      f"effective_n={sm.effective_n:.1f} of {sm.n_psu} states, "
      f"df_survey={sm.df_survey}")

# per-cell raking cost: never-treated vs treated state, 2024
never_state = int(micro_diff.loc[micro_diff.g == 0, "state"].iloc[0])
treated_state = int(micro_diff.loc[micro_diff.g == 2020, "state"].iloc[0])
for label, st in [("never-treated", never_state), ("2020-cohort", treated_state)]:
    diag = bd.as_balance_diagnostic(adj_cell[(st, 2024)], res_adapter)
    print(f"\n{label} state {st}, 2024 raking cell:")
    print(f"  balance_design_effect: {diag['balance_design_effect']:.2f}")
    print(f"  balance_kish_ess:      {diag['balance_kish_ess']:.0f}")
    print(f"  balance_asmd_max_post: {diag['balance_asmd_max_post']:.4f}")
    print(f"  att (full panel fit):  {diag['att'] * 100:+.2f}pp")

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.

The 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.

11. When do you need calibration? A checklist#

Your estimand	Drift type	What you need
Descriptive (levels, trends)	any	Calibrated weights, per wave (balance tutorial)
Causal DiD	common across arms	Design weights suffice for the point estimate; calibration still improves descriptives and SEs
Causal DiD	correlated with treatment timing	Calibration at the granularity of your comparison units (this tutorial)

Three caveats that decide whether raking is enough:

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.
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) - a model-based route diff-diff tracks on its roadmap.
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.

Summary#

Key takeaways:

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).
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.
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.
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).
No estimator choice rescues distorted data: CS, Sun-Abraham, and ImputationDiD were all equally biased before raking and all recovered after.
Monitor demographic shares by cohort and per-cell raking design effects alongside pre-trends; they are the early-warning system for composition drift.

Quick reference - the whole workflow:

# 1. rake each comparison cell to population margins (balance)
adj = (Sample.from_frame(cell_df, id_column="id", weight_column="design_weight")
       .set_target(acs_sample).adjust(method="rake", variables=["age_band", "educ_cat"]))

# 2. hand the calibrated weights to diff-diff (native seam)
design = SurveyDesign(weights="w_raked", strata="stratum", psu="psu", fpc="fpc")
panel, stage2 = aggregate_survey(micro, by=["state", "year"], outcomes="smoker",
                                 survey_design=design)
result = CallawaySantAnna().fit(panel, outcome="smoker_mean", unit="state",
                                time="year", first_treat="g", survey_design=stage2)

Related tutorials: Tutorial 16: Survey DiD (SurveyDesign in depth: replicate weights, subpopulations, DEFF) - Tutorial 17: Brand Awareness Survey - Tutorial 22: HAD Survey-Weighted Workflow - balance quickstart and the balance x diff-diff BRFSS tutorial.

References:

Callaway, B. & Sant’Anna, P. H. C. (2021). “Difference-in-Differences with Multiple Time Periods.” Journal of Econometrics, 225(2), 200-230.
Sant’Anna, P. H. C. & Xu, Q. (2023). “Difference-in-Differences with Compositional Changes.” arXiv:2304.13925.
Deville, J.-C. & Särndal, C.-E. (1992). “Calibration Estimators in Survey Sampling.” JASA, 87(418), 376-382.
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.
Groves, R. M. & Peytcheva, E. (2008). “The Impact of Nonresponse Rates on Nonresponse Bias: A Meta-Analysis.” Public Opinion Quarterly, 72(2), 167-189.
Solon, G., Haider, S. J. & Wooldridge, J. M. (2015). “What Are We Weighting For?” Journal of Human Resources, 50(2), 301-316.
Lumley, T. (2004). “Analysis of Complex Survey Samples.” Journal of Statistical Software, 9(8).
Sarig, T., Galili, T. & Eilat, R. (2023). “balance - a Python package for balancing biased data samples.” arXiv:2307.06024.