Add read_kilosort4_motion function (#4518)

Co-authored-by: Chris Halcrow <57948917+chrishalcrow@users.noreply.github.com>

Author: alejoe91

.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
 
 

.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):
         """

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",

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.

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
 ):

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

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