Imputation DiD (Borusyak et al. 2024)#

Efficient imputation estimator for staggered Difference-in-Differences.

This module implements the methodology from Borusyak, Jaravel & Spiess (2024), “Revisiting Event-Study Designs: Robust and Efficient Estimation”, Review of Economic Studies.

The estimator:

Runs OLS on untreated observations to estimate unit + time fixed effects
Imputes counterfactual Y(0) for treated observations
Aggregates imputed treatment effects with researcher-chosen weights

Inference uses the conservative clustered variance estimator from Theorem 3.

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

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

ImputationDiD#

Main estimator class for imputation DiD estimation.

class diff_diff.ImputationDiD[source]

Bases: ImputationDiDBootstrapMixin, _ImputationAggregationMixin, BaseEstimator

Borusyak-Jaravel-Spiess (2024) imputation DiD estimator.

This is the efficient estimator for staggered Difference-in-Differences under parallel trends. It produces shorter confidence intervals than Callaway-Sant’Anna (~50% shorter) and Sun-Abraham (2-3.5x shorter) under homogeneous treatment effects.

The estimation procedure: 1. Run OLS on untreated observations to estimate unit + time fixed effects 2. Impute counterfactual Y(0) for treated observations 3. Aggregate imputed treatment effects with researcher-chosen weights

Inference uses the conservative clustered variance estimator from Theorem 3 of the paper.

Parameters:

anticipation (int, default=0) – Number of periods before treatment where effects may occur. Must be a non-negative integer; bool is rejected.
alpha (float, default=0.05) – Significance level for confidence intervals.
cluster (str, optional) – Column name for cluster-robust standard errors. If None, clusters at the unit level by default.
vcov_type (str, default="hc1") – Variance estimator family. Permanently narrow to {"hc1"} per the IF-based variance contract (Theorem 3): analytical-sandwich families {classical, hc2, hc2_bm} and conley are rejected at __init__ with methodology-rooted messages. cluster= invokes per-cluster IF summation; survey_design= invokes TSL on the combined IF. See REGISTRY.md for the cross-estimator IF-vs-sandwich taxonomy.
n_bootstrap (int, default=0) – Number of bootstrap iterations. If 0, uses analytical inference (conservative variance from Theorem 3).
bootstrap_weights (str, default="rademacher") – Type of bootstrap weights: “rademacher”, “mammen”, or “webb”.
seed (int, optional) – Random seed for reproducibility.
rank_deficient_action (str, default="warn") – Action when design matrix is rank-deficient: - “warn”: Issue warning and drop linearly dependent columns - “error”: Raise ValueError - “silent”: Drop columns silently
horizon_max (int, optional) – Maximum event-study horizon. If set, event study effects are only computed for abs(h) <= horizon_max.
aux_partition (str, default="cohort_horizon") – Controls the auxiliary model partition for Theorem 3 variance: - “cohort_horizon”: Groups by cohort x relative time (finest) - “cohort”: Groups by cohort only (coarser) - “horizon”: Groups by relative time only (coarser) Coarser partitions pool more observations per auxiliary group and are typically – but not necessarily – more conservative: on a balanced panel with uniform weights “cohort” coincides exactly with the default, and per-horizon “horizon” SEs can be smaller than the default’s under leave_one_out=True (measured against Stata did_imputation avgeffectsby(); see the REGISTRY ImputationDiD partition note).
pretrends (bool, default=False) – If True, event study includes pre-treatment horizons for visual pre-trends assessment. Pre-period effects should be ~0 under parallel trends. Only affects event_study aggregation; overall ATT and group aggregation are unchanged.
leave_one_out (bool, default=False) – If True, apply the Borusyak-Jaravel-Spiess (2024) Supplementary Appendix A.9 leave-one-out finite-sample refinement to the conservative variance. The non-LOO auxiliary aggregate tau_tilde_g is built from the fitted tau_hat_it and thus partially overfits to the noise epsilon_it, biasing the variance downward. LOO recomputes each unit’s group aggregate excluding that unit – implemented efficiently by rescaling each treated auxiliary residual by 1 / (1 - v_ig**2 / sum_j v_jg**2) (App. A.9), which is exactly equivalent to the direct leave-one-out at the per-unit cluster sum. Yields a larger, less-downward-biased SE (Prop. A8: unbiased for an upper bound). Default False preserves R didimputation parity; the refinement is an option in the authors’ Stata did_imputation. LOO is undefined for a group with a single positive-weight unit (App. A.9 footnote 51): such groups fall back to the non-LOO residual with a UserWarning. The Prop. A8 direction (LOO >= non-LOO) is guaranteed at the default unit clustering; coarser cluster= / analytical survey_design= / n_bootstrap compositions apply the same rescale but are a library extension beyond the paper’s derivation. Replicate-weight survey designs raise NotImplementedError (their variance bypasses the influence-function path where the rescale lives).
df_convention ({"residual", "cluster", "normal"}, default "residual") – Degrees-of-freedom convention for the PRETRENDS lead regression’s per-lead t/p/CI — the one ImputationDiD surface running the shared clustered CR1 sandwich (pretrends=True, surfaced fit-time via the deprecated aggregate="event_study"/"all" or post-fit via results.aggregate('event_study')). "residual" (default) uses the lead regression’s residual df (n − k_kept − absorbed [time, unit] rank) — the 3.9 fix: previously silent normal-theory z on plain clustered fits; "cluster" uses G − 1; "normal" deliberately uses z. The full-design survey df keeps precedence on survey fits. Everything else — the BJS Theorem-3 overall/event-study inference and the joint pretrend Wald F (which keeps its cluster-robust F(q, G − 1) reference) — is knob-independent; an explicitly non-default value on a configuration that never surfaces the per-lead inference warns at fit time. The default flips to "cluster" at v4.

