Fix a cross-band interpolation bug, and allow time_vector in interpolate_motion

src/spikeinterface/core/baserecording.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+
 import warnings
 from pathlib import Path
 
@@ -7,14 +8,9 @@
 
 from .base import BaseSegment
 from .baserecordingsnippets import BaseRecordingSnippets
-from .core_tools import (
-    convert_bytes_to_str,
-    convert_seconds_to_str,
-)
-from .recording_tools import write_binary_recording
-
-
+from .core_tools import convert_bytes_to_str, convert_seconds_to_str
 from .job_tools import split_job_kwargs
+from .recording_tools import write_binary_recording
 
 
 class BaseRecording(BaseRecordingSnippets):
@@ -921,11 +917,11 @@ def time_to_sample_index(self, time_s):
                 sample_index = time_s * self.sampling_frequency
             else:
                 sample_index = (time_s - self.t_start) * self.sampling_frequency
-            sample_index = round(sample_index)
+            sample_index = np.round(sample_index).astype(int)
         else:
             sample_index = np.searchsorted(self.time_vector, time_s, side="right") - 1
 
-        return int(sample_index)
+        return sample_index
 
     def get_num_samples(self) -> int:
         """Returns the number of samples in this signal segment

src/spikeinterface/sortingcomponents/motion/motion_interpolation.py

@@ -6,6 +6,8 @@
 from spikeinterface.preprocessing.basepreprocessor import BasePreprocessor, BasePreprocessorSegment
 from spikeinterface.preprocessing.filter import fix_dtype
 
+from .motion_utils import ensure_time_bin_edges, ensure_time_bins
+
 
 def correct_motion_on_peaks(peaks, peak_locations, motion, recording) -> np.ndarray:
     """
