<jupyter_start><jupyter_text>Sort Data<jupyter_code>def sortData(recording, storeDir, sorter_list = ['MS4','KS2', 'SC'], dthr = 5 , ks2Path = '/home/alexgonzalez/Documents/MATLAB/Kilosort2/' ):
    
    valid_sorters = ['MS4','KS2', 'SC', 'HS','TDC','Klusta']
    
    # the timeout is the greater between 3 mins or 10 percent of the recording in seconds 
    timeout_dur = int(np.max([0.1*recording.get_num_frames()/recording.get_sampling_frequency(),180]))
    out = {}
    
    for sorter in sorter_list:
        if sorter in valid_sorters:
            t0 = time.time()
            if sorter is 'MS4':
                params = ss.Mountainsort4Sorter.default_params()
                params['detect_threshold'] = dthr
                params['curation'] = True
                params['whiten'] = False
                algorithm = ss.run_mountainsort4
            elif sorter is 'IC':
                params = ss.IronClustSorter.default_params()
                params['detect_threshold'] = dthr
                params['filter'] = False
            elif sorter is 'KS2':
                params = ss.Kilosort2Sorter.default_params()
                params['detect_threshold'] = dthr
                params['car'] = False
                ss.Kilosort2Sorter.set_kilosort2_path(ks2Path)
                algorithm = ss.run_kilosort2
            elif sorter is 'SC':
                params = ss.SpykingcircusSorter.default_params()
                params['filter'] = False
                params['detect_threshold'] = dthr
                algorithm = ss.run_spykingcircus
            elif sorter is 'HS':
                params  = ss.HerdingspikesSorter.default_params()
                params['pca_whiten'] = False
                params['filter'] = False
                algorithm = ss.run_herdingspikes
            elif sorter is 'TDC':
                params = ss.TridesclousSorter.default_params()
                params['relative_threshold'] = dthr
                algorithm = ss.run_tridesclous
            elif sorter is 'Klusta':
                params = ss.KlustaSorter.default_params()
                params['num_starting_clusters'] = 20
                algorithm = ss.run_klusta

            try:
                with timeout(seconds=timeout_dur):
                    print()
                    print('Sorting with {}'.format(sorter))
                    out[sorter] = algorithm(recording=recording, **params, output_folder=storeDir+'_'+sorter)
                    t1 = time.time()
                    print('Time to sort using {0} = {1:0.2f}s'.format(sorter,t1-t0))
                    print('Found {} clusters'.format( len(out[sorter].get_unit_ids()) ))

            except TimeoutError:
                print()
                print('Sorter {} Timeout.'.format(sorter))
                out[sorter] = None
                continue
            except:
                print()
                print('Error running sorter {}'.format(sorter))
                print("Error", sys.exc_info()[0],sys.exc_info()[1],sys.exc_info()[2].tb_lineno)
                out[sorter] = None
                continue
        else:
            print()
            print('Invalid Sorter'.format(sorter))
            out[sorter] = None
        
    return out

def get_signals_mad(signals, mask=None):
    
    if signals.ndim==1:
        if mask is None:
            return mad(signals)
        else:
            return mad(signals[mask])
        
    elif signals.ndim>1:
        nChans, nSamps = signals.shape
        
    sig_mad = np.zeros(nChans)
    if mask is None:
        for ch in range(nChans):
            sig_mad[ch] = mad(signals[ch])
    elif mask.shape[0]==nChans:
        for ch in range(nChans):
            sig_mad[ch] = mad(signals[ch,mask[ch]])
    elif mask.ndim==1 and len(mask)==nSamps:
        for ch in range(nChans):
            sig_madp[ch] = mad(signals[ch],mask)
    else:
        print('Mask does not match the data. Ignoring,')
        for ch in range(nChans):
            sig_mad[ch] = mad(signals[ch])
            
    return sig_mad

def get_cluster_snr(spike_train, hp_signals, sig_mad=None, mask = None):
    """
    Inputs:
    spike_train -> array of samples indicating the peak of the detected spike
    hp_signals -> high passed signal for spike detection
    sig_mad -> the median absolute deviation for each channel, if not passed, it will be computed
    mask -> mask of valid values for the hp_signal
    
    Returns:
    snr -> float indicating the SNR (higher number the better.)
    amp_mad -> measure of the stability of the amplitudes (lower the number the better)
    """
    amps = hp_signals[:,spike_train] # amplitdes of the spikes
    median_amps = np.abs(np.median(amps,axis=1)) # median amplitude of spikes for all channels
    
    max_amp_ch = np.argmax(median_amps) # chan with max median amp
    max_amp = np.max(median_amps) # value of the max median amp
    
    amp_mad = get_signals_mad(amps[max_amp_ch]) # deviation of the median amp
    
    if sig_mad is None:
        sig_mad = mad(hp_signals[max_amp_ch], mask[max_am_ch])
    else:
        sig_mad = sig_mad[max_amp_ch]
        
        
    snr = max_amp/sig_mad # signal to noise ratio
    
    return snr, max_amp_ch, max_amp, amp_mad

