Fast condition checkup iterating on multidimensional numpy array

Question

I have a big ndimensional array. I want to iterate on it to check whether a condition is locally satisfied. Next snippet explains my problem.

a = np.random.randint(2, size=(60,80,3,3))

test = np.array([[1,0,0],[0,1,0],[0,0,0]])

for i in xrange(a.shape[0]):
    for j in xrange(b.shape[1]):
        if (a[i,j] == test).all():
            # Do something with indices i and j

The code is obviously very slow. I tried using numpy.where but it's not working as it looks for equality at each of the four indices.

EDIT: I do also need to store the indices (i,j) which satisfy the condition

Phillip · Accepted Answer · 2014-10-20 13:59:48Z

1

np.apply_over_axes(np.prod, a == test, [3,2]) == 1

gives you an array of size (60,80,1,1) which is True whereever the condition holds. A shorter, preferrable version found by the thread starter is

(a == test).all(axis=(2,3))

Both are equivalent, but the latter avoids the boolean → integer → boolean conversion. Use np.where on that array to get the indices (i, j).

edited Oct 20, 2014 at 13:59

answered Oct 20, 2014 at 10:45

Phillip

13.7k33 silver badges43 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

fmonegaglia Over a year ago

Wow, seems to work. Do you think it is quite the same as np.where(a == test).all(axis=(2,3)), 1, 0 ? I found this to work either, are at least looks like to.

Phillip Over a year ago

You mean np.where(a == test).all(axis=(2,3))? Yes, that's equivalent. And looks much better.

Collectives™ on Stack Overflow

Fast condition checkup iterating on multidimensional numpy array

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related