numpy cross-correlation - vectorizing numpy cross-correlation - vectorizing numpy numpy

numpy cross-correlation - vectorizing


Here's one vectorized approach based on 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]

Runtime test and verify results -

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 loopIn [13]: %timeit vectorized_approach(allData)100 loops, best of 3: 15.1 ms per loopIn [14]: np.allclose(vectorized_approach(allData),np.asarray(org_approach(allData)))Out[14]: True