Address AI review round 3: normalize np.inf never-treated encoding

first_treat=np.inf (documented as valid for never-treated units) passed
all `> 0` checks in sun_abraham.py and staggered.py, causing np.inf to
be treated as a cohort with -inf relative times. Fixed by normalizing
np.inf to 0 immediately after computing _never_treated in both estimators.
Added regression tests verifying equivalence between encodings.

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

Author: igerber

METHODOLOGY_REVIEW.md

@@ -357,6 +357,12 @@ variables appear to the left of the `|` separator.
    panels these are identical; for unbalanced panels the new formula correctly reflects
    actual sample composition at each event-time. Added unbalanced panel test.
 
+7. **Normalize `np.inf` never-treated encoding** (`sun_abraham.py`, `fit()`):
+   `first_treat=np.inf` (documented as valid for never-treated) was included in
+   `treatment_groups` and `_rel_time` via `> 0` checks, producing `-inf` event times.
+   Fixed by normalizing `np.inf` to `0` immediately after computing `_never_treated`.
+   Same fix applied to `staggered.py` (`CallawaySantAnna`).
+
 **Outstanding Concerns:**
 - **Inference distribution**: Cohort-level p-values use t-distribution (via
   `LinearRegression.get_inference()`), while aggregated event study and overall ATT

diff_diff/staggered.py

@@ -415,6 +415,7 @@ def _precompute_structures(
             cohort_masks[g] = (unit_cohorts == g)
 
         # Never-treated mask
+        # np.inf was normalized to 0 in fit(), so the np.inf check is defensive only
         never_treated_mask = (unit_cohorts == 0) | (unit_cohorts == np.inf)
 
         # Pre-compute covariate matrices by time period if needed
@@ -645,6 +646,8 @@ def fit(
 
         # Never-treated indicator (first_treat = 0 or 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
 
         # Get unique units
         unit_info = df.groupby(unit).agg({

diff_diff/sun_abraham.py

@@ -508,6 +508,8 @@ def fit(
 
         # Never-treated indicator
         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
 
         # Get unique units
         unit_info = (
@@ -1057,6 +1059,7 @@ def _run_bootstrap(
                 df_b[time] - df_b[first_treat],
                 np.nan
             )
+            # np.inf was normalized to 0 in fit(), so the np.inf check is defensive only
             df_b["_never_treated"] = (
                 (df_b[first_treat] == 0) | (df_b[first_treat] == np.inf)
             )

tests/test_staggered.py

@@ -102,6 +102,28 @@ def test_zero_treatment_effect(self):
         # Effect should be close to zero
         assert abs(results.overall_att) < 3 * results.overall_se
 
+    def test_never_treated_inf_encoding(self):
+        """Test that first_treat=np.inf is handled as never-treated, not as a cohort."""
+        data = generate_staggered_data(n_units=200, n_periods=10, n_cohorts=3, seed=42)
+
+        cs = CallawaySantAnna(n_bootstrap=0)
+        results_zero = cs.fit(
+            data.copy(), outcome="outcome", unit="unit", time="time", first_treat="first_treat"
+        )
+
+        # Re-encode never-treated from 0 to np.inf
+        data_inf = data.copy()
+        data_inf.loc[data_inf["first_treat"] == 0, "first_treat"] = np.inf
+
+        results_inf = cs.fit(
+            data_inf, outcome="outcome", unit="unit", time="time", first_treat="first_treat"
+        )
+
+        # Results should be identical
+        assert np.isclose(results_inf.overall_att, results_zero.overall_att), (
+            f"ATT differs: inf={results_inf.overall_att}, zero={results_zero.overall_att}"
+        )
+
     def test_event_study_aggregation(self):
         """Test event study aggregation."""
         data = generate_staggered_data()

tests/test_sun_abraham.py

@@ -1366,3 +1366,39 @@ def test_iw_weights_unbalanced_panel(self):
             assert abs(w - expected_w) < 1e-10, (
                 f"Weight for cohort {g} at e={affected_e}: got {w}, expected {expected_w}"
             )
+
+    def test_never_treated_inf_encoding(self):
+        """Test that first_treat=np.inf is handled as never-treated, not as a cohort."""
+        data = generate_staggered_data(n_units=200, n_periods=10, n_cohorts=3, seed=42)
+
+        # Run with first_treat=0 as baseline
+        sa = SunAbraham(n_bootstrap=0)
+        results_zero = sa.fit(
+            data.copy(), outcome="outcome", unit="unit", time="time", first_treat="first_treat"
+        )
+
+        # Re-encode never-treated from 0 to np.inf
+        data_inf = data.copy()
+        data_inf.loc[data_inf["first_treat"] == 0, "first_treat"] = np.inf
+
+        results_inf = sa.fit(
+            data_inf, outcome="outcome", unit="unit", time="time", first_treat="first_treat"
+        )
+
+        # np.inf must not appear as a cohort in weights
+        for e, weights in results_inf.cohort_weights.items():
+            assert np.inf not in weights, (
+                f"np.inf found as cohort key in weights at e={e}"
+            )
+
+        # No ±inf in event study periods
+        for e in results_inf.event_study_effects.keys():
+            assert np.isfinite(e), f"Non-finite event time {e} in event study"
+
+        # 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}"
+        )