Операция широковещательной маски на массиве

Question

Я пытаюсь улучшить производительность довольно простой операции маскирования для трех массивов, учитывая расстояние элементов в одном из их столбцов до того же столбца четвертого массива.Все массивы имеют одинаковую форму.

Можно ли улучшить производительность этой операции с помощью широковещания?

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

Answer 1

В качестве альтернативы вещанию я получаю примерно 10-кратное ускорение с помощью numba.

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]