add tests and fixes

Author: alejoe91

src/spikeinterface/core/basesorting.py

@@ -350,6 +350,7 @@ def get_unit_spike_trains(
             )
 
         segment_index = self._check_segment_index(segment_index)
+        segment = self.segments[segment_index]
         if use_cache:
             # TODO: speed things up
             ordered_spike_vector, slices = self.to_reordered_spike_vector(
@@ -372,8 +373,9 @@ def get_unit_spike_trains(
                 spike_trains[unit_id] = spike_frames
         else:
             spike_trains = segment.get_unit_spike_trains(
-                unit_ids=unit_ids, start_frame=start_frame, end_frame=end_frame, return_times=return_times
+                unit_ids=unit_ids, start_frame=start_frame, end_frame=end_frame
             )
+        return spike_trains
 
     def get_unit_spike_trains_in_seconds(
         self,
@@ -453,10 +455,10 @@ def get_unit_spike_trains_in_seconds(
             use_cache=use_cache,
         )
         for unit_id in unit_ids:
-            spike_frames = spike_frames[unit_id]
+            spike_frames_unit = spike_frames[unit_id]
             t_start = segment._t_start if segment._t_start is not None else 0
-            spike_times[unit_id] = spike_frames / self.get_sampling_frequency()
-        return t_start + spike_times
+            spike_times[unit_id] = spike_frames_unit / self.get_sampling_frequency() + t_start
+        return spike_times
 
     def register_recording(self, recording, check_spike_frames: bool = True):
         """

src/spikeinterface/core/tests/test_basesorting.py

@@ -310,6 +310,37 @@ def test_select_periods():
     np.testing.assert_array_equal(sliced_sorting.to_spike_vector(), sliced_sorting_array.to_spike_vector())
 
 
+@pytest.mark.parametrize("use_cache", [False, True])
+def test_get_unit_spike_trains(use_cache):
+    sampling_frequency = 10_000.0
+    duration = 1.0
+    num_samples = int(sampling_frequency * duration)
+    num_units = 10
+    sorting = generate_sorting(durations=[duration], sampling_frequency=sampling_frequency, num_units=num_units)
+
+    all_spike_trains = sorting.get_unit_spike_trains(unit_ids=sorting.unit_ids, use_cache=use_cache)
+    assert isinstance(all_spike_trains, dict)
+    assert set(all_spike_trains.keys()) == set(sorting.unit_ids)
+    for unit_id in sorting.unit_ids:
+        spiketrain = sorting.get_unit_spike_train(segment_index=0, unit_id=unit_id, use_cache=use_cache)
+        assert np.array_equal(all_spike_trains[unit_id], spiketrain)
+
+    # test with times
+    spike_trains_times = sorting.get_unit_spike_trains_in_seconds(
+        unit_ids=sorting.unit_ids, return_times=True, use_cache=use_cache
+    )
+    assert isinstance(spike_trains_times, dict)
+    assert set(spike_trains_times.keys()) == set(sorting.unit_ids)
+    for unit_id in sorting.unit_ids:
+        spiketrain = sorting.get_unit_spike_train(
+            segment_index=0, unit_id=unit_id, use_cache=use_cache, return_times=True
+        )
+        spiketrain_times = sorting.get_unit_spike_train_in_seconds(
+            segment_index=0, unit_id=unit_id, use_cache=use_cache
+        )
+        assert np.allclose(spiketrain_times, spiketrain)
+
+
 if __name__ == "__main__":
     import tempfile