Просто нарежьте два элемента вдоль последней оси и выполните те же операции векторизованным способом, чтобы получить маску, и, наконец, внесите маску в указатель во входной массив, чтобы присвоить 0s
-
mask = (np.abs(arr[...,0]) < threshold) & (np.abs(arr[...,1]) < threshold)
arr[mask] = 0
Обратите внимание, что arr[...,0]
- это еще один способ поставить arr[:,:,0]
, и он предназначен для нарезки общего ndarray вдоль последней оси.Аналогично, для arr[...,1]
.
В качестве альтернативы предварительно вычислите абсолютные значения и используйте их для сравнения с threshold
и найдите совпадения all
с последней осью, чтобы получить ту же маску -
ab = np.abs(arr)
mask = (ab < threshold).all(-1)
Или используйте тот же метод нарезки после вычисления абсолютных значений -
mask = (ab[...,0] < threshold) & (ab[...,1] < threshold)
Для больших массивов мы также можем использовать numexpr
module -
import numexpr as ne
m0 = ne.evaluate('abs(arr)<threshold')
mask = m0[...,0] & m0[...,1]
Сроки -
In [209]: arr = np.random.rand(1080,1920,2)
In [210]: threshold = 1
In [211]: %timeit (np.abs(arr[...,0])<threshold) & (np.abs(arr[...,1])<threshold)
100 loops, best of 3: 10.2 ms per loop
In [212]: %timeit np.abs(arr).all(1)
10 loops, best of 3: 34.5 ms per loop
In [213]: %%timeit
...: ab = np.abs(arr)
...: (ab[...,0] < threshold) & (ab[...,1] < threshold)
...:
100 loops, best of 3: 11 ms per loop
In [214]: %%timeit
...: m0 = ne.evaluate('abs(arr)<threshold')
...: m0[...,0] & m0[...,1]
...:
100 loops, best of 3: 4.79 ms per loop