Repository: SpikeInterface/spikeinterface

Issue:
"hard" merge of units in sorting analyzer fails on extensions computations
Hi all, 

I am encountering an error during the extension computation (amplitude_scalings + spike_amplitudes + spike_locations) when applying curation (with merges) on an analyzer. 

This error happens only when merging mode is set to `"hard"` in the call of `apply_curation()`. I'm not very familiar with the internal spikeinterface code for computing extensions so I'm struggling to debug this and I would love a bit of help. Note that when applying the curation to the sorting object, it works fine. 

You'll find here the analyzer, sorting and recording object together with the curation_dict.json file that I'm using so you can reproduce (hopefully) this behavior : https://instituteicm-my.sharepoint.com/:f:/g/personal/anthony_pinto_icm-institute_org/EhU1nQWpaNBFtVovmQaDUeEB5RJnkYReYfWu4ChMQ088OA?e=Aw9XSe

Here is also the code that you need to be able to reproduce : 

```python
import spikeinterface as si
import json
from spikeinterface.curation import apply_curation

raw_sorting = si.load_extractor(r'path_to_sorting')
raw_analyzer = si.load_sorting_analyzer(r'path_to_analyzer')
recording = si.load_extractor(r'path_to_recording')
raw_analyzer._recording = recording
with open(r'path_to_curation_dict', 'r') as f:
    curation_dict = json.load(f)

print(recording), print(raw_sorting), print(raw_analyzer), print(curation_dict)

curated_analyzer = apply_curation(raw_analyzer, curation_dict, new_id_strategy='take_first', merging_mode='hard')
``` 

Help would be really appreciated ! Thank you !

Cheers, 
Anthony

Discussion:
@alejoe91: Hi @Antho2422 

What's the error? :)@Antho2422: @alejoe91 oops I completly forgot this point ! 

Here it is :  
```python
def compute(self, traces, peaks):
370 --> waveforms = traces[peaks["sample_index"][:, None] + np.arange(-self.nbefore, self.nafter)]
    return waveforms

IndexError: index 33028 is out of bounds for axis 0 with size 33028
``` @alejoe91: I cannot reproduce the issue. Can you make sure you are using the main branch?

