使用遮罩沿numpy 4D數組刪除全零的空元素

Question

給定一個示例numpy數組，如下所示：

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)

考慮到第一個元素全為零，如何獲取形狀為（3,2,3,3）的numpy數組？ 我的數據集是形狀（m，x，y，z）中較大的一個，我需要返回非零（mn，x，y，z）數組，其中n是（x，y，z）形狀的數組與全零。

到目前為止，我已經嘗試過了：

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]]]]
'''

但是應用a[~mask]給我一個扁平的數組：

[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,)

我需要的是這樣的：

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]]]])

獎勵：我需要將其應用於單獨的/鏡像（m，x，y，z）形狀的數組，所以也許我需要使用屏蔽方法？

Answer 1

使用除第一個軸以外的all軸來創建用於索引的布爾數組：

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]]]])