Manipolazione di array veloce basata sull'inclusione di elementi nella matrice binaria

Question

Per un grande insieme di punti distribuiti casualmente in un reticolo 2D, voglio estrarre in modo efficiente un subarray, che contiene solo gli elementi che, approssimati come indici, sono assegnati a non- valori zero in una matrice binaria 2D separata. Attualmente, il mio script è il seguente:

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 

Tuttavia, ho il sospetto che ci sia un modo migliore per farlo. Lo script precedente può richiedere un tempo piuttosto lungo per eseguire 10000 iterazioni.

Answer 1

Sei corretto. Il calcolo di cui sopra, si può essere fatto in modo più efficiente, senza un ciclo for in Python usando 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) 

questo si traduce in un ~ 150x speed-up.

Answer 2

Sembra che si possa ottenere un aumento notevole delle prestazioni se si lavora con matrici indicizzate linearmente. Ecco un'implementazione vettorializzare per risolvere il nostro caso, simile a @rth's answer, ma utilizzando l'indicizzazione lineare -

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

Così, in breve abbiamo:

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

test Runtime -

Questa sezione confronta la approccio proposto in questa soluzione rispetto allo other solution che utilizza l'indicizzazione di colonne.

Caso # 1 (datasizes originale):

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 

Caso # 2 (datasizes più grande):

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 

Così, l'incremento delle prestazioni con questa indicizzazione lineare sembra essere su 20% - 30%.