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

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

你是对的。使用高级 numpy索引，可以更有效地完成上述计算，而无需在 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 倍的加速。