Move interpolators tests to sortingcomponents

Author: alejoe91

pyproject.toml

@@ -123,8 +123,6 @@ test_core = [
     "pytest",
     "pytest-dependency",
     "psutil",
-    "scipy", # for motion interpolators
-    "s3fs", # for streaming tests
 
     # for github test : probeinterface and neo from master
     # for release we need pypi, so this need to be commented

src/spikeinterface/core/tests/test_loading.py

@@ -176,13 +176,15 @@ def test_load_motion(tmp_path, generate_motion_object):
     assert motion == motion_loaded
 
 
+@pytest.mark.streaming_extractors
 @pytest.mark.skipif(not HAVE_S3, reason="s3fs not installed")
 def test_remote_recording():
     s3_path = "s3://spikeinterface-sorting-analyzer-test/recording_for_analyzer_short.zarr/"
     rec = load(s3_path)
     assert isinstance(rec, ZarrRecordingExtractor)
 
 
+@pytest.mark.streaming_extractors
 @pytest.mark.skipif(not HAVE_S3, reason="s3fs not installed")
 def test_remote_analyzer():
     s3_path = "s3://spikeinterface-sorting-analyzer-test/analyzer_remote_test.zarr/"

src/spikeinterface/core/tests/test_motion_core.py

@@ -29,8 +29,8 @@ def make_fake_motion():
     return motion
 
 
-def test_Motion(tmp_path):
-
+def test_motion_object(tmp_path):
+    """Basic tests for Motion object representation, saving and loading."""
     temporal_bins_s = np.arange(0.0, 10.0, 1.0)
     spatial_bins_um = np.array([100.0, 200.0])
 
@@ -43,37 +43,13 @@ def test_Motion(tmp_path):
     # serialize with pickle before interpolation fit
     motion2 = pickle.loads(pickle.dumps(motion))
     assert motion2.interpolators is None
-    # serialize with pickle after interpolation fit
-    motion2.make_interpolators()
-    assert motion2.interpolators is not None
-    motion2 = pickle.loads(pickle.dumps(motion2))
-    assert motion2.interpolators is not None
-
-    # to/from dict
-    motion2 = Motion.from_dict(motion.to_dict())
-    assert motion == motion2
-    assert motion2.interpolators is None
-
-    # do interpolate
-    displacement = motion.get_displacement_at_time_and_depth([2, 4.4, 11], [120.0, 80.0, 150.0])
-    # print(displacement)
-    assert displacement.shape[0] == 3
-    # check clip
-    assert displacement[2] == 20.0
-
-    # interpolate grid
-    displacement = motion.get_displacement_at_time_and_depth([2, 4.4, 11, 15, 19], [150.0, 80.0], grid=True)
-    assert displacement.shape == (2, 5)
-    assert displacement[0, 2] == 20.0
 
     # save/load to folder
     folder = tmp_path / "motion_saved"
-    if folder.exists():
-        shutil.rmtree(folder)
     motion.save(folder)
     motion2 = Motion.load(folder)
     assert motion == motion2
 
 
 if __name__ == "__main__":
-    test_Motion()
+    test_motion_object()

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

@@ -1,17 +1,19 @@
 import warnings
 
 import numpy as np
-import spikeinterface.core as sc
+import pickle
 
+from spikeinterface import NumpyRecording, load
+from spikeinterface.core import generate_ground_truth_recording
 from spikeinterface.core.motion import Motion
+from spikeinterface.core.testing import check_recordings_equal
 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
-from spikeinterface.core import generate_ground_truth_recording
 
 
 def make_fake_motion(rec):
@@ -36,12 +38,49 @@ def make_fake_motion(rec):
     return motion
 
 
+def test_motion_object_interpolators(tmp_path):
+    rec, sorting = make_dataset()
+    motion = make_fake_motion(rec)
+
+    # serialize with pickle after interpolation fit
+    motion.make_interpolators()
+    assert motion.interpolators is not None
+    motion2 = pickle.loads(pickle.dumps(motion))
+    assert motion2.interpolators is not None
+
+    # to/from dict
+    motion2 = Motion.from_dict(motion.to_dict())
+    assert motion == motion2
+    # serialization to/from dict loses interpolators
+    assert motion2.interpolators is None
+    motion2.make_interpolators()
+
+    # save/load to folder
+    folder = tmp_path / "motion_saved"
+    motion.save(folder)
+    motion3 = Motion.load(folder)
+    # save/load to folder loses interpolators
+    assert motion3.interpolators is None
+
+    # do interpolate
+    displacement = motion.get_displacement_at_time_and_depth([2, 4.4, 11], [120.0, 80.0, 150.0])
+    displacement2 = motion2.get_displacement_at_time_and_depth([2, 4.4, 11], [120.0, 80.0, 150.0])
+    displacement3 = motion3.get_displacement_at_time_and_depth([2, 4.4, 11], [120.0, 80.0, 150.0])
+    # print(displacement)
+    assert displacement.shape[0] == 3
+    # check interpolators reload
+    np.testing.assert_array_equal(displacement, displacement2)
+    np.testing.assert_array_equal(displacement, displacement3)
+
+    # interpolate grid
+    displacement = motion.get_displacement_at_time_and_depth([2, 4.4, 11, 15, 19], [150.0, 80.0], grid=True)
+    assert displacement.shape == (2, 5)
+
+
 def test_correct_motion_on_peaks():
     rec, sorting = make_dataset()
     peaks = sorting.to_spike_vector()
-    print(peaks.dtype)
     motion = make_fake_motion(rec)
-    # print(motion)
 
     # fake locations
     peak_locations = np.zeros((peaks.size), dtype=[("x", "float32"), ("y", "float")])
@@ -98,7 +137,7 @@ def test_interpolation_simple():
     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 = NumpyRecording(traces, sampling_frequency=1)
     rec.set_dummy_probe_from_locations(np.c_[np.zeros(num_chans_drifted), np.arange(num_chans_drifted)])
 
     true_motion = Motion(np.arange(n_samples)[:, None], 0.5 + np.arange(n_samples), np.zeros(1))
@@ -157,14 +196,14 @@ def test_cross_band_interpolation():
     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 = NumpyRecording(traces_lfp, sampling_frequency=fs_lfp)
     rec_lfp.set_dummy_probe_from_locations(geom)
 
     # same for AP
     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)
+    rec_ap = NumpyRecording(traces_ap, sampling_frequency=fs_ap)
     rec_ap.set_dummy_probe_from_locations(geom)
 
     # set times for both, and silence the warning
@@ -227,7 +266,9 @@ def test_InterpolateMotionRecording():
 
     # test dump.load when multi segments
     rec2.dump("rec_motion_interp.pickle")
-    rec3 = sc.load("rec_motion_interp.pickle")
+    rec3 = load("rec_motion_interp.pickle")
+
+    check_recordings_equal(rec2, rec3)
 
     # import matplotlib.pyplot as plt
     # import spikeinterface.widgets as sw