Question

对于2D晶格中的大量随机分布点，我想要有效地提取子阵列，该子阵列仅包含近似为索引的元素，这些元素在单独的2D二进制矩阵中被分配给非零值。目前，我的脚本如下：

lat_len = 100 # lattice length
input = np.random.random(size=(1000,2)) * lat_len
binary_matrix = np.random.choice(2, lat_len * lat_len).reshape(lat_len, -1)

def landed(input):
    output = []
    input_as_indices = np.floor(input)
    for i in range(len(input)):
        if binary_matrix[input_as_indices[i,0], input_as_indices[i,1]] == 1:
            output.append(input[i])
    output = np.asarray(output)
    return output

然而，我怀疑必须有更好的方法来做到这一点。上面的脚本可能需要很长时间才能运行10000次迭代。

Answer 1

你是对的。使用advanced numpy indexing，

在python中没有for循环可以更有效地完成上面的计算

def landed2(input):
    idx = np.floor(input).astype(np.int)
    mask = binary_matrix[idx[:,0], idx[:,1]] == 1
    return input[mask]

res1 = landed(input)
res2 = landed2(input)
np.testing.assert_allclose(res1, res2)

这导致加速度达到~150倍。

Answer 2

如果使用线性索引数组，似乎可以显着提升性能。这是一个解决我们案例的矢量化实现，类似于@rth's answer，但使用线性索引 -

# Get floor-ed indices
idx = np.floor(input).astype(np.int)

# Calculate linear indices 
lin_idx = idx[:,0]*lat_len + idx[:,1]

# Index raveled/flattened version of binary_matrix with lin_idx
# to extract and form the desired output
out = input[binary_matrix.ravel()[lin_idx] ==1]

因此，简而言之，我们有：

out = input[binary_matrix.ravel()[idx[:,0]*lat_len + idx[:,1]] ==1]

运行时测试 -

本节将此解决方案中提议的方法与使用行列索引的other solution进行比较。

案例＃1（原始数据）：

In [62]: lat_len = 100 # lattice length
    ...: input = np.random.random(size=(1000,2)) * lat_len
    ...: binary_matrix = np.random.choice(2, lat_len * lat_len).
                                             reshape(lat_len, -1)
    ...: 

In [63]: idx = np.floor(input).astype(np.int)

In [64]: %timeit input[binary_matrix[idx[:,0], idx[:,1]] == 1]
10000 loops, best of 3: 121 µs per loop

In [65]: %timeit input[binary_matrix.ravel()[idx[:,0]*lat_len + idx[:,1]] ==1]
10000 loops, best of 3: 103 µs per loop

案例＃2（更大的数据量）：

In [75]: lat_len = 1000 # lattice length
    ...: input = np.random.random(size=(100000,2)) * lat_len
    ...: binary_matrix = np.random.choice(2, lat_len * lat_len).
                                             reshape(lat_len, -1)
    ...: 

In [76]: idx = np.floor(input).astype(np.int)

In [77]: %timeit input[binary_matrix[idx[:,0], idx[:,1]] == 1]
100 loops, best of 3: 18.5 ms per loop

In [78]: %timeit input[binary_matrix.ravel()[idx[:,0]*lat_len + idx[:,1]] ==1]
100 loops, best of 3: 13.1 ms per loop

因此，这种线性索引的性能提升似乎是 20% - 30% 。

基于二进制矩阵中元素包含的快速数组操作

2 个答案: