match: expose matched index pairs as NearestNeighborMatch.matched_indexes_ (#621)

causalml/match.py

@@ -97,6 +97,15 @@ class NearestNeighborMatch:
             seed
         n_jobs (int): The number of parallel jobs to run for neighbors search.
             None means 1 unless in a joblib.parallel_backend context. -1 means using all processors
+
+    Fitted attributes (populated after :meth:`match` or :meth:`match_by_group`):
+        matched_indexes_ (pandas.DataFrame): two-column dataframe with the
+            (from, to) pairs of original data indices produced by the last
+            call. ``from`` corresponds to the matching-source group
+            (treatment if ``treatment_to_control`` else control); ``to``
+            is the matched counterpart from the opposite group. Useful for
+            joining matched pairs back to upstream metadata or auditing the
+            matching outcome (see uber/causalml#621).
     """
 
     def __init__(
@@ -179,10 +188,15 @@ def match(self, data, treatment_col, score_cols):
             match_from_scaled = pd.concat([match_from_scaled] * self.ratio, axis=0)
 
             cond = (distances / np.sqrt(len(score_cols))) < sdcal
+            # Capture the full (from, to) pair mapping before deduplicating
+            # the from-side, so ``matched_indexes_`` (set below) can expose
+            # every matched pair — not just the unique from-indices.
+            from_pairs = np.array(match_from_scaled.loc[cond].index)
+            to_pairs = np.array(match_to_scaled.iloc[indices[cond]].index)
             # Deduplicate the indices of the treatment group
-            from_idx_matched = np.unique(match_from_scaled.loc[cond].index)
+            from_idx_matched = np.unique(from_pairs)
             # XXX: Should we deduplicate the indices of the control group too?
-            to_idx_matched = np.array(match_to_scaled.iloc[indices[cond]].index)
+            to_idx_matched = to_pairs
         else:
             assert len(score_cols) == 1, (
                 "Matching on multiple columns is only supported using the "
@@ -199,6 +213,13 @@ def match(self, data, treatment_col, score_cols):
 
             from_idx_matched = []
             to_idx_matched = []
+            # Track each accepted (from, to) pair so ``matched_indexes_`` can
+            # report the actual row-aligned pairing. `from_idx_matched` only
+            # holds the first time a from-index is accepted, so it cannot
+            # be repeated `ratio` times to recover pairs (the loop may
+            # accept 1..ratio matches per from depending on the caliper).
+            from_pairs_list = []
+            to_pairs_list = []
             match_to["unmatched"] = True
 
             for from_idx in from_indices:
@@ -214,8 +235,22 @@ def match(self, data, treatment_col, score_cols):
                         if i == 0:
                             from_idx_matched.append(from_idx)
                         to_idx_matched.append(to_idx)
+                        from_pairs_list.append(from_idx)
+                        to_pairs_list.append(to_idx)
                         match_to.loc[to_idx, "unmatched"] = False
 
+        # Persist the (from_idx, to_idx) pair mapping so callers can join the
+        # matched pairs back onto the original data without re-running the
+        # matching. Each row corresponds to one matched pair; with ratio>1 a
+        # single ``from`` index can appear multiple times against distinct
+        # ``to`` indices. See uber/causalml#621.
+        if self.replace:
+            self.matched_indexes_ = pd.DataFrame({"from": from_pairs, "to": to_pairs})
+        else:
+            self.matched_indexes_ = pd.DataFrame(
+                {"from": np.array(from_pairs_list), "to": np.array(to_pairs_list)}
+            )
+
         return data.loc[
             np.concatenate([np.array(from_idx_matched), np.array(to_idx_matched)])
         ]

tests/test_match.py

@@ -78,6 +78,59 @@ def test_nearest_neighbor_match_control_to_treatment(generate_unmatched_data):
     assert 2 * sum(matched[TREATMENT_COL] == 0) == sum(matched[TREATMENT_COL] != 0)
 
 
+def test_nearest_neighbor_match_exposes_matched_indexes(generate_unmatched_data):
+    """Regression test for uber/causalml#621.
+
+    The user requested access to the matched index pairs as an attribute.
+    Previously the (from, to) pairs were computed inside ``match`` and
+    discarded; only the joined dataframe was returned. This test asserts:
+
+    1. ``matched_indexes_`` is set after ``match()`` runs.
+    2. It has the documented schema (``from``, ``to``).
+    3. The exposed pairs are consistent with the returned matched dataframe
+       — every index in ``matched_indexes_['from']`` and ``['to']`` shows up
+       in the matched dataframe's index.
+    4. With ``replace=True, ratio=2``, a single ``from`` index can appear
+       multiple times (one row per matched control), proving the attribute
+       captures the *pair* mapping rather than the deduplicated set.
+    """
+    df, features = generate_unmatched_data()
+
+    # Replacement path with ratio=2 — exercises the NearestNeighbors branch.
+    psm = NearestNeighborMatch(replace=True, ratio=2, random_state=RANDOM_SEED)
+    matched = psm.match(data=df, treatment_col=TREATMENT_COL, score_cols=[SCORE_COL])
+
+    assert hasattr(
+        psm, "matched_indexes_"
+    ), "matched_indexes_ must be set after match() — see uber/causalml#621"
+    assert isinstance(psm.matched_indexes_, pd.DataFrame)
+    assert set(psm.matched_indexes_.columns) == {"from", "to"}
+    assert len(psm.matched_indexes_) > 0
+
+    # The matched dataframe is the union of unique from + all to indices, so
+    # every from / to in the pair table must appear in the matched index.
+    matched_idx = set(matched.index)
+    assert set(psm.matched_indexes_["from"].unique()).issubset(matched_idx)
+    assert set(psm.matched_indexes_["to"].unique()).issubset(matched_idx)
+
+    # ratio=2 with replace=True: at least one `from` should have two `to`
+    # matches, proving the attribute captures the pair mapping (not just
+    # the deduplicated unique-from set).
+    counts_per_from = psm.matched_indexes_["from"].value_counts()
+    assert (counts_per_from >= 2).any(), (
+        "With replace=True, ratio=2 the pair table must show at least one "
+        "from-index paired against multiple to-indices"
+    )
+
+    # No-replacement path is also exercised by the existing tests; here we
+    # just confirm the attribute is populated for that path too.
+    psm2 = NearestNeighborMatch(replace=False, ratio=1, random_state=RANDOM_SEED)
+    psm2.match(data=df, treatment_col=TREATMENT_COL, score_cols=[SCORE_COL])
+    assert hasattr(psm2, "matched_indexes_")
+    assert set(psm2.matched_indexes_.columns) == {"from", "to"}
+    assert len(psm2.matched_indexes_) > 0
+
+
 def test_match_optimizer(generate_unmatched_data):
     df, features = generate_unmatched_data()