@@ -54,6 +56,7 @@ def interpolate_motion_on_traces(
     segment_index=None,
     channel_inds=None,
     interpolation_time_bin_centers_s=None,
+    interpolation_time_bin_edges_s=None,
     spatial_interpolation_method="kriging",
     spatial_interpolation_kwargs={},
     dtype=None,
@@ -119,17 +122,26 @@ def interpolate_motion_on_traces(
     total_num_chans = channel_locations.shape[0]
 
     # -- determine the blocks of frames that will land in the same interpolation time bin
-    time_bins = interpolation_time_bin_centers_s
-    if time_bins is None:
-        time_bins = motion.temporal_bins_s[segment_index]
-    bin_s = time_bins[1] - time_bins[0]
-    bins_start = time_bins[0] - 0.5 * bin_s
-    # nearest bin center for each frame?
-    bin_inds = (times - bins_start) // bin_s
-    bin_inds = bin_inds.astype(int)
+    if interpolation_time_bin_centers_s is None and interpolation_time_bin_edges_s is None:
+        bin_centers_s = motion.temporal_bin_edges_s[segment_index]
+        bin_edges_s = motion.temporal_bin_edges_s[segment_index]
+    else:
+        bin_centers_s, bin_edges_s = ensure_time_bins(interpolation_time_bin_centers_s, interpolation_time_bin_edges_s)
+
+    # nearest interpolation bin:
+    # seachsorted(b, t, side="right") == i means that b[i-1] <= t < b[i]
+    # hence the -1. doing it with "left" is not as nice -- we want t==b[0]
+    # to lead to i=1 (rounding down).
+    # time_bins are bin centers, but we want to snap to the nearest center.
+    # idea is to get the left bin edges and bin the interp times.
+    # this is like subtracting bin_dt_s/2, but allows non-equally-spaced bins.
+    # it's fine to use the first bin center for the first left edge
+    bin_inds = np.searchsorted(bin_edges_s, times, side="right") - 1
+
     # the time bins may not cover the whole set of times in the recording,
     # so we need to clip these indices to the valid range
-    np.clip(bin_inds, 0, time_bins.size, out=bin_inds)
+    n_bins = bin_edges_s.shape[0] - 1
+    np.clip(bin_inds, 0, n_bins - 1, out=bin_inds)
 
     # -- what are the possibilities here anyway?
     bins_here = np.arange(bin_inds[0], bin_inds[-1] + 1)
@@ -138,7 +150,7 @@ def interpolate_motion_on_traces(
     interp_times = np.empty(total_num_chans)
     current_start_index = 0
     for bin_ind in bins_here:
-        bin_time = time_bins[bin_ind]
+        bin_time = bin_centers_s[bin_ind]
         interp_times.fill(bin_time)
         channel_motions = motion.get_displacement_at_time_and_depth(
             interp_times,
@@ -297,6 +309,7 @@ def __init__(
         p=1,
         num_closest=3,
         interpolation_time_bin_centers_s=None,
+        interpolation_time_bin_edges_s=None,
         interpolation_time_bin_size_s=None,
         dtype=None,
         **spatial_interpolation_kwargs,
@@ -363,9 +376,14 @@ def __init__(
 
         # handle manual interpolation_time_bin_centers_s
         # the case where interpolation_time_bin_size_s is set is handled per-segment below
-        if interpolation_time_bin_centers_s is None:
+        if interpolation_time_bin_centers_s is None and interpolation_time_bin_edges_s is None:
             if interpolation_time_bin_size_s is None:
                 interpolation_time_bin_centers_s = motion.temporal_bins_s
+                interpolation_time_bin_edges_s = motion.temporal_bin_edges_s
+        else:
+            interpolation_time_bin_centers_s, interpolation_time_bin_edges_s = ensure_time_bins(
+                interpolation_time_bin_centers_s, interpolation_time_bin_edges_s
+            )
 
         for segment_index, parent_segment in enumerate(recording._recording_segments):
             # finish the per-segment part of the time bin logic
@@ -375,8 +393,13 @@ def __init__(
                 t_start, t_end = parent_segment.sample_index_to_time(np.array([0, s_end]))
                 halfbin = interpolation_time_bin_size_s / 2.0
                 segment_interpolation_time_bins_s = np.arange(t_start + halfbin, t_end, interpolation_time_bin_size_s)
+                segment_interpolation_time_bin_edges_s = np.arange(
+                    t_start, t_end + halfbin, interpolation_time_bin_size_s
+                )
+                assert segment_interpolation_time_bin_edges_s.shape == (segment_interpolation_time_bins_s.shape[0] + 1,)
             else:
                 segment_interpolation_time_bins_s = interpolation_time_bin_centers_s[segment_index]
+                segment_interpolation_time_bin_edges_s = interpolation_time_bin_edges_s[segment_index]
 
             rec_segment = InterpolateMotionRecordingSegment(
                 parent_segment,
@@ -387,6 +410,7 @@ def __init__(
                 channel_inds,
                 segment_index,
                 segment_interpolation_time_bins_s,
+                segment_interpolation_time_bin_edges_s,
                 dtype=dtype_,
             )
             self.add_recording_segment(rec_segment)
@@ -420,6 +444,7 @@ def __init__(
         channel_inds,
         segment_index,
         interpolation_time_bin_centers_s,
+        interpolation_time_bin_edges_s,
         dtype="float32",
     ):
         BasePreprocessorSegment.__init__(self, parent_recording_segment)
@@ -429,13 +454,11 @@ def __init__(
         self.channel_inds = channel_inds
         self.segment_index = segment_index
         self.interpolation_time_bin_centers_s = interpolation_time_bin_centers_s
+        self.interpolation_time_bin_edges_s = interpolation_time_bin_edges_s
         self.dtype = dtype
         self.motion = motion
 
     def get_traces(self, start_frame, end_frame, channel_indices):
-        if self.time_vector is not None:
-            raise NotImplementedError("InterpolateMotionRecording does not yet support recordings with time_vectors.")
-
         if start_frame is None:
             start_frame = 0
         if end_frame is None:
@@ -453,7 +476,7 @@ def get_traces(self, start_frame, end_frame, channel_indices):
             channel_inds=self.channel_inds,
             spatial_interpolation_method=self.spatial_interpolation_method,
             spatial_interpolation_kwargs=self.spatial_interpolation_kwargs,
-            interpolation_time_bin_centers_s=self.interpolation_time_bin_centers_s,
+            interpolation_time_bin_edges_s=self.interpolation_time_bin_edges_s,
         )
 
         if channel_indices is not None:

src/spikeinterface/sortingcomponents/motion/motion_utils.py

@@ -1,5 +1,5 @@
-import warnings
 import json
+import warnings
 from pathlib import Path
 
 import numpy as np
@@ -54,6 +54,7 @@ def __init__(self, displacement, temporal_bins_s, spatial_bins_um, direction="y"
         self.direction = direction
         self.dim = ["x", "y", "z"].index(direction)
         self.check_properties()
+        self.temporal_bin_edges_s = [ensure_time_bin_edges(tbins) for tbins in self.temporal_bins_s]
 
     def check_properties(self):
         assert all(d.ndim == 2 for d in self.displacement)
@@ -576,3 +577,24 @@ def make_3d_motion_histograms(
         motion_histograms = np.log2(1 + motion_histograms)
 
     return motion_histograms, temporal_bin_edges, spatial_bin_edges
+
+
+def ensure_time_bins(time_bin_centers_s=None, time_bin_edges_s=None):
+    if time_bin_centers_s is None and time_bin_edges_s is None:
+        raise ValueError("Need at least one of time_bin_centers_s or time_bin_edges_s.")
+
+    if time_bin_centers_s is None:
+        assert time_bin_edges_s.ndim == 1 and time_bin_edges_s.size >= 2
+        time_bin_centers_s = 0.5 * (time_bin_edges_s[1:] + time_bin_edges_s[:-1])
+
+    if time_bin_edges_s is None:
+        time_bin_edges_s = np.empty(time_bin_centers_s.shape[0] + 1, dtype=time_bin_centers_s.dtype)
+        time_bin_edges_s[[0, -1]] = time_bin_centers_s[[0, -1]]
+        if time_bin_centers_s.size > 2:
+            time_bin_edges_s[1:-1] = 0.5 * (time_bin_centers_s[1:] + time_bin_centers_s[:-1])
+
+    return time_bin_centers_s, time_bin_edges_s
+
+
+def ensure_time_bin_edges(time_bin_centers_s=None, time_bin_edges_s=None):
+    return ensure_time_bins(time_bin_centers_s, time_bin_edges_s)[1]

src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py

@@ -1,16 +1,14 @@
-from pathlib import Path
+import warnings
 
 import numpy as np
-import pytest
 import spikeinterface.core as sc
-from spikeinterface import download_dataset
+from spikeinterface.sortingcomponents.motion import Motion
 from spikeinterface.sortingcomponents.motion.motion_interpolation import (
     InterpolateMotionRecording,
     correct_motion_on_peaks,
     interpolate_motion,
     interpolate_motion_on_traces,
 )
-from spikeinterface.sortingcomponents.motion import Motion
 from spikeinterface.sortingcomponents.tests.common import make_dataset
 
 
@@ -67,18 +65,20 @@ def test_interpolate_motion_on_traces():
     times = rec.get_times()[0:30000]
 
     for method in ("kriging", "idw", "nearest"):
-        traces_corrected = interpolate_motion_on_traces(
-            traces,
-            times,
-            channel_locations,
-            motion,
-            channel_inds=None,
-            spatial_interpolation_method=method,
-            # spatial_interpolation_kwargs={},
-            spatial_interpolation_kwargs={"force_extrapolate": True},
-        )
-        assert traces.shape == traces_corrected.shape
-        assert traces.dtype == traces_corrected.dtype
+        for interpolation_time_bin_centers_s in (None, np.linspace(*times[[0, -1]], num=3)):
+            traces_corrected = interpolate_motion_on_traces(
+                traces,
+                times,
+                channel_locations,
+                motion,
+                channel_inds=None,
+                spatial_interpolation_method=method,
+                interpolation_time_bin_centers_s=interpolation_time_bin_centers_s,
+                # spatial_interpolation_kwargs={},
+                spatial_interpolation_kwargs={"force_extrapolate": True},
+            )
+            assert traces.shape == traces_corrected.shape
+            assert traces.dtype == traces_corrected.dtype
 
 
 def test_interpolation_simple():
@@ -115,6 +115,66 @@ def test_interpolation_simple():
     assert np.all(traces_corrected[:, 2:] == 0)
 
 
+def test_cross_band_interpolation():
+    """Simple version of using LFP to interpolate AP data
+
+    This also tests the time vector implementation in interpolation.
+    The idea is to have two recordings which are all 0s with a 1 that
+    moves from one channel to another after 3s. They're at different
+    sampling frequencies. motion estimation in one sampling frequency
+    applied to the other should still lead to perfect correction.
+    """
+    from spikeinterface.sortingcomponents.motion import estimate_motion
+
+    # sampling freqs and timing for AP and LFP recordings
+    fs_lfp = 50.0
+    fs_ap = 300.0
+    t_start = 10.0
+    total_duration = 5.0
+    nt_lfp = int(fs_lfp * total_duration)
+    nt_ap = int(fs_ap * total_duration)
+    t_switch = 3
+
+    # because interpolation uses bin centers logic, there will be a half
+    # bin offset at the change point in the AP recording.
+    halfbin_ap_lfp = int(0.5 * (fs_ap / fs_lfp))
+
+    # channel geometry
+    nc = 10
+    geom = np.c_[np.zeros(nc), np.arange(nc)]
+
+    # make an LFP recording which drifts a bit
+    traces_lfp = np.zeros((nt_lfp, nc))
+    traces_lfp[: int(t_switch * fs_lfp), 5] = 1.0
+    traces_lfp[int(t_switch * fs_lfp) :, 6] = 1.0
+    rec_lfp = sc.NumpyRecording(traces_lfp, sampling_frequency=fs_lfp)
+    rec_lfp.set_dummy_probe_from_locations(geom)
+
+    # same for AP
+    traces_ap = np.zeros((nt_ap, nc))
+    traces_ap[: int(t_switch * fs_ap) - halfbin_ap_lfp, 5] = 1.0
+    traces_ap[int(t_switch * fs_ap) - halfbin_ap_lfp :, 6] = 1.0
+    rec_ap = sc.NumpyRecording(traces_ap, sampling_frequency=fs_ap)
+    rec_ap.set_dummy_probe_from_locations(geom)
+
+    # set times for both, and silence the warning
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=UserWarning)
+        rec_lfp.set_times(t_start + np.arange(nt_lfp) / fs_lfp)
+        rec_ap.set_times(t_start + np.arange(nt_ap) / fs_ap)
+
+    # estimate motion
+    motion = estimate_motion(rec_lfp, method="dredge_lfp", rigid=True)
+
+    # nearest to keep it simple
+    rec_corrected = interpolate_motion(rec_ap, motion, spatial_interpolation_method="nearest", num_closest=2)
+    traces_corrected = rec_corrected.get_traces()
+    target = np.zeros((nt_ap, nc - 2))
+    target[:, 4] = 1
+    ii, jj = np.nonzero(traces_corrected)
+    assert np.array_equal(traces_corrected, target)
+
+
 def test_InterpolateMotionRecording():
     rec, sorting = make_dataset()
     motion = make_fake_motion(rec)
@@ -147,6 +207,7 @@ def test_InterpolateMotionRecording():
 
 if __name__ == "__main__":
     # test_correct_motion_on_peaks()
-    # test_interpolate_motion_on_traces()
-    test_interpolation_simple()
-    test_InterpolateMotionRecording()
+    test_interpolate_motion_on_traces()
+    # test_interpolation_simple()
+    # test_InterpolateMotionRecording()
+    test_cross_band_interpolation()