Staggered DiD Module (staggered)

The staggered module implements difference-in-differences estimation for settings with staggered treatment adoption, based on Lee and Wooldridge (2025).

Overview

In staggered settings, different units begin treatment at different times (cohorts). This module provides:

• Data transformations: Cohort-specific demeaning and detrending

• Control group selection: Never-treated or not-yet-treated units

• Effect estimation: (g,r)-specific, cohort, and overall effects

• Multiple estimators: RA (regression adjustment), IPW, IPWRA, PSM

• Randomization inference: Bootstrap and permutation tests

Key Concepts

Cohort (g)

The period when a unit first receives treatment. Units that are never treated can be encoded as gvar=0, gvar=NaN, or gvar=inf (all internally mapped to \(\infty\)).

(g, r) Effect

The treatment effect \(\tau_{gr}\) for cohort \(g\) at calendar time \(r\), where \(r \geq g\). This is the ATT for units first treated in period \(g\), evaluated at time \(r\).

Control Group Strategies

Lee and Wooldridge (2025) establishes that under no anticipation and conditional parallel trends, both never-treated and not-yet-treated units provide valid counterfactuals:

• never_treated: Only units with \(D_\infty = 1\) (never treated during observation). Required when using aggregate='cohort' or aggregate='overall'.

• not_yet_treated: Never-treated plus cohorts h > r (units first treated after period r). Uses more control observations, potentially improving efficiency.

• all_others: All units not in the treated cohort, including units that were already treated in earlier periods. This option is primarily intended for replication and diagnostics; it may introduce forbidden comparisons under the no-anticipation assumption.

The theoretical justification shows that the cohort assignments are unconfounded with respect to the transformed potential outcome conditional on covariates.

Aggregation Levels

• none: Returns all \((g, r)\)-specific effects

• cohort: Averages effects within each cohort: \(\tau_g = \frac{1}{T-g+1} \sum_{r=g}^{T} \tau_{gr}\)

• overall: Cohort-share weighted average: \(\tau_\omega = \sum_g \omega_g \tau_g\) where \(\omega_g = N_g/N_{treat}\)

All Units Eventually Treated

When no units remain untreated through period \(T\) (no never-treated group), treatment effects are defined relative to \(Y_t(T)\) instead of \(Y_t(\infty)\). Effects can be estimated for cohorts \(g \in \{S, \ldots, T-1\}\); the final cohort (\(g = T\)) serves as the control for all earlier cohorts in period \(T\). See Methodological Notes for theoretical details.

Transformations

lwdid.staggered.transform_staggered_demean(data, y, ivar, tvar, gvar, never_treated_values=None, exclude_pre_periods=0)

Apply cohort-specific demeaning transformation for staggered designs.

Computes the transformed outcome for each treatment cohort g and post-treatment calendar time r:

\[\dot{Y}_{irg} = Y_{ir} - \bar{Y}_{i,pre(g)}\]

where \(\bar{Y}_{i,pre(g)}\) is the mean of unit i’s outcomes over all periods strictly before g.

The transformation is applied to all units (treated, not-yet-treated, and never-treated) for each cohort, enabling flexible control group selection during estimation.

Parameters:

• data (pd.DataFrame) – Long-format panel data with columns for outcome, unit identifier, time period, and cohort assignment.

• y (str) – Outcome variable column name.

• ivar (str) – Unit identifier column name.

• tvar (str) – Time variable column name (must be numeric or coercible to numeric).

• gvar (str) – Cohort variable column name indicating first treatment period.

• never_treated_values (list or None, optional) – Values in gvar indicating never-treated units. Default recognizes NaN, 0, and np.inf as never-treated indicators.

• exclude_pre_periods (int, default 0) –

  Number of periods immediately before treatment to exclude from pre-treatment mean calculation. Used for robustness checks when no-anticipation assumption may be violated.

  For cohort g, the pre-treatment mean is computed using periods {T_min, …, g-1-exclude_pre_periods} instead of {T_min, …, g-1}.

Returns:

Copy of input data with additional columns named ydot_g{g}_r{r} containing the transformed outcome for each (cohort, period) pair.

Return type:

pd.DataFrame

Raises:

• ValueError – If required columns are missing, no valid cohorts exist, or any cohort has no pre-treatment periods (cohort <= T_min).

• InsufficientPrePeriodsError – If any cohort has fewer than 1 pre-treatment period remaining after applying exclude_pre_periods.

Warns:

UserWarning – If no never-treated units are found in the data.

See also

transform_staggered_detrend

Apply linear detrending transformation.

get_cohorts

Extract treatment cohorts from data.

Notes

The pre-treatment mean is computed once per cohort and remains fixed across all post-treatment periods. This ensures consistency with the identification strategy where each cohort defines its own baseline.

lwdid.staggered.transform_staggered_detrend(data, y, ivar, tvar, gvar, never_treated_values=None, exclude_pre_periods=0)

Apply cohort-specific linear detrending transformation for staggered designs.

Fits unit-specific linear trends using pre-treatment data and computes out-of-sample residuals for post-treatment periods:

\[\ddot{Y}_{irg} = Y_{ir} - (\hat{A}_{ig} + \hat{B}_{ig} \cdot r)\]

where \(\hat{A}_{ig}\) and \(\hat{B}_{ig}\) are OLS estimates from regressing \(Y_{it}\) on a constant and \(t\) for periods \(t < g\).

This transformation removes unit-specific heterogeneous linear trends, relaxing the parallel trends assumption to allow treatment assignment correlated with trend slopes.

Parameters:

• data (pd.DataFrame) – Long-format panel data with columns for outcome, unit identifier, time period, and cohort assignment.

• y (str) – Outcome variable column name.

• ivar (str) – Unit identifier column name.

• tvar (str) – Time variable column name (must be numeric or coercible to numeric).

• gvar (str) – Cohort variable column name indicating first treatment period.

• never_treated_values (list or None, optional) – Values in gvar indicating never-treated units. Default recognizes NaN, 0, and np.inf as never-treated indicators.

• exclude_pre_periods (int, default 0) –

  Number of periods immediately before treatment to exclude from pre-treatment trend estimation. Used for robustness checks when no-anticipation assumption may be violated.

  For cohort g, the trend is estimated using periods {T_min, …, g-1-exclude_pre_periods} instead of {T_min, …, g-1}.

Returns:

Copy of input data with additional columns named ycheck_g{g}_r{r} containing the detrended outcome for each (cohort, period) pair.

Return type:

pd.DataFrame

Raises:

• ValueError – If required columns are missing, no valid cohorts exist, or any cohort has fewer than two pre-treatment periods.

• InsufficientPrePeriodsError – If any cohort has fewer than 2 pre-treatment periods remaining after applying exclude_pre_periods.

Warns:

UserWarning – If no never-treated units are found in the data.

See also

transform_staggered_demean

Apply demeaning transformation.

get_cohorts

Extract treatment cohorts from data.

Notes

Detrending requires at least two pre-treatment periods per cohort to identify both intercept and slope parameters. Units with insufficient pre-treatment data receive NaN values for transformed outcomes.

The detrending approach is appropriate when treatment and control groups may have different pre-treatment trends but those trends are linear. For more complex trend patterns, consider using only periods where trends appear approximately parallel.

lwdid.staggered.transform_staggered_demeanq(data, y, ivar, tvar, gvar, season_var, Q=4, never_treated_values=None, exclude_pre_periods=0)

Apply cohort-specific seasonal demeaning for staggered designs.

Computes transformed outcome for each cohort g and post-treatment calendar time r:

\[\dot{Y}_{irg} = Y_{ir} - \hat{\mu}_{ig} - \sum_{q=2}^{Q} \hat{\gamma}_{qg} D_q\]

where parameters are estimated from pre-treatment periods (t < g).

Parameters:

• data (pd.DataFrame) – Long-format panel data with columns for outcome, unit identifier, time period, cohort assignment, and seasonal indicator.

• y (str) – Outcome variable column name.

• ivar (str) – Unit identifier column name.

• tvar (str) – Time variable column name (must be numeric or coercible to numeric).

• gvar (str) – Cohort variable column name indicating first treatment period.

• season_var (str) – Seasonal indicator column name. Values should be integers from 1 to Q representing seasonal periods (e.g., quarters 1-4, months 1-12).

• Q (int, default 4) – Number of seasonal periods per cycle. Common values: - 4: Quarterly data (default) - 12: Monthly data - 52: Weekly data

• never_treated_values (list or None, optional) – Values in gvar indicating never-treated units. Default recognizes NaN, 0, and np.inf as never-treated indicators.

• exclude_pre_periods (int, default 0) –

  Number of periods immediately before treatment to exclude from pre-treatment seasonal estimation. Used for robustness checks when no-anticipation assumption may be violated.

  For cohort g, the seasonal parameters are estimated using periods {T_min, …, g-1-exclude_pre_periods} instead of {T_min, …, g-1}.

Returns:

Copy of input data with additional columns named ydot_g{g}_r{r} containing the seasonally-adjusted transformed outcome for each (cohort, period) pair.

Return type:

pd.DataFrame

Raises:

• ValueError – If required columns are missing, no valid cohorts exist, or any cohort has insufficient pre-treatment periods for seasonal estimation.

• InsufficientPrePeriodsError – If any cohort has fewer than Q+1 pre-treatment periods remaining after applying exclude_pre_periods.

Warns:

UserWarning – If no never-treated units are found in the data.

See also

transform_staggered_demean

Demeaning without seasonal adjustment.

transform_staggered_detrendq

Seasonal detrending transformation.

Notes

The seasonal parameters are estimated once per cohort using all periods strictly before g. This ensures consistency with the identification strategy where each cohort defines its own baseline.

Minimum required pre-treatment observations per unit is Q + 1 to ensure at least one residual degree of freedom for OLS estimation.

lwdid.staggered.transform_staggered_detrendq(data, y, ivar, tvar, gvar, season_var, Q=4, never_treated_values=None, exclude_pre_periods=0)

Apply cohort-specific seasonal detrending for staggered designs.

Computes transformed outcome for each cohort g and post-treatment calendar time r:

\[\ddot{Y}_{irg} = Y_{ir} - \hat{\alpha}_{ig} - \hat{\beta}_{ig} r - \sum_{q=2}^{Q} \hat{\gamma}_{qg} D_q\]

where parameters are estimated from pre-treatment periods (t < g).

Parameters:

• data (pd.DataFrame) – Long-format panel data with columns for outcome, unit identifier, time period, cohort assignment, and seasonal indicator.

• y (str) – Outcome variable column name.

• ivar (str) – Unit identifier column name.

• tvar (str) – Time variable column name (must be numeric or coercible to numeric).

• gvar (str) – Cohort variable column name indicating first treatment period.

• season_var (str) – Seasonal indicator column name. Values should be integers from 1 to Q representing seasonal periods (e.g., quarters 1-4, months 1-12).

• Q (int, default 4) – Number of seasonal periods per cycle. Common values: - 4: Quarterly data (default) - 12: Monthly data - 52: Weekly data

• never_treated_values (list or None, optional) – Values in gvar indicating never-treated units. Default recognizes NaN, 0, and np.inf as never-treated indicators.

• exclude_pre_periods (int, default 0) –

  Number of periods immediately before treatment to exclude from pre-treatment seasonal trend estimation. Used for robustness checks when no-anticipation assumption may be violated.

  For cohort g, the seasonal trend parameters are estimated using periods {T_min, …, g-1-exclude_pre_periods} instead of {T_min, …, g-1}.

Returns:

Copy of input data with additional columns named ycheck_g{g}_r{r} containing the seasonally-adjusted detrended outcome for each (cohort, period) pair.

