Repository: catalystneuro/roiextractors

Issue:
Unify Volumetric/MultiPlane Imaging API
# Current Behavior

Right now, the only imaging extractor that supports 3D imaging is BrukerTiff.  It does so by taking a list of `SinglePlaneImagingExtractors` and combines them a la `MultiImagingExtractor`.  Note that each `SinglePlaneImagingExtractor` is itself a `MultiImagingExtractor` that combines multiple 'chunked' files corresponding to a single depth plane.  Each `SinglePlaneImagingExtractor` instance only supports 1 channel and 1 plane, which are passed via `stream_name` during the `__init__`.  Each `MultiPlaneImagingExtractor` instance supports 1 channel (but all planes) which is passed via `stream_name` during the `__init__`.

# Issues

- 'streams' are a low-cohesion concept to represent optical channels and depth planes.
  - They are referred to by string names, which can be different in the `MultiPlaneImagingExtractor` (which only needs to specify the channel) and the `SinglePlaneImagingExtractor` (which needs to specify some formatted combination of channel and plane).
  - They are supposed to be analogous to 'AP', 'LF', etc. streams in spikeinterface, but those streams refer to collections of channels that undergo different types of pre-processing.  In contrast, each stream as they are implemented in BrukerTiff just refer to 1 channel and 1 plane.  There may be a relevant parallel in ophys (ex. functional vs anatomical streams), but specifying a plane/channel combo is not that parallel.
  - optical channels and depth planes are fully independent concepts (1 plane + many channels, 1 channel + many planes, many channels + many planes are all possible) with rich associated metadata that can guide the API: channel name, channel wavelength, plane depth, etc.  By lumping these two axes together as 'streams' we lose the ability to interact with them separately.
- Combining multiple `SinglePlaneImagingExtractors` with `MultiImagingExtractor` is a bit hacky and confusing.
  - `MultiImagingExtractor` was designed to concatenate `ImagingExtractors` along the `frames` axis, but here we want to combine multiple `IEs` along a separate depth axis.
  - Using `MultiImagingExtractor` for this purpose leads to wonky stuff like `start_frames = [0, 0, ..., 0]`.
- BrukerTiff specifies channel during the `__init__` but all other IE's specify channel during the call to `get_video`.
  - Not sure which is better, but we should be consistent.

# Proposed Solutions

- Specify channel and plane separately by index (or ID?) and then the extractor handles the logic of how to load that channel/plane combo in their file structure (Bruker, ScanImage, etc.)
- Add a `MultiPlaneImagingExtractor` base class to handle the combination of IE's along the depth axis.  Could come with extra metadata properties like `get_plane_depth`.
- We need to specify plane at `__init__` time so that SinglePlaneImagingExtractors can be appropriately combined into MultiPlaneImagingExtractors.  So, to keep it consistent we should specify channel also at `__init__` time.  This will require refactoring of all the existing IE's but I think it is the most convenient implementation agnostic to file structure. It also simplifies `get_video` to only take `start_frame` and `end_frame`, which I think is a good thing.

### Do you have any interest in helping implement the feature?

