diff --git a/.github/scripts/check_kilosort4_releases.py b/.github/scripts/check_kilosort4_releases.py index 8bd0163e3a..a143d29b40 100644 --- a/.github/scripts/check_kilosort4_releases.py +++ b/.github/scripts/check_kilosort4_releases.py @@ -15,10 +15,8 @@ def get_pypi_versions(package_name): response.raise_for_status() data = response.json() versions = list(sorted(data["releases"].keys())) - # Filter out versions that are less than 4.0.16 and different from 4.0.26 and 4.0.27 - # (buggy - https://github.com/MouseLand/Kilosort/releases/tag/v4.0.26) - versions = [ver for ver in versions if parse(ver) >= parse("4.0.16") and - parse(ver) not in [parse("4.0.26"), parse("4.0.27")]] + # Filter out versions that are less than 4.1.1, since they don't support numpy 2.0 + versions = [ver for ver in versions if parse(ver) >= parse("4.1.1")] return versions diff --git a/.github/scripts/test_kilosort4_ci.py b/.github/scripts/test_kilosort4_ci.py index b075704628..75bac8f03b 100644 --- a/.github/scripts/test_kilosort4_ci.py +++ b/.github/scripts/test_kilosort4_ci.py @@ -28,7 +28,8 @@ import spikeinterface.full as si from spikeinterface.core.testing import check_sortings_equal -from spikeinterface.sorters.external.kilosort4 import Kilosort4Sorter +from spikeinterface.sorters.external.kilosort4 import Kilosort4Sorter, read_kilosort4_motion +from spikeinterface.core.motion import Motion from probeinterface.io import write_prb from spikeinterface.extractors import read_kilosort_as_analyzer @@ -669,6 +670,43 @@ def monkeypatch_filter_function(self, X, ops=None, ibatch=None): assert np.allclose(results["ks"]["st"], results["si"]["st"], rtol=0, atol=1) assert np.array_equal(results["ks"]["clus"], results["si"]["clus"]) + def test_read_kilosort4_motion(self, recording_and_paths, tmp_path): + """ + Test that read_kilosort4_motion returns a Motion object whose displacement + equals dshift (not dshift + yblk), and that temporal/spatial bins are correct. + """ + recording, _ = recording_and_paths + sorter_output_dir = tmp_path / "ks4_motion_output" / "sorter_output" + + si.run_sorter( + "kilosort4", + recording, + folder=tmp_path / "ks4_motion_output", + remove_existing_folder=True, + ) + + ops = np.load(sorter_output_dir / "ops.npy", allow_pickle=True).item() + yblk = ops["yblk"] + dshift = ops["dshift"] + + # without recording: temporal bins estimated from batch count + motion = read_kilosort4_motion(sorter_output_dir) + assert isinstance(motion, Motion) + assert motion.displacement[0].shape == dshift.shape + np.testing.assert_array_equal(motion.displacement[0], dshift) + np.testing.assert_array_equal(motion.spatial_bins_um, yblk) + assert motion.temporal_bins_s[0].shape[0] == dshift.shape[0] + # displacement must be relative (not offset by spatial bin position) + assert not np.allclose(motion.displacement[0], dshift + yblk) + + # with recording: temporal bins bounded by recording times + motion_rec = read_kilosort4_motion(sorter_output_dir, recording=recording) + assert isinstance(motion_rec, Motion) + np.testing.assert_array_equal(motion_rec.displacement[0], dshift) + assert motion_rec.temporal_bins_s[0].shape[0] == dshift.shape[0] + assert motion_rec.temporal_bins_s[0][0] >= recording.get_start_time() + assert motion_rec.temporal_bins_s[0][-1] <= recording.get_end_time() + ##### Helpers ###### def _get_kilosort_native_settings(self, recording, paths, param_key, param_value): """ diff --git a/pyproject.toml b/pyproject.toml index a578310c57..2fd1bf8dae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,8 +20,7 @@ classifiers = [ dependencies = [ - "numpy>=1.20;python_version<'3.13'", - "numpy>=2.0.0;python_version>='3.13'", + "numpy>=2.0.0", "threadpoolctl>=3.0.0", "tqdm", "zarr>=2.18,<3", diff --git a/src/spikeinterface/core/motion.py b/src/spikeinterface/core/motion.py index de57b4f9a4..6a99ca29c8 100644 --- a/src/spikeinterface/core/motion.py +++ b/src/spikeinterface/core/motion.py @@ -14,7 +14,8 @@ class Motion: Parameters ---------- displacement : numpy array 2d or list of - Motion estimate in um. + Motion estimate in um, relative to the spatial_bins_um. + The first dimension is temporal bins, the second dimension is spatial bins. List is the number of segment. For each semgent : @@ -93,6 +94,7 @@ def get_displacement_at_time_and_depth(self, times_s, locations_um, segment_inde Parameters ---------- times_s: np.array + Times at which to evaluate the motion, in seconds. This should be a one-dimensional array. locations_um: np.array Either this is a one-dimensional array (a vector of positions along self.dimension), or else a 2d array with the 2 or 3 spatial dimensions indexed along axis=1. diff --git a/src/spikeinterface/extractors/phykilosortextractors.py b/src/spikeinterface/extractors/phykilosortextractors.py index 408cdb32b2..92ae2a0437 100644 --- a/src/spikeinterface/extractors/phykilosortextractors.py +++ b/src/spikeinterface/extractors/phykilosortextractors.py @@ -390,7 +390,7 @@ def read_kilosort_as_analyzer(folder_path, unwhiten=True, gain_to_uV=None, offse # kilosort occasionally contains a few spikes just beyond the recording end point, which can lead # to errors later. To avoid this, we pad the recording with an extra second of blank time. - duration = sorting.segments[0]._all_spikes[-1] / sampling_frequency + 1 + duration = sorting.segments[0]._all_spike_times[-1] / sampling_frequency + 1 if (phy_path / "probe.prb").is_file(): probegroup = read_prb(phy_path / "probe.prb") @@ -415,8 +415,11 @@ 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) - sorting_analyzer = create_sorting_analyzer(sorting, recording, sparse=True, sparsity=sparsity) + sorting_analyzer = create_sorting_analyzer( + sorting, recording, sparse=True, sparsity=sparsity, main_channel_indices=main_channel_indices + ) # first compute random spikes. These do nothing, but are needed for si-gui to run sorting_analyzer.compute("random_spikes") @@ -487,6 +490,17 @@ 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): + """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) + return main_channel_indices + + def _make_templates( sorting_analyzer, kilosort_output_path, mask, sampling_frequency, gain_to_uV, offset_to_uV, unwhiten=True ): diff --git a/src/spikeinterface/sorters/external/kilosort4.py b/src/spikeinterface/sorters/external/kilosort4.py index e54fb34cda..bd856dc13b 100644 --- a/src/spikeinterface/sorters/external/kilosort4.py +++ b/src/spikeinterface/sorters/external/kilosort4.py @@ -1,7 +1,10 @@ import warnings +from pathlib import Path from packaging import version -from spikeinterface.core import write_binary_recording +import numpy as np + +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 @@ -77,7 +80,7 @@ def is_installed(cls): @classmethod def get_sorter_version(cls): - """kilosort.__version__ <4.0.10 is always '4'""" + """kilosort.__version__ < 4.0.10 is always '4'""" return importlib_version("kilosort") @classmethod @@ -122,12 +125,11 @@ def initialize_folder(cls, recording, output_folder, verbose, remove_existing_fo @classmethod def check_sorter_version(cls): kilosort_version = version.parse(cls.get_sorter_version()) - if kilosort_version < version.parse("4.0.16"): - raise Exception( - f"""SpikeInterface only supports kilosort versions 4.0.16 and above. You are running version {kilosort_version}. To install the latest version, run: - >>> pip install kilosort --upgrade - """ - ) + if kilosort_version < version.parse("4.1.1"): + raise Exception(f"""SpikeInterface only supports kilosort versions 4.1.1 and above (which support numpy>=2). + You are running version {kilosort_version}. To install the latest version, run: + >>> pip install kilosort --upgrade + """) @classmethod def _setup_recording(cls, recording, sorter_output_folder, params, verbose): @@ -177,7 +179,6 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose): import time import torch - import numpy as np import logging if version.parse(cls.get_sorter_version()) < version.parse("4.0.16"): @@ -461,6 +462,11 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose): if (sorter_output_folder / "recording.dat").is_file(): (sorter_output_folder / "recording.dat").unlink() + # close logger + for handler in logger.handlers.copy(): + logger.removeHandler(handler) + handler.close() + @classmethod def _get_result_from_folder(cls, sorter_output_folder): return KilosortBase._get_result_from_folder(sorter_output_folder) @@ -469,7 +475,6 @@ def _get_result_from_folder(cls, sorter_output_folder): def _setup_json_probe_map(cls, recording, sorter_output_folder): """Create a JSON probe map file for Kilosort4.""" from kilosort.io import save_probe - import numpy as np groups = recording.get_channel_groups() positions = np.array(recording.get_channel_locations()) @@ -492,3 +497,47 @@ def _setup_json_probe_map(cls, recording, sorter_output_folder): "n_chan": n_chan, } save_probe(probe, str(sorter_output_folder / "chanMap.json")) + + +def read_kilosort4_motion(sorter_output_folder: str | Path, recording: BaseRecording | None = None) -> Motion: + """Reads the motion information from a Kilosort4 output folder and returns a Motion object. + + Parameters + ---------- + sorter_output_folder: str or Path + The path to the Kilosort4 output folder. + recording: BaseRecording, optional + The recording object. If provided, the temporal bins will be estimated based on the recording's + start and end times. If not provided, the temporal bins will be estimated based on the number + of batches in the ops file. + + Returns + ------- + Motion + A Motion object containing the displacement, temporal bins, and spatial bins. + + """ + sorter_output_folder = Path(sorter_output_folder) + ops_file = sorter_output_folder / "ops.npy" + if not ops_file.is_file(): + raise FileNotFoundError("'ops.npy' file not found!") + ops = np.load(ops_file, allow_pickle=True).item() + yblk = ops.get("yblk") + dshift = ops.get("dshift") + if yblk is None or dshift is None: + raise Exception("'yblk' and 'dshift' fields not found in ops file!") + displacement = dshift + spatial_bins_um = yblk + # estimate temporal bins + batch_size = ops["batch_size"] + fs = ops["fs"] + t_bin = batch_size / fs + if recording is not None: + t_start = recording.get_start_time() + t_end = recording.get_end_time() + temporal_bins_s = np.linspace(t_start + t_bin / 2, t_end - t_bin / 2, displacement.shape[0]) + else: + temporal_bins_s = np.arange(displacement.shape[0]) * t_bin + t_bin / 2 + + motion = Motion(displacement=displacement, temporal_bins_s=temporal_bins_s, spatial_bins_um=spatial_bins_um) + return motion diff --git a/src/spikeinterface/sorters/sorterlist.py b/src/spikeinterface/sorters/sorterlist.py index 3cd34d7f5f..7122fdf976 100644 --- a/src/spikeinterface/sorters/sorterlist.py +++ b/src/spikeinterface/sorters/sorterlist.py @@ -6,7 +6,7 @@ from .external.kilosort2 import Kilosort2Sorter from .external.kilosort2_5 import Kilosort2_5Sorter from .external.kilosort3 import Kilosort3Sorter -from .external.kilosort4 import Kilosort4Sorter +from .external.kilosort4 import Kilosort4Sorter, read_kilosort4_motion from .external.pykilosort import PyKilosortSorter from .external.klusta import KlustaSorter from .external.mountainsort4 import Mountainsort4Sorter