def get_sorter_cluster_snr(sorter, hp_signals, sig_mad=None, mask=None):
        cluster_ids = sorter.get_unit_ids()
        n_clusters = len(cluster_ids)
        
        cl_snrs = {}
        cl_amp_mad = {}
        for cl in cluster_ids:
            spike_train = sorter.get_unit_spike_train(cl)
            cl_snrs[cl], cl_amp_mad[cl] = get_cluster_snr(spike_train,hp_signals,sig_mad, mask)
        return cl_snrs, cl_amp_mad

def coeff_variation_ln(isi):
    sln = np.std(np.log(isi))
    return np.sqrt(np.exp(sln**2) - 1 )

def plot_isi_dist(spk_train,fs = 32000,dt = 2, ax=None):
    if ax is None:
        fig, ax = plt.subplots()
    
    xlims = [0, 50]
    
    isi = np.diff(spk_train)/fs*1000
    bins = np.arange(0, xlims[1], dt)
    ax = sns.distplot(isi,bins=bins, norm_hist=False, kde=True, kde_kws = {'bw':dt, 'lw': 3, 'clip':xlims})
    ax.set_xlim(xlims)
    ax.axvline(0,color='k',linestyle='--',linewidth=2, alpha=0.5)
    ax.set_xlabel('ISI [ms]')
    
    return ax

storeDir = Path(dataRoot,'PreProcessed',subject,session,'tt_{}'.format(tt))
storeDir.mkdir(parents=True, exist_ok=True)

tt_masks = create_chan_masks(tt_info['Raw']['ClippedSegs'],rawSignals)
tt_mad = get_signals_mad(spk_signals,tt_masks)

tt_info['Spk']['mad'] = tt_mad

spk_masked_signals = se.NumpyRecordingExtractor(timeseries=spk_signals*tt_masks, geom=tt_info['tt_geom'], sampling_frequency=fs)
import spikesorters as ss
ss.IronClustSorter.default_params()
sortDict = sortData(spk_masked_signals, storeDir=str(storeDir / 'tmp'))
#ss = sortData(spk_masked_signals, sorter_list = ['KS2', 'HS'], storeDir=str(storeDir / 'tmp'))
#sortDict['KS2']=ss['KS2']
#del ss
cluster_stats = pd.DataFrame()
cluster_num = 0
for sorter_id, sort in sortDict.items():
    units = sort.get_unit_ids()
    if len(units)>0:
        for unit in units:
            key = sorter_id + '_' + str(unit)
            cluster_stats.at[key,'cl_num'] = cluster_num
            cluster_stats.at[key,'sorter'] = sorter_id
            cluster_stats.at[key,'sorter_cl_num'] = unit
            
            unit_spk_train = sort.get_unit_spike_train(unit)
            cluster_stats.at[key,'fr'] = len(unit_spk_train)/tt_info['n_samps']*tt_info['fs']
            
            n_spikes = len(unit_spk_train)
            isi = np.diff(unit_spk_train)/tt_info['fs']*1000
            c,b=np.histogram(isi,np.arange(50))
            isi_viol = np.sum(c[b[:-1]<=2])/n_spikes
            cv = coeff_variation_ln(isi)
            
            cluster_stats.at[key, 'isi_viol'] = isi_viol
            cluster_stats.at[key, 'cv'] = cv
            cluster_stats.at[key, 'n_spikes'] = n_spikes
            
            cluster_stats.at[key, 'snr'] , cluster_stats.at[key, 'amp_ch'] ,  cluster_stats.at[key, 'amp_med_ch'],  cluster_stats.at[key, 'amp_mad_ch']  \
            = get_cluster_snr(unit_spk_train, spk_signals, sig_mad = tt_info['Spk']['mad'],mask = tt_masks)
            
            cluster_num +=1
            
cluster_stats
snr_thr = 5
fr_low_thr = 0.5
fr_high_thr = 90
isi_viol_thr = 0.5

