{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using custom 'good/mua' labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Kilosort4 determines if a unit is 'good' or 'mua' by computing a few metrics based on its correlogram (see methods in [the Kilosort paper](https://www.nature.com/articles/s41592-024-02232-7) for more details). This tutorial demonstrates how to change these labels using your own criteria."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After sorting our data, we need to load the results and determine where the labels need to be saved. You could also use the same variables as they're returned by `run_kilosort` if you're doing this in a single script using the API, but note that `st` will have 3 columns in that case instead of just 1 column as used in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import shutil\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from kilosort.run_kilosort import load_sorting\n",
    "\n",
    "# Path to load existing sorting results from.\n",
    "results_dir = Path('c:/users/jacob/.kilosort/.test_data/kilosort4/')\n",
    "# Paths where new labels will be saved for use with Phy.\n",
    "save_1 = results_dir / 'cluster_KSLabel.tsv'\n",
    "save_2 = results_dir / 'cluster_group.tsv'\n",
    "# Save a backup of KS4's original labels before overwriting (recommended).\n",
    "shutil.copyfile(save_1, results_dir / 'cluster_KSLabel_backup.tsv')\n",
    "\n",
    "# Load sorting results\n",
    "ops, st, clu, similar_templates, is_ref, est_contam_rate, kept_spikes = \\\n",
    "    load_sorting(results_dir)\n",
    "\n",
    "cluster_labels = np.unique(clu)  # integer label for each cluster\n",
    "fs = ops['fs']                   # sampling rate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will also load Kilosort4's default labels. For this tutorial we'll use those as a starting point, requiring that units satisfy the old criteria *and* the new criteria. You could also just ignore the old labels and only use your own."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Option 1: Use existing labels as a starting point.\n",
    "#           KS4 assigns \"good\" where is_ref is True, and \"mua\" otherwise.\n",
    "label_good = is_ref.copy()\n",
    "\n",
    "# Option 2: Ignore KS4's labels and only use your own criteria.\n",
    "# label_good = np.ones(cluster_labels.size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some examples of other criteria you might want to use would be only labeling units 'good' if they have a firing rate above 1Hz and a contamination rate below 0.2. Whatever criteria you want to use, the process is the same: create a boolean array with shape (n_clusters,) that is True if the criteria is met, and False otherwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "contam_good = est_contam_rate < 0.2   # this already has shape (n_clusters,)\n",
    "fr_good = np.zeros(cluster_labels.size, dtype=bool)\n",
    "for i, c in enumerate(cluster_labels):\n",
    "    # Get all spikes assigned to this cluster\n",
    "    spikes = st[clu == c]\n",
    "    # Compute est. firing rate using your preferred method.\n",
    "    # Note that this formula will not work well for units that drop in and out.\n",
    "    fr = spikes.size / (spikes.max()/fs - spikes.min()/fs)\n",
    "    if fr >= 1:\n",
    "        fr_good[i] = True\n",
    "\n",
    "# Update labels, requiring that all criteria hold for each cluster.\n",
    "label_good = np.logical_and(label_good, contam_good, fr_good)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another example would be to only assign \"good\" to units with a presence ratio above some fraction, say 0.5. This will restrict the \"good\" label to units that are detected for at least half of the recording.\n",
    "\n",
    "This involves binning the data into large chunks to determine which periods of time each unit is active during. We want the bins to be large enough to not penalize units with low firing rates, but still small enough to capture periods when a unit is not detected. We recommend setting the number of bins such that each bin is around 5 minutes as a starting point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Formula adapted from https://github.com/AllenInstitute/ecephys_spike_sorting/\n",
    "\n",
    "def presence_ratio(spike_train, num_bins, min_time, max_time, min_spike_pct=0.05):\n",
    "    h, b = np.histogram(spike_train, np.linspace(min_time, max_time, num_bins))\n",
    "    min_spikes = h.mean()*min_spike_pct\n",
    "\n",
    "    # NOTE: Allen Institute formula leaves off the -1 to force the ratio to\n",
    "    #       never reach 1.0. We've included it here because without it the ratio\n",
    "    #       is biased too much for a small number of bins.\n",
    "    return np.sum(h > min_spikes) / (num_bins - 1)\n",
    "\n",
    "# Compute presence ratio for each cluster\n",
    "presence = np.zeros(cluster_labels.size)\n",
    "min_time = st.min()\n",
    "max_time = st.max()\n",
    "for i, c in enumerate(cluster_labels):\n",
    "    spikes = st[clu == c]\n",
    "    presence[i] = presence_ratio(spikes, 10, min_time, max_time)\n",
    "\n",
    "presence_good = presence >= 0.5\n",
    "# Update labels with the additional criteria.\n",
    "label_good = np.logical_and(label_good, presence_good)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After we're finished changing labels, we need to save them again in the format expected by Phy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert True/False to 'good'/'mua'\n",
    "ks_labels = ['good' if b else 'mua' for b in label_good]\n",
    "\n",
    "# Write to two .tsv files.\n",
    "with open(save_1, 'w') as f:\n",
    "    f.write(f'cluster_id\\tKSLabel\\n')\n",
    "    for i, p in enumerate(ks_labels):\n",
    "        f.write(f'{i}\\t{p}\\n')\n",
    "shutil.copyfile(save_1, save_2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.18 ('kilosort4')",
   "language": "python",
   "name": "python3"
  },
  "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.9.18"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c450ffd893003221542b409439a3c6dd3d0fae98f595b9e4724fd711cbdc16b9"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}