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]
