fix: stabilize nearest-neighbor matching ties

Author: brycewang-stanford

.tierd_campaign/CAMPAIGN.md

@@ -457,3 +457,15 @@ noisy/forest-dependent — anchored on β₁ + null structure instead.
   `proximal`): 17/17 green.
 - **P2 tally: 20 batches, 114 tests, 47 estimators.** worklist → 182
   (`scripts/tierd_classify.py` exact).
+
+### 2026-06-09 (cont.) — matching tie-break reproducibility follow-up
+- Fixed the deferred nearest-neighbor reproducibility lane: `sp.match`
+  Euclidean/propensity NN ties are now resolved by source DataFrame index
+  instead of `argpartition` / incidental row order. Lower-index pool units win
+  exact equal-distance ties; lower-index target units break
+  without-replacement assignment ties. This is a correctness/reproducibility
+  fix, not a new estimator batch.
+- Added a shuffled-row regression test for equal-distance controls and tightened
+  the LaLonde guard from the old cross-backend PSM band to the pinned
+  `1963.43` ATT. CHANGELOG/MIGRATION document that only exact-tie designs can
+  move.

CHANGELOG.md

@@ -184,12 +184,18 @@ All notable changes to StatsPAI will be documented in this file.
     outside its own enum. (Also surfaced and corrected a latent duplicate
     `harvest_did` registration.)
 