Return type:

pd.DataFrame

Raises:

• ValueError – If required columns are missing, no valid cohorts exist, or any cohort has insufficient pre-treatment periods for seasonal detrending.

• InsufficientPrePeriodsError – If any cohort has fewer than Q+2 pre-treatment periods remaining after applying exclude_pre_periods.

Warns:

UserWarning – If no never-treated units are found in the data.

See also

transform_staggered_detrend

Detrending without seasonal adjustment.

transform_staggered_demeanq

Seasonal demeaning transformation.

Notes

The seasonal and trend parameters are estimated once per cohort using all periods strictly before g. Time is centered at the pre-treatment mean for numerical stability.

Minimum required pre-treatment observations per unit is Q + 2 to ensure at least one residual degree of freedom for OLS estimation (intercept + slope + Q-1 seasonal dummies = Q+1 parameters).

Note

All four transformation methods (demean, detrend, demeanq, detrendq) are available through the main lwdid() function in staggered mode. The seasonal transformations (demeanq, detrendq) require the season_var and Q parameters.

lwdid.staggered.get_cohorts(data, gvar, ivar, never_treated_values=None)

Extract valid treatment cohorts from panel data.

Identifies all distinct first-treatment periods present in the data, excluding units that are never treated. Never-treated status is determined by missing values (NaN) or explicit indicator values specified by the user.

Parameters:

• data (pd.DataFrame) – Panel data in long format containing unit and cohort identifiers.

• gvar (str) – Column name for the cohort variable indicating first treatment period.

• ivar (str) – Column name for the unit identifier.

• never_treated_values (list or None, optional) – Values in gvar that indicate never-treated units. Default is [0, np.inf]. NaN values are always treated as never-treated regardless of this parameter.

Returns:

Sorted list of unique treatment cohort values, excluding any never-treated indicators.

Return type:

list of int

Raises:

KeyError – If gvar or ivar columns are not found in data.

See also

get_valid_periods_for_cohort

Determine post-treatment periods for a cohort.

transform_staggered_demean

Apply demeaning transformation.

Notes

In staggered adoption designs, a cohort comprises all units first treated in the same time period. The cohort identifier g also defines the pre- treatment window (periods 1 through g-1) used for baseline transformations.

lwdid.staggered.get_valid_periods_for_cohort(cohort, T_max)

Determine valid post-treatment periods for a treatment cohort.

For treatment cohort g, returns the set of calendar times {g, g+1, …, T_max} during which the cohort is exposed to treatment and cohort-time specific average treatment effects on the treated can be estimated.

Parameters:

• cohort (int) – Treatment cohort identifier representing the first treatment period.

• T_max (int) – Maximum calendar time period available in the panel data.

Returns:

Consecutive integers from cohort through T_max inclusive, representing all periods where the cohort experiences treatment.

Return type:

list of int

See also

get_cohorts

Extract treatment cohorts from data.

transform_staggered_demean

Apply demeaning transformation.

Notes

The event-time relative to treatment onset ranges from 0 (instantaneous effect at treatment start) to T_max - cohort (longest exposure duration). Each calendar time r corresponds to event-time e = r - g for cohort g.

Control Groups

class lwdid.staggered.ControlGroupStrategy(value)

Enumeration of control group selection strategies.

Defines which units are eligible to serve as controls when estimating treatment effects for a given cohort-period pair. The choice of strategy affects both the valid comparison group and the identifying assumptions required.

Variables:

• NEVER_TREATED (str) – Use only units that never receive treatment throughout the observation window. Provides a stable control group whose composition does not vary across periods. Required for aggregated effect estimation where controls must be consistent.

• NOT_YET_TREATED (str) – Use never-treated units plus units not yet treated at the current period. Expands the control pool by including future treatment cohorts as temporary controls, improving estimation efficiency under conditional parallel trends and no anticipation.

• ALL_OTHERS (str) – Use all units not in the focal treatment cohort, including already-treated units from earlier cohorts. May induce forbidden comparisons that violate identification assumptions. Provided primarily for replication and diagnostic purposes.

• AUTO (str) – Automatically select based on data availability. Prefers the not-yet-treated strategy when sufficient controls are available, falling back to never-treated only when necessary.

See also

get_valid_control_units

Apply strategy to select control units.

count_control_units_by_strategy

Compare control counts across strategies.

NEVER_TREATED = 'never_treated'

NOT_YET_TREATED = 'not_yet_treated'

ALL_OTHERS = 'all_others'

AUTO = 'auto'

lwdid.staggered.get_valid_control_units(data, gvar, ivar, cohort, period, strategy=ControlGroupStrategy.NOT_YET_TREATED, never_treated_values=None, is_pre_treatment=False)

Determine valid control units for a specific cohort-period pair.

For estimating the ATT of cohort g in period r, identifies which units can serve as valid controls based on the selected strategy and the fundamental strict inequality criterion.

Parameters:

• data (pd.DataFrame) – Panel dataset containing unit and treatment timing information.

• gvar (str) – Column name indicating first treatment period for each unit.

• ivar (str) – Column name containing unit identifiers.

• cohort (int or float) – Treatment cohort (first treatment period g) of the treated group.

• period (int or float) – Calendar time period r for which to identify controls. For post-treatment: must satisfy period >= cohort. For pre-treatment: must satisfy period < cohort.

• strategy (ControlGroupStrategy, default NOT_YET_TREATED) – Strategy for selecting control units.

• never_treated_values (list, optional) – Values in gvar indicating never-treated status. Defaults to [0, np.inf].

• is_pre_treatment (bool, default False) – If True, selects control units for pre-treatment period estimation (parallel trends testing). For pre-treatment periods t < g, the control group includes all units not yet treated at period t.

Returns:

Boolean Series indexed by unit ID where True indicates valid control.

Return type:

pd.Series

Raises:

• ValueError – If data is empty, or period constraints are violated.

• KeyError – If required columns are not found.

• TypeError – If gvar column is not numeric.

See also

get_all_control_masks

Batch computation for multiple cohort-period pairs.

get_all_control_masks_pre

Batch computation for pre-treatment periods.

validate_control_group

Validate control group size requirements.

Notes

The strict inequality criterion (gvar > period) is fundamental:

• Units with gvar == period are beginning treatment in period r and thus belong to the treatment group, not the control group.

• This ensures valid controls have not yet been exposed to treatment.

For post-treatment estimation (period >= cohort):

The treatment cohort is automatically excluded because cohort units have gvar == cohort <= period, failing the gvar > period criterion.

For pre-treatment estimation (period < cohort):

The treatment cohort is correctly included as controls because period < cohort implies gvar (== cohort) > period. At pre-treatment periods, these units are not yet treated and serve as valid comparisons for parallel trends assessment.

lwdid.staggered.get_all_control_masks(data, gvar, ivar, cohorts, T_max, T_min=None, strategy=ControlGroupStrategy.NOT_YET_TREATED, never_treated_values=None)

Compute control group masks for all cohort-period combinations.

Efficiently generates control masks for multiple cohort-period pairs by pre-computing shared data structures. This batch approach avoids redundant groupby operations when estimating effects across many (cohort, period) combinations.

Parameters:

• data (pd.DataFrame) – Panel dataset containing unit and treatment timing information.

• gvar (str) – Column name indicating first treatment period for each unit.

• ivar (str) – Column name containing unit identifiers.

• cohorts (list of int or float) – Treatment cohorts for which to generate control masks.

• T_max (int or float) – Maximum time period to consider (inclusive).

• T_min (int or float, optional) – Minimum time period. Reserved for future extension.

• strategy (ControlGroupStrategy, default NOT_YET_TREATED) – Strategy for selecting control units.

• never_treated_values (list, optional) – Values in gvar indicating never-treated status. Defaults to [0, np.inf].

Returns:

Dictionary mapping (cohort, period) tuples to boolean Series indexed by unit ID. True indicates valid control status.

Return type:

dict

Raises:

ValueError – If the input data is empty.

See also

get_valid_control_units

Single cohort-period control mask.

get_all_control_masks_pre

Batch computation for pre-treatment periods.

Notes

For each cohort g, masks are generated for post-treatment periods {g, g+1, …, T_max}. The never-treated mask is computed once and reused across all cohort-period pairs, while not-yet-treated masks vary by period due to the strict inequality criterion.

lwdid.staggered.get_all_control_masks_pre(data, gvar, ivar, cohorts, T_min, strategy=ControlGroupStrategy.NOT_YET_TREATED, never_treated_values=None)

Compute control group masks for all pre-treatment cohort-period combinations.

Efficiently generates control masks for pre-treatment periods by pre-computing shared data structures. Used for parallel trends testing and event study visualization where pre-treatment effects should be approximately zero under the identifying assumptions.

Parameters:

• data (pd.DataFrame) – Panel dataset containing unit and treatment timing information.

• gvar (str) – Column name indicating first treatment period for each unit.

• ivar (str) – Column name containing unit identifiers.

• cohorts (list of int or float) – Treatment cohorts for which to generate control masks.

• T_min (int or float) – Minimum time period in the data (inclusive).

• strategy (ControlGroupStrategy, default NOT_YET_TREATED) – Strategy for selecting control units.

• never_treated_values (list, optional) – Values in gvar indicating never-treated status. Defaults to [0, np.inf].

Returns:

Dictionary mapping (cohort, period) tuples to boolean Series indexed by unit ID. True indicates valid control status.

Return type:

dict

Raises:

ValueError – If the input data is empty.

See also

get_valid_control_units

Single cohort-period control mask.

get_all_control_masks

Batch computation for post-treatment periods.

Notes

For each cohort g, masks are generated for pre-treatment periods {T_min, T_min+1, …, g-1}. At pre-treatment period t < g, the focal treatment cohort (gvar == g) is correctly included as controls because these units are not yet treated.

lwdid.staggered.validate_control_group(control_mask, cohort, period, min_control_units=1, aggregate_type=None, has_never_treated=True, strategy=None)

Validate whether a control group meets estimation requirements.

Checks control group suitability for treatment effect estimation, including minimum size requirements and aggregation constraints.

Parameters:

• control_mask (pd.Series) – Boolean Series indexed by unit ID indicating control group membership.

• cohort (int or float) – Treatment cohort being estimated.

• period (int or float) – Time period being estimated.

• min_control_units (int, default 1) – Minimum number of control units required for estimation.

• aggregate_type (str, optional) – Type of aggregation (‘cohort’ or ‘overall’). Aggregated effects require never-treated units because not-yet-treated controls vary across periods and cannot form a consistent comparison group.

• has_never_treated (bool, default True) – Whether the data contains any never-treated units.

• strategy (ControlGroupStrategy, optional) – Control group strategy being used. Generates warnings when aggregated estimation uses non-recommended strategies.

Return type:

tuple[bool, str]

Returns:

• is_valid (bool) – True if the control group passes all validation checks.

• message (str) – Descriptive message indicating success or failure reason.

See also

get_valid_control_units

Generate control group masks.

has_never_treated_units

Check for never-treated unit availability.

Notes

Validation checks are applied in priority order:

1. Non-empty control group (required for any estimation)

2. Minimum size requirement (ensures sufficient degrees of freedom)

3. Aggregation constraints: cohort-level and overall effects require never-treated units because they aggregate across multiple periods, and not-yet-treated units transition out of the control group as they become treated

lwdid.staggered.identify_never_treated_units(data, gvar, ivar, never_treated_values=None)

Identify units that never receive treatment.

Creates a boolean mask indicating which units are classified as never-treated based on their treatment timing variable values.

Parameters:

• data (pd.DataFrame) – Panel dataset containing unit and treatment timing information.

