StaggeredDifferenceInDifferences#

class causalpy.experiments.staggered_did.StaggeredDifferenceInDifferences[source]#

A class to analyse data from staggered adoption Difference-in-Differences settings.

This class implements the Borusyak, Jaravel, and Spiess (BJS, 2024) imputation estimator for staggered adoption settings. It fits a model on untreated observations only (pre-treatment periods for eventually-treated units plus all periods for never-treated units), then predicts counterfactual outcomes for all observations. Treatment effects are computed as the difference between observed and predicted outcomes for treated observations.

Parameters:

  • data (DataFrame) – A pandas dataframe with panel data (unit x time observations).

  • formula (str) – A statistical model formula. Recommended: “y ~ 1 + C(unit) + C(time)” for unit and time fixed effects.

  • unit_variable_name (str) – Name of the column identifying units.

  • time_variable_name (str) – Name of the column identifying time periods.

  • treated_variable_name (str) – Name of the column indicating treatment status (0/1). Defaults to “treated”.

  • treatment_time_variable_name (str | None) – Name of the column containing unit-level treatment time (G_i). If None, treatment time is inferred from the treated_variable_name column.

  • never_treated_value (Any) – Value indicating never-treated units in treatment_time column. Defaults to np.inf.

  • model (PyMCModel | RegressorMixin | None) – A model for the untreated outcome. Defaults to LinearRegression.

  • event_window (tuple[int, int] | None) – Tuple (min_event_time, max_event_time) to restrict event-time aggregation. If None, uses all available event-times.

  • reference_event_time (int) – Event-time whose effect is normalised to zero. Used by the ETWFE estimator, where the corresponding column is omitted from the effect surface. Must satisfy -n_leads <= reference_event_time <= -1. Defaults to -1. Unused (reserved) by the imputation estimator.

  • estimator (Literal['imputation', 'etwfe']) – Which estimator to run. "imputation" (default) is the Borusyak-Jaravel-Spiess fit-on-untreated-then-impute approach. "etwfe" is Wooldridge’s extended two-way fixed effects (Mundlak) estimator: a saturated regression on the full sample with one treatment effect per (cohort, event-time) cell.

  • conditioning (Optional[Literal['mundlak', 'dummy']]) – How the two-way effects are conditioned in the ETWFE estimator. None (default) resolves to "mundlak" for PyMC models and "dummy" for scikit-learn models. "mundlak" is rejected for scikit-learn models (see Notes). Only valid when estimator="etwfe".

  • n_leads (int) – Number of pre-treatment lead terms the ETWFE estimator should estimate. Defaults to 0 (post-treatment cells only). Only valid when estimator="etwfe".

  • max_event_time (int | None) – Largest event-time given its own column in the ETWFE effect surface. Treated observations beyond it are top-binned into that column. If None (default), every observed treated event-time gets its own column. Only valid when estimator="etwfe".

  • covariates (list[str] | None) – Names of additional covariate columns to include additively in the ETWFE design. The formula’s right-hand side is ignored by the ETWFE estimator, so covariates must be supplied here. Only valid when estimator="etwfe".

  • se_type (Literal['cluster', 'classical']) – Standard error type for the OLS ETWFE path. "cluster" (default) is a cluster-by-unit sandwich estimator. Only valid when estimator="etwfe".

  • **kwargs (Any) – Additional keyword arguments forwarded to BaseExperiment.

data_#

Augmented data with G (treatment time), event_time, y_hat0 (counterfactual), and tau_hat (treatment effect) columns.

Type:

pd.DataFrame

att_group_time_#

Group-time ATT estimates: ATT(g, t) for each cohort g and calendar time t. Includes an identified column; non-identified cells have NaN estimates.

Type:

pd.DataFrame

att_event_time_#

Event-time ATT estimates: ATT(e) for each event-time e = t - G. Includes an identified column; non-identified cells have NaN estimates.

Type:

pd.DataFrame

non_identified_periods_#

Calendar periods with no untreated observations.

Type:

set

non_identified_cohorts_#

Treatment cohorts with at least one non-identified post-treatment ATT(g, t).

Type:

set

att_#

ETWFE only. The aggregated average-over-the-treated ATT. On the PyMC path this is the posterior of the in-model att deterministic; on the OLS path it is the point estimate w'b.

Type:

xarray.DataArray or float

att_se_#

