Repository: catalystneuro/roiextractors

Issue:
Add MultiImagingExtractor
We have a MultiSegmentationExtractor, but no MultiImagingExtractor. @weiglszonja needs this for the Fee lab project in order to splice together a series of separate ImagingExtractors into one continuous ImagingExtractor to be written to the same TwoPhotonSeries

Comments:

Patch:
diff --git a/src/roiextractors/multiimagingextractor.py b/src/roiextractors/multiimagingextractor.py
new file mode 100644
index 0000000..8dbc144
--- /dev/null
+++ b/src/roiextractors/multiimagingextractor.py
@@ -0,0 +1,112 @@
+from collections import defaultdict
+from typing import Tuple, List, Iterable, Optional
+
+import numpy as np
+
+from .extraction_tools import ArrayType, NumpyArray, check_get_frames_args
+from .imagingextractor import ImagingExtractor
+
+
+class MultiImagingExtractor(ImagingExtractor):
+    """
+    This class is used to combine multiple ImagingExtractor objects by frames.
+    """
+
+    extractor_name = "MultiImagingExtractor"
+    installed = True
+    installation_mesg = ""
+
+    def __init__(self, imaging_extractors: List[ImagingExtractor]):
+        """
+        Parameters
+        ----------
+        imaging_extractors: list of ImagingExtractor
+            list of imaging extractor objects
+        """
+        super().__init__()
+        assert isinstance(imaging_extractors, list), "Enter a list of ImagingExtractor objects as argument"
+        assert all(isinstance(IX, ImagingExtractor) for IX in imaging_extractors)
+        self._imaging_extractors = imaging_extractors
+
+        # Checks that properties are consistent between extractors
+        self._check_consistency_between_imaging_extractors()
+
+        self._start_frames, self._end_frames = [], []
+        num_frames = 0
+        for imaging_extractor in self._imaging_extractors:
+            self._start_frames.append(num_frames)
+            num_frames = num_frames + imaging_extractor.get_num_frames()
+            self._end_frames.append(num_frames)
+        self._num_frames = num_frames
+
+        if any((getattr(imaging_extractor, "_times") is not None for imaging_extractor in self._imaging_extractors)):
+            times = self._get_times()
+            self.set_times(times=times)
+
+    def _check_consistency_between_imaging_extractors(self):
+        properties_to_check = dict(
+            get_sampling_frequency="The sampling frequency",
+            get_image_size="The size of a frame",
+            get_num_channels="The number of channels",
+            get_channel_names="The name of the channels",
+        )
+        for method, property_message in properties_to_check.items():
+            values = [getattr(extractor, method, None)() for extractor in self._imaging_extractors]
+            unique_values = set(tuple(v) if isinstance(v, Iterable) else v for v in values)
+            assert (
+                len(unique_values) == 1
+            ), f"{property_message} is not consistent over the files (found {unique_values})."
+
+    def _get_times(self):
+        frame_indices = np.array([*range(self._start_frames[0], self._end_frames[-1])])
+        times = self.frame_to_time(frames=frame_indices)
+
+        for extractor_index, extractor in enumerate(self._imaging_extractors):
+            if getattr(extractor, "_times") is not None:
+                to_replace = [*range(self._start_frames[extractor_index], self._end_frames[extractor_index])]
+                times[to_replace] = extractor._times
+
+        return times
+
+    @check_get_frames_args
+    def get_frames(self, frame_idxs: ArrayType, channel: Optional[int] = 0) -> NumpyArray:
+        extractor_indices = np.searchsorted(self._end_frames, frame_idxs, side="right")
+        relative_frame_indices = frame_idxs - np.array(self._start_frames)[extractor_indices]
+        # Match frame_idxs to imaging extractors
+        extractors_dict = defaultdict(list)
+        for extractor_index, frame_index in zip(extractor_indices, relative_frame_indices):
+            extractors_dict[extractor_index].append(frame_index)
+
+        frames_to_concatenate = []
+        # Extract frames for each extractor and concatenate
+        for extractor_index, frame_indices in extractors_dict.items():
+            frames_for_each_extractor = self._get_frames_from_an_imaging_extractor(
+                extractor_index=extractor_index,
+                frame_idxs=frame_indices,
+            )
+            if len(frame_indices) == 1:
+                frames_for_each_extractor = frames_for_each_extractor[np.newaxis, ...]
+            frames_to_concatenate.append(frames_for_each_extractor)
+
+        frames = np.concatenate(frames_to_concatenate, axis=0).squeeze()
+        return frames
+
+    def _get_frames_from_an_imaging_extractor(self, extractor_index: int, frame_idxs: ArrayType) -> NumpyArray:
+        imaging_extractor = self._imaging_extractors[extractor_index]
+        frames = imaging_extractor.get_frames(frame_idxs=frame_idxs)
+        return frames
+
+    def get_image_size(self) -> Tuple:
+        return self._imaging_extractors[0].get_image_size()
+
+    def get_num_frames(self) -> int:
+        return self._num_frames
+
+    def get_sampling_frequency(self) -> float:
+        return self._imaging_extractors[0].get_sampling_frequency()
+
+    def get_channel_names(self) -> list:
+        return self._imaging_extractors[0].get_channel_names()
+
+    def get_num_channels(self) -> int:
+        return self._imaging_extractors[0].get_num_channels()


