Interactive notebook

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

Imputation DiD (Borusyak, Jaravel & Spiess 2024)#

This tutorial demonstrates the ImputationDiD estimator, which implements the efficient imputation approach from Borusyak, Jaravel & Spiess (2024), “Revisiting Event-Study Designs: Robust and Efficient Estimation”, Review of Economic Studies.

When to use ImputationDiD:

Staggered adoption settings where treatment effects may be homogeneous across cohorts and time — produces ~50% shorter CIs than Callaway-Sant’Anna
When you want to use all untreated observations (never-treated + not-yet-treated) for maximum efficiency
As a complement to Callaway-Sant’Anna or Sun-Abraham: if all three agree, results are robust; if they disagree, investigate heterogeneity

[ ]:

import numpy as np
import warnings
warnings.filterwarnings('ignore')

from diff_diff import (
    ImputationDiD, CallawaySantAnna, SunAbraham,
    generate_staggered_data, plot_event_study
)

# For nicer plots (optional)
try:
    import matplotlib.pyplot as plt
    plt.style.use('seaborn-v0_8-whitegrid')
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False
    print("matplotlib not installed - visualization examples will be skipped")

Basic Usage#

The imputation estimator follows a simple three-step process:

Estimate unit and time fixed effects using only untreated observations
Impute counterfactual Y(0) for treated observations
Aggregate imputed treatment effects with researcher-chosen weights

[ ]:

# Generate staggered adoption data with known treatment effect
data = generate_staggered_data(n_units=300, n_periods=10, treatment_effect=2.0, seed=42)

# Fit the imputation estimator
est = ImputationDiD()
results = est.fit(data, outcome='outcome', unit='unit', time='period', first_treat='first_treat')
results.print_summary()

Event Study with Pre-Trend Diagnostics#

Event study aggregation estimates treatment effects at each relative time horizon. Setting pretrends=True adds pre-period coefficients (negative horizons) to the event study, enabling a diagnostic check of the parallel trends assumption.

Under parallel trends, pre-period coefficients should cluster around zero — indicating no differential trends before treatment. The reference period (h = -1) is normalized to zero by construction.

[ ]:

# Fit with event study aggregation and pre-period coefficients
est = ImputationDiD(pretrends=True)
results_es = est.fit(data, outcome='outcome', unit='unit', time='period',
                     first_treat='first_treat', aggregate='event_study')

# Plot event study — pre-period region is automatically shaded
if HAS_MATPLOTLIB:
    plot_event_study(results_es, title='Imputation DiD Event Study (with Pre-Trends)')
else:
    print("Install matplotlib to see visualizations: pip install matplotlib")

[ ]:

# View event study effects as a table
results_es.to_dataframe(level='event_study')

Formal Pre-Trend Test#

The event study plot above gives a visual diagnostic — do pre-period coefficients look close to zero? For a statistical check, pretrend_test() runs a Wald F-test on whether all pre-treatment leads are jointly zero (Equation 9 in the paper). This complements the plot: the eye spots patterns, the F-test quantifies evidence consistent with parallel trends.

Note: pretrend_test() does not require pretrends=True — it runs its own internal lead regression on untreated observations, independent of the treatment effect estimator (Proposition 9). This avoids the pre-testing problem identified by Roth (2022).

[ ]:

# Run pre-trend test
pt = results.pretrend_test(n_leads=3)
print(f"F-statistic: {pt['f_stat']:.3f}")
print(f"P-value:     {pt['p_value']:.4f}")
print(f"Leads tested: {pt['n_leads']}")
print(f"\nConclusion: {'Fail to reject' if pt['p_value'] > 0.05 else 'Reject'} parallel trends at 5% level")

Comparison with Other Estimators#

Under homogeneous treatment effects, ImputationDiD, Callaway-Sant’Anna, and Sun-Abraham should produce similar point estimates. The key difference is efficiency — ImputationDiD produces shorter confidence intervals.

[ ]:

# Fit all three estimators on the same data
imp = ImputationDiD().fit(data, outcome='outcome', unit='unit',
                          time='period', first_treat='first_treat')
cs = CallawaySantAnna().fit(data, outcome='outcome', unit='unit',
                            time='period', first_treat='first_treat')
sa = SunAbraham().fit(data, outcome='outcome', unit='unit',
                      time='period', first_treat='first_treat')

print("Estimator Comparison (True effect = 2.0)")
print("=" * 55)
print(f"{'Estimator':<25} {'ATT':>8} {'SE':>8} {'CI Width':>10}")
print("-" * 55)

for name, r in [("ImputationDiD", imp), ("CallawaySantAnna", cs), ("SunAbraham", sa)]:
    ci_width = r.overall_conf_int[1] - r.overall_conf_int[0]
    print(f"{name:<25} {r.overall_att:>8.3f} {r.overall_se:>8.3f} {ci_width:>10.3f}")

Group Aggregation#

Group aggregation estimates average treatment effects by treatment cohort (groups defined by first treatment period).

[ ]:

# Fit with group aggregation
results_grp = ImputationDiD().fit(data, outcome='outcome', unit='unit',
                                   time='period', first_treat='first_treat',
                                   aggregate='group')
results_grp.to_dataframe(level='group')

Advanced Features#

Anticipation#

If treatment effects begin before the official treatment date, use the anticipation parameter to account for this.

[ ]:

# Account for 1 period of anticipation
est_antic = ImputationDiD(anticipation=1)
results_antic = est_antic.fit(data, outcome='outcome', unit='unit',
                               time='period', first_treat='first_treat')
print(f"ATT (no anticipation):    {results.overall_att:.3f}")
print(f"ATT (1-period anticipation): {results_antic.overall_att:.3f}")

Auxiliary Model Partition#

The aux_partition parameter controls the auxiliary model partition for the conservative variance estimator (Theorem 3). Finer partitions give tighter SEs but may overfit with few observations per group.

[ ]:

# Compare different partition choices
for partition in ['cohort_horizon', 'cohort', 'horizon']:
    r = ImputationDiD(aux_partition=partition).fit(
        data, outcome='outcome', unit='unit',
        time='period', first_treat='first_treat')
    print(f"aux_partition='{partition}': ATT={r.overall_att:.3f}, SE={r.overall_se:.3f}")

Summary#

Feature	ImputationDiD	CallawaySantAnna	SunAbraham
Approach	Impute Y(0) via FE model	Group-time ATT(g,t)	Saturated regression
Efficiency	Most efficient under homogeneity	Less efficient	Least efficient
Robustness	Requires homogeneity for efficiency	Fully robust to heterogeneity	Robust to heterogeneity
Control group	All untreated (always)	Never-treated or not-yet-treated	Never-treated
Best for	Homogeneous effects, maximum power	Heterogeneous effects, flexible	Robustness check

Reference: Borusyak, K., Jaravel, X., & Spiess, J. (2024). Revisiting Event-Study Designs: Robust and Efficient Estimation. Review of Economic Studies, 91(6), 3253-3285.