3-sort_spikes_v0.100.8.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spike sort\n",
    "\n",
    "Notebook within the chronic ephys processing pipeline\n",
    "- 1-preprocess_acoustics\n",
    "- 2-curate_acoustics\n",
    "- **3-sort_spikes**\n",
    "- 4-curate_spikes\n",
    "\n",
    "Use the environment **spikeproc** to run this notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": "markdown",
   "metadata": {},
   "source": [
    "#### Set `dmin` and `dminx`\n",
    "**Setting these appropriately will greatly reduce sort time**\n",
    "- The default value for dmin is the median distance between contacts -- if contacts are irregularly spaced, like in a modular Neuropixels 2.0 setup, will need to specify a value\n",
    "- The default for dminx is 32um (designed for Neuropixels probes)\n",
    "\n",
    "Support documentation [here](https://kilosort.readthedocs.io/en/latest/parameters.html#dmin-and-dminx)"
   ]
  },
  {
   "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, 'dmin':15, 'dminx':32} # kilosort 4, neuropixels (set dmin and dminx to true pitch)\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 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_p5y10_23':[\n",
    "        {'sess_par_list':['2024-05-16'], # sessions (will process all epochs within)\n",
    "         'probe':{'probe_type':'neuropixels-2.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": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "___________ z_p5y10_23 2024-05-16 1246_g0 ___________\n",
      "prep..\n",
      "sort..\n",
      "========================================\n",
      "Loading recording with SpikeInterface...\n",
      "number of samples: 368121306\n",
      "number of channels: 384\n",
      "number of segments: 1\n",
      "sampling rate: 30000.0\n",
      "dtype: int16\n",
      "========================================\n",
      "Preprocessing filters computed in  2077.34s; total  2077.34s\n",
      "\n",
      "computing drift\n",
      "Re-computing universal templates from data.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 6136/6136 [15:52:26<00:00,  9.31s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "drift computed in  59153.07s; total  61230.41s\n",
      "\n",
      "Extracting spikes using templates\n",
      "Re-computing universal templates from data.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 6136/6136 [16:23:41<00:00,  9.62s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "19999820 spikes extracted in  61300.19s; total  122530.61s\n",
      "\n",
      "First clustering\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 107/107 [19:49<00:00, 11.11s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1217 clusters found, in  1197.43s; total  123728.03s\n",
      "\n",
      "Extracting spikes using cluster waveforms\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 6136/6136 [14:52:18<00:00,  8.73s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37481401 spikes extracted in  53546.84s; total  177274.87s\n",
      "\n",
      "Final clustering\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 107/107 [26:00<00:00, 14.58s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "829 clusters found, in  1560.73s; total  178835.74s\n",
      "\n",
      "Merging clusters\n",
      "742 units found, in  140.57s; total  178976.31s\n",
      "\n",
      "Saving to phy and computing refractory periods\n",
      "338 units found with good refractory periods\n",
      "\n",
      "Total runtime: 179044.84s = 49:2984:4 h:m:s\n",
      "kilosort4 run time 179047.45s\n",
      "bandpass..\n",
      "waveform..\n",
      "metrics..\n",
      "An exception occurred: [Errno 13] Permission denied: '/tmp/spikeinterface_cache/tmpnq7ml3iz'\n",
      "['z_p5y10_23', '2024-05-16', 'sglx', '1246_g0', 'FAIL']\n",
      "___________ z_p5y10_23 2024-05-16 1611_g0 ___________\n",
      "prep..\n",
      "sort..\n",
      "========================================\n",
      "Loading recording with SpikeInterface...\n",
      "number of samples: 158273746\n",
      "number of channels: 384\n",
      "number of segments: 1\n",
      "sampling rate: 30000.0\n",
      "dtype: int16\n",
      "========================================\n",
      "Preprocessing filters computed in  878.17s; total  878.28s\n",
      "\n",
      "computing drift\n",
      "Re-computing universal templates from data.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2638/2638 [7:23:47<00:00, 10.09s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "drift computed in  27475.47s; total  28353.88s\n",
      "\n",
      "Extracting spikes using templates\n",
      "Re-computing universal templates from data.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 48%|████████████████████████████████████████████████████████████████████████████▉                                                                                  | 1276/2638 [3:13:08<3:26:09,  9.08s/it]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "File \u001b[0;32m<timed exec>:93\u001b[0m\n",
      "File \u001b[0;32m/mnt/cube/lo/envs/spikeproc/lib/python3.8/site-packages/spikeinterface/sorters/runsorter.py:175\u001b[0m, in \u001b[0;36mrun_sorter\u001b[0;34m(sorter_name, recording, output_folder, remove_existing_folder, delete_output_folder, verbose, raise_error, docker_image, singularity_image, delete_container_files, with_output, **sorter_params)\u001b[0m\n\u001b[1;32m    168\u001b[0m             container_image \u001b[38;5;241m=\u001b[39m singularity_image\n\u001b[1;32m    169\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m run_sorter_container(\n\u001b[1;32m    170\u001b[0m         container_image\u001b[38;5;241m=\u001b[39mcontainer_image,\n\u001b[1;32m    171\u001b[0m         mode\u001b[38;5;241m=\u001b[39mmode,\n\u001b[1;32m    172\u001b[0m         \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mcommon_kwargs,\n\u001b[1;32m    173\u001b[0m     )\n\u001b[0;32m--> 175\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrun_sorter_local\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcommon_kwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/mnt/cube/lo/envs/spikeproc/lib/python3.8/site-packages/spikeinterface/sorters/runsorter.py:225\u001b[0m, in \u001b[0;36mrun_sorter_local\u001b[0;34m(sorter_name, recording, output_folder, remove_existing_folder, delete_output_folder, verbose, raise_error, with_output, **sorter_params)\u001b[0m\n\u001b[1;32m    223\u001b[0m SorterClass\u001b[38;5;241m.\u001b[39mset_params_to_folder(recording, output_folder, sorter_params, verbose)\n\u001b[1;32m    224\u001b[0m SorterClass\u001b[38;5;241m.\u001b[39msetup_recording(recording, output_folder, verbose\u001b[38;5;241m=\u001b[39mverbose)\n\u001b[0;32m--> 225\u001b[0m \u001b[43mSorterClass\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_from_folder\u001b[49m\u001b[43m(\u001b[49m\u001b[43moutput_folder\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mraise_error\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mverbose\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    226\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m with_output:\n\u001b[1;32m    227\u001b[0m     sorting \u001b[38;5;241m=\u001b[39m SorterClass\u001b[38;5;241m.\u001b[39mget_result_from_folder(output_folder, register_recording\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m, sorting_info\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
      "File \u001b[0;32m/mnt/cube/lo/envs/spikeproc/lib/python3.8/site-packages/spikeinterface/sorters/basesorter.py:258\u001b[0m, in \u001b[0;36mBaseSorter.run_from_folder\u001b[0;34m(cls, output_folder, raise_error, verbose)\u001b[0m\n\u001b[1;32m    255\u001b[0m t0 \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mperf_counter()\n\u001b[1;32m    257\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 258\u001b[0m     \u001b[43mSorterClass\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_from_folder\u001b[49m\u001b[43m(\u001b[49m\u001b[43msorter_output_folder\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msorter_params\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mverbose\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    259\u001b[0m     t1 \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mperf_counter()\n\u001b[1;32m    260\u001b[0m     run_time \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mfloat\u001b[39m(t1 \u001b[38;5;241m-\u001b[39m t0)\n",
      "File \u001b[0;32m/mnt/cube/lo/envs/spikeproc/lib/python3.8/site-packages/spikeinterface/sorters/external/kilosort4.py:260\u001b[0m, in \u001b[0;36mKilosort4Sorter._run_from_folder\u001b[0;34m(cls, sorter_output_folder, params, verbose)\u001b[0m\n\u001b[1;32m    235\u001b[0m     bfile \u001b[38;5;241m=\u001b[39m BinaryFiltered(\n\u001b[1;32m    236\u001b[0m         ops[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfilename\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[1;32m    237\u001b[0m         n_chan_bin,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    252\u001b[0m         file_object\u001b[38;5;241m=\u001b[39mfile_object,\n\u001b[1;32m    253\u001b[0m     )\n\u001b[1;32m    255\u001b[0m \u001b[38;5;66;03m# TODO: don't think we need to do this actually\u001b[39;00m\n\u001b[1;32m    256\u001b[0m \u001b[38;5;66;03m# Save intermediate `ops` for use by GUI plots\u001b[39;00m\n\u001b[1;32m    257\u001b[0m \u001b[38;5;66;03m# io.save_ops(ops, results_dir)\u001b[39;00m\n\u001b[1;32m    258\u001b[0m \n\u001b[1;32m    259\u001b[0m \u001b[38;5;66;03m# Sort spikes and save results\u001b[39;00m\n\u001b[0;32m--> 260\u001b[0m st, tF, _, _ \u001b[38;5;241m=\u001b[39m \u001b[43mdetect_spikes\u001b[49m\u001b[43m(\u001b[49m\u001b[43mops\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbfile\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtic0\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtic0\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mprogress_bar\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprogress_bar\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    261\u001b[0m clu, Wall \u001b[38;5;241m=\u001b[39m cluster_spikes(st, tF, ops, device, bfile, tic0\u001b[38;5;241m=\u001b[39mtic0, progress_bar\u001b[38;5;241m=\u001b[39mprogress_bar)\n\u001b[1;32m    262\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m params[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mskip_kilosort_preprocessing\u001b[39m\u001b[38;5;124m\"\u001b[39m]:\n",
      "File \u001b[0;32m/mnt/cube/lo/envs/spikeproc/lib/python3.8/site-packages/kilosort/run_kilosort.py:392\u001b[0m, in \u001b[0;36mdetect_spikes\u001b[0;34m(ops, device, bfile, tic0, progress_bar)\u001b[0m\n\u001b[1;32m    390\u001b[0m tic \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime()\n\u001b[1;32m    391\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124mExtracting spikes using templates\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m--> 392\u001b[0m st0, tF, ops \u001b[38;5;241m=\u001b[39m \u001b[43mspikedetect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43mops\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbfile\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mprogress_bar\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprogress_bar\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    393\u001b[0m tF \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mfrom_numpy(tF)\n\u001b[1;32m    394\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlen\u001b[39m(st0)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m spikes extracted in \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtime\u001b[38;5;241m.\u001b[39mtime()\u001b[38;5;241m-\u001b[39mtic \u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m .2f\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124ms; \u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m+\u001b[39m \n\u001b[1;32m    395\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtotal \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtime\u001b[38;5;241m.\u001b[39mtime()\u001b[38;5;241m-\u001b[39mtic0 \u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m .2f\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124ms\u001b[39m\u001b[38;5;124m'\u001b[39m)\n",
      "File \u001b[0;32m/mnt/cube/lo/envs/spikeproc/lib/python3.8/site-packages/kilosort/spikedetect.py:233\u001b[0m, in \u001b[0;36mrun\u001b[0;34m(ops, bfile, device, progress_bar)\u001b[0m\n\u001b[1;32m    229\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m ibatch \u001b[38;5;129;01min\u001b[39;00m tqdm(np\u001b[38;5;241m.\u001b[39marange(bfile\u001b[38;5;241m.\u001b[39mn_batches), miniters\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m200\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m progress_bar \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m, \n\u001b[1;32m    230\u001b[0m                     mininterval\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m60\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m progress_bar \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[1;32m    231\u001b[0m     X \u001b[38;5;241m=\u001b[39m bfile\u001b[38;5;241m.\u001b[39mpadded_batch_to_torch(ibatch, ops)\n\u001b[0;32m--> 233\u001b[0m     xy, imax, amp, adist \u001b[38;5;241m=\u001b[39m \u001b[43mtemplate_match\u001b[49m\u001b[43m(\u001b[49m\u001b[43mX\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mops\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43miC\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43miC2\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mweigh\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    234\u001b[0m     yct \u001b[38;5;241m=\u001b[39m yweighted(yc, iC, adist, xy, device\u001b[38;5;241m=\u001b[39mdevice)\n\u001b[1;32m    235\u001b[0m     nsp \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(xy)\n",
      "File \u001b[0;32m/mnt/cube/lo/envs/spikeproc/lib/python3.8/site-packages/kilosort/spikedetect.py:156\u001b[0m, in \u001b[0;36mtemplate_match\u001b[0;34m(X, ops, iC, iC2, weigh, device)\u001b[0m\n\u001b[1;32m    154\u001b[0m Amaxs[:,\u001b[38;5;241m-\u001b[39mnt:] \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m    155\u001b[0m Amaxs  \u001b[38;5;241m=\u001b[39m max_pool1d(Amaxs\u001b[38;5;241m.\u001b[39munsqueeze(\u001b[38;5;241m0\u001b[39m), (\u001b[38;5;241m2\u001b[39m\u001b[38;5;241m*\u001b[39mnt0\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m1\u001b[39m), stride \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m, padding \u001b[38;5;241m=\u001b[39m nt0)\u001b[38;5;241m.\u001b[39msqueeze(\u001b[38;5;241m0\u001b[39m)\n\u001b[0;32m--> 156\u001b[0m xy \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlogical_and\u001b[49m\u001b[43m(\u001b[49m\u001b[43mAmaxs\u001b[49m\u001b[38;5;241;43m==\u001b[39;49m\u001b[43mAs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mAs\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m>\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mops\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mTh_universal\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnonzero\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    157\u001b[0m imax \u001b[38;5;241m=\u001b[39m imaxs[xy[:,\u001b[38;5;241m0\u001b[39m], xy[:,\u001b[38;5;241m1\u001b[39m]]\n\u001b[1;32m    158\u001b[0m amp \u001b[38;5;241m=\u001b[39m As[xy[:,\u001b[38;5;241m0\u001b[39m], xy[:,\u001b[38;5;241m1\u001b[39m]]\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "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",
    "            \n",
    "            for this_epoch in sess_epochs:\n",
    "                \n",
    "                # set output directory\n",
    "                epoch_struct = et.sgl_struct(sess_par,this_epoch,ephys_software=sess_par['ephys_software'])\n",
    "                sess_par['epoch'] = this_epoch\n",
    "                sort_folder = 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, this_epoch+'_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,this_epoch,'___________')\n",
    "                        # prepare recording for sorting\n",
    "                        print('prep..')\n",
    "                        if sess_par['ephys_software'] == 'sglx':\n",
    "                            # load recording\n",
    "                            rec_path = epoch_struct['folders']['sglx']\n",
    "                            this_rec = si.read_spikeglx(folder_path=rec_path,stream_name='imec0.ap')\n",
    "                            # save probe map prior to re-ordering for sorting\n",
    "                            probe_df = this_rec.get_probe().to_dataframe()\n",
    "                            probe_df.to_pickle(os.path.join(epoch_struct['folders']['derived'],'probe_map_df.pickle'))\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_p = 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_p = this_rec.set_probe(p,group_mode='by_probe') # assign new Probe object to probe\n",
    "                            else:\n",
    "                                this_rec_p = si.highpass_spatial_filter(recording=this_rec)\n",
    "                        elif sess_par['ephys_software'] =='oe':\n",
    "                            # load recording\n",
    "                            rec_path = [f.path for f in os.scandir(epoch_struct['folders']['oe']) if f.is_dir()][0]\n",
    "                            this_rec = si.read_openephys(folder_path=rec_path)\n",
    "                            # add probe\n",
    "                            this_probe = pm.make_probes(this_sess_config['probe']['probe_type'],this_sess_config['probe']['probe_model']) # neuronexus, Buzsaki64\n",
    "                            this_rec_p = this_rec.set_probe(this_probe,group_mode='by_shank')\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",
    "                        torch.cuda.empty_cache()\n",
    "                        this_sort = si.run_sorter(sorter_name=this_sess_config['sorter'],recording=this_rec_p,output_folder=sort_folder,\n",
    "                                             remove_existing_folder=True,delete_output_folder=False,delete_container_files=False,\n",
    "                                             verbose=verbose,raise_error=raise_error,**sort_params)\n",
    "                        torch.cuda.empty_cache()\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_folder,'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, this_epoch+'_spikesort_'+this_sess_config['sort']+'.log'), 'w') as f:\n",
    "                            f.write(sess_par['bird']+' '+sess_par['sess']+' '+this_epoch+' 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'],this_epoch,'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, this_epoch+'_spikesort_'+this_sess_config['sort']+'.log'), 'w') as f:\n",
    "                            f.write(sess_par['bird']+' '+sess_par['sess']+' '+this_epoch+' sort failed\\n')\n",
    "                            f.write(traceback.format_exc())\n",
    "                        sort_summary = [this_bird,this_sess,sess_par['ephys_software'],this_epoch,'FAIL']\n",
    "                else:\n",
    "                    sort_summary = [this_bird,this_sess,sess_par['ephys_software'],this_epoch,'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
}