wooldridge: CI R6 P1 fixes — plot cohort_share suppresses CI bands + Results metadata for cohort_trends/aggregation weights

Author: igerber

diff_diff/wooldridge.py

@@ -1375,6 +1375,7 @@ def _fit_ols(
             _bm_artifacts=bm_artifacts,
             _df_one_way=df_one_way,
             cohort_trend_coefs=cohort_trend_coefs,
+            cohort_trends=self.cohort_trends,
         )
 
         # 9. Optional multiplier bootstrap (overrides analytic SE for overall ATT).

diff_diff/wooldridge_results.py

@@ -100,6 +100,23 @@ class WooldridgeDiDResults:
     # or wait for the deferred bootstrap-cohort-share follow-up.
     _bootstrap_used: bool = field(default=False, repr=False)
 
+    # Model-surface metadata for self-describing reporting.
+    # ``cohort_trends`` records whether the fit was produced under the
+    # Section 8 / Eq. 8.1 heterogeneous-cohort-trends design (paper
+    # W2025 ``dg_i · t`` interactions on the OLS path). False on the
+    # default ``cohort_trends=False`` fit and on logit/Poisson paths
+    # (which reject ``cohort_trends=True`` at the constructor).
+    # ``active_aggregation_weights`` records the most recent
+    # ``aggregate(weights=...)`` scheme that wrote to the ``overall_*``
+    # / ``event_study_effects`` / ``group_effects`` / ``calendar_effects``
+    # fields ("cell" by default; flips to "cohort_share" after an opt-in
+    # cohort-share aggregation). Surfaced in ``summary()`` so downstream
+    # consumers can tell which estimand the printed rows represent
+    # without inspecting ``cohort_trend_coefs`` (which is legitimately
+    # empty on all-treated panels per the last-cohort drop rule).
+    cohort_trends: bool = field(default=False, repr=False)
+    active_aggregation_weights: str = field(default="cell", repr=False)
+
     # ------------------------------------------------------------------ #
     # Internal — used by aggregate() for delta-method SEs                 #
     # ------------------------------------------------------------------ #
@@ -195,6 +212,14 @@ def aggregate(self, type: str, weights: str = "cell") -> "WooldridgeDiDResults":
                 f"jwdid_estat-style cell-count weighting."
             )
 
+        # Record the active aggregation weighting scheme on the Results
+        # object so downstream reporting (``summary()`` / ``to_dataframe()``
+        # / ``__repr__``) can label the persisted ``overall_*`` /
+        # ``event_study_effects`` / ``group_effects`` / ``calendar_effects``
+        # fields with the right estimand. Stamped AFTER the raise-paths
+        # above so invalid combinations don't pollute the metadata.
+        self.active_aggregation_weights = weights
+
         gt = self.group_time_effects
         cell_weights = self._gt_weights
         n_g_per_cohort = self._n_g_per_cohort
@@ -556,6 +581,8 @@ def summary(self, aggregation: str = "simple") -> str:
             f"Observations:    {self.n_obs}",
             f"Treated units:   {self.n_treated_units}",
             f"Control units:   {self.n_control_units}",
+            f"Cohort trends:   {self.cohort_trends}",
+            f"Aggregation w:   {self.active_aggregation_weights}",
             "-" * 70,
         ]
 
@@ -714,7 +741,10 @@ def plot_event_study(self, weights: str = "cell", **kwargs) -> None:
             paper W2025 Eq. 7.6 cohort-share-by-exposure weights
             (post-treatment ``k >= 0`` only); inference fields are
             fail-closed to NaN per the Section 7.5 conditional-on-shares
-            contract documented in REGISTRY.
+            contract documented in REGISTRY, and the plot **suppresses
+            error bars / CI bands** to honor the fail-closed contract
+            (the conditional-on-shares SE would build a misleading
+            normal-theory CI in the plotter).
         **kwargs
             Forwarded to ``diff_diff.visualization.plot_event_study``.
 
@@ -738,7 +768,17 @@ def plot_event_study(self, weights: str = "cell", **kwargs) -> None:
         from diff_diff.visualization import plot_event_study  # type: ignore
 
         effects = {k: v["att"] for k, v in (self.event_study_effects or {}).items()}
-        se = {k: v["se"] for k, v in (self.event_study_effects or {}).items()}
+        if weights == "cohort_share":
+            # Honor the fail-closed inference contract per paper Section
+            # 7.5: the conditional-on-shares SE understates unconditional
+            # uncertainty, so passing the finite SE into the plotter
+            # would let it render a normal-theory CI that contradicts
+            # the NaN inference fields the aggregate() helper produces.
+            # Pass NaN SEs so the plotter suppresses error bars / CI
+            # bands. Locked by ``test_plot_event_study_cohort_share_suppresses_error_bars``.
+            se = {k: float("nan") for k in (self.event_study_effects or {})}
+        else:
+            se = {k: v["se"] for k, v in (self.event_study_effects or {}).items()}
         plot_event_study(effects=effects, se=se, alpha=self.alpha, **kwargs)
 
     # --- Inference-field aliases (balance/external-adapter compatibility) ---

