[英]Faster way to threshold a 4-D numpy array
我有一个大小为 (98,359,256,269) 的 4D numpy 数组,我想设置阈值。 现在,我有两个单独的列表,分别保存前 2 维和后 2 维的坐标。 (前 2 个维度的 mag_ang 和后 2 个维度的索引)。
指数大小:(61821,2)
mag_ang 的大小:(35182,2)
目前,我的代码如下所示:
inner_points = []
for k in indices:
x = k[0]
y = k[1]
for i,ctr in enumerate(mag_ang):
mag = ctr[0]
ang = ctr[1]
if X[mag][ang][x][y] > 10:
inner_points.append((y,x))
这段代码有效,但速度很慢,我想知道是否有更多的pythonic/更快的方法来做到这一点?
直接使用numpy
。 如果indices
和mag_ang
是 numpy arrays 两列,每列对应适当的坐标:
(x, y), (mag, ang) = indices.T, mag_ang.T
index_matrix = np.meshgrid(mag, ang, x, y).T.reshape(-1,4)
inner_mag, inner_ang, inner_x, inner_y = np.where(X[index_matrix] > 10)
现在你的inner...
变量为每个坐标保存 arrays 。 要获得一个单一的 par 列表,您可以 zip inner_y
和inner_x
。
(编辑:添加了第二种替代方法)
使用 numpy 多数组索引:
import time
import numpy as np
n_mag, n_ang, n_x, n_y = 10, 12, 5, 6
shape = n_mag, n_ang, n_x, n_y
X = np.random.random_sample(shape) * 20
nb_indices = 100 # 61821
indices = np.c_[np.random.randint(0, n_x, nb_indices), np.random.randint(0, n_y, nb_indices)]
nb_mag_ang = 50 # 35182
mag_ang = np.c_[np.random.randint(0, n_mag, nb_mag_ang), np.random.randint(0, n_ang, nb_mag_ang)]
# original method
inner_points = []
start = time.time()
for x, y in indices:
for mag, ang in mag_ang:
if X[mag][ang][x][y] > 10:
inner_points.append((y, x))
end = time.time()
print(end - start)
# faster method 1:
inner_points_faster1 = []
start = time.time()
for x, y in indices:
if np.any(X[mag_ang[:, 0], mag_ang[:, 1], x, y] > 10):
inner_points_faster1.append((y, x))
end = time.time()
print(end - start)
# faster method 2:
start = time.time()
# note: depending on the real size of mag_ang and indices, you may wish to do this the other way round ?
found = X[:, :, indices[:, 0], indices[:, 1]][mag_ang[:, 0], mag_ang[:, 1], :] > 10
# 'found' shape is (nb_mag_ang x nb_indices)
assert found.shape == (nb_mag_ang, nb_indices)
matching_indices_mask = found.any(axis=0)
inner_points_faster2 = indices[matching_indices_mask, :]
end = time.time()
print(end - start)
# finally assert equality of findings
inner_points = np.unique(np.array(inner_points))
inner_points_faster1 = np.unique(np.array(inner_points_faster1))
inner_points_faster2 = np.unique(inner_points_faster2)
assert np.array_equal(inner_points, inner_points_faster1)
assert np.array_equal(inner_points, inner_points_faster2)
产量
0.04685807228088379
0.0
0.0
(当然,如果你增加形状,第二个和第三个时间不会为零)
最后说明:这里我在最后使用“唯一”,但最好提前为indices
和mag_ang
arrays 做这件事(除非你确定它们已经是唯一的)
这里有一些利用broadcasting
的 vecorized 方式 -
thresh = 10
mask = X[mag_ang[:,0],mag_ang[:,1],indices[:,0,None],indices[:,1,None]]>thresh
r = np.where(mask)[0]
inner_points_out = indices[r][:,::-1]
对于较大的 arrays,我们可以先比较然后索引来获取掩码 -
mask = (X>thresh)[mag_ang[:,0],mag_ang[:,1],indices[:,0,None],indices[:,1,None]]
如果您只对indices
的唯一坐标感兴趣,请直接使用掩码 -
inner_points_out = indices[mask.any(1)][:,::-1]
对于大型 arrays,我们还可以通过numexpr
模块利用多核。
因此,首先导入模块 -
import numexpr as ne
然后,在前面列出的计算中将 ( (X>thresh)
替换为ne.evaluate('X>thresh')
。
使用np.where
inner = np.where(X > 10)
a, b, x, y = zip(*inner)
inner_points = np.vstack([y, x]).T
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.