Codex Wave A R1: fix sparse haversine cutoff > πR and DiD absorb-time leak

Address Wave A codex review round 1 findings.

P0 #1 — Sparse haversine wrong at cutoff > π·R_earth
  - `_compute_spatial_bartlett_meat_sparse` computed
    `chord_radius = 2·sin(cutoff/(2R))` directly. That function is
    monotone only on [0, πR]; for cutoff > π·R (~20,015 km, half-Earth
    circumference) it shrinks again, and the kd-tree silently drops
    pairs that still have positive Bartlett weight. The dense path
    saturates at πR via `_haversine_km`'s `np.clip(a, 0, 1)`; the
    sparse path now mirrors that by clamping the arc to π radians
    before the chord computation:
        arc_radians = min(cutoff / R_earth, π)
        chord_radius = 2 · sin(arc_radians / 2)
    At cutoff >= πR, chord_radius reaches 2 (sphere diameter), so the
    kd-tree captures all pairs.

P0 #2 — DiD + Conley + absorb leaks demeaned time into the helper
  - `DifferenceInDifferences.fit()` built `_conley_time_arr`,
    `_conley_unit_arr`, and `_conley_coords_arr` from `working_data`,
    but `absorb` demeans columns in `working_data` (including `time`
    when `time` is listed in `absorb`, or whenever `time` appears in
    `vars_to_demean`). The Conley helper would then partition the
    within-period spatial sandwich on residualized floats instead of
    the true pre/post periods. Now reads from the original `data`
    frame, mirroring `TwoWayFixedEffects.fit` which has the same
    FWL-composability contract: meat is computed on demeaned scores
    but the kernel grid uses raw coords + time/unit.

P2 — Regression coverage
  - `test_sparse_haversine_cutoff_above_half_earth_circumference`:
    forces sparse path with cutoff = 25,000 km (> π·R_earth ≈ 20,015 km);
    asserts dense/sparse parity at atol=1e-10.
  - `test_sparse_haversine_cutoff_at_exactly_half_earth_circumference`:
    boundary test at cutoff = π·R_earth exactly.
  - `test_did_conley_with_absorb_uses_raw_time_labels`: monkeypatches
    `_compute_conley_vcov` to capture the `time` arg, asserts the
    helper sees raw integer 0/1 labels (not absorb-demeaned floats).

P3 — Stale docstrings reconciled
  - `diff_diff/estimators.py` DiD class docstring `vcov_type` entry:
    "rejected at fit-time" → Wave A support description.
  - `diff_diff/estimators.py` MultiPeriodDiD class docstring `vcov_type`
    entry: dropped the "cluster= raises" restriction (Wave A #119 lifts
    it via the combined kernel).
  - `diff_diff/linalg.py` `compute_robust_vcov` docstring: `vcov_type`
    entry updated for the combined spatial + cluster kernel.
  - `diff_diff/linalg.py` `LinearRegression.__init__` docstring:
    Conley contract updated for combined cluster kernel + DiD support.

Full Wave A regression: 514 passed (up from 511 with the 3 new tests).

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

Author: igerber

diff_diff/conley.py

@@ -298,11 +298,19 @@ def _compute_spatial_bartlett_meat_sparse(
                 np.sin(lat_r),
             ]
         )
