Address Codex review: MPD+survey reject, NaN unit guard, MPD coords guard

- Conley + survey_design front-door reject on MultiPeriodDiD.fit() so
  the user gets NotImplementedError instead of silent BRR/TSL SEs. The
  previous bypass: MPD passed return_vcov=not _use_survey_vcov to
  solve_ols, so the conley+weights guard inside _compute_robust_vcov_numpy
  never fired, and compute_survey_vcov silently overwrote the vcov.
- solve_ols also gains a top-level Conley-only validator that catches
  bypass cases for direct compute_robust_vcov callers. Scoped to Conley
  only — hc2_bm + replicate-weight silently routing to BRR is a
  long-standing intentional contract preserved by other tests.
- _validate_conley_kwargs now rejects NaN / pd.NA in conley_unit. The
  prior code accepted them and np.unique + boolean mask silently dropped
  those rows from the per-unit serial HAC sum.
- MultiPeriodDiD.fit() adds an explicit conley_coords / conley_cutoff_km
  None-guard so missing kwargs raise ValueError instead of a raw
  TypeError on `self.conley_coords[0]`.
- TwoWayFixedEffects class docstring updated from Phase 1
  "Conley rejected" / "Phase 2 will add Driscoll-Kraay product kernel"
  to the shipped block-decomposed contract.
- Regression tests in tests/test_conley_vcov.py:
  * TestConleyValidatorHelpers: NaN-float and pd.NA-object unit cases
  * TestConleyEstimatorIntegration: MPD + survey_design (pweight TSL and
    stratified PSU) raises; MPD missing conley_coords raises ValueError.

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

Author: igerber

diff_diff/conley.py