• gvar (str) – Column name indicating first treatment period for each unit.

• ivar (str) – Column name containing unit identifiers.

• never_treated_values (list, optional) – Values in gvar indicating never-treated status. Defaults to [0, np.inf]. Units with NaN in gvar are also classified as never-treated regardless of this parameter.

Returns:

Boolean Series indexed by ivar (unit ID). True indicates never-treated status.

Return type:

pd.Series

Raises:

• ValueError – If the input data is empty.

• KeyError – If gvar or ivar column is not found in the data.

See also

has_never_treated_units

Check presence of never-treated units.

get_valid_control_units

Select control units for estimation.

Notes

Never-treated units are identified through three mechanisms:

1. Missing values (NaN) in gvar, representing units with no recorded treatment date.

2. Zero values, a common coding convention for never-treated.

3. Infinity values, representing treatment dates beyond the observation window.

lwdid.staggered.has_never_treated_units(data, gvar, ivar, never_treated_values=None)

Check whether the data contains any never-treated units.

A convenience function for quickly determining if a never-treated control group is available for estimation. This is particularly useful for deciding whether aggregated effects can be estimated.

Parameters:

• data (pd.DataFrame) – Panel dataset containing unit and treatment timing information.

• gvar (str) – Column name indicating first treatment period for each unit.

• ivar (str) – Column name containing unit identifiers.

• never_treated_values (list, optional) – Values in gvar indicating never-treated status. Defaults to [0, np.inf].

Returns:

True if at least one never-treated unit exists.

Return type:

bool

See also

identify_never_treated_units

Get full mask of never-treated units.

validate_control_group

Validate control group for aggregation.

lwdid.staggered.count_control_units_by_strategy(data, gvar, ivar, cohort, period, never_treated_values=None)

Count available control units under different selection strategies.

A diagnostic function to help users understand data structure and make informed decisions about control group selection.

Parameters:

• data (pd.DataFrame) – Panel dataset containing unit and treatment timing information.

• gvar (str) – Column name indicating first treatment period for each unit.

• ivar (str) – Column name containing unit identifiers.

• cohort (int or float) – Treatment cohort of interest.

• period (int or float) – Time period of interest.

• never_treated_values (list, optional) – Values in gvar indicating never-treated status. Defaults to [0, np.inf].

Returns:

Dictionary with keys:

• 'never_treated': Count of never-treated units.

• 'not_yet_treated_only': Count of units treated in future periods (excluding never-treated).

• 'not_yet_treated_total': Total valid controls under the not-yet-treated strategy.

• 'treatment_cohort': Count of units in the treatment cohort.

Return type:

dict

Raises:

ValueError – If the input data is empty.

See also

get_valid_control_units

Generate control masks for estimation.

ControlGroupStrategy

Available control group selection strategies.

Notes

The not-yet-treated count uses strict inequality (gvar > period) to exclude units beginning treatment in the current period.

Estimation

class lwdid.staggered.CohortTimeEffect(cohort, period, event_time, att, se, ci_lower, ci_upper, t_stat, pvalue, n_treated, n_control, n_total, df_resid=0, df_inference=0)

Container for a single cohort-time treatment effect estimate.

Stores the ATT for treatment cohort g at calendar time r, along with inference statistics from cross-sectional regression.

Variables:

• cohort (int) – Treatment cohort identifier (first treatment period g).

• period (int) – Calendar time period r.

• event_time (int) – Event time relative to treatment onset (e = r - g).

• att (float) – Estimated average treatment effect on the treated.

• se (float) – Standard error of the ATT estimate.

• ci_lower (float) – Lower bound of the confidence interval.

• ci_upper (float) – Upper bound of the confidence interval.

• t_stat (float) – t-statistic for testing H0: ATT = 0.

• pvalue (float) – Two-sided p-value from t-distribution.

• n_treated (int) – Number of treated units in estimation sample.

• n_control (int) – Number of control units in estimation sample.

• n_total (int) – Total sample size (n_treated + n_control).

• df_resid (int) – Residual degrees of freedom. Default is 0.

• df_inference (int) – Degrees of freedom for inference. Default is 0.

cohort: int

period: int

event_time: int

att: float

se: float

ci_lower: float

ci_upper: float

t_stat: float

pvalue: float

n_treated: int

n_control: int

n_total: int

df_resid: int = 0

df_inference: int = 0

lwdid.staggered.estimate_cohort_time_effects(data_transformed, gvar, ivar, tvar, controls=None, vce=None, cluster_var=None, control_strategy='not_yet_treated', never_treated_values=None, min_obs=3, min_treated=1, min_control=1, alpha=0.05, estimator='ra', transform_type='demean', propensity_controls=None, trim_threshold=0.01, se_method='analytical', n_neighbors=1, with_replacement=True, caliper=None, return_diagnostics=False, match_order=None)

Estimate treatment effects for all valid cohort-period pairs.

Iterates over all treatment cohorts and their post-treatment periods, estimating the ATT for each (cohort, period) combination using the specified estimation method. For each pair, the estimation sample consists of the treatment cohort units plus valid control units determined by the control group strategy.

Parameters:

• data_transformed (pd.DataFrame) – Panel data containing transformed outcome columns generated by transform_staggered_demean or transform_staggered_detrend. Must include columns for gvar, ivar, tvar, and transformed outcomes named ‘ydot_g{g}_r{r}’ or ‘ycheck_g{g}_r{r}’.

• gvar (str) – Name of the cohort variable column indicating first treatment period.

• ivar (str) – Name of the unit identifier column.

• tvar (str) – Name of the time variable column.

• controls (list of str, optional) – Names of time-invariant control variable columns.

• vce (str, optional) – Variance estimation type: None (homoskedastic), ‘hc3’ (heteroskedasticity-robust), or ‘cluster’ (cluster-robust).

• cluster_var (str, optional) – Name of the cluster variable column. Required when vce=’cluster’.

• control_strategy (str, default='not_yet_treated') – Control group selection strategy: ‘never_treated’ uses only never-treated units; ‘not_yet_treated’ includes units first treated after the current period; ‘all_others’ uses all units not in the treatment cohort (including already-treated units); ‘auto’ selects based on data availability.

• never_treated_values (list, optional) – Values in gvar indicating never-treated units. Defaults to [0, np.inf] and NaN values.

• min_obs (int, default=3) – Minimum total sample size required for estimation.

• min_treated (int, default=1) – Minimum number of treated units required.

• min_control (int, default=1) – Minimum number of control units required.

• alpha (float, default=0.05) – Significance level for confidence interval construction.

• estimator (str, default='ra') – Estimation method: ‘ra’ (regression adjustment), ‘ipwra’ (inverse probability weighted regression adjustment), or ‘psm’ (propensity score matching).

• transform_type (str, default='demean') – Transformation type applied to the data: ‘demean’ or ‘detrend’. Determines the column prefix for transformed outcomes.

• propensity_controls (list of str, optional) – Control variables for the propensity score model. If None, uses the same variables as controls.

• trim_threshold (float, default=0.01) – Propensity score trimming threshold for IPWRA and PSM. Observations with extreme propensity scores are excluded.

• se_method (str, default='analytical') – Standard error method for IPWRA: ‘analytical’ or ‘bootstrap’.

• n_neighbors (int, default=1) – Number of nearest neighbors for PSM matching.

• with_replacement (bool, default=True) – Whether PSM matching allows replacement.

• caliper (float, optional) – Maximum propensity score distance for PSM matching.

• return_diagnostics (bool, default=False) – Reserved for future use. Currently has no effect.

• match_order (str, optional) – Reserved for future use. Currently has no effect.

Returns:

Estimation results for all valid (cohort, period) pairs, sorted by cohort and period.

Return type:

list of CohortTimeEffect

Raises:

ValueError – If required columns are missing, no valid treatment cohorts exist, or parameter values are invalid.

See also

transform_staggered_demean

Demeaning transformation for staggered designs.

transform_staggered_detrend

Detrending transformation for staggered designs.

aggregate_to_cohort

Aggregate cohort-time effects to cohort-level effects.

aggregate_to_overall

Aggregate effects to a single overall estimate.

Notes

The estimation proceeds separately for each (cohort, period) pair. For treatment cohort g in calendar time r, the control group consists of never-treated units and units first treated after period r. This rolling selection ensures that control units satisfy no anticipation at time r.

Under the ‘not_yet_treated’ strategy, the control pool varies by period: earlier periods have more controls available since fewer cohorts have been treated. The ‘never_treated’ strategy uses a fixed control pool across all periods, which may be more restrictive but avoids potential confounding from units that eventually receive treatment.

Sample size requirements (min_obs, min_treated, min_control) are checked before estimation. Pairs that fail these checks are silently skipped and reported in the warning message.

lwdid.staggered.run_ols_regression(data, y, d, controls=None, vce=None, cluster_var=None, alpha=0.05)

Estimate the ATT via OLS regression on cross-sectional data.

Regresses the transformed outcome on a constant and treatment indicator, optionally including control variables. When controls are included and sample sizes permit, interactions with the treatment indicator (centered at treated-group mean) are added to allow heterogeneous covariate effects.

Parameters:

• data (pd.DataFrame) – Cross-sectional data with one row per unit.

• y (str) – Name of the dependent variable column (transformed outcome).

• d (str) – Name of the treatment indicator column.

• controls (list of str, optional) – Names of control variable columns.

• vce ({None, 'hc0', 'hc1', 'hc2', 'hc3', 'hc4', 'robust', 'cluster'}, optional) –

  Variance estimation method (case-insensitive):

  • None: Homoskedastic OLS standard errors. Enables exact t-based inference under normality assumption.

  • ’hc0’: Basic heteroskedasticity-robust. No finite-sample adjustment.

  • ’hc1’ or ‘robust’: HC1 with degrees-of-freedom correction n/(n-k).

  • ’hc2’: Leverage-adjusted using (1 - h_ii)^{-1}.

  • ’hc3’: Small-sample adjusted using (1 - h_ii)^{-2}. Recommended for small samples with few treated or control units.

  • ’hc4’: Adaptive leverage correction using δ_i = min(4, n·h_ii/k). Recommended when extreme leverage exists.

  • ’cluster’: Cluster-robust. Requires cluster_var.

• cluster_var (str, optional) – Cluster variable column. Required when vce=’cluster’.

• alpha (float, default=0.05) – Significance level for confidence interval construction.

Returns:

Estimation results with keys: ‘att’, ‘se’, ‘ci_lower’, ‘ci_upper’, ‘t_stat’, ‘pvalue’, ‘nobs’, ‘df_resid’, ‘df_inference’.

Return type:

dict

Raises:

ValueError – If required columns are missing, sample size is insufficient, or the design matrix is singular.

See also

estimate_cohort_time_effects

Estimate ATT for all cohort-period pairs.

_compute_hc_variance

HC variance computation details.

Notes

HC standard error ordering (typical):

• SE(HC0) ≤ SE(HC1) due to degrees-of-freedom correction

• SE(HC2) ≤ SE(HC3) due to stronger leverage adjustment

• SE(HC4) adapts based on leverage; may exceed HC3 for high-leverage obs

lwdid.staggered.results_to_dataframe(results)

Convert a list of CohortTimeEffect objects to a pandas DataFrame.

Parameters:

results (list of CohortTimeEffect) – Estimation results from estimate_cohort_time_effects.

Returns:

DataFrame with columns: cohort, period, event_time, att, se, ci_lower, ci_upper, t_stat, pvalue, n_treated, n_control, n_total. Returns an empty DataFrame with appropriate columns if the input list is empty.

Return type:

pd.DataFrame

See also

estimate_cohort_time_effects

