wooldridge: CI R2 P1 fix — plot_event_study always re-aggregates + cohort_share→cell regression

Author: igerber

diff_diff/wooldridge_results.py

@@ -694,7 +694,7 @@ def to_dataframe(self, aggregation: str = "event") -> pd.DataFrame:
         return pd.DataFrame(rows)
 
     def plot_event_study(self, weights: str = "cell", **kwargs) -> None:
-        """Event study plot. Calls ``aggregate('event', weights=weights)`` if needed.
+        """Event study plot. Always calls ``aggregate('event', weights=weights)``.
 
         Parameters
         ----------
@@ -707,19 +707,24 @@ def plot_event_study(self, weights: str = "cell", **kwargs) -> None:
             contract documented in REGISTRY.
         **kwargs
             Forwarded to ``diff_diff.visualization.plot_event_study``.
+
+        Notes
+        -----
+        The wrapper unconditionally re-aggregates the event study under
+        the requested ``weights`` scheme. This avoids the stale-cache
+        hazard where a prior ``plot_event_study(weights="cohort_share")``
+        call would leave the cached ``event_study_effects`` restricted
+        to ``k >= 0`` (per the Eq. 7.6 scope), and a subsequent
+        ``plot_event_study()`` (default ``weights="cell"``) call would
+        silently reuse the cohort-share-keyed cache instead of restoring
+        the full event range including pre-period placebo leads.
         """
-        # Recompute under the active weighting scheme if the cached
-        # event_study_effects was built under a different scheme — or
-        # has not been built yet. Aggregating under "cell" then under
-        # "cohort_share" (or vice versa) replaces ``event_study_effects``
-        # in place per the existing aggregate() contract.
-        if self.event_study_effects is None:
-            self.aggregate("event", weights=weights)
-        elif weights == "cohort_share":
-            # Force re-aggregation so the cohort-share contract is
-            # honored from a wrapper call that may have been preceded
-            # by an aggregate("event", weights="cell") at fit time.
-            self.aggregate("event", weights=weights)
+        # Always re-aggregate under the requested weighting scheme. The
+        # aggregate() method replaces ``event_study_effects`` in place
+        # per the existing contract, so this is cheap and avoids
+        # cohort_share→cell (or any cross-scheme) stale-cache bugs.
+        self.aggregate("event", weights=weights)
+
         from diff_diff.visualization import plot_event_study  # type: ignore
 
         effects = {k: v["att"] for k, v in (self.event_study_effects or {}).items()}

tests/test_methodology_wooldridge.py

@@ -1469,6 +1469,57 @@ def test_plot_event_study_propagates_weights_kwarg(self) -> None:
             "leads."
         )
 
+    def test_plot_event_study_cohort_share_to_cell_round_trip_restores_placebo_leads(
+        self,
+    ) -> None:
+        """CI R2 P1 fix: ``plot_event_study()`` reverse direction (cohort_share → cell).
+
+        Codex caught a stale-cache hazard: my earlier fix re-aggregated
+        on ``weights="cohort_share"`` but skipped re-aggregation on the
+        default ``weights="cell"`` path when the cached ``event_study_effects``
+        was already populated. A user calling ``plot_event_study(weights=
+        "cohort_share")`` (which restricts to k>=0) and then
+        ``plot_event_study()`` (default cell weights) would silently
+        plot the stale cohort-share-keyed data. The fix unconditionally
+        re-aggregates on every call.
+
+        This test exercises the reverse direction by:
+        1. First call: ``plot_event_study(weights="cohort_share")`` —
+           caches cohort-share keys (k >= 0 only)
+        2. Second call: ``plot_event_study()`` (default cell weights) —
+           must restore the full event range including k < 0 placebo leads
+        """
+        from unittest.mock import patch
+
+        rng = np.random.default_rng(_BASE_SEED_SECTION8 + 16)
+        panel = _make_three_cohort_four_period_panel(rng, n_per_cohort=80, sigma=0.05)
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=UserWarning)
+            res = WooldridgeDiD(method="ols", control_group="never_treated").fit(
+                panel, outcome="y", unit="unit", time="time", cohort="cohort"
+            )
+        # Step 1: plot under cohort_share — caches k>=0 keys
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=UserWarning)
+            with patch("diff_diff.visualization.plot_event_study"):
+                res.plot_event_study(weights="cohort_share")
+        assert res.event_study_effects is not None
+        cohort_share_keys = sorted(res.event_study_effects.keys())
+        assert all(k >= 0 for k in cohort_share_keys), (
+            f"DGP precondition: cohort_share path must restrict to k>=0; "
+            f"got {cohort_share_keys}"
+        )
+        # Step 2: plot under default weights="cell" — must restore k<0 leads
+        with patch("diff_diff.visualization.plot_event_study"):
+            res.plot_event_study()  # default weights="cell"
+        assert res.event_study_effects is not None
+        cell_keys = sorted(res.event_study_effects.keys())
+        assert any(k < 0 for k in cell_keys), (
+            f"plot_event_study() (cell weights) after a cohort_share "
+            f"call must restore k<0 placebo leads; got {cell_keys}. "
+            f"Stale cohort_share cache was reused — CI R2 P1 fix regressed."
+        )
+
     def test_cohort_trends_true_plus_bootstrap_preserves_bootstrap_se(self) -> None:
         """R5 P1 fix: ``cohort_trends=True`` + ``n_bootstrap > 0`` runs cleanly.