tests/test_methodology_wooldridge.py

@@ -1546,6 +1546,86 @@ def test_cohort_trends_true_plus_aggregate_calendar(self) -> None:
                 finite_count += 1
         assert finite_count >= 1
 
+    def test_plot_event_study_cohort_share_suppresses_error_bars(self) -> None:
+        """CI R6 P1 fix: ``plot_event_study(weights="cohort_share")`` passes NaN SEs.
+
+        Honors the Section 7.5 fail-closed inference contract: the
+        conditional-on-shares SE that ``aggregate()`` returns is NOT a
+        valid input for a normal-theory CI band, so the plot helper
+        receives NaN SEs and therefore suppresses error bars. Locked
+        by inspecting the ``se`` kwarg the helper was called with.
+        """
+        from unittest.mock import patch
+
+        rng = np.random.default_rng(_BASE_SEED_SECTION8 + 19)
+        panel = _make_three_cohort_four_period_panel(rng, n_per_cohort=80, sigma=0.05)
+        res = WooldridgeDiD(method="ols").fit(
+            panel, outcome="y", unit="unit", time="time", cohort="cohort"
+        )
+        # cohort_share path: plot helper must receive NaN SEs
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=UserWarning)
+            with patch("diff_diff.visualization.plot_event_study") as mock_plot:
+                res.plot_event_study(weights="cohort_share")
+        assert mock_plot.call_count == 1
+        se_arg = mock_plot.call_args.kwargs["se"]
+        assert se_arg, "plot_event_study should be called with a non-empty se dict"
+        assert all(np.isnan(v) for v in se_arg.values()), (
+            f"weights='cohort_share' must pass NaN SEs to the plot helper "
+            f"to suppress error bars; got {se_arg}"
+        )
+        # cell path: plot helper must receive FINITE SEs (control)
+        with patch("diff_diff.visualization.plot_event_study") as mock_plot:
+            res.plot_event_study(weights="cell")
+        assert mock_plot.call_count == 1
+        se_arg_cell = mock_plot.call_args.kwargs["se"]
+        assert se_arg_cell, "cell path should also pass a non-empty se dict"
+        assert any(np.isfinite(v) and v > 0 for v in se_arg_cell.values()), (
+            f"weights='cell' must pass finite SEs to the plot helper for "
+            f"normal-theory CI bands; got {se_arg_cell}"
+        )
+
+    def test_results_metadata_records_cohort_trends_and_active_aggregation_weights(
+        self,
+    ) -> None:
+        """CI R6 P1 fix: Results object surfaces cohort_trends + active_aggregation_weights.
+
+        Downstream consumers need to know which model produced
+        ``overall_*`` / ``event_study_effects`` / ``group_effects`` /
+        ``calendar_effects`` (cohort_trends on/off) and which weighting
+        scheme is currently cached (``"cell"`` vs ``"cohort_share"``).
+        ``cohort_trend_coefs`` is NOT a reliable proxy for cohort_trends
+        because the all-treated drop rule legitimately empties the dict
+        on G=1-after-drop panels.
+        """
+        rng = np.random.default_rng(_BASE_SEED_SECTION8 + 20)
+        panel = _make_heterogeneous_trends_panel(rng, n_per_cohort=80, sigma=0.05)
+        # Default fit: cohort_trends=False, default aggregation weights "cell"
+        res_default = WooldridgeDiD(method="ols").fit(
+            panel, outcome="y", unit="unit", time="time", cohort="cohort"
+        )
+        assert res_default.cohort_trends is False
+        assert res_default.active_aggregation_weights == "cell"
+        # cohort_trends=True fit: metadata reflects it
+        res_trends = WooldridgeDiD(method="ols", cohort_trends=True).fit(
+            panel, outcome="y", unit="unit", time="time", cohort="cohort"
+        )
+        assert res_trends.cohort_trends is True
+        assert res_trends.active_aggregation_weights == "cell"
+        # Opt-in cohort_share aggregation flips active_aggregation_weights
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=UserWarning)
+            res_trends.aggregate("simple", weights="cohort_share")
+        assert res_trends.active_aggregation_weights == "cohort_share"
+        # Surface in summary()
+        summary_text = res_trends.summary("simple")
+        assert (
+            "Cohort trends:   True" in summary_text
+        ), f"summary() should surface cohort_trends; got:\n{summary_text}"
+        assert "Aggregation w:   cohort_share" in summary_text, (
+            f"summary() should surface active_aggregation_weights; got:\n" f"{summary_text}"
+        )
+
     def test_plot_event_study_propagates_weights_kwarg(self) -> None:
         """CI R1 P1 fix: ``plot_event_study(weights=...)`` propagates through aggregate().