Estimate ATT for all cohort-period pairs.

CohortTimeEffect

Container class for individual effect estimates.

Aggregation

class lwdid.staggered.CohortEffect(cohort, att, se, ci_lower, ci_upper, t_stat, pvalue, n_periods, n_units, n_control, df_resid=0, df_inference=0)

Cohort-specific time-averaged treatment effect estimate.

Stores the ATT estimate for a single treatment cohort, computed by averaging the transformed outcome across all post-treatment periods and regressing on cohort membership with never-treated controls.

Variables:

• cohort (int) – Treatment cohort identifier (first treatment period).

• att (float) – Point estimate of the cohort-level ATT.

• se (float) – Standard error of the ATT estimate.

• ci_lower (float) – Lower bound of the confidence interval.

• ci_upper (float) – Upper bound of the confidence interval.

• t_stat (float) – t-statistic for the null hypothesis of zero effect.

• pvalue (float) – Two-sided p-value.

• n_periods (int) – Number of post-treatment periods included in the average.

• n_units (int) – Number of treated units in this cohort.

• n_control (int) – Number of never-treated control units.

• df_resid (int) – Residual degrees of freedom from the aggregation regression.

• df_inference (int) – Degrees of freedom used for inference (equals df_resid for homoskedastic/HC variance, or number of clusters minus one for cluster-robust variance).

See also

aggregate_to_cohort

Function that produces CohortEffect instances.

cohort_effects_to_dataframe

Convert results to DataFrame format.

cohort: int

att: float

se: float

ci_lower: float

ci_upper: float

t_stat: float

pvalue: float

n_periods: int

n_units: int

n_control: int

df_resid: int = 0

df_inference: int = 0

class lwdid.staggered.OverallEffect(att, se, ci_lower, ci_upper, t_stat, pvalue, cohort_weights, n_treated, n_control, df_resid=0, df_inference=0)

Overall cohort-size-weighted treatment effect estimate.

Stores the aggregate ATT estimate across all treatment cohorts, with weights proportional to cohort sizes. The weight for cohort g is \(w(g) = N_g / \sum_{g'} N_{g'}\), where \(N_g\) is the number of treated units in cohort g.

Variables:

• att (float) – Point estimate of the overall weighted ATT.

• se (float) – Standard error of the ATT estimate.

• ci_lower (float) – Lower bound of the confidence interval.

• ci_upper (float) – Upper bound of the confidence interval.

• t_stat (float) – t-statistic for the null hypothesis of zero effect.

• pvalue (float) – Two-sided p-value.

• cohort_weights (dict[int, float]) – Mapping from cohort identifiers to weights (weights sum to one).

• n_treated (int) – Total number of ever-treated units.

• n_control (int) – Number of never-treated control units.

• df_resid (int) – Residual degrees of freedom from the aggregation regression.

• df_inference (int) – Degrees of freedom used for inference (equals df_resid for homoskedastic/HC variance, or number of clusters minus one for cluster-robust variance).

See also

aggregate_to_overall

Function that produces OverallEffect instances.

att: float

se: float

ci_lower: float

ci_upper: float

t_stat: float

pvalue: float

cohort_weights: dict[int, float]

n_treated: int

n_control: int

df_resid: int = 0

df_inference: int = 0

class lwdid.staggered.EventTimeEffect(event_time, att, se, ci_lower, ci_upper, t_stat, pvalue, df_inference, n_cohorts, cohort_contributions, weight_sum, alpha=0.05)

Event-time aggregated treatment effect estimate (WATT).

Represents the weighted average treatment effect at event time r, computed as WATT(r) = Σ_{g∈G_r} w(g,r) · ATT(g, g+r), where weights are proportional to cohort sizes.

Variables:

• event_time (int) – Relative time since treatment (r = period - cohort). Negative values indicate pre-treatment periods, positive values indicate post-treatment.

• att (float) – Point estimate of the weighted ATT at this event time.

• se (float) – Standard error: SE(WATT(r)) = sqrt(Σ [w(g,r)]² × [SE(ATT)]²).

• ci_lower (float) – Lower bound of confidence interval using t-distribution.

• ci_upper (float) – Upper bound of confidence interval using t-distribution.

• t_stat (float) – t-statistic: WATT(r) / SE(WATT(r)).

• pvalue (float) – Two-sided p-value from t-distribution.

• df_inference (int) – Degrees of freedom for t-distribution inference. Uses min(df_g) across contributing cohorts as conservative choice.

• n_cohorts (int) – Number of cohorts contributing to this event time estimate.

• cohort_contributions (dict[int, float]) – Mapping from cohort g to its weighted contribution w(g,r) × ATT(g, g+r).

• weight_sum (float) – Sum of normalized weights (should equal 1.0 for validation).

• alpha (float) – Significance level used for confidence interval (default 0.05).

See also

aggregate_to_event_time

Function that produces EventTimeEffect instances.

event_time_effects_to_dataframe

Convert results to DataFrame format.

event_time: int

att: float

se: float

ci_lower: float

ci_upper: float

t_stat: float

pvalue: float

df_inference: int

n_cohorts: int

cohort_contributions: dict[int, float]

weight_sum: float

alpha: float = 0.05

lwdid.staggered.aggregate_to_cohort(data_transformed, gvar, ivar, tvar, cohorts, T_max, transform_type='demean', vce=None, cluster_var=None, alpha=0.05, controls=None, never_treated_values=None)

Aggregate period-specific effects to cohort-level effects.

For each cohort, estimates the time-averaged ATT via cross-sectional OLS regression of time-averaged transformed outcomes on cohort membership indicators with never-treated controls.

Parameters:

• data_transformed (pd.DataFrame) – Panel data with transformation columns named ‘ydot_g{g}_r{r}’ (demean) or ‘ycheck_g{g}_r{r}’ (detrend).

• gvar (str) – Cohort variable column name.

• ivar (str) – Unit identifier column name.

• tvar (str) – Time variable column name.

• cohorts (list of int) – Treatment cohorts to estimate.

• T_max (int) – Maximum time period.

• transform_type ({'demean', 'detrend'}, default='demean') – Transformation type determining column prefix.

• vce ({None, 'robust', 'hc0', 'hc1', 'hc2', 'hc3', 'hc4', 'cluster'}, optional) – Variance-covariance estimator.

• cluster_var (str, optional) – Cluster variable for cluster-robust standard errors.

• alpha (float, default=0.05) – Significance level for confidence intervals.

• controls (list of str, optional) – Time-invariant control variables to include.

• never_treated_values (list, optional) – Deprecated and ignored. Never-treated units are detected automatically.

Returns:

Cohort effect estimates sorted by cohort.

Return type:

list of CohortEffect

Raises:

NoNeverTreatedError – If no never-treated control units exist.

See also

CohortEffect

Result container for cohort-level estimates.

aggregate_to_overall

Aggregate across cohorts to overall effect.

cohort_effects_to_dataframe

Convert results to DataFrame format.

Notes

For each cohort g, the time-averaged transformed outcome is:

\[\bar{Y}_{ig} = \frac{1}{T - g + 1} \sum_{r=g}^{T} \dot{Y}_{irg}\]

where \(\dot{Y}_{irg}\) is the demeaned or detrended outcome for unit i in period r. The cohort ATT is estimated by regressing \(\bar{Y}_{ig}\) on cohort membership indicator \(D_{ig}\) using never-treated controls:

\[\bar{Y}_{ig} = \alpha + \tau_g D_{ig} + \varepsilon_i\]

The coefficient \(\tau_g\) estimates the time-averaged treatment effect for cohort g.

lwdid.staggered.aggregate_to_overall(data_transformed, gvar, ivar, tvar, transform_type='demean', vce=None, cluster_var=None, alpha=0.05, controls=None, never_treated_values=None)

Estimate overall cohort-size-weighted treatment effect.

Estimates the weighted average ATT across all treated cohorts via cross-sectional OLS regression on aggregated transformed outcomes with cohort weights proportional to cohort sizes and never-treated controls.

Parameters:

• data_transformed (pd.DataFrame) – Panel data with transformation columns named ‘ydot_g{g}_r{r}’ (demean) or ‘ycheck_g{g}_r{r}’ (detrend).

• gvar (str) – Cohort variable column name.

• ivar (str) – Unit identifier column name.

• tvar (str) – Time variable column name.

• transform_type ({'demean', 'detrend'}, default='demean') – Transformation type determining column prefix.

• vce ({None, 'robust', 'hc0', 'hc1', 'hc2', 'hc3', 'hc4', 'cluster'}, optional) – Variance-covariance estimator.

• cluster_var (str, optional) – Cluster variable for cluster-robust standard errors.

• alpha (float, default=0.05) – Significance level for confidence intervals.

• controls (list of str, optional) – Time-invariant control variables to include.

• never_treated_values (list, optional) – Deprecated and ignored. Never-treated units are detected automatically.

Returns:

Overall effect estimate with inference statistics and cohort weights.

Return type:

OverallEffect

Raises:

• NoNeverTreatedError – If no never-treated control units exist.

• ValueError – If sample size insufficient or no valid cohorts exist.

See also

OverallEffect

Result container for overall effect estimate.

aggregate_to_cohort

Estimate cohort-specific effects.

construct_aggregated_outcome

Helper function for outcome construction.

Notes

Cohort weights are proportional to cohort sizes:

\[w(g) = \frac{N_g}{\sum_{g'} N_{g'}}\]

For treated units, the aggregated outcome equals their cohort’s time-averaged transformation. For never-treated units, it is a weighted average across all cohorts’ transformations:

\[\bar{Y}_i = \sum_g w(g) \cdot \bar{Y}_{ig}\]

The overall ATT is estimated by regressing \(\bar{Y}_i\) on an ever-treated indicator with never-treated controls.

lwdid.staggered.aggregate_to_event_time(cohort_time_effects, cohort_sizes, alpha=0.05, event_time_range=None, df_strategy='conservative', verbose=False)

Aggregate cohort-time ATT estimates to event-time weighted ATT (WATT).

Computes WATT(r) = Σ_{g∈G_r} w(g,r) · ATT(g, g+r), where weights are proportional to cohort sizes. Uses t-distribution for proper inference.

Parameters:

• cohort_time_effects (pd.DataFrame or list) – Cohort-time specific ATT estimates. DataFrame must contain columns: ‘cohort’, ‘period’, ‘att’, ‘se’. Optional: ‘df_inference’. If list, must contain objects with these attributes.

• cohort_sizes (dict[int, int]) – Mapping from cohort identifier to number of treated units in that cohort. Used to compute weights w(g,r) = N_g / Σ N_g’.

• alpha (float, default=0.05) – Significance level for confidence intervals.

• event_time_range (tuple[int, int] or None, optional) – If provided, only compute WATT for event times in [min, max]. If None, compute for all available event times.

• df_strategy ({'conservative', 'weighted', 'fallback'}, default='conservative') – Strategy for selecting degrees of freedom: - ‘conservative’: min(df_g) across contributing cohorts - ‘weighted’: weighted average of df_g - ‘fallback’: n_cohorts - 1 (ignores individual df)

• verbose (bool, default=False) – If True, log diagnostic information about aggregation.

Returns:

Event-time aggregated effects sorted by event_time.

Return type:

list[EventTimeEffect]

Raises:

ValueError – If cohort_time_effects is empty or cohort_sizes contains invalid values.

Warns:

UserWarning – If weight sum deviates from 1.0, if cohorts are excluded due to missing data, or if df information is unavailable.

See also

EventTimeEffect

Result container for event-time estimates.

event_time_effects_to_dataframe

Convert results to DataFrame format.

_compute_event_time_weights