results_

Estimation results after calling fit().

Type:: ImputationDiDResults

is_fitted_

Whether the model has been fitted.

Type:: bool

Examples

Basic usage:

>>> from diff_diff import ImputationDiD, generate_staggered_data
>>> data = generate_staggered_data(n_units=200, seed=42)
>>> est = ImputationDiD()
>>> results = est.fit(data, outcome='outcome', unit='unit',
...                   time='time', first_treat='first_treat')
>>> results.print_summary()

With a post-fit event study (M-021):

>>> est = ImputationDiD()
>>> results = est.fit(data, outcome='outcome', unit='unit',
...                   time='time', first_treat='first_treat')
>>> es = results.aggregate('event_study')
>>> from diff_diff import plot_event_study
>>> plot_event_study(es)

Notes

The imputation estimator uses ALL untreated observations (never-treated + not-yet-treated periods of eventually-treated units) to estimate the counterfactual model. There is no control_group parameter because this is fundamental to the method’s efficiency.

References

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

Methods

`fit`(data, outcome, unit, time, first_treat)	Fit the imputation DiD estimator.
`get_params`([deep])	Get estimator parameters (sklearn-compatible).
`set_params`(**params)	Set estimator parameters (sklearn-compatible, transactional).

__init__(anticipation=0, alpha=0.05, cluster=None, vcov_type='hc1', n_bootstrap=0, bootstrap_weights='rademacher', seed=None, rank_deficient_action='warn', horizon_max=None, aux_partition='cohort_horizon', pretrends=False, leave_one_out=False, df_convention='residual')[source]

Parameters:

anticipation (int)
alpha (float)
cluster (str | None)
vcov_type (str)
n_bootstrap (int)
bootstrap_weights (str)
seed (int | None)
rank_deficient_action (str)
horizon_max (int | None)
aux_partition (str)
pretrends (bool)
leave_one_out (bool)
df_convention (str)

anticipation: int

alpha: float

n_bootstrap: int

bootstrap_weights: str

seed: int | None

rank_deficient_action: str

horizon_max: int | None

aux_partition: str

pretrends: bool

leave_one_out: bool

df_convention: str

results_: ImputationDiDResults | None

fit(data, outcome, unit, time, first_treat, covariates=None, aggregate=<not supplied>, balance_e=<not supplied>, survey_design=None)[source]

Fit the imputation DiD estimator.

Parameters:

