Exposing Cf amplitude in st[:,3]

kilosort/io.py

@@ -226,7 +226,7 @@ def save_to_phy(st, clu, tF, Wall, probe, ops, imin, results_dir=None,
     Parameters
     ----------
     st : np.ndarray
-        3-column array of peak time (in samples), template, and amplitude for
+        4-column array of peak time (in samples), template, amplitude, and threshold amplitude for
         each spike.
     clu : np.ndarray
         1D vector of cluster ids indicating which spike came from which cluster,

kilosort/run_kilosort.py

@@ -124,8 +124,8 @@ def run_kilosort(settings, probe=None, probe_name=None, filename=None,
     ops : dict
         Dictionary storing settings and results for all algorithmic steps.
     st : np.ndarray
-        3-column array of peak time (in samples), template, and amplitude for
-        each spike.
+        4-column array of peak time (in samples), template, amplitude, and 
+        threshold amplitude for each spike.
     clu : np.ndarray
         1D vector of cluster ids indicating which spike came from which cluster,
         same shape as `st[:,0]`.
@@ -620,8 +620,8 @@ def detect_spikes(ops, device, bfile, tic0=np.nan, progress_bar=None,
     Returns
     -------
     st : np.ndarray
-        3-column array of peak time (in samples), template, and amplitude for
-        each spike.
+        4-column array of peak time (in samples), template, amplitude, and threshold
+        amplitude for each spike.
     clu : np.ndarray
         1D vector of cluster ids indicating which spike came from which cluster,
         same shape as `st`.
@@ -690,8 +690,8 @@ def cluster_spikes(st, tF, ops, device, bfile, tic0=np.nan, progress_bar=None,
     Parameters
     ----------
     st : np.ndarray
-        3-column array of peak time (in samples), template, and amplitude for
-        each spike.
+        4-column array of peak time (in samples), template, amplitud, and threshold
+        amplitude for each spike.
     tF : torch.Tensor
         PC features for each spike, with shape
         (n_spikes, nearest_chans, n_pcs)
@@ -760,8 +760,8 @@ def save_sorting(ops, results_dir, st, clu, tF, Wall, imin, tic0=np.nan,
     results_dir : pathlib.Path
         Directory where results should be saved.
     st : np.ndarray
-        3-column array of peak time (in samples), template, and amplitude for
-        each spike.
+        4-column array of peak time (in samples), template, amplitude, and thresold
+        amplitude for each spike.
     clu : np.ndarray
         1D vector of cluster ids indicating which spike came from which cluster,
         same shape as `st[:,0]`.

kilosort/template_matching.py

@@ -66,7 +66,7 @@ def extract(ops, bfile, U, device=torch.device('cuda'), progress_bar=None):
 
     tiwave = torch.arange(-(nt//2), nt//2+1, device=device) 
     ctc = prepare_matching(ops, U)
-    st = np.zeros((10**6, 3), 'float64')
+    st = np.zeros((10**6, 4), 'float64')
     tF  = torch.zeros((10**6, nC , ops['settings']['n_pcs']))
     k = 0
     prog = tqdm(
@@ -81,7 +81,7 @@ def extract(ops, bfile, U, device=torch.device('cuda'), progress_bar=None):
                 log_performance(logger, 'debug', f'Batch {ibatch}')
 
             X = bfile.padded_batch_to_torch(ibatch, ops)
-            stt, amps, Xres = run_matching(ops, X, U, ctc, device=device)
+            stt, amps, th_amps, Xres = run_matching(ops, X, U, ctc, device=device)
             xfeat = Xres[iCC[:, iU[stt[:,1:2]]],stt[:,:1] + tiwave] @ ops['wPCA'].T
             xfeat += amps * Ucc[:,stt[:,1]]
 
@@ -92,6 +92,7 @@ def extract(ops, bfile, U, device=torch.device('cuda'), progress_bar=None):
                 stt = stt[~neg_spikes,:]
                 xfeat = xfeat[:,~neg_spikes,:]
                 amps = amps[~neg_spikes,:]
+                th_amps = th_amps[~neg_spikes,:]
 
             nsp = len(stt) 
             if k+nsp>st.shape[0]:                     
@@ -101,7 +102,8 @@ def extract(ops, bfile, U, device=torch.device('cuda'), progress_bar=None):
             stt = stt.double()
             st[k:k+nsp,0] = ((stt[:,0]-nt) + ibatch * (ops['batch_size'])).cpu().numpy() - nt//2 + ops['nt0min']
             st[k:k+nsp,1] = stt[:,1].cpu().numpy()
-            st[k:k+nsp,2] = amps[:,0].cpu().numpy()
+            st[k:k+nsp,2] = amps.cpu().numpy().squeeze()
+            st[k:k+nsp,3] = th_amps.cpu().numpy().squeeze()
 
             tF[k:k+nsp]  = xfeat.transpose(0,1).cpu()
 
@@ -180,13 +182,15 @@ def run_matching(ops, X, U, ctc, device=torch.device('cuda')):
 
     st = torch.zeros((100000,2), dtype = torch.int64, device = device)
     amps = torch.zeros((100000,1), dtype = torch.float, device = device)
+    th_amps = torch.zeros((100000,1), dtype = torch.float, device = device)
     k = 0
 
     Xres = X.clone()
     lam = 20
 
     for t in range(100):
         # Cf = 2 * B - nm.unsqueeze(-1) 
+        # Cf is shape (n_units, n_times)
         Cf = torch.relu(B)**2 /nm.unsqueeze(-1)
         #a = 1 + lam
         #b = torch.relu(B) + lam * mu.unsqueeze(-1)
@@ -204,6 +208,7 @@ def run_matching(ops, X, U, ctc, device=torch.device('cuda')):
         cnd2 = torch.abs(Cmax[0,0] - Cfmax) < 1e-9
         xs = torch.nonzero(cnd1 * cnd2)
 
+
         if len(xs)==0:
             #print('iter %d'%t)
             break
@@ -218,6 +223,7 @@ def run_matching(ops, X, U, ctc, device=torch.device('cuda')):
         st[k:k+nsp, 1] = iY[:,0]
         amps[k:k+nsp] = B[iY,iX] / nm[iY]
         amp = amps[k:k+nsp]
+        th_amps[k:k+nsp] = Cmax[0,0,iX[:,0],None]**.5
 
         k+= nsp
 
@@ -230,8 +236,9 @@ def run_matching(ops, X, U, ctc, device=torch.device('cuda')):
 
     st = st[:k]
     amps = amps[:k]
+    th_amps = th_amps[:k]
 
-    return  st, amps, Xres
+    return  st, amps, th_amps, Xres
 
 
 def merging_function(ops, Wall, clu, st, r_thresh=0.5, mode='ccg', device=torch.device('cuda')):