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

Question

对于二维晶格中大量随机分布的点，我想高效地提取一个子数组，该子数组仅包含近似为索引的元素，这些元素被分配给单独的二维二进制矩阵中的非零值。目前，我的脚本如下:

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 倍的加速。