data (pd.DataFrame) – Panel data with unit and time identifiers.
outcome (str) – Name of outcome variable column.
unit (str) – Name of unit identifier column.
time (str) – Name of time period column.
first_treat (str) – Name of column indicating when unit was first treated. Use 0 (or np.inf) for never-treated units.
covariates (list of str, optional) – List of covariate column names.
aggregate (str, optional) – DEPRECATED (3.9, removed in 4.0; row M-021): aggregate as a post-fit step instead — results.aggregate('event_study') / .aggregate('group') / .aggregate('simple') / .aggregate('total'). Supplying ANY value (None included) warns; the deprecated path still works and returns exactly the numbers it always did (fit-time mode: None/”simple” overall only, “event_study”, “group”, or “all”).
balance_e (int, optional) – DEPRECATED (3.9, removed in 4.0; row M-118): moves onto results.aggregate('event_study', balance_e=...). Restricts the event study to cohorts observed at every relative time in [-balance_e, max_h] (the balanced-window rule).
survey_design (SurveyDesign, optional) – Survey design specification for design-based inference. Supports pweight only (aweight/fweight raise ValueError). Supports strata, PSU, and FPC for design-based variance via compute_survey_if_variance(). Strata enters survey df for t-distribution inference. Both analytical (n_bootstrap=0) and bootstrap inference are supported.

Returns:

Object containing all estimation results.

Return type:

ImputationDiDResults

Raises:

ValueError – If required columns are missing or data validation fails.

summary()[source]

Get summary of estimation results.

Return type:: str

print_summary()[source]

Print summary to stdout.

Return type:: None

get_params(deep=True)

Get estimator parameters (sklearn-compatible).

Parameters:: deep (bool, default True) – Accepted for sklearn compatibility (sklearn.base.clone calls get_params(deep=False)) and ignored: no diff-diff estimator nests another estimator.
Returns:: Estimator parameters suitable for passing to __init__. Keys follow the __init__ signature order. Where a class stores a parameter’s raw value under a different attribute (_PARAM_ATTR_ALIASES), the raw value is returned so that type(est)(**est.get_params()) reconstructs the same configuration.
Return type:: Dict[str, Any]

set_params(**params)

Set estimator parameters (sklearn-compatible, transactional).

Validates the merged configuration by constructing a throwaway instance of type(self) - so a rejected call raises the same error __init__ would and leaves this estimator unchanged - then adopts the probe’s parameter attributes and the class’s declared derived-config attributes. Fitted state is never touched.

Parameters:

**params (Any) – Estimator parameters. Unknown names raise ValueError before any validation or mutation.
self (TSelf)
**params

Return type:

self

ImputationDiDResults#

Results container for imputation DiD estimation.

class diff_diff.ImputationDiDResults[source]

Bases: BaseResults, AggregationMixin

Results from Borusyak-Jaravel-Spiess (2024) imputation DiD estimation.

treatment_effects

Unit-level treatment effects with columns: unit, time, tau_hat, weight.

Type:: pd.DataFrame

overall_att

Overall average treatment effect on the treated.

Type:: float

overall_se

Standard error of overall ATT.

Type:: float

overall_t_stat

T-statistic for overall ATT.

Type:: float

overall_p_value

P-value for overall ATT.

Type:: float

overall_conf_int

Confidence interval for overall ATT.

Type:: tuple

event_study_effects

Dictionary mapping relative time h to effect dict with keys: ‘effect’, ‘se’, ‘t_stat’, ‘p_value’, ‘conf_int’, ‘n_obs’.

Type:: dict, optional

group_effects

Dictionary mapping cohort g to effect dict.

Type:: dict, optional

groups

List of treatment cohorts.

Type:: list

time_periods

List of all time periods.

Type:: list

n_obs

Total number of observations.

Type:: int

n_treated_obs

Number of treated observations (\(|\Omega_1|\)).

Type:: int

n_untreated_obs

Number of untreated observations (\(|\Omega_0|\)).

Type:: int

n_treated_units

Number of ever-treated units.

Type:: int

n_control_units

Number of units contributing to Omega_0.

Type:: int

alpha

Significance level used.

Type:: float

pretrend_results

