From 507b6b3cf19d0f10069e2415f134dca7fb709b47 Mon Sep 17 00:00:00 2001
From: Charlie Windolf <charliexwindolf@gmail.com>
Date: Sun, 3 Nov 2024 19:16:26 -0500
Subject: [PATCH 01/17] Address time bin issue arising in LFP-based reg, which
 AP-based reg doesn't trigger

---
 .../motion/motion_interpolation.py              | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
index a5e6ded519..975f43919d 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
@@ -3,7 +3,8 @@
 import numpy as np
 from spikeinterface.core.core_tools import define_function_from_class
 from spikeinterface.preprocessing import get_spatial_interpolation_kernel
-from spikeinterface.preprocessing.basepreprocessor import BasePreprocessor, BasePreprocessorSegment
+from spikeinterface.preprocessing.basepreprocessor import (
+    BasePreprocessor, BasePreprocessorSegment)
 from spikeinterface.preprocessing.filter import fix_dtype
 
 
@@ -122,14 +123,18 @@ def interpolate_motion_on_traces(
     time_bins = interpolation_time_bin_centers_s
     if time_bins is None:
         time_bins = motion.temporal_bins_s[segment_index]
+
+    # 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.
+    # time_bins are bin centers, so subtract half the bin length. this leads
+    # to snapping to the nearest bin center.
     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)
+    bin_inds = np.searchsorted(time_bins - bin_s / 2, 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)
+    np.clip(bin_inds, 0, time_bins.size - 1, out=bin_inds)
 
     # -- what are the possibilities here anyway?
     bins_here = np.arange(bin_inds[0], bin_inds[-1] + 1)

From 4e38ac18be65051d30d15f3d25bada943af3e31f Mon Sep 17 00:00:00 2001
From: Charlie Windolf <charliexwindolf@gmail.com>
Date: Mon, 4 Nov 2024 11:06:46 -0500
Subject: [PATCH 02/17] Fix LFP-based AP interp bug and allow time_vector in
 interpolation

---
 .../motion/motion_interpolation.py            | 18 +++--
 .../motion/tests/test_motion_interpolation.py | 78 ++++++++++++++++---
 2 files changed, 77 insertions(+), 19 deletions(-)

diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
index 975f43919d..4fd42a8b39 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
@@ -126,11 +126,16 @@ def interpolate_motion_on_traces(
 
     # 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.
-    # time_bins are bin centers, so subtract half the bin length. this leads
-    # to snapping to the nearest bin center.
-    bin_s = time_bins[1] - time_bins[0]
-    bin_inds = np.searchsorted(time_bins - bin_s / 2, times, side="right") - 1
+    # 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.
+    bin_left = np.zeros_like(time_bins)
+    # it's fine to use the first bin center for the first left edge
+    bin_left[0] = time_bins[0]
+    bin_left[1:] = 0.5 * (time_bins[1:] + time_bins[:-1])
+    bin_inds = np.searchsorted(bin_left, 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
@@ -438,9 +443,6 @@ def __init__(
         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:
diff --git a/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
index e022f0cc6c..69f681a1be 100644
--- a/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
@@ -1,16 +1,11 @@
-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.motion_interpolation import (
-    InterpolateMotionRecording,
-    correct_motion_on_peaks,
-    interpolate_motion,
-    interpolate_motion_on_traces,
-)
 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.tests.common import make_dataset
 
 
@@ -115,6 +110,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)
@@ -148,5 +203,6 @@ def test_InterpolateMotionRecording():
 if __name__ == "__main__":
     # test_correct_motion_on_peaks()
     # test_interpolate_motion_on_traces()
-    test_interpolation_simple()
-    test_InterpolateMotionRecording()
+    # test_interpolation_simple()
+    # test_InterpolateMotionRecording()
+    test_cross_band_interpolation()

From 726170b1526b954b5a26edd70d3162e476ed9f53 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 4 Nov 2024 16:28:44 +0000
Subject: [PATCH 03/17] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../sortingcomponents/motion/motion_interpolation.py       | 3 +--
 .../motion/tests/test_motion_interpolation.py              | 7 +++++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
index 4fd42a8b39..810264d9e4 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
@@ -3,8 +3,7 @@
 import numpy as np
 from spikeinterface.core.core_tools import define_function_from_class
 from spikeinterface.preprocessing import get_spatial_interpolation_kernel
-from spikeinterface.preprocessing.basepreprocessor import (
-    BasePreprocessor, BasePreprocessorSegment)
+from spikeinterface.preprocessing.basepreprocessor import BasePreprocessor, BasePreprocessorSegment
 from spikeinterface.preprocessing.filter import fix_dtype
 
 
diff --git a/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
index 69f681a1be..88af619220 100644
--- a/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
@@ -4,8 +4,11 @@
 import spikeinterface.core as sc
 from spikeinterface.sortingcomponents.motion import Motion
 from spikeinterface.sortingcomponents.motion.motion_interpolation import (
-    InterpolateMotionRecording, correct_motion_on_peaks, interpolate_motion,
-    interpolate_motion_on_traces)
+    InterpolateMotionRecording,
+    correct_motion_on_peaks,
+    interpolate_motion,
+    interpolate_motion_on_traces,
+)
 from spikeinterface.sortingcomponents.tests.common import make_dataset
 
 

From e791fe18671c2998fde9d44295c54a5781ca2e46 Mon Sep 17 00:00:00 2001
From: Charlie Windolf <charliexwindolf@gmail.com>
Date: Mon, 11 Nov 2024 10:43:48 -0500
Subject: [PATCH 04/17] Cache bin edges in 2 places as discussed with Sam

---
 .../motion/motion_interpolation.py            | 39 +++++++++++++------
 .../sortingcomponents/motion/motion_utils.py  | 24 +++++++++++-
 .../motion/tests/test_motion_interpolation.py | 28 ++++++-------
 3 files changed, 66 insertions(+), 25 deletions(-)

diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
index 4fd42a8b39..89696f5041 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
@@ -7,6 +7,8 @@
     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:
     """
@@ -55,6 +57,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,
@@ -120,9 +123,11 @@ 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]
+    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]
@@ -131,15 +136,13 @@ def interpolate_motion_on_traces(
     # 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.
-    bin_left = np.zeros_like(time_bins)
     # it's fine to use the first bin center for the first left edge
-    bin_left[0] = time_bins[0]
-    bin_left[1:] = 0.5 * (time_bins[1:] + time_bins[:-1])
-    bin_inds = np.searchsorted(bin_left, times, side="right") - 1
+    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 - 1, 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)
@@ -148,7 +151,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,
@@ -307,6 +310,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,
@@ -373,9 +377,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
@@ -385,8 +394,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,
@@ -397,6 +411,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)
@@ -430,6 +445,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)
@@ -439,6 +455,7 @@ 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
 
@@ -460,7 +477,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:
diff --git a/src/spikeinterface/sortingcomponents/motion/motion_utils.py b/src/spikeinterface/sortingcomponents/motion/motion_utils.py
index 635624cca8..ec0a55a8f8 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_utils.py
+++ b/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]
diff --git a/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
index 69f681a1be..07cb5b8ab6 100644
--- a/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
@@ -62,18 +62,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():
@@ -202,7 +204,7 @@ def test_InterpolateMotionRecording():
 
 if __name__ == "__main__":
     # test_correct_motion_on_peaks()
-    # test_interpolate_motion_on_traces()
+    test_interpolate_motion_on_traces()
     # test_interpolation_simple()
     # test_InterpolateMotionRecording()
     test_cross_band_interpolation()

From d8f39b5a70dd83f4e1fff71d41036692fba20b38 Mon Sep 17 00:00:00 2001
From: Charlie Windolf <charliexwindolf@gmail.com>
Date: Mon, 11 Nov 2024 12:30:32 -0500
Subject: [PATCH 05/17] Sorry if this is shoe-horning in a change...

---
 src/spikeinterface/core/baserecording.py | 17 +++++++----------
 1 file changed, 7 insertions(+), 10 deletions(-)

diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py
index 5e2e9e4014..b95bfb1ad0 100644
--- a/src/spikeinterface/core/baserecording.py
+++ b/src/spikeinterface/core/baserecording.py
@@ -1,20 +1,17 @@
 from __future__ import annotations
+
 import warnings
 from pathlib import Path
 
 import numpy as np
-from probeinterface import Probe, ProbeGroup, read_probeinterface, select_axes, write_probeinterface
+from probeinterface import (Probe, ProbeGroup, read_probeinterface,
+                            select_axes, write_probeinterface)
 
 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 +918,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

From 0a201e17a0b3de283f06c5456010fb20fd8cd209 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 11 Nov 2024 17:31:00 +0000
Subject: [PATCH 06/17] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 src/spikeinterface/core/baserecording.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py
index b95bfb1ad0..6d4509db12 100644
--- a/src/spikeinterface/core/baserecording.py
+++ b/src/spikeinterface/core/baserecording.py
@@ -4,8 +4,7 @@
 from pathlib import Path
 
 import numpy as np
-from probeinterface import (Probe, ProbeGroup, read_probeinterface,
-                            select_axes, write_probeinterface)
+from probeinterface import Probe, ProbeGroup, read_probeinterface, select_axes, write_probeinterface
 
 from .base import BaseSegment
 from .baserecordingsnippets import BaseRecordingSnippets

From ad00beb182967ebe68f59ebfd7f1abad3002a10e Mon Sep 17 00:00:00 2001
From: Charlie Windolf <charliexwindolf@gmail.com>
Date: Thu, 21 Nov 2024 20:44:20 +0000
Subject: [PATCH 07/17] Update `interpolate_motion_on_traces` docstring

---
 .../sortingcomponents/motion/motion_interpolation.py           | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
index d207dced08..7c6f0ba71a 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
@@ -83,6 +83,9 @@ def interpolate_motion_on_traces(
     interpolation_time_bin_centers_s : None or np.array
         Manually specify the time bins which the interpolation happens
         in for this segment. If None, these are the motion estimate's time bins.
+    interpolation_time_bin_edges_s : None or np.array
+        If present, interpolation chunks will be the time bins defined by these edges
+        rather than interpolation_time_bin_centers_s or the motion's bins.
     spatial_interpolation_method : "idw" | "kriging", default: "kriging"
         The spatial interpolation method used to interpolate the channel locations:
             * idw : Inverse Distance Weighing

From 28527d2a8bcd03a3c6d9036ac9be180ad4dc6334 Mon Sep 17 00:00:00 2001
From: Charlie Windolf <charliexwindolf@gmail.com>
Date: Thu, 21 Nov 2024 20:46:30 +0000
Subject: [PATCH 08/17] Should be centers!

---
 .../sortingcomponents/motion/motion_interpolation.py            | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
index 7c6f0ba71a..f0fff5c039 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
@@ -126,7 +126,7 @@ def interpolate_motion_on_traces(
 
     # -- determine the blocks of frames that will land in the same interpolation time bin
     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_centers_s = motion.temporal_bin_centers_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)

From b3b3fcf5be7b67451f54cac753844cbb06b8d4a3 Mon Sep 17 00:00:00 2001
From: Charlie Windolf <charliexwindolf@gmail.com>
Date: Thu, 21 Nov 2024 20:47:34 +0000
Subject: [PATCH 09/17] Typo

---
 .../sortingcomponents/motion/motion_interpolation.py            | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
index f0fff5c039..9e32e189d9 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
@@ -132,7 +132,7 @@ def interpolate_motion_on_traces(
         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]
+    # searchsorted(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.

From b02860e463262a3b9522ebf9badc45c860cf8d29 Mon Sep 17 00:00:00 2001
From: Charlie Windolf <charliexwindolf@gmail.com>
Date: Thu, 21 Nov 2024 20:51:11 +0000
Subject: [PATCH 10/17] Clarify comments

---
 .../sortingcomponents/motion/motion_interpolation.py        | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
index 9e32e189d9..8f96579228 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
@@ -131,14 +131,10 @@ def interpolate_motion_on_traces(
     else:
         bin_centers_s, bin_edges_s = ensure_time_bins(interpolation_time_bin_centers_s, interpolation_time_bin_edges_s)
 
-    # nearest interpolation bin:
+    # bin the frame times according to the interpolation time bins.
     # searchsorted(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,

From df2484002b562861f0f64bc053b5da619253e983 Mon Sep 17 00:00:00 2001
From: Charlie Windolf <charliexwindolf@gmail.com>
Date: Thu, 21 Nov 2024 20:57:08 +0000
Subject: [PATCH 11/17] Rename variables for clarity

---
 .../motion/motion_interpolation.py            | 25 ++++++++++---------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
index 8f96579228..2bd3493650 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
@@ -126,30 +126,30 @@ def interpolate_motion_on_traces(
 
     # -- determine the blocks of frames that will land in the same interpolation time bin
     if interpolation_time_bin_centers_s is None and interpolation_time_bin_edges_s is None:
-        bin_centers_s = motion.temporal_bin_centers_s[segment_index]
-        bin_edges_s = motion.temporal_bin_edges_s[segment_index]
+        interpolation_time_bin_centers_s = motion.temporal_bin_centers_s[segment_index]
+        interpolation_time_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)
+        interpolation_time_bin_centers_s, interpolation_time_bin_edges_s = ensure_time_bins(interpolation_time_bin_centers_s, interpolation_time_bin_edges_s)
 
     # bin the frame times according to the interpolation time bins.
     # searchsorted(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).
-    bin_inds = np.searchsorted(bin_edges_s, times, side="right") - 1
+    interpolation_bin_inds = np.searchsorted(interpolation_time_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
-    n_bins = bin_edges_s.shape[0] - 1
-    np.clip(bin_inds, 0, n_bins - 1, out=bin_inds)
+    n_bins = interpolation_time_bin_edges_s.shape[0] - 1
+    np.clip(interpolation_bin_inds, 0, n_bins - 1, out=interpolation_bin_inds)
 
     # -- what are the possibilities here anyway?
-    bins_here = np.arange(bin_inds[0], bin_inds[-1] + 1)
+    interpolation_bins_here = np.arange(interpolation_bin_inds[0], interpolation_bin_inds[-1] + 1)
 
     # inperpolation kernel will be the same per temporal bin
     interp_times = np.empty(total_num_chans)
     current_start_index = 0
-    for bin_ind in bins_here:
-        bin_time = bin_centers_s[bin_ind]
+    for interp_bin_ind in interpolation_bins_here:
+        bin_time = bin_centers_s[interp_bin_ind]
         interp_times.fill(bin_time)
         channel_motions = motion.get_displacement_at_time_and_depth(
             interp_times,
@@ -177,16 +177,17 @@ def interpolate_motion_on_traces(
         # ax.set_title(f"bin_ind {bin_ind} - {bin_time}s - {spatial_interpolation_method}")
         # plt.show()
 
+        # quick search logic to find frames corresponding to this interpolation bin in the recording
         # quickly find the end of this bin, which is also the start of the next
         next_start_index = current_start_index + np.searchsorted(
-            bin_inds[current_start_index:], bin_ind + 1, side="left"
+            interpolation_bin_inds[current_start_index:], interp_bin_ind + 1, side="left"
         )
-        in_bin = slice(current_start_index, next_start_index)
+        frames_in_bin = slice(current_start_index, next_start_index)
 
         # here we use a simple np.matmul even if dirft_kernel can be super sparse.
         # because the speed for a sparse matmul is not so good when we disable multi threaad (due multi processing
         # in ChunkRecordingExecutor)
-        np.matmul(traces[in_bin], drift_kernel, out=traces_corrected[in_bin])
+        np.matmul(traces[frames_in_bin], drift_kernel, out=traces_corrected[frames_in_bin])
         current_start_index = next_start_index
 
     return traces_corrected

From 91fb7320eafbcdaa8fe728e36cbc7fa686f32ba8 Mon Sep 17 00:00:00 2001
From: Charlie Windolf <charliexwindolf@gmail.com>
Date: Thu, 21 Nov 2024 20:59:13 +0000
Subject: [PATCH 12/17] Note on clipping behavior

---
 .../sortingcomponents/motion/motion_interpolation.py        | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
index 2bd3493650..8aed8085bf 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
@@ -64,7 +64,11 @@ def interpolate_motion_on_traces(
     """
     Apply inverse motion with spatial interpolation on traces.
 
-    Traces can be full traces, but also waveforms snippets.
+    Traces can be full traces, but also waveforms snippets. Times used for looking up
+    displacements are controlled by interpolation_time_bin_edges_s or 
+    interpolation_time_bin_centers_s, or fall back to the Motion object's time bins
+    by default; times in the recording outside these time bins use the closest edge
+    bin's displacement value during interpolation.
 
     Parameters
     ----------

From b80bad71e2c1ee316048beb79a9169dd8f68c6ff Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 21 Nov 2024 21:01:47 +0000
Subject: [PATCH 13/17] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../sortingcomponents/motion/motion_interpolation.py        | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
index 8aed8085bf..14471f77fc 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
@@ -65,7 +65,7 @@ def interpolate_motion_on_traces(
     Apply inverse motion with spatial interpolation on traces.
 
     Traces can be full traces, but also waveforms snippets. Times used for looking up
-    displacements are controlled by interpolation_time_bin_edges_s or 
+    displacements are controlled by interpolation_time_bin_edges_s or
     interpolation_time_bin_centers_s, or fall back to the Motion object's time bins
     by default; times in the recording outside these time bins use the closest edge
     bin's displacement value during interpolation.
@@ -133,7 +133,9 @@ def interpolate_motion_on_traces(
         interpolation_time_bin_centers_s = motion.temporal_bin_centers_s[segment_index]
         interpolation_time_bin_edges_s = motion.temporal_bin_edges_s[segment_index]
     else:
-        interpolation_time_bin_centers_s, interpolation_time_bin_edges_s = ensure_time_bins(interpolation_time_bin_centers_s, interpolation_time_bin_edges_s)
+        interpolation_time_bin_centers_s, interpolation_time_bin_edges_s = ensure_time_bins(
+            interpolation_time_bin_centers_s, interpolation_time_bin_edges_s
+        )
 
     # bin the frame times according to the interpolation time bins.
     # searchsorted(b, t, side="right") == i means that b[i-1] <= t < b[i]

From c89060314e233b89e4e9112e74c7643545806d22 Mon Sep 17 00:00:00 2001
From: Charlie Windolf <charliexwindolf@gmail.com>
Date: Thu, 21 Nov 2024 16:09:45 -0500
Subject: [PATCH 14/17] Fix variable typo; add docstring

---
 .../motion/motion_interpolation.py               |  5 +++--
 .../sortingcomponents/motion/motion_utils.py     | 16 ++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
index 14471f77fc..e87f83751c 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
@@ -3,7 +3,8 @@
 import numpy as np
 from spikeinterface.core.core_tools import define_function_from_class
 from spikeinterface.preprocessing import get_spatial_interpolation_kernel
-from spikeinterface.preprocessing.basepreprocessor import BasePreprocessor, BasePreprocessorSegment
+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
@@ -155,7 +156,7 @@ def interpolate_motion_on_traces(
     interp_times = np.empty(total_num_chans)
     current_start_index = 0
     for interp_bin_ind in interpolation_bins_here:
-        bin_time = bin_centers_s[interp_bin_ind]
+        bin_time = interpolation_time_bin_centers_s[interp_bin_ind]
         interp_times.fill(bin_time)
         channel_motions = motion.get_displacement_at_time_and_depth(
             interp_times,
diff --git a/src/spikeinterface/sortingcomponents/motion/motion_utils.py b/src/spikeinterface/sortingcomponents/motion/motion_utils.py
index ec0a55a8f8..680d75f221 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_utils.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_utils.py
@@ -580,6 +580,22 @@ def make_3d_motion_histograms(
 
 
 def ensure_time_bins(time_bin_centers_s=None, time_bin_edges_s=None):
+    """Ensure that both bin edges and bin centers are present
+
+    If either of the inputs are None but not both, the missing is reconstructed
+    from the present. Going from edges to centers is done by taking midpoints.
+    Going from centers to edges is done by taking midpoints and padding with the
+    left and rightmost centers.
+
+    Parameters
+    ----------
+    time_bin_centers_s : None or np.array
+    time_bin_edges_s : None or np.array
+
+    Returns
+    -------
+    time_bin_centers_s, time_bin_edges_s
+    """
     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.")
 

From 6d2e47911a5e64737f2c98f9cdc199e9f6d306fc Mon Sep 17 00:00:00 2001
From: Charlie Windolf <charliexwindolf@gmail.com>
Date: Thu, 21 Nov 2024 16:14:00 -0500
Subject: [PATCH 15/17] Variable names in tests

---
 .../motion/tests/test_motion_interpolation.py | 45 +++++++++----------
 1 file changed, 21 insertions(+), 24 deletions(-)

diff --git a/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
index 8542b62524..c97c8324ba 100644
--- a/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
@@ -4,11 +4,8 @@
 import spikeinterface.core as sc
 from spikeinterface.sortingcomponents.motion import Motion
 from spikeinterface.sortingcomponents.motion.motion_interpolation import (
-    InterpolateMotionRecording,
-    correct_motion_on_peaks,
-    interpolate_motion,
-    interpolate_motion_on_traces,
-)
+    InterpolateMotionRecording, correct_motion_on_peaks, interpolate_motion,
+    interpolate_motion_on_traces)
 from spikeinterface.sortingcomponents.tests.common import make_dataset
 
 
@@ -84,26 +81,26 @@ def test_interpolate_motion_on_traces():
 def test_interpolation_simple():
     # a recording where a 1 moves at 1 chan per second. 30 chans 10 frames.
     # there will be 9 chans of drift, so we add 9 chans of padding to the bottom
-    nt = nc0 = 10  # these need to be the same for this test
-    nc1 = nc0 + nc0 - 1
-    traces = np.zeros((nt, nc1), dtype="float32")
-    traces[:, :nc0] = np.eye(nc0)
+    n_samples = num_chans_orig = 10  # these need to be the same for this test
+    num_chans_drifted = num_chans_orig + num_chans_orig - 1
+    traces = np.zeros((n_samples, num_chans_drifted), dtype="float32")
+    traces[:, :num_chans_orig] = np.eye(num_chans_orig)
     rec = sc.NumpyRecording(traces, sampling_frequency=1)
-    rec.set_dummy_probe_from_locations(np.c_[np.zeros(nc1), np.arange(nc1)])
+    rec.set_dummy_probe_from_locations(np.c_[np.zeros(num_chans_drifted), np.arange(num_chans_drifted)])
 
-    true_motion = Motion(np.arange(nt)[:, None], 0.5 + np.arange(nt), np.zeros(1))
+    true_motion = Motion(np.arange(n_samples)[:, None], 0.5 + np.arange(n_samples), np.zeros(1))
     rec_corrected = interpolate_motion(rec, true_motion, spatial_interpolation_method="nearest")
     traces_corrected = rec_corrected.get_traces()
-    assert traces_corrected.shape == (nc0, nc0)
-    assert np.array_equal(traces_corrected[:, 0], np.ones(nt))
-    assert np.array_equal(traces_corrected[:, 1:], np.zeros((nt, nc0 - 1)))
+    assert traces_corrected.shape == (num_chans_orig, num_chans_orig)
+    assert np.array_equal(traces_corrected[:, 0], np.ones(n_samples))
+    assert np.array_equal(traces_corrected[:, 1:], np.zeros((n_samples, num_chans_orig - 1)))
 
     # let's try a new version where we interpolate too slowly
     rec_corrected = interpolate_motion(
         rec, true_motion, spatial_interpolation_method="nearest", num_closest=2, interpolation_time_bin_size_s=2
     )
     traces_corrected = rec_corrected.get_traces()
-    assert traces_corrected.shape == (nc0, nc0)
+    assert traces_corrected.shape == (num_chans_orig, num_chans_orig)
     # what happens with nearest here?
     # well... due to rounding towards the nearest even number, the motion (which at
     # these time bin centers is 0.5, 2.5, 4.5, ...) flips the signal's nearest
@@ -131,8 +128,8 @@ def test_cross_band_interpolation():
     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)
+    num_samples_lfp = int(fs_lfp * total_duration)
+    num_samples_ap = int(fs_ap * total_duration)
     t_switch = 3
 
     # because interpolation uses bin centers logic, there will be a half
@@ -140,18 +137,18 @@ def test_cross_band_interpolation():
     halfbin_ap_lfp = int(0.5 * (fs_ap / fs_lfp))
 
     # channel geometry
-    nc = 10
-    geom = np.c_[np.zeros(nc), np.arange(nc)]
+    num_chans = 10
+    geom = np.c_[np.zeros(num_chans), np.arange(num_chans)]
 
     # make an LFP recording which drifts a bit
-    traces_lfp = np.zeros((nt_lfp, nc))
+    traces_lfp = np.zeros((num_samples_lfp, num_chans))
     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 = np.zeros((num_samples_ap, num_chans))
     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)
@@ -160,8 +157,8 @@ def test_cross_band_interpolation():
     # 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)
+        rec_lfp.set_times(t_start + np.arange(num_samples_lfp) / fs_lfp)
+        rec_ap.set_times(t_start + np.arange(num_samples_ap) / fs_ap)
 
     # estimate motion
     motion = estimate_motion(rec_lfp, method="dredge_lfp", rigid=True)
@@ -169,7 +166,7 @@ def test_cross_band_interpolation():
     # 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 = np.zeros((num_samples_ap, num_chans - 2))
     target[:, 4] = 1
     ii, jj = np.nonzero(traces_corrected)
     assert np.array_equal(traces_corrected, target)

From b4c91a0d941e97be68908b11974d18f802ec74a5 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 21 Nov 2024 21:15:30 +0000
Subject: [PATCH 16/17] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../sortingcomponents/motion/motion_interpolation.py       | 3 +--
 .../motion/tests/test_motion_interpolation.py              | 7 +++++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
index e87f83751c..b3a4c9a207 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
@@ -3,8 +3,7 @@
 import numpy as np
 from spikeinterface.core.core_tools import define_function_from_class
 from spikeinterface.preprocessing import get_spatial_interpolation_kernel
-from spikeinterface.preprocessing.basepreprocessor import (
-    BasePreprocessor, BasePreprocessorSegment)
+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
diff --git a/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
index c97c8324ba..e4ba870325 100644
--- a/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
@@ -4,8 +4,11 @@
 import spikeinterface.core as sc
 from spikeinterface.sortingcomponents.motion import Motion
 from spikeinterface.sortingcomponents.motion.motion_interpolation import (
-    InterpolateMotionRecording, correct_motion_on_peaks, interpolate_motion,
-    interpolate_motion_on_traces)
+    InterpolateMotionRecording,
+    correct_motion_on_peaks,
+    interpolate_motion,
+    interpolate_motion_on_traces,
+)
 from spikeinterface.sortingcomponents.tests.common import make_dataset
 
 

From 38e0adac18c2c862f9193e5a5d643b394a6d3d52 Mon Sep 17 00:00:00 2001
From: Charlie Windolf <charliexwindolf@gmail.com>
Date: Fri, 22 Nov 2024 10:28:34 -0500
Subject: [PATCH 17/17] Typo

---
 .../sortingcomponents/motion/motion_interpolation.py            | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
index b3a4c9a207..fc8ccb788b 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_interpolation.py
@@ -130,7 +130,7 @@ def interpolate_motion_on_traces(
 
     # -- determine the blocks of frames that will land in the same interpolation time bin
     if interpolation_time_bin_centers_s is None and interpolation_time_bin_edges_s is None:
-        interpolation_time_bin_centers_s = motion.temporal_bin_centers_s[segment_index]
+        interpolation_time_bin_centers_s = motion.temporal_bins_s[segment_index]
         interpolation_time_bin_edges_s = motion.temporal_bin_edges_s[segment_index]
     else:
         interpolation_time_bin_centers_s, interpolation_time_bin_edges_s = ensure_time_bins(