Address CI AI review: NaN first_treat + StaggeredTripleDifference coverage

Two new P1 findings from CI AI re-review:

1. ContinuousDiD.fit() still accepted NaN `first_treat` values. NaN
   survives preprocessing but satisfies neither the treated (g > 0)
   nor never-treated (g == 0) mask, so affected units were silently
   excluded from both pools — same silent-failure shape as the
   already-rejected `first_treat < 0`. Reject NaN explicitly with
   ValueError + row count.

2. `StaggeredTripleDifference.fit()` silently recoded
   `first_treat=np.inf -> 0` via `.replace([np.inf, float("inf")], 0)`.
   Its sibling `staggered.py:1508-1519` already surfaces this with a
   UserWarning; this PR mirrors that contract so the estimator no
   longer shifts units between treated and never-treated pools
   without signaling. REGISTRY gets a matching **Note:** under the
   StaggeredTripleDifference section.

Regression tests:
- test_nan_first_treat_raises_with_row_count (ContinuousDiD) on a
  4-unit x 3-period panel with 3 NaN rows.
- test_inf_first_treat_works (StaggeredTripleDifference) upgraded
  from silent-success to `pytest.warns(UserWarning, match=row-count)`.

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

Author: igerber

diff_diff/continuous_did.py

@@ -235,6 +235,19 @@ def fit(
         # existing `.replace([np.inf, float("inf")], 0)` semantics on the
         # next line).
         first_treat_vals = df[first_treat].values
+        # Reject NaN first_treat explicitly. NaN survives preprocessing but
+        # satisfies neither the treated (g > 0) nor never-treated (g == 0)
+        # mask, so affected units would be silently excluded from the
+        # estimator (same silent-failure shape as `first_treat < 0`).
+        nan_mask = pd.isna(df[first_treat])
+        n_nan_first_treat = int(nan_mask.sum())
+        if n_nan_first_treat > 0:
+            raise ValueError(
+                f"{n_nan_first_treat} row(s) have NaN '{first_treat}' "
+                f"values. Valid values are 0 (never-treated) or a positive "
+                f"treatment period; such units would otherwise be silently "
+                f"excluded from both treated and control pools."
+            )
         inf_mask = np.isposinf(first_treat_vals)
         n_inf_first_treat = int(inf_mask.sum())
         if n_inf_first_treat > 0:

diff_diff/staggered_triple_diff.py

@@ -284,6 +284,19 @@ def fit(
 
         if first_treat != "first_treat":
             df["first_treat"] = df[first_treat]
+        # Surface the inf → 0 recategorization the same way StaggeredDiD does
+        # (see `staggered.py:1508-1519`). Silently recoding inf would shift
+        # units between treated and never-treated pools with no signal
+        # (axis-E silent coercion under the Phase 2 audit).
+        _inf_mask = np.isposinf(df["first_treat"].values)
+        if _inf_mask.any():
+            n_inf_rows = int(_inf_mask.sum())
+            warnings.warn(
+                f"{n_inf_rows} row(s) have first_treat=inf; recoding to 0 "
+                f"(never-treated). Use first_treat=0 to suppress this warning.",
+                UserWarning,
+                stacklevel=2,
+            )
         df["first_treat"] = df["first_treat"].replace([np.inf, float("inf")], 0)
 
         precomputed = self._precompute_structures(

docs/methodology/REGISTRY.md

@@ -1692,6 +1692,7 @@ Balanced panel. Key variables:
 - `Q_i` (`eligibility`): binary, time-invariant eligibility indicator
 - Treatment: `D_{i,t} = 1{t >= S_i AND Q_i = 1}` (absorbing)
 - Covariates `X_i`: time-invariant (first observation per unit used)
+- **Note:** `first_treat=inf` (R-style never-enabled marker) is accepted and normalized to `0` internally. The recoding now emits a `UserWarning` reporting the affected row count so the reclassification is not silent (axis-E silent coercion under the Phase 2 audit, mirroring the StaggeredDiD behavior). Pass `first_treat=0` directly to avoid the warning.
 
 *Estimator equation (Equation 4.1 in paper, as implemented):*
 

tests/test_continuous_did.py

@@ -751,6 +751,29 @@ def test_negative_first_treat_raises_with_row_count(self):
         ):
             est.fit(data, "outcome", "unit", "period", "first_treat", "dose")
 
+    def test_nan_first_treat_raises_with_row_count(self):
+        """NaN `first_treat` must raise ValueError with the row count. Without
+        this guard, NaN rows survive preprocessing but match neither the
+        treated (g > 0) nor never-treated (g == 0) mask, so the affected
+        units would be silently excluded."""
+        rows = []
+        for unit in range(4):
+            # Unit 0 has NaN first_treat across all 3 periods (3 NaN rows).
+            ft = np.nan if unit == 0 else 0.0
+            for t in range(1, 4):
+                rows.append({
+                    "unit": unit, "period": t, "outcome": float(unit + t),
+                    "first_treat": ft, "dose": 0.0,
+                })
+        data = pd.DataFrame(rows)
+        est = ContinuousDiD()
+
+        with pytest.raises(
+            ValueError,
+            match=r"3 row\(s\) have NaN 'first_treat' values",
+        ):
+            est.fit(data, "outcome", "unit", "period", "first_treat", "dose")
+
     def test_positive_inf_warning_silent_when_no_inf(self):
         """+inf warning is gated on +inf rows only; panels with only valid
         non-negative values (including just 0 and positive periods) must

tests/test_staggered_triple_diff.py

@@ -397,12 +397,23 @@ def test_missing_column_raises(self, simple_data):
             est.fit(simple_data, "outcome", "unit", "period", "nonexistent", "eligibility")
 
     def test_inf_first_treat_works(self):
-        """Never-enabled units encoded as inf should work."""
+        """Never-enabled units encoded as inf should be recoded to 0, and the
+        recoding must surface a UserWarning with the affected row count
+        (axis-E silent coercion, mirroring the StaggeredDiD behavior)."""
         data = generate_staggered_ddd_data(n_units=100, seed=33)
         data["first_treat"] = data["first_treat"].astype(float)
-        data.loc[data["first_treat"] == 0, "first_treat"] = np.inf
+        inf_mask = data["first_treat"] == 0
+        n_inf_rows = int(inf_mask.sum())
+        data.loc[inf_mask, "first_treat"] = np.inf
         est = StaggeredTripleDifference()
-        res = est.fit(data, "outcome", "unit", "period", "first_treat", "eligibility")
+
+        with pytest.warns(
+            UserWarning,
+            match=rf"{n_inf_rows} row\(s\) have first_treat=inf; recoding to 0",
+        ):
+            res = est.fit(
+                data, "outcome", "unit", "period", "first_treat", "eligibility"
+            )
         assert np.isfinite(res.overall_att)
 
     def test_survey_design_invalid_type_raises(self, simple_data):