CODEWITHSUNDEEP
pythonnumpynull
weiji14
weiji14

Reputation: 517

Remove empty elements with all zeros along a numpy 4D array using mask

Given a sample numpy array like so:

a = np.array([[[[0,0,0], [0,0,0], [0,0,0]],
               [[0,0,0], [0,0,0], [0,0,0]]],
              [[[0,1,2], [1,1,1], [1,1,1]],
               [[1,1,1], [1,2,2], [1,1,1]]],
              [[[0,1,2], [1,1,1], [1,1,1]],
               [[1,1,1], [1,2,2], [1,1,1]]],
              [[[0,1,2], [1,1,1], [1,1,1]],
               [[1,1,1], [1,2,2], [1,1,1]]]])
#a.shape = (4, 2, 3, 3)

How can I get it to return a numpy array with shape (3,2,3,3) considering that the first element is all zeros? My dataset is a bigger one of shape (m, x, y, z) and I'll need to return non-zero (m-n, x,y,z) arrays where n are the (x,y,z) shaped arrays with all zeros.

So far I tried this:

mask = np.equal(a, np.zeros(shape=(2,3,3)))

'''
Returns:
        [[[[ True  True  True]
   [ True  True  True]
   [ True  True  True]]

  [[ True  True  True]
   [ True  True  True]
   [ True  True  True]]]


 [[[ True False False]
   [False False False]
   [False False False]]

  [[False False False]
   [False False False]
   [False False False]]]


 [[[ True False False]
   [False False False]
   [False False False]]

  [[False False False]
   [False False False]
   [False False False]]]


 [[[ True False False]
   [False False False]
   [False False False]]

  [[False False False]
   [False False False]
   [False False False]]]]
'''

But applying a[~mask] gives me a flattened array:

[1 2 1 1 1 1 1 1 1 1 1 1 2 2 1 1 1 1 2 1 1 1 1 1 1 1 1 1 1 2 2 1 1 1 1 2 1
 1 1 1 1 1 1 1 1 1 2 2 1 1 1] (51,)

What I need is something like this:

np.array([[[[0,1,2], [1,1,1], [1,1,1]],
           [[1,1,1], [1,2,2], [1,1,1]]],
          [[[0,1,2], [1,1,1], [1,1,1]],
           [[1,1,1], [1,2,2], [1,1,1]]],
          [[[0,1,2], [1,1,1], [1,1,1]],
           [[1,1,1], [1,2,2], [1,1,1]]]])

Bonus: I need to apply this to a separate/mirror (m, x, y, z) shaped array so maybe I'll need a masked approach?

Upvotes: 0

Views: 518

Answers (1)

akuiper
akuiper

Reputation: 214967

Use all over axises other than the first axis to create the boolean array for indexing:

a[~(a == 0).all(axis=(1,2,3))]

#array([[[[0, 1, 2],
#         [1, 1, 1],
#         [1, 1, 1]],

#        [[1, 1, 1],
#         [1, 2, 2],
#         [1, 1, 1]]],


#       [[[0, 1, 2],
#         [1, 1, 1],
#         [1, 1, 1]],

#        [[1, 1, 1],
#         [1, 2, 2],
#         [1, 1, 1]]],


#       [[[0, 1, 2],
#         [1, 1, 1],
#         [1, 1, 1]],

#        [[1, 1, 1],
#         [1, 2, 2],
#         [1, 1, 1]]]])

Upvotes: 2

Related Questions