Yes.

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/catalystneuro/roiextractors/blob/master/.github/CODE_OF_CONDUCT.rst)
- [X] Have you ensured this bug was not already [reported](https://github.com/catalystneuro/roiextractors/issues)?

Comments:

Patch:
diff --git a/src/roiextractors/extractorlist.py b/src/roiextractors/extractorlist.py
index 091265c..7241de8 100644
--- a/src/roiextractors/extractorlist.py
+++ b/src/roiextractors/extractorlist.py
@@ -25,6 +25,7 @@ from .extractors.memmapextractors import MemmapImagingExtractor
 from .extractors.miniscopeimagingextractor import MiniscopeImagingExtractor
 from .multisegmentationextractor import MultiSegmentationExtractor
 from .multiimagingextractor import MultiImagingExtractor
+from .volumetricimagingextractor import VolumetricImagingExtractor
 
 imaging_extractor_full_list = [
     NumpyImagingExtractor,
@@ -39,6 +40,7 @@ imaging_extractor_full_list = [
     SbxImagingExtractor,
     NumpyMemmapImagingExtractor,
     MemmapImagingExtractor,
+    VolumetricImagingExtractor,
 ]
 
 segmentation_extractor_full_list = [
diff --git a/src/roiextractors/volumetricimagingextractor.py b/src/roiextractors/volumetricimagingextractor.py
new file mode 100644
index 0000000..2abf0c1
--- /dev/null
+++ b/src/roiextractors/volumetricimagingextractor.py
@@ -0,0 +1,169 @@
+"""Base class definition for volumetric imaging extractors."""
+
+from typing import Tuple, List, Iterable, Optional
+import numpy as np
+
+from .extraction_tools import ArrayType, DtypeType
+from .imagingextractor import ImagingExtractor
+
+
+class VolumetricImagingExtractor(ImagingExtractor):
+    """Class to combine multiple ImagingExtractor objects by depth plane."""
+
+    extractor_name = "VolumetricImaging"
+    installed = True
+    installatiuon_mesage = ""
+
+    def __init__(self, imaging_extractors: List[ImagingExtractor]):
+        """Initialize a VolumetricImagingExtractor object from a list of ImagingExtractors.
+
+        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(imaging_extractor, ImagingExtractor) for imaging_extractor in imaging_extractors)
+        self._check_consistency_between_imaging_extractors(imaging_extractors)
+        self._imaging_extractors = imaging_extractors
+        self._num_planes = len(imaging_extractors)
+
+    def _check_consistency_between_imaging_extractors(self, imaging_extractors: List[ImagingExtractor]):
+        """Check that essential properties are consistent between extractors so that they can be combined appropriately.
+
+        Parameters
+        ----------
+        imaging_extractors: list of ImagingExtractor
+            list of imaging extractor objects
+
+        Raises
+        ------
+        AssertionError
+            If any of the properties are not consistent between extractors.
+
+        Notes
+        -----
+        This method checks the following properties:
+            - sampling frequency
+            - image size
+            - number of channels
+            - channel names
+            - data type
+            - num_frames
+        """
+        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",
+            get_dtype="The data type",
+            get_num_frames="The number of frames",
+        )
+        for method, property_message in properties_to_check.items():
+            values = [getattr(extractor, method)() for extractor in 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_video(self, start_frame: Optional[int] = None, end_frame: Optional[int] = None) -> np.ndarray:
+        """Get the video frames.
+
+        Parameters
+        ----------
+        start_frame: int, optional
+            Start frame index (inclusive).
+        end_frame: int, optional
+            End frame index (exclusive).
+
+        Returns
+        -------
+        video: numpy.ndarray
+            The 3D video frames (num_frames, num_rows, num_columns, num_planes).
+        """
+        if start_frame is None:
+            start_frame = 0
+        elif start_frame < 0:
+            start_frame = self.get_num_frames() + start_frame
+        elif start_frame >= self.get_num_frames():
+            raise ValueError(
+                f"start_frame {start_frame} is greater than or equal to the number of frames {self.get_num_frames()}"
+            )
+        if end_frame is None:
+            end_frame = self.get_num_frames()
+        elif end_frame < 0:
+            end_frame = self.get_num_frames() + end_frame
+        elif end_frame > self.get_num_frames():
+            raise ValueError(f"end_frame {end_frame} is greater than the number of frames {self.get_num_frames()}")
+        if end_frame <= start_frame:
+            raise ValueError(f"end_frame {end_frame} is less than or equal to start_frame {start_frame}")
+
+        video = np.zeros((end_frame - start_frame, *self.get_image_size()), self.get_dtype())
+        for i, imaging_extractor in enumerate(self._imaging_extractors):
+            video[..., i] = imaging_extractor.get_video(start_frame, end_frame)
+        return video
+
+    def get_frames(self, frame_idxs: ArrayType) -> np.ndarray:
+        """Get specific video frames from indices (not necessarily continuous).
+
+        Parameters
+        ----------
+        frame_idxs: array-like
+            Indices of frames to return.
+
+        Returns
+        -------
+        frames: numpy.ndarray
+            The 3D video frames (num_rows, num_columns, num_planes).
+        """
+        if isinstance(frame_idxs, int):
+            frame_idxs = [frame_idxs]
+        for frame_idx in frame_idxs:
+            if frame_idx < -1 * self.get_num_frames() or frame_idx >= self.get_num_frames():
+                raise ValueError(f"frame_idx {frame_idx} is out of bounds")
+
+        # Note np.all([]) returns True so not all(np.diff(frame_idxs) == 1) returns False if frame_idxs is a single int
+        if not all(np.diff(frame_idxs) == 1):
+            frames = np.zeros((len(frame_idxs), *self.get_image_size()), self.get_dtype())
+            for i, imaging_extractor in enumerate(self._imaging_extractors):
+                frames[..., i] = imaging_extractor.get_frames(frame_idxs)
+            return frames
+        else:
+            return self.get_video(start_frame=frame_idxs[0], end_frame=frame_idxs[-1] + 1)
+
+    def get_image_size(self) -> Tuple:
+        """Get the size of a single frame.
+
+        Returns
+        -------
+        image_size: tuple
+            The size of a single frame (num_rows, num_columns, num_planes).
+        """
+        image_size = (*self._imaging_extractors[0].get_image_size(), self.get_num_planes())
+        return image_size
+
+    def get_num_planes(self) -> int:
+        """Get the number of depth planes.
+
+        Returns
+        -------
+        _num_planes: int
+            The number of depth planes.
+        """
+        return self._num_planes
+
+    def get_num_frames(self) -> int:
+        return self._imaging_extractors[0].get_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()
+
+    def get_dtype(self) -> DtypeType:
+        return self._imaging_extractors[0].get_dtype()


Test patch:
diff --git a/src/roiextractors/testing.py b/src/roiextractors/testing.py
index d172089..47d694d 100644
--- a/src/roiextractors/testing.py
+++ b/src/roiextractors/testing.py
@@ -53,6 +53,7 @@ def generate_dummy_imaging_extractor(
     num_channels: int = 1,
     sampling_frequency: float = 30,
     dtype: DtypeType = "uint16",
+    channel_names: Optional[list] = None,
 ):
     """Generate a dummy imaging extractor for testing.
 
@@ -78,7 +79,8 @@ def generate_dummy_imaging_extractor(
     ImagingExtractor
         An imaging extractor with random data fed into `NumpyImagingExtractor`.
     """
-    channel_names = [f"channel_num_{num}" for num in range(num_channels)]
+    if channel_names is None:
+        channel_names = [f"channel_num_{num}" for num in range(num_channels)]
 
     size = (num_frames, num_rows, num_columns, num_channels)
     video = generate_dummy_video(size=size, dtype=dtype)
diff --git a/tests/test_volumetricimagingextractor.py b/tests/test_volumetricimagingextractor.py
new file mode 100644
index 0000000..62424da
--- /dev/null
+++ b/tests/test_volumetricimagingextractor.py
@@ -0,0 +1,121 @@
+import pytest
+import numpy as np
+from roiextractors.testing import generate_dummy_imaging_extractor
+from roiextractors import VolumetricImagingExtractor
+
+num_frames = 10
+
+
+@pytest.fixture(scope="module", params=[1, 2])
+def imaging_extractors(request):
+    num_channels = request.param
+    return [generate_dummy_imaging_extractor(num_channels=num_channels, num_frames=num_frames) for _ in range(3)]
+
+
+@pytest.fixture(scope="module")
+def volumetric_imaging_extractor(imaging_extractors):
+    return VolumetricImagingExtractor(imaging_extractors)
+
+
+@pytest.mark.parametrize(
+    "params",
+    [
+        [dict(sampling_frequency=1), dict(sampling_frequency=2)],
+        [dict(num_rows=1), dict(num_rows=2)],
+        [dict(num_channels=1), dict(num_channels=2)],
+        [dict(channel_names=["a"], num_channels=1), dict(channel_names=["b"], num_channels=1)],
+        [dict(dtype=np.int16), dict(dtype=np.float32)],
+        [dict(num_frames=1), dict(num_frames=2)],
+    ],
+)
+def test_check_consistency_between_imaging_extractors(params):
+    imaging_extractors = [generate_dummy_imaging_extractor(**param) for param in params]
+    with pytest.raises(AssertionError):
+        VolumetricImagingExtractor(imaging_extractors=imaging_extractors)
+
+
+@pytest.mark.parametrize("start_frame, end_frame", [(None, None), (0, num_frames), (3, 7), (-2, -1)])
+def test_get_video(volumetric_imaging_extractor, start_frame, end_frame):
+    video = volumetric_imaging_extractor.get_video(start_frame=start_frame, end_frame=end_frame)
+    expected_video = []
+    for extractor in volumetric_imaging_extractor._imaging_extractors:
+        expected_video.append(extractor.get_video(start_frame=start_frame, end_frame=end_frame))
+    expected_video = np.array(expected_video)
+    expected_video = np.moveaxis(expected_video, 0, -1)
+    assert np.all(video == expected_video)
+
+
+@pytest.mark.parametrize("start_frame, end_frame", [(num_frames + 1, None), (None, num_frames + 1), (2, 1)])
+def test_get_video_invalid(volumetric_imaging_extractor, start_frame, end_frame):
+    with pytest.raises(ValueError):
+        volumetric_imaging_extractor.get_video(start_frame=start_frame, end_frame=end_frame)
+
+
+@pytest.mark.parametrize("frame_idxs", [0, [0, 1, 2], [0, num_frames - 1], [-3, -1]])
+def test_get_frames(volumetric_imaging_extractor, frame_idxs):
+    frames = volumetric_imaging_extractor.get_frames(frame_idxs=frame_idxs)
+    expected_frames = []
+    for extractor in volumetric_imaging_extractor._imaging_extractors:
+        expected_frames.append(extractor.get_frames(frame_idxs=frame_idxs))
+    expected_frames = np.array(expected_frames)
+    expected_frames = np.moveaxis(expected_frames, 0, -1)
+    assert np.all(frames == expected_frames)
+
+
+@pytest.mark.parametrize("frame_idxs", [num_frames, [0, num_frames], [-num_frames - 1, -1]])
+def test_get_frames_invalid(volumetric_imaging_extractor, frame_idxs):
+    with pytest.raises(ValueError):
+        volumetric_imaging_extractor.get_frames(frame_idxs=frame_idxs)
+
+
+@pytest.mark.parametrize("num_rows, num_columns, num_planes", [(1, 2, 3), (2, 1, 3), (3, 2, 1)])
+def test_get_image_size(num_rows, num_columns, num_planes):
+    imaging_extractors = [
+        generate_dummy_imaging_extractor(num_rows=num_rows, num_columns=num_columns) for _ in range(num_planes)
+    ]
+    volumetric_imaging_extractor = VolumetricImagingExtractor(imaging_extractors=imaging_extractors)
+    assert volumetric_imaging_extractor.get_image_size() == (num_rows, num_columns, num_planes)
+
+
+@pytest.mark.parametrize("num_planes", [1, 2, 3])
+def test_get_num_planes(num_planes):
+    imaging_extractors = [generate_dummy_imaging_extractor() for _ in range(num_planes)]
+    volumetric_imaging_extractor = VolumetricImagingExtractor(imaging_extractors=imaging_extractors)
+    assert volumetric_imaging_extractor.get_num_planes() == num_planes
+
+
+@pytest.mark.parametrize("num_frames", [1, 2, 3])
+def test_get_num_frames(num_frames):
+    imaging_extractors = [generate_dummy_imaging_extractor(num_frames=num_frames)]
+    volumetric_imaging_extractor = VolumetricImagingExtractor(imaging_extractors=imaging_extractors)
+    assert volumetric_imaging_extractor.get_num_frames() == num_frames
+
+
+@pytest.mark.parametrize("sampling_frequency", [1, 2, 3])
+def test_get_sampling_frequency(sampling_frequency):
+    imaging_extractors = [generate_dummy_imaging_extractor(sampling_frequency=sampling_frequency)]
+    volumetric_imaging_extractor = VolumetricImagingExtractor(imaging_extractors=imaging_extractors)
+    assert volumetric_imaging_extractor.get_sampling_frequency() == sampling_frequency
+
+
+@pytest.mark.parametrize("channel_names", [["Channel 1"], [" Channel 1 ", "Channel 2"]])
+def test_get_channel_names(channel_names):
+    imaging_extractors = [
+        generate_dummy_imaging_extractor(channel_names=channel_names, num_channels=len(channel_names))
+    ]
+    volumetric_imaging_extractor = VolumetricImagingExtractor(imaging_extractors=imaging_extractors)
+    assert volumetric_imaging_extractor.get_channel_names() == channel_names
+
+
+@pytest.mark.parametrize("num_channels", [1, 2, 3])
+def test_get_num_channels(num_channels):
+    imaging_extractors = [generate_dummy_imaging_extractor(num_channels=num_channels)]
+    volumetric_imaging_extractor = VolumetricImagingExtractor(imaging_extractors=imaging_extractors)
+    assert volumetric_imaging_extractor.get_num_channels() == num_channels
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.int16, np.uint8])
+def test_get_dtype(dtype):
+    imaging_extractors = [generate_dummy_imaging_extractor(dtype=dtype)]
+    volumetric_imaging_extractor = VolumetricImagingExtractor(imaging_extractors=imaging_extractors)
+    assert volumetric_imaging_extractor.get_dtype() == dtype