删除多维数组中的 nans

Question

我有一个多维数组。示例（二维）：

x = np.array([[     1.,      1.,  np.nan,  np.nan],
              [     2.,  np.nan,      2.,  np.nan],
              [ np.nan,      3.,  np.nan,  np.nan]])

是否有一种简单、有效的方法可以沿轴“压缩”/“挤压”/“推出”nans？我的意思是，这样输出（这里：axis=0）就会变成：

np.array([[  1.,  1.,  np.nan,  np.nan],
          [  2.,  3.,      2.,  np.nan]])

还应该适用于超过 2 个维度。

Answer 1

你可以在非nan元素的掩码上使用argsort；使用稳定的排序算法（如合并排序）来保留非nan 元素的原始顺序：

mask = np.isnan(x)
cut = np.min(np.count_nonzero(mask, axis=0))
x[np.argsort(~mask, axis=0, kind='mergesort')[cut:], np.arange(x.shape[1])]

输出：

array([[  1.,   1.,  nan,  nan],
       [  2.,   3.,   2.,  nan]])

ND 版本：

import numpy as np

def nan_bouncer(x, axis=0):
    if axis != 0:
        x = np.moveaxis(x, axis, 0)
    mask = np.isnan(x)
    cut = np.min(np.count_nonzero(mask, axis=0))
    idx = tuple(np.ogrid[tuple(map(slice, x.shape[1:]))])
    res = x[(np.argsort(~mask, axis=0, kind='mergesort')[cut:],) + idx] 
    return res if axis == 0 else np.moveaxis(res, 0, axis)

#demo
data = np.random.randint(0, 3, (3, 4, 4)).astype(float)
data /= data / data

print(data)
print(nan_bouncer(data))
print(nan_bouncer(data, 2))

样本输出：

[[[ nan   1.   2.   1.]
  [  2.  nan  nan   2.]
  [  2.   1.   1.   2.]
  [  1.   1.   2.  nan]]

 [[ nan  nan   2.   1.]
  [  2.   2.  nan   1.]
  [  2.   2.   2.   2.]
  [  2.   2.  nan   1.]]

 [[  1.   1.  nan  nan]
  [  1.   1.   2.   1.]
  [  2.  nan   2.   1.]
  [  1.   1.   1.   2.]]]


[[[ nan  nan  nan  nan]
  [  2.  nan  nan   2.]
  [  2.  nan   1.   2.]
  [  1.   1.  nan  nan]]

 [[ nan   1.   2.   1.]
  [  2.   2.  nan   1.]
  [  2.   1.   2.   2.]
  [  2.   2.   2.   1.]]

 [[  1.   1.   2.   1.]
  [  1.   1.   2.   1.]
  [  2.   2.   2.   1.]
  [  1.   1.   1.   2.]]]


[[[ nan   1.   2.   1.]
  [ nan  nan   2.   2.]
  [  2.   1.   1.   2.]
  [ nan   1.   1.   2.]]

 [[ nan  nan   2.   1.]
  [ nan   2.   2.   1.]
  [  2.   2.   2.   2.]
  [ nan   2.   2.   1.]]

 [[ nan  nan   1.   1.]
  [  1.   1.   2.   1.]
  [ nan   2.   2.   1.]
  [  1.   1.   1.   2.]]]