From 670e88ad514ebaf7330827387bf5c53317638d3e Mon Sep 17 00:00:00 2001 From: chrishalcrow Date: Fri, 31 Jul 2026 11:27:03 +0100 Subject: [PATCH 1/3] 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/3] 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/3] 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