-- **`sp.replicate('lalonde_1986')` classic-track 1:1 NN PSM golden number
-  refreshed 2012.5 → 1963.4.** The original pin was stale relative to the
-  current deterministic `sp.match(method='nearest')` output (a tie-handling
-  change on the binary covariates moved it ~2.5%); no estimator numerics
-  changed. A new regression guard (`tests/test_tierD_lalonde_psm_guard.py`)
-  pins all three classic LaLonde numbers so future drift fails loudly.
+- **⚠️ Reproducibility fix: `sp.match(method='nearest')` now resolves exact
+  nearest-neighbor ties by source DataFrame index.** The Euclidean/propensity
+  nearest-neighbor path previously relied on `argpartition` / incidental row
+  order for equal-distance controls (and target order when matching without
+  replacement), so exact ties on discrete or binary covariates could move the
+  ATT across environments. Equal-distance ties now select lower-index pool
+  units first, with lower-index target units used as the without-replacement
+  fallback. A shuffled-row regression test guards the policy, and
+  `tests/test_tierD_lalonde_psm_guard.py` is tightened from the old
+  cross-backend band to the pinned LaLonde PSM ATT (`1963.43`). Existing results
+  change only when previous data contained exact equal-distance ties. See
+  [MIGRATION.md](MIGRATION.md#matching-nearest-tie-break).
 
 ### Known issues
 

MIGRATION.md

@@ -41,6 +41,30 @@ numeric figure uses an HR or CI E-value.
 
 ---
 
+<a id="matching-nearest-tie-break"></a>
+
+## Unreleased — ⚠️ `sp.match` nearest-neighbor tie-breaking stabilised
+
+**What changed.** `sp.match(method='nearest')` now resolves exact equal-distance
+nearest-neighbor ties by the source DataFrame index. Previously the
+Euclidean/propensity nearest-neighbor path delegated tie selection to
+`argpartition` and incidental row order, so ties on discrete or binary
+covariates could move the ATT across environments. Lower-index control units
+are now selected first; when matching without replacement and multiple treated
+units have the same best distance, lower-index treated units are assigned first.
+
+**Who is affected.** Only users whose matching data contain exact
+equal-distance ties. Continuous covariates without exact ties are unchanged.
+For tied designs, results are now deterministic across row order and backend as
+long as the DataFrame index preserves unit identity.
+
+**Action required.** None for code. If you previously recorded a nearest-match
+estimate on tied discrete covariates, re-run it once and treat the new value as
+the stable pin. The bundled LaLonde 1:1 NN PSM guard is now pinned at `1963.43`
+instead of allowing the old cross-backend tie band.
+
+---
+
 <a id="blp-maxiter-fix"></a>
 
 ## Unreleased — ⚠️ `sp.blp` functionality fix (was non-functional)

src/statspai/matching/match.py

@@ -315,6 +315,7 @@ def fit(self) -> CausalResult:
         T = clean[self.treat].values.astype(int)
         Y = clean[self.y].values.astype(float)
         X = clean[self.covariates].values.astype(float)
+        row_order = self._stable_index_order(clean.index)
 
         idx_t = np.where(T == 1)[0]
         idx_c = np.where(T == 0)[0]
@@ -346,7 +347,7 @@ def fit(self) -> CausalResult:
             att, se, balance, extra_info = self._fit_exact(Y, X, T, idx_t, idx_c)
             method_label = 'Matching (Exact)'
         else:
-            att, se, balance = self._fit_nearest(Y, X, T, idx_t, idx_c)
+            att, se, balance = self._fit_nearest(Y, X, T, idx_t, idx_c, row_order)
             dist_name = self.distance.capitalize()
             bc_tag = ', BC' if self.bias_correction else ''
             method_label = f'Matching ({dist_name}{bc_tag})'
@@ -398,23 +399,41 @@ def fit(self) -> CausalResult:
     # Nearest-neighbor matching (propensity / mahalanobis / euclidean)
     # ==================================================================
 
-    def _fit_nearest(self, Y, X, T, idx_t, idx_c):
+    def _fit_nearest(self, Y, X, T, idx_t, idx_c, row_order=None):
         """Nearest-neighbor matching with configurable distance metric."""
+        if row_order is None:
+            row_order = np.arange(len(T), dtype=float)
+
         # For propensity distance, estimate PS once with actual treatment
         pscore = self._logit_propensity(X, T, poly=self.ps_poly) if self.distance == 'propensity' else None
 
         # Build distance matrix
         dist_mat = self._compute_distance_matrix(X, idx_t, idx_c, pscore)
 
         if self.estimand == 'ATT':
-            matches, weights = self._nn_match_from_dist(dist_mat, self.caliper)
+            matches, weights = self._nn_match_from_dist(
+                dist_mat,
+                self.caliper,
+                target_order=row_order[idx_t],
+                pool_order=row_order[idx_c],
+            )
             att = self._compute_effect(Y, idx_t, idx_c, X, matches, weights)
             se = self._ai_se(Y, X, T, idx_t, idx_c, matches, weights)
         else:
             # ATE: match both directions, reuse the same propensity scores
             dist_ct = self._compute_distance_matrix(X, idx_c, idx_t, pscore)
-            m_tc, w_tc = self._nn_match_from_dist(dist_mat, self.caliper)
-            m_ct, w_ct = self._nn_match_from_dist(dist_ct, self.caliper)
+            m_tc, w_tc = self._nn_match_from_dist(
+                dist_mat,
+                self.caliper,
+                target_order=row_order[idx_t],
+                pool_order=row_order[idx_c],
+            )
+            m_ct, w_ct = self._nn_match_from_dist(
+                dist_ct,
+                self.caliper,
+                target_order=row_order[idx_c],
+                pool_order=row_order[idx_t],
+            )
             att_part = self._compute_effect(Y, idx_t, idx_c, X, m_tc, w_tc)
             atc_part = self._compute_effect(Y, idx_c, idx_t, X, m_ct, w_ct)
             n_t, n_c = len(idx_t), len(idx_c)
@@ -427,6 +446,29 @@ def _fit_nearest(self, Y, X, T, idx_t, idx_c):
 
         return att, se, balance
 
+    @staticmethod
+    def _stable_index_order(index):
+        """Numeric rank of DataFrame index labels for deterministic tie-breaking."""
+        labels = np.asarray(pd.Index(index))
+        n = len(labels)
+        try:
+            order = np.argsort(labels, kind='mergesort')
+        except TypeError:
+            order = np.array(
+                sorted(
+                    range(n),
+                    key=lambda i: (
+                        type(labels[i]).__name__,
+                        repr(labels[i]),
+                        i,
+                    ),
+                ),
+                dtype=int,
+            )
+        ranks = np.empty(n, dtype=float)
+        ranks[order] = np.arange(n, dtype=float)
+        return ranks
+
     def _compute_distance_matrix(self, X, idx_from, idx_to, pscore=None):
         """Compute distance matrix between two groups."""
         X_from = X[idx_from]
@@ -713,7 +755,7 @@ def _logit_propensity(X, T, poly=1):
     # NN matching helpers
     # ==================================================================
 
-    def _nn_match_from_dist(self, dist, caliper=None):
+    def _nn_match_from_dist(self, dist, caliper=None, target_order=None, pool_order=None):
         """
         k-NN matching from a precomputed distance matrix.
 
@@ -722,20 +764,38 @@ def _nn_match_from_dist(self, dist, caliper=None):
         order of their minimum distance (best match first) so the greedy
         assignment favours the closest pairs.
 
+        Equal-distance ties are resolved by the source DataFrame index:
+        lower-index pool units are selected first, and without-replacement
+        target processing falls back to target index order. This makes
+        matching deterministic across BLAS/NumPy backends and independent of
+        incidental row order when index labels preserve unit identity.
+
         References: Cunningham (2021, Ch. 5) discusses with- vs.
         without-replacement matching and the bias–variance trade-off.
         """
         n_target = dist.shape[0]
         matches = [None] * n_target
         weights = [None] * n_target
+        if target_order is None:
+            target_order = np.arange(n_target, dtype=float)
+        if pool_order is None:
+            pool_order = np.arange(dist.shape[1], dtype=float)
+
+        def _nearest_indices(d, k):
+            finite = np.isfinite(d)
+            if not np.any(finite):
+                return np.array([], dtype=int)
+            candidates = np.where(finite)[0]
+            order = np.lexsort((pool_order[candidates], d[candidates]))
+            return candidates[order[:k]]
 
         # Without replacement: process treated units greedily by best
         # minimum distance so each control is used at most once.
         if not self.replace:
             used = set()
             # Sort treated units by their minimum distance to any control
             min_dists = np.min(dist, axis=1)
-            order = np.argsort(min_dists)
+            order = np.lexsort((target_order, min_dists))
 
             for i in order:
                 d = dist[i].copy()
@@ -751,7 +811,7 @@ def _nn_match_from_dist(self, dist, caliper=None):
                     weights[i] = np.array([])
                     continue
 
-                idx = np.argpartition(d, k)[:k]
+                idx = _nearest_indices(d, k)
                 matches[i] = idx
                 weights[i] = np.ones(k) / k
                 used.update(idx.tolist())
@@ -770,7 +830,7 @@ def _nn_match_from_dist(self, dist, caliper=None):
                 weights[i] = np.array([])
                 continue
 
-            idx = np.argpartition(d, k)[:k]
+            idx = _nearest_indices(d, k)
             matches[i] = idx
             weights[i] = np.ones(k) / k
 

tests/test_matching.py

@@ -172,6 +172,33 @@ def test_basic(self, selection_bias_data):
         assert abs(result.estimate - 2.0) < 1.5
 
 
+class TestNearestTieBreaking:
+    """Nearest-neighbor ties should be deterministic and index-anchored."""
+
+    def test_equal_distance_ties_use_source_index_not_row_order(self):
+        df = pd.DataFrame(
+            {
+                'unit': ['treated_0', 'control_high', 'control_low',
+                         'treated_1', 'control_1'],
+                'y': [10.0, 100.0, 0.0, 30.0, 25.0],
+                'treat': [1, 0, 0, 1, 0],
+                'x': [0.0, 0.0, 0.0, 1.0, 1.0],
+            },
+            # control_low and control_high are exact ties for treated_0.
+            # The lower source index must win regardless of row order.
+            index=[100, 20, 5, 200, 30],
+        )
+        shuffled = df.loc[[20, 100, 30, 5, 200]]
+
+        r1 = match(df, y='y', treat='treat', covariates=['x'],
+                   distance='euclidean', method='nearest')
+        r2 = match(shuffled, y='y', treat='treat', covariates=['x'],
+                   distance='euclidean', method='nearest')
+
+        assert r1.estimate == pytest.approx(7.5, abs=1e-12)
+        assert r2.estimate == pytest.approx(7.5, abs=1e-12)
+
+
 class TestExactMatching:
     """distance='exact'."""
 

tests/test_tierD_lalonde_psm_guard.py

@@ -2,19 +2,17 @@
 
 Part of the P1 campaign (see ``.tierd_campaign/CAMPAIGN.md``). The
 ``sp.replicate('lalonde_1986')`` registry ships *golden numbers* for the classic
-track, but no test guarded the live 1:1 NN propensity-score-matching ATT — so a
-tie-handling change in ``sp.match`` silently drifted it from the originally
-pinned $2,012.5 to the current deterministic $1,963.4 (a 2.5% move) without any
-red flag. This test is that missing guard: it pins all three classic LaLonde
-numbers to their current deterministic values on the bundled real data, so any
-future drift fails loudly.
+track, but no test guarded the live 1:1 NN propensity-score-matching ATT. This
+test is that missing guard: it pins all three classic LaLonde numbers to their
+current deterministic values on the bundled real data, including the stable
+index-anchored nearest-neighbor tie-break in ``sp.match``.
 
 The naive and covariate-adjusted OLS values reproduce R ``MatchIt`` to the
-dollar; the PSM value is the current deterministic ``sp.match(method='nearest')``
-output (sensitive to tie-breaking on the binary covariates, hence the value the
-registry pin was refreshed to).
+dollar; the PSM value is the deterministic ``sp.match(method='nearest')`` output
+after exact-distance ties on the binary covariates are resolved by source index.
 
-Purely additive — no estimator numerics changed (campaign red line).
+This guard accompanies the nearest-neighbor tie-break fix; only exact-tie
+matching designs can move.
 """
 
 import pytest
@@ -52,16 +50,7 @@ def test_psm_att_is_deterministic_and_pinned(lalonde):
         for _ in range(3)
     ]
     assert len(set(round(v, 6) for v in vals)) == 1  # deterministic within a run
-    # The exact PSM ATT depends on how ties on the binary covariates are broken,
-    # which is linear-algebra-backend dependent and deterministic only *within*
-    # an environment: tight BLAS builds land at $1963.4, the GitHub
-    # ubuntu-latest OpenBLAS build at $1813.4. We keep the strict within-run
-    # determinism check (the env-independent anti-regression guard) and bound a
-    # band that still catches any real algorithm change (matching breaking ->
-    # ~0, negative, or wildly off) while tolerating the documented cross-backend
-    # tie-break range. Making sp.match tie-breaks backend-deterministic is
-    # tracked as a separate Tier-D follow-up (see .tierd_campaign/CAMPAIGN.md).
-    assert 1750.0 <= vals[0] <= 2050.0, f"PSM ATT {vals[0]:.1f} outside tie-break band"
+    assert vals[0] == pytest.approx(1963.43, abs=0.1)
 
 
 def test_psm_recovers_experimental_benchmark(lalonde):