Joint Peristimulus Time Histogram in Python

jspth.py

def jspth(spikes1, spikes2):
    '''Joint Peristimulus Time Histogram
    
    Parameters
    ----------
    spikes1, spikes2: ndarray, shape (n_time, n_trials)
    
    Returns
    -------
    joint_histogram : ndarray, shape (n_time, n_time)
    
    '''
    n_time = spikes1.shape[0]
    joint_histogram = np.zeros((n_time, n_time))
    for single_trial_spikes1, single_trial_spikes2 in zip(spikes1.T, spikes2.T):
        joint_histogram += (single_trial_spikes1[:, np.newaxis]
                            * single_trial_spikes2[:, np.newaxis].T)
    
    return joint_histogram