-        # Chord on the unit sphere matching arc-length cutoff_km.
-        # Multiply by (1 + 1e-12) to absorb chord-projection float roundoff
-        # so that pairs whose arc-distance is just under cutoff aren't
-        # excluded by a chord-roundoff false-negative in the kd-tree query.
-        chord_radius = 2.0 * np.sin(cutoff / (2.0 * _CONLEY_EARTH_RADIUS_KM))
+        # Chord on the unit sphere matching arc-length cutoff_km. Clamp the
+        # arc to π radians (half-Earth circumference) before computing the
+        # chord: the function `arc -> 2 sin(arc/2)` is monotonically
+        # increasing only on [0, π] and reaches its global max of 2 (the
+        # unit-sphere diameter) at arc = π. Without the clamp, cutoffs
+        # above π·R (~20,015 km) would compute a SMALLER chord radius than
+        # at cutoff = π·R and silently drop antipodal pairs that still have
+        # positive Bartlett weight. The dense path saturates at π·R via
+        # `_haversine_km`'s `np.clip(a, 0, 1)` (arcsin caps at π/2), so the
+        # clamp mirrors that saturation. The 1+1e-12 epsilon then absorbs
+        # chord-projection float roundoff at sub-cutoff distances.
+        arc_radians = min(cutoff / _CONLEY_EARTH_RADIUS_KM, np.pi)
+        chord_radius = 2.0 * np.sin(arc_radians / 2.0)
         chord_radius *= 1.0 + 1e-12
         tree = cKDTree(xyz)
         neighbors = tree.query_ball_tree(tree, r=chord_radius, p=2.0)

diff_diff/estimators.py

@@ -69,15 +69,18 @@ class DifferenceInDifferences:
           with ``cluster=``, Pustejovsky-Tipton (2018) CR2 cluster-robust.
           (Note: ``MultiPeriodDiD`` does NOT yet support ``cluster=`` with
           ``"hc2_bm"`` — see ``MultiPeriodDiD`` docstring and REGISTRY.md.)
-        - ``"conley"``: Conley 1999 spatial-HAC sandwich. **Accepted by the
-          constructor for sklearn-style API symmetry but rejected at
-          fit-time on ``DifferenceInDifferences``** because DiD is
-          intrinsically a two-period panel design, and Phase 1's cross-
-          sectional Conley does not handle the time dimension. The
-          supported Phase 1 path for Conley is direct
-          ``compute_robust_vcov`` / ``LinearRegression`` on a single-period
-          regression. Phase 2 will add the space-time product kernel
-          (Driscoll-Kraay) and lift the rejection.
+        - ``"conley"``: Conley 1999 spatial-HAC sandwich. Pass
+          ``conley_coords=(lat_col, lon_col)``, ``conley_cutoff_km=<float>``,
+          and ``conley_lag_cutoff=<int>`` on the constructor; pass
+          ``unit=<col>`` as a fit-time kwarg to :meth:`fit` (NOT on
+          ``__init__``; unused unless Conley is set; not part of
+          ``get_params()`` / ``set_params()``). The block-decomposed panel
+          sandwich (matches R ``conleyreg`` with ``lag_cutoff > 0``) sums
+          within-period spatial pairs plus within-unit Bartlett serial
+          pairs (lag=0 excluded). Explicit ``cluster=<col>`` enables the
+          combined spatial + cluster product kernel; ``survey_design=``
+          and ``inference='wild_bootstrap'`` both raise
+          ``NotImplementedError``.
     alpha : float, default=0.05
         Significance level for confidence intervals.
     inference : str, default="analytical"