Populated by pretrend_test().

Type:: dict, optional

bootstrap_results

Bootstrap inference results.

Type:: ImputationBootstrapResults, optional

event_study_df

Scalar survey df governing the event-study rows’ t-inference (the FINAL replicate-override value on replicate fits; leads included on survey fits). None on non-survey fits, on bootstrapped fits, when no fit-time event-study surface was built, and for the replicate-undefined 0 sentinel.

Type:: float or None

Methods

`aggregate`(type[, weights, balance_e])	Re-aggregate this fit without refitting.
`summary`([alpha])	Generate formatted summary of estimation results.
`print_summary`([alpha])	Print summary to stdout.
`to_dataframe`([level])	Convert results to DataFrame.
`pretrend_test`([n_leads])	Run a pre-trend test (Equation 9 of Borusyak et al. 2024).

treatment_effects: DataFrame

overall_att: float

overall_se: float

overall_t_stat: float

overall_p_value: float

overall_conf_int: Tuple[float, float]

event_study_effects: Dict[int, Dict[str, Any]] | None

group_effects: Dict[Any, Dict[str, Any]] | None

groups: List[Any]

time_periods: List[Any]

n_obs: int

n_treated_obs: int

n_untreated_obs: int

n_treated_units: int

n_control_units: int

alpha: float = 0.05

anticipation: int = 0

pretrend_results: Dict[str, Any] | None = None

bootstrap_results: ImputationBootstrapResults | None = None

survey_metadata: Any | None = None

vcov_type: str = 'hc1'

cluster_name: str | None = None

n_clusters: int | None = None

leave_one_out: bool = False

df_convention: str | None = None

event_study_df: float | None = None

property att: float

property se: float

property conf_int: Tuple[float, float]

property p_value: float

property t_stat: float

__repr__()[source]

Concise string representation.

Return type:: str

property coef_var: float

SE / abs(overall ATT). NaN when ATT is 0 or SE non-finite.

Type:: Coefficient of variation

summary(alpha=None)[source]

Generate formatted summary of estimation results.

Parameters:: alpha (float, optional) – Accepted for signature uniformity. The stored intervals were computed at fit time; a value different from the stored alpha raises ValueError rather than silently recomputing or relabeling (bootstrap percentile intervals cannot be reconstructed from the reported SE). Re-fit at the desired alpha instead.
Returns:: Formatted summary.
Return type:: str

print_summary(alpha=None)[source]

Print summary to stdout.

Parameters:: alpha (float | None)
Return type:: None

to_dataframe(level='observation')[source]

Convert results to DataFrame.

Parameters:: level (str, default="observation") – Level of aggregation: - “observation”: Unit-level treatment effects - “event_study”: Event study effects by relative time - “group”: Group (cohort) effects
Returns:: Results as DataFrame.
Return type:: pd.DataFrame

to_dict()[source]

Convert results to a dictionary.

Provides flat headline aliases (att/se/t_stat/p_value/ conf_int_lower/conf_int_upper) plus variance-estimator metadata (vcov_type, optional cluster_name/n_clusters, optional n_bootstrap, inference_method).

Returns:: Dictionary containing the headline overall ATT and inference metadata. Per-cohort / per-horizon detail is exposed via to_dataframe().
Return type:: Dict[str, Any]

pretrend_test(n_leads=None)[source]

Run a pre-trend test (Equation 9 of Borusyak et al. 2024).

Adds pre-treatment lead indicators to the Step 1 OLS and tests their joint significance via a Wald F-test (cluster-robust, or design-based survey VCV when survey_design was provided at fit).

Parameters:: n_leads (int, optional) – Number of pre-treatment leads to include. If None, uses all available pre-treatment periods minus one (for the reference period).
Returns:: Dictionary with keys: ‘f_stat’, ‘p_value’, ‘df’, ‘n_leads’, ‘lead_coefficients’.
Return type:: dict

property is_significant: bool: Check if overall ATT is significant.

property significance_stars: str: Significance stars for overall ATT.

