Skip to content

Make numpy.searchsorted match NumPy when first argument is unsorted #7005

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sklam merged 3 commits into from
Jun 9, 2022
Merged

Make numpy.searchsorted match NumPy when first argument is unsorted #7005

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
55cb103
Make numpy.searchsorted match NumPy when first argument is unsorted
brandonwillard May 6, 2021
a6bff25
Simplify searchsorted unordered test values
brandonwillard May 10, 2021
24fe303
Added another unordered searchsorted test with a larger array
brandonwillard Jun 3, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
Make numpy.searchsorted match NumPy when first argument is unsorted
  • Loading branch information
@brandonwillard
brandonwillard committed Jun 3, 2022
commit 55cb103714bec562e4a76b78e30cd5a52f0885d6
54 changes: 46 additions & 8 deletions numba/np/arraymath.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3740,17 +3740,43 @@ def bincount_impl(a, weights=None, minlength=0):


def _searchsorted(func):
def searchsorted_inner(a, v):
n = len(a)
def searchsorted_inner(a, v, v_last, lo, hi, n):
Comment thread
brandonwillard marked this conversation as resolved.
Resolved
"""Perform inner loop of searchsorted (i.e. a binary search).

This is loosely based on the NumPy implementation in [1]_.

Parameters
----------
a: 1-D array_like
The input array.
v: array_like
The current value to insert into `a`.
v_last: array_like
The previous value inserted into `a`.
lo: int
The initial/previous "low" value of the binary search.
hi: int
The initial/previous "high" value of the binary search.
n: int
The length of `a`.


.. [1] https://github.com/numpy/numpy/blob/809e8d26b03f549fd0b812a17b8a166bcd966889/numpy/core/src/npysort/binsearch.cpp#L173
""" # noqa: E501
if np.isnan(v):
# Find the first nan (i.e. the last from the end of a,
# since there shouldn't be many of them in practice)
for i in range(n, 0, -1):
if not np.isnan(a[i - 1]):
return i
return 0
lo = 0
hi = n

if v_last < v:
hi = n
else:
lo = 0
hi = hi + 1 if hi < n else n

@stuartarchibald stuartarchibald May 10, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The tests don't appear to exercise the true branch?

@brandonwillard brandonwillard May 10, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All the tests, or just the unordered one?

@brandonwillard brandonwillard May 10, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I threw in some print statements and all the branches appear to be visited under the current tests.

@esc esc Jun 2, 2022

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I manually checked the coverage of the branches and both are covered by the unordered tests.

@esc esc Jun 2, 2022

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I found the branches a bit confusing. From what I understand the values of lo and hi are initialized anyway in the searchsorted_imp and not v_last < v is only needed when the array to insert into isn't sorted. When reading the code it seems odd that one branch initializes only hi and the other both hi and lo. Perhaps, in the interest of comprehensibility write the following instead:

        if not v_last < v:
            lo = 0
            hi = hi + 1 if hi < n else n

And perhaps add a comment to describe why this is needed?

@esc esc Jun 3, 2022

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I just saw the numpy implementation, it seems like this code matches that implementation, so probably best to leave as is.


while hi > lo:
mid = (lo + hi) >> 1
if func(a[mid], (v)):
Expand Down Expand Up @@ -3782,24 +3808,36 @@ def searchsorted(a, v, side='left'):
if isinstance(v, types.Array):
# N-d array and output
def searchsorted_impl(a, v, side='left'):
n = len(a)
lo = 0
hi = n
out = np.empty(v.shape, np.intp)
v_last = v.flat[0]
for view, outview in np.nditer((v, out)):
index = loop_impl(a, view.item())
outview.itemset(index)
lo = loop_impl(a, view.item(), v_last, lo, hi, n)
v_last = view.item()
outview.itemset(lo)
return out

elif isinstance(v, types.Sequence):
# 1-d sequence and output
def searchsorted_impl(a, v, side='left'):
n = len(a)
lo = 0
hi = n
out = np.empty(len(v), np.intp)
v_last = v[0]
for i in range(len(v)):
out[i] = loop_impl(a, v[i])
lo = loop_impl(a, v[i], v_last, lo, hi, n)
out[i] = lo
v_last = v[i]
return out
else:
# Scalar value and output
# Note: NaNs come last in Numpy-sorted arrays
def searchsorted_impl(a, v, side='left'):
return loop_impl(a, v)
n = len(a)
return loop_impl(a, v, v, 0, n, n)

return searchsorted_impl

Expand Down
11 changes: 11 additions & 0 deletions numba/tests/test_np_functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1097,6 +1097,17 @@ def nonconst_side(a, v, side='left'):
with self.assertTypingError():
cfunc([1,2], 1, side='right')

# Test unordered values
a = np.array([0.29769574, 0.71649186, 0.20475563])
v = np.array(
[
[0.18847123, 0.39659508],
[0.56220006, 0.57428752],
[0.86720994, 0.44522637],
]
)
check(a, v)

def test_digitize(self):
pyfunc = digitize
cfunc = jit(nopython=True)(pyfunc)
Expand Down