"""
Goodman-Bacon Decomposition for Two-Way Fixed Effects.
Implements the decomposition from Goodman-Bacon (2021) that shows how
TWFE estimates with staggered treatment timing can be written as a
weighted average of all possible 2x2 DiD comparisons.
Reference:
Goodman-Bacon, A. (2021). Difference-in-differences with variation
in treatment timing. Journal of Econometrics, 225(2), 254-277.
"""
import warnings
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple
import numpy as np
import pandas as pd
from diff_diff.results import _format_survey_block
from diff_diff.utils import within_transform as _within_transform_util
@dataclass
class Comparison2x2:
"""
A single 2x2 DiD comparison in the Bacon decomposition.
Attributes
----------
treated_group : Any
The timing group used as "treated" in this comparison.
control_group : Any
The timing group used as "control" in this comparison.
comparison_type : str
Type of comparison: "treated_vs_never", "earlier_vs_later",
or "later_vs_earlier".
estimate : float
The 2x2 DiD estimate for this comparison.
weight : float
The weight assigned to this comparison in the TWFE average.
n_treated : int
Number of treated observations in this comparison.
n_control : int
Number of control observations in this comparison.
time_window : Tuple[float, float]
The (start, end) time period for this comparison.
"""
treated_group: Any
control_group: Any
comparison_type: str
estimate: float
weight: float
n_treated: int
n_control: int
time_window: Tuple[float, float]
def __repr__(self) -> str:
return (
f"Comparison2x2({self.treated_group} vs {self.control_group}, "
f"type={self.comparison_type}, ={self.estimate:.4f}, "
f"weight={self.weight:.4f})"
)
@dataclass
class BaconDecompositionResults:
"""
Results from Goodman-Bacon decomposition of TWFE.
This decomposition shows that the TWFE estimate equals a weighted
average of all possible 2x2 DiD comparisons between timing groups.
Attributes
----------
twfe_estimate : float
The overall TWFE coefficient (should equal weighted sum of 2x2 estimates).
comparisons : List[Comparison2x2]
List of all 2x2 comparisons with their estimates and weights.
total_weight_treated_vs_never : float
Total weight on treated vs never-treated comparisons.
total_weight_earlier_vs_later : float
Total weight on earlier vs later treated comparisons.
total_weight_later_vs_earlier : float
Total weight on later vs earlier treated comparisons (forbidden).
weighted_avg_treated_vs_never : float
Weighted average effect from treated vs never-treated comparisons.
weighted_avg_earlier_vs_later : float
Weighted average effect from earlier vs later comparisons.
weighted_avg_later_vs_earlier : float
Weighted average effect from later vs earlier comparisons.
n_timing_groups : int
Number of distinct treatment timing groups.
n_never_treated : int
Number of never-treated units.
timing_groups : List[Any]
List of treatment timing cohorts.
"""
twfe_estimate: float
comparisons: List[Comparison2x2]
total_weight_treated_vs_never: float
total_weight_earlier_vs_later: float
total_weight_later_vs_earlier: float
weighted_avg_treated_vs_never: Optional[float]
weighted_avg_earlier_vs_later: Optional[float]
weighted_avg_later_vs_earlier: Optional[float]
n_timing_groups: int
n_never_treated: int
timing_groups: List[Any]
n_obs: int = 0
decomposition_error: float = field(default=0.0)
# Survey design metadata (SurveyMetadata instance from diff_diff.survey)
survey_metadata: Optional[Any] = field(default=None)
def __repr__(self) -> str:
return (
f"BaconDecompositionResults(TWFE={self.twfe_estimate:.4f}, "
f"n_comparisons={len(self.comparisons)}, "
f"n_groups={self.n_timing_groups})"
)
def summary(self) -> str:
"""
Generate a formatted summary of the decomposition.
Returns
-------
str
Formatted summary table.
"""
lines = [
"=" * 85,
"Goodman-Bacon Decomposition of Two-Way Fixed Effects".center(85),
"=" * 85,
"",
f"{'Total observations:':10}",
f"{'Treatment timing groups:':10}",
f"{'Never-treated units:':10}",
f"{'Total 2x2 comparisons:':10}",
"",
]
# Add survey design info
if self.survey_metadata is not None:
sm = self.survey_metadata
lines.extend(_format_survey_block(sm, 85))
lines.extend(
[
"-" * 85,
"TWFE Decomposition".center(85),
"-" * 85,
"",
f"{'TWFE Estimate:':12.4f}",
f"{'Weighted Sum of 2x2 Estimates:':12.4f}",
f"{'Decomposition Error:':12.6f}",
"",
]
)
# Weight breakdown by comparison type
lines.extend(
[
"-" * 85,
"Weight Breakdown by Comparison Type".center(85),
"-" * 85,
f"{'Comparison Type':12} {'Avg Effect':>12} {'Contribution':>12}",
"-" * 85,
]
)
# Treated vs Never-treated
if self.total_weight_treated_vs_never > 0:
contrib = self.total_weight_treated_vs_never * (self.weighted_avg_treated_vs_never or 0)
lines.append(
f"{'Treated vs Never-treated':12.4f} "
f"{self.weighted_avg_treated_vs_never or 0:>12.4f} "
f"{contrib:>12.4f}"
)
# Earlier vs Later
if self.total_weight_earlier_vs_later > 0:
contrib = self.total_weight_earlier_vs_later * (self.weighted_avg_earlier_vs_later or 0)
lines.append(
f"{'Earlier vs Later treated':12.4f} "
f"{self.weighted_avg_earlier_vs_later or 0:>12.4f} "
f"{contrib:>12.4f}"
)
# Later vs Earlier (forbidden)
if self.total_weight_later_vs_earlier > 0:
contrib = self.total_weight_later_vs_earlier * (self.weighted_avg_later_vs_earlier or 0)
lines.append(
f"{'Later vs Earlier (forbidden)':12.4f} "
f"{self.weighted_avg_later_vs_earlier or 0:>12.4f} "
f"{contrib:>12.4f}"
)
lines.extend(
[
"-" * 85,
f"{'Total':12.4f} "
f"{'':>12} {self._weighted_sum():>12.4f}",
"-" * 85,
"",
]
)
# Warning about forbidden comparisons
if self.total_weight_later_vs_earlier > 0.01:
pct = self.total_weight_later_vs_earlier * 100
lines.extend(
[
"WARNING: {:.1f}% of weight is on 'forbidden' comparisons where".format(pct),
"already-treated units serve as controls. This can bias TWFE",
"when treatment effects are heterogeneous over time.",
"",
"Consider using Callaway-Sant'Anna or other robust estimators.",
"",
]
)
lines.append("=" * 85)
return "\n".join(lines)
def print_summary(self) -> None:
"""Print the summary to stdout."""
print(self.summary())
def _weighted_sum(self) -> float:
"""Calculate weighted sum of 2x2 estimates."""
return sum(c.weight * c.estimate for c in self.comparisons)
def _total_weight(self) -> float:
"""Calculate total weight (should be 1.0)."""
return sum(c.weight for c in self.comparisons)
def to_dataframe(self) -> pd.DataFrame:
"""
Convert comparisons to a DataFrame.
Returns
-------
pd.DataFrame
DataFrame with one row per 2x2 comparison.
"""
rows = []
for c in self.comparisons:
rows.append(
{
"treated_group": c.treated_group,
"control_group": c.control_group,
"comparison_type": c.comparison_type,
"estimate": c.estimate,
"weight": c.weight,
"n_treated": c.n_treated,
"n_control": c.n_control,
"time_start": c.time_window[0],
"time_end": c.time_window[1],
}
)
return pd.DataFrame(rows)
def weight_by_type(self) -> Dict[str, float]:
"""
Get total weight by comparison type.
Returns
-------
Dict[str, float]
Dictionary mapping comparison type to total weight.
"""
return {
"treated_vs_never": self.total_weight_treated_vs_never,
"earlier_vs_later": self.total_weight_earlier_vs_later,
"later_vs_earlier": self.total_weight_later_vs_earlier,
}
def effect_by_type(self) -> Dict[str, Optional[float]]:
"""
Get weighted average effect by comparison type.
Returns
-------
Dict[str, Optional[float]]
Dictionary mapping comparison type to weighted average effect.
"""
return {
"treated_vs_never": self.weighted_avg_treated_vs_never,
"earlier_vs_later": self.weighted_avg_earlier_vs_later,
"later_vs_earlier": self.weighted_avg_later_vs_earlier,
}
class BaconDecomposition:
"""
Goodman-Bacon (2021) decomposition of Two-Way Fixed Effects estimator.
This class decomposes a TWFE estimate into a weighted average of all
possible 2x2 DiD comparisons, revealing the implicit comparisons that
drive the TWFE estimate and their relative importance.
The decomposition identifies three types of comparisons:
1. **Treated vs Never-treated**: Uses never-treated units as controls.
These are "clean" comparisons without bias concerns.
2. **Earlier vs Later treated**: Units treated earlier are compared to
units treated later, using the later group as controls before they
are treated. These are valid comparisons.
3. **Later vs Earlier treated**: Units treated later are compared to
units treated earlier, using the earlier group as controls AFTER
they are already treated. These are "forbidden comparisons" that
can introduce bias when treatment effects vary over time.
Parameters
----------
weights : str, default="approximate"
Weight calculation method:
- "approximate": Fast simplified formula using group shares and
treatment variance. Good for diagnostic purposes where relative
weights are sufficient to identify problematic comparisons.
- "exact": Variance-based weights from Goodman-Bacon (2021) Theorem 1.
Use for publication-quality decompositions where the weighted sum
must closely match the TWFE estimate.
Attributes
----------
weights : str
The weight calculation method.
results_ : BaconDecompositionResults
Decomposition results after calling fit().
is_fitted_ : bool
Whether the model has been fitted.
Examples
--------
Basic usage:
>>> import pandas as pd
>>> from diff_diff import BaconDecomposition
>>>
>>> # Panel data with staggered treatment
>>> data = pd.DataFrame({
... 'unit': [...],
... 'time': [...],
... 'outcome': [...],
... 'first_treat': [...] # 0 for never-treated
... })
>>>
>>> bacon = BaconDecomposition()
>>> results = bacon.fit(data, outcome='outcome', unit='unit',
... time='time', first_treat='first_treat')
>>> results.print_summary()
Visualizing the decomposition:
>>> from diff_diff import plot_bacon
>>> plot_bacon(results)
Notes
-----
The key insight from Goodman-Bacon (2021) is that TWFE with staggered
treatment timing implicitly makes comparisons using already-treated
units as controls. When treatment effects are dynamic (changing over
time since treatment), these "forbidden comparisons" can bias the
TWFE estimate, potentially even reversing its sign.
The decomposition helps diagnose this issue by showing:
- How much weight is on each type of comparison
- Whether forbidden comparisons contribute significantly to the estimate
- How the 2x2 estimates vary across comparison types
If forbidden comparisons have substantial weight and different estimates
than clean comparisons, consider using robust estimators like
Callaway-Sant'Anna that avoid these problematic comparisons.
References
----------
Goodman-Bacon, A. (2021). Difference-in-differences with variation in
treatment timing. Journal of Econometrics, 225(2), 254-277.
See Also
--------
CallawaySantAnna : Robust estimator for staggered DiD
TwoWayFixedEffects : The TWFE estimator being decomposed
"""
def __init__(self, weights: str = "approximate"):
"""
Initialize BaconDecomposition.
Parameters
----------
weights : str, default="approximate"
Weight calculation method:
- "approximate": Fast simplified formula (default)
- "exact": Variance-based weights from Goodman-Bacon (2021)
"""
if weights not in ("approximate", "exact"):
raise ValueError(f"weights must be 'approximate' or 'exact', got '{weights}'")
self.weights = weights
self.results_: Optional[BaconDecompositionResults] = None
self.is_fitted_: bool = False
def fit(
self,
data: pd.DataFrame,
outcome: str,
unit: str,
time: str,
first_treat: str,
survey_design=None,
) -> BaconDecompositionResults:
"""
Perform the Goodman-Bacon decomposition.
Parameters
----------
data : pd.DataFrame
Panel data with unit and time identifiers.
outcome : str
Name of outcome variable column.
unit : str
Name of unit identifier column.
time : str
Name of time period column.
first_treat : str
Name of column indicating when unit was first treated.
Use 0 (or np.inf) for never-treated units.
survey_design : SurveyDesign, optional
Survey design specification for weighted estimation.
When provided, all means and group shares use survey weights.
The decomposition remains diagnostic (no survey vcov needed).
Returns
-------
BaconDecompositionResults
Object containing decomposition results.
Raises
------
ValueError
If required columns are missing or data validation fails.
"""
# Validate inputs
required_cols = [outcome, unit, time, first_treat]
missing = [c for c in required_cols if c not in data.columns]
if missing:
raise ValueError(f"Missing columns: {missing}")
# Resolve survey design if provided
from diff_diff.survey import _resolve_survey_for_fit
resolved_survey, survey_weights, survey_weight_type, survey_metadata = (
_resolve_survey_for_fit(survey_design, data, "analytical")
)
# Reject replicate-weight designs Bacon decomposition is a
# diagnostic that does not compute replicate-based variance
if resolved_survey is not None and resolved_survey.uses_replicate_variance:
raise NotImplementedError(
"BaconDecomposition does not support replicate-weight survey "
"designs. Use a TSL-based survey design (strata/psu/fpc)."
)
# Validate within-unit constancy for exact survey weights only.
# The exact-weight path collapses to per-unit weights via groupby().first(),
# which requires constant survey columns within units. The approximate path
# uses observation-level weighted means and does not need this constraint.
if resolved_survey is not None and self.weights == "exact":
from diff_diff.survey import _validate_unit_constant_survey
_validate_unit_constant_survey(data, unit, survey_design)
# Create working copy
df = data.copy()
# Ensure numeric types
df[time] = pd.to_numeric(df[time])
df[first_treat] = pd.to_numeric(df[first_treat])
# Check for balanced panel
periods_per_unit = df.groupby(unit)[time].count()
if periods_per_unit.nunique() > 1:
warnings.warn(
"Unbalanced panel detected. Bacon decomposition assumes "
"balanced panels. Results may be inaccurate.",
UserWarning,
stacklevel=2,
)
# Get unique time periods and timing groups
time_periods = sorted(df[time].unique())
# Identify never-treated and timing groups
# Never-treated: first_treat = 0 or inf
never_treated_mask = (df[first_treat] == 0) | (df[first_treat] == np.inf)
timing_groups = sorted([g for g in df[first_treat].unique() if g > 0 and g != np.inf])
# Get unit-level treatment timing
unit_info = df.groupby(unit).agg({first_treat: "first"}).reset_index()
n_never_treated = ((unit_info[first_treat] == 0) | (unit_info[first_treat] == np.inf)).sum()
# Create treatment indicator (D_it = 1 if treated at time t)
# Use unique internal name to avoid conflicts with user data
_TREAT_COL = "__bacon_treated_internal__"
df[_TREAT_COL] = (~never_treated_mask) & (df[time] >= df[first_treat])
# First, compute TWFE estimate for reference
twfe_estimate = self._compute_twfe(
df, outcome, unit, time, _TREAT_COL, weights=survey_weights
)
# Perform decomposition
comparisons = []
# 1. Treated vs Never-treated comparisons
if n_never_treated > 0:
for g in timing_groups:
comp = self._compute_treated_vs_never(
df,
outcome,
unit,
time,
first_treat,
g,
time_periods,
weights=survey_weights,
)
if comp is not None:
comparisons.append(comp)
# 2. Timing group comparisons (earlier vs later and later vs earlier)
for i, g_early in enumerate(timing_groups):
for g_late in timing_groups[i + 1 :]:
# Earlier vs Later: g_early treated, g_late as control
comp_early = self._compute_timing_comparison(
df,
outcome,
unit,
time,
first_treat,
g_early,
g_late,
time_periods,
"earlier_vs_later",
weights=survey_weights,
)
if comp_early is not None:
comparisons.append(comp_early)
# Later vs Earlier: g_late treated, g_early as control (forbidden)
comp_late = self._compute_timing_comparison(
df,
outcome,
unit,
time,
first_treat,
g_late,
g_early,
time_periods,
"later_vs_earlier",
weights=survey_weights,
)
if comp_late is not None:
comparisons.append(comp_late)
# Recompute exact weights if requested
if self.weights == "exact":
self._recompute_exact_weights(
comparisons,
df,
outcome,
unit,
time,
first_treat,
time_periods,
weights=survey_weights,
)
if not comparisons:
raise ValueError(
"No valid 2x2 comparisons remain after filtering. "
"All cells have zero effective weight or insufficient data. "
"Check subpopulation/domain definition."
)
# Normalize weights to sum to 1
total_weight = sum(c.weight for c in comparisons)
if total_weight > 0:
for c in comparisons:
c.weight = c.weight / total_weight
# Calculate weight totals and weighted averages by type
weight_by_type = {"treated_vs_never": 0.0, "earlier_vs_later": 0.0, "later_vs_earlier": 0.0}
weighted_sum_by_type = {
"treated_vs_never": 0.0,
"earlier_vs_later": 0.0,
"later_vs_earlier": 0.0,
}
for c in comparisons:
weight_by_type[c.comparison_type] += c.weight
weighted_sum_by_type[c.comparison_type] += c.weight * c.estimate
# Calculate weighted averages
avg_by_type = {}
for ctype in weight_by_type:
if weight_by_type[ctype] > 0:
avg_by_type[ctype] = weighted_sum_by_type[ctype] / weight_by_type[ctype]
else:
avg_by_type[ctype] = None
# Calculate decomposition error
weighted_sum = sum(c.weight * c.estimate for c in comparisons)
decomp_error = abs(twfe_estimate - weighted_sum)
self.results_ = BaconDecompositionResults(
twfe_estimate=twfe_estimate,
comparisons=comparisons,
total_weight_treated_vs_never=weight_by_type["treated_vs_never"],
total_weight_earlier_vs_later=weight_by_type["earlier_vs_later"],
total_weight_later_vs_earlier=weight_by_type["later_vs_earlier"],
weighted_avg_treated_vs_never=avg_by_type["treated_vs_never"],
weighted_avg_earlier_vs_later=avg_by_type["earlier_vs_later"],
weighted_avg_later_vs_earlier=avg_by_type["later_vs_earlier"],
n_timing_groups=len(timing_groups),
n_never_treated=n_never_treated,
timing_groups=timing_groups,
n_obs=len(df),
decomposition_error=decomp_error,
survey_metadata=survey_metadata,
)
self.is_fitted_ = True
return self.results_
def _compute_twfe(
self,
df: pd.DataFrame,
outcome: str,
unit: str,
time: str,
treat_col: str = "__bacon_treated_internal__",
weights: Optional[np.ndarray] = None,
) -> float:
"""Compute TWFE estimate using within-transformation."""
# Apply two-way within transformation (weighted if survey weights provided)
df_dm = _within_transform_util(
df,
[outcome, treat_col],
unit,
time,
suffix="_within",
weights=weights,
)
# Extract within-transformed values
y_within = df_dm[f"{outcome}_within"].values
d_within = df_dm[f"{treat_col}_within"].values
# OLS on demeaned data: beta = sum(w * d * y) / sum(w * d^2)
w = weights if weights is not None else np.ones(len(y_within))
d_var = np.sum(w * d_within**2)
if d_var > 0:
beta = np.sum(w * d_within * y_within) / d_var
else:
beta = 0.0
return beta
def _recompute_exact_weights(
self,
comparisons: List[Comparison2x2],
df: pd.DataFrame,
outcome: str,
unit: str,
time: str,
first_treat: str,
time_periods: List[Any],
weights: Optional[np.ndarray] = None,
) -> None:
"""
Recompute weights using exact variance-based formula from Theorem 1.
This modifies comparison weights in-place to use the exact formula
from Goodman-Bacon (2021) which accounts for within-group variance
of the treatment indicator in each 2x2 comparison window.
When survey weights are provided, uses weighted unit counts and
within-group variance of the treatment indicator.
"""
n_total_obs = len(df)
w_arr = weights if weights is not None else np.ones(n_total_obs)
# Store weights as a column for safe label-based subsetting
df = df.copy()
df["_sw"] = w_arr
w_total = np.sum(w_arr)
n_total_units = df[unit].nunique()
for comp in comparisons:
# Get data for this specific comparison
if comp.comparison_type == "treated_vs_never":
pre_periods = [t for t in time_periods if t < comp.treated_group]
post_periods = [t for t in time_periods if t >= comp.treated_group]
# Get units in each group
units_treated = df[df[first_treat] == comp.treated_group][unit].unique()
units_control = df[(df[first_treat] == 0) | (df[first_treat] == np.inf)][
unit
].unique()
elif comp.comparison_type == "earlier_vs_later":
g_early = comp.treated_group
g_late = comp.control_group
pre_periods = [t for t in time_periods if t < g_early]
post_periods = [t for t in time_periods if g_early 0 else 0.0
# Create treatment indicator for the 2x2
T_pre = len(pre_periods)
T_post = len(post_periods)
T_window = T_pre + T_post
# Variance of D within the 2x2 for treated group
# D = 0 in pre, D = 1 in post for treated units
# D = 0 for all periods for control units in this window
D_k = T_post / T_window # proportion treated for treated group
# Within-comparison variance of treatment (weighted)
# Var(D) = n_k/(n_k+n_l) * D_k * (1-D_k) for the 2x2
var_D_22 = n_k_share * D_k * (1 - D_k)
# Exact weight: proportional to sample share * variance
# Scale by weighted unit share to account for subsample
# Use survey-weighted unit mass when weights present
if weights is not None:
# Sum of per-unit weights for treated + control units in this 2x2
unit_w_k = (
df_22.loc[treated_mask_22, "_sw"]
.groupby(df_22.loc[treated_mask_22, unit])
.first()
.sum()
)
unit_w_l = (
df_22.loc[~treated_mask_22, "_sw"]
.groupby(df_22.loc[~treated_mask_22, unit])
.first()
.sum()
)
unit_share = (unit_w_k + unit_w_l) / w_total
else:
unit_share = (n_k + n_l) / n_total_units
comp.weight = sample_share * var_D_22 * unit_share
def _compute_treated_vs_never(
self,
df: pd.DataFrame,
outcome: str,
unit: str,
time: str,
first_treat: str,
treated_group: Any,
time_periods: List[Any],
weights: Optional[np.ndarray] = None,
) -> Optional[Comparison2x2]:
"""
Compute 2x2 DiD comparing treated group to never-treated.
This is a "clean" comparison using the full sample of a treated
cohort versus never-treated units.
"""
# Get treated and never-treated units
never_mask = (df[first_treat] == 0) | (df[first_treat] == np.inf)
treated_mask = df[first_treat] == treated_group
df_treated = df[treated_mask]
df_never = df[never_mask]
if len(df_treated) == 0 or len(df_never) == 0:
return None
# Time window: all periods
t_min = min(time_periods)
t_max = max(time_periods)
# Pre and post periods for this group
pre_periods = [t for t in time_periods if t < treated_group]
post_periods = [t for t in time_periods if t >= treated_group]
if not pre_periods or not post_periods:
return None
# Compute 2x2 DiD estimate using weighted means if survey weights provided
w = weights if weights is not None else np.ones(len(df))
y = df[outcome].values
treated_pre_mask = treated_mask & df[time].isin(pre_periods)
treated_post_mask = treated_mask & df[time].isin(post_periods)
never_pre_mask = never_mask & df[time].isin(pre_periods)
never_post_mask = never_mask & df[time].isin(post_periods)
# Guard against empty cells (unbalanced/filtered panels)
# Also check positive weight mass for survey/subpopulation designs
if not (
np.any(treated_pre_mask)
and np.any(treated_post_mask)
and np.any(never_pre_mask)
and np.any(never_post_mask)
):
return None
if (
np.sum(w[treated_pre_mask])