Shift times using writable copies of OOM arrays

Recordings loaded from disk may have time vectors
that are out-of-memory (OOM) arrays: read-only
mem-maps or zarr.Array. These can't be shifted in
place, and need to be materialized as as numpy
arrays before shifting.

Author: grahamfindlay

src/spikeinterface/core/tests/test_time_handling.py

@@ -375,6 +375,33 @@ def test_save_and_load_time_shift(self, request, fixture_name, tmp_path):
                 times_recording.get_times(segment_index=idx), loaded_recording.get_times(segment_index=idx)
             )
 
+    @pytest.mark.parametrize("save_format", ["binary", "zarr"])
+    def test_shift_times_after_load(self, request, save_format, tmp_path):
+        """
+        Shift times on a recording loaded from disk as a read-only np.memmap
+        (binary folder) and a lazy zarr.Array (zarr). Neither supports an in-place
+        `+=`, so `shift_times` must shift a writable copy. 
+        """
+        _, times_recording, all_times = self._get_fixture_data(request, "time_vector_recording")
+
+        folder = tmp_path / "rec"
+        times_recording.save(format=save_format, folder=folder)
+        load_path = folder.with_suffix(".zarr") if save_format == "zarr" else folder
+        loaded = si.load(load_path)
+
+        # Confirm we are actually exercising a non-writeable / non-ndarray path.
+        for idx in range(loaded.get_num_segments()):
+            tv = loaded.segments[idx].time_vector
+            assert not (isinstance(tv, np.ndarray) and tv.flags.writeable)
+
+        shift = 123.456
+        loaded.shift_times(shift)  
+
+        for idx in range(loaded.get_num_segments()):
+            assert np.allclose(
+                loaded.get_times(segment_index=idx), all_times[idx] + shift, rtol=0, atol=1e-8
+            )
+
     def _store_all_times(self, recording):
         """
         Convenience function to store original times of all segments to a dict.

src/spikeinterface/core/time_series.py

@@ -273,13 +273,12 @@ def shift_times(self, shift: int | float, segment_index: int | None = None) -> N
             rs = self.segments[segment_index]
 
             if self.has_time_vector(segment_index=segment_index):
-                if rs.time_vector.flags.writeable:
-                    # If the time_vector is writeable, shift in-place to avoid a copy.
-                    rs.time_vector += shift
+                if isinstance(rs.time_vector, np.ndarray) and rs.time_vector.flags.writeable:
+                    # If this is an in-memory numpy array
+                    rs.time_vector += shift  # in-place, no copy
                 else:
-                    # If the time_vector is a memmap from `np.load(..., mmap_mode='r')`,
-                    # in-place modification would error, so we shift a writable copy.
-                    rs.time_vector = rs.time_vector + shift
+                    # If this is a read-only memmap or zarr.Array
+                    rs.time_vector = np.asarray(rs.time_vector) + shift
             else:
                 new_start_time = 0 + shift if rs.t_start is None else rs.t_start + shift
                 rs.t_start = new_start_time