Bombcell integration

examples/get_started/example_bombcell_unit_labelling.ipynb

@@ -0,0 +1,262 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Bombcell unit labelling\n",
+    "\n",
+    "With this notebook you can:\n",
+    "- load a SortingAnalyzer\n",
+    "- compute required extensions\n",
+    "- label units based on quality thresholds\n",
+    "- generating and save summary plots\n",
+    "- save metrics and results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "import spikeinterface as si\n",
+    "from spikeinterface.curation import (\n",
+    "    bombcell_get_default_thresholds,\n",
+    "    bombcell_label_units,\n",
+    "    save_thresholds,\n",
+    "    load_thresholds,\n",
+    ")\n",
+    "from spikeinterface.widgets import plot_unit_labelling_all"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### load a SortingAnalyzer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Change this to your analyzer path - you need to have already generated a sorting analyzer. see quickstart.py for how to do this\n",
+    "analyzer_path = \"/Users/jf5479/Downloads/M25_D18/kilosort4_sa\"\n",
+    "output_folder = Path(analyzer_path) / \"bombcell\"\n",
+    "\n",
+    "analyzer = si.load_sorting_analyzer(analyzer_path)\n",
+    "analyzer"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### compute required extensions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Templates (required for template_metrics)\n",
+    "if not analyzer.has_extension(\"templates\"):\n",
+    "    analyzer.compute(\"templates\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Template metrics\n",
+    "if not analyzer.has_extension(\"template_metrics\"):\n",
+    "    analyzer.compute(\"template_metrics\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Quality metrics (and dependencies)\n",
+    "if not analyzer.has_extension(\"spike_amplitudes\"):\n",
+    "    analyzer.compute(\"spike_amplitudes\")\n",
+    "\n",
+    "if not analyzer.has_extension(\"noise_levels\"):\n",
+    "    analyzer.compute(\"noise_levels\")\n",
+    "\n",
+    "if not analyzer.has_extension(\"quality_metrics\"):\n",
+    "    analyzer.compute(\"quality_metrics\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### get metrics"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "qm = analyzer.get_extension(\"quality_metrics\").get_data()\n",
+    "tm = analyzer.get_extension(\"template_metrics\").get_data()\n",
+    "\n",
+    "print(f\"Quality metrics: {list(qm.columns)}\")\n",
+    "print(f\"Template metrics: {list(tm.columns)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### set labelling thresholds"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Use default thresholds\n",
+    "thresholds = bombcell_get_default_thresholds()\n",
+    "\n",
+    "# Or load from file:\n",
+    "# thresholds = load_thresholds(\"my_thresholds.json\")\n",
+    "\n",
+    "thresholds"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Optionally modify thresholds\n",
+    "# thresholds[\"amplitude_median\"][\"min\"] = 50  # stricter\n",
+    "# thresholds[\"rp_contamination\"][\"max\"] = 0.05  # stricter"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Optionally set and load thresholds from a JSON file \n",
+    "# Load thresholds from saved JSON\n",
+    "thresholds = load_thresholds(output_folder / \"thresholds.json\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The JSON file format looks like:\n",
+    "```json\n",
+    "{\n",
+    "    \"amplitude_median\": {\"min\": 40, \"max\": null},\n",
+    "    \"num_positive_peaks\": {\"min\": null, \"max\": 2},\n",
+    "    \"peak_to_trough_duration\": {\"min\": 0.0001, \"max\": 0.00115}\n",
+    "}\n",
+    "```\n",
+    "`null` in JSON becomes `np.nan` (threshold disabled)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### label units"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "unit_type, unit_type_string = bombcell_label_units(\n",
+    "    quality_metrics=qm,\n",
+    "    template_metrics=tm,\n",
+    "    thresholds=thresholds,\n",
+    "    label_non_somatic=True,\n",
+    "    split_non_somatic_good_mua=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### generate summary plots"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plots = plot_unit_labelling_all(\n",
+    "    analyzer,\n",
+    "    unit_type,\n",
+    "    unit_type_string,\n",
+    "    quality_metrics=qm,\n",
+    "    template_metrics=tm,\n",
+    "    thresholds=thresholds,\n",
+    "    save_folder=output_folder,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### save labelling thresholds"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "save_thresholds(thresholds, output_folder / \"thresholds.json\")\n",
+    "\n",
+    "print(f\"Results saved to: {output_folder.absolute()}\")\n",
+    "print(\"\\nFiles:\")\n",
+    "for f in sorted(output_folder.glob(\"*\")):\n",
+    "    print(f\"  - {f.name}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

src/spikeinterface/curation/__init__.py

@@ -21,5 +21,18 @@
 from .sortingview_curation import apply_sortingview_curation
 
 # automated curation
+from .unit_labelling import (
+    WAVEFORM_METRICS,
+    SPIKE_QUALITY_METRICS,
+    NON_SOMATIC_METRICS,
+    bombcell_get_default_thresholds,
+    bombcell_label_units,
+    apply_thresholds,
+    get_labelling_summary,
+    save_thresholds,
+    load_thresholds,
+    save_labelling_results,
+)
+
 from .model_based_curation import auto_label_units, load_model
 from .train_manual_curation import train_model, get_default_classifier_search_spaces

src/spikeinterface/curation/default_thresholds.json

@@ -0,0 +1,74 @@
+{
+    "num_positive_peaks": {
+        "min": null,
+        "max": 2
+    },
+    "num_negative_peaks": {
+        "min": null,
+        "max": 1
+    },
+    "peak_to_trough_duration": {
+        "min": 0.0001,
+        "max": 0.00115
+    },
+    "waveform_baseline_flatness": {
+        "min": null,
+        "max": 0.5
+    },
+    "peak_after_to_trough_ratio": {
+        "min": null,
+        "max": 0.8
+    },
+    "exp_decay": {
+        "min": 0.01,
+        "max": 0.1
+    },
+    "amplitude_median": {
+        "min": 40,
+        "max": null
+    },
+    "snr_bombcell": {
+        "min": 5,
+        "max": null
+    },
+    "amplitude_cutoff": {
+        "min": null,
+        "max": 0.2
+    },
+    "num_spikes": {
+        "min": 300,
+        "max": null
+    },
+    "rp_contamination": {
+        "min": null,
+        "max": 0.1
+    },
+    "presence_ratio": {
+        "min": 0.7,
+        "max": null
+    },
+    "drift_ptp": {
+        "min": null,
+        "max": 100
+    },
+    "peak_before_to_trough_ratio": {
+        "min": null,
+        "max": 3
+    },
+    "peak_before_width": {
+        "min": 150,
+        "max": null
+    },
+    "trough_width": {
+        "min": 200,
+        "max": null
+    },
+    "peak_before_to_peak_after_ratio": {
+        "min": null,
+        "max": 3
+    },
+    "main_peak_to_trough_ratio": {
+        "min": null,
+        "max": 0.8
+    }
+}