Test patch:
diff --git a/tests/test_internals/test_multiimagingextractor.py b/tests/test_internals/test_multiimagingextractor.py
new file mode 100644
index 0000000..e1d4dbc
--- /dev/null
+++ b/tests/test_internals/test_multiimagingextractor.py
@@ -0,0 +1,150 @@
+import unittest
+
+import numpy as np
+from hdmf.testing import TestCase
+from numpy.testing import assert_array_equal
+from parameterized import parameterized, param
+
+from roiextractors.multiimagingextractor import MultiImagingExtractor
+from roiextractors.testing import generate_dummy_imaging_extractor
+
+
+class TestMultiImagingExtractor(TestCase):
+    extractors = None
+
+    @classmethod
+    def setUpClass(cls):
+        cls.extractors = [
+            generate_dummy_imaging_extractor(num_frames=10, num_rows=3, num_columns=4, sampling_frequency=20.0)
+            for _ in range(3)
+        ]
+        cls.multi_imaging_extractor = MultiImagingExtractor(imaging_extractors=cls.extractors)
+
+    def test_get_image_size(self):
+        assert self.multi_imaging_extractor.get_image_size() == self.extractors[0].get_image_size()
+
+    def test_get_num_frames(self):
+        assert self.multi_imaging_extractor.get_num_frames() == 30
+
+    def test_get_sampling_frequency(self):
+        assert self.multi_imaging_extractor.get_sampling_frequency() == 20.0
+
+    def test_get_channel_names(self):
+        assert self.multi_imaging_extractor.get_channel_names() == ["channel_num_0"]
+
+    def test_get_num_channels(self):
+        assert self.multi_imaging_extractor.get_num_channels() == 1
+
+    def test_get_frames_assertion(self):
+        with self.assertRaisesWith(exc_type=AssertionError, exc_msg="'frame_idxs' exceed number of frames"):
+            self.multi_imaging_extractor.get_frames(frame_idxs=[31])
+
+    def test_get_non_consecutive_frames(self):
+        test_frames = self.multi_imaging_extractor.get_frames(frame_idxs=[8, 10, 12, 15, 20, 29])
+        expected_frames = np.concatenate(
+            (
+                self.extractors[0].get_frames(frame_idxs=[8])[np.newaxis, ...],
+                self.extractors[1].get_frames(frame_idxs=[0, 2, 5]),
+                self.extractors[2].get_frames(frame_idxs=[0, 9]),
+            ),
+            axis=0,
+        )
+        assert_array_equal(test_frames, expected_frames)
+
+    def test_get_consecutive_frames(self):
+        test_frames = self.multi_imaging_extractor.get_frames(frame_idxs=np.arange(16, 22))
+        expected_frames = np.concatenate(
+            (
+                self.extractors[1].get_frames(frame_idxs=np.arange(6, 10)),
+                self.extractors[2].get_frames(frame_idxs=[0, 1]),
+            ),
+            axis=0,
+        )
+
+        assert_array_equal(test_frames, expected_frames)
+
+    def test_get_all_frames(self):
+        test_frames = self.multi_imaging_extractor.get_frames(frame_idxs=np.arange(0, 30))
+        expected_frames = np.concatenate(
+            [extractor.get_frames(np.arange(0, 10)) for extractor in self.extractors],
+            axis=0,
+        )
+
+        assert_array_equal(test_frames, expected_frames)
+
+    def test_get_video(self):
+        test_frames = self.multi_imaging_extractor.get_video()
+        expected_frames = np.concatenate(
+            [self.extractors[i].get_video() for i in range(3)],
+            axis=0,
+        )
+        assert_array_equal(test_frames, expected_frames)
+
+    def test_set_incorrect_times(self):
+        with self.assertRaisesWith(
+            exc_type=AssertionError,
+            exc_msg="'times' should have the same length of the number of frames!",
+        ):
+            self.multi_imaging_extractor.set_times(times=np.arange(0, 10) / 30.0)
+
+        self.assertEqual(self.multi_imaging_extractor._times, None)
+
+    def test_set_times(self):
+        self.extractors[1].set_times(np.arange(0, 10) / 30.0)
+        multi_imaging_extractor = MultiImagingExtractor(imaging_extractors=self.extractors)
+
+        dummy_times = np.arange(0, 30) / 20.0
+        to_replace = [*range(multi_imaging_extractor._start_frames[1], multi_imaging_extractor._end_frames[1])]
+
+        dummy_times[to_replace] = self.extractors[1]._times
+        assert_array_equal(multi_imaging_extractor._times, dummy_times)
+
+        self.multi_imaging_extractor.set_times(times=dummy_times)
+
+        assert_array_equal(self.multi_imaging_extractor._times, dummy_times)
+
+    @parameterized.expand(
+        [
+            param(
+                rows=3,
+                columns=4,
+                sampling_frequency=15.0,
+                num_channels=1,
+                expected_error_msg="The sampling frequency is not consistent over the files (found {20.0, 15.0}).",
+            ),
+            param(
+                rows=3,
+                columns=5,
+                sampling_frequency=20.0,
+                num_channels=1,
+                expected_error_msg="The size of a frame is not consistent over the files (found {(3, 4), (3, 5)}).",
+            ),
+            param(
+                rows=3,
+                columns=4,
+                sampling_frequency=20.0,
+                num_channels=2,
+                expected_error_msg="The number of channels is not consistent over the files (found {1, 2}).",
+            ),
+        ],
+    )
+    def test_inconsistent_property_assertion(self, rows, columns, sampling_frequency, num_channels, expected_error_msg):
+        inconsistent_extractors = [
+            self.extractors[0],
+            generate_dummy_imaging_extractor(
+                num_frames=1,
+                num_rows=rows,
+                num_columns=columns,
+                sampling_frequency=sampling_frequency,
+                num_channels=num_channels,
+            ),
+        ]
+        with self.assertRaisesWith(
+            exc_type=AssertionError,
+            exc_msg=expected_error_msg,
+        ):
+            MultiImagingExtractor(imaging_extractors=inconsistent_extractors)
+
+
+if __name__ == "__main__":
+    unittest.main()