R-style NA propagation for avg_att (P1) and scipy_lstsq tolerance (P1)

- avg_att is now NaN if ANY post-period effect is unidentified, matching
  R's default NA propagation semantics (mean(c(1,2,NA)) returns NA)
- Add explicit cond=1e-07 to scipy_lstsq calls for consistency with
  Rust backend and QR rank tolerance
- Document avg_att NA behavior in REGISTRY.md
- Add test for avg_att NaN propagation when period effect is unidentified

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

Author: igerber

diff_diff/estimators.py

@@ -935,29 +935,25 @@ def fit(  # type: ignore[override]
             effect_indices.append(idx)
 
         # Compute average treatment effect
-        # Only average over identified (non-NaN) period effects
+        # R-style NA propagation: if ANY period effect is NaN, average is undefined
         effect_arr = np.array(effect_values)
-        identified_effects = effect_arr[~np.isnan(effect_arr)]
 
-        if len(identified_effects) == 0:
-            # All period effects are NaN - cannot compute average
+        if np.any(np.isnan(effect_arr)):
+            # Some period effects are NaN (unidentified) - cannot compute valid average
+            # This follows R's default behavior where mean(c(1, 2, NA)) returns NA
             avg_att = np.nan
             avg_se = np.nan
             avg_t_stat = np.nan
             avg_p_value = np.nan
             avg_conf_int = (np.nan, np.nan)
         else:
-            # Average ATT = mean of identified period-specific effects
-            avg_att = float(np.mean(identified_effects))
+            # All effects identified - compute average normally
+            avg_att = float(np.mean(effect_arr))
 
             # Standard error of average: need to account for covariance
-            # Only use identified effects in the variance calculation
-            identified_mask = ~np.isnan(effect_arr)
-            identified_indices = [idx for idx, m in zip(effect_indices, identified_mask) if m]
-            n_identified = len(identified_indices)
-
-            sub_vcov = vcov[np.ix_(identified_indices, identified_indices)]
-            avg_var = np.sum(sub_vcov) / (n_identified ** 2)
+            n_post = len(post_periods)
+            sub_vcov = vcov[np.ix_(effect_indices, effect_indices)]
+            avg_var = np.sum(sub_vcov) / (n_post ** 2)
 
             if np.isnan(avg_var) or avg_var < 0:
                 # Vcov has NaN (dropped columns) - propagate NaN

diff_diff/linalg.py

@@ -619,8 +619,9 @@ def _solve_ols_numpy(
         X_reduced = X[:, kept_cols]
 
         # Solve the reduced system (now full-rank)
+        # Use cond=1e-07 for consistency with Rust backend and QR rank tolerance
         coefficients_reduced = scipy_lstsq(
-            X_reduced, y, lapack_driver="gelsd", check_finite=False
+            X_reduced, y, lapack_driver="gelsd", check_finite=False, cond=1e-07
         )[0]
 
         # Expand coefficients to full size with NaN for dropped columns
@@ -638,7 +639,8 @@ def _solve_ols_numpy(
             vcov = _expand_vcov_with_nan(vcov_reduced, k, kept_cols)
     else:
         # Full-rank case: proceed normally
-        coefficients = scipy_lstsq(X, y, lapack_driver="gelsd", check_finite=False)[0]
+        # Use cond=1e-07 for consistency with Rust backend and QR rank tolerance
+        coefficients = scipy_lstsq(X, y, lapack_driver="gelsd", check_finite=False, cond=1e-07)[0]
 
         # Compute residuals and fitted values
         fitted = X @ coefficients

docs/methodology/REGISTRY.md

@@ -101,6 +101,7 @@ where E_i is treatment time for unit i, and δ_e are event-study coefficients.
 - Never-treated units: event-time indicators are all zero
 - Endpoint binning: distant event times can be binned
 - Rank-deficient design matrix (collinearity): warns and sets NA for dropped coefficients (R-style, matches `lm()`)
+- Average ATT (`avg_att`) is NA if any post-period effect is unidentified (R-style NA propagation)
 
 **Reference implementation(s):**
 - R: `fixest::feols()` with `i(event_time, ref=-1)`

tests/test_estimators.py

@@ -1667,6 +1667,46 @@ def test_rank_deficient_design_warns_and_sets_nan(self, multi_period_data):
         assert results.vcov is not None
         assert np.any(np.isnan(results.vcov)), "Vcov should have NaN for dropped column"
 
+        # avg_att should still be computed because all period effects are identified
+        assert not np.isnan(results.avg_att), "avg_att should be valid when all period effects are identified"
+
+    def test_avg_att_nan_when_period_effect_nan(self, multi_period_data):
+        """Test that avg_att is NaN if any period effect is NaN (R-style NA propagation)."""
+        import warnings
+
+        # Remove all treated observations in period 3 to make that interaction
+        # unidentified (column of zeros)
+        data_no_treated_period3 = multi_period_data[
+            ~((multi_period_data["treated"] == 1) & (multi_period_data["period"] == 3))
+        ].copy()
+
+        did = MultiPeriodDiD()
+
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            results = did.fit(
+                data_no_treated_period3,
+                outcome="outcome",
+                treatment="treated",
+                time="period",
+                post_periods=[3, 4, 5]
+            )
+
+        # Should have warning about rank deficiency (treated:period_3 is all zeros)
+        rank_warnings = [x for x in w if "Rank-deficient" in str(x.message)
+                        or "collinear" in str(x.message).lower()]
+        assert len(rank_warnings) > 0, "Expected warning about rank deficiency"
+
+        # The treated×period_3 interaction should have NaN coefficient (unidentified)
+        pe_3 = results.period_effects[3]
+        assert np.isnan(pe_3.effect), "Period 3 effect should be NaN (unidentified)"
+
+        # avg_att should be NaN because one period effect is NaN (R-style NA propagation)
+        assert np.isnan(results.avg_att), "avg_att should be NaN when any period effect is NaN"
+        assert np.isnan(results.avg_se), "avg_se should be NaN when avg_att is NaN"
+        assert np.isnan(results.avg_t_stat), "avg_t_stat should be NaN when avg_att is NaN"
+        assert np.isnan(results.avg_p_value), "avg_p_value should be NaN when avg_att is NaN"
+
 
 class TestSyntheticDiD:
     """Tests for SyntheticDiD estimator."""