3.1-sort_spikes_concatenate_v0.100.8.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spike sort\n",
    "\n",
    "This notebook is a modified version of the *3-sort_spikes* in the chronic ephys processing pipeline\n",
    "\n",
    "This notebook allows you to concatenate multiple recordings to be spike sorted together. *Be careful, this is only recommended for consecutive recordings on the same date.*\n",
    "\n",
    "Use the environment **spikeproc** to run this notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import pickle\n",
    "os.environ[\"NPY_MATLAB_PATH\"] = '/mnt/cube/chronic_ephys/code/npy-matlab'\n",
    "os.environ[\"KILOSORT2_PATH\"] = '/mnt/cube/chronic_ephys/code/Kilosort2'\n",
    "os.environ[\"KILOSORT3_PATH\"] = '/mnt/cube/chronic_ephys/code/Kilosort'\n",
    "import spikeinterface.full as si\n",
    "import sys\n",
    "import traceback\n",
    "import torch\n",
    "sys.path.append('/mnt/cube/lo/envs/ceciestunepipe/')\n",
    "from ceciestunepipe.file import bcistructure as et\n",
    "from ceciestunepipe.mods import probe_maps as pm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'batch_size': 60000,\n",
       " 'nblocks': 1,\n",
       " 'Th_universal': 9,\n",
       " 'Th_learned': 8,\n",
       " 'do_CAR': True,\n",
       " 'invert_sign': False,\n",
       " 'nt': 61,\n",
       " 'artifact_threshold': None,\n",
       " 'nskip': 25,\n",
       " 'whitening_range': 32,\n",
       " 'binning_depth': 5,\n",
       " 'sig_interp': 20,\n",
       " 'nt0min': None,\n",
       " 'dmin': None,\n",
       " 'dminx': None,\n",
       " 'min_template_size': 10,\n",
       " 'template_sizes': 5,\n",
       " 'nearest_chans': 10,\n",
       " 'nearest_templates': 100,\n",
       " 'templates_from_data': True,\n",
       " 'n_templates': 6,\n",
       " 'n_pcs': 6,\n",
       " 'Th_single_ch': 6,\n",
       " 'acg_threshold': 0.2,\n",
       " 'ccg_threshold': 0.25,\n",
       " 'cluster_downsampling': 20,\n",
       " 'cluster_pcs': 64,\n",
       " 'duplicate_spike_bins': 15,\n",
       " 'do_correction': True,\n",
       " 'keep_good_only': False,\n",
       " 'save_extra_kwargs': False,\n",
       " 'skip_kilosort_preprocessing': False,\n",
       " 'scaleproc': None}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "si.get_default_sorter_params('kilosort4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# non default spike sorting parameters\n",
    "sort_params_dict_ks3 = {'minFR':0.001, 'minfr_goodchannels':0.001} # kilosort 3\n",
    "sort_params_dict_ks4_npx = {'nblocks':5, 'Th_universal':8, 'Th_learned':7} # kilosort 4, neuropixels\n",
    "sort_params_dict_ks4_nnx64 = {'nblocks':0, 'nearest_templates':64,\n",
    "                              'Th_universal':8, 'Th_learned':7} # kilosort 4, neuronexus 64 chan\n",
    "\n",
    "# waveform extraction parameters\n",
    "wave_params_dict = {'ms_before':1, 'ms_after':2, 'max_spikes_per_unit':500,\n",
    "                    'sparse':True, 'num_spikes_for_sparsity':100, 'method':'radius',\n",
    "                    'radius_um':40, 'n_components':5, 'mode':'by_channel_local'}\n",
    "\n",
    "# print stuff\n",
    "verbose = True\n",
    "\n",
    "# errors break sorting\n",
    "raise_error = False\n",
    "\n",
    "# restrict sorting to a specific GPU\n",
    "restrict_to_gpu = 1 # 0 1 2 None\n",
    "\n",
    "# use specific GPU if specified\n",
    "if restrict_to_gpu is not None:\n",
    "    os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "    os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"{}\".format(restrict_to_gpu)\n",
    "\n",
    "# parallel processing params\n",
    "job_kwargs = dict(n_jobs=28,chunk_duration=\"1s\",progress_bar=False)\n",
    "si.set_global_job_kwargs(**job_kwargs)\n",
    "\n",
    "# force processing of previous failed sorts\n",
    "skip_failed = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "bird_rec_dict = {\n",
    "    'z_y19o20_21':[\n",
    "        {'sess_par_list':['2021-10-27'], # sessions (will process all epochs within)\n",
    "         'probe':{'probe_type':'neuropixels-1.0'}, # probe specs\n",
    "         'sort':'sort_0', # label for this sort instance\n",
    "         'sorter':'kilosort4', # sort method\n",
    "         'sort_params':sort_params_dict_ks4_npx, # non-default sort params\n",
    "         'wave_params':wave_params_dict, # waveform extraction params\n",
    "         'ephys_software':'sglx' # sglx or oe\n",
    "        },\n",
    "    ],\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run sorts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# store sort summaries\n",
    "sort_summary_all = []\n",
    "\n",
    "# loop through all birds / recordings\n",
    "for this_bird in bird_rec_dict.keys():\n",
    "    # get session configurations\n",
    "    sess_all = bird_rec_dict[this_bird]\n",
    "    \n",
    "    # loop through session configurations\n",
    "    for this_sess_config in sess_all:\n",
    "        \n",
    "        # loop through sessions\n",
    "        for this_sess in this_sess_config['sess_par_list']:\n",
    "            log_dir = os.path.join('/mnt/cube/chronic_ephys/log', this_bird, this_sess)\n",
    "            \n",
    "            # build session parameter dictionary\n",
    "            sess_par = {'bird':this_bird,\n",
    "                        'sess':this_sess,\n",
    "                        'ephys_software':this_sess_config['ephys_software'],\n",
    "                        'sorter':this_sess_config['sorter'],\n",
    "                        'sort':this_sess_config['sort']}\n",
    "            # get epochs\n",
    "            sess_epochs = et.list_ephys_epochs(sess_par)\n",
    "            concat_epochs = '-'.join(sess_epochs) # process all epochs together -- edit to concatenate only some if desired\n",
    "            \n",
    "            # set output directory\n",
    "            epoch_struct = et.sgl_struct(sess_par,concat_epochs,ephys_software=sess_par['ephys_software'])\n",
    "            sort_path = epoch_struct['folders']['derived'] + '/{}/{}/'.format(sess_par['sorter'],sess_par['sort'])\n",
    "\n",
    "            # get spike sort log\n",
    "            try:\n",
    "                with open(os.path.join(log_dir,concat_epochs+'_spikesort_'+this_sess_config['sort']+'.log'), 'r') as f:\n",
    "                    log_message=f.readline() # read the first line of the log file\n",
    "                if log_message[:-1] == sess_par['bird']+' '+sess_par['sess']+' sort complete without error':\n",
    "                    print(sess_par['bird'],sess_par['sess'],'already exists -- skipping sort')\n",
    "                    run_proc = False\n",
    "                elif log_message[:-1] == sess_par['bird']+' '+sess_par['sess']+' sort failed':\n",
    "                    if skip_failed:\n",
    "                        print(sess_par['bird'],sess_par['sess'],'previously failed -- skipping sort')\n",
    "                        run_proc = False\n",
    "                    else:\n",
    "                        run_proc = True\n",
    "                else: # uninterpretable log file\n",
    "                    run_proc = True\n",
    "            except: # no existing log file\n",
    "                run_proc = True\n",
    "\n",
    "            # run sort\n",
    "            if run_proc:\n",
    "                try:\n",
    "                    print('___________',this_bird,this_sess,concat_epochs,'___________')\n",
    "                    # prepare recording for sorting\n",
    "                    print('prep..')\n",
    "                    if sess_par['ephys_software'] == 'sglx':\n",
    "                        # load recordings\n",
    "                        rec_path = '/'.join(epoch_struct['folders']['sglx'].split('/')[:-1])\n",
    "                        epoch_list = [e for e in os.listdir(rec_path) if os.path.isdir(os.path.join(rec_path,e))]\n",
    "                        recording_list = []\n",
    "                        for this_epoch in epoch_list: # load all epochs in a session -- edit to load only some if desired\n",
    "                            this_rec = si.read_spikeglx(folder_path=os.path.join(rec_path,this_epoch), stream_name='imec0.ap')\n",
    "                            probe_df = this_rec.get_probe().to_dataframe()\n",
    "                            # ibl destriping\n",
    "                            this_rec = si.highpass_filter(recording=this_rec)\n",
    "                            this_rec = si.phase_shift(recording=this_rec)\n",
    "                            bad_good_channel_ids = si.detect_bad_channels(recording=this_rec)\n",
    "                            if len(bad_good_channel_ids[0]) > 0:\n",
    "                                this_rec = si.interpolate_bad_channels(recording=this_rec,bad_channel_ids=bad_good_channel_ids[0])\n",
    "                            if this_sess_config['probe']['probe_type'] == 'neuropixels-2.0':\n",
    "                                # highpass by shank\n",
    "                                split_rec = this_rec.split_by(property='group',outputs='list') # split recording by shank\n",
    "                                split_rec = [si.highpass_spatial_filter(recording=r,n_channel_pad=min(r.get_num_channels(),60)) for r in split_rec]\n",
    "                                this_rec = si.aggregate_channels(split_rec) # recombine shanks\n",
    "                                # stack shanks\n",
    "                                p,_ = pm.stack_shanks(probe_df) # make new Probe object with shanks stacked\n",
    "                                this_rec = this_rec.set_probe(p,group_mode='by_probe') # assign new Probe object to probe\n",
    "                            else:\n",
    "                                this_rec = si.highpass_spatial_filter(recording=this_rec)\n",
    "                            recording_list.append(this_rec)\n",
    "                        \n",
    "                        # concatenate recordings\n",
    "                        this_rec_p = si.concatenate_recordings(recording_list)\n",
    "                        # save probe map prior to re-ordering for sorting\n",
    "                        probe_df.to_pickle(os.path.join(epoch_struct['folders']['derived'],'probe_map_df.pickle'))\n",
    "                    else:\n",
    "                        raise TypeError('Only implemented for sglx')\n",
    "                    # set sort params\n",
    "                    this_rec_p = si.concatenate_recordings([this_rec_p])\n",
    "                    sort_params = si.get_default_sorter_params(this_sess_config['sorter'])\n",
    "                    for this_param in this_sess_config['sort_params'].keys():\n",
    "                        sort_params[this_param] = this_sess_config['sort_params'][this_param]\n",
    "                    # run sort\n",
    "                    print('sort..')\n",
    "                    this_sort = si.run_sorter(sorter_name=this_sess_config['sorter'],recording=this_rec_p,output_folder=sort_path,\n",
    "                                         remove_existing_folder=True,delete_output_folder=False,delete_container_files=False,\n",
    "                                         verbose=verbose,raise_error=raise_error,**sort_params)\n",
    "                    # bandpass recording before waveform extraction\n",
    "                    print('bandpass..')\n",
    "                    this_rec_pf = si.bandpass_filter(recording=this_rec_p)\n",
    "                    # extract waveforms\n",
    "                    print('waveform..')\n",
    "                    wave_params = this_sess_config['wave_params']\n",
    "                    wave = si.extract_waveforms(this_rec_pf,this_sort,folder=os.path.join(sort_path,'waveforms'),\n",
    "                                                ms_before=wave_params['ms_before'],ms_after=wave_params['ms_after'],\n",
    "                                                max_spikes_per_unit=wave_params['max_spikes_per_unit'],\n",
    "                                                sparse=wave_params['sparse'],num_spikes_for_sparsity=wave_params['num_spikes_for_sparsity'],\n",
    "                                                method=wave_params['method'],radius_um=wave_params['radius_um'],overwrite=True,**job_kwargs)\n",
    "                    # compute metrics\n",
    "                    print('metrics..')\n",
    "                    loc = si.compute_unit_locations(waveform_extractor=wave)\n",
    "                    cor = si.compute_correlograms(waveform_or_sorting_extractor=wave)\n",
    "                    sim = si.compute_template_similarity(waveform_extractor=wave)\n",
    "                    amp = si.compute_spike_amplitudes(waveform_extractor=wave,**job_kwargs)\n",
    "                    pca = si.compute_principal_components(waveform_extractor=wave,n_components=wave_params['n_components'],\n",
    "                                                          mode=wave_params['mode'],**job_kwargs)\n",
    "                    qms = si.get_quality_metric_list()\n",
    "                    metric_names = []\n",
    "                    bad_metrics = []\n",
    "                    for qm in qms:\n",
    "                        try:\n",
    "                            si.compute_quality_metrics(waveform_extractor=wave,verbose=False,metric_names=[qm],**job_kwargs)\n",
    "                            metric_names.append(qm)\n",
    "                        except:\n",
    "                            bad_metrics.append(qm)\n",
    "                    met = si.compute_quality_metrics(waveform_extractor=wave,verbose=verbose,metric_names=metric_names,**job_kwargs)\n",
    "\n",
    "                    # mark complete\n",
    "                    print('COMPLETE!!')\n",
    "\n",
    "                    # log complete sort\n",
    "                    if not os.path.exists(log_dir): os.makedirs(log_dir)\n",
    "                    with open(os.path.join(log_dir,concat_epochs+'_spikesort_'+this_sess_config['sort']+'.log'), 'w') as f:\n",
    "                        f.write(sess_par['bird']+' '+sess_par['sess']+' sort complete without error\\n\\n')\n",
    "                        f.write('Sort method: '+this_sess_config['sorter']+'\\n\\n')\n",
    "                        f.write('Sort params: '+str(sort_params)+'\\n\\n')\n",
    "                        f.write('Computed quality metrics: '+str(metric_names)+'\\n\\n')\n",
    "                        f.write('Failed quality metrics: '+str(bad_metrics)+'\\n')\n",
    "                    sort_summary = [this_bird,this_sess,sess_par['ephys_software'],concat_epochs,'COMPLETE']\n",
    "\n",
    "                except Exception as e:\n",
    "                    # mark exception\n",
    "                    print(\"An exception occurred:\", e)\n",
    "\n",
    "                    # log failed sort\n",
    "                    if not os.path.exists(log_dir): os.makedirs(log_dir)\n",
    "                    with open(os.path.join(log_dir,concat_epochs+'_spikesort_'+this_sess_config['sort']+'.log'), 'w') as f:\n",
    "                        f.write(sess_par['bird']+' '+sess_par['sess']+' sort failed\\n')\n",
    "                        f.write(traceback.format_exc())\n",
    "                    sort_summary = [this_bird,this_sess,sess_par['ephys_software'],concat_epochs,'FAIL']\n",
    "            else:\n",
    "                sort_summary = [this_bird,this_sess,sess_par['ephys_software'],concat_epochs,'EXISTS']\n",
    "\n",
    "            # report and store sort summary\n",
    "            print(sort_summary)\n",
    "            sort_summary_all.append(sort_summary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "spikeproc",
   "language": "python",
   "name": "spikeproc"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}