@@ -514,24 +517,30 @@ def fit(
         # backward compatibility (see `_resolve_effective_vcov_type`).
         _fit_vcov_type = self._resolve_effective_vcov_type(effective_cluster_ids)
 
-        # Build Conley coord/time/unit arrays when applicable. The same
-        # `working_data` (post-absorb demean if absorb was used) supplies
-        # the rows; the validator above already enforced that `unit` is
-        # set and conley_coords/conley_cutoff_km/conley_lag_cutoff are
-        # consistent. The time column for the panel block-decomposed form
-        # is the `time` argument passed to fit() (i.e. the post-treatment
-        # period indicator, which is integer 0/1 for the standard DiD
-        # design — `_compute_conley_vcov` normalizes to dense codes
-        # internally so non-numeric labels still work).
+        # Build Conley coord/time/unit arrays when applicable. CRITICAL:
+        # read from the ORIGINAL `data` frame, NOT `working_data` — `absorb`
+        # demeans `time` (and any column listed in `absorb`) in working_data,
+        # so reading `working_data[time]` would silently partition the
+        # within-period spatial sandwich on residualized floats instead of
+        # the true pre/post periods (Codex Wave A R1 P0). Coords are likewise
+        # read from raw `data` for symmetry with TwoWayFixedEffects
+        # (`twfe.py::TwoWayFixedEffects.fit`) which has the same FWL-
+        # composability contract: the meat is computed on demeaned scores
+        # but the kernel grid uses the original space (coords) and time/unit
+        # indexing. `_compute_conley_vcov` normalizes time labels to dense
+        # codes 0..T-1 internally, so non-numeric `time` labels (datetime64,
+        # pd.Period, strings) still work on the MultiPeriodDiD path; DiD's
+        # binary `time` column is integer 0/1 by convention and is unaffected
+        # by the normalization.
         if _fit_vcov_type == "conley":
             _conley_coords_arr: Optional[np.ndarray] = np.column_stack(
                 [
-                    working_data[self.conley_coords[0]].values.astype(np.float64),
-                    working_data[self.conley_coords[1]].values.astype(np.float64),
+                    data[self.conley_coords[0]].values.astype(np.float64),
+                    data[self.conley_coords[1]].values.astype(np.float64),
                 ]
             )
-            _conley_time_arr: Optional[np.ndarray] = np.asarray(working_data[time].values)
-            _conley_unit_arr: Optional[np.ndarray] = working_data[unit].values
+            _conley_time_arr: Optional[np.ndarray] = np.asarray(data[time].values)
+            _conley_unit_arr: Optional[np.ndarray] = data[unit].values
         else:
             _conley_coords_arr = None
             _conley_time_arr = None
@@ -1123,9 +1132,12 @@ class MultiPeriodDiD(DifferenceInDifferences):
           auto-derived from the ``time`` column at fit-time and normalized
           to dense panel-period codes ``0..T-1`` so ``conley_lag_cutoff``
           always counts panel periods (works for int / datetime64 /
-          ``pd.Period`` / string encodings). Restrictions: ``cluster=``,
-          ``survey_design=``, and ``inference="wild_bootstrap"`` raise on
-          this path (Phase 5 / follow-up).
+          ``pd.Period`` / string encodings). Explicit ``cluster=<col>``
+          enables the combined spatial + cluster product kernel
+          (Wave A #119; cluster must be constant within each unit across
+          periods). Restrictions: ``survey_design=`` and
+          ``inference="wild_bootstrap"`` raise on this path
+          (Phase 5 / follow-up).
     alpha : float, default=0.05
         Significance level for confidence intervals.
     conley_coords, conley_cutoff_km, conley_metric, conley_kernel, conley_lag_cutoff

diff_diff/linalg.py

@@ -1290,9 +1290,13 @@ def compute_robust_vcov(
       bandwidth). Two operating modes: cross-sectional (default) and panel
       block-decomposed (pass the three co-required kwargs ``conley_time`` /
       ``conley_unit`` / ``conley_lag_cutoff``, matching R ``conleyreg`` with
-      ``lag_cutoff > 0``). Combining with ``cluster_ids`` or ``weights``
-      raises ``NotImplementedError`` (combined spatial-cluster kernel and
-      Bertanha-Imbens weighted-Conley deferred to follow-up PRs).
+      ``lag_cutoff > 0``). Combining with ``cluster_ids`` applies the
+      combined spatial + cluster product kernel
+      ``K_total[i, j] = K_space(d_ij/h) · 1{c_i = c_j}`` (Wave A #119; on
+      the panel path the cluster must be constant within each unit across
+      periods). Combining with ``weights`` still raises
+      ``NotImplementedError`` (Bertanha-Imbens 2014 weighted-Conley
+      deferred to a follow-up PR).
 
     Parameters
     ----------
@@ -2485,13 +2489,15 @@ class LinearRegression:
         pass the three co-required kwargs ``conley_time`` / ``conley_unit`` /
         ``conley_lag_cutoff``). Requires ``conley_coords`` (n × 2 array) and
         a positive ``conley_cutoff_km``. Combining ``vcov_type="conley"``
