Szybka manipulacja szykami w oparciu o włączenie elementów do macierzy binarnej

Question

Dla dużego zestawu losowo rozmieszczonych punktów w sieci 2D, chcę wydajnie wyodrębnić podbarwę, która zawiera tylko te elementy, które w przybliżeniu są indeksami wartości zerowe w osobnej macierzy binarnej 2D. Obecnie mój skrypt jest następujący:

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 

Podejrzewam jednak, że musi istnieć lepszy sposób na zrobienie tego. Powyższy skrypt może trwać dość długo w przypadku 10000 iteracji.

Answer 1

Masz rację. Obliczenia powyżej, można to zrobić bardziej efektywnie bez pętli for w Pythonie przy użyciu 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) 

Skutkuje to ~ 150x przyspieszyć.

Answer 2

Wygląda na to, że można znacznie zwiększyć wydajność, jeśli pracujemy z liniowo indeksowanymi tablicami. Oto wektorowy realizacja rozwiązać naszą sprawę, podobnie jak @rth's answer, ale przy użyciu indeksowania liniowy -

# 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] 

Tak więc, krótko mówiąc mamy:

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

Runtime testy -

Ta sekcja porównuje proponowane podejście w tym rozwiązaniu przeciwko other solution, które wykorzystuje indeksowanie wiersza-kolumny.

Case # 1 (datasizes oryginalne):

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 

Case # 2 (Większe 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 

Zatem zwiększenie wydajności z tego liniowego indeksowania wydaje się być o 20% - 30%.