Fix CI review Round 2: NaN overall for trends, reject controls+heterogeneity

P1: trends_linear + L_max>=2 overall_* is now NaN (R does not compute
    an aggregate in trends_lin mode). Cumulated effects available via
    results.linear_trends_effects[l].

P1: heterogeneity + controls now raises ValueError (matching R's
    predict_het which disallows controls). REGISTRY documents
    heterogeneity as partial implementation (post-treatment only,
    no placebo regressions or joint null test).

P3: Added fit() docstrings for heterogeneity and design2 parameters.
P3: Updated to_dataframe() error text with new level names.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: igerber

diff_diff/chaisemartin_dhaultfoeuille.py

@@ -539,6 +539,17 @@ def fit(
             Requires ``L_max >= 1`` and time-invariant values per group.
         honest_did : bool, default=False
             **Reserved for Phase 3** (HonestDiD integration on placebos).
+        heterogeneity : str, optional
+            Column name for a time-invariant covariate to test for
+            heterogeneous effects (Web Appendix Section 1.5, Lemma 7).
+            Partial implementation: post-treatment regressions only
+            (no placebo regressions or joint null test). Cannot be
+            combined with ``controls``. Requires ``L_max >= 1``.
+        design2 : bool, default=False
+            If ``True``, identify and report switch-in/switch-out
+            (Design-2) groups. Convenience wrapper (descriptive summary,
+            not full paper re-estimation). Requires
+            ``drop_larger_lower=False`` to retain 2-switch groups.
         survey_design : Any, optional
             **Not supported in any phase.** Survey design integration is
             handled as a separate effort after all three phases ship.
@@ -2105,20 +2116,17 @@ def fit(
             linear_trends_effects = cumulated if cumulated else None
 
         # When trends_linear=True and L_max>=2, suppress cost_benefit_delta
-        # (which is computed on second-differences) and set overall_* from
-        # the cumulated level effects instead. This prevents the results
-        # surface from labeling a second-difference aggregate as delta^{fd}
-        # (a level-effect estimand).
+        # and NaN out the overall_* surface. R's did_multiplegt_dyn with
+        # trends_lin=TRUE does not compute an aggregate "average total
+        # effect" - users should access cumulated level effects via
+        # results.linear_trends_effects[l] instead.
         if _is_trends_linear and L_max is not None and L_max >= 2:
             cost_benefit_result = None
-            if linear_trends_effects:
-                max_h = max(linear_trends_effects.keys())
-                lt = linear_trends_effects[max_h]
-                effective_overall_att = lt["effect"]
-                effective_overall_se = lt["se"]
-                effective_overall_t = lt["t_stat"]
-                effective_overall_p = lt["p_value"]
-                effective_overall_ci = lt["conf_int"]
+            effective_overall_att = float("nan")
+            effective_overall_se = float("nan")
+            effective_overall_t = float("nan")
+            effective_overall_p = float("nan")
+            effective_overall_ci = (float("nan"), float("nan"))
 
         # ------------------------------------------------------------------
         # Heterogeneity testing (Web Appendix Section 1.5, Lemma 7)
@@ -2130,6 +2138,14 @@ def fit(
                 raise ValueError(
                     f"heterogeneity column {het_col!r} not found in data."
                 )
+            # R's predict_het disallows controls; our partial implementation
+            # follows this restriction to avoid inconsistent behavior.
+            if controls is not None:
+                raise ValueError(
+                    "heterogeneity cannot be combined with controls. "
+                    "R's did_multiplegt_dyn disallows predict_het with "
+                    "controls; remove one of the two options."
+                )
             # Extract per-group covariate (must be time-invariant)
             het_per_group = data.groupby(group)[het_col].nunique()
             het_varying = het_per_group[het_per_group > 1]

diff_diff/chaisemartin_dhaultfoeuille_results.py

@@ -1054,7 +1054,8 @@ def to_dataframe(self, level: str = "overall") -> pd.DataFrame:
         else:
             raise ValueError(
                 f"Unknown level: {level!r}. Use 'overall', 'joiners_leavers', "
-                f"'per_period', 'event_study', 'normalized', or 'twfe_weights'."
+                f"'per_period', 'event_study', 'normalized', 'twfe_weights', "
+                f"'heterogeneity', or 'linear_trends'."
             )
 
 

docs/methodology/REGISTRY.md

@@ -615,7 +615,7 @@ Alternative: Multiplier bootstrap clustered at group via the `n_bootstrap` param
 
 - **Note (Phase 3 state-set trends):** Implements state-set-specific trends from Web Appendix Section 1.4 (Assumptions 13-14). Restricts the control pool for each switcher to groups in the same set (e.g., same state in county-level data). The restriction applies in BOTH `_compute_multi_horizon_dids()` (point estimates) and `_compute_per_group_if_multi_horizon()` (influence functions) to ensure IF consistency. Cohort structure stays as `(D_{g,1}, F_g, S_g)` triples (does not incorporate set membership). Set membership must be time-invariant per group. Activated via `trends_nonparam="state_column"` in `fit()`.
 
-- **Note (Phase 3 heterogeneity testing):** Implements the heterogeneity test from Web Appendix Section 1.5 (Assumption 15, Lemma 7). Saturated OLS regression of `S_g * (Y_{g, F_g-1+l} - Y_{g, F_g-1})` on a time-invariant covariate `X_g` plus cohort indicator dummies. Standard OLS inference is valid - the paper shows no need to account for DID estimation error. Results stored in `results.heterogeneity_effects`. Activated via `heterogeneity="covariate_column"` in `fit()`.
+- **Note (Phase 3 heterogeneity testing - partial implementation):** Partial implementation of the heterogeneity test from Web Appendix Section 1.5 (Assumption 15, Lemma 7). Computes post-treatment saturated OLS regressions of `S_g * (Y_{g, F_g-1+l} - Y_{g, F_g-1})` on a time-invariant covariate `X_g` plus cohort indicator dummies. Standard OLS inference is valid (paper shows no DID error correction needed). **Deviation from R `predict_het`:** R's full `predict_het` option additionally computes placebo regressions and a joint null test, and disallows combination with `controls`. This implementation provides only post-treatment regressions. Combination with `controls` is rejected (matching R). Results stored in `results.heterogeneity_effects`. Activated via `heterogeneity="covariate_column"` in `fit()`.
 
 - **Note (Phase 3 Design-2 switch-in/switch-out):** Convenience wrapper for Web Appendix Section 1.6 (Assumption 16). Identifies groups with exactly 2 treatment changes (join then leave), reports switch-in and switch-out mean effects. This is a descriptive summary, not a full re-estimation with specialized control pools as described in the paper. The paper notes Design-2 can be implemented by "running the command on a restricted subsample and using `trends_nonparam` for the entry-timing grouping." Activated via `design2=True` in `fit()`, requires `drop_larger_lower=False` to retain 2-switch groups.
 

tests/test_chaisemartin_dhaultfoeuille.py

@@ -2546,24 +2546,31 @@ def test_trends_with_covariates(self):
             df, "outcome", "group", "period", "treatment",
             controls=["X1"], L_max=2, trends_linear=True,
         )
-        assert np.isfinite(r.overall_att)
+        # overall_att is NaN for trends + L_max>=2 (no aggregate)
+        assert np.isnan(r.overall_att)
         assert r.covariate_residuals is not None
         assert r.linear_trends_effects is not None
 
     def test_trends_linear_lmax2_overall_surface(self):
-        """Overall surface under trends_linear + L_max>=2 uses cumulated level effects."""
+        """Under trends_linear + L_max>=2, overall_* is NaN (no aggregate).
+
+        R's did_multiplegt_dyn with trends_lin=TRUE does not compute an
+        aggregate average total effect. Cumulated level effects are
+        available via results.linear_trends_effects[l].
+        """
         df = self._make_panel_with_trends()
         r = ChaisemartinDHaultfoeuille(seed=1).fit(
             df, "outcome", "group", "period", "treatment",
             L_max=3, trends_linear=True,
         )
-        # overall_att should equal the cumulated level effect at max horizon
-        assert r.linear_trends_effects is not None
-        max_h = max(r.linear_trends_effects.keys())
-        cum_effect = r.linear_trends_effects[max_h]["effect"]
-        assert r.overall_att == pytest.approx(cum_effect, abs=1e-10)
-        # cost_benefit_delta should be suppressed (not computed on second-diffs)
+        # overall_* should be NaN (not computed in trends mode)
+        assert np.isnan(r.overall_att)
+        assert np.isnan(r.overall_se)
+        # cost_benefit_delta suppressed
         assert r.cost_benefit_delta is None
+        # Cumulated effects still available
+        assert r.linear_trends_effects is not None
+        assert len(r.linear_trends_effects) >= 1
 
     def test_cumulated_se_nan_propagation(self):
         """Cumulated SE is NaN when a component horizon has NaN SE."""
@@ -2755,6 +2762,16 @@ def test_heterogeneity_missing_column(self):
                 L_max=1, heterogeneity="nonexistent",
             )
 
+    def test_heterogeneity_rejects_controls(self):
+        """heterogeneity + controls raises ValueError (matching R predict_het)."""
+        df = self._make_panel_with_het()
+        df["X1"] = np.random.RandomState(42).normal(0, 1, len(df))
+        with pytest.raises(ValueError, match="cannot be combined with controls"):
+            ChaisemartinDHaultfoeuille(seed=1).fit(
+                df, "outcome", "group", "period", "treatment",
+                L_max=1, heterogeneity="het_x", controls=["X1"],
+            )
+
 
 class TestDesign2:
     """Design-2 switch-in/switch-out separation (ROADMAP item 3e)."""