@@ -230,13 +230,25 @@ def _validate_conley_kwargs(
             raise ValueError(
                 f"conley_time must be a 1-D array of length {n}; got shape {time_arr.shape}."
             )
-        if not np.isfinite(time_arr).all():
+        # Use pandas.isna for object/categorical dtypes; np.isfinite catches
+        # NaN/inf for numeric dtypes only.
+        import pandas as _pd
+
+        if _pd.isna(time_arr).any():
+            raise ValueError("conley_time contains NaN or missing values.")
+        if time_arr.dtype.kind in "fc" and not np.isfinite(time_arr).all():
             raise ValueError("conley_time contains NaN or inf values.")
         unit_arr = np.asarray(unit)
         if unit_arr.ndim != 1 or unit_arr.shape[0] != n:
             raise ValueError(
                 f"conley_unit must be a 1-D array of length {n}; got shape {unit_arr.shape}."
             )
+        # conley_unit may be any hashable (int, string, categorical). Use
+        # pandas.isna for cross-dtype NA detection. NaN unit IDs would silently
+        # drop those rows from the per-unit serial HAC sum at
+        # `np.unique(unit_arr) + mask_u = unit_arr == u_val`.
+        if _pd.isna(unit_arr).any():
+            raise ValueError("conley_unit contains NaN or missing values.")
         if not (isinstance(lag_cutoff, (int, np.integer)) and int(lag_cutoff) >= 0):
             raise ValueError(
                 f"conley_lag_cutoff must be a non-negative integer; got {lag_cutoff!r}."

diff_diff/estimators.py

@@ -1380,6 +1380,12 @@ def fit(  # type: ignore[override]
         # array extraction (and conley_coords resolution from column names)
         # happens just below at the solve_ols call.
         if self.vcov_type == "conley":
+            if self.conley_coords is None or self.conley_cutoff_km is None:
+                raise ValueError(
+                    "MultiPeriodDiD(vcov_type='conley') requires "
+                    "conley_coords=(lat_col, lon_col) and conley_cutoff_km "
+                    "on the constructor."
+                )
             if unit is None:
                 raise ValueError(
                     "MultiPeriodDiD(vcov_type='conley') requires unit= at "
@@ -1393,6 +1399,15 @@ def fit(  # type: ignore[override]
                     "within-period only, no serial component). See R "
                     "conleyreg's `lag_cutoff` argument for the convention."
                 )
+            if survey_design is not None:
+                raise NotImplementedError(
+                    "MultiPeriodDiD(vcov_type='conley', survey_design=...) "
+                    "is not supported: Conley + survey weights / replicate "
+                    "vcov is deferred to a follow-up PR (Bertanha-Imbens 2014 "
+                    "territory). Use vcov_type='hc1' for survey-aware "
+                    "cluster-robust without spatial HAC, or drop survey_design= "
+                    "for panel Conley."
+                )
         # Pre-compute non_ref_periods (needed for absorb demeaning)
         non_ref_periods = [p for p in all_periods if p != reference_period]
 

diff_diff/linalg.py

@@ -660,6 +660,19 @@ def solve_ols(
                 "Clean your data or set check_finite=False to skip this check."
             )
 
+    # Front-door Conley-specific validation. Runs BEFORE the routing/backend
+    # branching so `return_vcov=False` cannot bypass the conley+weights /
+    # conley+cluster_ids guards. Conley needs this because survey-replicate
+    # paths in MultiPeriodDiD pass `return_vcov=False` + `weights=survey_w`
+    # and would otherwise silently return survey SEs under a Conley request.
+    # The full `_validate_vcov_args` is NOT called here because other vcov
+    # types (e.g. `hc2_bm` + replicate weights) intentionally fall through
+    # to the survey-vcov path with the analytical validator skipped — that
+    # legacy contract is preserved.
+    if vcov_type == "conley":
+        if cluster_ids is not None or weights is not None:
+            _validate_vcov_args(vcov_type, cluster_ids, weights)
+
     # WLS transformation: apply sqrt(w) scaling to X and y
     # This happens BEFORE routing to Rust or NumPy backends — they receive
     # pre-transformed X_w, y_w and solve standard OLS.

diff_diff/twfe.py

@@ -67,19 +67,24 @@ class TwoWayFixedEffects(DifferenceInDifferences):
     ``TODO.md`` under Methodology/Correctness; also documented in
     ``docs/methodology/REGISTRY.md``.
 
-    **Conley spatial-HAC (``vcov_type="conley"``) is rejected at fit-time
-    in Phase 1.** TwoWayFixedEffects is intrinsically a multi-period panel
-    estimator and Phase 1's cross-sectional Conley does not handle the
-    time dimension — applying it over (unit, time) rows would treat same-
-    unit cross-time pairs as ``d_ij = 0 → K = 1``, mishandling the space-
-    time HAC. The supported Phase 1 path for Conley is direct
-    ``compute_robust_vcov`` / ``LinearRegression`` on a single-period
-    regression. The ``conley_*`` kwargs are inherited from
-    ``DifferenceInDifferences.__init__`` for sklearn-style API symmetry
-    (``get_params`` / ``set_params`` round-trip), but
-    ``TwoWayFixedEffects(vcov_type="conley").fit(...)`` raises
-    ``NotImplementedError``. Phase 2 will add the space-time product kernel
-    (Driscoll-Kraay) and lift the rejection.
+    **Conley spatial-HAC (``vcov_type="conley"``) is supported via the
+    block-decomposed panel sandwich (matches R ``conleyreg`` with
+    ``lag_cutoff > 0``).** Pass ``conley_coords=(lat_col, lon_col)``,
+    ``conley_cutoff_km=<float>``, and ``conley_lag_cutoff=<int>`` on the
+    constructor; the ``time`` / ``unit`` arrays are auto-derived from the
+    estimator's ``time`` and ``unit`` column-name arguments at fit-time.
+    The sandwich sums within-period spatial pairs plus within-unit
+    Bartlett serial pairs (lag=0 excluded to avoid double-counting); this
+    is NOT a multiplicative product kernel. FWL composability: the
+    within-transformed scores ``S = X_demeaned * residuals_demeaned`` form
+    the same meat as the full-dummy-expansion design. The temporal kernel
+    is hardcoded Bartlett regardless of ``conley_kernel`` (matches
+    ``conleyreg::time_dist``). Restrictions:
+    ``inference="wild_bootstrap"`` + Conley raises (incompatible inference
+    modes); explicit ``cluster=...`` + Conley raises (combined spatial-
+    cluster kernel deferred to a follow-up PR; auto-cluster on the Conley
+    path is silently dropped); ``survey_design=`` + Conley raises (Phase 5
+    follow-up).
 
     Warning: TWFE can be biased with staggered treatment timing
     and heterogeneous treatment effects. Consider using

tests/test_conley_vcov.py

@@ -433,6 +433,43 @@ def test_panel_time_nan_raises(self):
                 lag_cutoff=1,
             )
 
+    def test_panel_unit_nan_float_raises(self):
+        """NaN unit IDs would silently drop those rows from the per-unit
+        serial HAC sum at `np.unique(unit_arr) + mask_u = unit_arr == u_val`.
+        Closes Codex P1.
+        """
+        n = 4
+        unit = np.array([1.0, 2.0, np.nan, 3.0])
+        with pytest.raises(ValueError, match="conley_unit contains NaN"):
+            _validate_conley_kwargs(
+                coords=np.zeros((n, 2)),
+                cutoff=100.0,
+                metric="euclidean",
+                kernel="bartlett",
+                n=n,
+                time=np.array([1.0, 2.0, 1.0, 2.0]),
+                unit=unit,
+                lag_cutoff=1,
+            )
+
+    def test_panel_unit_pd_na_object_raises(self):
+        """Object-dtype unit IDs (mixed string + pd.NA) must also raise."""
+        import pandas as pd
+
+        n = 4
+        unit = np.array(["A", "B", pd.NA, "C"], dtype=object)
+        with pytest.raises(ValueError, match="conley_unit contains NaN"):
+            _validate_conley_kwargs(
+                coords=np.zeros((n, 2)),
+                cutoff=100.0,
+                metric="euclidean",
+                kernel="bartlett",
+                n=n,
+                time=np.array([1.0, 2.0, 1.0, 2.0]),
+                unit=unit,
+                lag_cutoff=1,
+            )
+
 
 # ---------------------------------------------------------------------------
 # TestConleyDirectHelper — _compute_conley_vcov correctness
@@ -1062,6 +1099,122 @@ def test_multi_period_did_conley_missing_lag_cutoff_raises(self):
                 reference_period=1,
             )
 
+    def test_multi_period_did_conley_with_survey_design_raises(self):
+        """MPD + vcov_type='conley' + survey_design raises NotImplementedError.
+
+        Closes Codex P0: previously, MPD passed return_vcov=False to solve_ols
+        when _use_survey_vcov=True, bypassing the conley + weights guard, and
+        then overwrote vcov with compute_survey_vcov — silently returning
+        survey SEs under a Conley request.
+        """
+        import pandas as pd
+
+        from diff_diff import MultiPeriodDiD
+        from diff_diff.survey import SurveyDesign
+
+        rng = np.random.default_rng(seed=29)
+        n_units = 24
+        rows = []
+        for u in range(n_units):
+            lat = rng.uniform(-30, 30)
+            lon = rng.uniform(-100, 100)
+            for t in range(3):
+                rows.append(
+                    {
+                        "unit": u,
+                        "time": t,
+                        "y": rng.normal(),
+                        "treated": int(u < 12),
+                        "lat": lat,
+                        "lon": lon,
+                        "weight": 1.0 + 0.1 * rng.random(),
+                        "stratum": u % 4,
+                        "psu": u // 6,
+                    }
+                )
+        df_mp = pd.DataFrame(rows)
+        # Pure pweight (no PSU / strata) — would route through analytical conley
+        # path; the guard must fire before solve_ols.
+        sd_tsl = SurveyDesign(weights="weight", weight_type="pweight")
+        with pytest.raises(NotImplementedError, match="survey_design"):
+            MultiPeriodDiD(
+                vcov_type="conley",
+                conley_coords=("lat", "lon"),
+                conley_cutoff_km=2000.0,
+                conley_lag_cutoff=1,
+            ).fit(
+                df_mp,
+                outcome="y",
+                treatment="treated",
+                time="time",
+                post_periods=[2],
+                unit="unit",
+                reference_period=1,
+                survey_design=sd_tsl,
+            )
+        # Stratified PSU survey design — would route through Taylor TSL path
+        # and was the canonical bypass case the codex reviewer flagged.
+        sd_psu = SurveyDesign(
+            weights="weight",
+            strata="stratum",
+            psu="psu",
+            weight_type="pweight",
+            nest=True,
+        )
+        with pytest.raises(NotImplementedError, match="survey_design"):
+            MultiPeriodDiD(
+                vcov_type="conley",
+                conley_coords=("lat", "lon"),
+                conley_cutoff_km=2000.0,
+                conley_lag_cutoff=1,
+            ).fit(
+                df_mp,
+                outcome="y",
+                treatment="treated",
+                time="time",
+                post_periods=[2],
+                unit="unit",
+                reference_period=1,
+                survey_design=sd_psu,
+            )
+
+    def test_multi_period_did_conley_missing_coords_raises(self):
+        """MPD + vcov_type='conley' without conley_coords raises a clean
+        ValueError instead of a raw TypeError on `self.conley_coords[0]`.
+        Closes Codex P2 #1.
+        """
+        import pandas as pd
+
+        from diff_diff import MultiPeriodDiD
+
+        rng = np.random.default_rng(seed=31)
+        n_units = 10
+        rows = []
+        for u in range(n_units):
+            for t in range(2):
+                rows.append(
+                    {
+                        "unit": u,
+                        "time": t,
+                        "y": rng.normal(),
+                        "treated": int(u < 5),
+                    }
+                )
+        df_mp = pd.DataFrame(rows)
+        with pytest.raises(ValueError, match="conley_coords.*conley_cutoff_km"):
+            MultiPeriodDiD(
+                vcov_type="conley",
+                conley_lag_cutoff=1,
+            ).fit(
+                df_mp,
+                outcome="y",
+                treatment="treated",
+                time="time",
+                post_periods=[1],
+                unit="unit",
+                reference_period=0,
+            )
+
 
 class TestConleyTWFE:
     """TwoWayFixedEffects + vcov_type='conley' uses the Phase 2 block-decomposed