ETWFE only. Standard error of att_ on the OLS path; None on the PyMC path, where att_ carries its own posterior.

Type:

float or None

tau_surface_#

ETWFE only. Long-form (cohort, event_time, att, ...) table covering every estimated cell, including lead cells.

Type:

pd.DataFrame

att_weights_#

ETWFE only. The average-over-the-treated weight matrix w_gk = N_gk / sum(N_gk), cohorts x event-times. Lead columns are zero.

Type:

pd.DataFrame

event_time_grid_#

ETWFE only. The event-times actually estimated, reference omitted.

Type:

np.ndarray

etwfe_formula_#

ETWFE only, OLS path. The generated saturated patsy formula.

Type:

str

estimator, conditioning, n_leads, se_type

The resolved configuration, echoed back. conditioning is None for the imputation estimator and the resolved "mundlak" / "dummy" value for ETWFE.

Notes

Estimate extraction

The Borusyak-Jaravel-Spiess imputation estimator fits the untreated outcome model using only observations that are not yet treated or never treated. It predicts each treated observation’s untreated potential outcome, subtracts that prediction from the observed outcome, and averages the resulting one-sided contrasts into group-time and event-time ATTs. Bayesian aggregation retains posterior uncertainty in mu; OLS aggregation uses point predictions and standard-error approximations.

Like Interrupted Time Series, this fit-predict-subtract procedure is a reduced-form estimator. The corresponding structural contrast is a saturated regression as in Wooldridge’s extended two-way fixed effects (ETWFE) framework, which CausalPy does not currently implement.

This estimator requires the following identifying assumptions:

  1. Absorbing treatment: Once a unit receives treatment, it must remain treated in all subsequent periods. Treatment cannot be reversed or temporarily suspended. This is validated at runtime.

  2. Parallel trends: In the absence of treatment, treated and control units would have followed parallel outcome trajectories.

  3. No anticipation: Units do not change their behavior in anticipation of future treatment.

  4. Untreated support at each calendar period: The time fixed effect \(\gamma_t\) for calendar period \(t\) is identified only if at least one unit is untreated in that period. Without never-treated units, post-treatment effects for the last-treated cohort (and any calendar periods where every unit is already treated) are not identified. CausalPy warns when this condition fails and marks the affected ATT(g, t) and ATT(e) cells as non-identified in the output tables.

Panel Balance: This implementation supports both balanced and unbalanced panel data. While balanced panels (where each unit is observed in every time period) are common in staggered DiD applications, the imputation-based approach of Borusyak et al. (2024) can accommodate unbalanced panels. The key requirement is that treatment timing is well-defined for each unit, not that all units are observed in all periods. Unit and observation counts in the summary output are computed without assuming balanced panels.

ETWFE and the formula argument: the estimator="etwfe" path builds its own saturated design and uses only the left-hand side of formula. A UserWarning names any right-hand-side term beyond 1, 0, C(unit) and C(time). This keeps the canonical "y ~ 1 + C(unit) + C(time)" call working when a user simply flips estimator=.

Mundlak conditioning requires PyMC: with free unit dummies the Mundlak unit mean is exactly collinear with them, so a pseudo-inverse would silently drop it and “Mundlak OLS” would be numerically identical to the dummy fit. Genuine Mundlak conditioning needs partial pooling, i.e. the PyMC path.

ETWFE covariates enter additively. Wooldridge’s centred-covariate by (g, k) interactions are not implemented; this is future work.

The Mundlak time coefficient ``g_t`` must not be interpreted. Under conditioning="mundlak" the Mundlak time mean dbar_time is a deterministic function of the calendar period t alone, so it lies exactly in the span of the time effects beta_t. g_t is therefore identified only by its prior: its posterior carries no information from the data, and reading it as “the effect of average exposure in a period” is a mistake. This is a property of the Mundlak device, not a defect of the implementation. The ATT is unaffected. tau – and hence att_ – is identified off within-cell variation, which is orthogonal to any function of t alone, so the collinearity between dbar_time and beta_t moves posterior mass between two nuisance parameters without touching the estimand. g_u is better behaved, because the unit intercepts are only partially pooled and shrinkage identifies it, but it too is a nuisance parameter.

References

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

Wooldridge, J. M. (2021). Two-Way Fixed Effects, the Two-Way Mundlak Regression, and Difference-in-Differences Estimators. Working paper.

