-
-
Notifications
You must be signed in to change notification settings - Fork 10.9k
MAINT: Speed up numpy.nonzero. #18368
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.
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
Changes from 1 commit
c404910
a2b2a35
07dc717
9b5c1c4
5c79ab3
0015523
36801dc
399d1b9
8435e97
2dbe866
aa422cb
a302d4a
85ae692
3480e1e
3582137
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
…mensions
- Loading branch information
There are no files selected for viewing
Original file line number | Diff line number | Diff line change | ||
---|---|---|---|---|
|
@@ -1285,6 +1285,14 @@ def test_nonzero_onedim(self): | |||
idxs = np.nonzero(x)[0] | ||||
assert_equal(np.array_equal(np.where(x != 0)[0], idxs), True) | ||||
|
||||
for _ in range(iters): # Check for slices | ||||
for dtype in types: | ||||
x = ((2**33)*np.random.randn(100)).astype(dtype) | ||||
x_view = x[:50] | ||||
x_view[[0,2,8,22,24]] = 0 | ||||
idxs = np.nonzero(x_view)[0] | ||||
assert_equal(np.array_equal(np.where(x_view != 0)[0], idxs), True) | ||||
|
||||
|
||||
def test_nonzero_twodim(self): | ||||
x = np.array([[0, 1, 0], [2, 0, 3]]) | ||||
|
@@ -1312,6 +1320,8 @@ def test_nonzero_twodim(self): | |||
types = [bool, np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, np.uint64, np.float32, np.float64] | ||||
idxs = np.arange(100) | ||||
iters = 10 | ||||
dim_permutation = np.arange(2) # Shuffle dimensions | ||||
|
||||
for _ in range(iters): | ||||
for dtype in types: | ||||
x = ((2**33)*np.random.randn(10, 10)).astype(dtype) | ||||
|
@@ -1327,10 +1337,90 @@ def test_nonzero_twodim(self): | |||
assert_equal(nzs, 0) | ||||
|
||||
|
||||
for _ in range(iters): # This time randomly permute dimensions | ||||
for dtype in types: | ||||
x = ((2**33)*np.random.randn(10, 10)).astype(dtype) | ||||
np.random.shuffle(idxs) | ||||
num_zeros = np.random.randint(0, 101) | ||||
x.flat[idxs[:num_zeros]] = 0 | ||||
np.random.shuffle(dim_permutation) | ||||
x = np.transpose(x, dim_permutation) | ||||
idxs_0, idxs_1 = np.nonzero(x) | ||||
assert_equal(np.count_nonzero(x), len(idxs_0)) | ||||
nzs = 0 | ||||
for i,j in zip(idxs_0, idxs_1): | ||||
if x[i,j] == 0: | ||||
nzs += 1 | ||||
assert_equal(nzs, 0) | ||||
|
||||
# Repeat the above logic for fortran contiguous arrays | ||||
for _ in range(iters): | ||||
for dtype in types: | ||||
x = ((2**33)*np.random.randn(10, 10)).astype(dtype) | ||||
x = np.asfortranarray(x) | ||||
np.random.shuffle(idxs) | ||||
num_zeros = np.random.randint(0, 101) | ||||
x.flat[idxs[:num_zeros]] = 0 | ||||
idxs_0, idxs_1 = np.nonzero(x) | ||||
assert_equal(np.count_nonzero(x), len(idxs_0)) | ||||
nzs = 0 | ||||
for i,j in zip(idxs_0, idxs_1): | ||||
if x[i,j] == 0: | ||||
nzs += 1 | ||||
assert_equal(nzs, 0) | ||||
|
||||
|
||||
for _ in range(iters): # This time randomly permute dimensions | ||||
for dtype in types: | ||||
x = ((2**33)*np.random.randn(10, 10)).astype(dtype) | ||||
x = np.asfortranarray(x) | ||||
np.random.shuffle(idxs) | ||||
num_zeros = np.random.randint(0, 101) | ||||
x.flat[idxs[:num_zeros]] = 0 | ||||
np.random.shuffle(dim_permutation) | ||||
x = np.transpose(x, dim_permutation) | ||||
idxs_0, idxs_1 = np.nonzero(x) | ||||
assert_equal(np.count_nonzero(x), len(idxs_0)) | ||||
nzs = 0 | ||||
for i,j in zip(idxs_0, idxs_1): | ||||
if x[i,j] == 0: | ||||
nzs += 1 | ||||
assert_equal(nzs, 0) | ||||
|
||||
|
||||
for i in range(iters): # check for slices | ||||
x = ((2**33)*np.random.randn(10, 10)).astype(np.int32) | ||||
x.flat[[2,6,9,15,21]] = 0 | ||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||
x_view = x[:5,:5] | ||||
|
||||
idxs_0, idxs_1 = np.nonzero(x_view) | ||||
assert_equal(np.count_nonzero(x_view), len(idxs_0)) | ||||
nzs = 0 | ||||
for i,j in zip(idxs_0, idxs_1): | ||||
if x_view[i,j] == 0: | ||||
nzs += 1 | ||||
assert_equal(nzs, 0) | ||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. You can just use |
||||
|
||||
y = np.asfortranarray(x) | ||||
y_view = y[:5,:5] | ||||
|
||||
idxs_0, idxs_1 = np.nonzero(y_view) | ||||
assert_equal(np.count_nonzero(y_view), len(idxs_0)) | ||||
nzs = 0 | ||||
for i,j in zip(idxs_0, idxs_1): | ||||
if y_view[i,j] == 0: | ||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||
nzs += 1 | ||||
assert_equal(nzs, 0) | ||||
|
||||
|
||||
|
||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Excessive white space. There should be one empty line between |
||||
|
||||
def test_nonzero_threedim(self): | ||||
types = [bool, np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, np.uint64, np.float32, np.float64] | ||||
idxs = np.arange(300) | ||||
iters = 10 | ||||
dim_permutation = np.arange(3) # Shuffle dimensions | ||||
|
||||
for _ in range(iters): | ||||
for dtype in types: | ||||
x = ((2**33)*np.random.randn(10, 10, 3)).astype(dtype) | ||||
|
@@ -1346,6 +1436,83 @@ def test_nonzero_threedim(self): | |||
assert_equal(nzs, 0) | ||||
|
||||
|
||||
for _ in range(iters): # This time randomly permute dimensions | ||||
for dtype in types: | ||||
x = ((2**33)*np.random.randn(10, 10, 3)).astype(dtype) | ||||
np.random.shuffle(idxs) | ||||
num_zeros = np.random.randint(0, 301) | ||||
x.flat[idxs[:num_zeros]] = 0 | ||||
np.random.shuffle(dim_permutation) | ||||
x = np.transpose(x, dim_permutation) | ||||
idxs_0, idxs_1, idxs_2 = np.nonzero(x) | ||||
assert_equal(np.count_nonzero(x), len(idxs_0)) | ||||
nzs = 0 | ||||
for i,j,k in zip(idxs_0, idxs_1, idxs_2): | ||||
if x[i,j,k] == 0: | ||||
nzs += 1 | ||||
assert_equal(nzs, 0) | ||||
|
||||
# Repeat the above for fortran contiguous arrays | ||||
for _ in range(iters): | ||||
for dtype in types: | ||||
x = ((2**33)*np.random.randn(10, 10, 3)).astype(dtype) | ||||
x = np.asfortranarray(x) | ||||
np.random.shuffle(idxs) | ||||
num_zeros = np.random.randint(0, 301) | ||||
x.flat[idxs[:num_zeros]] = 0 | ||||
idxs_0, idxs_1, idxs_2 = np.nonzero(x) | ||||
assert_equal(np.count_nonzero(x), len(idxs_0)) | ||||
nzs = 0 | ||||
for i,j,k in zip(idxs_0, idxs_1, idxs_2): | ||||
if x[i,j,k] == 0: | ||||
nzs += 1 | ||||
assert_equal(nzs, 0) | ||||
|
||||
|
||||
for _ in range(iters): # This time randomly permute dimensions | ||||
for dtype in types: | ||||
x = ((2**33)*np.random.randn(10, 10, 3)).astype(dtype) | ||||
x = np.asfortranarray(x) | ||||
np.random.shuffle(idxs) | ||||
num_zeros = np.random.randint(0, 301) | ||||
x.flat[idxs[:num_zeros]] = 0 | ||||
np.random.shuffle(dim_permutation) | ||||
x = np.transpose(x, dim_permutation) | ||||
idxs_0, idxs_1, idxs_2 = np.nonzero(x) | ||||
assert_equal(np.count_nonzero(x), len(idxs_0)) | ||||
nzs = 0 | ||||
for i,j,k in zip(idxs_0, idxs_1, idxs_2): | ||||
if x[i,j,k] == 0: | ||||
nzs += 1 | ||||
assert_equal(nzs, 0) | ||||
|
||||
|
||||
for i in range(iters): # Check for slices | ||||
x = ((2**33)*np.random.randn(10, 10)).astype(np.int32) | ||||
x.flat[[2,6,9,15,21]] = 0 | ||||
x_view = x[:5,:5] | ||||
|
||||
idxs_0, idxs_1 = np.nonzero(x_view) | ||||
assert_equal(np.count_nonzero(x_view), len(idxs_0)) | ||||
nzs = 0 | ||||
for i,j in zip(idxs_0, idxs_1): | ||||
if x_view[i,j] == 0: | ||||
nzs += 1 | ||||
assert_equal(nzs, 0) | ||||
|
||||
y = np.asfortranarray(x) | ||||
y_view = y[:5,:5] | ||||
|
||||
idxs_0, idxs_1 = np.nonzero(y_view) | ||||
assert_equal(np.count_nonzero(y_view), len(idxs_0)) | ||||
nzs = 0 | ||||
for i,j in zip(idxs_0, idxs_1): | ||||
if y_view[i,j] == 0: | ||||
nzs += 1 | ||||
assert_equal(nzs, 0) | ||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. And one more comment that got losts. These tests seem very repitetive, can you try to use |
||||
|
||||
|
||||
|
||||
def test_sparse(self): | ||||
# test special sparse condition boolean code path | ||||
for i in range(20): | ||||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Its fine to implement nonzero in Python here, but please refactor it. You can use
np.nditer
if that helps. Or flatten andnp.unravel_index
(andnp.ravel()
) if the n dimensions are annoying to deal with.