![Image](https://github.com/user-attachments/assets/83205b43-fc6a-4fed-b8e3-36673f837470)@Antho2422: @alejoe91 I just tested with the last update of spikeinterface on branch master an unfortunately I have the same issue....@alejoe91: This was running the code you provided using the main branch, so I'm not able to reproduce it@Antho2422: Weird, I just updated my environment and a collegue was able to reproduce the issue aswell... I will create a new environment from scratch and test again :)
@alejoe91: Which OS + Python version are you using? Could you reproduce on the data you shared or on another dataset?
@Antho2422: I'm on windows and using Python3.9. I only tried on the data I shared today but I remember having the same issue on other data (not 100% sure tho). 

Here is the same issue when creating a new venv in python3.9 then installing spikeinterface directly from the git. 

![Image](https://github.com/user-attachments/assets/1b500d2c-0c86-4fbd-877f-9b95582b2fdc)@alejoe91: Can you try to run it with 1 job? Might be Windows related, but it's strange!@Antho2422: I already ran it with one job only same issue. I think it is not Python related aswell because I just tried with Python3.11 and I have the same issue. Another collegue managed to pass the line that was causing problem (**On windows**) (and got another issue). So maybe this has something to do with a dependency ? I will investigate and compare environments.@Antho2422: @alejoe91 okay, I managed to pass the line that was causing this issue

Previously I was using the last version of the code and it was not working, I had to switch to spikeinterface 101.2 to make it work. I though you were using the last version of the code here and I was also using the last version (up to date with the git) https://github.com/SpikeInterface/spikeinterface/issues/3637#issuecomment-2607041926 but I may have missunderstood. 

I still have the same error as my collegue now which is this one : `TypeError: ComputeQualityMetrics._set_params() got an unexpected keyword argument 'metric_params'`? I don't know but this is maybe triggered by the fact that I am using 101.2 ? 

![Image](https://github.com/user-attachments/assets/d829e15e-875c-4f34-9937-7719b6524914)@alejoe91: Can you try to install from source and from main, not 0.101.2?@Antho2422: @alejoe91 I did, I was getting the error when using directly source code from main. Switching to 0.101.2 fixed it@zm711: I'm a little booked up for the next couple days. But if we can't sort this I can try to reproduce this on my Windows machine on Friday. With some of the updates to formats I've had some minor hiccups with curation and qc stuff but I haven't played with the `apply_curation` function.@Antho2422: I think the issue is not OS dependent. @alejoe91 are you sure you were using the main branch here https://github.com/SpikeInterface/spikeinterface/issues/3637#issuecomment-2607041926 ? I'm a bit intrigued by the warning (I don't have it on main)@alejoe91: Yep, reinstalled from source to be sure

![Image](https://github.com/user-attachments/assets/832d0a6f-4a9a-45db-9ef5-21e4d42a1459)

![Image](https://github.com/user-attachments/assets/44438cf2-be53-463c-add3-87c99a0ac710)

And added a print statement for you ;)

![Image](https://github.com/user-attachments/assets/92485acf-eb2b-4ee6-aa22-2e734b12cbd0)


Are you sure you are using the main branch??@Antho2422: haha okay I trust you @alejoe91 ! Thank you

I am on main I'm 100% sure, the code I'm using is here and as you can see I have the updates from the last commit (except the one you just made a few minutes ago) 

![Image](https://github.com/user-attachments/assets/524df1cc-dd09-43d7-9486-a66a62b261e2)

![Image](https://github.com/user-attachments/assets/294d1082-8595-41a5-a563-c4bc5cef6055)@Antho2422: The only reason I can imagine right now would be some package that as a different version in my env, here is the output of pip freeze 

[non_working_env.txt](https://github.com/user-attachments/files/18508775/non_working_env.txt)

To create my environement I installed spikeinterface by doing `pip install spikeinterface[full]` and then updated it to last main version by doing `pip install git+https://github.com/SpikeInterface/spikeinterface.git`. I'm using Python3.11 for this test but Python3.9 was not working aswell@alejoe91: Can you try from the anaconda prompt to do:
```
pip uninstall -y spikeinterface
pip install git+https://github.com/SpikeInterface/spikeinterface.git
```
?@Antho2422: Same problem...@alejoe91: let's see if we can reproduce it on windows!@samuelgarcia: I think, there are spikes too near the border and we are not safe enought after the merging in mode "hard".
But this is strange.@Antho2422: That is my guess as well, when chunking the recording, some peaks might be near the border of the chunk, when adding nafter here `traces[peaks["sample_index"][:, None] + np.arange(-self.nbefore, self.nafter)]` the index of the peaks goes over the samples of the chunk

But the question is why is it working in @alejoe91 setup ? @samuelgarcia: I think this depend on the numpy version installed that behave strictly or not.
we need to fix anyway@Antho2422: @alejoe91 what numpy version do you have at the moment ? I'm on `1.26.4`@alejoe91: same@chrishalcrow: Hello, I managed to reproduce the issue! I haven't solved it yet but I've found something worth looking at. If I run

```
raw_analyzer.compute_several_extensions({"spike_amplitudes":{},"spike_locations":{},"amplitude_scalings":{}})
```

it's fine. But the `all_nodes` list (a variable used in `run_node_pipeline` in `sortinganalyzer.py`) is in a different order than when the merging takes place. For the `compute_several_extensions`, all_nodes is equal to 

```
all_nodes = [<spikeinterface.core.node_pipeline.SpikeRetriever object at 0x1766bb0e0>, 
<spikeinterface.postprocessing.spike_amplitudes.SpikeAmplitudeNode object at 0x313befc20>, 
<spikeinterface.core.node_pipeline.SpikeRetriever object at 0x175b3fec0>, 
<spikeinterface.core.node_pipeline.ExtractDenseWaveforms object at 0x175b3e0c0>, 
<spikeinterface.sortingcomponents.peak_localization.LocalizeCenterOfMass object at 0x175b3daf0>, 
<spikeinterface.core.node_pipeline.SpikeRetriever object at 0x175b3fe30>, 
<spikeinterface.postprocessing.amplitude_scalings.AmplitudeScalingNode object at 0x175b3ea50>]
```

while during the merging, all_nodes is equal to:

```
all_nodes = [<spikeinterface.core.node_pipeline.SpikeRetriever object at 0x1766ef950>, 
<spikeinterface.postprocessing.amplitude_scalings.AmplitudeScalingNode object at 0x175cf2f00>, 
<spikeinterface.core.node_pipeline.SpikeRetriever object at 0x175cf32c0>, 
<spikeinterface.postprocessing.spike_amplitudes.SpikeAmplitudeNode object at 0x175cf2f30>, 
<spikeinterface.core.node_pipeline.SpikeRetriever object at 0x177eacbf0>, 
<spikeinterface.core.node_pipeline.ExtractDenseWaveforms object at 0x1760ea5a0>, 
<spikeinterface.sortingcomponents.peak_localization.LocalizeCenterOfMass object at 0x176025670>]
```
 
If I change the order:
```
raw_analyzer.compute_several_extensions({"amplitude_scalings":{},"spike_amplitudes":{},"spike_locations":{}})
```
I get the familiar error...

```
...
File [~/fromgit/spikeinterface/src/spikeinterface/core/node_pipeline.py:370](http://localhost:65144/lab/tree/~/fromgit/spikeinterface/src/spikeinterface/core/node_pipeline.py#line=369), in ExtractDenseWaveforms.compute(self, traces, peaks)
    369 def compute(self, traces, peaks):
--> 370     waveforms = traces[peaks["sample_index"][:, None] + np.arange(-self.nbefore, self.nafter)]
    371     return waveforms

IndexError: index 33028 is out of bounds for axis 0 with size 33028
```

So I suspect it's to do with the order of computing the extensions. @samuelgarcia is there a reason that `amplitude_scalings` should be computed after `spike_amplitudes`? Or maybe the nodes need to be reordered during `compute_several_extensions`?@Antho2422: Thanks @chrishalcrow ! It would help me a loooot if we get this solved asap :)

Maybe it's related with the fact that list ordering can be handled different from one OS to another ? (it could explain why it works on @alejoe91 computer) 

I don't know !@samuelgarcia: @chrishalcrow spike_ampitudes and amplitude_scaling should be independant and the order should not count.

But what you point is really important.
In short :  there is mechanim to make the first special (the first spike reteiver) in the nodepipeline.
My guess is that the others spikeretreiver behave differently and the border are not handle correctly. I will also try to have a look.

@chrishalcrow: Hey @Antho2422 . Maybe while we try and figure this out, you could only compute `amplitude_scalings` at the very end. I guess you don't need these to help figure out if you want to merge or not. So I suggest

Load your sorting analyzer
Delete the amplitude scalings extension
Do your merging
Compute amplitude scalings (if you really need them!)

And we'll let you know when this is fixed!@Antho2422: Thank you @samuelgarcia @chrishalcrow !

I will do this workaround for the moment :)

Please let me know when this is fixed !

Cheers@alejoe91: Ok guys I found a lead. Turns out that this only fails if `amplitude_scalings` is first in the list. I think this is because the first node is used throughout the pipeline to estimate some indices, and the `amplitude_scaling` spike retriever is trickier because it also retrieves spikes within the margin (which are used for proper fitting at the borders).

Hopefully will find qa fix soon!@alejoe91: Should be fixed here: #3642@Antho2422: Awesome ! Will test asap ! Thank you guys !

Solution patch:
diff --git a/examples/tutorials/curation/plot_1_automated_curation.py b/examples/tutorials/curation/plot_1_automated_curation.py
index e88b0973d..ae0135320 100644
--- a/examples/tutorials/curation/plot_1_automated_curation.py
+++ b/examples/tutorials/curation/plot_1_automated_curation.py
@@ -209,8 +209,8 @@ plt.legend(); plt.grid(True); plt.show()
 # trained on real data.
 #
 # For example, the following classifiers are trained on Neuropixels data from 11 mice recorded in
-# V1,SC and ALM: https://huggingface.co/AnoushkaJain3/noise_neural_classifier/ and
-# https://huggingface.co/AnoushkaJain3/sua_mua_classifier/ . One will classify units into
+# V1,SC and ALM: https://huggingface.co/SpikeInterface/UnitRefine_noise_neural_classifier/ and
+# https://huggingface.co/SpikeInterface/UnitRefine_sua_mua_classifier/. One will classify units into
 # `noise` or `not-noise` and the other will classify the `not-noise` units into single
 # unit activity (sua) units and multi-unit activity (mua) units.
 #
@@ -221,8 +221,8 @@ plt.legend(); plt.grid(True); plt.show()
 
 # Apply the noise/not-noise model
 noise_neuron_labels = sc.auto_label_units(
-    sorting_analyzer = sorting_analyzer,
-    repo_id = "AnoushkaJain3/noise_neural_classifier",
+    sorting_analyzer=sorting_analyzer,
+    repo_id="SpikeInterface/UnitRefine_noise_neural_classifier",
     trust_model=True,
 )
 
@@ -231,8 +231,8 @@ analyzer_neural = sorting_analyzer.remove_units(noise_units.index)
 
 # Apply the sua/mua model
 sua_mua_labels = sc.auto_label_units(
-    sorting_analyzer = analyzer_neural,
-    repo_id = "AnoushkaJain3/sua_mua_classifier",
+    sorting_analyzer=analyzer_neural,
+    repo_id="SpikeInterface/UnitRefine_sua_mua_classifier",
     trust_model=True,
 )
 
diff --git a/src/spikeinterface/core/globals.py b/src/spikeinterface/core/globals.py
index e9974adff..85ad0bfab 100644
--- a/src/spikeinterface/core/globals.py
+++ b/src/spikeinterface/core/globals.py
@@ -98,7 +98,7 @@ def is_set_global_dataset_folder() -> bool:
 
 ########################################
 _default_job_kwargs = dict(
-    pool_engine="thread", n_jobs=1, chunk_duration="1s", progress_bar=True, mp_context=None, max_threads_per_worker=1
+    pool_engine="process", n_jobs=1, chunk_duration="1s", progress_bar=True, mp_context=None, max_threads_per_worker=1
 )
 
 global global_job_kwargs
diff --git a/src/spikeinterface/core/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py
index 53c2445c7..9496e03ba 100644
--- a/src/spikeinterface/core/node_pipeline.py
+++ b/src/spikeinterface/core/node_pipeline.py
@@ -191,7 +191,7 @@ class SpikeRetriever(PeakSource):
         self._dtype = spike_peak_dtype
 
         self.include_spikes_in_margin = include_spikes_in_margin
-        if include_spikes_in_margin is not None:
+        if include_spikes_in_margin:
             self._dtype = spike_peak_dtype + [("in_margin", "bool")]
 
         self.peaks = sorting_to_peaks(sorting, extremum_channel_inds, self._dtype)
@@ -228,12 +228,6 @@ class SpikeRetriever(PeakSource):
         # get local peaks
         sl = self.segment_slices[segment_index]
         peaks_in_segment = self.peaks[sl]
-        # if self.include_spikes_in_margin:
-        #     i0, i1 = np.searchsorted(
-        #         peaks_in_segment["sample_index"], [start_frame - max_margin, end_frame + max_margin]
-        #     )
-        # else:
-        #     i0, i1 = np.searchsorted(peaks_in_segment["sample_index"], [start_frame, end_frame])
         i0, i1 = peak_slice
 
         local_peaks = peaks_in_segment[i0:i1]
@@ -435,21 +429,59 @@ class ExtractSparseWaveforms(WaveformsNode):
         return sparse_wfs
 
 
-def find_parent_of_type(list_of_parents, parent_type, unique=True):
+def find_parent_of_type(list_of_parents, parent_type):
+    """
+    Find a single parent of a given type(s) in a list of parents.
+    If multiple parents of the given type are found, the first parent is returned.
+
+    Parameters
+    ----------
+    list_of_parents : list of PipelineNode
+        List of parents to search through.
+    parent_type : type | tuple of types
+        The type of parent to search for.
+
+    Returns
+    -------
+    parent : PipelineNode or None
+        The parent of the given type. Returns None if no parent of the given type is found.
+    """
     if list_of_parents is None:
         return None
 
+    parents = find_parents_of_type(list_of_parents, parent_type)
+
+    if len(parents) > 0:
+        return parents[0]
+    else:
+        return None
+
+
+def find_parents_of_type(list_of_parents, parent_type):
+    """
+    Find all parents of a given type(s) in a list of parents.
+
+    Parameters
+    ----------
+    list_of_parents : list of PipelineNode
+        List of parents to search through.
+    parent_type : type | tuple of types
+        The type(s) of parents to search for.
+
+    Returns
+    -------
+    parents : list of PipelineNode
+        List of parents of the given type(s). Returns an empty list if no parents of the given type(s) are found.
+    """
+    if list_of_parents is None:
+        return []
+
     parents = []
     for parent in list_of_parents:
         if isinstance(parent, parent_type):
             parents.append(parent)
 
-    if unique and len(parents) == 1:
-        return parents[0]
-    elif not unique and len(parents) > 1:
-        return parents[0]
-    else:
-        return None
+    return parents
 
 
 def check_graph(nodes):
@@ -471,7 +503,7 @@ def check_graph(nodes):
             assert parent in nodes, f"Node {node} has parent {parent} that was not passed in nodes"
             assert (
                 nodes.index(parent) < i
-            ), f"Node are ordered incorrectly: {node} before {parent} in the pipeline definition."
+            ), f"Nodes are ordered incorrectly: {node} before {parent} in the pipeline definition."
 
     return nodes
 
@@ -607,12 +639,16 @@ def _compute_peak_pipeline_chunk(segment_index, start_frame, end_frame, worker_c
     skip_after_n_peaks_per_worker = worker_ctx["skip_after_n_peaks_per_worker"]
 
     recording_segment = recording._recording_segments[segment_index]
-    node0 = nodes[0]
-
-    if isinstance(node0, (SpikeRetriever, PeakRetriever)):
-        # in this case PeakSource could have no peaks and so no need to load traces just skip
-        peak_slice = i0, i1 = node0.get_peak_slice(segment_index, start_frame, end_frame, max_margin)
-        load_trace_and_compute = i0 < i1
+    retrievers = find_parents_of_type(nodes, (SpikeRetriever, PeakRetriever))
+    # get peak slices once for all retrievers
+    peak_slice_by_retriever = {}
+    for retriever in retrievers:
+        peak_slice = i0, i1 = retriever.get_peak_slice(segment_index, start_frame, end_frame, max_margin)
+        peak_slice_by_retriever[retriever] = peak_slice
+
+    if len(peak_slice_by_retriever) > 0:
+        # in this case the retrievers could have no peaks, so we test if any spikes are in the chunk
+        load_trace_and_compute = any(i0 < i1 for i0, i1 in peak_slice_by_retriever.values())
     else:
         # PeakDetector always need traces
         load_trace_and_compute = True
@@ -646,7 +682,8 @@ def _compute_peak_pipeline_chunk(segment_index, start_frame, end_frame, worker_c
                 node_output = node.compute(trace_detection, start_frame, end_frame, segment_index, max_margin)
                 # set sample index to local
                 node_output[0]["sample_index"] += extra_margin
-            elif isinstance(node, PeakSource):
+            elif isinstance(node, (PeakRetriever, SpikeRetriever)):
+                peak_slice = peak_slice_by_retriever[node]
                 node_output = node.compute(traces_chunk, start_frame, end_frame, segment_index, max_margin, peak_slice)
             else:
                 # TODO later when in master: change the signature of all nodes (or maybe not!)
diff --git a/src/spikeinterface/sortingcomponents/motion/motion_utils.py b/src/spikeinterface/sortingcomponents/motion/motion_utils.py
index 680d75f22..5c0264649 100644
--- a/src/spikeinterface/sortingcomponents/motion/motion_utils.py
+++ b/src/spikeinterface/sortingcomponents/motion/motion_utils.py
@@ -587,10 +587,14 @@ def ensure_time_bins(time_bin_centers_s=None, time_bin_edges_s=None):
     Going from centers to edges is done by taking midpoints and padding with the
     left and rightmost centers.
 
+    To handle multi segment, this function is working both:
+      * array/array input
+      * list[array]/list[array] input
+
     Parameters
     ----------
-    time_bin_centers_s : None or np.array
-    time_bin_edges_s : None or np.array
+    time_bin_centers_s : None or np.array or list[np.array]
+    time_bin_edges_s : None or np.array or list[np.array]
 
     Returns
     -------
@@ -600,17 +604,34 @@ def ensure_time_bins(time_bin_centers_s=None, time_bin_edges_s=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 isinstance(time_bin_edges_s, list):
+            # multi segment cas
+            time_bin_centers_s = []
+            for be in time_bin_edges_s:
+                bc, _ = ensure_time_bins(time_bin_centers_s=None, time_bin_edges_s=be)
+                time_bin_centers_s.append(bc)
+        else:
+            # simple segment
+            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])
+        if isinstance(time_bin_centers_s, list):
+            # multi segment cas
+            time_bin_edges_s = []
+            for bc in time_bin_centers_s:
+                _, be = ensure_time_bins(time_bin_centers_s=bc, time_bin_edges_s=None)
+                time_bin_edges_s.append(be)
+        else:
+            # simple segment
+            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]


Test patch:
diff --git a/src/spikeinterface/core/tests/test_globals.py b/src/spikeinterface/core/tests/test_globals.py
index cc8ff1007..ba31f6f08 100644
--- a/src/spikeinterface/core/tests/test_globals.py
+++ b/src/spikeinterface/core/tests/test_globals.py
@@ -37,7 +37,7 @@ def test_global_tmp_folder(create_cache_folder):
 
 def test_global_job_kwargs():
     job_kwargs = dict(
-        pool_engine="thread",
+        pool_engine="process",
         n_jobs=4,
         chunk_duration="1s",
         progress_bar=True,
@@ -47,7 +47,7 @@ def test_global_job_kwargs():
     global_job_kwargs = get_global_job_kwargs()
 
     assert global_job_kwargs == dict(
-        pool_engine="thread",
+        pool_engine="process",
         n_jobs=1,
         chunk_duration="1s",
         progress_bar=True,
@@ -62,7 +62,7 @@ def test_global_job_kwargs():
     set_global_job_kwargs(**partial_job_kwargs)
     global_job_kwargs = get_global_job_kwargs()
     assert global_job_kwargs == dict(
-        pool_engine="thread",
+        pool_engine="process",
         n_jobs=2,
         chunk_duration="1s",
         progress_bar=True,
diff --git a/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py b/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
index e4ba87032..807b8e6c9 100644
--- a/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
+++ b/src/spikeinterface/sortingcomponents/motion/tests/test_motion_interpolation.py
@@ -10,18 +10,26 @@ from spikeinterface.sortingcomponents.motion.motion_interpolation import (
     interpolate_motion_on_traces,
 )
 from spikeinterface.sortingcomponents.tests.common import make_dataset
-
+from spikeinterface.core import generate_ground_truth_recording
 
 def make_fake_motion(rec):
     # make a fake motion object
-    duration = rec.get_total_duration()
+
     locs = rec.get_channel_locations()
-    temporal_bins = np.arange(0.5, duration - 0.49, 0.5)
     spatial_bins = np.arange(locs[:, 1].min(), locs[:, 1].max(), 100)
-    displacement = np.zeros((temporal_bins.size, spatial_bins.size))
-    displacement[:, :] = np.linspace(-30, 30, temporal_bins.size)[:, None]
 
-    motion = Motion([displacement], [temporal_bins], spatial_bins, direction="y")
+    displacement = []
+    temporal_bins = []
+    for segment_index in range(rec.get_num_segments()):
+        duration = rec.get_duration(segment_index=segment_index)
+        seg_time_bins = np.arange(0.5, duration - 0.49, 0.5)
+        seg_disp = np.zeros((seg_time_bins.size, spatial_bins.size))
+        seg_disp[:, :] = np.linspace(-30, 30, seg_time_bins.size)[:, None]
+        
+        temporal_bins.append(seg_time_bins)
+        displacement.append(seg_disp)
+
+    motion = Motion(displacement, temporal_bins, spatial_bins, direction="y")
 
     return motion
 
@@ -176,7 +184,27 @@ def test_cross_band_interpolation():
 
 
 def test_InterpolateMotionRecording():
-    rec, sorting = make_dataset()
+    # rec, sorting = make_dataset()
+
+    # 2 segments
+    rec, sorting = generate_ground_truth_recording(
+        durations=[30.0],
+        sampling_frequency=30000.0,
+        num_channels=32,
+        num_units=10,
+        generate_probe_kwargs=dict(
+            num_columns=2,
+            xpitch=20,
+            ypitch=20,
+            contact_shapes="circle",
+            contact_shape_params={"radius": 6},
+        ),
+        generate_sorting_kwargs=dict(firing_rates=6.0, refractory_period_ms=4.0),
+        noise_kwargs=dict(noise_levels=5.0, strategy="on_the_fly"),
+        seed=2205,
+    )
+
+
     motion = make_fake_motion(rec)
 
     rec2 = InterpolateMotionRecording(rec, motion, border_mode="force_extrapolate")
@@ -187,15 +215,19 @@ def test_InterpolateMotionRecording():
 
     rec2 = InterpolateMotionRecording(rec, motion, border_mode="remove_channels")
     assert rec2.channel_ids.size == 24
-    for ch_id in (0, 1, 14, 15, 16, 17, 30, 31):
+    for ch_id in ("0", "1", "14", "15", "16", "17", "30", "31"):
         assert ch_id not in rec2.channel_ids
 
     traces = rec2.get_traces(segment_index=0, start_frame=0, end_frame=30000)
     assert traces.shape == (30000, 24)
 
-    traces = rec2.get_traces(segment_index=0, start_frame=0, end_frame=30000, channel_ids=[3, 4])
+    traces = rec2.get_traces(segment_index=0, start_frame=0, end_frame=30000, channel_ids=["3", "4"])
     assert traces.shape == (30000, 2)
 
+    # test dump.load when multi segments
+    rec2.dump("rec_motion_interp.pickle")
+    rec3 = sc.load("rec_motion_interp.pickle")
+
     # import matplotlib.pyplot as plt
     # import spikeinterface.widgets as sw
     # fig, ax = plt.subplots()
@@ -207,7 +239,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()
+    test_InterpolateMotionRecording()
+    # test_cross_band_interpolation()