Codex CI R8 P1: cluster-aware sparse density gate + sparse+cluster regressions

Two related P1s on R8 (both about the sparse+cluster interaction).

P1 (Performance) — density gate is now cluster-aware
  The R6 density gate counted ALL spatial in-range pairs before the
  cluster mask was applied, but the actual CSR matrix stores only
  within-cluster pairs (the inner row-loop drops cross-cluster
  neighbors before populating row/col/data lists). On large clustered
  panels — exactly the use case sparse+cluster is supposed to win —
  spatial density could be ~100% while within-cluster density is tiny,
  and we'd spuriously fall back to dense O(n²) work.

  Fix: when the unrefined spatial density exceeds the threshold AND
  `cluster_codes is not None`, refine by counting only within-cluster
  in-range pairs (per-cluster kd-trees, summed). The refinement runs
  only when needed (the gate is already cleared if spatial density is
  below threshold), so the per-cluster sub-tree builds are amortized.

  Stores the projected coordinates as `tree_coords` once on the
  haversine branch (was `xyz` locally before; renaming makes the
  per-cluster build symmetric across metrics).

P1 (Documentation/Tests) — sparse+cluster regression coverage
  Sparse coverage previously exercised unclustered dense-vs-sparse
  parity; cluster coverage previously exercised dense/estimator
  behavior — but the combined sparse+cluster path was load-bearing
  and untested. Adds three regressions:
    - test_sparse_with_cluster_matches_dense: cross-sectional 600-row
      panel, 8 clusters, dense vs forced-sparse parity at atol=1e-10.
    - test_sparse_with_cluster_panel_matches_dense: 3-period block-
      decomposed panel, 200 units, 5 time-invariant clusters, dense
      vs forced-sparse parity at atol=1e-10.
    - test_sparse_density_gate_cluster_aware: tight spatial cluster
      (100% global density) + 50 small clusters → within-cluster
      density << 30%. Locks the refinement contract: NO density
      warning, sparse path executes, result matches dense.

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

Author: igerber

diff_diff/conley.py