Weight computation helper.

_select_degrees_of_freedom

Degrees of freedom selection.

Notes

The weighted average treatment effect at event time r is:

\[\text{WATT}(r) = \sum_{g \in G_r} w(g, r) \cdot \widehat{\tau}_{g, g+r}\]

where \(G_r\) is the set of cohorts observed at event time r, and weights \(w(g, r) = N_g / \sum_{g' \in G_r} N_{g'}\) are proportional to cohort sizes. The standard error assumes independence across cohorts:

\[\text{SE}(\text{WATT}(r)) = \sqrt{\sum_{g \in G_r} w(g, r)^2 \cdot \text{SE}(\widehat{\tau}_{g, g+r})^2}\]

Confidence intervals use t-distribution rather than normal distribution, which provides proper inference for small samples.

lwdid.staggered.construct_aggregated_outcome(data, gvar, ivar, tvar, weights, cohorts, T_max, transform_type='demean', never_treated_values=None)

Construct aggregated outcome variable for overall effect estimation.

Computes unit-level aggregated outcomes with treatment-status-specific construction: treated units use their own cohort’s time-averaged transformation, while never-treated units use a cohort-weighted mixture.

Parameters:

• data (pd.DataFrame) – Panel data with transformation columns.

• gvar (str) – Cohort variable column name.

• ivar (str) – Unit identifier column name.

• tvar (str) – Time variable column name.

• weights (dict[int, float]) – Cohort weights (must sum to one, keys must match cohorts).

• cohorts (list of int) – Treatment cohorts to include.

• T_max (int) – Maximum time period.

• transform_type ({'demean', 'detrend'}, default='demean') – Transformation type determining column prefix.

• never_treated_values (list, optional) – Deprecated and ignored. Never-treated units are detected automatically.

Returns:

Aggregated outcome indexed by unit ID.

Return type:

pd.Series

Raises:

ValueError – If cohorts list is empty or weights keys do not match cohorts.

See also

aggregate_to_overall

Uses this function to construct outcomes.

_compute_cohort_aggregated_variable

Computes cohort-specific averages.

Notes

The aggregated outcome is constructed differently by treatment status:

• Treated units: Use their own cohort’s time-averaged transformation, \(\bar{Y}_i = \bar{Y}_{ig}\) where g is the unit’s cohort.

• Never-treated units: Use a weighted average across all cohorts, \(\bar{Y}_i = \sum_g w(g) \cdot \bar{Y}_{ig}\).

When some cohorts have missing data for a never-treated unit, weights are renormalized to available cohorts. This preserves unbiasedness if missingness is unrelated to potential outcomes.

lwdid.staggered.cohort_effects_to_dataframe(results)

Convert CohortEffect objects to a pandas DataFrame.

Parameters:

results (list of CohortEffect) – Cohort effect estimates from aggregate_to_cohort.

Returns:

Tabular representation with one row per cohort.

Return type:

pd.DataFrame

Raises:

TypeError – If results is None.

IPW Estimator

Inverse probability weighting estimates treatment effects by weighting observations based on their propensity scores.

class lwdid.staggered.IPWResult(att, se, ci_lower, ci_upper, t_stat, pvalue, propensity_scores, weights, propensity_model_coef, n_treated, n_control, df_resid=0, df_inference=0, weights_cv=0.0, diagnostics=None)

Result container for inverse probability weighting ATT estimation.

Variables:

• att (float) – ATT point estimate.

• se (float) – Standard error.

• ci_lower (float) – Lower bound of confidence interval.

• ci_upper (float) – Upper bound of confidence interval.

• t_stat (float) – t-statistic.

• pvalue (float) – Two-sided p-value.

• propensity_scores (np.ndarray) – Estimated propensity scores.

• weights (np.ndarray) – IPW weights w = p/(1-p) for control units.

• propensity_model_coef (dict[str, float]) – Propensity score model coefficients.

• n_treated (int) – Number of treated units.

• n_control (int) – Number of control units.

• df_resid (int) – Residual degrees of freedom.

• df_inference (int) – Degrees of freedom for inference.

• weights_cv (float) – Coefficient of variation of IPW weights.

• diagnostics (dict[str, Any] or None) – Additional diagnostic information.

att: float

se: float

ci_lower: float

ci_upper: float

t_stat: float

pvalue: float

propensity_scores: ndarray

weights: ndarray

propensity_model_coef: dict[str, float]

n_treated: int

n_control: int

df_resid: int = 0

df_inference: int = 0

weights_cv: float = 0.0

diagnostics: dict[str, Any] | None = None

lwdid.staggered.estimate_ipw(data, y, d, propensity_controls, trim_threshold=0.01, alpha=0.05, vce=None, return_diagnostics=False, gvar_col=None, ivar_col=None, cohort_g=None, period_r=None)

Estimate ATT using inverse probability weighting.

IPW estimates the average treatment effect on the treated by reweighting control observations to match the covariate distribution of treated units.

Parameters:

• data (pd.DataFrame) – Dataset containing outcome, treatment, and control variables.

• y (str) – Outcome variable column name.

• d (str) – Treatment indicator column name (1=treated, 0=control).

• propensity_controls (list[str]) – Variables for propensity score model.

• trim_threshold (float, default=0.01) – Propensity score trimming threshold. Observations with propensity scores outside [trim_threshold, 1-trim_threshold] are excluded.

• alpha (float, default=0.05) – Significance level for confidence intervals.

• vce (str, optional) – Variance estimator type (‘robust’, ‘hc0’, ‘hc1’, ‘hc2’, ‘hc3’).

• return_diagnostics (bool, default=False) – Whether to return additional diagnostic information.

• gvar_col (str, optional) – Column name for cohort indicator (for staggered designs).

• ivar_col (str, optional) – Column name for unit identifier (for staggered designs).

• cohort_g (int, optional) – Cohort value (for staggered designs).

• period_r (int, optional) – Period value (for staggered designs).

Returns:

Estimation results including ATT, standard error, and diagnostics.

Return type:

IPWResult

Raises:

ValueError – If no propensity controls are specified, no treated or control units exist, or all observations are trimmed.

Notes

The IPW estimator for ATT is:

\[\hat{\tau}_{IPW} = \frac{1}{N_1} \sum_{i:D_i=1} Y_i - \frac{1}{N_1} \sum_{i:D_i=0} \frac{\hat{p}(X_i)}{1-\hat{p}(X_i)} Y_i\]

where \(\hat{p}(X_i)\) is the estimated propensity score.

Inference uses the normal distribution based on influence function asymptotics, consistent with IPWRA and PSM estimators in this package.

IPWRA Estimator

The doubly robust IPWRA estimator combines regression adjustment and inverse probability weighting, providing consistent estimates when either the outcome model or the propensity score model is correctly specified.

class lwdid.staggered.IPWRAResult(att, se, ci_lower, ci_upper, t_stat, pvalue, propensity_scores, weights, outcome_model_coef, propensity_model_coef, n_treated, n_control)

Result container for doubly robust IPWRA ATT estimation.

Variables:

• att (float) – ATT point estimate.

• se (float) – Standard error.

• ci_lower (float) – Lower bound of confidence interval.

• ci_upper (float) – Upper bound of confidence interval.

• t_stat (float) – t-statistic.

• pvalue (float) – Two-sided p-value.

• propensity_scores (np.ndarray) – Estimated propensity scores.

• weights (np.ndarray) – IPW weights w = p/(1-p) for control units.

• outcome_model_coef (dict[str, float]) – Outcome model coefficients.

• propensity_model_coef (dict[str, float]) – Propensity score model coefficients.

• n_treated (int) – Number of treated units.

• n_control (int) – Number of control units.

att: float

se: float

ci_lower: float

ci_upper: float

t_stat: float

pvalue: float

propensity_scores: ndarray

weights: ndarray

outcome_model_coef: dict[str, float]

propensity_model_coef: dict[str, float]

n_treated: int

n_control: int

lwdid.staggered.estimate_ipwra(data, y, d, controls, propensity_controls=None, trim_threshold=0.01, se_method='analytical', n_bootstrap=200, seed=None, alpha=0.05, return_diagnostics=False, gvar_col=None, ivar_col=None, cohort_g=None, period_r=None)

Estimate ATT using inverse probability weighted regression adjustment.

IPWRA is a doubly robust estimator that combines propensity score weighting with outcome regression. The estimator remains consistent if either the propensity score model or the outcome regression model is correctly specified, providing protection against model misspecification.

The IPWRA-ATT estimator is computed as:

\[\hat{\tau} = \frac{1}{N_1} \sum_{D_i=1} [Y_i - \hat{m}_0(X_i)] - \frac{\sum_{D_i=0} w_i (Y_i - \hat{m}_0(X_i))} {\sum_{D_i=0} w_i}\]

where \(\hat{m}_0(X)\) is the estimated outcome regression for controls, and \(w_i = \hat{p}(X_i) / (1 - \hat{p}(X_i))\) are IPW weights.

Parameters:

• data (pd.DataFrame) – Cross-sectional data with one row per unit.

• y (str) – Outcome variable column name (typically transformed outcome).

• d (str) – Treatment indicator column name (1=treated, 0=control).

• controls (list[str]) – Control variables for the outcome regression model.

• propensity_controls (list[str], optional) – Control variables for the propensity score model. If None, uses the same variables as controls.

• trim_threshold (float, default=0.01) – Propensity score trimming threshold. Scores are clipped to [trim_threshold, 1-trim_threshold] to prevent extreme weights.

• se_method (str, default='analytical') – Standard error computation method: ‘analytical’ uses influence functions, ‘bootstrap’ uses nonparametric bootstrap.

• n_bootstrap (int, default=200) – Number of bootstrap replications when se_method=’bootstrap’.

• seed (int, optional) – Random seed for bootstrap resampling.

• alpha (float, default=0.05) – Significance level for confidence intervals.

• return_diagnostics (bool, default=False) – Whether to return additional diagnostic information.

• gvar_col (str, optional) – Cohort indicator column name (for staggered designs).

• ivar_col (str, optional) – Unit identifier column name (for staggered designs).

• cohort_g (int, optional) – Cohort value (for staggered designs).

• period_r (int, optional) – Period value (for staggered designs).

Returns:

Estimation results containing ATT estimate, standard error, confidence interval, propensity scores, and model coefficients.

Return type:

IPWRAResult

Raises:

ValueError – If required columns are missing, sample sizes are insufficient, or model estimation fails to converge.

See also

estimate_ipw

Pure IPW estimator without outcome regression.

estimate_psm

Propensity score matching estimator.

lwdid.staggered.estimate_propensity_score(data, d, controls, trim_threshold=0.01)

Estimate propensity scores using logistic regression.

Fits a logit model P(D=1|X) without regularization, using maximum likelihood estimation for unbiased coefficient estimates.

Parameters:

• data (pd.DataFrame) – Dataset containing treatment and control variables.

• d (str) – Treatment indicator column name.

• controls (list[str]) – Covariate column names for the propensity score model.

• trim_threshold (float, default=0.01) – Trimming threshold for extreme propensity scores.

Return type:

tuple[ndarray, dict[str, float]]

Returns:

• propensity_scores (np.ndarray) – Estimated propensity scores, trimmed to [trim_threshold, 1-trim_threshold].

• coefficients (dict[str, float]) – Dictionary mapping variable names to estimated coefficients, including ‘_intercept’ for the constant term.

Raises:

ValueError – If the logistic regression fails to converge.

lwdid.staggered.estimate_outcome_model(data, y, d, controls, sample_weights=None)

Estimate the outcome regression model on control units.

