-
Notifications
You must be signed in to change notification settings - Fork 23
/
ccg.py
190 lines (141 loc) · 6.05 KB
/
ccg.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
# -*- coding: utf-8 -*-
"""Cross-correlograms."""
#------------------------------------------------------------------------------
# Imports
#------------------------------------------------------------------------------
import numpy as np
from phylib.utils._types import _as_array
from phylib.io.array import _index_of, _unique
#------------------------------------------------------------------------------
# Cross-correlograms
#------------------------------------------------------------------------------
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 _diff_shifted(arr, steps=1):
arr = _as_array(arr)
return arr[steps:] - arr[:len(arr) - steps]
def _create_correlograms_array(n_clusters, winsize_bins):
return np.zeros((n_clusters, n_clusters, winsize_bins // 2 + 1),
dtype=np.int32)
def _symmetrize_correlograms(correlograms):
"""Return the symmetrized version of the CCG arrays."""
n_clusters, _, n_bins = correlograms.shape
assert n_clusters == _
# We symmetrize c[i, j, 0].
# This is necessary because the algorithm in correlograms()
# is sensitive to the order of identical spikes.
correlograms[..., 0] = np.maximum(correlograms[..., 0],
correlograms[..., 0].T)
sym = correlograms[..., 1:][..., ::-1]
sym = np.transpose(sym, (1, 0, 2))
return np.dstack((sym, correlograms))
def firing_rate(spike_clusters, cluster_ids=None, bin_size=None, duration=None):
"""Compute the average number of spikes per cluster per bin."""
# Take the cluster order into account.
if cluster_ids is None:
cluster_ids = _unique(spike_clusters)
else:
cluster_ids = _as_array(cluster_ids)
# Like spike_clusters, but with 0..n_clusters-1 indices.
spike_clusters_i = _index_of(spike_clusters, cluster_ids)
assert bin_size > 0
bc = np.bincount(spike_clusters_i)
# Handle the case where the last cluster(s) are empty.
if len(bc) < len(cluster_ids):
n = len(cluster_ids) - len(bc)
bc = np.concatenate((bc, np.zeros(n, dtype=bc.dtype)))
assert bc.shape == (len(cluster_ids),)
return bc * np.c_[bc] * (bin_size / (duration or 1.))
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 *all* 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.
sample_rate : float
Sampling rate.
symmetrize : boolean (True)
Whether the output matrix should be symmetrized or not.
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 order 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=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
if symmetrize:
return _symmetrize_correlograms(correlograms)
else:
return correlograms