-        with ``cluster_ids``, ``weights``, or ``survey_design`` raises
-        ``NotImplementedError`` (combined spatial-cluster kernel +
-        Bertanha-Imbens weighted-Conley deferred to follow-up PRs). The
-        panel DiD estimator rejects at fit-time (DiD.fit() has no unit
-        column declaration; use MultiPeriodDiD or TwoWayFixedEffects, which
-        thread their `unit`/`time` columns into the block-decomposed
-        sandwich via ``conley_lag_cutoff`` on the constructor).
+        with ``cluster_ids`` applies the combined spatial + cluster product
+        kernel (Wave A #119; cluster must be constant within each unit on
+        the panel path). Combining with ``weights`` / ``survey_design``
+        still raises ``NotImplementedError`` (Bertanha-Imbens 2014
+        weighted-Conley deferred to a follow-up PR). The DiD / MPD / TWFE
+        estimators all support panel Conley by passing ``unit`` at fit-time
+        (DiD as a fit-time kwarg; MPD/TWFE via the existing ``unit``
+        argument), threading ``conley_time`` / ``conley_unit`` into the
+        block-decomposed sandwich.
     conley_coords : ndarray of shape (n, 2), optional
         Required when ``vcov_type="conley"``. Two-column array of
         ``[lat, lon]`` (degrees, for ``conley_metric="haversine"``) or

tests/test_conley_vcov.py

@@ -1124,6 +1124,54 @@ def test_did_conley_matches_mpd_post_periods_1(self, two_period_panel):
         np.testing.assert_allclose(res_did.att, res_mpd.att, atol=1e-10)
         np.testing.assert_allclose(res_did.se, res_mpd.se, atol=1e-10)
 
