From 8d812c011e42bcb2164344d29f5601529538abc3 Mon Sep 17 00:00:00 2001 From: jbbqqf Date: Sat, 9 May 2026 20:18:18 +0200 Subject: [PATCH] match: expose matched index pairs as NearestNeighborMatch.matched_indexes_ (#621) Capture the (from, to) pair mapping computed inside ``match()`` and expose it as a fitted attribute, so downstream callers can join matched pairs back onto upstream metadata and audit the matching outcome without re-running the algorithm. The attribute is a two-column DataFrame (``from``, ``to``) populated in both the replacement (NearestNeighbors) and no-replacement (loop) branches. Adds a regression test covering replace=True/ratio=2 and replace=False that fails on master with AttributeError and passes on this branch. --- causalml/match.py | 39 +++++++++++++++++++++++++++++++-- tests/test_match.py | 53 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 2 deletions(-) diff --git a/causalml/match.py b/causalml/match.py index 395a0777..8f96524c 100644 --- a/causalml/match.py +++ b/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)]) ] diff --git a/tests/test_match.py b/tests/test_match.py index e5785e44..f53b0e5a 100644 --- a/tests/test_match.py +++ b/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()