cross-correlazione numpy - vettorizzazione

Question

Ho un gran numero di cross-correlazioni da calcolare e sto cercando il modo più veloce per farlo. Sto assumendo che la soluzione del problema sia più utile che con i loop

Ho un array 3D etichettato come elettrodo x cronometro x trial (forma: 64x256x913). Voglio calcolare la massima cross-correlazione dei tempi per ogni coppia di elettrodi, per ogni prova.

In particolare: per ogni prova, voglio prendere ciascuna combinazione di elettrodi di coppia e calcolare il valore di cross-correlazione massimo per ogni coppia. Ciò si tradurrà in 4096 (64 * 64) valori di cross-correlazione massimi in una singola riga/vettore. Questo sarà fatto per ogni prova, impilando ciascuna delle righe/vettori l'uno sopra l'altro risultando in un array 2D finale di forma 913 * 4096 contenente i valori di massima correlazione incrociata

Questo è un sacco di calcoli ma voglio prova a trovare il metodo più veloce per farlo. Ho preso in giro alcuni proto-codici usando elenchi come contenitori che potrebbero aiutare a spiegare un po 'meglio il problema. Ci possono essere degli errori logici, ma in entrambi i casi il codice non viene eseguito sul mio computer perché c'è così tanto da calcolare che Python si blocca. Eccolo:

#allData is a 64x256x913 array 

all_experiment_trials = [] 
for trial in range(allData.shape[2]): 
    all_trial_electrodes = [] 
    for electrode in range(allData.shape[0]): 
     for other_electrode in range(allData.shape[0]): 
      if electrode == other_electrode: 
       pass 
      else: 
       single_xcorr = max(np.correlate(allData[electrode,:,trial], allData[other_electrode,:,trial], "full")) 
       all_trial_electrodes.append(single_xcorr) 
    all_experiment_trials.append(all_trial_electrodes) 

Ovviamente i loop sono veramente lenti per questo tipo di cose. Esiste una soluzione vettoriale per questo usando array numpy?

Ho controllato le cose come correlate2d() e simili, ma non credo che lavorano veramente nel mio caso come io non sto moltiplicando 2 matrici insieme

Answer 1

Ecco un approccio vettorizzati sulla base di np.einsum -

def vectorized_approach(allData): 
    # Get shape 
    M,N,R = allData.shape 

    # Valid mask based on condition: "if electrode == other_electrode" 
    valid_mask = np.mod(np.arange(M*M),M+1)!=0 

    # Elementwise multiplications across all elements in axis=0, 
    # and then summation along axis=1 
    out = np.einsum('ijkl,ijkl->lij',allData[:,None,:,:],allData[None,:,:,:]) 

    # Use valid mask to skip columns and have the final output 
    return out.reshape(R,-1)[:,valid_mask] 

prova Runtime e verificare i risultati -

In [10]: allData = np.random.rand(20,80,200) 

In [11]: def org_approach(allData): 
    ...:  all_experiment_trials = [] 
    ...:  for trial in range(allData.shape[2]): 
    ...:   all_trial_electrodes = [] 
    ...:   for electrode in range(allData.shape[0]): 
    ...:    for other_electrode in range(allData.shape[0]): 
    ...:     if electrode == other_electrode: 
    ...:      pass 
    ...:     else: 
    ...:      single_xcorr = max(np.correlate(allData[electrode,:,trial], allData[other_electrode,:,trial])) 
    ...:      all_trial_electrodes.append(single_xcorr) 
    ...:   all_experiment_trials.append(all_trial_electrodes) 
    ...:  return all_experiment_trials 
    ...: 

In [12]: %timeit org_approach(allData) 
1 loops, best of 3: 1.04 s per loop 

In [13]: %timeit vectorized_approach(allData) 
100 loops, best of 3: 15.1 ms per loop 

In [14]: np.allclose(vectorized_approach(allData),np.asarray(org_approach(allData))) 
Out[14]: True