@@ -457,11 +457,13 @@ def _compute_spatial_bartlett_meat_sparse(
         query_r = 2.0 * np.sin(arc_radians / 2.0)
         query_r *= 1.0 + 1e-12
         tree = cKDTree(xyz)
+        tree_coords = xyz  # used for per-cluster sub-trees in the density gate
     elif metric == "euclidean":
         # Small relative epsilon for symmetry with the haversine branch.
         # Bartlett's <=-vs-< boundary is moot since kernel is exactly 0 at u=1.
         query_r = cutoff * (1.0 + 1e-12)
         tree = cKDTree(coords)
+        tree_coords = coords
     else:
         raise ValueError(
             "sparse Conley path requires metric in {'haversine', 'euclidean'}; "
@@ -475,6 +477,25 @@ def _compute_spatial_bartlett_meat_sparse(
     # memory than dense float64). Codex CI R6 P2.
     n_pairs_in_range = int(tree.count_neighbors(tree, r=query_r, p=2.0))
     density = n_pairs_in_range / float(n * n)
+    if density > density_threshold and cluster_codes is not None:
+        # Cluster-aware refinement: the actual CSR nnz reflects within-
+        # cluster pairs only (the inner loop drops cross-cluster neighbors
+        # before adding to the sparse matrix). On large clustered panels,
+        # spatial density can be ~100% while within-cluster density is
+        # tiny — without refinement we'd spuriously fall back to dense
+        # exactly when sparse helps most. Refine by counting per-cluster.
+        # Codex CI R8 P1: the unrefined density check was cluster-blind.
+        refined_pairs = 0
+        for c in np.unique(cluster_codes):
+            in_c = cluster_codes == c
+            n_c = int(in_c.sum())
+            if n_c <= 1:
+                # Singleton cluster contributes only the diagonal self-pair.
+                refined_pairs += n_c
+                continue
+            sub_tree = cKDTree(tree_coords[in_c])
+            refined_pairs += int(sub_tree.count_neighbors(sub_tree, r=query_r, p=2.0))
+        density = refined_pairs / float(n * n)
     if density > density_threshold:
         warnings.warn(
             f"Conley sparse path: neighbor density {density:.1%} exceeds "

tests/test_conley_vcov.py

@@ -3125,6 +3125,145 @@ def test_sparse_density_gate_does_not_trigger_below_threshold(self):
                 len(density_warnings) == 0
             ), "Density gate should not have triggered at low density"
 
+    def test_sparse_with_cluster_matches_dense(self):
+        """Sparse + combined cluster kernel matches the dense path bit-for-bit
+        at atol=1e-10 (cross-sectional). Codex CI R8 P1: load-bearing
+        sparse+cluster regression that was missing prior."""
+        rng = np.random.default_rng(seed=181)
+        n = 600
+        coords = rng.uniform(0.0, 100.0, size=(n, 2))
+        cluster_ids = rng.integers(0, 8, size=n)  # 8 clusters
+        X = np.column_stack([np.ones(n), rng.standard_normal(n), rng.standard_normal(n)])
+        y = X @ np.array([1.0, 1.5, -0.5]) + rng.standard_normal(n) * 0.4
+        coefs, *_ = np.linalg.lstsq(X, y, rcond=None)
+        residuals = y - X @ coefs
+        bread = X.T @ X
+        V_dense = _compute_conley_vcov(
+            X,
+            residuals,
+            coords,
+            15.0,
+            "euclidean",
+            "bartlett",
+            bread,
+            cluster_ids=cluster_ids,
+            _conley_sparse=False,
+        )
+        V_sparse = _compute_conley_vcov(
+            X,
+            residuals,
+            coords,
+            15.0,
+            "euclidean",
+            "bartlett",
+            bread,
+            cluster_ids=cluster_ids,
+            _conley_sparse=True,
+        )
+        np.testing.assert_allclose(V_sparse, V_dense, atol=1e-10, rtol=1e-10)
+
+    def test_sparse_with_cluster_panel_matches_dense(self):
+        """Sparse + cluster + panel block-decomposed sandwich matches the
+        dense path on a 3-period panel with time-invariant cluster.
+        Codex CI R8 P1."""
+        rng = np.random.default_rng(seed=191)
+        n_units = 200
+        T = 3
+        unit = np.repeat(np.arange(n_units), T)
+        time = np.tile(np.arange(T), n_units)
+        n = n_units * T
+        # Time-invariant coords + cluster (per-unit)
+        unit_coords = rng.uniform(0.0, 100.0, size=(n_units, 2))
+        coords = unit_coords[unit]
+        cluster_per_unit = rng.integers(0, 5, size=n_units)
+        cluster_ids = cluster_per_unit[unit]
+        X = np.column_stack([np.ones(n), rng.standard_normal(n), rng.standard_normal(n)])
+        y = X @ np.array([1.0, 1.5, -0.3]) + rng.standard_normal(n) * 0.4
+        coefs, *_ = np.linalg.lstsq(X, y, rcond=None)
+        residuals = y - X @ coefs
+        bread = X.T @ X
+        V_dense = _compute_conley_vcov(
+            X,
+            residuals,
+            coords,
+            15.0,
+            "euclidean",
+            "bartlett",
+            bread,
+            time=time,
+            unit=unit,
+            lag_cutoff=1,
+            cluster_ids=cluster_ids,
+            _conley_sparse=False,
+        )
+        V_sparse = _compute_conley_vcov(
+            X,
+            residuals,
+            coords,
+            15.0,
+            "euclidean",
+            "bartlett",
+            bread,
+            time=time,
+            unit=unit,
+            lag_cutoff=1,
+            cluster_ids=cluster_ids,
+            _conley_sparse=True,
+        )
+        np.testing.assert_allclose(V_sparse, V_dense, atol=1e-10, rtol=1e-10)
+
+    def test_sparse_density_gate_cluster_aware(self):
+        """High global spatial density + many small clusters: within-cluster
+        density is low, so the sparse path should NOT spuriously fall back
+        to dense. Tests the cluster-aware refinement of the density gate.
+        Codex CI R8 P1."""
+        # Tight spatial cluster of points → 100% global spatial density,
+        # but split into many small disjoint clusters so post-mask density
+        # is well below 30%.
+        rng = np.random.default_rng(seed=199)
+        n = 200
+        coords = rng.uniform(0.0, 5.0, size=(n, 2))  # tight cluster
+        # 50 clusters of 4 points each → within-cluster nnz = 4*4 = 16 per
+        # cluster, total = 50*16 = 800 = 800/(200*200) = 2% density << 30%.
+        cluster_ids = np.repeat(np.arange(50), 4)
+        X = np.column_stack([np.ones(n), rng.standard_normal(n)])
+        y = X @ np.array([1.0, 1.5]) + rng.standard_normal(n) * 0.4
+        coefs, *_ = np.linalg.lstsq(X, y, rcond=None)
+        residuals = y - X @ coefs
+        bread = X.T @ X
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            V_sparse = _compute_conley_vcov(
+                X,
+                residuals,
+                coords,
+                1e6,  # cutoff > data span → 100% global density
+                "euclidean",
+                "bartlett",
+                bread,
+                cluster_ids=cluster_ids,
+                _conley_sparse=True,
+            )
+            density_warnings = [msg for msg in w if "exceeds threshold" in str(msg.message)]
+            assert len(density_warnings) == 0, (
+                "Density gate spuriously fell back to dense even though "
+                "within-cluster density is low — cluster-aware refinement "
+                "is not working."
+            )
+        # Result must equal the dense path (sparse path executed correctly)
+        V_dense = _compute_conley_vcov(
+            X,
+            residuals,
+            coords,
+            1e6,
+            "euclidean",
+            "bartlett",
+            bread,
+            cluster_ids=cluster_ids,
+            _conley_sparse=False,
+        )
+        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