__init__(treatment_effects, overall_att, overall_se, overall_t_stat, overall_p_value, overall_conf_int, event_study_effects, group_effects, groups, time_periods, n_obs, n_treated_obs, n_untreated_obs, n_treated_units, n_control_units, alpha=0.05, anticipation=0, pretrend_results=None, bootstrap_results=None, _estimator_ref=None, survey_metadata=None, vcov_type='hc1', cluster_name=None, n_clusters=None, leave_one_out=False, df_convention=None, _aggregation_kit=None, event_study_df=None)

Parameters:

treatment_effects (DataFrame)
overall_att (float)
overall_se (float)
overall_t_stat (float)
overall_p_value (float)
overall_conf_int (Tuple[float, float])
event_study_effects (Dict[int, Dict[str, Any]] | None)
group_effects (Dict[Any, Dict[str, Any]] | None)
groups (List[Any])
time_periods (List[Any])
n_obs (int)
n_treated_obs (int)
n_untreated_obs (int)
n_treated_units (int)
n_control_units (int)
alpha (float)
anticipation (int)
pretrend_results (Dict[str, Any] | None)
bootstrap_results (ImputationBootstrapResults | None)
_estimator_ref (Any | None)
survey_metadata (Any | None)
vcov_type (str)
cluster_name (str | None)
n_clusters (int | None)
leave_one_out (bool)
df_convention (str | None)
_aggregation_kit (Any | None)
event_study_df (float | None)

Return type:

None

ImputationBootstrapResults#

Bootstrap inference results.

class diff_diff.ImputationBootstrapResults[source]

Bases: object

Results from ImputationDiD bootstrap inference.

Bootstrap is a library extension beyond Borusyak et al. (2024), which proposes only analytical inference via the conservative variance estimator. Provided for consistency with CallawaySantAnna and SunAbraham.

n_bootstrap

Number of bootstrap iterations.

Type:: int

weight_type

Type of bootstrap weights: “rademacher”, “mammen”, or “webb”.

Type:: str

alpha

Significance level used for confidence intervals.

Type:: float

overall_att_se

Bootstrap standard error for overall ATT.

Type:: float

overall_att_ci

Bootstrap confidence interval for overall ATT.

Type:: tuple

overall_att_p_value

Bootstrap p-value for overall ATT.

Type:: float

event_study_ses

Bootstrap SEs for event study effects.

Type:: dict, optional

event_study_cis

Bootstrap CIs for event study effects.

Type:: dict, optional

event_study_p_values

Bootstrap p-values for event study effects.

Type:: dict, optional

group_ses

Bootstrap SEs for group effects.

Type:: dict, optional

group_cis

Bootstrap CIs for group effects.

Type:: dict, optional

group_p_values

Bootstrap p-values for group effects.

Type:: dict, optional

bootstrap_distribution

Full bootstrap distribution of overall ATT.

Type:: np.ndarray, optional

n_bootstrap: int

weight_type: str

alpha: float

overall_att_se: float

overall_att_ci: Tuple[float, float]

overall_att_p_value: float

event_study_ses: Dict[int, float] | None = None

event_study_cis: Dict[int, Tuple[float, float]] | None = None

event_study_p_values: Dict[int, float] | None = None

group_ses: Dict[Any, float] | None = None

group_cis: Dict[Any, Tuple[float, float]] | None = None

group_p_values: Dict[Any, float] | None = None

bootstrap_distribution: ndarray | None = None

__init__(n_bootstrap, weight_type, alpha, overall_att_se, overall_att_ci, overall_att_p_value, event_study_ses=None, event_study_cis=None, event_study_p_values=None, group_ses=None, group_cis=None, group_p_values=None, bootstrap_distribution=None)

Parameters:

n_bootstrap (int)
weight_type (str)
alpha (float)
overall_att_se (float)
overall_att_ci (Tuple[float, float])
overall_att_p_value (float)
event_study_ses (Dict[int, float] | None)
event_study_cis (Dict[int, Tuple[float, float]] | None)
event_study_p_values (Dict[int, float] | None)
group_ses (Dict[Any, float] | None)
group_cis (Dict[Any, Tuple[float, float]] | None)
group_p_values (Dict[Any, float] | None)
bootstrap_distribution (ndarray | None)

