Fix NearestNeighborMatch crash when the control pool runs below ratio#934
Merged
jeongyoonlee merged 1 commit intoJul 8, 2026
Merged
Conversation
match() with replace=False called np.argpartition(dist, self.ratio) on the pool of still-unmatched controls. As controls get matched and removed from the pool, that pool shrinks, and once it holds ratio or fewer units argpartition raised "kth out of bounds". A balanced 1:1 match hits this on the very last treatment unit, so replace=False matching failed on balanced (and control-scarce) datasets. Cap the partition index at the last available position and stop early when no unmatched controls remain. The abundant-control path is unchanged; a regression test covers the balanced 1:1 and control-scarce cases. Signed-off-by: Arpit Jain <arpitjain099@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
NearestNeighborMatch.match()withreplace=Falsecrashes on datasets where the pool of still-unmatched controls shrinks toratioor fewer units. It callsnp.argpartition(dist, self.ratio)on the remaining controls, and once that pool holdsratioor fewer elements argpartition raisesValueError: kth(=...) out of bounds. A balanced 1:1 dataset hits this on the very last treatment unit, so plain 1:1 matching without replacement fails on any balanced (or control-scarce) input.The fix caps the partition index at the last available position and stops early once no unmatched controls remain. The abundant-control path is untouched, so existing behavior and the current tests are unchanged.
Test plan
tests/test_match.pycovering the balanced 1:1 case (previously raisedkth out of bounds) and the control-scarce case.pytest tests/test_match.pypasses locally.