fix: read function displacement, add test for read_kilosort4_motion function

Author: alejoe91

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

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/sorters/external/kilosort4.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 from packaging import version
 
+import numpy as np
 
 from spikeinterface.core import write_binary_recording, Motion, BaseRecording
 from spikeinterface.sorters.basesorter import BaseSorter, get_job_kwargs
@@ -171,7 +172,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"):
@@ -468,7 +468,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())
@@ -520,7 +519,7 @@ def read_kilosort4_motion(sorter_output_folder: str | Path, recording: BaseRecor
     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 + yblk
+    displacement = dshift
     spatial_bins_um = yblk
     # estimate temporal bins
     batch_size = ops["batch_size"]
@@ -529,7 +528,7 @@ def read_kilosort4_motion(sorter_output_folder: str | Path, recording: BaseRecor
     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)
+        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