阵列上的广播掩码操作

Posted 2021-04-03

考虑到其中一列中的元素到第四个数组的同一列的距离，我试图在三个数组上改进相当简单的屏蔽操作的性能。所有阵列都具有相同的形状。

可以通过广播改善此操作的性能吗？

# Random data with proper shape
x1, x2, x3, x4 = np.random.uniform(1., 10., (4, 10, 1000))

# This is the operation I' trying to 
dist = 0.01
for x in (x2, x3, x4):
    # Mask of the distance between the column '-6' of x1 versus arrays
    # x2, x3, x4
    msk = abs(x1[-6] - x[-6]) > dist

    # If the distance in this array is larger than the maximum allowed (dist),
    # mask with values from 'x1'.
    x[:, msk] = x1[:, msk]

答案

作为广播的替代方案，我使用numba大约加速10倍。

np.random.seed(0)
xs = np.random.uniform(0, 10, (4, 10, 1000))
x1, x2, x3, x4 = xs.copy()

from numba import jit


@jit(nopython=True)
def modified(xs):
    dist = .01
    for i in range(1, 4):
        for j in range(1000):
            if abs(xs[i, -6, j] - xs[0, -6, j]) > dist:
                for k in range(10):
                    xs[i, k, j] = xs[0, k, j]