curated_sorts = {}
for sorter_id, sort in sortDict.items():
    t0 = time.time()
    print('{} Number of units before automatic curation = {}'.format(sorter_id,len(sort.get_unit_ids())))
    # fr thresholding
    if len(sort.get_unit_ids())>0:
        sort = st.curation.threshold_firing_rate(sort,threshold=fr_low_thr, threshold_sign='less')
    if len(sort.get_unit_ids())>0:
        sort = st.curation.threshold_firing_rate(sort,threshold=fr_high_thr, threshold_sign='greater')
    
    #isi thresholding
    if len(sort.get_unit_ids())>0:
        sort = st.curation.threshold_isi_violations(sort,threshold=isi_viol_thr,threshold_sign='less')
    
    # snr thrsholding
    if len(sort.get_unit_ids())>0:
        cl_snr_dict, _ = get_sorter_cluster_snr(sort, spk_signals, sig_mad = tt_info['Spk']['mad'],mask = tt_masks)
        cl_snr = np.array([[cl, cl_snr_dict[cl]] for cl in cl_snr_dict.keys()])
        bad_units = cl_snr[cl_snr[:,1]>snr_thr,0]
        if len(bad_units)>0:
            sort.exclude_units(bad_units)

    print('{} Number of units after automatic curation = {}'.format(sorter_id,len(sort.get_unit_ids())))
    curated_sorts[sorter_id]=sort
    t1 = time.time()
    print('Tim
plot_isi_dist(unit_spk_train)
def get_waveforms(spk_train, hp_data, wf_lims = [-12,20], n_wf=5000):
    n_spikes = len(spk_train)
    n_chans = hp_data.shape[0]
    
    sampled_spikes = spk_train[np.random.randint(n_spikes,size=n_wf)]
    wf_samps = np.arange(wf_lims[0],wf_lims[1])
    
    mwf = np.zeros( (n_chans, len(wf_samps)) )
    for samp_spk in sampled_spikes:
        mwf += hp_data[:,wf_samps+samp_spk]
    mwf /= n_wf
    
    return mwf

wf = get_waveforms(unit_spk_train, spk_masked_signals.get_traces())
amps = pd.DataFrame(spk_masked_signals.get_traces()[:,unit_spk_train].T,columns=['ch'+str(ch) for ch in range(4)]) 
ax = sns.boxplot(data=amps)
#ax = sns.stripplot(data=amps, color='0.5',alpha=0.01,ax=ax)
sns.catplot(data=amps,kind='bar')
mad(amps['ch0'])
n_spikes = len(unit_spk_train)
spk_samples = np.random.randint(n_spikes,size=1000)
spk_masked_signals.get_traces()[:,unit_spk_train[spk_samples]].mean(axis=1)
spk_masked_signals.get_traces()[:,unit_spk_train].mean(axis=1)
spkTrains = {}

for sorter_id, sort in sortDict.items():
    units = sort.get_unit_ids()
    if len(units)>0:
        for unit in units:
            key = sorter_id + '_' + str(unit)
            spkTrains[key] = sort.get_unit_spike_train(unit)

mc
mc = {}
for sorter_id, sort  in sortDict.items():
    t0 = time.time()
    temp = st.validation.MetricCalculator(sort,spk_masked_signals)
    temp.compute_metrics (metric_names=['firing_rate'])
    mc[sorter_id] = temp.get_metrics_df()
    t1 = time.time()
    print('Time to compute metrics for {0} is {1:0.2f} :'.format(sorter_id,t1-t0))

    
mc
plot_isi_dist(spkTrains['KS2_8'])
keys = list(spkTrains.keys())
nKeys = len(keys)
    
def acf(x, length=1000):
    return np.array([1]+[np.corrcoef(x[:-i], x[i:])[0,1] for i in range(1, length)])

plt.plot(acf(np.diff(spkTrains['KS2_8'])))
from spectrum import xcorr

fs*0.01
keys

win = 20 # in ms
q = int(fs*win/1000)
x = np.zeros(nSamps)
x[spkTrains[keys[1]]] = 1
x2 = scipy.signal.decimate(x,q,ftype='fir')

y = np.zeros(nSamps)
y[spkTrains[keys[2]]] = 1
y2 = scipy.signal.decimate(y,q,ftype='fir')
xc = np.correlate(x2,y2,mode='full')
plt.plot(xc)
c,l = xcorr(x2,y2,maxlags=100, norm='coeff')
# Mountainsort4 spike sorting
t0 = time.time()
default_ms4_params = ss.Mountainsort4Sorter.default_params()
default_ms4_params['detect_threshold'] = 4
default_ms4_params['curation'] = False
default_ms4_params['whiten'] = False

sorting_MS4 = ss.run_mountainsort4(recording=hp_masked_recording, **default_ms4_params, output_folder=str(storeDir / 'tmp_MS4'))
t1 = time.time()

print('Time to sort using MS4 = {0:0.2f}s'.format(t1-t0))
print('Units found with MS4:', sorting_MS4.get_unit_ids())
ms4_mc = st.validation.MetricCalculator(sorting_MS4,recording)
ms4_mc.compute_metrics()
print(ms4_mc.get_metrics_df())
# SpyKING Circus spike sorting
t0 = time.time()
SC_params = ss.SpykingcircusSorter.default_params()
SC_params['filter'] = False
SC_params['detect_threshold'] = 4

sorting_SC = ss.run_spykingcircus(recording=recording, **SC_params, output_folder=str(storeDir / 'tmp_SC'))
t1 = time.time()

print('Time to sort using SC = {0:0.2f}s'.format(t1-t0))
print('Units found with Spyking Circus:', sorting_SC.get_unit_ids())
SC_mc = st.validation.MetricCalculator(sorting_SC,recording)
SC_mc.compute_metrics()
print(SC_mc.get_metrics_df())
# Kilosort2 spike sorting (KILOSORT2_PATH and NPY_MATLAB_PATH can be set as environment variables)
#import os
#os.environ["KILOSORT2_PATH"] = "/home/alexgonzalez/Documents/MATLAB/Kilosort2/"
#os.environ["NPY_MATLAB_PATH"] = "/home/alexgonzalez/Documents/MATLAB/npy-matlab/"
t0 = time.time()
ks_params =  ss.Kilosort2Sorter.default_params()
ks_params['detect_threshold'] = 4
ks_params['car'] = False

ss.Kilosort2Sorter.set_kilosort2_path('/home/alexgonzalez/Documents/MATLAB/Kilosort2/')
sorting_KS2 = ss.run_kilosort2(recording, **ks_params, output_folder= str(storeDir / 'tmp_KS2'))
t1 = time.time()
print('Time to sort using KS2 = {0:0.2f}s'.format(t1-t0))
print('Units found with Kilosort2', sorting_KS2.get_unit_ids())

# Tridesclous spike sorting
ss = reload(ss)
TDC_params =  ss.TridesclousSorter.default_params()
TDC_params['relative_threshold'] = 4

shutil.rmtree(str(storeDir / 'tmp_TDC'),ignore_errors=True,onerror=None)

to = time.time()
sorting_TDC = ss.run_tridesclous(recording, **TDC_params,output_folder=str(storeDir / 'tmp_TDC'))
t1 = time.time()
print('Time to sort using TDC = {0:0.2f}s'.format(t1-t0))
print('Units found with Tridesclous:', sorting_TDC.get_unit_ids())

# herding spikes
t0 = time.time()
sorting_HS = ss.run_herdingspikes(recording,output_folder=str(storeDir / 'tmp_HS'))
t1 = time.time()
print('Time to sort using HerdingSpikes = {0:0.2f}s'.format(t1-t0))
print('Units found with HerdingSpikes:', sorting_HS.get_unit_ids())

HS_mc = st.validation.MetricCalculator(sorting_HS,recording)
HS_mc.compute_metrics()
print(HS_mc.get_metrics_df())
# Klusta
t0 = time.time()
#klusta_params = ss.KlustaSorter.default_params()
sorting_klusta = ss.run_klusta(recording,output_folder=str(storeDir / 'tmp_klusta'))
t1 = time.time()
print('Time to sort using Klusta = {0:0.2f}s'.format(t1-t0))
print('Units found with Klusta:', sorting_klusta.get_unit_ids())

resting
# compare SC KS MS4
import spikecomparison as sc
import spikewidgets as sw
mcmp = sc.compare_multiple_sorters(sorting_list=[sorting_KS2,sorting_MS4,sorting_SC,sorting_hs],name_list=['KS2','MS4','SC','HS'],verbose=True)
sw.plot_multicomp_graph(mcmp)
agr_2 = mcmp.get_agreement_sorting(minimum_matching=2)
agr_2.get_unit_ids()
comp1 = sc.compare_two_sorters(sorting1=sorting_KS2,sorting2=sorting_MS4,sorting1_name='KS2',sorting2_name='MS4')
comp2 = sc.compare_two_sorters(sorting1=sorting_SC,sorting2=sorting_hs,sorting1_name='SC',sorting2_name='HS')
comp3 = sc.compare_two_sorters(sorting1=sorting_SC,sorting2=sorting_MS4,sorting1_name='SC',sorting2_name='MS4')
comp4 = sc.compare_two_sorters(sorting1=sorting_HS,sorting2=sorting_MS4,sorting1_name='HS',sorting2_name='MS4')

sw.plot_agreement_matrix(comp1)
sw.plot_agreement_matrix(comp2)

sw.plot_agreement_matrix(comp3)

sw.plot_agreement_matrix(comp4)
w_wf = sw.plot_unit_waveforms(recording, sorting_MS4, max_spikes_per_unit=100)

        
int(np.max([0.1*recording.get_num_frames()/recording.get_sampling_frequency(),180]))
180/190<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13095087763606517, 'fraction_numerical': 0.01843667129688092, 'mean_word_length': 3.4156676888764985, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 20, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '55469', 'license': 'MIT', 'n_tokens_mistral': 6050, 'n_tokens_neox': 5818, 'n_words': 981, 'path': 'Notebooks/spikeinterface_examples/preprocess_and_sort.ipynb', 'repo_name': 'alexgonzl/TreeMazeAnalyses', 'size': '15838'}