Mundlak, Y. (1978). On the Pooling of Time Series and Cross Section Data. Econometrica, 46(1), 69-85.

Examples

>>> import causalpy as cp
>>> from causalpy.data.simulate_data import generate_staggered_did_data
>>> df = generate_staggered_did_data(n_units=30, n_time_periods=15, seed=42)
>>> result = cp.StaggeredDifferenceInDifferences(
...     df,
...     formula="y ~ 1 + C(unit) + C(time)",
...     unit_variable_name="unit",
...     time_variable_name="time",
...     treated_variable_name="treated",
...     treatment_time_variable_name="treatment_time",
...     model=cp.pymc_models.LinearRegression(
...         sample_kwargs={
...             "tune": 100,
...             "draws": 200,
...             "chains": 2,
...             "progressbar": False,
...         }
...     ),
... )

The same call switched onto Wooldridge’s extended two-way fixed effects estimator. Only estimator, conditioning and the model class change; the ATT is then available as att_, a posterior of the in-model att deterministic, and the full effect surface as tau_surface_:

>>> result = cp.StaggeredDifferenceInDifferences(
...     df,
...     formula="y ~ 1 + C(unit) + C(time)",
...     unit_variable_name="unit",
...     time_variable_name="time",
...     treated_variable_name="treated",
...     treatment_time_variable_name="treatment_time",
...     estimator="etwfe",
...     conditioning="mundlak",
...     n_leads=4,
...     model=cp.pymc_models.ETWFERegression(
...         sample_kwargs={
...             "tune": 500,
...             "draws": 500,
...             "chains": 4,
...             "progressbar": False,
...         }
...     ),
... )
>>> float(result.att_.mean())
>>> result.tau_surface_.head()
>>> fig, axes = result.plot_tau_surface()

Methods

StaggeredDifferenceInDifferences.algorithm()

Run the experiment algorithm for the selected estimator.

StaggeredDifferenceInDifferences.effect_summary(*)

Generate a decision-ready summary of causal effects for Staggered Difference-in-Differences.

StaggeredDifferenceInDifferences.fit(*args, ...)

Fit the underlying model.

StaggeredDifferenceInDifferences.generate_report(*)

Generate a self-contained HTML report for this experiment.

StaggeredDifferenceInDifferences.get_plot_data([...])

Get event-time plotting data.

StaggeredDifferenceInDifferences.input_validation()

Validate the input data and parameters.

StaggeredDifferenceInDifferences.plot(*[, ...])

Plot the staggered difference-in-differences event study.

StaggeredDifferenceInDifferences.plot_group_time(*)

Plot cohort-specific ATT(g, t) trajectories.

StaggeredDifferenceInDifferences.plot_tau_surface([...])

Plot the ETWFE effect surface, one panel per adoption cohort.

StaggeredDifferenceInDifferences.print_coefficients([...])

Ask the model to print its coefficients.

StaggeredDifferenceInDifferences.set_maketables_options(*)

Set optional maketables rendering options for this experiment.

StaggeredDifferenceInDifferences.summary([...])

Print summary of main results.

Attributes

att_

ETWFE-only results.

att_se_

idata

Return fitted InferenceData when the model backend supports it.

supports_bayes

supports_ols

supports_pymc_forecast

y_pred

Model predictions.

labels

data

__init__(data, formula, unit_variable_name, time_variable_name, treated_variable_name='treated', treatment_time_variable_name=None, never_treated_value=inf, model=None, event_window=None, reference_event_time=-1, estimator='imputation', conditioning=None, n_leads=0, max_event_time=None, covariates=None, se_type='cluster', **kwargs)[source]#
Parameters:

  • data (DataFrame)

  • formula (str)

  • unit_variable_name (str)

  • time_variable_name (str)

  • treated_variable_name (str)

  • treatment_time_variable_name (str | None)

  • never_treated_value (Any)

  • model (PyMCModel | RegressorMixin | None)

  • event_window (tuple[int, int] | None)

  • reference_event_time (int)

  • estimator (Literal['imputation', 'etwfe'])

  • conditioning (Literal['mundlak', 'dummy'] | None)

  • n_leads (int)

  • max_event_time (int | None)

  • covariates (list[str] | None)

  • se_type (Literal['cluster', 'classical'])

  • kwargs (Any)

Return type:

None

classmethod __new__(*args, **kwargs)#