+    def test_did_conley_with_absorb_uses_raw_time_labels(self, two_period_panel, monkeypatch):
+        """DiD + Conley + absorb=[<unit>] must feed the Conley helper the
+        ORIGINAL time/unit/coord columns from `data`, not the absorb-demeaned
+        `working_data` (in which time has been residualized to floats).
+        Otherwise the within-period spatial sandwich silently partitions on
+        per-unit demeaned floats instead of the true pre/post periods.
+        Codex Wave A R1 P0 #2.
+        """
+        import diff_diff.linalg as linalg_module
+        from diff_diff import DifferenceInDifferences
+
+        df = two_period_panel.copy()
+        captured: dict = {"time_arg": None, "unit_arg": None}
+        orig = linalg_module._compute_conley_vcov
+
+        def _spy(*args, **kwargs):
+            captured["time_arg"] = kwargs.get("time")
+            captured["unit_arg"] = kwargs.get("unit")
+            return orig(*args, **kwargs)
+
+        monkeypatch.setattr(linalg_module, "_compute_conley_vcov", _spy)
+        DifferenceInDifferences(
+            vcov_type="conley",
+            conley_coords=("lat", "lon"),
+            conley_cutoff_km=2000.0,
+            conley_lag_cutoff=1,
+        ).fit(
+            df,
+            outcome="y",
+            treatment="treated",
+            time="time",
+            unit="unit",
+            absorb=["unit"],
+        )
+        assert captured["time_arg"] is not None
+        # Raw labels are integer 0/1 (the binary post-treatment indicator);
+        # demeaned values would be floats from absorb's within-unit
+        # demeaning. np.unique on raw labels yields exactly 2 distinct
+        # values; on demeaned floats it would yield ~n_units distinct.
+        time_arg = np.asarray(captured["time_arg"])
+        uniques = np.unique(time_arg)
+        assert len(uniques) == 2, (
+            f"Expected 2 unique time labels (raw 0/1), got {len(uniques)}: "
+            f"{uniques[:5]} — absorb is leaking demeaned time into the "
+            "Conley helper."
+        )
+        assert set(uniques.tolist()) == {0, 1}, f"Expected raw integer labels 0/1, got {uniques}"
+
     def test_multi_period_did_with_conley_panel(self):
         """Phase 2 MultiPeriodDiD + vcov_type='conley' uses the block-decomposed
         sandwich (matches R conleyreg). Verifies that finite SEs are produced
@@ -2631,6 +2679,88 @@ def test_helper_direct_matches_dense_meat_haversine(self):
         meat_sparse = _compute_spatial_bartlett_meat_sparse(S, coords, cutoff, "haversine")
         np.testing.assert_allclose(meat_sparse, meat_dense, atol=1e-10, rtol=1e-10)
 
+    def test_sparse_haversine_cutoff_above_half_earth_circumference(self):
+        """Sparse haversine path with conley_cutoff_km > π·R_earth (~20,015 km)
+        must include all geometrically eligible pairs. Without the arc-radians
+        clamp, the chord-radius formula 2·sin(arc/2) shrinks for arc > π and
+        the kd-tree silently drops pairs that still have positive Bartlett
+        weight. The dense path saturates at π·R via _haversine_km's clip;
+        the sparse path matches via the clamp. Codex Wave A R1 P0 #1.
+        """
+        rng = np.random.default_rng(seed=101)
+        n = 200
+        lats = rng.uniform(-90.0, 90.0, size=n)
+        lons = rng.uniform(-180.0, 180.0, size=n)
+        coords = np.column_stack([lats, lons])
+        X = np.column_stack([np.ones(n), rng.standard_normal(n), rng.standard_normal(n)])
+        y = X @ np.array([1.0, 2.0, -0.5]) + rng.standard_normal(n) * 0.4
+        coefs, *_ = np.linalg.lstsq(X, y, rcond=None)
+        residuals = y - X @ coefs
+        bread = X.T @ X
+        # Cutoff well above half-Earth circumference (~20,015 km). Without
+        # the clamp, the sparse path drops antipodal pairs and the meat
+        # diverges from the dense path.
+        cutoff_km = 25_000.0  # > π·R_earth ≈ 20015 km
+        V_dense = _compute_conley_vcov(
+            X,
+            residuals,
+            coords,
+            cutoff_km,
+            "haversine",
+            "bartlett",
+            bread,
+            _conley_sparse=False,
+        )
+        V_sparse = _compute_conley_vcov(
+            X,
+            residuals,
+            coords,
+            cutoff_km,
+            "haversine",
+            "bartlett",
+            bread,
+            _conley_sparse=True,
+        )
+        np.testing.assert_allclose(V_sparse, V_dense, atol=1e-10, rtol=1e-10)
+
+    def test_sparse_haversine_cutoff_at_exactly_half_earth_circumference(self):
+        """Cutoff = π·R_earth: chord radius = 2 (sphere diameter); all
+        pairs are included. Bartlett at u=1 returns 0, so the antipodal
+        pair contributes zero — but pairs at all other distances
+        contribute. Sparse and dense paths must agree."""
+        rng = np.random.default_rng(seed=103)
+        n = 150
+        lats = rng.uniform(-90.0, 90.0, size=n)
+        lons = rng.uniform(-180.0, 180.0, size=n)
+        coords = np.column_stack([lats, lons])
+        X = np.column_stack([np.ones(n), rng.standard_normal(n)])
+        y = X @ np.array([1.0, 1.5]) + rng.standard_normal(n) * 0.5
+        coefs, *_ = np.linalg.lstsq(X, y, rcond=None)
+        residuals = y - X @ coefs
+        bread = X.T @ X
+        cutoff_km = float(np.pi * _CONLEY_EARTH_RADIUS_KM)  # ≈ 20015.16 km
+        V_dense = _compute_conley_vcov(
+            X,
+            residuals,
+            coords,
+            cutoff_km,
+            "haversine",
+            "bartlett",
+            bread,
+            _conley_sparse=False,
+        )
+        V_sparse = _compute_conley_vcov(
+            X,
+            residuals,
+            coords,
+            cutoff_km,
+            "haversine",
+            "bartlett",
+            bread,
+            _conley_sparse=True,
+        )
+        np.testing.assert_allclose(V_sparse, V_dense, atol=1e-10, rtol=1e-10)
+
     def test_panel_block_decomposed_sparse_matches_dense(self):
         """Panel block-decomposed sandwich produces the same vcov whether
         the spatial component is computed dense or sparse. The serial