Add empty-treatment guard, last_cohort aggregate/bootstrap tests, fix n_obs

Address rerun review findings:
- Guard against empty treatment_groups before max() in last_cohort branch
- Fix n_obs to report effective sample size (n_units * n_periods) after trimming
- Add last_cohort tests for aggregate=event_study/all and bootstrap
- Add all-never-treated + last_cohort regression test
- Update influence_functions docstring to match dict contract

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

Author: igerber

diff_diff/efficient_did.py

@@ -415,6 +415,11 @@ def fit(
         # Control group logic
         if self.control_group == "last_cohort":
             # Always reclassify last cohort as pseudo-control when requested
+            if not treatment_groups:
+                raise ValueError(
+                    "No treated cohorts found. control_group='last_cohort' requires "
+                    "at least 2 treatment cohorts."
+                )
             last_g = max(treatment_groups)
             treatment_groups = [g for g in treatment_groups if g != last_g]
             if not treatment_groups:
@@ -900,7 +905,7 @@ def fit(
             overall_conf_int=overall_ci,
             groups=treatment_groups,
             time_periods=time_periods,
-            n_obs=len(df),
+            n_obs=n_units * len(time_periods),
             n_treated_units=n_treated_units,
             n_control_units=n_control_units,
             alpha=self.alpha,

diff_diff/efficient_did_results.py

@@ -105,8 +105,10 @@ class EfficientDiDResults:
         ``{(g, t): ndarray}`` — diagnostic: weight vector per target.
     omega_condition_numbers : dict, optional
         ``{(g, t): float}`` — diagnostic: Omega* condition numbers.
-    influence_functions : ndarray, optional
-        Stored EIF matrix for bootstrap / manual SE computation.
+    influence_functions : dict, optional
+        ``{(g, t): ndarray(n_units,)}`` — per-unit EIF values for each
+        group-time cell.  Only populated when ``store_eif=True`` in
+        :meth:`~EfficientDiD.fit` (used internally by ``hausman_pretest``).
     bootstrap_results : EDiDBootstrapResults, optional
         Bootstrap inference results.
     estimation_path : str

tests/test_efficient_did.py

@@ -1026,6 +1026,62 @@ def test_control_group_invalid_raises(self):
         with pytest.raises(ValueError, match="control_group"):
             EfficientDiD(control_group="invalid")
 
+    def test_last_cohort_no_treated_raises(self):
+        """All-never-treated data with last_cohort should raise."""
+        df = _make_staggered_panel(n_per_group=0, n_control=100, groups=())
+        with pytest.raises(ValueError, match="No treated cohorts"):
+            EfficientDiD(control_group="last_cohort").fit(df, "y", "unit", "time", "first_treat")
+
+    def test_last_cohort_aggregate_event_study(self):
+        """last_cohort with aggregate='event_study' should produce finite results."""
+        df = _make_staggered_panel(
+            n_per_group=60,
+            n_control=0,
+            groups=(3, 5, 7),
+            effects={3: 2.0, 5: 1.5, 7: 1.0},
+        )
+        result = EfficientDiD(control_group="last_cohort").fit(
+            df, "y", "unit", "time", "first_treat", aggregate="event_study"
+        )
+        assert result.event_study_effects is not None
+        assert 7 not in result.groups
+        for e, d in result.event_study_effects.items():
+            assert np.isfinite(d["effect"])
+
+    def test_last_cohort_aggregate_all(self):
+        """last_cohort with aggregate='all' should produce finite results."""
+        df = _make_staggered_panel(
+            n_per_group=60,
+            n_control=0,
+            groups=(3, 5, 7),
+            effects={3: 2.0, 5: 1.5, 7: 1.0},
+        )
+        result = EfficientDiD(control_group="last_cohort").fit(
+            df, "y", "unit", "time", "first_treat", aggregate="all"
+        )
+        assert result.event_study_effects is not None
+        assert result.group_effects is not None
+        assert 7 not in result.groups
+        for g, d in result.group_effects.items():
+            assert g != 7
+            assert np.isfinite(d["effect"])
+
+    def test_last_cohort_bootstrap(self, ci_params):
+        """last_cohort with bootstrap should produce finite inference."""
+        n_boot = ci_params.bootstrap(99)
+        df = _make_staggered_panel(
+            n_per_group=60,
+            n_control=0,
+            groups=(3, 5, 7),
+            effects={3: 2.0, 5: 1.5, 7: 1.0},
+        )
+        result = EfficientDiD(control_group="last_cohort", n_bootstrap=n_boot, seed=42).fit(
+            df, "y", "unit", "time", "first_treat"
+        )
+        assert np.isfinite(result.overall_se)
+        assert result.overall_se > 0
+        assert 7 not in result.groups
+
 
 class TestBalanceE:
     """Test balance_e event study balancing."""