Get only non-zero sub arrays from a N dimensional Numpy array

Question

I have a numpy array 'arr' that is of shape (1756020, 28, 28, 4). Basically 'arr' has 1756020 small arrays of shape (28,28,4). Out of the 1756020 arrays 967210 are 'all zero' and 788810 has all non-zero values. I want to remove all the 967210 'all zero' small arrays. I wrote a if else loop using the condition arr[i]==0.any() but it takes a lot of time. Is there a better way to do it?

Answer 1

One way to vectorise your logic is to use numpy.any with a tuple argument for axis containing non-tested dimensions.

# set up 4d array of ones
A = np.ones((5, 3, 3, 4))

# make second of shape (3, 3, 4) = 0
A[1] = 0  # or A[1, ...] = 0; or A[1, :, :, :] = 0

# find out which are non-zero
res = np.any(A, axis=(1, 2, 3))

print(res)

[True False True True True]

This feature is available in numpy v0.17 upwards. As per the docs:

axis : None or int or tuple of ints, optional

If this is a tuple of ints, a reduction is performed on multiple axes, instead of a single axis or all the axes as before.