Return type:

None

Convenience Function#

diff_diff.imputation_did(data, outcome, unit, time, first_treat, covariates=None, aggregate=<not supplied>, balance_e=<not supplied>, survey_design=None, vcov_type='hc1', **kwargs)[source]#

Convenience function for imputation DiD estimation.

Deprecated since version 3.9: imputation_did() is deprecated and will be removed in 4.0 (row M-070). Construct the estimator instead: ImputationDiD(...).fit(data, ...).

This is a shortcut for creating an ImputationDiD estimator and calling fit().

Parameters:

data (pd.DataFrame) – Panel data.
outcome (str) – Outcome variable column name.
unit (str) – Unit identifier column name.
time (str) – Time period column name.
first_treat (str) – Column indicating first treatment period (0 for never-treated).
covariates (list of str, optional) – Covariate column names.
aggregate (str, optional) – DEPRECATED (3.9, removed in 4.0; row M-021): forwarded to fit(), which warns — aggregate post-fit via results.aggregate('event_study') instead. A plain wrapper call (kwarg not supplied) never fires the aggregate warning; since 3.9 every wrapper call fires the M-070 wrapper-deprecation warning.
balance_e (int, optional) – DEPRECATED (3.9, removed in 4.0; row M-118): forwarded to fit(), which warns — moves onto results.aggregate('event_study', balance_e=...).
survey_design (SurveyDesign, optional) – Survey design specification for design-based inference. Supports pweight only (aweight/fweight raise ValueError). Supports strata, PSU, and FPC for design-based variance. Strata enters survey df for t-distribution inference. Both analytical (n_bootstrap=0) and bootstrap inference are supported.
vcov_type (str, default="hc1") – Variance estimator family. ImputationDiD permanently accepts {"hc1"} only — analytical-sandwich families {classical, hc2, hc2_bm} are rejected at __init__ because the Theorem 3 per-unit IF aggregation has no single design matrix on which hat-matrix leverage or Bell-McCaffrey Satterthwaite DOF can be defined. cluster= invokes per-cluster IF summation; survey_design= invokes TSL on the combined IF.
**kwargs – Additional keyword arguments passed to ImputationDiD constructor.

Returns:

Estimation results.

Return type:

ImputationDiDResults

Examples

>>> from diff_diff import imputation_did, generate_staggered_data
>>> data = generate_staggered_data(seed=42)
>>> results = imputation_did(data, 'outcome', 'unit', 'time', 'first_treat')
>>> results.print_summary()
>>> results.aggregate('event_study').summary()  # post-fit aggregation

Example Usage#

Basic usage:

from diff_diff import ImputationDiD, generate_staggered_data

data = generate_staggered_data(n_units=200, seed=42)
est = ImputationDiD()
results = est.fit(data, outcome='outcome', unit='unit',
                  time='period', first_treat='first_treat')
results.print_summary()

Event study with visualization:

from diff_diff import ImputationDiD, plot_event_study

est = ImputationDiD()
results = est.fit(data, outcome='outcome', unit='unit',
                  time='period', first_treat='first_treat')
es = results.aggregate('event_study')   # post-fit (M-021)
total = results.aggregate('total')      # estimator-owned total (3.10)
plot_event_study(es)

Pre-trend test:

results = est.fit(data, outcome='outcome', unit='unit',
                  time='period', first_treat='first_treat')
pt = results.pretrend_test(n_leads=3)
print(f"F-stat: {pt['f_stat']:.3f}, p-value: {pt['p_value']:.4f}")

Comparison with other estimators:

from diff_diff import ImputationDiD, CallawaySantAnna, SunAbraham

# All three should agree under homogeneous effects
imp = ImputationDiD().fit(data, ...)
cs = CallawaySantAnna().fit(data, ...)
sa = SunAbraham().fit(data, ...)

print(f"Imputation ATT: {imp.overall_att:.3f} (SE: {imp.overall_se:.3f})")
print(f"CS ATT: {cs.overall_att:.3f} (SE: {cs.overall_se:.3f})")
print(f"SA ATT: {sa.overall_att:.3f} (SE: {sa.overall_se:.3f})")