Question

我写了一个阈值函数TH（arr，threshold），该函数接受矢量[u，v]的2D数组，如果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

如何在2D numpy数组中对向量[u，v]设置阈值？

1 个答案: