基於元素包含在二進制矩陣中的快速數組操作

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蟒環，

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（原始datasizes）：

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（放大datasizes）：

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%。