Address AI review round 4: move np.inf normalization before treatment_groups

The np.inf → 0 normalization was placed after treatment_groups was
computed, so np.inf passed the `g > 0` filter and leaked into treatment
cohorts. Reorder in both sun_abraham.py and staggered.py so
normalization precedes treatment_groups derivation. Add results.groups
assertion and all-never-treated ValueError test.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: igerber

diff_diff/staggered.py

@@ -640,15 +640,15 @@ def fit(
         # This avoids hardcoding column names in internal methods
         df['first_treat'] = df[first_treat]
 
-        # Identify groups and time periods
-        time_periods = sorted(df[time].unique())
-        treatment_groups = sorted([g for g in df[first_treat].unique() if g > 0])
-
-        # Never-treated indicator (first_treat = 0 or inf)
+        # Never-treated indicator (must precede treatment_groups to exclude np.inf)
         df['_never_treated'] = (df[first_treat] == 0) | (df[first_treat] == np.inf)
         # Normalize np.inf → 0 so all downstream `> 0` checks exclude never-treated
         df.loc[df[first_treat] == np.inf, first_treat] = 0
 
+        # Identify groups and time periods
+        time_periods = sorted(df[time].unique())
+        treatment_groups = sorted([g for g in df[first_treat].unique() if g > 0])
+
         # Get unique units
         unit_info = df.groupby(unit).agg({
             first_treat: 'first',

diff_diff/sun_abraham.py

@@ -502,15 +502,15 @@ def fit(
         df[time] = pd.to_numeric(df[time])
         df[first_treat] = pd.to_numeric(df[first_treat])
 
-        # Identify groups and time periods
-        time_periods = sorted(df[time].unique())
-        treatment_groups = sorted([g for g in df[first_treat].unique() if g > 0])
-
-        # Never-treated indicator
+        # Never-treated indicator (must precede treatment_groups to exclude np.inf)
         df["_never_treated"] = (df[first_treat] == 0) | (df[first_treat] == np.inf)
         # Normalize np.inf → 0 so all downstream `> 0` checks exclude never-treated
         df.loc[df[first_treat] == np.inf, first_treat] = 0
 
+        # Identify groups and time periods
+        time_periods = sorted(df[time].unique())
+        treatment_groups = sorted([g for g in df[first_treat].unique() if g > 0])
+
         # Get unique units
         unit_info = (
             df.groupby(unit)

tests/test_sun_abraham.py

@@ -1395,10 +1395,31 @@ def test_never_treated_inf_encoding(self):
         for e in results_inf.event_study_effects.keys():
             assert np.isfinite(e), f"Non-finite event time {e} in event study"
 
+        # np.inf must not appear in results.groups
+        assert np.inf not in results_inf.groups, (
+            f"np.inf found in results.groups: {results_inf.groups}"
+        )
+
         # Results should be identical to first_treat=0 encoding
         assert np.isclose(results_inf.overall_att, results_zero.overall_att), (
             f"ATT differs: inf={results_inf.overall_att}, zero={results_zero.overall_att}"
         )
         assert np.isclose(results_inf.overall_se, results_zero.overall_se), (
             f"SE differs: inf={results_inf.overall_se}, zero={results_zero.overall_se}"
         )
+
+    def test_all_never_treated_inf_raises(self):
+        """Test that all-never-treated data with np.inf encoding raises ValueError."""
+        data = generate_staggered_data(n_units=100, n_periods=10, n_cohorts=3, seed=42)
+        # Set ALL units to never-treated via np.inf
+        data["first_treat"] = np.inf
+
+        sa = SunAbraham(n_bootstrap=0)
+        with pytest.raises(ValueError, match="No treated units found"):
+            sa.fit(
+                data,
+                outcome="outcome",
+                unit="unit",
+                time="time",
+                first_treat="first_treat",
+            )