From 670e88ad514ebaf7330827387bf5c53317638d3e Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 31 Jul 2026 11:27:03 +0100 Subject: [PATCH 1/7] add compute main channel ids to kilosort run sorter --- .../extractors/phykilosortextractors.py | 4 ++-- src/spikeinterface/sorters/external/kilosort4.py | 15 ++++++++++++++- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/src/spikeinterface/extractors/phykilosortextractors.py b/src/spikeinterface/extractors/phykilosortextractors.py index 92ae2a0437..335b80e1d0 100644 --- a/src/spikeinterface/extractors/phykilosortextractors.py +++ b/src/spikeinterface/extractors/phykilosortextractors.py @@ -415,7 +415,7 @@ def read_kilosort_as_analyzer(folder_path, unwhiten=True, gain_to_uV=None, offse ) sparsity = _make_sparsity_from_templates(sorting, recording, phy_path) - main_channel_indices = _make_main_channel_indices_from_templates(sorting, recording, phy_path) + main_channel_indices = _make_main_channel_indices_from_templates(phy_path) sorting_analyzer = create_sorting_analyzer( sorting, recording, sparse=True, sparsity=sparsity, main_channel_indices=main_channel_indices @@ -490,7 +490,7 @@ def _make_sparsity_from_templates(sorting, recording, kilosort_output_path): return ChannelSparsity(mask, unit_ids=unit_ids, channel_ids=channel_ids) -def _make_main_channel_indices_from_templates(sorting, recording, kilosort_output_path): +def _make_main_channel_indices_from_templates(kilosort_output_path): """Constructs the `main_channel_indices` from kilosort output, by finding the channel containing the largest peak-to-peak value.""" diff --git a/src/spikeinterface/sorters/external/kilosort4.py b/src/spikeinterface/sorters/external/kilosort4.py index bd856dc13b..7796f556f7 100644 --- a/src/spikeinterface/sorters/external/kilosort4.py +++ b/src/spikeinterface/sorters/external/kilosort4.py @@ -1,4 +1,5 @@ import warnings +import csv from pathlib import Path from packaging import version @@ -7,7 +8,6 @@ from spikeinterface.core import write_binary_recording, Motion, BaseRecording from spikeinterface.sorters.basesorter import BaseSorter, get_job_kwargs from .kilosortbase import KilosortBase -from spikeinterface.sorters.basesorter import get_job_kwargs from importlib.metadata import version as importlib_version @@ -457,6 +457,19 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose): save_preprocessed_copy=save_preprocessed_copy, ) + if (results_dir / "templates.npy").is_file(): + # Note: these are the whitened templates + templates = np.load(results_dir / "templates.npy") + # main channel indices are the argmax of the ptp of the templates + main_channel_indices = np.argmax(np.ptp(templates, axis=1), axis=1) + main_channel_ids = recording.channel_ids[main_channel_indices] + # save main_channel_ids + with open(results_dir / "cluster_main_channel_id.tsv", "w", newline="", encoding="utf-8") as f: + writer = csv.writer(f, delimiter="\t") + writer.writerow(["cluster_id", "main_channel_id"]) + for unit_index, item in enumerate(main_channel_ids): + writer.writerow([unit_index, item]) + if params["delete_recording_dat"]: # only delete dat file if it was created by the wrapper if (sorter_output_folder / "recording.dat").is_file(): From e80696f634068335c5d4c27360e56ec222c5a886 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Wed, 16 Sep 2026 08:56:54 +0100 Subject: [PATCH 2/7] move main_channel_id to kilosort load --- .../extractors/phykilosortextractors.py | 30 +++++++++++++++---- .../sorters/external/kilosort4.py | 13 -------- .../sorters/external/kilosortbase.py | 18 +++++++++-- 3 files changed, 41 insertions(+), 20 deletions(-) diff --git a/src/spikeinterface/extractors/phykilosortextractors.py b/src/spikeinterface/extractors/phykilosortextractors.py index 335b80e1d0..1d7a2aaab0 100644 --- a/src/spikeinterface/extractors/phykilosortextractors.py +++ b/src/spikeinterface/extractors/phykilosortextractors.py @@ -66,6 +66,7 @@ def __init__( keep_good_only: bool = False, remove_empty_units: bool = False, load_all_cluster_properties: bool = True, + channel_ids: list | np.ndarray | None = None, ): try: import pandas as pd @@ -228,6 +229,12 @@ def __init__( values_ = cluster_info[prop_name].values self.set_property(key=prop_name, values=values_) + if channel_ids is not None: + main_channel_indices = _make_main_channel_indices_from_templates(phy_folder) + if main_channel_indices is not None: + main_channel_ids = channel_ids[main_channel_indices] + self.set_property(key="main_channel_id", values=main_channel_ids) + self.annotate(phy_folder=str(phy_folder.resolve())) self.add_sorting_segment(PhySortingSegment(spike_times_clean, spike_clusters_clean)) @@ -333,13 +340,20 @@ class KiloSortSortingExtractor(BasePhyKilosortSortingExtractor): The loaded Sorting object. """ - def __init__(self, folder_path: Path | str, keep_good_only: bool = False, remove_empty_units: bool = True): + def __init__( + self, + folder_path: Path | str, + keep_good_only: bool = False, + remove_empty_units: bool = True, + channel_ids: list | np.ndarray | None = None, + ): BasePhyKilosortSortingExtractor.__init__( self, folder_path, exclude_cluster_groups=None, keep_good_only=keep_good_only, remove_empty_units=remove_empty_units, + channel_ids=channel_ids, ) self._kwargs = {"folder_path": str(Path(folder_path).absolute()), "keep_good_only": keep_good_only} @@ -494,10 +508,16 @@ def _make_main_channel_indices_from_templates(kilosort_output_path): """Constructs the `main_channel_indices` from kilosort output, by finding the channel containing the largest peak-to-peak value.""" - templates = np.load(kilosort_output_path / "templates.npy") - # main channel indices are the argmax of the ptp of the templates, which is the channel with - # the largest peak-to-peak amplitude - main_channel_indices = np.argmax(np.ptp(templates, axis=1), axis=1) + templates_filepath = kilosort_output_path / "templates.npy" + + if templates_filepath.is_file(): + templates = np.load(kilosort_output_path / "templates.npy") + # main channel indices are the argmax of the ptp of the templates, which is the channel with + # the largest peak-to-peak amplitude + main_channel_indices = np.argmax(np.ptp(templates, axis=1), axis=1) + else: + main_channel_indices = None + return main_channel_indices diff --git a/src/spikeinterface/sorters/external/kilosort4.py b/src/spikeinterface/sorters/external/kilosort4.py index 7796f556f7..313c553657 100644 --- a/src/spikeinterface/sorters/external/kilosort4.py +++ b/src/spikeinterface/sorters/external/kilosort4.py @@ -457,19 +457,6 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose): save_preprocessed_copy=save_preprocessed_copy, ) - if (results_dir / "templates.npy").is_file(): - # Note: these are the whitened templates - templates = np.load(results_dir / "templates.npy") - # main channel indices are the argmax of the ptp of the templates - main_channel_indices = np.argmax(np.ptp(templates, axis=1), axis=1) - main_channel_ids = recording.channel_ids[main_channel_indices] - # save main_channel_ids - with open(results_dir / "cluster_main_channel_id.tsv", "w", newline="", encoding="utf-8") as f: - writer = csv.writer(f, delimiter="\t") - writer.writerow(["cluster_id", "main_channel_id"]) - for unit_index, item in enumerate(main_channel_ids): - writer.writerow([unit_index, item]) - if params["delete_recording_dat"]: # only delete dat file if it was created by the wrapper if (sorter_output_folder / "recording.dat").is_file(): diff --git a/src/spikeinterface/sorters/external/kilosortbase.py b/src/spikeinterface/sorters/external/kilosortbase.py index 4cdb51e21e..99ebbab758 100644 --- a/src/spikeinterface/sorters/external/kilosortbase.py +++ b/src/spikeinterface/sorters/external/kilosortbase.py @@ -9,7 +9,7 @@ from spikeinterface.sorters.utils import ShellScript, get_matlab_shell_name, get_bash_path from spikeinterface.sorters.basesorter import get_job_kwargs from spikeinterface.extractors.extractor_classes import KiloSortSortingExtractor -from spikeinterface.core import write_binary_recording +from spikeinterface.core import write_binary_recording, load from spikeinterface.preprocessing.zero_channel_pad import TracePaddedRecording @@ -254,6 +254,20 @@ def _get_result_from_folder(cls, sorter_output_folder): params_file = sorter_output_folder / "spikeinterface_params.json" with params_file.open("r") as f: sorter_params = json.load(f)["sorter_params"] + + recording_file = sorter_output_folder.parent / "spikeinterface_recording.json" + if recording_file.is_file(): + try: + # TODO: load the channel ids without loading the recording + recording = load(recording_file) + channel_ids = recording.channel_ids + except: + channel_ids = None + else: + channel_ids = None + keep_good_only = sorter_params.get("keep_good_only", False) - sorting = KiloSortSortingExtractor(folder_path=sorter_output_folder, keep_good_only=keep_good_only) + sorting = KiloSortSortingExtractor( + folder_path=sorter_output_folder, keep_good_only=keep_good_only, channel_ids=channel_ids + ) return sorting From 256d5481fa33ffbdf55a77e9b1c4193f80fce123 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Wed, 16 Sep 2026 08:58:33 +0100 Subject: [PATCH 3/7] remove csv import from kilosort4 --- src/spikeinterface/sorters/external/kilosort4.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/spikeinterface/sorters/external/kilosort4.py b/src/spikeinterface/sorters/external/kilosort4.py index 313c553657..ea5e596b79 100644 --- a/src/spikeinterface/sorters/external/kilosort4.py +++ b/src/spikeinterface/sorters/external/kilosort4.py @@ -1,5 +1,4 @@ import warnings -import csv from pathlib import Path from packaging import version From cd2bc8adec7f8766668cde98f08a405371bd00f9 Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Thu, 17 Sep 2026 11:40:54 +0100 Subject: [PATCH 4/7] use load_recording_from_folder and take account of bad channels --- .../extractors/phykilosortextractors.py | 19 +++++++++++++------ .../sorters/external/kilosortbase.py | 14 +++----------- 2 files changed, 16 insertions(+), 17 deletions(-) diff --git a/src/spikeinterface/extractors/phykilosortextractors.py b/src/spikeinterface/extractors/phykilosortextractors.py index 8c57d3a9c5..95cc7c9c64 100644 --- a/src/spikeinterface/extractors/phykilosortextractors.py +++ b/src/spikeinterface/extractors/phykilosortextractors.py @@ -36,6 +36,8 @@ class BasePhyKilosortSortingExtractor(BaseSorting): If True, empty units are removed from the sorting extractor. load_all_cluster_properties : bool, default: True If True, all cluster properties are loaded from the tsv/csv files. + channel_ids : list | np.ndarray | None, default None + The channel_ids of the recording passed to `run_sorter` Notes ----- @@ -134,6 +136,17 @@ def __init__( cluster_info = pd.DataFrame({"cluster_id": unique_unit_ids}) cluster_info["group"] = ["unsorted"] * len(unique_unit_ids) + # we need to add main_channel_ids before selecting good units etc. + if channel_ids is not None: + # if the user has a non-trivial channel_map (from passing bad channel ids or otherwise) + # we need to remap the channel_ids + channel_map = np.load(phy_folder / "channel_map.npy").reshape(-1).astype("int64", copy=False) + channel_ids = np.asarray(channel_ids)[channel_map] + main_channel_indices = _make_main_channel_indices_from_templates(phy_folder) + if main_channel_indices is not None: + main_channel_ids = channel_ids[main_channel_indices] + cluster_info["main_channel_id"] = main_channel_ids + if exclude_cluster_groups is not None: if isinstance(exclude_cluster_groups, str): cluster_info = cluster_info.query(f"group != '{exclude_cluster_groups}'") @@ -230,12 +243,6 @@ def __init__( values_ = cluster_info[prop_name].values self.set_property(key=prop_name, values=values_) - if channel_ids is not None: - main_channel_indices = _make_main_channel_indices_from_templates(phy_folder) - if main_channel_indices is not None: - main_channel_ids = channel_ids[main_channel_indices] - self.set_property(key="main_channel_id", values=main_channel_ids) - self.annotate(phy_folder=str(phy_folder.resolve())) self.add_sorting_segment(PhySortingSegment(spike_times_clean, spike_clusters_clean)) diff --git a/src/spikeinterface/sorters/external/kilosortbase.py b/src/spikeinterface/sorters/external/kilosortbase.py index 99ebbab758..48a6de8de2 100644 --- a/src/spikeinterface/sorters/external/kilosortbase.py +++ b/src/spikeinterface/sorters/external/kilosortbase.py @@ -7,7 +7,7 @@ import numpy as np from spikeinterface.sorters.utils import ShellScript, get_matlab_shell_name, get_bash_path -from spikeinterface.sorters.basesorter import get_job_kwargs +from spikeinterface.sorters.basesorter import get_job_kwargs, BaseSorter from spikeinterface.extractors.extractor_classes import KiloSortSortingExtractor from spikeinterface.core import write_binary_recording, load from spikeinterface.preprocessing.zero_channel_pad import TracePaddedRecording @@ -255,16 +255,8 @@ def _get_result_from_folder(cls, sorter_output_folder): with params_file.open("r") as f: sorter_params = json.load(f)["sorter_params"] - recording_file = sorter_output_folder.parent / "spikeinterface_recording.json" - if recording_file.is_file(): - try: - # TODO: load the channel ids without loading the recording - recording = load(recording_file) - channel_ids = recording.channel_ids - except: - channel_ids = None - else: - channel_ids = None + recording = BaseSorter.load_recording_from_folder(sorter_output_folder.parent, with_warnings=False) + channel_ids = recording.channel_ids keep_good_only = sorter_params.get("keep_good_only", False) sorting = KiloSortSortingExtractor( From e3a662458b96a09bfd782febe5859d1e55debd52 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 17 Sep 2026 15:54:40 +0200 Subject: [PATCH 5/7] fix: remove unused imports and us cls --- src/spikeinterface/sorters/external/kilosortbase.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/spikeinterface/sorters/external/kilosortbase.py b/src/spikeinterface/sorters/external/kilosortbase.py index 48a6de8de2..6d5af8e7e2 100644 --- a/src/spikeinterface/sorters/external/kilosortbase.py +++ b/src/spikeinterface/sorters/external/kilosortbase.py @@ -7,9 +7,9 @@ import numpy as np from spikeinterface.sorters.utils import ShellScript, get_matlab_shell_name, get_bash_path -from spikeinterface.sorters.basesorter import get_job_kwargs, BaseSorter +from spikeinterface.sorters.basesorter import get_job_kwargs from spikeinterface.extractors.extractor_classes import KiloSortSortingExtractor -from spikeinterface.core import write_binary_recording, load +from spikeinterface.core import write_binary_recording from spikeinterface.preprocessing.zero_channel_pad import TracePaddedRecording @@ -255,7 +255,7 @@ def _get_result_from_folder(cls, sorter_output_folder): with params_file.open("r") as f: sorter_params = json.load(f)["sorter_params"] - recording = BaseSorter.load_recording_from_folder(sorter_output_folder.parent, with_warnings=False) + recording = cls.load_recording_from_folder(sorter_output_folder.parent, with_warnings=False) channel_ids = recording.channel_ids keep_good_only = sorter_params.get("keep_good_only", False) From 137cb90d0fa9c7b0781d53df485e34a334fe0520 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 17 Sep 2026 15:56:09 +0200 Subject: [PATCH 6/7] Apply suggestion from @alejoe91 --- src/spikeinterface/sorters/external/kilosortbase.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/spikeinterface/sorters/external/kilosortbase.py b/src/spikeinterface/sorters/external/kilosortbase.py index 6d5af8e7e2..06957c3a7c 100644 --- a/src/spikeinterface/sorters/external/kilosortbase.py +++ b/src/spikeinterface/sorters/external/kilosortbase.py @@ -256,7 +256,7 @@ def _get_result_from_folder(cls, sorter_output_folder): sorter_params = json.load(f)["sorter_params"] recording = cls.load_recording_from_folder(sorter_output_folder.parent, with_warnings=False) - channel_ids = recording.channel_ids + channel_ids = recording.channel_ids if recording is not None else None keep_good_only = sorter_params.get("keep_good_only", False) sorting = KiloSortSortingExtractor( From 4e0e3e9c05c65014c17f92cb4e747fb53d539508 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 17 Sep 2026 16:30:40 +0200 Subject: [PATCH 7/7] Revert cls change Co-authored-by: Alessio Buccino --- src/spikeinterface/sorters/external/kilosortbase.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/sorters/external/kilosortbase.py b/src/spikeinterface/sorters/external/kilosortbase.py index 06957c3a7c..5ca1f544a0 100644 --- a/src/spikeinterface/sorters/external/kilosortbase.py +++ b/src/spikeinterface/sorters/external/kilosortbase.py @@ -7,7 +7,7 @@ import numpy as np from spikeinterface.sorters.utils import ShellScript, get_matlab_shell_name, get_bash_path -from spikeinterface.sorters.basesorter import get_job_kwargs +from spikeinterface.sorters.basesorter import get_job_kwargs, BaseSorter from spikeinterface.extractors.extractor_classes import KiloSortSortingExtractor from spikeinterface.core import write_binary_recording from spikeinterface.preprocessing.zero_channel_pad import TracePaddedRecording @@ -255,7 +255,7 @@ def _get_result_from_folder(cls, sorter_output_folder): with params_file.open("r") as f: sorter_params = json.load(f)["sorter_params"] - recording = cls.load_recording_from_folder(sorter_output_folder.parent, with_warnings=False) + recording = BaseSorter.load_recording_from_folder(sorter_output_folder.parent, with_warnings=False) channel_ids = recording.channel_ids if recording is not None else None keep_good_only = sorter_params.get("keep_good_only", False)