Fits a linear regression E(Y|X, D=0) using control observations only, then generates predicted values for all units. Optionally uses weighted least squares (WLS) when sample_weights are provided.

Parameters:

• data (pd.DataFrame) – Dataset containing outcome, treatment, and control variables.

• y (str) – Outcome variable column name.

• d (str) – Treatment indicator column name.

• controls (list[str]) – Covariate column names for the outcome model.

• sample_weights (np.ndarray, optional) – Weights for weighted least squares estimation. If provided, must have the same length as the data. Only weights for control units (D=0) are used in fitting. For IPWRA with ATT targeting, weights should be w_i = p(X_i) / (1 - p(X_i)) to reweight controls to the treated covariate distribution.

Return type:

tuple[ndarray, dict[str, float]]

Returns:

• predictions (np.ndarray) – Predicted outcome values for all units based on control regression.

• coefficients (dict[str, float]) – Dictionary mapping variable names to estimated coefficients, including ‘_intercept’ for the constant term.

Raises:

ValueError – If the design matrix is singular and cannot be inverted.

Notes

When sample_weights are provided, the outcome model is estimated using WLS:

\[\hat{\beta} = (X'WX)^{-1} X'WY\]

where W is the diagonal weight matrix. For IPWRA with ATT targeting, the weights for control units should be w_i = p(X_i) / (1 - p(X_i)), which reweights the control sample to match the treated covariate distribution.

PSM Estimator

class lwdid.staggered.PSMResult(att, se, ci_lower, ci_upper, t_stat, pvalue, propensity_scores, match_counts, matched_control_ids, n_treated, n_control, n_matched, caliper, n_dropped)

Result container for propensity score matching ATT estimation.

Variables:

• att (float) – ATT point estimate from nearest neighbor matching.

• se (float) – Standard error (heteroskedasticity-robust or bootstrap).

• ci_lower (float) – Lower bound of confidence interval.

• ci_upper (float) – Upper bound of confidence interval.

• t_stat (float) – t-statistic for testing ATT=0.

• pvalue (float) – Two-sided p-value.

• propensity_scores (np.ndarray) – Estimated propensity scores for all observations.

• match_counts (np.ndarray) – Number of matches for each treated unit.

• matched_control_ids (list[list[int]]) – List of matched control unit indices for each treated unit.

• n_treated (int) – Number of treated units.

• n_control (int) – Number of control units.

• n_matched (int) – Number of unique control units used in matching.

• caliper (float or None) – Caliper value used for matching, if any.

• n_dropped (int) – Number of treated units dropped due to caliper constraint.

att: float

se: float

ci_lower: float

ci_upper: float

t_stat: float

pvalue: float

propensity_scores: ndarray

match_counts: ndarray

matched_control_ids: list[list[int]]

n_treated: int

n_control: int

n_matched: int

caliper: float | None

n_dropped: int

lwdid.staggered.estimate_psm(data, y, d, propensity_controls, n_neighbors=1, with_replacement=True, caliper=None, caliper_scale='sd', trim_threshold=0.01, se_method='abadie_imbens', n_bootstrap=200, seed=None, alpha=0.05, match_order='data', return_diagnostics=False, gvar_col=None, ivar_col=None, cohort_g=None, period_r=None)

Estimate ATT using propensity score matching.

Uses nearest-neighbor matching on estimated propensity scores to find comparable control units for each treated unit. The ATT is estimated as the average difference between treated outcomes and matched control outcomes.

The PSM-ATT estimator is:

\[\hat{\tau} = \frac{1}{N_1} \sum_{D_i=1} \left[ Y_i - \frac{1}{k} \sum_{j \in M(i)} Y_j \right]\]

where M(i) is the set of k nearest-neighbor matches for unit i.

Parameters:

• data (pd.DataFrame) – Cross-sectional data with one row per unit.

• y (str) – Outcome variable column name (typically transformed outcome).

• d (str) – Treatment indicator column name (1=treated, 0=control).

• propensity_controls (list[str]) – Covariate column names for propensity score model.

• n_neighbors (int, default=1) – Number of control matches per treated unit (k).

• with_replacement (bool, default=True) – Whether to match with replacement.

• caliper (float, optional) – Maximum propensity score distance for valid matches.

• caliper_scale (str, default='sd') – Scale for caliper: ‘sd’ (standard deviation units) or ‘absolute’.

• trim_threshold (float, default=0.01) – Propensity score trimming threshold.

• se_method (str, default='abadie_imbens') – Standard error method: ‘abadie_imbens’ uses heteroskedasticity-robust variance estimation, ‘bootstrap’ uses nonparametric bootstrap.

• n_bootstrap (int, default=200) – Number of bootstrap replications when se_method=’bootstrap’.

• seed (int, optional) – Random seed for bootstrap resampling.

• alpha (float, default=0.05) – Significance level for confidence intervals.

• match_order (str, default='data') – Order in which to process treated units for matching.

• return_diagnostics (bool, default=False) – Whether to return additional diagnostic information.

• gvar_col (str, optional) – Cohort indicator column name (for staggered designs).

• ivar_col (str, optional) – Unit identifier column name (for staggered designs).

• cohort_g (int, optional) – Cohort value (for staggered designs).

• period_r (int, optional) – Period value (for staggered designs).

Returns:

Estimation results containing ATT estimate, standard error, confidence interval, and matching diagnostics.

Return type:

PSMResult

Raises:

ValueError – If required columns are missing, sample sizes are insufficient, or no valid matches can be found.

See also

estimate_ipwra

Doubly robust IPWRA estimator.

estimate_ipw

Pure IPW estimator.

Inference Distribution by Estimator

The reference distribution used for constructing confidence intervals and computing p-values varies by estimator. The following table summarizes the inference approach for each estimator in the staggered module:

Summary Table:

• RA (Regression Adjustment): t-distribution with df = N_treated + N_control - k

• IPW (Inverse Probability Weighting): Normal distribution (asymptotic inference)

• IPWRA (Doubly Robust): Normal distribution (asymptotic inference)

• PSM (Propensity Score Matching): Normal distribution (asymptotic inference)

Detailed Explanation:

RA (Regression Adjustment)

The regression adjustment estimator uses the t-distribution for inference, following Lee and Wooldridge (2026). Under classical linear model assumptions (normality and homoskedasticity), this provides exact finite-sample inference. With heteroskedasticity-robust standard errors (HC0-HC4), the t-distribution provides a conservative approximation that improves upon normal approximations in small samples.

IPW, IPWRA, and PSM

The IPW, IPWRA, and PSM estimators use the normal distribution for asymptotic inference because:

1. These estimators rely on influence function-based variance estimation

2. Asymptotic theory justifies normal approximations for these methods

3. For large samples, the normal distribution provides valid inference

For small samples where exact inference is desired, consider using RA with vce=None instead of IPW/IPWRA/PSM.

Practical Recommendations:

1. Small samples (N < 50): Use RA with vce=None for exact t-based inference, or use randomization inference (ri=True) for assumption-free testing.

2. Moderate samples (50 ≤ N < 200): Use RA or IPWRA with HC3 standard errors (vce='hc3').

3. Large samples (N ≥ 200): All estimators with asymptotic inference are appropriate; IPWRA is recommended when functional form assumptions are uncertain due to its double robustness property.

See Methodological Notes for detailed theoretical foundations.

Randomization Inference

class lwdid.staggered.StaggeredRIResult(p_value, ri_method, ri_reps, ri_valid, ri_failed, observed_stat, permutation_stats)

Container for staggered randomization inference results.

This dataclass stores the output from permutation or bootstrap-based randomization inference procedures, including the p-value, replication diagnostics, and the full distribution of resampled statistics.

Variables:

• p_value (float) – Two-sided p-value for the null hypothesis that the ATT equals zero.

• ri_method (str) – Resampling method used, either ‘permutation’ or ‘bootstrap’.

• ri_reps (int) – Number of requested replications.

• ri_valid (int) – Number of valid replications that produced non-missing statistics.

• ri_failed (int) – Number of replications that failed due to estimation errors.

• observed_stat (float) – Observed ATT estimate being tested.

• permutation_stats (NDArray[np.float64]) – Array of ATT statistics from valid replications, excluding NaN values from failed replications.

See also

randomization_inference_staggered

Main function that produces this result.

ri_overall_effect

Convenience wrapper for overall effect inference.

ri_cohort_effect

Convenience wrapper for cohort-specific inference.

p_value: float

ri_method: str

ri_reps: int

ri_valid: int

ri_failed: int

observed_stat: float

permutation_stats: ndarray[tuple[Any, ...], dtype[float64]]

__repr__()

Return a concise string representation of the result.

Return type:

str

lwdid.staggered.randomization_inference_staggered(data, gvar, ivar, tvar, y, observed_att, target='overall', target_cohort=None, target_period=None, ri_method='permutation', rireps=1000, seed=None, rolling='demean', controls=None, vce=None, cluster_var=None, n_never_treated=0)

Perform randomization inference for staggered DiD estimation.

Tests the null hypothesis that the average treatment effect on the treated equals zero by permuting or bootstrapping treatment cohort assignments and computing the empirical distribution of the test statistic.

Parameters:

• data (pd.DataFrame) – Panel data in long format with unit, time, cohort, and outcome columns.

• gvar (str) – Column name for the treatment cohort variable. Units with missing, zero, or infinite values are treated as never-treated.

• ivar (str) – Column name for the unit identifier.

• tvar (str) – Column name for the time period variable.

• y (str) – Column name for the outcome variable.

• observed_att (float) – Observed ATT estimate to be tested against the null hypothesis.

• target ({'overall', 'cohort', 'cohort_time'}, default 'overall') –

  Aggregation level for the target effect:

  • ’overall’: Overall weighted average effect across all cohorts

  • ’cohort’: Cohort-specific average effect (requires target_cohort)

  • ’cohort_time’: Effect for a specific (g, r) pair (requires both target_cohort and target_period)

• target_cohort (int, optional) – Target cohort for ‘cohort’ or ‘cohort_time’ targets.

• target_period (int, optional) – Target time period for ‘cohort_time’ target.

• ri_method ({'permutation', 'bootstrap'}, default 'permutation') –

  Resampling method for generating the null distribution:

  • ’permutation’: Without-replacement permutation preserving cohort sizes (Fisher exact randomization inference)

  • ’bootstrap’: With-replacement sampling from unit cohort assignments

• rireps (int, default 1000) – Number of resampling replications.

• seed (int, optional) – Random seed for reproducibility of the resampling procedure.

• rolling ({'demean', 'detrend'}, default 'demean') –

  Transformation method for removing pre-treatment variation:

  • ’demean’: Subtract pre-treatment mean from each unit

  • ’detrend’: Remove unit-specific linear time trend

• controls (list of str, optional) – Column names for control variables to include in estimation.

• vce (str, optional) – Variance-covariance estimator type for standard errors.

• cluster_var (str, optional) – Column name for clustering standard errors.

• n_never_treated (int, default 0) – Number of never-treated units. Required for overall effect inference to ensure a consistent reference group across permutations.

Returns:

Dataclass containing the p-value, replication counts, observed statistic, and array of permutation statistics.

Return type:

StaggeredRIResult

Raises:

RandomizationError – If input parameters are invalid, if there are insufficient units for inference, or if too few replications produce valid estimates.

See also

ri_overall_effect

Convenience wrapper for overall effect inference.

ri_cohort_effect

Convenience wrapper for cohort-specific inference.

Notes

The permutation procedure shuffles cohort assignments while preserving the marginal cohort distribution, generating the null distribution under the sharp null hypothesis of no treatment effect.

The two-sided p-value is computed as:

\[p = \frac{1}{R} \sum_{r=1}^{R} \mathbf{1}\{|\hat{\tau}^{(r)}| \geq |\hat{\tau}^{obs}|\}\]

where \(R\) is the number of valid replications and \(\hat{\tau}^{(r)}\) is the ATT from replication \(r\).

A minimum of 50 valid replications (or 10% of rireps, whichever is larger) is required for reliable p-value computation.

lwdid.staggered.ri_overall_effect(data, gvar, ivar, tvar, y, observed_att, rolling='demean', ri_method='permutation', rireps=1000, seed=None, vce=None, cluster_var=None)

Perform randomization inference for the overall weighted ATT.

This is a convenience wrapper around randomization_inference_staggered for testing the aggregate effect across all treatment cohorts. The overall effect is a cohort-share-weighted average of cohort-specific ATTs.

Parameters:

• data (pd.DataFrame) – Panel data in long format with unit, time, cohort, and outcome columns.

• gvar (str) – Column name for the treatment cohort variable.

• ivar (str) – Column name for the unit identifier.

• tvar (str) – Column name for the time period.

• y (str) – Column name for the outcome variable.

• observed_att (float) – Observed overall ATT estimate to test.

• rolling ({'demean', 'detrend'}, default 'demean') – Transformation method for pre-treatment variation removal.

• ri_method ({'permutation', 'bootstrap'}, default 'permutation') – Resampling method for null distribution generation.

• rireps (int, default 1000) – Number of resampling replications.

• seed (int, optional) – Random seed for reproducibility.

• vce (str, optional) – Variance-covariance estimator type.

• cluster_var (str, optional) – Column name for clustering standard errors.

Returns:

Randomization inference results including p-value and diagnostics.

Return type:

StaggeredRIResult

See also

randomization_inference_staggered

Full-featured inference function.

ri_cohort_effect

Inference for cohort-specific effects.

lwdid.staggered.ri_cohort_effect(data, gvar, ivar, tvar, y, target_cohort, observed_att, rolling='demean', ri_method='permutation', rireps=1000, seed=None, vce=None, cluster_var=None)

Perform randomization inference for a cohort-specific ATT.

This is a convenience wrapper around randomization_inference_staggered for testing the average effect for a specific treatment cohort. The cohort-specific ATT averages effects across all post-treatment periods for units first treated in the target cohort.

Parameters:

• data (pd.DataFrame) – Panel data in long format with unit, time, cohort, and outcome columns.

• gvar (str) – Column name for the treatment cohort variable.

• ivar (str) – Column name for the unit identifier.

• tvar (str) – Column name for the time period.

• y (str) – Column name for the outcome variable.

• target_cohort (int) – Treatment cohort (first treatment period) to test.

• observed_att (float) – Observed cohort-specific ATT estimate to test.

• rolling ({'demean', 'detrend'}, default 'demean') – Transformation method for pre-treatment variation removal.

• ri_method ({'permutation', 'bootstrap'}, default 'permutation') – Resampling method for null distribution generation.

• rireps (int, default 1000) – Number of resampling replications.

• seed (int, optional) – Random seed for reproducibility.

• vce (str, optional) – Variance-covariance estimator type.

• cluster_var (str, optional) – Column name for clustering standard errors.

Returns:

Randomization inference results including p-value and diagnostics.

Return type:

StaggeredRIResult

See also

randomization_inference_staggered

Full-featured inference function.

ri_overall_effect

Inference for overall weighted effect.

Pre-treatment Dynamics

The pre-treatment dynamics module implements estimation and testing for pre-treatment periods, following Lee and Wooldridge (2025).

Pre-treatment Transformations

lwdid.staggered.transform_staggered_demean_pre(data, y, ivar, tvar, gvar, never_treated_values=None)

Apply cohort-specific rolling demeaning for pre-treatment periods.

For each cohort g and pre-treatment period t < g, computes:

\[\dot{Y}_{itg} = Y_{it} - \frac{1}{g-t-1} \sum_{q=t+1}^{g-1} Y_{iq}\]

The anchor point t = g-1 is set to 0 by convention, as the future window {g, …, g-1} is empty.

Parameters:

• data (pd.DataFrame) – Long-format panel data with columns for outcome, unit identifier, time period, and cohort assignment.

• y (str) – Outcome variable column name.

• ivar (str) – Unit identifier column name.

• tvar (str) – Time variable column name (must be numeric or coercible to numeric).

• gvar (str) – Cohort variable column name indicating first treatment period.

• never_treated_values (sequence or None, optional) – Values in gvar indicating never-treated units. Default recognizes NaN, 0, and np.inf as never-treated indicators.

Returns:

Copy of input data with additional columns named ydot_pre_g{g}_t{t} containing the transformed outcome for each (cohort, period) pair.

Return type:

pd.DataFrame

Raises:

ValueError – If required columns are missing or no valid cohorts exist.

See also

transform_staggered_detrend_pre

Apply linear detrending transformation.

transform_staggered_demean

Post-treatment demeaning transformation.

Notes

The transformation is applied to ALL units (treated, not-yet-treated, and never-treated) for each cohort, enabling flexible control group selection during estimation.

lwdid.staggered.transform_staggered_detrend_pre(data, y, ivar, tvar, gvar, never_treated_values=None)

Apply cohort-specific rolling detrending for pre-treatment periods.

For each cohort g and pre-treatment period t ≤ g-3, computes:

\[\ddot{Y}_{itg} = Y_{it} - \hat{Y}_{itg}\]

where Ŷ_{itg} is the OLS-fitted value from regressing Y_{iq} on q for q ∈ {t+1, …, g-1}.

The anchor points are: - t = g-1: Set to 0.0 (empty future window) - t = g-2: Set to NaN (only 1 future period, need 2 for OLS)

Parameters:

• data (pd.DataFrame) – Long-format panel data with columns for outcome, unit identifier, time period, and cohort assignment.

• y (str) – Outcome variable column name.

• ivar (str) – Unit identifier column name.

• tvar (str) – Time variable column name (must be numeric or coercible to numeric).

• gvar (str) – Cohort variable column name indicating first treatment period.

• never_treated_values (sequence or None, optional) – Values in gvar indicating never-treated units. Default recognizes NaN, 0, and np.inf as never-treated indicators.

Returns:

Copy of input data with additional columns named ycheck_pre_g{g}_t{t} containing the detrended outcome for each (cohort, period) pair.

Return type:

pd.DataFrame

Raises:

ValueError – If required columns are missing or no valid cohorts exist.

Warns:

UserWarning – If a cohort has fewer than 3 pre-treatment periods (detrending requires at least 2 future periods for OLS).

See also

transform_staggered_demean_pre

Apply demeaning transformation.

transform_staggered_detrend

Post-treatment detrending transformation.

Notes

Detrending requires at least 2 future pre-treatment periods to estimate both intercept and slope. For t = g-2, only one future period exists, so the detrended value is NaN. For t = g-1 (anchor), the value is 0.

The transformation is applied to ALL units for each cohort.

lwdid.staggered.get_pre_treatment_periods_for_cohort(cohort, T_min)

Return list of pre-treatment periods for a treatment cohort.

Generates the sequence {T_min, T_min+1, …, g-1} representing all calendar periods strictly before the first treatment period g.

Parameters:

• cohort (int) – Treatment cohort identifier (first treatment period g).

• T_min (int) – Minimum time period in the panel data.

Returns:

Pre-treatment periods in ascending order.

Return type:

list of int

See also

get_valid_periods_for_cohort

Get post-treatment periods for a cohort.

Pre-treatment Estimation

class lwdid.staggered.PreTreatmentEffect(cohort, period, event_time, att, se, ci_lower, ci_upper, t_stat, pvalue, n_treated, n_control, is_anchor=False, rolling_window_size=0, df_inference=0)

Container for a pre-treatment period ATT estimate.

Stores the ATT for treatment cohort g at pre-treatment period t < g, along with inference statistics. Under parallel trends, these estimates should be approximately zero.

Variables:

• cohort (int) – Treatment cohort identifier (first treatment period g).

• period (int) – Calendar time period t (where t < g).

• event_time (int) – Event time relative to treatment onset (e = t - g, always negative).

• att (float) – Estimated pre-treatment ATT (should be ~0 under parallel trends).

• se (float) – Standard error of the ATT estimate.

• ci_lower (float) – Lower bound of confidence interval.

• ci_upper (float) – Upper bound of confidence interval.

• t_stat (float) – t-statistic for testing H0: ATT = 0.

• pvalue (float) – Two-sided p-value.

• n_treated (int) – Number of treated units in estimation sample.

• n_control (int) – Number of control units in estimation sample.

• is_anchor (bool) – True if this is the anchor point (t = g-1, e = -1). Anchor points have ATT = 0, SE = 0 by convention.

• rolling_window_size (int) – Number of future periods used in transformation {t+1, …, g-1}.

Notes

The dataclass is frozen (immutable) to ensure estimation results cannot be accidentally modified after creation.

cohort: int

period: int

event_time: int

att: float

se: float

ci_lower: float

ci_upper: float

t_stat: float

pvalue: float

n_treated: int

n_control: int

is_anchor: bool = False

rolling_window_size: int = 0

df_inference: int = 0

lwdid.staggered.estimate_pre_treatment_effects(data_transformed, gvar, ivar, tvar, controls=None, vce=None, cluster_var=None, control_strategy='not_yet_treated', never_treated_values=None, min_obs=3, min_treated=1, min_control=1, alpha=0.05, estimator='ra', transform_type='demean', propensity_controls=None, trim_threshold=0.01, se_method='analytical', n_neighbors=1, with_replacement=True, caliper=None)

Estimate pre-treatment effects for all valid cohort-period pairs.

Iterates over all treatment cohorts and their pre-treatment periods, estimating the ATT for each (cohort, period) combination. Under the parallel trends assumption, these estimates should be approximately zero.

Parameters:

• data_transformed (pd.DataFrame) – Panel data containing pre-treatment transformed outcome columns generated by transform_staggered_demean_pre or transform_staggered_detrend_pre. Must include columns for gvar, ivar, tvar, and transformed outcomes named ‘ydot_pre_g{g}_t{t}’ or ‘ycheck_pre_g{g}_t{t}’.

• gvar (str) – Name of the cohort variable column indicating first treatment period.

• ivar (str) – Name of the unit identifier column.

• tvar (str) – Name of the time variable column.

• controls (list of str, optional) – Names of time-invariant control variable columns.

• vce (str, optional) – Variance estimation type: None (homoskedastic), ‘hc3’ (heteroskedasticity-robust), or ‘cluster’ (cluster-robust).

• cluster_var (str, optional) – Name of the cluster variable column. Required when vce=’cluster’.

• control_strategy (str, default='not_yet_treated') – Control group selection strategy: ‘never_treated’ uses only never-treated units; ‘not_yet_treated’ includes units first treated after the current period; ‘all_others’ uses all units not in the treatment cohort (including already-treated units).

• never_treated_values (list, optional) – Values in gvar indicating never-treated units. Defaults to [0, np.inf] and NaN values.

• min_obs (int, default=3) – Minimum total sample size required for estimation.

• min_treated (int, default=1) – Minimum number of treated units required.

• min_control (int, default=1) – Minimum number of control units required.

• alpha (float, default=0.05) – Significance level for confidence interval construction.

• estimator (str, default='ra') – Estimation method: ‘ra’ (regression adjustment), ‘ipwra’ (inverse probability weighted regression adjustment), or ‘psm’ (propensity score matching).

• transform_type (str, default='demean') – Transformation type applied to the data: ‘demean’ or ‘detrend’. Determines the column prefix for transformed outcomes.

• propensity_controls (list of str, optional) – Control variables for the propensity score model. If None, uses the same variables as controls.

• trim_threshold (float, default=0.01) – Propensity score trimming threshold for IPWRA and PSM.

• se_method (str, default='analytical') – Standard error method for IPWRA: ‘analytical’ or ‘bootstrap’.

• n_neighbors (int, default=1) – Number of nearest neighbors for PSM matching.

• with_replacement (bool, default=True) – Whether PSM matching allows replacement.

• caliper (float, optional) – Maximum propensity score distance for PSM matching.

Returns:

Estimation results for all valid (cohort, period) pairs, sorted by cohort and then by event_time (descending, so anchor point comes first).

Return type:

list of PreTreatmentEffect

Raises:

ValueError – If required columns are missing, no valid treatment cohorts exist, or parameter values are invalid.

See also

transform_staggered_demean_pre

Pre-treatment demeaning transformation.

transform_staggered_detrend_pre

Pre-treatment detrending transformation.

test_parallel_trends

Statistical test for parallel trends assumption.

Notes

The anchor point (t = g-1, event time e = -1) is handled specially:

• ATT is set to exactly 0.0 (by construction of the transformation)

• SE is set to 0.0

• is_anchor flag is set to True

For pre-treatment periods, the control group is defined as: control = {units with gvar > t} ∪ {never-treated units}.

lwdid.staggered.pre_treatment_effects_to_dataframe(results)

Convert a list of PreTreatmentEffect objects to a pandas DataFrame.

Parameters:

results (list of PreTreatmentEffect) – Estimation results from estimate_pre_treatment_effects.

Returns:

DataFrame with columns: cohort, period, event_time, att, se, ci_lower, ci_upper, t_stat, pvalue, n_treated, n_control, is_anchor, rolling_window_size. Returns an empty DataFrame with appropriate columns if the input list is empty.

Return type:

pd.DataFrame

See also

estimate_pre_treatment_effects

Estimate pre-treatment ATT.

PreTreatmentEffect

Container class for individual effect estimates.

Parallel Trends Testing

class lwdid.staggered.ParallelTrendsTestResult(individual_tests, joint_f_stat, joint_pvalue, joint_df1, joint_df2, n_pre_periods, excluded_periods=<factory>, reject_null=False, alpha=0.05)

Container for parallel trends test results.

Stores both individual t-tests for each pre-treatment period and the joint F-test for the null hypothesis that all pre-treatment ATT estimates equal zero.

Variables:

• individual_tests (pd.DataFrame) – DataFrame with columns: event_time, att, se, t_stat, pvalue. Contains test results for each pre-treatment period (excluding anchor point).

• joint_f_stat (float) – F-statistic for joint test H0: all pre-treatment ATT = 0.

• joint_pvalue (float) – P-value for joint F-test.

• joint_df1 (int) – Numerator degrees of freedom (number of pre-treatment periods included in the test).

• joint_df2 (int) – Denominator degrees of freedom.

• n_pre_periods (int) – Number of pre-treatment periods included in test.

• excluded_periods (list) – Event times excluded due to missing SE, anchor point, or other issues.

• reject_null (bool) – True if joint test rejects H0 at the specified alpha level.

• alpha (float) – Significance level used for the test.

Notes

Rejection of the null hypothesis suggests potential violation of the parallel trends assumption. However, failure to reject does not prove parallel trends holds - it may simply reflect low power.

individual_tests: DataFrame

joint_f_stat: float

joint_pvalue: float

joint_df1: int

joint_df2: int

n_pre_periods: int

excluded_periods: list

reject_null: bool = False

alpha: float = 0.05

lwdid.staggered.run_parallel_trends_test(pre_treatment_effects, alpha=0.05, test_type='f', min_pre_periods=2)

Test the parallel trends assumption using pre-treatment ATT estimates.

Performs both individual t-tests for each pre-treatment period and a joint test for the null hypothesis that all pre-treatment ATT estimates equal zero.

Parameters:

• pre_treatment_effects (list of PreTreatmentEffect) – Pre-treatment effect estimates from estimate_pre_treatment_effects.

• alpha (float, default=0.05) – Significance level for hypothesis tests.

• test_type (str, default='f') – Type of joint test: ‘f’ for F-test, ‘wald’ for Wald/chi-squared test.

• min_pre_periods (int, default=2) – Minimum number of pre-treatment periods required for joint test. If fewer periods are available, a warning is issued.

Returns:

Test results including individual t-tests and joint test.

Return type:

ParallelTrendsTestResult

Raises:

ValueError – If pre_treatment_effects is empty or test_type is invalid.

See also

estimate_pre_treatment_effects

Estimate pre-treatment ATT.

PreTreatmentEffect

Container for pre-treatment effect estimates.

Notes

The anchor point at event time e = -1 is excluded from testing because it is exactly zero by construction of the rolling transformation.

Test interpretation guidelines:

• reject_null = True: Evidence against parallel trends. The pre-treatment ATT estimates are jointly significantly different from zero, suggesting potential violation of the identifying assumption.

• reject_null = False: No evidence against parallel trends. This does not prove the assumption holds; the test may lack sufficient power to detect violations, particularly with few pre-treatment periods or small sample sizes.

lwdid.staggered.summarize_parallel_trends_test(test_result)

Generate a human-readable summary of parallel trends test results.

Produces a formatted text report containing joint test statistics, individual period-specific test results, and interpretation guidance.

Parameters:

test_result (ParallelTrendsTestResult) – Test results from run_parallel_trends_test.

Returns:

Multi-line formatted summary including: - Joint test F-statistic, p-value, and degrees of freedom - Individual t-tests for each pre-treatment event time - List of excluded event times (anchor points and missing data) - Plain-language interpretation of the test outcome

Return type:

str

See also

run_parallel_trends_test

Compute parallel trends test statistics.

ParallelTrendsTestResult

Container for test results.

Examples

Basic Staggered Estimation

import pandas as pd
from lwdid import lwdid

# Load Castle Law data
data = pd.read_csv('castle.csv')

# Create gvar from effyear (NaN = never treated -> 0)
data['gvar'] = data['effyear'].fillna(0).astype(int)

# Run staggered DiD
results = lwdid(
    data=data,
    y='lhomicide',        # Log homicide rate
    ivar='sid',           # State ID (integer)
    tvar='year',          # Year
    gvar='gvar',          # First treatment year
    rolling='demean',     # Demeaning transformation
    control_group='never_treated',  # Use only never-treated as controls
    aggregate='overall',  # Get overall weighted effect
    vce='hc3'            # HC3 standard errors
)

# View results
print(results.summary())
print(f"Overall ATT: {results.att_overall:.4f}")
print(f"SE: {results.se_overall:.4f}")

Cohort-Specific Effects

# Get cohort-specific effects
results = lwdid(
    data=data,
    y='lhomicide',
    ivar='sid',
    tvar='year',
    gvar='gvar',
    rolling='demean',
    aggregate='cohort',   # Aggregate within cohorts
)

# Access cohort effects
print(results.att_by_cohort)
# Columns: cohort, att, se, ci_lower, ci_upper, t_stat, pvalue, n_units, n_periods

All (g, r) Effects

# Get all cohort-time specific effects
results = lwdid(
    data=data,
    y='lhomicide',
    ivar='sid',
    tvar='year',
    gvar='gvar',
    rolling='demean',
    aggregate='none',     # No aggregation
)

# Access all (g,r) effects
print(results.att_by_cohort_time)
# Columns: cohort, period, event_time, att, se, ci_lower, ci_upper, t_stat, pvalue, n_treated, n_control

Event Study Plot

# Generate event study plot
results = lwdid(
    data=data,
    y='lhomicide',
    ivar='sid',
    tvar='year',
    gvar='gvar',
    aggregate='none',
)

# Plot event study
fig = results.plot_event_study(
    title='Castle Doctrine Effect',
    ylabel='Effect on Log Homicide Rate'
)
fig.savefig('event_study.png', dpi=300)

Event Time Aggregation (WATT)

from lwdid.staggered import aggregate_to_event_time

# Get all cohort-time specific effects
results = lwdid(
    data=data,
    y='lhomicide',
    ivar='sid',
    tvar='year',
    gvar='gvar',
    aggregate='none',
)

# Aggregate to event time using WATT (Weighted ATT)
# WATT(r) = Σ w(g,r) × ATT(g, g+r) where w(g,r) = N_g / Σ N_g'
watt_effects = aggregate_to_event_time(
    cohort_time_effects=results.att_by_cohort_time,
    cohort_sizes=results.cohort_sizes,
    alpha=0.05,
    df_strategy='conservative',  # Use min(df) across cohorts
)

# Access event-time aggregated effects
for e in watt_effects:
    print(f"Event time {e.event_time}: WATT={e.att:.4f}, SE={e.se:.4f}, "
          f"CI=[{e.ci_lower:.4f}, {e.ci_upper:.4f}], p={e.pvalue:.4f}")

# Or use plot_event_study with weighted aggregation
fig, ax = results.plot_event_study(
    aggregation='weighted',  # Use WATT aggregation
    title='Event Study with WATT Aggregation',
)

Pre-treatment Dynamics and Parallel Trends Testing

# Estimate with pre-treatment dynamics for parallel trends assessment
results = lwdid(
    data=data,
    y='lhomicide',
    ivar='sid',
    tvar='year',
    gvar='gvar',
    rolling='demean',
    aggregate='cohort',
    include_pretreatment=True,    # Enable pre-treatment estimation
    pretreatment_test=True,       # Run parallel trends test
    pretreatment_alpha=0.05,      # Significance level
)

# View summary with pre-treatment results
print(results.summary())

# Access pre-treatment ATT estimates
print(results.att_pre_treatment)
# Columns: cohort, period, event_time, att, se, ci_lower, ci_upper,
#          t_stat, pvalue, n_treated, n_control, is_anchor, rolling_window_size

# Access parallel trends test results
pt = results.parallel_trends_test
print(f"Joint F-stat: {pt.joint_f_stat:.4f}")
print(f"P-value: {pt.joint_pvalue:.4f}")
print(f"Reject H0: {pt.reject_null}")

# Plot event study with pre-treatment effects
fig, ax = results.plot_event_study(
    include_pre_treatment=True,
    title='Event Study with Pre-treatment Effects',
    pre_treatment_color='gray',
    post_treatment_color='blue',
)
fig.savefig('event_study_pretreatment.png', dpi=300)

Low-Level API Usage

For more control, the staggered module can be used directly:

from lwdid.staggered import (
    transform_staggered_demean,
    estimate_cohort_time_effects,
    aggregate_to_overall,
    ControlGroupStrategy
)

# Step 1: Transform data
data_transformed = transform_staggered_demean(
    data=data,
    y='lhomicide',
    ivar='sid',
    tvar='year',
    gvar='gvar'
)

# Step 2: Estimate (g,r) effects
effects = estimate_cohort_time_effects(
    data_transformed=data_transformed,
    gvar='gvar',
    ivar='sid',
    tvar='year',
    control_strategy='never_treated',
    vce='hc3'
)

# Step 3: Aggregate to overall effect
overall = aggregate_to_overall(
    data_transformed=data_transformed,
    gvar='gvar',
    ivar='sid',
    tvar='year',
    cohorts=[2005, 2006, 2007, 2008, 2009],
    T_max=2010,
    vce='hc3'
)

print(f"Overall ATT: {overall.att:.4f} (SE: {overall.se:.4f})")

See Also

• lwdid.lwdid() - Main estimation function with staggered support

• User Guide - Comprehensive usage guide

• Methodological Notes - Theoretical foundations

• Examples - Examples including Castle Law analysis