如何在二维 numpy 数组中阈值向量 [u,v]？

Question

我写了一个阈值函数 TH(arr, threshold)，它接受向量 [u,v] 的二维数组，如果 u 和 v 的绝对值都低于指定阈值，则将它们设置为 0。

该函数由 2 个 for 循环组成，可以完成这项工作，但需要大量计算（我在大型数据集上运行它）。

例子：

[u, v] --> 输出（阈值 = 1）

[2, 2] --> [2, 2]

[2, .1] --> [2, .1]

[.1,.1] --> [0, 0]

我可以使用哪些其他方法/函数来更有效地解决此问题（使用列表理解或其他方法）？

这里有一些代码：

import numpy as np
import time
start = time.time()

def TH(arr, threshold):
    for idx, value in enumerate(arr):
        for i, item in enumerate(value):
            if np.abs(item[0]) < threshold and np.abs(item[1]) < threshold:
                arr[idx][i][0] = 0.0
                arr[idx][i][1] = 0.0
    return arr

a = np.array([[[.5,.8], [3,4], [3,.1]],
              [[0,2], [.5,.5], [.3,3]],
              [[.4,.4], [.1,.1], [.5,5]]])

a = TH(a, threshold = 1)
print(a)

end = time.time()
print("Run time: ", end-start)

输出：

[[[0.  0. ]
  [3.  4. ]
  [3.  0.1]]

 [[0.  2. ]
  [0.  0. ]
  [0.3 3. ]]

 [[0.  0. ]
  [0.  0. ]
  [0.5 5. ]]]

Run time:  0.0009984970092773438

Answer 1

简单地将沿最后一个轴的两个元素切片，并以矢量化的方式执行相同的操作以获得掩码，最后将掩码索引到输入数组中以分配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