From d6be16576d27366b539554cdc055404d9046c50d Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Thu, 14 Sep 2023 16:53:08 +0200
Subject: [PATCH 01/22] update with streams

---
 .../suite2p/suite2psegmentationextractor.py   | 155 ++++++++++++------
 tests/test_suite2psegmentationextractor.py    |  45 +++++
 2 files changed, 152 insertions(+), 48 deletions(-)
 create mode 100644 tests/test_suite2psegmentationextractor.py

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index fd64ce54..42e2d259 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -1,10 +1,11 @@
 import shutil
 from pathlib import Path
 from typing import Optional
+from warnings import warn
 
 import numpy as np
 
-from ...extraction_tools import PathType, IntType
+from ...extraction_tools import PathType
 from ...extraction_tools import _image_mask_extractor
 from ...multisegmentationextractor import MultiSegmentationExtractor
 from ...segmentationextractor import SegmentationExtractor
@@ -14,82 +15,140 @@ class Suite2pSegmentationExtractor(SegmentationExtractor):
     extractor_name = "Suite2pSegmentationExtractor"
     installed = True  # check at class level if installed or not
     is_writable = False
-    mode = "file"
+    mode = "folder"
     installation_mesg = ""  # error message when not installed
 
+    @classmethod
+    def get_streams(cls, folder_path: PathType):
+        folder_path = Path(folder_path)
+        stream_paths = [f for f in folder_path.iterdir() if f.is_dir()]
+        chan_1_streams = [f"chan1_{stream_path.stem}" for stream_path in stream_paths]
+        streams = dict(channel_streams=["chan1"], plane_streams=dict(chan1=chan_1_streams))
+
+        chan_2_streams = []
+        for stream_path in stream_paths:
+            if list(stream_path.glob("F_chan2.npy")):
+                chan_2_streams.append(f"chan2_{stream_path.stem}")
+
+        if chan_2_streams:
+            streams["channel_streams"].append("chan2")
+            streams["plane_streams"].update(chan2=chan_2_streams)
+
+        return streams
+
     def __init__(
         self,
-        folder_path: Optional[PathType] = None,
-        combined: bool = False,
-        plane_no: IntType = 0,
-        file_path: Optional[PathType] = None,
+        folder_path: PathType,
+        stream_name: Optional[str] = None,
+        combined: Optional[bool] = None,  # TODO: to be removed
+        plane_no: Optional[int] = None,  # TODO: to be removed
     ):
         """
-        Creating SegmentationExtractor object out of suite 2p data type.
+        The SegmentationExtractor for the suite2p format.
+
         Parameters
         ----------
         folder_path: str or Path
-            ~/suite2p folder location on disk
-        combined: bool
-            if the plane is a combined plane as in the Suite2p pipeline
-        plane_no: int
-            the plane for which to extract segmentation for.
-        file_path: str or Path [Deprecated]
-            ~/suite2p folder location on disk
-
+            The path to the 'suite2p' folder.
+        stream_name: str, optional
+            The name of the stream to load, to determine which streams are available use Suite2pSegmentationExtractor.get_streams(folder_path).
         """
-        from warnings import warn
 
-        if file_path is not None:
+        if combined:
+            warning_string = "Keyword argument 'combined' is deprecated and will be removed on or after Nov, 2023. "
+            warn(
+                message=warning_string,
+                category=DeprecationWarning,
+            )
+        if plane_no:
             warning_string = (
-                "The keyword argument 'file_path' is being deprecated on or after August, 2022 in favor of 'folder_path'. "
-                "'folder_path' takes precence over 'file_path'."
+                "Keyword argument 'plane_no' is deprecated and will be removed on or after Nov, 2023 in favor of 'stream_name'."
+                "Specify which stream you wish to load with the 'stream_name' keyword argument."
             )
             warn(
                 message=warning_string,
                 category=DeprecationWarning,
             )
-            folder_path = file_path if folder_path is None else folder_path
 
-        SegmentationExtractor.__init__(self)
-        self.combined = combined
-        self.plane_no = plane_no
+        streams = self.get_streams(folder_path=folder_path)
+        if stream_name is None:
+            if len(streams["channel_streams"]) > 1:
+                raise ValueError(
+                    "More than one channel is detected! Please specify which stream you wish to load with the `stream_name` argument. "
+                    "To see what streams are available, call `Suite2pSegmentationExtractor.get_streams(folder_path=...)`."
+                )
+            channel_stream_name = streams["channel_streams"][0]
+            stream_name = streams["plane_streams"][channel_stream_name][0]
+
+        channel_stream_name = stream_name.split("_")[0]
+        if channel_stream_name not in streams["channel_streams"]:
+            raise ValueError(
+                f"The selected stream '{channel_stream_name}' is not a valid stream name. To see what streams are available, "
+                f"call `Suite2pSegmentationExtractor.get_streams(folder_path=...)`."
+            )
+
+        plane_stream_names = streams["plane_streams"][channel_stream_name]
+        if stream_name is not None and stream_name not in plane_stream_names:
+            raise ValueError(
+                f"The selected stream '{stream_name}' is not in the available plane_streams '{plane_stream_names}'!"
+            )
+        self.stream_name = stream_name
+
+        super().__init__()
+
         self.folder_path = Path(folder_path)
 
-        self.stat = self._load_npy("stat.npy")
-        self._roi_response_raw = self._load_npy("F.npy", mmap_mode="r").T
-        self._roi_response_neuropil = self._load_npy("Fneu.npy", mmap_mode="r").T
-        self._roi_response_deconvolved = self._load_npy("spks.npy", mmap_mode="r").T
+        self._options = self._load_npy(file_name="ops.npy").item()
+
+        fluorescence_traces_file_name = "F.npy" if channel_stream_name == "chan1" else "F_chan2.npy"
+        neuropil_traces_file_name = "Fneu.npy" if channel_stream_name == "chan1" else "Fneu_chan2.npy"
+
+        self._sampling_frequency = self._options["fs"]
+        self._num_frames = self._options["nframes"]
+        self._image_size = (self._options["Ly"], self._options["Lx"])
+
+        self.stat = self._load_npy(file_name="stat.npy")
+        self._roi_response_raw = self._load_npy(file_name=fluorescence_traces_file_name, mmap_mode="r").T
+        self._roi_response_neuropil = self._load_npy(file_name=neuropil_traces_file_name, mmap_mode="r").T
+        self._roi_response_deconvolved = self._load_npy(file_name="spks.npy", mmap_mode="r").T
         self.iscell = self._load_npy("iscell.npy", mmap_mode="r")
-        self.ops = self._load_npy("ops.npy").item()
 
-        self._channel_names = [f"OpticalChannel{i}" for i in range(self.ops["nchannels"])]
-        self._sampling_frequency = self.ops["fs"] * [2 if self.combined else 1][0]
-        self._raw_movie_file_location = self.ops.get("filelist", [None])[0]
-        self._image_correlation = self._summary_image_read("Vcorr")
-        self._image_mean = self._summary_image_read("meanImg")
+        channel_name = (
+            "OpticalChannel"
+            if len(streams["channel_streams"]) == 1
+            else channel_stream_name.capitalize()
+        )
+        self._channel_names = [channel_name]
+
+        self._image_correlation = self._correlation_image_read()
+        image_mean_name = "meanImg" if channel_stream_name == "chan1" else f"meanImg_chan2"
+        self._image_mean = self._options[image_mean_name]
 
-    def _load_npy(self, filename, mmap_mode=None):
-        file_path = self.folder_path / f"plane{self.plane_no}" / filename
+    def _load_npy(self, file_name: str, mmap_mode=None):
+        plane_stream_name = self.stream_name.split("_")[-1]
+        file_path = self.folder_path / plane_stream_name / file_name
         return np.load(file_path, mmap_mode=mmap_mode, allow_pickle=mmap_mode is None)
 
+    def get_num_frames(self) -> int:
+        return self._num_frames
+
     def get_accepted_list(self):
         return list(np.where(self.iscell[:, 0] == 1)[0])
 
     def get_rejected_list(self):
         return list(np.where(self.iscell[:, 0] == 0)[0])
 
