Python/Numpy - 被掩盖的数组非常慢

Question

有没有什么办法可以加快numpy中的掩码数组的速度？我之前有一个效率非常低的函数，我把它改写成使用掩码数组（这样我可以只掩盖某些行，而不是像之前那样复制和删除行）。但是，我很震惊地发现，使用掩码的函数竟然慢了10倍，因为掩码数组的速度实在太慢了。

举个例子，看看下面的代码（对我来说，掩码的速度慢了超过6倍）：

import timeit
import numpy as np
import numpy.ma as ma

def test(row):
   return row[0] + row[1]

a = np.arange(1000).reshape(500, 2)
t = timeit.Timer('np.apply_along_axis(test, 1, a)','from __main__ import test, a, np')
print round(t.timeit(100), 6)

b = ma.array(a)
t = timeit.Timer('ma.apply_along_axis(test, 1, b)','from __main__ import test, b, ma')
print round(t.timeit(100), 6)

Answer 1

常见解决方法

我知道的最有效的方法是手动处理掩码。这里有一个简单的基准测试，用来计算沿某个轴的 masked mean（掩码均值）。截至2021年（np.版本 1.19.2），手动实现的速度是原来的3倍。

值得注意的是，

• np.nanmean 的速度和 ma.mean 一样慢。不过，我没有找到简单的解决办法，因为 0 * nan -> nan，而且 np.where 的速度也比较慢。
• opencv 通常在其例程中有一个 掩码参数。不过，切换库在大多数情况下可能不太合适。

基准测试

benchmark manual (np.sum(..values..)/np.sum(..counts..))
    time for 100x np_mean: 0.15721

benchmark ma.mean
    time for 100x ma_mean: 0.580072

benchmark np.nanmean
    time for 100x nan_mean: 0.609166


np_mean[:5]: [0.74468436 0.75447124 0.75628326 0.74990387 0.74708414]
ma_mean[:5]: [0.7446843592460088 0.7544712410870448 0.7562832614361736
 0.7499038657880674 0.747084143818861]
nan_mean[:5]: [0.74468436 0.75447124 0.75628326 0.74990387 0.74708414]
np_mean == ma_mean:  True
np_mean == nan_mean:  True
np.__version__: 1.19.2

代码

import timeit
import numpy as np
import numpy.ma as ma

np.random.seed(0)

arr = np.random.rand(1000, 1000)
msk = arr > .5  # POSITIV mask: only emelemts > .5 are processed

print('\nbenchmark manual (np.sum(..values..)/np.sum(..counts..))')
np_mean = np.sum(arr * msk, axis=0)/np.sum(msk, axis=0)
t = timeit.Timer('np_mean = np.sum(arr * msk, axis=0)/np.sum(msk, axis=0)', globals=globals())
print('\ttime for 100x np_mean:', round(t.timeit(100), 6))

print('\nbenchmark ma.mean')
ma_arr = ma.masked_array(arr, mask=~msk)
ma_mean = ma.mean(ma_arr, axis=0)
t = timeit.Timer('ma_mean = ma.mean(ma_arr, axis=0)', globals=globals())
print('\ttime for 100x ma_mean:', round(t.timeit(100), 6))

print('\nbenchmark np.nanmean')
nan_arr = arr.copy()
nan_arr[~msk] = np.nan
nan_mean = np.nanmean(nan_arr, axis=0)
t = timeit.Timer('nan_mean = np.nanmean(nan_arr, axis=0)', globals=globals())
print('\ttime for 100x nan_mean:', round(t.timeit(100), 6))

print('\n')
print('np_mean[:5]:', np_mean[:5])
print('ma_mean[:5]:', ma_mean[:5])
print('nan_mean[:5]:', nan_mean[:5])
print('np_mean == ma_mean: ', (np_mean == ma_mean).all())
print('np_mean == nan_mean: ', (np_mean == nan_mean).all())

print('np.__version__:', np.__version__)

手动版本只在数组中没有 nans 的情况下有效。如果 arr 包含 nans：只需通过 msk = np.isnan(arr) 构造掩码，然后用 arr = np.nan_to_num(arr, copy=False, nan=0) 替换 arr 中的 nans。

Answer 2

我不知道为什么被屏蔽的数组功能运行得这么慢，但听起来你是用这个屏蔽来选择行（而不是单个值），你可以从被屏蔽的行创建一个普通数组，然后使用np函数：

b.mask = np.zeros(500)
b.mask[498] = True
t = timeit.Timer('c=b.view(np.ndarray)[~b.mask[:,0]]; np.apply_along_axis(test, 1, c)','from __main__ import test, b, ma, np')
print round(t.timeit(100), 6)

更好的办法是根本不使用被屏蔽的数组；只需将你的数据和一个一维的屏蔽数组作为独立的变量来维护：

a = np.arange(1000).reshape(500, 2)
mask = np.ones(a.shape[0], dtype=bool)
mask[498] = False
out = np.apply_along_axis(test, 1, a[mask])