def _increment(arr, indices): """Increment some indices in a 1D vector of non-negative integers. Repeated indices are taken into account.""" arr = _as_array(arr) indices = _as_array(indices) bbins = np.bincount(indices) arr[:len(bbins)] += bbins return arr
def correlograms(spike_times, spike_clusters, cluster_ids=None, sample_rate=1., bin_size=None, window_size=None, symmetrize=True, ): """Compute all pairwise cross-correlograms among the clusters appearing in `spike_clusters`. Parameters ---------- spike_times : array-like Spike times in seconds. spike_clusters : array-like Spike-cluster mapping. cluster_ids : array-like The list of unique clusters, in any order. That order will be used in the output array. bin_size : float Size of the bin, in seconds. window_size : float Size of the window, in seconds. Returns ------- correlograms : array A `(n_clusters, n_clusters, winsize_samples)` array with all pairwise CCGs. """ assert sample_rate > 0. assert np.all(np.diff(spike_times) >= 0), ("The spike times must be " "increasing.") # Get the spike samples. spike_times = np.asarray(spike_times, dtype=np.float64) spike_samples = (spike_times * sample_rate).astype(np.int64) spike_clusters = _as_array(spike_clusters) assert spike_samples.ndim == 1 assert spike_samples.shape == spike_clusters.shape # Find `binsize`. bin_size = np.clip(bin_size, 1e-5, 1e5) # in seconds binsize = int(sample_rate * bin_size) # in samples assert binsize >= 1 # Find `winsize_bins`. window_size = np.clip(window_size, 1e-5, 1e5) # in seconds winsize_bins = 2 * int(.5 * window_size / bin_size) + 1 assert winsize_bins >= 1 assert winsize_bins % 2 == 1 # Take the cluster oder into account. if cluster_ids is None: clusters = _unique(spike_clusters) else: clusters = _as_array(cluster_ids) n_clusters = len(clusters) # Like spike_clusters, but with 0..n_clusters-1 indices. spike_clusters_i = _index_of(spike_clusters, clusters) # Shift between the two copies of the spike trains. shift = 1 # At a given shift, the mask precises which spikes have matching spikes # within the correlogram time window. mask = np.ones_like(spike_samples, dtype=np.bool) correlograms = _create_correlograms_array(n_clusters, winsize_bins) # The loop continues as long as there is at least one spike with # a matching spike. while mask[:-shift].any(): # Number of time samples between spike i and spike i+shift. spike_diff = _diff_shifted(spike_samples, shift) # Binarize the delays between spike i and spike i+shift. spike_diff_b = spike_diff // binsize # Spikes with no matching spikes are masked. mask[:-shift][spike_diff_b > (winsize_bins // 2)] = False # Cache the masked spike delays. m = mask[:-shift].copy() d = spike_diff_b[m] # # Update the masks given the clusters to update. # m0 = np.in1d(spike_clusters[:-shift], clusters) # m = m & m0 # d = spike_diff_b[m] d = spike_diff_b[m] # Find the indices in the raveled correlograms array that need # to be incremented, taking into account the spike clusters. indices = np.ravel_multi_index((spike_clusters_i[:-shift][m], spike_clusters_i[+shift:][m], d), correlograms.shape) # Increment the matching spikes in the correlograms array. _increment(correlograms.ravel(), indices) shift += 1 # Remove ACG peaks. correlograms[np.arange(n_clusters), np.arange(n_clusters), 0] = 0 if symmetrize: return _symmetrize_correlograms(correlograms) else: return correlograms
def _diff_shifted(arr, steps=1): arr = _as_array(arr) return arr[steps:] - arr[:len(arr) - steps]