Elegant way to index 3D array and keep non-zero axis using Numpy

Question

Say I have the following 3D array:

L=np.arange(18).reshape((2,3,3))
L[:,:,1] = 0; L[:,[0,1],:] = 0

In []: L
Out[]: 
array([[[ 0,  0,  0],
        [ 0,  0,  0],
        [ 6,  0,  8]],

       [[ 0,  0,  0],
        [ 0,  0,  0],
        [15,  0,  17]]])

where zero columns in L[0,:] are always matched by corresponding zero columns in L[1,:] .

I want to now remove the middle columns where the sum along the axis equals 0 (ignoring rows of zero. My current clumsy approach is

l=np.nonzero(L.sum(axis=1))[1]

In []: L[:,:,l[:len(l)/2]]
Out[]: 
array([[[ 0,  0],
        [ 0,  0],
        [ 6,  8]],

       [[ 0,  0],
        [ 0,  0],
        [15, 17]]])

What is a less roundabout way of doing this?

Answer 1

We can look for all zeros along the first two axes and use that for masking out those from the third axis -

L[:,:,~(L==0).all(axis=(0,1))]

Alternatively, using any() to replace ~all() -

L[:,:,(L!=0).any(axis=(0,1))]

We can use the ellipsis notation ... to replace :,: and also skip the arg axis to give us a compact version -

L[...,~(L==0).all((0,1))]
L[...,(L!=0).any((0,1))]

More on how ellipsis works for NumPy arrays, here .

---

For the sum part of the question, it would be similar -

L[...,L.sum((0,1))!=0]