Repository: catalystneuro/roiextractors

Issue:
add stub option
A stub option should limit the conversion to the first ~100 frames so that conversions can be tested quickly without reading/writing all of the data.

Comments:

Patch:
diff --git a/src/roiextractors/imagingextractor.py b/src/roiextractors/imagingextractor.py
index e65a618..059d9d4 100644
--- a/src/roiextractors/imagingextractor.py
+++ b/src/roiextractors/imagingextractor.py
@@ -1,5 +1,6 @@
+"""Base class definitions for all ImagingExtractors."""
 from abc import ABC, abstractmethod
-from typing import Union
+from typing import Union, Optional, Tuple
 import numpy as np
 from copy import deepcopy
 
@@ -9,17 +10,13 @@ from .extraction_tools import (
     ArrayType,
     PathType,
     DtypeType,
-    IntType,
     FloatType,
     check_get_videos_args,
 )
-from typing import Tuple
 
 
 class ImagingExtractor(ABC, BaseExtractor):
-    """An abstract class that contains all the meta-data and input data from
-    the imaging data.
-    """
+    """Abstract class that contains all the meta-data and input data from the imaging data."""
 
     def __init__(self) -> None:
         BaseExtractor.__init__(self)
@@ -27,7 +24,7 @@ class ImagingExtractor(ABC, BaseExtractor):
         self._memmapped = False
 
     @abstractmethod
-    def get_frames(self, frame_idxs: ArrayType, channel: int = 0) -> np.ndarray:
+    def get_frames(self, frame_idxs: ArrayType, channel: Optional[int] = 0) -> np.ndarray:
         pass
 
     @abstractmethod
@@ -117,8 +114,8 @@ class ImagingExtractor(ABC, BaseExtractor):
         times: array-like
             The times in seconds for each frame
         """
-        assert len(times) == self.get_num_frames(), "'times' should have the same length of the number of frames"
-        self._times = times.astype("float64")
+        assert len(times) == self.get_num_frames(), "'times' should have the same length of the number of frames!"
+        self._times = np.array(times).astype("float64")
 
     def copy_times(self, extractor: BaseExtractor) -> None:
         """This function copies times from another extractor.
@@ -131,6 +128,10 @@ class ImagingExtractor(ABC, BaseExtractor):
         if extractor._times is not None:
             self.set_times(deepcopy(extractor._times))
 
+    def frame_slice(self, start_frame, end_frame):
+        """Return a new ImagingExtractor ranging from the start_frame to the end_frame."""
+        return FrameSliceImagingExtractor(parent_imaging=self, start_frame=start_frame, end_frame=end_frame)
+
     @staticmethod
     def write_imaging(imaging, save_path: PathType, overwrite: bool = False):
         """
@@ -147,3 +148,77 @@ class ImagingExtractor(ABC, BaseExtractor):
             If True and save_path is existing, it is overwritten
         """
         raise NotImplementedError
