-
Notifications
You must be signed in to change notification settings - Fork 20
Expand file tree
/
Copy pathtwfe.py
More file actions
542 lines (477 loc) · 21.8 KB
/
twfe.py
File metadata and controls
542 lines (477 loc) · 21.8 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
"""
Two-Way Fixed Effects estimator for panel Difference-in-Differences.
"""
import warnings
from typing import TYPE_CHECKING, List, Optional
import numpy as np
import pandas as pd
if TYPE_CHECKING:
from diff_diff.bacon import BaconDecompositionResults
from diff_diff.estimators import DifferenceInDifferences
from diff_diff.linalg import LinearRegression
from diff_diff.results import DiDResults
from diff_diff.utils import (
within_transform as _within_transform_util,
)
class TwoWayFixedEffects(DifferenceInDifferences):
"""
Two-Way Fixed Effects (TWFE) estimator for panel DiD.
Extends DifferenceInDifferences to handle panel data with unit
and time fixed effects.
Parameters
----------
robust : bool, default=True
Whether to use heteroskedasticity-robust standard errors.
cluster : str, optional
Column name for cluster-robust standard errors.
If None, automatically clusters at the unit level (the `unit`
parameter passed to `fit()`). This differs from
DifferenceInDifferences where cluster=None means no clustering.
alpha : float, default=0.05
Significance level for confidence intervals.
Notes
-----
This estimator uses the regression:
Y_it = α_i + γ_t + β*(D_i × Post_t) + X_it'δ + ε_it
where α_i are unit fixed effects and γ_t are time fixed effects.
Warning: TWFE can be biased with staggered treatment timing
and heterogeneous treatment effects. Consider using
more robust estimators (e.g., Callaway-Sant'Anna) for
staggered designs.
"""
def fit( # type: ignore[override]
self,
data: pd.DataFrame,
outcome: str,
treatment: str,
time: str,
unit: str,
covariates: Optional[List[str]] = None,
survey_design: object = None,
) -> DiDResults:
"""
Fit Two-Way Fixed Effects model.
Parameters
----------
data : pd.DataFrame
Panel data.
outcome : str
Name of outcome variable column.
treatment : str
Name of treatment indicator column.
time : str
Name of time period column.
unit : str
Name of unit identifier column.
covariates : list, optional
List of covariate column names.
survey_design : SurveyDesign, optional
Survey design specification for design-based inference. When provided,
uses Taylor Series Linearization for variance estimation and
applies sampling weights to the regression.
Returns
-------
DiDResults
Estimation results.
"""
# Validate unit column exists
if unit not in data.columns:
raise ValueError(f"Unit column '{unit}' not found in data")
# Check for staggered treatment timing and warn if detected
self._check_staggered_treatment(data, treatment, time, unit)
# Warn if time has more than 2 unique values (not a binary post indicator)
n_unique_time = data[time].nunique()
if n_unique_time > 2:
warnings.warn(
f"The '{time}' column has {n_unique_time} unique values. "
f"TwoWayFixedEffects expects a binary (0/1) post indicator. "
f"Multi-period time values produce 'treated * period_number' instead of "
f"'treated * post_indicator', which may not estimate the standard DiD ATT. "
f"Consider creating a binary post column: "
f"df['post'] = (df['{time}'] >= cutoff).astype(int)",
UserWarning,
stacklevel=2,
)
elif n_unique_time == 2:
unique_vals = set(data[time].unique())
if unique_vals != {0, 1} and unique_vals != {False, True}:
warnings.warn(
f"The '{time}' column has values {sorted(unique_vals)} instead of {{0, 1}}. "
f"The ATT estimate is mathematically correct (within-transformation "
f"absorbs the scaling), but 0/1 encoding is recommended for clarity. "
f"Consider: df['{time}'] = (df['{time}'] == {max(unique_vals)}).astype(int)",
UserWarning,
stacklevel=2,
)
# Resolve survey design if provided
from diff_diff.survey import _resolve_effective_cluster, _resolve_survey_for_fit
resolved_survey, survey_weights, survey_weight_type, survey_metadata = (
_resolve_survey_for_fit(survey_design, data, self.inference)
)
_uses_replicate_twfe = (
resolved_survey is not None and resolved_survey.uses_replicate_variance
)
if _uses_replicate_twfe and self.inference == "wild_bootstrap":
raise ValueError(
"Cannot use inference='wild_bootstrap' with replicate-weight "
"survey designs. Replicate weights provide their own variance "
"estimation."
)
# Use unit-level clustering if not specified (use local variable to avoid mutation)
cluster_var = self.cluster if self.cluster is not None else unit
# Create treatment × post interaction from raw data before demeaning.
# This must be within-transformed alongside the outcome and covariates
# so that the regression uses demeaned regressors (FWL theorem).
data = data.copy()
data["_treatment_post"] = data[treatment] * data[time]
# Demean outcome, covariates, AND interaction in a single pass
all_vars = [outcome] + (covariates or []) + ["_treatment_post"]
data_demeaned = _within_transform_util(
data,
all_vars,
unit,
time,
suffix="_demeaned",
weights=survey_weights,
)
# Extract variables for regression
y = data_demeaned[f"{outcome}_demeaned"].values
X_list = [data_demeaned["_treatment_post_demeaned"].values]
if covariates:
for cov in covariates:
X_list.append(data_demeaned[f"{cov}_demeaned"].values)
X = np.column_stack([np.ones(len(y))] + X_list)
# ATT is the coefficient on treatment_post (index 1)
att_idx = 1
# Degrees of freedom adjustment for fixed effects
n_units = data[unit].nunique()
n_times = data[time].nunique()
df_adjustment = n_units + n_times - 2
# Always use LinearRegression for initial fit (unified code path)
# For wild bootstrap, we don't need cluster SEs from the initial fit
cluster_ids = data[cluster_var].values
# When survey PSU is present, it overrides cluster for variance estimation
effective_cluster_ids = _resolve_effective_cluster(
resolved_survey, cluster_ids, self.cluster
)
# For survey variance: only inject user-explicit cluster as PSU.
# TWFE's default unit clustering should not override the documented
# no-PSU survey path (implicit per-observation PSUs).
if resolved_survey is not None and self.cluster is None:
survey_cluster_ids = None
else:
survey_cluster_ids = effective_cluster_ids
# Inject cluster as effective PSU for survey variance estimation
if resolved_survey is not None and survey_cluster_ids is not None:
from diff_diff.survey import _inject_cluster_as_psu, compute_survey_metadata
resolved_survey = _inject_cluster_as_psu(resolved_survey, survey_cluster_ids)
if resolved_survey.psu is not None and survey_metadata is not None:
raw_w = (
data[survey_design.weights].values.astype(np.float64)
if survey_design.weights
else np.ones(len(data), dtype=np.float64)
)
survey_metadata = compute_survey_metadata(resolved_survey, raw_w)
# Pass rank_deficient_action to LinearRegression
# If "error", let LinearRegression raise immediately
# If "warn" or "silent", suppress generic warning and use TWFE's context-specific
# error/warning messages (more informative for panel data)
# For replicate designs: pass survey_design=None to prevent LinearRegression
# from computing replicate vcov on already-demeaned data (demeaning depends
# on weights, so replicate refits must re-demean at the estimator level).
_lr_survey_twfe = None if _uses_replicate_twfe else resolved_survey
if self.rank_deficient_action == "error":
reg = LinearRegression(
include_intercept=False,
robust=True,
cluster_ids=survey_cluster_ids if self.inference != "wild_bootstrap" else None,
alpha=self.alpha,
rank_deficient_action="error",
weights=survey_weights,
weight_type=survey_weight_type,
survey_design=_lr_survey_twfe,
).fit(X, y, df_adjustment=df_adjustment)
else:
# Suppress generic warning, TWFE provides context-specific messages below
with warnings.catch_warnings():
warnings.filterwarnings("ignore", message="Rank-deficient design matrix")
reg = LinearRegression(
include_intercept=False,
robust=True,
cluster_ids=(
survey_cluster_ids if self.inference != "wild_bootstrap" else None
),
alpha=self.alpha,
rank_deficient_action="silent",
weights=survey_weights,
weight_type=survey_weight_type,
survey_design=_lr_survey_twfe,
).fit(X, y, df_adjustment=df_adjustment)
coefficients = reg.coefficients_
residuals = reg.residuals_
fitted = reg.fitted_values_
r_squared = reg.r_squared()
assert coefficients is not None
att = coefficients[att_idx]
# Check for unidentified coefficients (collinearity)
# Build column names for informative error messages
column_names = ["intercept", "treatment×post"]
if covariates:
column_names.extend(covariates)
nan_mask = np.isnan(coefficients)
if np.any(nan_mask):
dropped_indices = np.where(nan_mask)[0]
dropped_names = [
column_names[i] if i < len(column_names) else f"column {i}" for i in dropped_indices
]
# Determine the source of collinearity for better error message
if att_idx in dropped_indices:
# Treatment coefficient is unidentified
raise ValueError(
f"Treatment effect cannot be identified due to collinearity. "
f"Dropped columns: {', '.join(dropped_names)}. "
"This can happen when: (1) treatment is perfectly collinear with "
"unit/time fixed effects, (2) all treated units are treated in all "
"periods, or (3) a covariate is collinear with the treatment indicator. "
"Check your data structure and model specification."
)
else:
# Only covariates are dropped - this is a warning, not an error
# The ATT can still be estimated
# Respect rank_deficient_action setting for warning
if self.rank_deficient_action == "warn":
warnings.warn(
f"Some covariates are collinear and were dropped: "
f"{', '.join(dropped_names)}. The treatment effect is still identified.",
UserWarning,
stacklevel=2,
)
# Get inference - replicate, bootstrap, or analytical
if _uses_replicate_twfe:
# Estimator-level replicate variance: re-do within-transform per replicate
from diff_diff.linalg import solve_ols
from diff_diff.survey import compute_replicate_refit_variance
from diff_diff.utils import safe_inference as _safe_inf
_all_vars_twfe = list(all_vars)
_covariates_twfe = list(covariates) if covariates else []
# Handle rank-deficient nuisance: refit only identified columns
_id_mask_twfe = ~np.isnan(coefficients)
_id_cols_twfe = np.where(_id_mask_twfe)[0]
def _refit_twfe(w_r):
# Drop zero-weight obs to prevent zero-sum demeaning
# (JK1/BRR half-samples zero entire clusters)
nz = w_r > 0
data_nz = data[nz].copy()
w_nz = w_r[nz]
data_dem_r = _within_transform_util(
data_nz, _all_vars_twfe, unit, time, suffix="_demeaned", weights=w_nz,
)
y_r = data_dem_r[f"{outcome}_demeaned"].values
X_list_r = [data_dem_r["_treatment_post_demeaned"].values]
for cov_ in _covariates_twfe:
X_list_r.append(data_dem_r[f"{cov_}_demeaned"].values)
X_r = np.column_stack([np.ones(len(y_r))] + X_list_r)
coef_r, _, _ = solve_ols(
X_r[:, _id_cols_twfe], y_r,
weights=w_nz, weight_type=survey_weight_type,
rank_deficient_action="silent", return_vcov=False,
)
return coef_r
from diff_diff.linalg import _expand_vcov_with_nan as _expand_twfe
vcov_reduced, _n_valid_rep_twfe = compute_replicate_refit_variance(
_refit_twfe, coefficients[_id_mask_twfe], resolved_survey
)
vcov = _expand_twfe(vcov_reduced, len(coefficients), _id_cols_twfe)
se = float(np.sqrt(max(vcov[att_idx, att_idx], 0.0)))
_df_rep = (
survey_metadata.df_survey
if survey_metadata and survey_metadata.df_survey
else 0 # rank-deficient replicate → NaN inference
)
if _n_valid_rep_twfe < resolved_survey.n_replicates:
_df_rep = _n_valid_rep_twfe - 1 if _n_valid_rep_twfe > 1 else 0
if survey_metadata is not None:
survey_metadata.df_survey = _df_rep if _df_rep > 0 else None
t_stat, p_value, conf_int = _safe_inf(att, se, alpha=self.alpha, df=_df_rep)
elif self.inference == "wild_bootstrap":
# Override with wild cluster bootstrap inference
se, p_value, conf_int, t_stat, vcov, _ = self._run_wild_bootstrap_inference(
X, y, residuals, cluster_ids, att_idx
)
else:
# Use analytical inference from LinearRegression
vcov = reg.vcov_
inference = reg.get_inference(att_idx)
se = inference.se
t_stat = inference.t_stat
p_value = inference.p_value
conf_int = inference.conf_int
# Count observations
treated_units = data[data[treatment] == 1][unit].unique()
n_treated = len(treated_units)
n_control = n_units - n_treated
# Determine inference method and bootstrap info
inference_method = "analytical"
n_bootstrap_used = None
n_clusters_used = None
if self._bootstrap_results is not None:
inference_method = "wild_bootstrap"
n_bootstrap_used = self._bootstrap_results.n_bootstrap
n_clusters_used = self._bootstrap_results.n_clusters
self.results_ = DiDResults(
att=att,
se=se,
t_stat=t_stat,
p_value=p_value,
conf_int=conf_int,
n_obs=len(y),
n_treated=n_treated,
n_control=n_control,
alpha=self.alpha,
coefficients={"ATT": float(att)},
vcov=vcov,
residuals=residuals,
fitted_values=fitted,
r_squared=r_squared,
inference_method=inference_method,
n_bootstrap=n_bootstrap_used,
n_clusters=n_clusters_used,
survey_metadata=survey_metadata,
)
self.is_fitted_ = True
return self.results_
def _within_transform(
self,
data: pd.DataFrame,
outcome: str,
unit: str,
time: str,
covariates: Optional[List[str]] = None,
) -> pd.DataFrame:
"""
Apply within transformation to remove unit and time fixed effects.
This implements the standard two-way within transformation:
y_it - y_i. - y_.t + y_..
Parameters
----------
data : pd.DataFrame
Panel data.
outcome : str
Outcome variable name.
unit : str
Unit identifier column.
time : str
Time period column.
covariates : list, optional
Covariate column names.
Returns
-------
pd.DataFrame
Data with demeaned variables.
"""
variables = [outcome] + (covariates or [])
return _within_transform_util(data, variables, unit, time, suffix="_demeaned")
def _check_staggered_treatment(
self,
data: pd.DataFrame,
treatment: str,
time: str,
unit: str,
) -> None:
"""
Check for staggered treatment timing and warn if detected.
Identifies if different units start treatment at different times,
which can bias TWFE estimates when treatment effects are heterogeneous.
Note: This check requires ``time`` to have actual period values (not
binary 0/1). With binary time, all treated units appear to start at
time=1, so staggering is undetectable.
"""
# Find first treatment time for each unit
treated_obs = data[data[treatment] == 1]
if len(treated_obs) == 0:
return # No treated observations
# Get first treatment time per unit
first_treat_times = treated_obs.groupby(unit)[time].min()
unique_treat_times = first_treat_times.unique()
if len(unique_treat_times) > 1:
n_groups = len(unique_treat_times)
warnings.warn(
f"Staggered treatment timing detected: {n_groups} treatment cohorts "
f"start treatment at different times. TWFE can be biased when treatment "
f"effects are heterogeneous across time. Consider using:\n"
f" - CallawaySantAnna estimator for robust estimates\n"
f" - TwoWayFixedEffects.decompose() to diagnose the decomposition\n"
f" - bacon_decompose() to see weight on 'forbidden' comparisons",
UserWarning,
stacklevel=3,
)
def decompose(
self,
data: pd.DataFrame,
outcome: str,
unit: str,
time: str,
first_treat: str,
weights: str = "approximate",
) -> "BaconDecompositionResults":
"""
Perform Goodman-Bacon decomposition of TWFE estimate.
Decomposes the TWFE estimate into a weighted average of all possible
2x2 DiD comparisons, revealing which comparisons drive the estimate
and whether problematic "forbidden comparisons" are involved.
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 each unit was first treated.
Use 0 (or np.inf) for never-treated units.
weights : str, default="approximate"
Weight calculation method:
- "approximate": Fast simplified formula (default). Good for
diagnostic purposes where relative weights are sufficient.
- "exact": Variance-based weights from Goodman-Bacon (2021)
Theorem 1. Use for publication-quality decompositions.
Returns
-------
BaconDecompositionResults
Decomposition results showing:
- TWFE estimate and its weighted-average breakdown
- List of all 2x2 comparisons with estimates and weights
- Total weight by comparison type (clean vs forbidden)
Examples
--------
>>> twfe = TwoWayFixedEffects()
>>> decomp = twfe.decompose(
... data, outcome='y', unit='id', time='t', first_treat='treat_year'
... )
>>> decomp.print_summary()
>>> # Check weight on forbidden comparisons
>>> if decomp.total_weight_later_vs_earlier > 0.2:
... print("Warning: significant forbidden comparison weight")
Notes
-----
This decomposition is essential for understanding potential TWFE bias
in staggered adoption designs. The three comparison types are:
1. **Treated vs Never-treated**: Clean comparisons using never-treated
units as controls. These are always valid.
2. **Earlier vs Later treated**: Uses later-treated units as controls
before they receive treatment. These are valid.
3. **Later vs Earlier treated**: Uses already-treated units as controls.
These "forbidden comparisons" can introduce bias when treatment
effects are dynamic (changing over time since treatment).
See Also
--------
bacon_decompose : Standalone decomposition function
BaconDecomposition : Class-based decomposition interface
CallawaySantAnna : Robust estimator that avoids forbidden comparisons
"""
from diff_diff.bacon import BaconDecomposition
decomp = BaconDecomposition(weights=weights)
return decomp.fit(data, outcome, unit, time, first_treat)