阵列上的广播掩码操作

时间:2018-12-14 21:31:34

标签: python performance numpy numpy-broadcasting

考虑到元素在它们中的一列到第四列的同一列之间的距离,我试图提高对三个数组的相当简单的屏蔽操作的性能。所有阵列都具有相同的形状。

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

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

1 个答案:

答案 0 :(得分:1)

作为广播的替代方法,使用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]