我编写了一个阈值函数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

解决方法:

只需沿最后一个轴对这两个元素进行切片,并以矢量化的方式执行相同的操作即可得到一个掩码,最后用掩码将其索引到输入数组中以分配0-

mask = (np.abs(arr[...,0]) < threshold) & (np.abs(arr[...,1]) < threshold)
arr[mask] = 0

请注意,arr […,0]是放置arr [：,：,0]的另一种方法,旨在沿最后一个轴切片通用ndarray.同样,对于arr […,1].

另外,也可以预先计算绝对值,并使用它们与阈值进行比较,并针对最后一个轴查找所有匹配项以获取相同的掩码-

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