Add Event Study (aka Dynamic Difference in Differences) functionality - #584
Add Event Study (aka Dynamic Difference in Differences) functionality#584drbenvincent wants to merge 42 commits into
Conversation
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #584 +/- ##
==========================================
+ Coverage 95.98% 96.09% +0.11%
==========================================
Files 104 107 +3
Lines 16277 17353 +1076
Branches 912 1010 +98
==========================================
+ Hits 15623 16676 +1053
- Misses 488 492 +4
- Partials 166 185 +19 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
The EventStudy class now requires a patsy-style formula to specify the outcome and fixed effects, removing the separate outcome_col argument. Design matrix construction uses patsy, and event-time dummies are appended. Input validation checks for formula presence, and tests and documentation are updated to reflect the new API and output format.
|
@cursor review |
Expanded documentation to explain the patsy formula syntax, the role of unit and time fixed effects, and how event-time dummies ($\beta_k$) are automatically constructed by the EventStudy class. Added details on the event window and reference event time parameters for clearer guidance.
Added a warning in the EventStudy class and documentation that the implementation only supports simultaneous treatment timing and does not support staggered adoption. Introduced a validation to raise a DataException if treated units have different treatment times. Added a corresponding test to ensure staggered adoption raises an error, and updated the notebook to clarify estimator limitations.
Enhanced the _bayesian_plot and _ols_plot methods in EventStudy to support configurable figure size and HDI probability. Updated docstrings to document new parameters and improved plot labeling for clarity.
Introduces event study support to effect_summary(), including parallel trends check and dynamic effect reporting. Updates event study class to allow HDI probability customization and reporting, and extends documentation with effect summary usage and interpretation.
The `generate_event_study_data` function now supports optional time-varying predictors generated as AR(1) processes, controlled by new parameters: `predictor_effects`, `ar_phi`, and `ar_scale`. Added the `generate_ar1_series` utility function. Updated docstrings and examples to reflect these changes. The event study PyMC notebook was updated with additional analysis and improved section headings.
Introduces integration tests for the EventStudy.effect_summary method using both PyMC and sklearn models. Tests verify the returned EffectSummary object, its table and text attributes, and key output elements.
juanitorduz
left a comment
juanitorduz
left a comment
There was a problem hiding this comment.
Very cool! This is super helpful in marketing (the effect of a marketing campaign depends on time!). We should try to find other real datasets to showcase this.
Minor comments.
Suggestion -> add more tests for the indidual methods and address @BugBot comments
Control units marked with np.inf or np.nan in the treat_time column are now correctly excluded from treated unit checks in EventStudy. Adds tests to ensure control units with np.inf, np.nan, or mixed markers are handled properly, preventing false staggered adoption errors.
Improved rounding logic in EventStudy.get_event_time_summary for PyMC and sklearn models, ensuring consistent application of the round_to parameter. Refactored and expanded test_event_study.py to cover more edge cases, input validation, and integration scenarios, including new tests for rounding, event window handling, and control units. Cleaned up and reorganized tests for clarity and maintainability.
Updates all experiment classes to filter out rows with NaN values after patsy design matrix construction, ensuring consistent shapes between data and design matrices. Adds comprehensive tests for NaN handling across all experiment classes to verify correct filtering and prevent shape mismatch errors.
Enhanced the EventStudy class to emit warnings for unbalanced panels, gaps in time periods, and excessive data loss due to NaN filtering. Added comprehensive unit tests to cover these edge cases and verify warning behavior. Updated interrogate badge to reflect increased coverage.
|
Done, along with some other citations View entire conversation on ReviewNB |
Replaced direct 'pip' calls with 'python -m pip' in CONTRIBUTING.md to ensure the correct pip is used within the conda environment. Added a note explaining the reason for this change.
| ... ) | ||
| >>> result = cp.EventStudy( | ||
| ... df, | ||
| ... formula="y ~ C(unit) + C(time)", |
There was a problem hiding this comment.
I think you should be able to encode the formula fully in patsy. Something like
"y ~ C(unit) + C(rel_time, Treatment(reference=-1))"
It just requires that you pre-processe the df['time'] column before running patsy.
Something like
df = df.sort_values(["unit", "time"])
df["t_idx"] = df.groupby("unit").cumcount()
df["event_idx"] = (
df[df["treated"] == 1]
.groupby("unit")["t_idx"]
.transform("min")
)
df["rel_time"] = df["t_idx"] - df["event_idx"]
There was a problem hiding this comment.
Sorry for the AI evaluation, but some major concerns with that proposal
The Suggestion
Use patsy's categorical handling to simplify the event-time dummy creation:
"y ~ C(unit) + C(rel_time, Treatment(reference=-1))"With pre-processing to compute rel_time before calling patsy.
Benefits
- Simpler code (removes ~50 lines of manual dummy creation)
- More idiomatic use of patsy
- Reference category handled automatically via
Treatment(reference=-1)
Concerns
1. Control Units Would Be Dropped
This is the most significant issue. Control units have NaN for treatment time, so rel_time would be NaN. Patsy drops rows with NaN values entirely.
Current behavior: Control units are kept in the estimation. Their event-time dummies are all 0, but they contribute to the time fixed effects estimation.
Proposed behavior: Control units would be excluded, fundamentally changing the model and breaking the parallel trends comparison.
2. Loss of Calendar Time Fixed Effects
The proposed formula replaces calendar time FEs with event-time dummies. But a proper TWFE event study needs both:
Y_it = α_i + λ_t + Σ β_k · 1{E_it = k} + ε_it
↑ ↑ ↑
unit time event-time
With only C(rel_time), we lose λ_t (calendar time effects). This changes the identifying assumptions of the model.
3. Event Window Handling
Currently, only event times within the window (e.g., -5 to +5) get dummies. Observations outside the window are kept but don't contribute to event-time indicators.
With patsy handling everything, we'd need to either:
- Filter out observations with
rel_timeoutside the window (loses data) - Let patsy create dummies for ALL unique
rel_timevalues (different model, many more parameters) - Bin/clip values outside the window to special categories
4. Minor Issues
- Coefficient naming: Changes from
event_time_0toC(rel_time, Treatment(reference=-1))[T.0] - Dynamic reference: Would need to build formula strings dynamically to inject
reference_event_time - Test updates: Existing tests check current coefficient naming
Possible Middle Ground
Keep the current model structure (which correctly handles control units and maintains both time FEs and event-time dummies) but simplify the internal implementation where possible.
Questions
- Am I missing something that would address the control unit issue?
- Is the loss of calendar time FEs intentional in your suggestion, or should both be included?
- Any thoughts on handling the event window bounds?
Happy to implement if we can resolve these concerns!
There was a problem hiding this comment.
Initial Concerns
The naive approach has some issues:
- Control units dropped: NaN
rel_timecauses patsy to drop entire rows - Loss of calendar time FEs: Replacing
C(time)withC(rel_time)loses λ_t - Event window handling: How to handle observations outside the window?
The Elegant Solution
Key insight: In patsy's Treatment coding, the reference category gets 0 for all dummies. This is exactly what we want for:
- The reference period (k=-1)
- Control units
- Observations outside the event window
Solution: Map all these cases to the reference category!
# Step 1: Compute relative time as usual
df["rel_time"] = df["time"] - df["treat_time"]
# Step 2: Create a categorical version where:
# - Control units → reference_event_time
# - Outside event window → reference_event_time
# - Inside window → actual rel_time
df["rel_time_cat"] = df["rel_time"].copy()
# Control units (NaN rel_time) get the reference category
df.loc[df["rel_time"].isna(), "rel_time_cat"] = reference_event_time
# Observations outside window also get reference category
outside_window = (df["rel_time"] < event_window[0]) | (df["rel_time"] > event_window[1])
df.loc[outside_window & df["rel_time"].notna(), "rel_time_cat"] = reference_event_time
# Step 3: Use patsy with BOTH time FEs and event-time dummies
formula = f"y ~ C(unit) + C(time) + C(rel_time_cat, Treatment(reference={reference_event_time}))"Why This Works
| Observation Type | rel_time_cat Value |
Result |
|---|---|---|
| Treated, k=-1 (reference) | -1 | 0 for all event dummies ✓ |
| Treated, k in window | actual k | Gets its own dummy ✓ |
| Treated, k outside window | -1 (reference) | 0 for all event dummies ✓ |
| Control unit | -1 (reference) | 0 for all event dummies ✓ |
Benefits
- Calendar time FEs preserved -
C(time)stays in the formula - Control units handled correctly - They contribute to FE estimation, get 0 for event dummies
- Event window handled naturally - Outside-window observations map to reference
- Much simpler code - Patsy creates all the dummies
- No rows dropped - No NaN values in the categorical column
Implementation Approach
The user could still provide a simple formula, and the class would augment it internally:
# User provides:
formula = "y ~ C(unit) + C(time)"
# Class internally transforms to:
augmented_formula = f"y ~ C(unit) + C(time) + C(rel_time_cat, Treatment(reference={ref}))"This significantly simplifies _build_design_matrix() while preserving all the correct model behavior.
There was a problem hiding this comment.
What do you think @NathanielF? Maybe we merge as is, create an issue for further work, then come back to it in a second iteration?
Resolve conflicts by combining main's refactored experiment structure (_build_design_matrices, _prepare_data, algorithm methods, default model classes) with event-study's NaN handling (return_type="dataframe" in dmatrices calls and index-based row filtering). Conflicting files: - README.md: keep both new method rows (Event Study + Staggered DiD) - 6 experiment files: take main's refactored methods, add PR's NaN handling - interrogate_badge.svg: accept main's deletion - references.bib: keep all new references from both sides Made-with: Cursor
- Restore _compute_statistics_rkink_ols stub that was dropped during auto-merge of reporting.py - Update pandas frequency string 'M' -> 'ME' in simulate_data.py (deprecated since pandas 2.2) Made-with: Cursor
Made-with: Cursor
Made-with: Cursor
NathanielF
left a comment
NathanielF
left a comment
There was a problem hiding this comment.
Sorry it took so long.
Automated triageRecommendation: Why:
Review focus:
Confidence: high |
drbenvincent
added
the
review:high
Sync with main (gallery.yaml, build_formula_matrices, formula NaN filtering), add EventStudy to gallery as event-study-pymc, implement unified plot() API, and set LinearRegression as the default model class. Co-authored-by: Cursor <[email protected]>
Align test_synthetic_data with main's rng-based internal helpers and remove the obsolete generate_time_series_data test. Co-authored-by: Cursor <[email protected]>
|
@juanitorduz — when you have a moment, could you take another look at this PR on the current head ( Your Dec review asked for more per-method tests and bugbot fixes. Since then we've:
CI is green on the latest push. Deferred items from the thread (real marketing datasets, demeaning FE for large N, patsy If you're happy with the current state, an approval would unblock merge. If anything still feels like a major sticking point, let us know and we'll address it. |
derwells
added
the
needs:maintainer-decision
5 participants
Closes #384
This PR adds Event Study / Dynamic DiD for simultaneous treatment timing (TWFE + event-time dummies). The staggered-adoption half of #384 landed on
mainin #621 (StaggeredDifferenceInDifferences); this PR is the complementary simultaneous-timing estimator and explicitly rejects staggered adoption.Summary
EventStudyexperiment class (PyMC + sklearn), witheffect_summary(), unifiedplot()API, andgenerate_event_study_data()simulatorevent-study-pymc.ipynb, gallery entry, glossary, reporting statistics tablebuild_formula_matrices(..., return_type="dataframe")+ row filtering across several experiment classes (bugbot finding)main(gallery.yaml, formula utils, PyMC 6 notebook conventions, etc.)PR checklist
Consider adding the time relative to treatment column in the notebook, not hidden in the experiment class.Since Dec review
mainand resolved conflictstest_event_study.py,test_nan_handling.py)np.infcontrol units,round_toinget_event_time_summary, NaN row/index alignmentDeferred (follow-up issues, not blocking this PR)
Treatment(reference=...)refactor for event-time dummies📚 Documentation preview 📚: https://causalpy--584.org.readthedocs.build/en/584/