+
+
+class FrameSliceImagingExtractor(ImagingExtractor):
+    """
+    Class to get a lazy frame slice.
+
+    Do not use this class directly but use `.frame_slice(...)` on an ImagingExtractor object.
+    """
+
+    extractor_name = "FrameSliceImagingExtractor"
+    installed = True
+    is_writable = True
+    installation_mesg = ""
+
+    def __init__(
+        self, parent_imaging: ImagingExtractor, start_frame: Optional[int] = None, end_frame: Optional[int] = None
+    ):
+        """
+        Initialize an ImagingExtractor whose frames subset the parent.
+
+        Subset is exclusive on the right bound, that is, the indexes of this ImagingExtractor range over
+        [0, ..., end_frame-start_frame-1], which is used to resolve the index mapping in `get_frames(frame_idxs=[...])`.
+
+        Parameters
+        ----------
+        parent_imaging : ImagingExtractor
+            The ImagingExtractor object to sebset the frames of.
+        start_frame : int, optional
+            The left bound of the frames to subset.
+            The default is the start frame of the parent.
+        end_frame : int, optional
+            The right bound of the frames, exlcusively, to subset.
+            The default is end frame of the parent.
+
+        """
+        self._parent_imaging = parent_imaging
+        self._start_frame = start_frame
+        self._end_frame = end_frame
+        self._num_frames = self._end_frame - self._start_frame
+
+        parent_size = self._parent_imaging.get_num_frames()
+        if start_frame is None:
+            start_frame = 0
+        else:
+            assert 0 <= start_frame < parent_size
+        if end_frame is None:
+            end_frame = parent_size
+        else:
+            assert 0 < end_frame <= parent_size
+        assert end_frame > start_frame, "'start_frame' must be smaller than 'end_frame'!"
+
+        super().__init__()
+        if getattr(self._parent_imaging, "_times") is not None:
+            self._times = self._parent_imaging._times[start_frame:end_frame]
+
+    def get_frames(self, frame_idxs: ArrayType, channel: Optional[int] = 0) -> np.ndarray:
+        assert max(frame_idxs) < self._num_frames, "'frame_idxs' range beyond number of available frames!"
+        mapped_frame_idxs = np.array(frame_idxs) + self._start_frame
+        return self._parent_imaging.get_frames(frame_idxs=mapped_frame_idxs, channel=channel)
+
+    def get_image_size(self) -> Tuple[int, int]:
+        return tuple(self._parent_imaging.get_image_size())
+
+    def get_num_frames(self) -> int:
+        return self._num_frames
+
+    def get_sampling_frequency(self) -> float:
+        return self._parent_imaging.get_sampling_frequency()
+
+    def get_channel_names(self) -> list:
+        return self._parent_imaging.get_channel_names()
+
+    def get_num_channels(self) -> int:
+        return self._parent_imaging.get_num_channels()


Test patch:
diff --git a/tests/test_internals/test_frame_slice_imaging.py b/tests/test_internals/test_frame_slice_imaging.py
new file mode 100644
index 0000000..de797d3
--- /dev/null
+++ b/tests/test_internals/test_frame_slice_imaging.py
@@ -0,0 +1,64 @@
+import unittest
+
+import numpy as np
+from hdmf.testing import TestCase
+from numpy.testing import assert_array_equal
+
+from roiextractors.testing import generate_dummy_imaging_extractor
+
+
+def test_frame_slicing_imaging_times():
+    num_frames = 10
+    timestamp_shift = 7.1
+    times = np.array(range(num_frames)) + timestamp_shift
+    start_frame, end_frame = 2, 7
+
+    toy_imaging_example = generate_dummy_imaging_extractor(num_frames=num_frames, rows=5, columns=4, num_channels=1)
+    toy_imaging_example.set_times(times=times)
+
+    frame_sliced_imaging = toy_imaging_example.frame_slice(start_frame=start_frame, end_frame=end_frame)
+    assert_array_equal(
+        frame_sliced_imaging.frame_to_time(
+            frames=np.array([idx for idx in range(frame_sliced_imaging.get_num_frames())])
+        ),
+        times[start_frame:end_frame],
+    )
+
+
+class TestFrameSliceImaging(TestCase):
+    @classmethod
+    def setUpClass(cls):
+        """Use a toy example of ten frames of a 5 x 4 grayscale image."""
+        cls.toy_imaging_example = generate_dummy_imaging_extractor(num_frames=10, rows=5, columns=4, num_channels=1)
+        cls.frame_sliced_imaging = cls.toy_imaging_example.frame_slice(start_frame=2, end_frame=7)
+
+    def test_get_image_size(self):
+        assert self.frame_sliced_imaging.get_image_size() == (5, 4)
+
+    def test_get_num_frames(self):
+        assert self.frame_sliced_imaging.get_num_frames() == 5
+
+    def test_get_sampling_frequency(self):
+        assert self.frame_sliced_imaging.get_sampling_frequency() == 30.0
+
+    def test_get_channel_names(self):
+        assert self.frame_sliced_imaging.get_channel_names() == ["channel_num_0"]
+
+    def test_get_num_channels(self):
+        assert self.frame_sliced_imaging.get_num_channels() == 1
+
+    def test_get_frames_assertion(self):
+        with self.assertRaisesWith(
+            exc_type=AssertionError, exc_msg="'frame_idxs' range beyond number of available frames!"
+        ):
+            self.frame_sliced_imaging.get_frames(frame_idxs=[6])
+
+    def test_get_frames(self):
+        assert_array_equal(
+            self.frame_sliced_imaging.get_frames(frame_idxs=[2, 4]),
+            self.toy_imaging_example.get_frames(frame_idxs=[4, 6]),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()