如何对二维 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]

我可以使用其他什么 methods/functions 来更有效地解决这个问题（使用列表理解或其他方法）？

这是一些代码：

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

或者，pre-compute 绝对值并使用它们与 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

如何对二维 numpy 数组中的向量 [u,v] 进行阈值处理？

How to threshold vectors [u,v] in a 2D numpy array?

python

arrays

numpy

threshold