-    def _summary_image_read(self, bstr="meanImg"):
-        img = None
-        if bstr in self.ops:
-            if bstr == "Vcorr" or bstr == "max_proj":
-                img = np.zeros((self.ops["Ly"], self.ops["Lx"]), np.float32)
-                img[
-                    (self.ops["Ly"] - self.ops["yrange"][-1]) : (self.ops["Ly"] - self.ops["yrange"][0]),
-                    self.ops["xrange"][0] : self.ops["xrange"][-1],
-                ] = self.ops[bstr]
-            else:
-                img = self.ops[bstr]
+    def _correlation_image_read(self):
+        correlation_image = self._options["Vcorr"]
+        if (self._options["yrange"][-1], self._options["xrange"][-1]) == self._image_size:
+            return correlation_image
+
+        img = np.zeros(self._image_size, correlation_image.dtype)
+        img[
+            (self._options["Ly"] - self._options["yrange"][-1]) : (self._options["Ly"] - self._options["yrange"][0]),
+            self._options["xrange"][0] : self._options["xrange"][-1],
+        ] = correlation_image
+
         return img
 
     @property
@@ -115,7 +174,7 @@ def get_roi_pixel_masks(self, roi_ids=None):
             pixel_mask.append(
                 np.vstack(
                     [
-                        self.ops["Ly"] - 1 - self.stat[i]["ypix"],
+                        self._options["Ly"] - 1 - self.stat[i]["ypix"],
                         self.stat[i]["xpix"],
                         self.stat[i]["lam"],
                     ]
@@ -130,7 +189,7 @@ def get_roi_pixel_masks(self, roi_ids=None):
         return [pixel_mask[i] for i in roi_idx_]
 
     def get_image_size(self):
-        return [self.ops["Ly"], self.ops["Lx"]]
+        return self._image_size
 
     @staticmethod
     def write_segmentation(segmentation_object: SegmentationExtractor, save_path: PathType, overwrite=True):
diff --git a/tests/test_suite2psegmentationextractor.py b/tests/test_suite2psegmentationextractor.py
new file mode 100644
index 00000000..9fe83083
--- /dev/null
+++ b/tests/test_suite2psegmentationextractor.py
@@ -0,0 +1,45 @@
+from hdmf.testing import TestCase
+
+from roiextractors import Suite2pSegmentationExtractor
+from tests.setup_paths import OPHYS_DATA_PATH
+
+
+class TestSuite2pSegmentationExtractor(TestCase):
+    @classmethod
+    def setUpClass(cls):
+        folder_path = str(OPHYS_DATA_PATH / "segmentation_datasets" / "suite2p")
+        cls.available_streams = dict(
+            channel_streams=["chan1", "chan2"],
+            plane_streams=dict(chan1=["chan1_plane0", "chan1_plane1", "chan1_combined"], chan2=["chan2_plane0", "chan2_plane1"])
+        )
+
+        cls.folder_path = folder_path
+
+        extractor = Suite2pSegmentationExtractor(folder_path=folder_path, stream_name="chan1_plane2")
+        cls.extractor = extractor
+
+    def test_stream_names(self):
+        self.assertEqual(Suite2pSegmentationExtractor.get_streams(folder_path=self.folder_path), self.available_streams)
+
+    def test_multi_stream_warns(self):
+        exc_msg = (
+            "More than one channel is detected! Please specify which stream you wish to load with the `stream_name` argument. "
+            "To see what streams are available, call `Suite2pSegmentationExtractor.get_streams(folder_path=...)`."
+            "This is going to raise ValueError in the future."
+        )
+        with self.assertRaisesWith(exc_type=ValueError, exc_msg=exc_msg):
+            Suite2pSegmentationExtractor(folder_path=self.folder_path)
+
+    def test_invalid_stream_raises(self):
+        exc_msg = (
+            "The selected stream 'plane0' is not a valid stream name. To see what streams are available, call `Suite2pSegmentationExtractor.get_streams(folder_path=...)`."
+        )
+        with self.assertRaisesWith(exc_type=ValueError, exc_msg=exc_msg):
+            Suite2pSegmentationExtractor(folder_path=self.folder_path, stream_name="plane0")
+
+    def test_incorrect_stream_raises(self):
+        exc_msg = (
+            "The selected stream 'chan1_plane2' is not in the available plane_streams '['chan1_plane0', 'chan1_plane1', 'chan1_combined']'!"
+        )
+        with self.assertRaisesWith(exc_type=ValueError, exc_msg=exc_msg):
+            Suite2pSegmentationExtractor(folder_path=self.folder_path, stream_name="chan1_plane2")

From 03fbc1e200bf7c9119111043159ef5f5d37306d0 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 14 Sep 2023 15:05:51 +0000
Subject: [PATCH 02/22] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../suite2p/suite2psegmentationextractor.py          |  6 +-----
 tests/test_suite2psegmentationextractor.py           | 12 +++++-------
 2 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index d7da33e9..9b3ba85b 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -122,11 +122,7 @@ def __init__(
         self._roi_response_deconvolved = self._load_npy(file_name="spks.npy", mmap_mode="r").T
         self.iscell = self._load_npy("iscell.npy", mmap_mode="r")
 
-        channel_name = (
-            "OpticalChannel"
-            if len(streams["channel_streams"]) == 1
-            else channel_stream_name.capitalize()
-        )
+        channel_name = "OpticalChannel" if len(streams["channel_streams"]) == 1 else channel_stream_name.capitalize()
         self._channel_names = [channel_name]
 
         self._image_correlation = self._correlation_image_read()
diff --git a/tests/test_suite2psegmentationextractor.py b/tests/test_suite2psegmentationextractor.py
index 9fe83083..3038a640 100644
--- a/tests/test_suite2psegmentationextractor.py
+++ b/tests/test_suite2psegmentationextractor.py
@@ -10,7 +10,9 @@ def setUpClass(cls):
         folder_path = str(OPHYS_DATA_PATH / "segmentation_datasets" / "suite2p")
         cls.available_streams = dict(
             channel_streams=["chan1", "chan2"],
-            plane_streams=dict(chan1=["chan1_plane0", "chan1_plane1", "chan1_combined"], chan2=["chan2_plane0", "chan2_plane1"])
+            plane_streams=dict(
+                chan1=["chan1_plane0", "chan1_plane1", "chan1_combined"], chan2=["chan2_plane0", "chan2_plane1"]
+            ),
         )
 
         cls.folder_path = folder_path
@@ -31,15 +33,11 @@ def test_multi_stream_warns(self):
             Suite2pSegmentationExtractor(folder_path=self.folder_path)
 
     def test_invalid_stream_raises(self):
-        exc_msg = (
-            "The selected stream 'plane0' is not a valid stream name. To see what streams are available, call `Suite2pSegmentationExtractor.get_streams(folder_path=...)`."
-        )
+        exc_msg = "The selected stream 'plane0' is not a valid stream name. To see what streams are available, call `Suite2pSegmentationExtractor.get_streams(folder_path=...)`."
         with self.assertRaisesWith(exc_type=ValueError, exc_msg=exc_msg):
             Suite2pSegmentationExtractor(folder_path=self.folder_path, stream_name="plane0")
 
     def test_incorrect_stream_raises(self):
-        exc_msg = (
-            "The selected stream 'chan1_plane2' is not in the available plane_streams '['chan1_plane0', 'chan1_plane1', 'chan1_combined']'!"
-        )
+        exc_msg = "The selected stream 'chan1_plane2' is not in the available plane_streams '['chan1_plane0', 'chan1_plane1', 'chan1_combined']'!"
         with self.assertRaisesWith(exc_type=ValueError, exc_msg=exc_msg):
             Suite2pSegmentationExtractor(folder_path=self.folder_path, stream_name="chan1_plane2")

From 8b04ae77831ae57ed176ee6d5c176366a79fd16b Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Thu, 14 Sep 2023 17:10:48 +0200
Subject: [PATCH 03/22] fix io test

---
 tests/test_io.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/tests/test_io.py b/tests/test_io.py
index f2692b2d..9932a767 100644
--- a/tests/test_io.py
+++ b/tests/test_io.py
@@ -138,11 +138,10 @@ def test_imaging_extractors_canonical_shape(self, extractor_class, extractor_kwa
         ),
         param(
             extractor_class=Suite2pSegmentationExtractor,
-            extractor_kwargs=dict(folder_path=str(OPHYS_DATA_PATH / "segmentation_datasets" / "suite2p")),
-        ),
-        param(
-            extractor_class=Suite2pSegmentationExtractor,
-            extractor_kwargs=dict(file_path=str(OPHYS_DATA_PATH / "segmentation_datasets" / "suite2p")),
+            extractor_kwargs=dict(
+                folder_path=str(OPHYS_DATA_PATH / "segmentation_datasets" / "suite2p"),
+                stream_name="chan1_plane0"
+            ),
         ),
     ]
 

From 1a35f8fc797ea795e2f972ba8dc0d30e38b3a4b9 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 14 Sep 2023 15:11:08 +0000
Subject: [PATCH 04/22] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/test_io.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/test_io.py b/tests/test_io.py
index 9932a767..43f567a5 100644
--- a/tests/test_io.py
+++ b/tests/test_io.py
@@ -139,8 +139,7 @@ def test_imaging_extractors_canonical_shape(self, extractor_class, extractor_kwa
         param(
             extractor_class=Suite2pSegmentationExtractor,
             extractor_kwargs=dict(
-                folder_path=str(OPHYS_DATA_PATH / "segmentation_datasets" / "suite2p"),
-                stream_name="chan1_plane0"
+                folder_path=str(OPHYS_DATA_PATH / "segmentation_datasets" / "suite2p"), stream_name="chan1_plane0"
             ),
         ),
     ]

From 28dabeb8deb35aa9150e26f46507a7d9cd3df482 Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Thu, 14 Sep 2023 17:27:22 +0200
Subject: [PATCH 05/22] fix extractor test

---
 tests/test_suite2psegmentationextractor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_suite2psegmentationextractor.py b/tests/test_suite2psegmentationextractor.py
index 3038a640..df86cad5 100644
--- a/tests/test_suite2psegmentationextractor.py
+++ b/tests/test_suite2psegmentationextractor.py
@@ -17,7 +17,7 @@ def setUpClass(cls):
 
         cls.folder_path = folder_path
 
-        extractor = Suite2pSegmentationExtractor(folder_path=folder_path, stream_name="chan1_plane2")
+        extractor = Suite2pSegmentationExtractor(folder_path=folder_path, stream_name="chan1_plane1")
         cls.extractor = extractor
 
     def test_stream_names(self):

From d9474fa2c0c0338626dfaae81e697f08b3226f0e Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Thu, 14 Sep 2023 17:28:45 +0200
Subject: [PATCH 06/22] pydocstyle

---
 .../extractors/suite2p/suite2psegmentationextractor.py          | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index 9b3ba85b..5f447999 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -62,7 +62,6 @@ def __init__(
             The name of the stream to load, to determine which streams are available use Suite2pSegmentationExtractor.get_streams(folder_path).
 
         """
-
         if combined:
             warning_string = "Keyword argument 'combined' is deprecated and will be removed on or after Nov, 2023. "
             warn(
@@ -159,7 +158,6 @@ def get_rejected_list(self):
     def _correlation_image_read(self):
         """Read correlation image from ops (settings) dict.
 
-
         Returns
         -------
         img : numpy.ndarray | None

From 2874854ca7203cf70be40531a6a9e9a45bba5659 Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Sat, 16 Sep 2023 17:26:46 +0200
Subject: [PATCH 07/22] fix tests

---
 .../extractors/suite2p/suite2psegmentationextractor.py |  4 +++-
 tests/test_suite2psegmentationextractor.py             | 10 +++-------
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index 5f447999..1a6a77fd 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -29,8 +29,10 @@ class Suite2pSegmentationExtractor(SegmentationExtractor):
 
     @classmethod
     def get_streams(cls, folder_path: PathType):
+        from natsort import natsorted
+
         folder_path = Path(folder_path)
-        stream_paths = [f for f in folder_path.iterdir() if f.is_dir()]
+        stream_paths = natsorted([f for f in folder_path.iterdir() if f.is_dir()])
         chan_1_streams = [f"chan1_{stream_path.stem}" for stream_path in stream_paths]
         streams = dict(channel_streams=["chan1"], plane_streams=dict(chan1=chan_1_streams))
 
diff --git a/tests/test_suite2psegmentationextractor.py b/tests/test_suite2psegmentationextractor.py
index df86cad5..f6822584 100644
--- a/tests/test_suite2psegmentationextractor.py
+++ b/tests/test_suite2psegmentationextractor.py
@@ -11,7 +11,7 @@ def setUpClass(cls):
         cls.available_streams = dict(
             channel_streams=["chan1", "chan2"],
             plane_streams=dict(
-                chan1=["chan1_plane0", "chan1_plane1", "chan1_combined"], chan2=["chan2_plane0", "chan2_plane1"]
+                chan1=["chan1_combined", "chan1_plane0", "chan1_plane1"], chan2=["chan2_plane0", "chan2_plane1"]
             ),
         )
 
@@ -24,11 +24,7 @@ def test_stream_names(self):
         self.assertEqual(Suite2pSegmentationExtractor.get_streams(folder_path=self.folder_path), self.available_streams)
 
     def test_multi_stream_warns(self):
-        exc_msg = (
-            "More than one channel is detected! Please specify which stream you wish to load with the `stream_name` argument. "
-            "To see what streams are available, call `Suite2pSegmentationExtractor.get_streams(folder_path=...)`."
-            "This is going to raise ValueError in the future."
-        )
+        exc_msg = "More than one channel is detected! Please specify which stream you wish to load with the `stream_name` argument. To see what streams are available, call `Suite2pSegmentationExtractor.get_streams(folder_path=...)`."
         with self.assertRaisesWith(exc_type=ValueError, exc_msg=exc_msg):
             Suite2pSegmentationExtractor(folder_path=self.folder_path)
 
@@ -38,6 +34,6 @@ def test_invalid_stream_raises(self):
             Suite2pSegmentationExtractor(folder_path=self.folder_path, stream_name="plane0")
 
     def test_incorrect_stream_raises(self):
-        exc_msg = "The selected stream 'chan1_plane2' is not in the available plane_streams '['chan1_plane0', 'chan1_plane1', 'chan1_combined']'!"
+        exc_msg = "The selected stream 'chan1_plane2' is not in the available plane_streams '['chan1_combined', 'chan1_plane0', 'chan1_plane1']'!"
         with self.assertRaisesWith(exc_type=ValueError, exc_msg=exc_msg):
             Suite2pSegmentationExtractor(folder_path=self.folder_path, stream_name="chan1_plane2")

From fc07661468debf821c7cc24563ca6c7e71fb8509 Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Mon, 18 Sep 2023 09:56:32 +0200
Subject: [PATCH 08/22] allow for incomplete input

---
 .../suite2p/suite2psegmentationextractor.py   | 53 ++++++++++++-------
 1 file changed, 33 insertions(+), 20 deletions(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index 1a6a77fd..7120e981 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -108,19 +108,20 @@ def __init__(
 
         self.folder_path = Path(folder_path)
 
-        self._options = self._load_npy(file_name="ops.npy").item()
+        options = self._load_npy(file_name="ops.npy")
+        self.options = options.item() if options is not None else options
+        self._sampling_frequency = self.options["fs"]
+        self._num_frames = self.options["nframes"]
+        self._image_size = (self.options["Ly"], self.options["Lx"])
+
+        self.stat = self._load_npy(file_name="stat.npy")
 
         fluorescence_traces_file_name = "F.npy" if channel_stream_name == "chan1" else "F_chan2.npy"
         neuropil_traces_file_name = "Fneu.npy" if channel_stream_name == "chan1" else "Fneu_chan2.npy"
+        self._roi_response_raw = self._load_npy(file_name=fluorescence_traces_file_name, mmap_mode="r", transpose=True)
+        self._roi_response_neuropil = self._load_npy(file_name=neuropil_traces_file_name, mmap_mode="r", transpose=True)
+        self._roi_response_deconvolved = self._load_npy(file_name="spks.npy", mmap_mode="r", transpose=True)
 
-        self._sampling_frequency = self._options["fs"]
-        self._num_frames = self._options["nframes"]
-        self._image_size = (self._options["Ly"], self._options["Lx"])
-
-        self.stat = self._load_npy(file_name="stat.npy")
-        self._roi_response_raw = self._load_npy(file_name=fluorescence_traces_file_name, mmap_mode="r").T
-        self._roi_response_neuropil = self._load_npy(file_name=neuropil_traces_file_name, mmap_mode="r").T
-        self._roi_response_deconvolved = self._load_npy(file_name="spks.npy", mmap_mode="r").T
         self.iscell = self._load_npy("iscell.npy", mmap_mode="r")
 
         channel_name = "OpticalChannel" if len(streams["channel_streams"]) == 1 else channel_stream_name.capitalize()
@@ -128,17 +129,19 @@ def __init__(
 
         self._image_correlation = self._correlation_image_read()
         image_mean_name = "meanImg" if channel_stream_name == "chan1" else f"meanImg_chan2"
-        self._image_mean = self._options[image_mean_name]
+        self._image_mean = self.options[image_mean_name] if image_mean_name in self.options else None
 
-    def _load_npy(self, file_name: str, mmap_mode=None):
-        """Load a .npy file with specified filename.
+    def _load_npy(self, file_name: str, mmap_mode=None, transpose: bool = False):
+        """Load a .npy file with specified filename. Returns None if file is missing.
 
         Parameters
         ----------
-        filename: str
+        file_name: str
             The name of the .npy file to load.
         mmap_mode: str
             The mode to use for memory mapping. See numpy.load for details.
+        transpose: bool, optional
+            Whether to transpose the loaded array.
 
         Returns
         -------
@@ -146,7 +149,14 @@ def _load_npy(self, file_name: str, mmap_mode=None):
         """
         plane_stream_name = self.stream_name.split("_")[-1]
         file_path = self.folder_path / plane_stream_name / file_name
-        return np.load(file_path, mmap_mode=mmap_mode, allow_pickle=mmap_mode is None)
+        if not file_path.exists():
+            return
+
+        data = np.load(file_path, mmap_mode=mmap_mode, allow_pickle=mmap_mode is None)
+        if transpose:
+            return data.T
+
+        return data
 
     def get_num_frames(self) -> int:
         return self._num_frames
@@ -165,14 +175,17 @@ def _correlation_image_read(self):
         img : numpy.ndarray | None
             The correlation image.
         """
-        correlation_image = self._options["Vcorr"]
-        if (self._options["yrange"][-1], self._options["xrange"][-1]) == self._image_size:
+        if "Vcorr" not in self.options:
+            return None
+
+        correlation_image = self.options["Vcorr"]
+        if (self.options["yrange"][-1], self.options["xrange"][-1]) == self._image_size:
             return correlation_image
 
         img = np.zeros(self._image_size, correlation_image.dtype)
         img[
-            (self._options["Ly"] - self._options["yrange"][-1]) : (self._options["Ly"] - self._options["yrange"][0]),
-            self._options["xrange"][0] : self._options["xrange"][-1],
+            (self.options["Ly"] - self.options["yrange"][-1]) : (self.options["Ly"] - self.options["yrange"][0]),
+            self.options["xrange"][0] : self.options["xrange"][-1],
         ] = correlation_image
 
         return img
@@ -201,7 +214,7 @@ def get_roi_pixel_masks(self, roi_ids=None):
             pixel_mask.append(
                 np.vstack(
                     [
-                        self._options["Ly"] - 1 - self.stat[i]["ypix"],
+                        self.stat[i]["ypix"],
                         self.stat[i]["xpix"],
                         self.stat[i]["lam"],
                     ]
@@ -289,7 +302,7 @@ def write_segmentation(segmentation_object: SegmentationExtractor, save_path: Pa
         for no, i in enumerate(stat):
             stat[no] = {
                 "med": roi_locs[no, :].tolist(),
-                "ypix": segmentation_object.get_image_size()[0] - 1 - pixel_masks[no][:, 0],
+                "ypix": pixel_masks[no][:, 0],
                 "xpix": pixel_masks[no][:, 1],
                 "lam": pixel_masks[no][:, 2],
             }

From 3ddda8970fecd71c4385e3153159e135b38837f2 Mon Sep 17 00:00:00 2001
From: Alessandra Trapani <alessandramaria.trapani@polimi.it>
Date: Mon, 2 Oct 2023 19:03:17 +0200
Subject: [PATCH 09/22] change from stream to channel+plane name

---
 .../suite2p/suite2psegmentationextractor.py   | 104 +++++++++++-------
 1 file changed, 64 insertions(+), 40 deletions(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index 7120e981..86d57e9b 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -9,7 +9,7 @@
 from pathlib import Path
 from typing import Optional
 from warnings import warn
-
+import os 
 import numpy as np
 
 from ...extraction_tools import PathType
@@ -28,29 +28,55 @@ class Suite2pSegmentationExtractor(SegmentationExtractor):
     installation_mesg = ""  # error message when not installed
 
     @classmethod
-    def get_streams(cls, folder_path: PathType):
-        from natsort import natsorted
+    def get_available_channels(cls, folder_path: PathType):
+        """Get the available channel names from the folder paths produced by Suite2p.
+        outputs
+        ----------
+        file_path : PathType
+            Path to Suite2p output path.
+        Returns
+        -------
+        channel_names: list
+            List of channel names.
+        """
 
-        folder_path = Path(folder_path)
-        stream_paths = natsorted([f for f in folder_path.iterdir() if f.is_dir()])
-        chan_1_streams = [f"chan1_{stream_path.stem}" for stream_path in stream_paths]
-        streams = dict(channel_streams=["chan1"], plane_streams=dict(chan1=chan_1_streams))
+        plane_names = cls.get_available_planes(folder_path=folder_path)
+
+        channel_names = ["chan1"]
+        file_path = folder_path / plane_names[0] / "F_chan2.npy"
 
-        chan_2_streams = []
-        for stream_path in stream_paths:
-            if list(stream_path.glob("F_chan2.npy")):
-                chan_2_streams.append(f"chan2_{stream_path.stem}")
+        if os.path.isfile(file_path):
+            channel_names.append("chan2")
 
-        if chan_2_streams:
-            streams["channel_streams"].append("chan2")
-            streams["plane_streams"].update(chan2=chan_2_streams)
+        return channel_names
+    
+    @classmethod
+    def get_available_planes(cls, folder_path: PathType):
+        """Get the available plane names from the folder produced by Suite2p.
+        outputs
+        ----------
+        file_path : PathType
+            Path to Suite2p output path.
+        Returns
+        -------
+        plane_names: list
+            List of plane names.
+        """
+
+        folder_path = Path(folder_path)
+        prefix = 'plane'
+        plane_names = []
+        for item in os.listdir(folder_path):
+            if item.startswith(prefix):   
+                plane_names.append(item)
 
-        return streams
+        return plane_names
 
     def __init__(
         self,
         folder_path: PathType,
-        stream_name: Optional[str] = None,
+        channel_name: Optional[str] = None,
+        plane_name: Optional[str] = None,
         combined: Optional[bool] = None,  # TODO: to be removed
         plane_no: Optional[int] = None,  # TODO: to be removed
     ):
@@ -60,7 +86,7 @@ def __init__(
         ----------
         folder_path: str or Path
             The path to the 'suite2p' folder.
-        stream_name: str, optional
+        plane_name: str, optional
             The name of the stream to load, to determine which streams are available use Suite2pSegmentationExtractor.get_streams(folder_path).
 
         """
@@ -72,37 +98,35 @@ def __init__(
             )
         if plane_no:
             warning_string = (
-                "Keyword argument 'plane_no' is deprecated and will be removed on or after Nov, 2023 in favor of 'stream_name'."
-                "Specify which stream you wish to load with the 'stream_name' keyword argument."
+                "Keyword argument 'plane_no' is deprecated and will be removed on or after Nov, 2023 in favor of 'plane_name'."
+                "Specify which stream you wish to load with the 'plane_name' keyword argument."
             )
             warn(
                 message=warning_string,
                 category=DeprecationWarning,
             )
 
-        streams = self.get_streams(folder_path=folder_path)
-        if stream_name is None:
-            if len(streams["channel_streams"]) > 1:
+        channel_names = self.get_available_channels(folder_path=folder_path)
+        if channel_name is None:
+            if len(channel_names) > 1:
                 raise ValueError(
-                    "More than one channel is detected! Please specify which stream you wish to load with the `stream_name` argument. "
-                    "To see what streams are available, call `Suite2pSegmentationExtractor.get_streams(folder_path=...)`."
+                    "More than one channel is detected! Please specify which channel you wish to load with the `channel_name` argument. "
+                    "To see what streams are available, call `Suite2pSegmentationExtractor.get_channel_names(folder_path=...)`."
                 )
-            channel_stream_name = streams["channel_streams"][0]
-            stream_name = streams["plane_streams"][channel_stream_name][0]
+            channel_name = channel_names["channel_streams"][0]
 
-        channel_stream_name = stream_name.split("_")[0]
-        if channel_stream_name not in streams["channel_streams"]:
+        if channel_name not in channel_names["channel_streams"]:
             raise ValueError(
-                f"The selected stream '{channel_stream_name}' is not a valid stream name. To see what streams are available, "
-                f"call `Suite2pSegmentationExtractor.get_streams(folder_path=...)`."
+                f"The selected channel '{channel_name}' is not a valid stream name. To see what channels are available, "
+                f"call `Suite2pSegmentationExtractor.get_channel_names(folder_path=...)`."
             )
 
-        plane_stream_names = streams["plane_streams"][channel_stream_name]
-        if stream_name is not None and stream_name not in plane_stream_names:
+        plane_names = self.get_available_planes(folder_path=folder_path)
+        if plane_name is not None and plane_name not in plane_names:
             raise ValueError(
-                f"The selected stream '{stream_name}' is not in the available plane_streams '{plane_stream_names}'!"
+                f"The selected plane '{plane_name}' is not in the available plane_streams '{plane_names['plane_streams']}'!"
             )
-        self.stream_name = stream_name
+        self.plane_name = plane_name
 
         super().__init__()
 
@@ -116,19 +140,19 @@ def __init__(
 
         self.stat = self._load_npy(file_name="stat.npy")
 
-        fluorescence_traces_file_name = "F.npy" if channel_stream_name == "chan1" else "F_chan2.npy"
-        neuropil_traces_file_name = "Fneu.npy" if channel_stream_name == "chan1" else "Fneu_chan2.npy"
+        fluorescence_traces_file_name = "F.npy" if channel_name == "chan1" else "F_chan2.npy"
+        neuropil_traces_file_name = "Fneu.npy" if channel_name == "chan1" else "Fneu_chan2.npy"
         self._roi_response_raw = self._load_npy(file_name=fluorescence_traces_file_name, mmap_mode="r", transpose=True)
         self._roi_response_neuropil = self._load_npy(file_name=neuropil_traces_file_name, mmap_mode="r", transpose=True)
         self._roi_response_deconvolved = self._load_npy(file_name="spks.npy", mmap_mode="r", transpose=True)
 
         self.iscell = self._load_npy("iscell.npy", mmap_mode="r")
 
-        channel_name = "OpticalChannel" if len(streams["channel_streams"]) == 1 else channel_stream_name.capitalize()
+        channel_name = "OpticalChannel" if len(channel_names["channel_streams"]) == 1 else channel_name.capitalize()
         self._channel_names = [channel_name]
 
         self._image_correlation = self._correlation_image_read()
-        image_mean_name = "meanImg" if channel_stream_name == "chan1" else f"meanImg_chan2"
+        image_mean_name = "meanImg" if channel_name == "chan1" else f"meanImg_chan2"
         self._image_mean = self.options[image_mean_name] if image_mean_name in self.options else None
 
     def _load_npy(self, file_name: str, mmap_mode=None, transpose: bool = False):
@@ -147,8 +171,8 @@ def _load_npy(self, file_name: str, mmap_mode=None, transpose: bool = False):
         -------
             The loaded .npy file.
         """
-        plane_stream_name = self.stream_name.split("_")[-1]
-        file_path = self.folder_path / plane_stream_name / file_name
+        plane_name = self.plane_name
+        file_path = self.folder_path / plane_name / file_name
         if not file_path.exists():
             return
 

From 960700713000e4610ae206a68fb191b9fbd3710b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 2 Oct 2023 17:03:36 +0000
Subject: [PATCH 10/22] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../extractors/suite2p/suite2psegmentationextractor.py    | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index 86d57e9b..3fce3be0 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -9,7 +9,7 @@
 from pathlib import Path
 from typing import Optional
 from warnings import warn
-import os 
+import os
 import numpy as np
 
 from ...extraction_tools import PathType
@@ -49,7 +49,7 @@ def get_available_channels(cls, folder_path: PathType):
             channel_names.append("chan2")
 
         return channel_names
-    
+
     @classmethod
     def get_available_planes(cls, folder_path: PathType):
         """Get the available plane names from the folder produced by Suite2p.
@@ -64,10 +64,10 @@ def get_available_planes(cls, folder_path: PathType):
         """
 
         folder_path = Path(folder_path)
-        prefix = 'plane'
+        prefix = "plane"
         plane_names = []
         for item in os.listdir(folder_path):
-            if item.startswith(prefix):   
+            if item.startswith(prefix):
                 plane_names.append(item)
 
         return plane_names

From 798994cae5a30b73a6fbfdd2f36869f79eb9b0f1 Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Tue, 3 Oct 2023 15:16:22 +0200
Subject: [PATCH 11/22] add tests

---
 tests/test_suite2psegmentationextractor.py | 99 +++++++++++++++++-----
 1 file changed, 80 insertions(+), 19 deletions(-)

diff --git a/tests/test_suite2psegmentationextractor.py b/tests/test_suite2psegmentationextractor.py
index f6822584..4c8cc457 100644
--- a/tests/test_suite2psegmentationextractor.py
+++ b/tests/test_suite2psegmentationextractor.py
@@ -1,4 +1,10 @@
+import shutil
+import tempfile
+from pathlib import Path
+
+import numpy as np
 from hdmf.testing import TestCase
+from numpy.testing import assert_array_equal
 
 from roiextractors import Suite2pSegmentationExtractor
 from tests.setup_paths import OPHYS_DATA_PATH
@@ -8,32 +14,87 @@ class TestSuite2pSegmentationExtractor(TestCase):
     @classmethod
     def setUpClass(cls):
         folder_path = str(OPHYS_DATA_PATH / "segmentation_datasets" / "suite2p")
-        cls.available_streams = dict(
-            channel_streams=["chan1", "chan2"],
-            plane_streams=dict(
-                chan1=["chan1_combined", "chan1_plane0", "chan1_plane1"], chan2=["chan2_plane0", "chan2_plane1"]
-            ),
-        )
+        cls.channel_names = ["chan1", "chan2"]
+        cls.plane_names = ["plane0", "plane1"]
 
-        cls.folder_path = folder_path
+        cls.folder_path = Path(folder_path)
 
-        extractor = Suite2pSegmentationExtractor(folder_path=folder_path, stream_name="chan1_plane1")
+        extractor = Suite2pSegmentationExtractor(folder_path=folder_path, channel_name="chan1", plane_name="plane0")
         cls.extractor = extractor
 
-    def test_stream_names(self):
-        self.assertEqual(Suite2pSegmentationExtractor.get_streams(folder_path=self.folder_path), self.available_streams)
+        cls.test_dir = Path(tempfile.mkdtemp())
 
-    def test_multi_stream_warns(self):
-        exc_msg = "More than one channel is detected! Please specify which stream you wish to load with the `stream_name` argument. To see what streams are available, call `Suite2pSegmentationExtractor.get_streams(folder_path=...)`."
-        with self.assertRaisesWith(exc_type=ValueError, exc_msg=exc_msg):
+        cls.first_channel_raw_traces = np.load(cls.folder_path / "plane0" / "F.npy").T
+        cls.second_channel_raw_traces = np.load(cls.folder_path / "plane0" / "F_chan2.npy").T
+
+    @classmethod
+    def tearDownClass(cls):
+        # remove the temporary directory and its contents
+        shutil.rmtree(cls.test_dir)
+
+    def test_channel_names(self):
+        self.assertEqual(Suite2pSegmentationExtractor.get_available_channels(folder_path=self.folder_path), self.channel_names)
+
+    def test_plane_names(self):
+        self.assertEqual(Suite2pSegmentationExtractor.get_available_planes(folder_path=self.folder_path), self.plane_names)
+
+    def test_multi_channel_warns(self):
+        exc_msg = "More than one channel is detected! Please specify which channel you wish to load with the `channel_name` argument. To see what channels are available, call `Suite2pSegmentationExtractor.get_available_channels(folder_path=...)`."
+        with self.assertWarnsWith(warn_type=UserWarning, exc_msg=exc_msg):
             Suite2pSegmentationExtractor(folder_path=self.folder_path)
 
-    def test_invalid_stream_raises(self):
-        exc_msg = "The selected stream 'plane0' is not a valid stream name. To see what streams are available, call `Suite2pSegmentationExtractor.get_streams(folder_path=...)`."
+    def test_multi_plane_warns(self):
+        exc_msg = "More than one plane is detected! Please specify which plane you wish to load with the `plane_name` argument. To see what planes are available, call `Suite2pSegmentationExtractor.get_available_planes(folder_path=...)`."
+        with self.assertWarnsWith(warn_type=UserWarning, exc_msg=exc_msg):
+            Suite2pSegmentationExtractor(folder_path=self.folder_path, channel_name="chan2")
+
+    def test_incorrect_plane_name_raises(self):
+        exc_msg = "The selected plane 'plane2' is not a valid plane name. To see what planes are available, call `Suite2pSegmentationExtractor.get_available_planes(folder_path=...)`."
         with self.assertRaisesWith(exc_type=ValueError, exc_msg=exc_msg):
-            Suite2pSegmentationExtractor(folder_path=self.folder_path, stream_name="plane0")
+            Suite2pSegmentationExtractor(folder_path=self.folder_path, plane_name="plane2")
 
-    def test_incorrect_stream_raises(self):
-        exc_msg = "The selected stream 'chan1_plane2' is not in the available plane_streams '['chan1_combined', 'chan1_plane0', 'chan1_plane1']'!"
+    def test_incorrect_channel_name_raises(self):
+        exc_msg = "The selected channel 'test' is not a valid channel name. To see what channels are available, call `Suite2pSegmentationExtractor.get_available_channels(folder_path=...)`."
         with self.assertRaisesWith(exc_type=ValueError, exc_msg=exc_msg):
-            Suite2pSegmentationExtractor(folder_path=self.folder_path, stream_name="chan1_plane2")
+            Suite2pSegmentationExtractor(folder_path=self.folder_path, channel_name="test")
+
+    def test_incomplete_extractor_load(self):
+        """Check extractor can be initialized when not all traces are available."""
+        # temporary directory for testing assertion when some of the files are missing
+        files_to_copy = ["stat.npy", "ops.npy", "iscell.npy", "Fneu.npy"]
+        (self.test_dir / "plane0").mkdir(exist_ok=True)
+        [shutil.copy(Path(self.folder_path) / "plane0" / file, self.test_dir / "plane0" / file) for file in files_to_copy]
+
+        extractor = Suite2pSegmentationExtractor(folder_path=self.test_dir)
+        traces_dict = extractor.get_traces_dict()
+        self.assertEqual(traces_dict["raw"], None)
+        self.assertEqual(traces_dict["dff"], None)
+        self.assertEqual(traces_dict["deconvolved"], None)
+
+    def test_image_size(self):
+        self.assertEqual(self.extractor.get_image_size(), (128, 128))
+
+    def test_num_frames(self):
+        self.assertEqual(self.extractor.get_num_frames(), 250)
+
+    def test_sampling_frequency(self):
+        self.assertEqual(self.extractor.get_sampling_frequency(), 10.0)
+
+    def test_channel_names(self):
+        self.assertEqual(self.extractor.get_channel_names(), ["Chan1"])
+
+    def test_num_channels(self):
+        self.assertEqual(self.extractor.get_num_channels(), 1)
+
+    def test_num_rois(self):
+        self.assertEqual(self.extractor.get_num_rois(), 15)
+
+    def test_extractor_first_channel_raw_traces(self):
+        assert_array_equal(self.extractor.get_traces(name="raw"), self.first_channel_raw_traces)
+
+    def test_extractor_second_channel(self):
+        extractor = Suite2pSegmentationExtractor(folder_path=self.folder_path, channel_name="chan2")
+        self.assertEqual(extractor.get_channel_names(), ["Chan2"])
+        traces = extractor.get_traces_dict()
+        self.assertEqual(traces["deconvolved"], None)
+        assert_array_equal(traces["raw"], self.second_channel_raw_traces)

From 7658e61e2a08cca929471d61c2848b8fc3a1f58f Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Tue, 3 Oct 2023 15:17:10 +0200
Subject: [PATCH 12/22] change to warn for multi channel multi plane

---
 .../suite2p/suite2psegmentationextractor.py   | 44 ++++++++++++-------
 1 file changed, 29 insertions(+), 15 deletions(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index 3fce3be0..ceb8ed16 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -30,7 +30,8 @@ class Suite2pSegmentationExtractor(SegmentationExtractor):
     @classmethod
     def get_available_channels(cls, folder_path: PathType):
         """Get the available channel names from the folder paths produced by Suite2p.
-        outputs
+
+        Parameters
         ----------
         file_path : PathType
             Path to Suite2p output path.
@@ -43,7 +44,7 @@ def get_available_channels(cls, folder_path: PathType):
         plane_names = cls.get_available_planes(folder_path=folder_path)
 
         channel_names = ["chan1"]
-        file_path = folder_path / plane_names[0] / "F_chan2.npy"
+        file_path = Path(folder_path) / plane_names[0] / "F_chan2.npy"
 
         if os.path.isfile(file_path):
             channel_names.append("chan2")
@@ -53,7 +54,8 @@ def get_available_channels(cls, folder_path: PathType):
     @classmethod
     def get_available_planes(cls, folder_path: PathType):
         """Get the available plane names from the folder produced by Suite2p.
-        outputs
+
+        Parameters
         ----------
         file_path : PathType
             Path to Suite2p output path.
@@ -109,22 +111,35 @@ def __init__(
         channel_names = self.get_available_channels(folder_path=folder_path)
         if channel_name is None:
             if len(channel_names) > 1:
-                raise ValueError(
+                # For backward compatibility maybe it is better to warn first
+                warn(
                     "More than one channel is detected! Please specify which channel you wish to load with the `channel_name` argument. "
-                    "To see what streams are available, call `Suite2pSegmentationExtractor.get_channel_names(folder_path=...)`."
+                    "To see what channels are available, call `Suite2pSegmentationExtractor.get_available_channels(folder_path=...)`.",
+                    UserWarning,
                 )
-            channel_name = channel_names["channel_streams"][0]
+            channel_name = channel_names[0]
 
-        if channel_name not in channel_names["channel_streams"]:
+        if channel_name not in channel_names:
             raise ValueError(
-                f"The selected channel '{channel_name}' is not a valid stream name. To see what channels are available, "
-                f"call `Suite2pSegmentationExtractor.get_channel_names(folder_path=...)`."
+                f"The selected channel '{channel_name}' is not a valid channel name. To see what channels are available, "
+                f"call `Suite2pSegmentationExtractor.get_available_channels(folder_path=...)`."
             )
 
         plane_names = self.get_available_planes(folder_path=folder_path)
-        if plane_name is not None and plane_name not in plane_names:
+        if plane_name is None:
+            if len(plane_names) > 1:
+                # For backward compatibility maybe it is better to warn first
+                warn(
+                    "More than one plane is detected! Please specify which plane you wish to load with the `plane_name` argument. "
+                    "To see what planes are available, call `Suite2pSegmentationExtractor.get_available_planes(folder_path=...)`.",
+                    UserWarning,
+                )
+            plane_name = plane_names[0]
+
+        if plane_name not in plane_names:
             raise ValueError(
-                f"The selected plane '{plane_name}' is not in the available plane_streams '{plane_names['plane_streams']}'!"
+                f"The selected plane '{plane_name}' is not a valid plane name. To see what planes are available, "
+                f"call `Suite2pSegmentationExtractor.get_available_planes(folder_path=...)`."
             )
         self.plane_name = plane_name
 
@@ -144,11 +159,11 @@ def __init__(
         neuropil_traces_file_name = "Fneu.npy" if channel_name == "chan1" else "Fneu_chan2.npy"
         self._roi_response_raw = self._load_npy(file_name=fluorescence_traces_file_name, mmap_mode="r", transpose=True)
         self._roi_response_neuropil = self._load_npy(file_name=neuropil_traces_file_name, mmap_mode="r", transpose=True)
-        self._roi_response_deconvolved = self._load_npy(file_name="spks.npy", mmap_mode="r", transpose=True)
+        self._roi_response_deconvolved = self._load_npy(file_name="spks.npy", mmap_mode="r", transpose=True) if channel_name == "chan1" else None
 
         self.iscell = self._load_npy("iscell.npy", mmap_mode="r")
 
-        channel_name = "OpticalChannel" if len(channel_names["channel_streams"]) == 1 else channel_name.capitalize()
+        channel_name = "OpticalChannel" if len(channel_names) == 1 else channel_name.capitalize()
         self._channel_names = [channel_name]
 
         self._image_correlation = self._correlation_image_read()
@@ -171,8 +186,7 @@ def _load_npy(self, file_name: str, mmap_mode=None, transpose: bool = False):
         -------
             The loaded .npy file.
         """
-        plane_name = self.plane_name
-        file_path = self.folder_path / plane_name / file_name
+        file_path = self.folder_path / self.plane_name / file_name
         if not file_path.exists():
             return
 

From 849a0f07a5361ead660e21a27ea5e5bde207ba04 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 3 Oct 2023 13:17:39 +0000
Subject: [PATCH 13/22] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../suite2p/suite2psegmentationextractor.py         |  4 +++-
 tests/test_suite2psegmentationextractor.py          | 13 ++++++++++---
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index ceb8ed16..8ed15a16 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -159,7 +159,9 @@ def __init__(
         neuropil_traces_file_name = "Fneu.npy" if channel_name == "chan1" else "Fneu_chan2.npy"
         self._roi_response_raw = self._load_npy(file_name=fluorescence_traces_file_name, mmap_mode="r", transpose=True)
         self._roi_response_neuropil = self._load_npy(file_name=neuropil_traces_file_name, mmap_mode="r", transpose=True)
-        self._roi_response_deconvolved = self._load_npy(file_name="spks.npy", mmap_mode="r", transpose=True) if channel_name == "chan1" else None
+        self._roi_response_deconvolved = (
+            self._load_npy(file_name="spks.npy", mmap_mode="r", transpose=True) if channel_name == "chan1" else None
+        )
 
         self.iscell = self._load_npy("iscell.npy", mmap_mode="r")
 
diff --git a/tests/test_suite2psegmentationextractor.py b/tests/test_suite2psegmentationextractor.py
index 4c8cc457..59853fb8 100644
--- a/tests/test_suite2psegmentationextractor.py
+++ b/tests/test_suite2psegmentationextractor.py
@@ -33,10 +33,14 @@ def tearDownClass(cls):
         shutil.rmtree(cls.test_dir)
 
     def test_channel_names(self):
-        self.assertEqual(Suite2pSegmentationExtractor.get_available_channels(folder_path=self.folder_path), self.channel_names)
+        self.assertEqual(
+            Suite2pSegmentationExtractor.get_available_channels(folder_path=self.folder_path), self.channel_names
+        )
 
     def test_plane_names(self):
-        self.assertEqual(Suite2pSegmentationExtractor.get_available_planes(folder_path=self.folder_path), self.plane_names)
+        self.assertEqual(
+            Suite2pSegmentationExtractor.get_available_planes(folder_path=self.folder_path), self.plane_names
+        )
 
     def test_multi_channel_warns(self):
         exc_msg = "More than one channel is detected! Please specify which channel you wish to load with the `channel_name` argument. To see what channels are available, call `Suite2pSegmentationExtractor.get_available_channels(folder_path=...)`."
@@ -63,7 +67,10 @@ def test_incomplete_extractor_load(self):
         # temporary directory for testing assertion when some of the files are missing
         files_to_copy = ["stat.npy", "ops.npy", "iscell.npy", "Fneu.npy"]
         (self.test_dir / "plane0").mkdir(exist_ok=True)
-        [shutil.copy(Path(self.folder_path) / "plane0" / file, self.test_dir / "plane0" / file) for file in files_to_copy]
+        [
+            shutil.copy(Path(self.folder_path) / "plane0" / file, self.test_dir / "plane0" / file)
+            for file in files_to_copy
+        ]
 
         extractor = Suite2pSegmentationExtractor(folder_path=self.test_dir)
         traces_dict = extractor.get_traces_dict()

From 115527874baf873ee34751e2c71ff5b51fedfe25 Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Tue, 3 Oct 2023 15:19:44 +0200
Subject: [PATCH 14/22] pydocstyle

---
 .../extractors/suite2p/suite2psegmentationextractor.py        | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index ceb8ed16..9f2ce962 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -35,12 +35,12 @@ def get_available_channels(cls, folder_path: PathType):
         ----------
         file_path : PathType
             Path to Suite2p output path.
+
         Returns
         -------
         channel_names: list
             List of channel names.
         """
-
         plane_names = cls.get_available_planes(folder_path=folder_path)
 
         channel_names = ["chan1"]
@@ -59,12 +59,12 @@ def get_available_planes(cls, folder_path: PathType):
         ----------
         file_path : PathType
             Path to Suite2p output path.
+
         Returns
         -------
         plane_names: list
             List of plane names.
         """
-
         folder_path = Path(folder_path)
         prefix = "plane"
         plane_names = []

From b442c8736ef0a5962af4f831be577ca33200c818 Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Thu, 5 Oct 2023 11:32:50 +0200
Subject: [PATCH 15/22] refactor get_available_planes and
 get_available_channels

---
 .../suite2p/suite2psegmentationextractor.py    | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index 845eef80..e51efc11 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -44,10 +44,10 @@ def get_available_channels(cls, folder_path: PathType):
         plane_names = cls.get_available_planes(folder_path=folder_path)
 
         channel_names = ["chan1"]
-        file_path = Path(folder_path) / plane_names[0] / "F_chan2.npy"
-
-        if os.path.isfile(file_path):
-            channel_names.append("chan2")
+        second_channel_paths = list((Path(folder_path) / plane_names[0]).glob("F_chan2.npy"))
+        if not second_channel_paths:
+            return channel_names
+        channel_names.append("chan2")
 
         return channel_names
 
@@ -65,13 +65,13 @@ def get_available_planes(cls, folder_path: PathType):
         plane_names: list
             List of plane names.
         """
+        from natsort import natsorted
+
         folder_path = Path(folder_path)
         prefix = "plane"
-        plane_names = []
-        for item in os.listdir(folder_path):
-            if item.startswith(prefix):
-                plane_names.append(item)
-
+        plane_paths = natsorted(folder_path.glob(pattern=prefix + "*"))
+        assert len(plane_paths), f"No planes found in '{folder_path}'."
+        plane_names = [plane_path.stem for plane_path in plane_paths]
         return plane_names
 
     def __init__(

From 1107f8625092f711eab28b3864316a2e9e5a5c14 Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Thu, 5 Oct 2023 11:37:14 +0200
Subject: [PATCH 16/22] fix io test for suite2p

---
 tests/test_io.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/test_io.py b/tests/test_io.py
index 43f567a5..e1cbb64b 100644
--- a/tests/test_io.py
+++ b/tests/test_io.py
@@ -139,7 +139,9 @@ def test_imaging_extractors_canonical_shape(self, extractor_class, extractor_kwa
         param(
             extractor_class=Suite2pSegmentationExtractor,
             extractor_kwargs=dict(
-                folder_path=str(OPHYS_DATA_PATH / "segmentation_datasets" / "suite2p"), stream_name="chan1_plane0"
+                folder_path=str(OPHYS_DATA_PATH / "segmentation_datasets" / "suite2p"),
+                channel_name="chan1",
+                plane_name="plane0",
             ),
         ),
     ]

From bcda76bcafe113e3115391b60a3b715d1813f85b Mon Sep 17 00:00:00 2001
From: Alessandra Trapani <alessandramaria.trapani@polimi.it>
Date: Thu, 5 Oct 2023 11:50:06 +0200
Subject: [PATCH 17/22] format code

---
 .../extractors/suite2p/suite2psegmentationextractor.py    | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index 86d57e9b..3fce3be0 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -9,7 +9,7 @@
 from pathlib import Path
 from typing import Optional
 from warnings import warn
-import os 
+import os
 import numpy as np
 
 from ...extraction_tools import PathType
@@ -49,7 +49,7 @@ def get_available_channels(cls, folder_path: PathType):
             channel_names.append("chan2")
 
         return channel_names
-    
+
     @classmethod
     def get_available_planes(cls, folder_path: PathType):
         """Get the available plane names from the folder produced by Suite2p.
@@ -64,10 +64,10 @@ def get_available_planes(cls, folder_path: PathType):
         """
 
         folder_path = Path(folder_path)
-        prefix = 'plane'
+        prefix = "plane"
         plane_names = []
         for item in os.listdir(folder_path):
-            if item.startswith(prefix):   
+            if item.startswith(prefix):
                 plane_names.append(item)
 
         return plane_names

From 1529a418652f44a349213497562bfd232518a633 Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Mon, 16 Oct 2023 15:16:42 +0200
Subject: [PATCH 18/22] add docstring for plane and channel name

---
 .../extractors/suite2p/suite2psegmentationextractor.py        | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index e51efc11..18809523 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -88,8 +88,10 @@ def __init__(
         ----------
         folder_path: str or Path
             The path to the 'suite2p' folder.
+        channel_name: str, optional
+            The name of the channel to load, to determine what channels are available use Suite2pSegmentationExtractor.get_available_channels(folder_path).
         plane_name: str, optional
-            The name of the stream to load, to determine which streams are available use Suite2pSegmentationExtractor.get_streams(folder_path).
+            The name of the plane to load, to determine what planes are available use Suite2pSegmentationExtractor.get_available_planes(folder_path).
 
         """
         if combined:

From 980d73808b981742617da9dc141cdd923497fd95 Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Tue, 17 Oct 2023 13:51:28 +0200
Subject: [PATCH 19/22] add channel name as class attribute

---
 .../extractors/suite2p/suite2psegmentationextractor.py         | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index 18809523..86d1cd65 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -121,7 +121,8 @@ def __init__(
                 )
             channel_name = channel_names[0]
 
-        if channel_name not in channel_names:
+        self.channel_name = channel_name
+        if self.channel_name not in channel_names:
             raise ValueError(
                 f"The selected channel '{channel_name}' is not a valid channel name. To see what channels are available, "
                 f"call `Suite2pSegmentationExtractor.get_available_channels(folder_path=...)`."

From 7e198f5a73b99775d5be4a66393a2b7987f98f52 Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Mon, 23 Oct 2023 15:09:48 +0200
Subject: [PATCH 20/22] add _image_masks to suite2p extractor

---
 .../suite2p/suite2psegmentationextractor.py   |  6 +++++
 tests/test_suite2psegmentationextractor.py    | 24 +++++++++++++++++--
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index 86d1cd65..2f102e49 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -174,6 +174,12 @@ def __init__(
         self._image_correlation = self._correlation_image_read()
         image_mean_name = "meanImg" if channel_name == "chan1" else f"meanImg_chan2"
         self._image_mean = self.options[image_mean_name] if image_mean_name in self.options else None
+        roi_indices = list(range(self.get_num_rois()))
+        self._image_masks = _image_mask_extractor(
+            self.get_roi_pixel_masks(),
+            roi_indices,
+            self.get_image_size(),
+        )
 
     def _load_npy(self, file_name: str, mmap_mode=None, transpose: bool = False):
         """Load a .npy file with specified filename. Returns None if file is missing.
diff --git a/tests/test_suite2psegmentationextractor.py b/tests/test_suite2psegmentationextractor.py
index 59853fb8..59dae738 100644
--- a/tests/test_suite2psegmentationextractor.py
+++ b/tests/test_suite2psegmentationextractor.py
@@ -7,6 +7,7 @@
 from numpy.testing import assert_array_equal
 
 from roiextractors import Suite2pSegmentationExtractor
+from roiextractors.extraction_tools import _image_mask_extractor
 from tests.setup_paths import OPHYS_DATA_PATH
 
 
@@ -27,6 +28,16 @@ def setUpClass(cls):
         cls.first_channel_raw_traces = np.load(cls.folder_path / "plane0" / "F.npy").T
         cls.second_channel_raw_traces = np.load(cls.folder_path / "plane0" / "F_chan2.npy").T
 
+        cls.image_size = (128, 128)
+        cls.num_rois = 15
+
+        pixel_masks = cls.extractor.get_roi_pixel_masks()
+        image_masks = np.zeros(shape=(*cls.image_size, cls.num_rois))
+        for roi_ind, pixel_mask in enumerate(pixel_masks):
+            for y, x, wt in pixel_mask:
+                image_masks[int(y), int(x), roi_ind] = wt
+        cls.image_masks = image_masks
+
     @classmethod
     def tearDownClass(cls):
         # remove the temporary directory and its contents
@@ -79,7 +90,7 @@ def test_incomplete_extractor_load(self):
         self.assertEqual(traces_dict["deconvolved"], None)
 
     def test_image_size(self):
-        self.assertEqual(self.extractor.get_image_size(), (128, 128))
+        self.assertEqual(self.extractor.get_image_size(), self.image_size)
 
     def test_num_frames(self):
         self.assertEqual(self.extractor.get_num_frames(), 250)
@@ -94,7 +105,7 @@ def test_num_channels(self):
         self.assertEqual(self.extractor.get_num_channels(), 1)
 
     def test_num_rois(self):
-        self.assertEqual(self.extractor.get_num_rois(), 15)
+        self.assertEqual(self.extractor.get_num_rois(), self.num_rois)
 
     def test_extractor_first_channel_raw_traces(self):
         assert_array_equal(self.extractor.get_traces(name="raw"), self.first_channel_raw_traces)
@@ -105,3 +116,12 @@ def test_extractor_second_channel(self):
         traces = extractor.get_traces_dict()
         self.assertEqual(traces["deconvolved"], None)
         assert_array_equal(traces["raw"], self.second_channel_raw_traces)
+
+    def test_extractor_image_masks(self):
+        """Test that the image masks are correctly extracted."""
+        assert_array_equal(self.extractor.get_roi_image_masks(), self.image_masks)
+
+    def test_extractor_image_masks_selected_rois(self):
+        """Test that the image masks are correctly extracted for a subset of ROIs."""
+        roi_indices = list(range(5))
+        assert_array_equal(self.extractor.get_roi_image_masks(roi_ids=roi_indices), self.image_masks[..., roi_indices])

From 636265d94df492467a0224445573b915c179b3c4 Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Mon, 23 Oct 2023 15:10:34 +0200
Subject: [PATCH 21/22] remove get_roi_image_masks override

---
 .../suite2p/suite2psegmentationextractor.py         | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
index 2f102e49..4ae952c0 100644
--- a/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
+++ b/src/roiextractors/extractors/suite2p/suite2psegmentationextractor.py
@@ -244,19 +244,6 @@ def roi_locations(self):
         """Returns the center locations (x, y) of each ROI."""
         return np.array([j["med"] for j in self.stat]).T.astype(int)
 
-    def get_roi_image_masks(self, roi_ids=None):
-        if roi_ids is None:
-            roi_idx_ = range(self.get_num_rois())
-        else:
-            roi_idx = [np.where(np.array(i) == self.get_roi_ids())[0] for i in roi_ids]
-            ele = [i for i, j in enumerate(roi_idx) if j.size == 0]
-            roi_idx_ = [j[0] for i, j in enumerate(roi_idx) if i not in ele]
-        return _image_mask_extractor(
-            self.get_roi_pixel_masks(roi_ids=roi_idx_),
-            list(range(len(roi_idx_))),
-            self.get_image_size(),
-        )
-
     def get_roi_pixel_masks(self, roi_ids=None):
         pixel_mask = []
         for i in range(self.get_num_rois()):

From 55589ca5292c2ef5076cfd29ea33fd7709152963 Mon Sep 17 00:00:00 2001
From: weiglszonja <wszonja@gmail.com>
Date: Tue, 31 Oct 2023 09:51:27 +0100
Subject: [PATCH 22/22] update CHANGELOG.md

---
 CHANGELOG.md | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1e28d995..974f6067 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,11 @@
 # Upcoming
 
+### Features
+
+* Updated `Suite2pSegmentationExtractor` to support multi channel and multi plane data. [PR #242](https://github.com/catalystneuro/roiextractors/pull/242)
+
+
+
 # v0.5.4
 
 ### Features