-
Notifications
You must be signed in to change notification settings - Fork 28
/
cluster.py
583 lines (433 loc) · 21.9 KB
/
cluster.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
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
import logging
from math import ceil
import numpy as np
import cupy as cp
from tqdm import tqdm
from .preprocess import my_min, my_sum
from .cptools import svdecon, zscore, ones
from .utils import Bunch, get_cuda
logger = logging.getLogger(__name__)
def getClosestChannels(probe, sigma, NchanClosest):
# this function outputs the closest channels to each channel,
# as well as a Gaussian-decaying mask as a function of pairwise distances
# sigma is the standard deviation of this Gaussian-mask
# compute distances between all pairs of channels
xc = cp.asarray(probe.xc, dtype=np.float32, order='F')
yc = cp.asarray(probe.yc, dtype=np.float32, order='F')
C2C = (xc[:, np.newaxis] - xc) ** 2 + (yc[:, np.newaxis] - yc) ** 2
C2C = cp.sqrt(C2C)
Nchan = C2C.shape[0]
# sort distances
isort = cp.argsort(C2C, axis=0)
# take NchanCLosest neighbors for each primary channel
iC = isort[:NchanClosest, :]
# in some cases we want a mask that decays as a function of distance between pairs of channels
# this is an awkward indexing to get the corresponding distances
ix = iC + cp.arange(0, Nchan ** 2, Nchan)
mask = cp.exp(-C2C.T.ravel()[ix] ** 2 / (2 * sigma ** 2))
# masks should be unit norm for each channel
mask = mask / cp.sqrt(1e-3 + cp.sum(mask ** 2, axis=0))
return iC, mask, C2C
def isolated_peaks_new(S1, params):
"""
takes a matrix of timepoints by channels S1
outputs threshold crossings that are relatively isolated from other peaks
outputs row, column and magnitude of the threshold crossing
"""
S1 = cp.asarray(S1)
# finding the local minimum in a sliding window within plus/minus loc_range extent
# across time and across channels
smin = my_min(S1, params.loc_range, [0, 1])
# the peaks are samples that achieve this local minimum, AND have negativities less
# than a preset threshold
peaks = (S1 < smin + 1e-3) & (S1 < params.spkTh)
# only take local peaks that are isolated from other local peaks
# if there is another local peak close by, this sum will be at least 2
sum_peaks = my_sum(peaks, params.long_range, [0, 1])
# set to 0 peaks that are not isolated, and multiply with the voltage values
peaks = peaks * (sum_peaks < 1.2) * S1
# exclude temporal buffers
peaks[:params.nt0, :] = 0
peaks[-params.nt0:, :] = 0
# find the non-zero peaks, and take their amplitudes
col, row = cp.nonzero(peaks.T)
# invert the sign of the amplitudes
mu = -peaks[row, col]
return row, col, mu
def get_SpikeSample(dataRAW, row, col, params):
"""
given a batch of data (time by channels), and some time (row) and channel (col) indices for
spikes, this function returns the 1D time clips of voltage around those spike times
"""
nT, nChan = dataRAW.shape
# times around the peak to consider
dt = cp.arange(params.nt0)
# the negativity is expected at nt0min, so we align the detected peaks there
dt = -params.nt0min + dt
# temporal indices (awkward way to index into full matrix of data)
indsT = row + dt[:, np.newaxis] + 1 # broadcasting
indsC = col
indsC[indsC < 0] = 0 # anything that's out of bounds just gets set to the limit
indsC[indsC >= nChan] = nChan - 1 # only needed for channels not time (due to time buffer)
indsT = cp.transpose(cp.atleast_3d(indsT), [0, 2, 1])
indsC = cp.transpose(cp.atleast_3d(indsC), [2, 0, 1])
# believe it or not, these indices grab just the right timesamples forour spikes
ix = indsT + indsC * nT
# grab the data and reshape it appropriately (time samples by channels by num spikes)
clips = dataRAW.T.ravel()[ix[:, 0, :]].reshape((dt.size, row.size), order='F') # HERE
return clips
def extractPCfromSnippets(proc, probe=None, params=None, Nbatch=None):
# extracts principal components for 1D snippets of spikes from all channels
# loads a subset of batches to find these snippets
NT = params.NT
nPCs = params.nPCs
Nchan = probe.Nchan
batchstart = np.arange(0, NT * Nbatch + 1, NT).astype(np.int64)
# extract the PCA projections
# initialize the covariance of single-channel spike waveforms
CC = cp.zeros(params.nt0, dtype=np.float32)
# from every 100th batch
for ibatch in range(0, Nbatch, 100):
offset = Nchan * batchstart[ibatch]
dat = proc.flat[offset:offset + NT * Nchan].reshape((-1, Nchan), order='F')
if dat.shape[0] == 0:
continue
# move data to GPU and scale it back to unit variance
dataRAW = cp.asarray(dat, dtype=np.float32) / params.scaleproc
# find isolated spikes from each batch
row, col, mu = isolated_peaks_new(dataRAW, params)
# for each peak, get the voltage snippet from that channel
c = get_SpikeSample(dataRAW, row, col, params)
# scale covariance down by 1,000 to maintain a good dynamic range
CC = CC + cp.dot(c, c.T) / 1e3
# the singular vectors of the covariance matrix are the PCs of the waveforms
U, Sv, V = svdecon(CC)
wPCA = U[:, :nPCs] # take as many as needed
# adjust the arbitrary sign of the first PC so its negativity is downward
# TODO: unclear - is 20 here the index into the spike waveform? Should this be hardcoded?
# - should it be nt0min instead?
wPCA[:, 0] = -wPCA[:, 0] * cp.sign(wPCA[20, 0])
return wPCA
def sortBatches2(ccb0):
# takes as input a matrix of nBatches by nBatches containing
# dissimilarities.
# outputs a matrix of sorted batches, and the sorting order, such that
# ccb1 = ccb0(isort, isort)
# put this matrix on the GPU
ccb0 = cp.asarray(ccb0, order='F')
# compute its svd on the GPU (this might also be fast enough on CPU)
u, s, v = svdecon(ccb0)
# HACK: consistency with MATLAB
u = u * cp.sign(u[0, 0])
v = v * cp.sign(u[0, 0])
# initialize the positions xs of the batch embeddings to be very small but proportional to
# the first PC
xs = .01 * u[:, 0] / cp.std(u[:, 0], ddof=1)
# 200 iterations of gradient descent should be enough
# TODO: move_to_config
niB = 200
# this learning rate should usually work fine, since it scales with the average gradient
# and ccb0 is z-scored
# TODO: move_to_config
eta = 1
for k in tqdm(range(niB), desc="Sorting %d batches" % ccb0.shape[0]):
# euclidian distances between 1D embedding positions
ds = (xs - xs[:, np.newaxis]) ** 2
# the transformed distances go through this function
W = cp.log(1 + ds)
# the error is the difference between ccb0 and W
err = ccb0 - W
# ignore the mean value of ccb0
err = err - cp.mean(err, axis=0)
# backpropagate the gradients
err = err / (1 + ds)
err2 = err * (xs[:, np.newaxis] - xs)
D = cp.mean(err2, axis=1) # one half of the gradients is along this direction
E = cp.mean(err2, axis=0) # the other half is along this direction
# we don't need to worry about the gradients for the diagonal because those are 0
# final gradients for the embedding variable
dx = -D + E.T
# take a gradient step
xs = xs - eta * dx
# sort the embedding positions xs
isort = cp.argsort(xs, axis=0)
# sort the matrix of dissimilarities
ccb1 = ccb0[isort, :][:, isort]
return ccb1, isort
def initializeWdata2(call, uprojDAT, Nchan, nPCs, Nfilt, iC):
# this function initializes cluster means for the fast kmeans per batch
# call are time indices for the spikes
# uprojDAT are features projections (Nfeatures by Nspikes)
# some more parameters need to be passed in from the main workspace
# pick random spikes from the sample
# WARNING: replace ceil by floor because this is a random index, and 0/1 indexing
# discrepancy between Python and MATLAB.
irand = np.floor(np.random.rand(Nfilt) * uprojDAT.shape[1]).astype(np.int32)
W = cp.zeros((nPCs, Nchan, Nfilt), dtype=np.float32)
for t in range(Nfilt):
ich = iC[:, call[irand[t]]] # the channels on which this spike lives
# for each selected spike, get its features
W[:, ich, t] = uprojDAT[:, irand[t]].reshape(W[:, ich, t].shape, order='F')
W = W.reshape((-1, Nfilt), order='F') # HERE
# add small amount of noise in case we accidentally picked the same spike twice
W = W + .001 * cp.random.normal(size=W.shape).astype(np.float32)
mu = cp.sqrt(cp.sum(W ** 2, axis=0)) # get the mean of the template
W = W / (1e-5 + mu) # and normalize the template
W = W.reshape((nPCs, Nchan, Nfilt), order='F') # HERE
nW = (W[0, ...] ** 2) # squared amplitude of the first PC feture
W = W.reshape((nPCs * Nchan, Nfilt), order='F') # HERE
# determine biggest channel according to the amplitude of the first PC
Wheights = cp.argmax(nW, axis=0)
return W, mu, Wheights, irand
def mexThSpkPC(Params, dataRAW, wPCA, iC):
code, constants = get_cuda('mexThSpkPC')
Nthreads = constants.Nthreads
maxFR = constants.maxFR
NT, Nchan, NchanNear, nt0, nt0min, spkTh, NrankPC = Params
NT = int(NT)
Nchan = int(Nchan)
# Input GPU arrays.
d_Params = cp.asarray(Params, dtype=np.float64, order='F')
d_data = cp.asarray(dataRAW, dtype=np.float32, order='F')
d_W = cp.asarray(wPCA, dtype=np.float32, order='F')
d_iC = cp.asarray(iC, dtype=np.int32, order='F')
# New GPU arrays.
d_dout = cp.zeros((Nchan, NT), dtype=np.float32, order='F')
d_dmax = cp.zeros((Nchan, NT), dtype=np.float32, order='F')
d_st = cp.zeros(maxFR, dtype=np.int32, order='F')
d_id = cp.zeros(maxFR, dtype=np.int32, order='F')
d_counter = cp.zeros(1, dtype=np.int32, order='F')
# filter the data with the temporal templates
Conv1D = cp.RawKernel(code, 'Conv1D')
Conv1D((Nchan,), (Nthreads,), (d_Params, d_data, d_W, d_dout))
# get the max of the data
max1D = cp.RawKernel(code, 'max1D')
max1D((Nchan,), (Nthreads,), (d_Params, d_dout, d_dmax))
# take max across nearby channels
maxChannels = cp.RawKernel(code, 'maxChannels')
maxChannels(
(int(NT // Nthreads),), (Nthreads,),
(d_Params, d_dout, d_dmax, d_iC, d_st, d_id, d_counter))
# move d_x to the CPU
minSize = 1
minSize = min(maxFR, int(d_counter[0]))
d_featPC = cp.zeros((NrankPC * NchanNear, minSize), dtype=np.float32, order='F')
d_id2 = cp.zeros(minSize, dtype=np.int32, order='F')
if (minSize > 0):
computeProjections = cp.RawKernel(code, 'computeProjections')
computeProjections(
(minSize,), (NchanNear, NrankPC), (d_Params, d_data, d_iC, d_st, d_id, d_W, d_featPC))
# TODO: check that the copy occurs on the GPU only
d_id2[:] = d_id[:minSize]
# Free memory.
# TODO: unclear - does this do something special for cupy objects?
# - all of these go out of scope in the next line anyway
del d_st, d_id, d_counter, d_Params, d_dmax, d_dout
# free_gpu_memory()
return d_featPC, d_id2
def extractPCbatch2(proc, params, probe, wPCA, ibatch, iC, Nbatch):
# this function finds threshold crossings in the data using
# projections onto the pre-determined principal components
# wPCA is number of time samples by number of PCs
# ibatch is a scalar indicating which batch to analyze
# iC is NchanNear by Nchan, indicating for each channel the nearest
# channels to it
nt0min = params.nt0min
spkTh = params.ThPre
nt0, NrankPC = wPCA.shape
NT, Nchan = params.NT, probe.Nchan
# starts with predefined PCA waveforms
wPCA = wPCA[:, :3]
NchanNear = iC.shape[0]
# batches start at these timepoints
batchstart = np.arange(0, NT * Nbatch + 1, NT).astype(np.int64)
offset = Nchan * batchstart[ibatch]
dat = proc.flat[offset:offset + NT * Nchan].reshape((-1, Nchan), order='F')
dataRAW = cp.asarray(dat, dtype=np.float32) / params.scaleproc
# another Params variable to take all our parameters into the C++ code
Params = [NT, Nchan, NchanNear, nt0, nt0min, spkTh, NrankPC]
# call a CUDA function to do the hard work
# returns a matrix of features uS, as well as the center channels for each spike
uS, idchan = mexThSpkPC(Params, dataRAW, wPCA, iC)
return uS, idchan
def mexClustering2(Params, uproj, W, mu, call, iMatch, iC):
code, _ = get_cuda('mexClustering2')
Nspikes = int(Params[0])
NrankPC = int(Params[1])
Nfilters = int(Params[2])
NchanNear = int(Params[6])
Nchan = int(Params[7])
d_Params = cp.asarray(Params, dtype=np.float64, order='F')
d_uproj = cp.asarray(uproj, dtype=np.float32, order='F')
d_W = cp.asarray(W, dtype=np.float32, order='F')
d_mu = cp.asarray(mu, dtype=np.float32, order='F')
d_call = cp.asarray(call, dtype=np.int32, order='F')
d_iC = cp.asarray(iC, dtype=np.int32, order='F')
d_iMatch = cp.asarray(iMatch, dtype=np.bool, order='F')
d_dWU = cp.zeros((NrankPC * Nchan, Nfilters), dtype=np.float32, order='F')
d_cmax = cp.zeros((Nspikes, Nfilters), dtype=np.float32, order='F')
d_id = cp.zeros(Nspikes, dtype=np.int32, order='F')
d_x = cp.zeros(Nspikes, dtype=np.float32, order='F')
d_nsp = cp.zeros(Nfilters, dtype=np.int32, order='F')
d_V = cp.zeros(Nfilters, dtype=np.float32, order='F')
# get list of cmaxes for each combination of neuron and filter
computeCost = cp.RawKernel(code, 'computeCost')
computeCost(
(Nfilters,), (1024,), (d_Params, d_uproj, d_mu, d_W, d_iMatch, d_iC, d_call, d_cmax))
# loop through cmax to find best template
bestFilter = cp.RawKernel(code, 'bestFilter')
bestFilter((40,), (256,), (d_Params, d_iMatch, d_iC, d_call, d_cmax, d_id, d_x))
# average all spikes for same template -- ORIGINAL
average_snips = cp.RawKernel(code, 'average_snips')
average_snips(
(Nfilters,), (NrankPC, NchanNear), (d_Params, d_iC, d_call, d_id, d_uproj, d_cmax, d_dWU))
count_spikes = cp.RawKernel(code, 'count_spikes')
count_spikes((7,), (256,), (d_Params, d_id, d_nsp, d_x, d_V))
del d_Params, d_V
return d_dWU, d_id, d_x, d_nsp, d_cmax
def mexDistances2(Params, Ws, W, iMatch, iC, Wh, mus, mu):
code, _ = get_cuda('mexDistances2')
Nspikes = int(Params[0])
Nfilters = int(Params[2])
d_Params = cp.asarray(Params, dtype=np.float64, order='F')
d_Ws = cp.asarray(Ws, dtype=np.float32, order='F')
d_W = cp.asarray(W, dtype=np.float32, order='F')
d_iMatch = cp.asarray(iMatch, dtype=np.bool, order='F')
d_iC = cp.asarray(iC, dtype=np.int32, order='F')
d_Wh = cp.asarray(Wh, dtype=np.int32, order='F')
d_mu = cp.asarray(mu, dtype=np.float32, order='F')
d_mus = cp.asarray(mus, dtype=np.float32, order='F')
d_cmax = cp.zeros(Nspikes * Nfilters, dtype=np.float32, order='F')
d_id = cp.zeros(Nspikes, dtype=np.int32, order='F')
d_x = cp.zeros(Nspikes, dtype=np.float32, order='F')
# get list of cmaxes for each combination of neuron and filter
computeCost = cp.RawKernel(code, 'computeCost')
computeCost(
(Nfilters,), (1024,), (d_Params, d_Ws, d_mus, d_W, d_mu, d_iMatch, d_iC, d_Wh, d_cmax))
# loop through cmax to find best template
bestFilter = cp.RawKernel(code, 'bestFilter')
bestFilter((40,), (256,), (d_Params, d_iMatch, d_Wh, d_cmax, d_mus, d_id, d_x))
del d_Params, d_cmax
return d_id, d_x
def clusterSingleBatches(ctx):
"""
outputs an ordering of the batches according to drift
for each batch, it extracts spikes as threshold crossings and clusters them with kmeans
the resulting cluster means are then compared for all pairs of batches, and a dissimilarity
score is assigned to each pair
the matrix of similarity scores is then re-ordered so that low dissimilaity is along
the diagonal
"""
Nbatch = ctx.intermediate.Nbatch
params = ctx.params
probe = ctx.probe
raw_data = ctx.raw_data
ir = ctx.intermediate
proc = ir.proc
if not params.reorder:
# if reordering is turned off, return consecutive order
iorig = np.arange(Nbatch)
return iorig, None, None
nPCs = params.nPCs
Nfilt = ceil(probe.Nchan / 2)
# extract PCA waveforms pooled over channels
wPCA = extractPCfromSnippets(proc, probe=probe, params=params, Nbatch=Nbatch)
Nchan = probe.Nchan
# TODO: move_to_config
niter = 10 # iterations for k-means. we won't run it to convergence to save time
nBatches = Nbatch
NchanNear = min(Nchan, 2 * 8 + 1)
# initialize big arrays on the GPU to hold the results from each batch
# this holds the unit norm templates
Ws = cp.zeros((nPCs, NchanNear, Nfilt, nBatches), dtype=np.float32, order='F')
# this holds the scalings
mus = cp.zeros((Nfilt, nBatches), dtype=np.float32, order='F')
# this holds the number of spikes for that cluster
ns = cp.zeros((Nfilt, nBatches), dtype=np.float32, order='F')
# this holds the center channel for each template
Whs = ones((Nfilt, nBatches), dtype=np.int32, order='F')
i0 = 0
# TODO: move_to_config
NrankPC = 3 # I am not sure if this gets used, but it goes into the function
# return an array of closest channels for each channel
iC = getClosestChannels(probe, params.sigmaMask, NchanNear)[0]
for ibatch in tqdm(range(nBatches), desc="Clustering spikes"):
enough_spikes = False
# extract spikes using PCA waveforms
uproj, call = extractPCbatch2(
proc, params, probe, wPCA, min(nBatches - 2, ibatch), iC, Nbatch)
if cp.sum(cp.isnan(uproj)) > 0:
break # I am not sure what case this safeguards against....
if uproj.shape[1] > Nfilt:
enough_spikes = True
# this initialize the k-means
W, mu, Wheights, irand = initializeWdata2(call, uproj, Nchan, nPCs, Nfilt, iC)
# Params is a whole bunch of parameters sent to the C++ scripts inside a float64 vector
Params = [uproj.shape[1], NrankPC, Nfilt, 0, W.shape[0], 0, NchanNear, Nchan]
for i in range(niter):
Wheights = Wheights.reshape((1, 1, -1), order='F')
iC = cp.atleast_3d(iC)
# we only compute distances to clusters on the same channels
# this tells us which spikes and which clusters might match
iMatch = cp.min(cp.abs(iC - Wheights), axis=0) < .1
# get iclust and update W
# CUDA script to efficiently compute distances for pairs in which iMatch is 1
dWU, iclust, dx, nsp, dV = mexClustering2(Params, uproj, W, mu, call, iMatch, iC)
dWU = dWU / (1e-5 + nsp.T) # divide the cumulative waveform by the number of spike
mu = cp.sqrt(cp.sum(dWU ** 2, axis=0)) # norm of cluster template
W = dWU / (1e-5 + mu) # unit normalize templates
W = W.reshape((nPCs, Nchan, Nfilt), order='F')
nW = W[0, ...] ** 2 # compute best channel from the square of the first PC feature
W = W.reshape((Nchan * nPCs, Nfilt), order='F')
Wheights = cp.argmax(nW, axis=0) # the new best channel of each cluster template
# carefully keep track of cluster templates in dense format
W = W.reshape((nPCs, Nchan, Nfilt), order='F')
W0 = cp.zeros((nPCs, NchanNear, Nfilt), dtype=np.float32, order='F')
for t in range(Nfilt):
W0[..., t] = W[:, iC[:, Wheights[t]], t].squeeze()
# I don't really know why this needs another normalization
W0 = W0 / (1e-5 + cp.sum(cp.sum(W0 ** 2, axis=0)[np.newaxis, ...], axis=1) ** .5)
# if a batch doesn't have enough spikes, it gets the cluster templates of the previous batc
if enough_spikes:
Ws[..., ibatch] = W0
mus[:, ibatch] = mu
ns[:, ibatch] = nsp
Whs[:, ibatch] = Wheights.astype(np.int32)
else:
logger.warning('Data batch #%d only had %d spikes.', ibatch, uproj.shape[1])
i0 = i0 + Nfilt
# anothr one of these Params variables transporting parameters to the C++ code
Params = [1, NrankPC, Nfilt, 0, W.shape[0], 0, NchanNear, Nchan]
# the total number of templates is the number of templates per batch times the number of batch
Params[0] = Ws.shape[2] * Ws.shape[3]
# initialize dissimilarity matrix
ccb = cp.zeros((nBatches, nBatches), dtype=np.float32, order='F')
for ibatch in tqdm(range(nBatches), desc="Computing distances"):
# for every batch, compute in parallel its dissimilarity to ALL other batches
Wh0 = Whs[:, ibatch] # this one is the primary batch
W0 = Ws[..., ibatch]
mu = mus[..., ibatch]
# embed the templates from the primary batch back into a full, sparse representation
W = cp.zeros((nPCs, Nchan, Nfilt), dtype=np.float32, order='F')
for t in range(Nfilt):
W[:, iC[:, Wh0[t]], t] = cp.atleast_3d(Ws[:, :, t, ibatch])
# pairs of templates that live on the same channels are potential "matches"
# TODO: move_to_config? is the 0.1 here important?
iMatch = cp.min(cp.abs(iC - Wh0.reshape((1, 1, -1), order='F')), axis=0) < .1
# compute dissimilarities for iMatch = 1
iclust, ds = mexDistances2(Params, Ws, W, iMatch, iC, Whs, mus, mu)
# ds are squared Euclidian distances
ds = ds.reshape((Nfilt, -1), order='F') # this should just be an Nfilt-long vector
ds = cp.maximum(0, ds)
# weigh the distances according to number of spikes in cluster
ccb[ibatch, :] = cp.mean(cp.sqrt(ds) * ns, axis=0) / cp.mean(ns, axis=0)
# ccb = cp.asnumpy(ccb)
# some normalization steps are needed: zscoring, and symmetrizing ccb
ccb0 = zscore(ccb, axis=0)
ccb0 = ccb0 + ccb0.T
# sort by manifold embedding algorithm
# iorig is the sorting of the batches
# ccbsort is the resorted matrix (useful for diagnosing drift)
ccbsort, iorig = sortBatches2(ccb0)
logger.info("Finished clustering.")
return Bunch(iorig=iorig, ccb0=ccb0, ccbsort=ccbsort)