1

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?

neither-nor
  • 1,245
  • 2
  • 17
  • 30
  • 1
    What if the first col of the first 2D slice were all zeros. Would the output change? – Divakar Sep 01 '17 at 21:31
  • 1
    @Divakar Sorry about the confusion. The output in your example technically shouldn't change, although the specific kind of arrays I'm working with wouldn't look like that: if the first column of [0,0,6] is set to all 0s, then the column that has [0,0,15] would also have all 0s. – neither-nor Sep 01 '17 at 22:31
  • 1
    I see. So, similar to the posted solution : `L[...,L.sum((0,1))!=0]`. – Divakar Sep 01 '17 at 22:33

1 Answers1

4

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]
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • Interesting! What do the `...` and `(0,1)` mean, respectively? – neither-nor Sep 01 '17 at 21:12
  • Cool. If I'm only interested in removing where the sum along axis=1 equals 0 (i.e., don't care about rows of zero), would the solution change? – neither-nor Sep 01 '17 at 21:20
  • 1
    @neither-nor Add a sample case for that one? – Divakar Sep 01 '17 at 21:22
  • 1
    @neither-nor: This code isn't looking for _rows_ or _columns_ of zero - it's looking for 2D slices of zeros, which is probably what you want. Of course, it's not taking advantage of your constraint on where zeros can actually be. – Eric Sep 02 '17 at 06:12