Ottimizzazione/rimozione del loop

Question

Ho il seguente codice che vorrei ottimizzare usando numpy, preferibilmente rimuovendo il loop. Non riesco a vedere come affrontarlo, quindi qualsiasi suggerimento sarebbe utile.

indici è una matrice numerica N (2) di numeri interi, N può essere di alcuni milioni. Quello che fa il codice è trovare gli indici ripetuti nella prima colonna. Per questi indici faccio tutte le combinazioni di due degli indici corrispondenti nella seconda colonna. Quindi li raccolgo insieme all'indice nella prima colonna.

index_sets = [] 
uniques, counts = np.unique(indices[:,0], return_counts=True) 
potentials = uniques[counts > 1] 
for p in potentials: 
    correspondents = indices[(indices[:,0] == p),1] 
    combs = np.vstack(list(combinations(correspondents, 2))) 
    combs = np.hstack((np.tile(p, (combs.shape[0], 1)), combs)) 
    index_sets.append(combs) 

Answer 1

Ecco una soluzione che è vettorizzata su N. Si noti che contiene ancora un ciclo for, ma è un ciclo su ciascun 'gruppo di moltiplicazioni-chiave', che è garantito essere un numero molto più piccolo (tipicamente poche decine al massimo).

Per N = 1.000.000, il tempo di esecuzione è un ordine di grandezza di un secondo sul mio PC.

import numpy_indexed as npi 
N = 1000000 
indices = np.random.randint(0, N/10, size=(N, 2)) 

def combinations(x): 
    """vectorized computation of combinations for an array of sequences of equal length 

    Parameters 
    ---------- 
    x : ndarray, [..., n_items] 

    Returns 
    ------- 
    ndarray, [..., n_items * (n_items - 1)/2, 2] 
    """ 
    return np.rollaxis(x[..., np.triu_indices(x.shape[-1], 1)], -2, x.ndim+1) 

def process(indices): 
    """process a subgroup of indices, all having equal multiplicity 

    Parameters 
    ---------- 
    indices : ndarray, [n, 2] 

    Returns 
    ------- 
    ndarray, [m, 3] 
    """ 
    keys, vals = npi.group_by(indices[:, 0], indices[:, 1]) 
    combs = combinations(vals) 
    keys = np.repeat(keys, combs.shape[1]) 
    return np.concatenate([keys[:, None], combs.reshape(-1, 2)], axis=1) 

index_groups = npi.group_by(npi.multiplicity(indices[:, 0])).split(indices) 
result = np.concatenate([process(ind) for ind in index_groups]) 

Disclaimer: io sono l'autore del pacchetto numpy_indexed.

Answer 2

alcuni miglioramenti potrebbero essere suggeriti:

• inizializzare di matrice di uscita, per il quale può pre-calcolare la stima del numero di righe necessarie per combinazioni memorizzazione corrispondente a ciascun gruppo. Sappiamo che con gli elementi N, il numero totale di combinazioni possibili sarebbe N*(N-1)/2, per darci le lunghezze di combinazione per ciascun gruppo. Inoltre, il numero totale di righe nell'array di output sarebbe la somma di tutte quelle lunghezze di intervallo.

• Pre-calcolare il più materiale possibile in modo vettorizzato prima di entrare in un ciclo.

• Utilizzare un ciclo per ottenere le combinazioni, che a causa del motivo irregolare non possono essere vettorializzate. Usare np.repeat per simulare la piastrellatura e farlo prima del ciclo per darci il primo elemento per ogni gruppo e quindi la prima colonna della matrice di output.

Così, con tutti quei miglioramenti in mente, un'implementazione sarebbe simile a questa -

# Remove rows with counts == 1 
_,idx, counts = np.unique(indices[:,0], return_index=True, return_counts=True) 
indices = np.delete(indices,idx[counts==1],axis=0) 

# Decide the starting indices of corresponding to start of new groups 
# charaterized by new elements along the sorted first column 
start_idx = np.unique(indices[:,0], return_index=True)[1] 
all_idx = np.append(start_idx,indices.shape[0]) 

# Get interval lengths that are required to store pairwise combinations 
# of each group for unique ID from column-0 
interval_lens = np.array([item*(item-1)/2 for item in np.diff(all_idx)]) 

# Setup output array and set the first column as a repeated array 
out = np.zeros((interval_lens.sum(),3),dtype=int) 
out[:,0] = np.repeat(indices[start_idx,0],interval_lens) 

# Decide the start-stop indices for storing into output array 
ssidx = np.append(0,np.cumsum(interval_lens)) 

# Finally run a loop gto store all the combinations into initialized o/p array 
for i in range(idx.size): 
    out[ssidx[i]:ssidx[i+1],1:] = \ 
    np.vstack(combinations(indices[all_idx[i]:all_idx[i+1],1],2)) 

Si prega di notare che l'array di uscita sarebbe un grande (M, 3) matrice di forma e non diviso in lista di array come prodotto dal codice originale. Se è ancora necessario in quanto tale, è possibile utilizzare np.split per lo stesso.

Inoltre, i test di runtime rapidi suggeriscono che non c'è molto miglioramento con il codice proposto. Quindi, probabilmente la maggior parte del runtime viene speso per ottenere le combinazioni. Pertanto, sembra che l'approccio alternativo a networkx sia particolarmente adatto per tali problemi di connessione potrebbe essere una soluzione migliore.