From 3d92fb12ea9c68d4f63c16bc7b94a4edc11eb030 Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@vathes.com>
Date: Thu, 10 Aug 2023 16:37:07 -0500
Subject: [PATCH 01/16] Update prairie_view_loader.py

---
 element_interface/prairie_view_loader.py | 172 ++++++++++++++++-------
 1 file changed, 120 insertions(+), 52 deletions(-)

diff --git a/element_interface/prairie_view_loader.py b/element_interface/prairie_view_loader.py
index 841e87a..ee306cd 100644
--- a/element_interface/prairie_view_loader.py
+++ b/element_interface/prairie_view_loader.py
@@ -1,65 +1,90 @@
 import pathlib
+from pathlib import Path
 import xml.etree.ElementTree as ET
 from datetime import datetime
-
 import numpy as np
 
 
-def get_prairieview_metadata(ome_tif_filepath: str) -> dict:
-    """Extract metadata for scans generated by Prairie View acquisition software.
+class PrairieViewMeta:
 
-    The Prairie View software generates one `.ome.tif` imaging file per frame
-    acquired. The metadata for all frames is contained in one .xml file. This
-    function locates the .xml file and generates a dictionary necessary to
-    populate the DataJoint `ScanInfo` and `Field` tables. Prairie View works
-    with resonance scanners with a single field. Prairie View does not support
-    bidirectional x and y scanning. ROI information is not contained in the
-    `.xml` file. All images generated using Prairie View have square dimensions(e.g. 512x512).
+    def __init__(self, prairieview_dir: str):
+        """Initialize PrairieViewMeta loader class
 
-    Args:
-        ome_tif_filepath: An absolute path to the .ome.tif image file.
+        Args:
+            prairieview_dir (str): string, absolute file path to directory containing PrairieView dataset
+        """
+        # ---- Search and verify CaImAn output file exists ----
+        # May return multiple xml files. Only need one that contains scan metadata.
+        self.prairieview_dir = Path(prairieview_dir)
 
-    Raises:
-        FileNotFoundError: No .xml file containing information about the acquired scan
-            was found at path in parent directory at `ome_tif_filepath`.
+        for file in self.prairieview_dir.glob("*.xml"):
+            xml_tree = ET.parse(file)
+            xml_root = xml_tree.getroot()
+            if xml_root.find(".//Sequence"):
+                self.xml_file = file
+                self._xml_root = xml_root
+                break
+        else:
+            raise FileNotFoundError(
+                f"No PrarieView metadata .xml file found at {prairieview_dir}"
+            )
 
-    Returns:
-        metainfo: A dict mapping keys to corresponding metadata values fetched from the
-            .xml file.
-    """
+        self._meta = None
 
-    # May return multiple xml files. Only need one that contains scan metadata.
-    xml_files_list = pathlib.Path(ome_tif_filepath).parent.glob("*.xml")
+    @property
+    def meta(self):
+        if self._meta is None:
+            self._meta = _extract_prairieview_metadata(self.xml_file)
+        return self._meta
 
-    for file in xml_files_list:
-        xml_tree = ET.parse(file)
-        xml_file = xml_tree.getroot()
-        if xml_file.find(".//Sequence"):
-            break
-    else:
-        raise FileNotFoundError(
-            f"No PrarieView metadata .xml file found at {pathlib.Path(ome_tif_filepath).parent}"
-        )
+    def get_prairieview_files(self, plane_idx=None, channel=None):
+        if plane_idx is None:
+            if self.meta['num_planes'] > 1:
+                raise ValueError(f"Please specify 'plane_idx' - Plane indices: {self.meta['plane_indices']}")
+            else:
+                plane_idx = self.meta['plane_indices'][0]
+        else:
+            assert plane_idx in self.meta['plane_indices'], f"Invalid 'plane_idx' - Plane indices: {self.meta['plane_indices']}"
+
+        if channel is None:
+            if self.meta['num_channels'] > 1:
+                raise ValueError(f"Please specify 'channel' - Channels: {self.meta['channels']}")
+            else:
+                plane_idx = self.meta['channels'][0]
+        else:
+            assert channel in self.meta['channels'], f"Invalid 'channel' - Channels: {self.meta['channels']}"
+
+        frames = self._xml_root.findall(f".//Sequence/Frame/[@index='{plane_idx}']/File/[@channel='{channel}']")
+        return [f.attrib['filename'] for f in frames]
+
+
+def _extract_prairieview_metadata(xml_filepath: str):
+    xml_filepath = Path(xml_filepath)
+    if not xml_filepath.exists():
+        raise FileNotFoundError(f"{xml_filepath} does not exist")
+    xml_tree = ET.parse(xml_filepath)
+    xml_root = xml_tree.getroot()
 
     bidirectional_scan = False  # Does not support bidirectional
     roi = 0
     n_fields = 1  # Always contains 1 field
-    recording_start_time = xml_file.find(".//Sequence/[@cycle='1']").attrib.get("time")
+    recording_start_time = xml_root.find(".//Sequence/[@cycle='1']").attrib.get("time")
 
     # Get all channels and find unique values
     channel_list = [
         int(channel.attrib.get("channel"))
-        for channel in xml_file.iterfind(".//Sequence/Frame/File/[@channel]")
+        for channel in xml_root.iterfind(".//Sequence/Frame/File/[@channel]")
     ]
-    n_channels = len(set(channel_list))
-    n_frames = len(xml_file.findall(".//Sequence/Frame"))
+    channels = set(channel_list)
+    n_channels = len(channels)
+    n_frames = len(xml_root.findall(".//Sequence/Frame"))
     framerate = 1 / float(
-        xml_file.findall('.//PVStateValue/[@key="framePeriod"]')[0].attrib.get("value")
+        xml_root.findall('.//PVStateValue/[@key="framePeriod"]')[0].attrib.get("value")
     )  # rate = 1/framePeriod
 
     usec_per_line = (
         float(
-            xml_file.findall(".//PVStateValue/[@key='scanLinePeriod']")[0].attrib.get(
+            xml_root.findall(".//PVStateValue/[@key='scanLinePeriod']")[0].attrib.get(
                 "value"
             )
         )
@@ -67,15 +92,15 @@ def get_prairieview_metadata(ome_tif_filepath: str) -> dict:
     )  # Convert from seconds to microseconds
 
     scan_datetime = datetime.strptime(
-        xml_file.attrib.get("date"), "%m/%d/%Y %I:%M:%S %p"
+        xml_root.attrib.get("date"), "%m/%d/%Y %I:%M:%S %p"
     )
 
     total_scan_duration = float(
-        xml_file.findall(".//Sequence/Frame")[-1].attrib.get("relativeTime")
+        xml_root.findall(".//Sequence/Frame")[-1].attrib.get("relativeTime")
     )
 
     pixel_height = int(
-        xml_file.findall(".//PVStateValue/[@key='pixelsPerLine']")[0].attrib.get(
+        xml_root.findall(".//PVStateValue/[@key='pixelsPerLine']")[0].attrib.get(
             "value"
         )
     )
@@ -83,7 +108,7 @@ def get_prairieview_metadata(ome_tif_filepath: str) -> dict:
     pixel_width = pixel_height
 
     um_per_pixel = float(
-        xml_file.find(
+        xml_root.find(
             ".//PVStateValue/[@key='micronsPerPixel']/IndexedValue/[@index='XAxis']"
         ).attrib.get("value")
     )
@@ -92,43 +117,45 @@ def get_prairieview_metadata(ome_tif_filepath: str) -> dict:
 
     # x and y coordinate values for the center of the field
     x_field = float(
-        xml_file.find(
+        xml_root.find(
             ".//PVStateValue/[@key='currentScanCenter']/IndexedValue/[@index='XAxis']"
         ).attrib.get("value")
     )
     y_field = float(
-        xml_file.find(
+        xml_root.find(
             ".//PVStateValue/[@key='currentScanCenter']/IndexedValue/[@index='YAxis']"
         ).attrib.get("value")
     )
+
     if (
-        xml_file.find(
+        xml_root.find(
             ".//Sequence/[@cycle='1']/Frame/PVStateShard/PVStateValue/[@key='positionCurrent']/SubindexedValues/[@index='ZAxis']"
         )
         is None
     ):
         z_fields = np.float64(
-            xml_file.find(
+            xml_root.find(
                 ".//PVStateValue/[@key='positionCurrent']/SubindexedValues/[@index='ZAxis']/SubindexedValue"
             ).attrib.get("value")
         )
         n_depths = 1
+        plane_indices = {0}
         assert z_fields.size == n_depths
         bidirection_z = False
-
     else:
         bidirection_z = (
-            xml_file.find(".//Sequence").attrib.get("bidirectionalZ") == "True"
+            xml_root.find(".//Sequence").attrib.get("bidirectionalZ") == "True"
         )
 
         # One "Frame" per depth in the .xml file. Gets number of frames in first sequence
         planes = [
             int(plane.attrib.get("index"))
-            for plane in xml_file.findall(".//Sequence/[@cycle='1']/Frame")
+            for plane in xml_root.findall(".//Sequence/[@cycle='1']/Frame")
         ]
-        n_depths = len(set(planes))
+        plane_indices = set(planes)
+        n_depths = len(plane_indices)
 
-        z_controllers = xml_file.findall(
+        z_controllers = xml_root.findall(
             ".//Sequence/[@cycle='1']/Frame/[@index='1']/PVStateShard/PVStateValue/[@key='positionCurrent']/SubindexedValues/[@index='ZAxis']/SubindexedValue"
         )
 
@@ -137,13 +164,13 @@ def get_prairieview_metadata(ome_tif_filepath: str) -> dict:
         # must change depths.
         if len(z_controllers) > 1:
             z_repeats = []
-            for controller in xml_file.findall(
+            for controller in xml_root.findall(
                 ".//Sequence/[@cycle='1']/Frame/[@index='1']/PVStateShard/PVStateValue/[@key='positionCurrent']/SubindexedValues/[@index='ZAxis']/"
             ):
                 z_repeats.append(
                     [
                         float(z.attrib.get("value"))
-                        for z in xml_file.findall(
+                        for z in xml_root.findall(
                             ".//Sequence/[@cycle='1']/Frame/PVStateShard/PVStateValue/[@key='positionCurrent']/SubindexedValues/[@index='ZAxis']/SubindexedValue/[@subindex='{0}']".format(
                                 controller.attrib.get("subindex")
                             )
@@ -163,7 +190,7 @@ def get_prairieview_metadata(ome_tif_filepath: str) -> dict:
         else:
             z_fields = [
                 z.attrib.get("value")
-                for z in xml_file.findall(
+                for z in xml_root.findall(
                     ".//Sequence/[@cycle='1']/Frame/PVStateShard/PVStateValue/[@key='positionCurrent']/SubindexedValues/[@index='ZAxis']/SubindexedValue/[@subindex='0']"
                 )
             ]
@@ -195,6 +222,47 @@ def get_prairieview_metadata(ome_tif_filepath: str) -> dict:
         fieldY=y_field,
         fieldZ=z_fields,
         recording_time=recording_start_time,
+        channels=list(channels),
+        plane_indices=list(plane_indices),
     )
 
     return metainfo
+
+
+def get_prairieview_metadata(ome_tif_filepath: str) -> dict:
+    """Extract metadata for scans generated by Prairie View acquisition software.
+
+    The Prairie View software generates one `.ome.tif` imaging file per frame
+    acquired. The metadata for all frames is contained in one .xml file. This
+    function locates the .xml file and generates a dictionary necessary to
+    populate the DataJoint `ScanInfo` and `Field` tables. Prairie View works
+    with resonance scanners with a single field. Prairie View does not support
+    bidirectional x and y scanning. ROI information is not contained in the
+    `.xml` file. All images generated using Prairie View have square dimensions(e.g. 512x512).
+
+    Args:
+        ome_tif_filepath: An absolute path to the .ome.tif image file.
+
+    Raises:
+        FileNotFoundError: No .xml file containing information about the acquired scan
+            was found at path in parent directory at `ome_tif_filepath`.
+
+    Returns:
+        metainfo: A dict mapping keys to corresponding metadata values fetched from the
+            .xml file.
+    """
+
+    # May return multiple xml files. Only need one that contains scan metadata.
+    xml_files_list = pathlib.Path(ome_tif_filepath).parent.glob("*.xml")
+
+    for file in xml_files_list:
+        xml_tree = ET.parse(file)
+        xml_file = xml_tree.getroot()
+        if xml_file.find(".//Sequence"):
+            break
+    else:
+        raise FileNotFoundError(
+            f"No PrarieView metadata .xml file found at {pathlib.Path(ome_tif_filepath).parent}"
+        )
+
+    return _extract_prairieview_metadata(file)

From da03b823139d316edbb5a32d31407f8982bbdcbb Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@vathes.com>
Date: Wed, 16 Aug 2023 18:33:03 -0500
Subject: [PATCH 02/16] support loading multiple caiman results for multi-plane

---
 element_interface/caiman_loader.py | 128 ++++++++++++++++++++++++++++-
 1 file changed, 124 insertions(+), 4 deletions(-)

diff --git a/element_interface/caiman_loader.py b/element_interface/caiman_loader.py
index 3726afd..3ae75c0 100644
--- a/element_interface/caiman_loader.py
+++ b/element_interface/caiman_loader.py
@@ -1,7 +1,7 @@
 import os
 import pathlib
 from datetime import datetime
-
+import re
 import caiman as cm
 import h5py
 import numpy as np
@@ -56,6 +56,119 @@ class CaImAn:
         segmentation_channel: hard-coded to 0
     """
 
+    def __init__(self, caiman_dir: str):
+        """Initialize CaImAn loader class
+
+        Args:
+            caiman_dir (str): string, absolute file path to CaIman directory
+
+        Raises:
+            FileNotFoundError: No CaImAn analysis output file found
+            FileNotFoundError: No CaImAn analysis output found, missing required fields
+        """
+        # ---- Search and verify CaImAn output file exists ----
+        caiman_dir = pathlib.Path(caiman_dir)
+        if not caiman_dir.exists():
+            raise FileNotFoundError("CaImAn directory not found: {}".format(caiman_dir))
+
+        caiman_subdirs = []
+        for fp in caiman_dir.glob("*.hdf5"):
+            with h5py.File(fp, "r") as h5f:
+                if all(s in h5f for s in _required_hdf5_fields):
+                    caiman_subdirs.append(fp.parent)
+
+        if not caiman_subdirs:
+            raise FileNotFoundError(
+                "No CaImAn analysis output file found at {}"
+                " containg all required fields ({})".format(
+                    caiman_dir, _required_hdf5_fields
+                )
+            )
+
+        self.planes = {}
+        for idx, caiman_subdir in enumerate(sorted(caiman_subdirs)):
+            pln_cm = _CaImAn(caiman_subdir.as_posix())
+            pln_idx_match = re.search(r"pln(\d+)_.*")
+            pln_idx = pln_idx_match.groups()[0] if pln_idx_match else idx
+            pln_cm.plane_idx = pln_idx
+            self.planes[pln_idx] = pln_cm
+
+        self._motion_correction = None
+        self._masks = None
+
+        self.creation_time = min(
+            [p.creation_time for p in self.planes.values()]
+        )  # ealiest file creation time
+        self.curation_time = max(
+            [p.curation_time for p in self.planes.values()]
+        )  # most recent curation time
+
+    @property
+    def motion_correction(self):
+        if self._motion_correction is None:
+            self._motion_correction = self.h5f["motion_correction"]
+        return self._motion_correction
+
+    @property
+    def masks(self):
+        if self._masks is None:
+            all_masks = []
+            for pln_idx, _caiman in sorted(self.planes.items()):
+                mask_count = len(all_masks)  # increment mask id from all "plane"
+                all_masks.extend([{**m, "mask_id": m["mask_id"] + mask_count} for m in _caiman.masks])
+
+            self._masks = all_masks
+        return self._masks
+
+    @property
+    def alignment_channel(self):
+        return 0  # hard-code to channel index 0
+
+    @property
+    def segmentation_channel(self):
+        return 0  # hard-code to channel index 0
+
+
+class _CaImAn:
+    """Parse the CaImAn output file
+
+    [CaImAn results doc](https://caiman.readthedocs.io/en/master/Getting_Started.html#result-variables-for-2p-batch-analysis)
+
+    Expecting the following objects:
+        - dims:
+        - dview:
+        - estimates:              Segmentations and traces
+        - mmap_file:
+        - params:                 Input parameters
+        - remove_very_bad_comps:
+        - skip_refinement:
+        - motion_correction:      Motion correction shifts and summary images
+
+    Example:
+        > output_dir = '<imaging_root_data_dir>/subject1/session0/caiman'
+
+        > loaded_dataset = caiman_loader.CaImAn(output_dir)
+
+    Attributes:
+        alignment_channel: hard-coded to 0
+        caiman_fp: file path with all required files:
+            "/motion_correction/reference_image",
+            "/motion_correction/correlation_image",
+            "/motion_correction/average_image",
+            "/motion_correction/max_image",
+            "/estimates/A",
+        cnmf: loaded caiman object; cm.source_extraction.cnmf.cnmf.load_CNMF(caiman_fp)
+        creation_time: file creation time
+        curation_time: file creation time
+        extract_masks: function to extract masks
+        h5f: caiman_fp read as h5py file
+        masks: dict result of extract_masks
+        motion_correction: h5f "motion_correction" property
+        params: cnmf.params
+        segmentation_channel: hard-coded to 0
+        plane_idx: N/A if `is3D` else hard-coded to 0
+    """
+
     def __init__(self, caiman_dir: str):
         """Initialize CaImAn loader class
 
@@ -89,13 +202,20 @@ def __init__(self, caiman_dir: str):
         self.params = self.cnmf.params
 
         self.h5f = h5py.File(self.caiman_fp, "r")
-        self.motion_correction = self.h5f["motion_correction"]
+        self.plane_idx = None if self.params.motion["is3D"] else 0
+        self._motion_correction = None
         self._masks = None
 
         # ---- Metainfo ----
         self.creation_time = datetime.fromtimestamp(os.stat(self.caiman_fp).st_ctime)
         self.curation_time = datetime.fromtimestamp(os.stat(self.caiman_fp).st_ctime)
 
+    @property
+    def motion_correction(self):
+        if self._motion_correction is None:
+            self._motion_correction = self.h5f["motion_correction"]
+        return self._motion_correction
+
     @property
     def masks(self):
         if self._masks is None:
@@ -139,7 +259,7 @@ def extract_masks(self) -> dict:
             else:
                 xpix, ypix = np.unravel_index(ind, self.cnmf.dims, order="F")
                 center_x, center_y = comp_contour["CoM"].astype(int)
-                center_z = 0
+                center_z = self.plane_idx
                 zpix = np.full(len(weights), center_z)
 
             masks.append(
@@ -161,7 +281,7 @@ def extract_masks(self) -> dict:
         return masks
 
 
-def _process_scanimage_tiff(scan_filenames, output_dir="./"):
+def _process_scanimage_tiff(scan_filenames, output_dir="./", split_depths=False):
     """
     Read ScanImage TIFF - reshape into volumetric data based on scanning depths/channels
     Save new TIFF files for each channel - with shape (frame x height x width x depth)

From 9bf69ff8f33af327e8e10450a1d32071b993183e Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@vathes.com>
Date: Sat, 19 Aug 2023 21:11:52 -0500
Subject: [PATCH 03/16] new caiman loader - handles multi-plane results

---
 element_interface/caiman_loader.py | 210 ++++++++++++++++++++++++++++-
 1 file changed, 204 insertions(+), 6 deletions(-)

diff --git a/element_interface/caiman_loader.py b/element_interface/caiman_loader.py
index 3ae75c0..0b90f38 100644
--- a/element_interface/caiman_loader.py
+++ b/element_interface/caiman_loader.py
@@ -93,9 +93,6 @@ def __init__(self, caiman_dir: str):
             pln_cm.plane_idx = pln_idx
             self.planes[pln_idx] = pln_cm
 
-        self._motion_correction = None
-        self._masks = None
-
         self.creation_time = min(
             [p.creation_time for p in self.planes.values()]
         )  # ealiest file creation time
@@ -103,19 +100,179 @@ def __init__(self, caiman_dir: str):
             [p.curation_time for p in self.planes.values()]
         )  # most recent curation time
 
+        # is this 3D CaImAn analyis or multiple 2D per-plane analysis
+        if len(self.planes) > 1:
+            # if more than one set of caiman result, likely to be multiple 2D per-plane
+            # assert that the "is3D" value are all False for each of the caiman result
+            assert all(p.params.motion["is3D"] is False for p in self.planes.values())
+            self.is3D = False
+            self.is_multiplane = True
+        else:
+            self.is3D = list(self.planes.values())[0].params.motion["is3D"]
+            self.is_multiplane = False
+
+        self._motion_correction = None
+        self._masks = None
+        self._ref_image = None
+        self._mean_image = None
+        self._max_proj_image = None
+        self._correlation_map = None
+
     @property
     def motion_correction(self):
         if self._motion_correction is None:
-            self._motion_correction = self.h5f["motion_correction"]
+            pass
         return self._motion_correction
 
+    def extract_rigid_mc(self):
+        # -- rigid motion correction --
+        rigid_correction = {}
+        for pln_idx, (plane, pln_cm) in enumerate(self.planes.items()):
+            if pln_idx == 0:
+                rigid_correction = {
+                    "x_shifts": pln_cm.motion_correction["shifts_rig"][:, 0],
+                    "y_shifts": pln_cm.motion_correction["shifts_rig"][:, 1],
+                }
+                rigid_correction["x_std"] = np.nanstd(
+                    rigid_correction["x_shifts"].flatten()
+                )
+                rigid_correction["y_std"] = np.nanstd(
+                    rigid_correction["y_shifts"].flatten()
+                )
+            else:
+                rigid_correction["x_shifts"] = np.vstack(
+                    [
+                        rigid_correction["x_shifts"],
+                        pln_cm.motion_correction["shifts_rig"][:, 0],
+                    ]
+                )
+                rigid_correction["x_std"] = np.nanstd(
+                    rigid_correction["x_shifts"].flatten()
+                )
+                rigid_correction["y_shifts"] = np.vstack(
+                    [
+                        rigid_correction["y_shifts"],
+                        pln_cm.motion_correction["shifts_rig"][:, 1],
+                    ]
+                )
+                rigid_correction["y_std"] = np.nanstd(
+                    rigid_correction["y_shifts"].flatten()
+                )
+
+        if not self.is_multiplane:
+            pln_cm = list(self.planes.values())[0]
+            rigid_correction["z_shifts"] = (
+                pln_cm.motion_correction["shifts_rig"][:, 2]
+                if self.is3D
+                else np.full_like(rigid_correction["x_shifts"], 0)
+            )
+            rigid_correction["z_std"] = (
+                np.nanstd(pln_cm.motion_correction["shifts_rig"][:, 2])
+                if self.is3D
+                else np.nan
+            )
+        else:
+            rigid_correction["z_shifts"] = np.full_like(rigid_correction["x_shifts"], 0)
+            rigid_correction["z_std"] = np.nan
+
+        rigid_correction["outlier_frames"] = None
+
+        return rigid_correction
+
+    def extract_pw_rigid_mc(self):
+        # -- piece-wise rigid motion correction --
+        nonrigid_correction, nonrigid_blocks = {}
+        for pln_idx, (plane, pln_cm) in enumerate(self.planes.items()):
+            if pln_idx == 0:
+                nonrigid_correction = {
+                    "block_height": (
+                        pln_cm.params.motion["strides"][0]
+                        + pln_cm.params.motion["overlaps"][0]
+                    ),
+                    "block_width": (
+                        pln_cm.params.motion["strides"][1]
+                        + pln_cm.params.motion["overlaps"][1]
+                    ),
+                    "block_depth": 1,
+                    "block_count_x": len(
+                        set(pln_cm.motion_correction["coord_shifts_els"][:, 0])
+                    ),
+                    "block_count_y": len(
+                        set(pln_cm.motion_correction["coord_shifts_els"][:, 2])
+                    ),
+                    "block_count_z": len(self.planes),
+                    "outlier_frames": None,
+                }
+            for b_id in range(len(pln_cm.motion_correction["x_shifts_els"][0, :])):
+                if b_id in nonrigid_blocks:
+                    nonrigid_blocks[b_id]["x_shifts"] = np.vstack(
+                        [
+                            nonrigid_blocks[b_id]["x_shifts"],
+                            pln_cm.motion_correction["x_shifts_els"][:, b_id],
+                        ]
+                    )
+                    nonrigid_blocks[b_id]["x_std"] = np.nanstd(
+                        nonrigid_blocks[b_id]["x_shifts"].flatten()
+                    )
+                    nonrigid_blocks[b_id]["y_shifts"] = np.vstack(
+                        [
+                            nonrigid_blocks[b_id]["y_shifts"],
+                            pln_cm.motion_correction["y_shifts_els"][:, b_id],
+                        ]
+                    )
+                    nonrigid_blocks[b_id]["y_std"] = np.nanstd(
+                        nonrigid_blocks[b_id]["y_shifts"].flatten()
+                    )
+                    nonrigid_blocks[b_id]["z_shifts"] = np.vstack(
+                        [
+                            nonrigid_blocks[b_id]["z_shifts"],
+                            np.full_like(
+                                pln_cm.motion_correction["x_shifts_els"][:, b_id],
+                                0,
+                            ),
+                        ]
+                    )
+                else:
+                    nonrigid_blocks[b_id] = {
+                        "block_id": b_id,
+                        "block_x": np.arange(
+                            *pln_cm.motion_correction["coord_shifts_els"][b_id, 0:2]
+                        ),
+                        "block_y": np.arange(
+                            *pln_cm.motion_correction["coord_shifts_els"][b_id, 2:4]
+                        ),
+                        "block_z": np.full_like(
+                            np.arange(
+                                *pln_cm.motion_correction["coord_shifts_els"][b_id, 0:2]
+                            ),
+                            pln_idx,
+                        ),
+                        "x_shifts": pln_cm.motion_correction["x_shifts_els"][:, b_id],
+                        "y_shifts": pln_cm.motion_correction["y_shifts_els"][:, b_id],
+                        "z_shifts": np.full_like(
+                            pln_cm.motion_correction["x_shifts_els"][:, b_id],
+                            0,
+                        ),
+                        "x_std": np.nanstd(
+                            pln_cm.motion_correction["x_shifts_els"][:, b_id]
+                        ),
+                        "y_std": np.nanstd(
+                            pln_cm.motion_correction["y_shifts_els"][:, b_id]
+                        ),
+                        "z_std": np.nan,
+                    }
+
+        return nonrigid_correction, nonrigid_blocks
+
     @property
     def masks(self):
         if self._masks is None:
             all_masks = []
-            for pln_idx, _caiman in sorted(self.planes.items()):
+            for pln_idx, pln_cm in sorted(self.planes.items()):
                 mask_count = len(all_masks)  # increment mask id from all "plane"
-                all_masks.extend([{**m, "mask_id": m["mask_id"] + mask_count} for m in _caiman.masks])
+                all_masks.extend(
+                    [{**m, "mask_id": m["mask_id"] + mask_count} for m in pln_cm.masks]
+                )
 
             self._masks = all_masks
         return self._masks
@@ -128,6 +285,47 @@ def alignment_channel(self):
     def segmentation_channel(self):
         return 0  # hard-code to channel index 0
 
+    # -- image property --
+
+    def _get_image(self, img_type):
+        if not self.is_multiplane:
+            pln_cm = list(self.planes.values())[0]
+            _img = (
+                pln_cm.motion_correction[img_type].transpose()
+                if self.is3D
+                else pln_cm.motion_correction[img_type][...][np.newaxis, ...]
+            )
+        else:
+            _img = np.dstack(
+                pln_cm.motion_correction[img_type][...]
+                for pln_cm in self.planes.values()
+            )
+        return _img
+
+    @property
+    def ref_image(self):
+        if self._ref_image is None:
+            self._ref_image = self._get_image("reference_image")
+        return self._ref_image
+
+    @property
+    def mean_image(self):
+        if self._mean_image is None:
+            self._mean_image = self._get_image("average_image")
+        return self._mean_image
+
+    @property
+    def max_proj_image(self):
+        if self._max_proj_image is None:
+            self._max_proj_image = self._get_image("max_image")
+        return self._max_proj_image
+
+    @property
+    def correlation_map(self):
+        if self._correlation_map is None:
+            self._correlation_map = self._get_image("correlation_image")
+        return self._correlation_map
+
 
 class _CaImAn:
     """Parse the CaImAn output file

From 05d7c3544665a9345b1732abc43d0311fadff143 Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@vathes.com>
Date: Sat, 19 Aug 2023 21:18:34 -0500
Subject: [PATCH 04/16] improve non-rigid motion correction loading

---
 element_interface/caiman_loader.py | 103 +++++++++++++++--------------
 1 file changed, 52 insertions(+), 51 deletions(-)

diff --git a/element_interface/caiman_loader.py b/element_interface/caiman_loader.py
index 0b90f38..2e73c3d 100644
--- a/element_interface/caiman_loader.py
+++ b/element_interface/caiman_loader.py
@@ -111,6 +111,11 @@ def __init__(self, caiman_dir: str):
             self.is3D = list(self.planes.values())[0].params.motion["is3D"]
             self.is_multiplane = False
 
+        if self.is_multiplane and self.is3D:
+            raise NotImplementedError(
+                f"Unable to load CaImAn results mixed between 3D and multi-plane analysis"
+            )
+
         self._motion_correction = None
         self._masks = None
         self._ref_image = None
@@ -183,6 +188,7 @@ def extract_pw_rigid_mc(self):
         # -- piece-wise rigid motion correction --
         nonrigid_correction, nonrigid_blocks = {}
         for pln_idx, (plane, pln_cm) in enumerate(self.planes.items()):
+            block_count = len(nonrigid_blocks)
             if pln_idx == 0:
                 nonrigid_correction = {
                     "block_height": (
@@ -204,63 +210,58 @@ def extract_pw_rigid_mc(self):
                     "outlier_frames": None,
                 }
             for b_id in range(len(pln_cm.motion_correction["x_shifts_els"][0, :])):
-                if b_id in nonrigid_blocks:
-                    nonrigid_blocks[b_id]["x_shifts"] = np.vstack(
-                        [
-                            nonrigid_blocks[b_id]["x_shifts"],
-                            pln_cm.motion_correction["x_shifts_els"][:, b_id],
-                        ]
-                    )
-                    nonrigid_blocks[b_id]["x_std"] = np.nanstd(
-                        nonrigid_blocks[b_id]["x_shifts"].flatten()
-                    )
-                    nonrigid_blocks[b_id]["y_shifts"] = np.vstack(
-                        [
-                            nonrigid_blocks[b_id]["y_shifts"],
-                            pln_cm.motion_correction["y_shifts_els"][:, b_id],
-                        ]
-                    )
-                    nonrigid_blocks[b_id]["y_std"] = np.nanstd(
-                        nonrigid_blocks[b_id]["y_shifts"].flatten()
-                    )
-                    nonrigid_blocks[b_id]["z_shifts"] = np.vstack(
-                        [
-                            nonrigid_blocks[b_id]["z_shifts"],
-                            np.full_like(
-                                pln_cm.motion_correction["x_shifts_els"][:, b_id],
-                                0,
-                            ),
-                        ]
-                    )
-                else:
-                    nonrigid_blocks[b_id] = {
-                        "block_id": b_id,
-                        "block_x": np.arange(
-                            *pln_cm.motion_correction["coord_shifts_els"][b_id, 0:2]
-                        ),
-                        "block_y": np.arange(
-                            *pln_cm.motion_correction["coord_shifts_els"][b_id, 2:4]
-                        ),
-                        "block_z": np.full_like(
+                b_id += block_count
+                nonrigid_blocks[b_id] = {
+                    "block_id": b_id,
+                    "block_x": np.arange(
+                        *pln_cm.motion_correction["coord_shifts_els"][b_id, 0:2]
+                    ),
+                    "block_y": np.arange(
+                        *pln_cm.motion_correction["coord_shifts_els"][b_id, 2:4]
+                    ),
+                    "block_z": (
+                        np.arange(
+                            *pln_cm.motion_correction["coord_shifts_els"][b_id, 4:6]
+                        )
+                        if self.is3D
+                        else np.full_like(
                             np.arange(
                                 *pln_cm.motion_correction["coord_shifts_els"][b_id, 0:2]
                             ),
                             pln_idx,
-                        ),
-                        "x_shifts": pln_cm.motion_correction["x_shifts_els"][:, b_id],
-                        "y_shifts": pln_cm.motion_correction["y_shifts_els"][:, b_id],
-                        "z_shifts": np.full_like(
+                        )
+                    ),
+                    "x_shifts": pln_cm.motion_correction["x_shifts_els"][:, b_id],
+                    "y_shifts": pln_cm.motion_correction["y_shifts_els"][:, b_id],
+                    "z_shifts": (
+                        pln_cm.motion_correction["z_shifts_els"][:, b_id]
+                        if self.is3D
+                        else np.full_like(
                             pln_cm.motion_correction["x_shifts_els"][:, b_id],
                             0,
-                        ),
-                        "x_std": np.nanstd(
-                            pln_cm.motion_correction["x_shifts_els"][:, b_id]
-                        ),
-                        "y_std": np.nanstd(
-                            pln_cm.motion_correction["y_shifts_els"][:, b_id]
-                        ),
-                        "z_std": np.nan,
-                    }
+                        )
+                    ),
+                    "x_std": np.nanstd(
+                        pln_cm.motion_correction["x_shifts_els"][:, b_id]
+                    ),
+                    "y_std": np.nanstd(
+                        pln_cm.motion_correction["y_shifts_els"][:, b_id]
+                    ),
+                    "z_std": (
+                        np.nanstd(pln_cm.motion_correction["z_shifts_els"][:, b_id])
+                        if self.is3D
+                        else np.nan
+                    ),
+                }
+
+        if not self.is_multiplane and self.is3D:
+            pln_cm = list(self.planes.values())[0]
+            nonrigid_correction["block_depth"] = (
+                pln_cm.params.motion["strides"][2] + pln_cm.params.motion["overlaps"][2]
+            )
+            nonrigid_correction["block_count_z"] = len(
+                set(pln_cm.motion_correction["coord_shifts_els"][:, 4])
+            )
 
         return nonrigid_correction, nonrigid_blocks
 

From b050265753080f2c7632b7b8a0a0734bbd239b19 Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@vathes.com>
Date: Tue, 22 Aug 2023 09:22:27 -0500
Subject: [PATCH 05/16] added routine to combine single-page tiffs into one
 bigtiff

---
 element_interface/caiman_loader.py       |  6 +-
 element_interface/prairie_view_loader.py | 74 +++++++++++++++++++-----
 2 files changed, 64 insertions(+), 16 deletions(-)

diff --git a/element_interface/caiman_loader.py b/element_interface/caiman_loader.py
index 2e73c3d..efdfc05 100644
--- a/element_interface/caiman_loader.py
+++ b/element_interface/caiman_loader.py
@@ -291,17 +291,17 @@ def segmentation_channel(self):
     def _get_image(self, img_type):
         if not self.is_multiplane:
             pln_cm = list(self.planes.values())[0]
-            _img = (
+            img_ = (
                 pln_cm.motion_correction[img_type].transpose()
                 if self.is3D
                 else pln_cm.motion_correction[img_type][...][np.newaxis, ...]
             )
         else:
-            _img = np.dstack(
+            img_ = np.dstack(
                 pln_cm.motion_correction[img_type][...]
                 for pln_cm in self.planes.values()
             )
-        return _img
+        return img_
 
     @property
     def ref_image(self):
diff --git a/element_interface/prairie_view_loader.py b/element_interface/prairie_view_loader.py
index ee306cd..746829b 100644
--- a/element_interface/prairie_view_loader.py
+++ b/element_interface/prairie_view_loader.py
@@ -1,12 +1,13 @@
+import os
 import pathlib
 from pathlib import Path
 import xml.etree.ElementTree as ET
 from datetime import datetime
 import numpy as np
+import tifffile
 
 
 class PrairieViewMeta:
-
     def __init__(self, prairieview_dir: str):
         """Initialize PrairieViewMeta loader class
 
@@ -37,25 +38,72 @@ def meta(self):
             self._meta = _extract_prairieview_metadata(self.xml_file)
         return self._meta
 
-    def get_prairieview_files(self, plane_idx=None, channel=None):
+    def get_prairieview_filenames(
+        self, plane_idx=None, channel=None, return_pln_chn=False
+    ):
+        """
+        Extract from metadata the set of tiff files specific to the specified "plane_idx" and "channel"
+        Args:
+            plane_idx: int - plane index
+            channel: int - channel
+            return_pln_chn: bool - if True, returns (filenames, plane_idx, channel), else returns `filenames`
+
+        Returns: List[str] - the set of tiff files specific to the specified "plane_idx" and "channel"
+        """
         if plane_idx is None:
-            if self.meta['num_planes'] > 1:
-                raise ValueError(f"Please specify 'plane_idx' - Plane indices: {self.meta['plane_indices']}")
+            if self.meta["num_planes"] > 1:
+                raise ValueError(
+                    f"Please specify 'plane_idx' - Plane indices: {self.meta['plane_indices']}"
+                )
             else:
-                plane_idx = self.meta['plane_indices'][0]
+                plane_idx = self.meta["plane_indices"][0]
         else:
-            assert plane_idx in self.meta['plane_indices'], f"Invalid 'plane_idx' - Plane indices: {self.meta['plane_indices']}"
+            assert (
+                plane_idx in self.meta["plane_indices"]
+            ), f"Invalid 'plane_idx' - Plane indices: {self.meta['plane_indices']}"
 
         if channel is None:
-            if self.meta['num_channels'] > 1:
-                raise ValueError(f"Please specify 'channel' - Channels: {self.meta['channels']}")
+            if self.meta["num_channels"] > 1:
+                raise ValueError(
+                    f"Please specify 'channel' - Channels: {self.meta['channels']}"
+                )
             else:
-                plane_idx = self.meta['channels'][0]
+                plane_idx = self.meta["channels"][0]
         else:
-            assert channel in self.meta['channels'], f"Invalid 'channel' - Channels: {self.meta['channels']}"
+            assert (
+                channel in self.meta["channels"]
+            ), f"Invalid 'channel' - Channels: {self.meta['channels']}"
+
+        frames = self._xml_root.findall(
+            f".//Sequence/Frame/[@index='{plane_idx}']/File/[@channel='{channel}']"
+        )
 
-        frames = self._xml_root.findall(f".//Sequence/Frame/[@index='{plane_idx}']/File/[@channel='{channel}']")
-        return [f.attrib['filename'] for f in frames]
+        fnames = [f.attrib["filename"] for f in frames]
+        return fnames if not return_pln_chn else (fnames, plane_idx, channel)
+
+    def write_single_tiff(
+        self, plane_idx=None, channel=None, output_prefix=None, output_dir="./"
+    ):
+        tiff_names, plane_idx, channel = self.get_prairieview_filenames(
+            plane_idx=plane_idx, channel=channel, return_pln_chn=True
+        )
+        combined_data = []
+        for input_file in tiff_names:
+            with tifffile.TiffFile(self.prairieview_dir / input_file) as tffl:
+                assert len(tffl.pages) == 1
+                combined_data.append(tffl.asarray())
+        combined_data = np.dstack(combined_data).transpose(
+            2, 0, 1
+        )  # (frame x height x width)
+
+        if output_prefix is None:
+            output_prefix = os.path.commonprefix(tiff_names)
+
+        tifffile.imwrite(
+            Path(output_dir) / f"{output_prefix}_pln{plane_idx}_chn{channel}",
+            combined_data,
+            metadata={"axes": "TXY", "'fps'": self.meta["frame_rate"]},
+        )
 
 
 def _extract_prairieview_metadata(xml_filepath: str):
@@ -159,7 +207,7 @@ def _extract_prairieview_metadata(xml_filepath: str):
             ".//Sequence/[@cycle='1']/Frame/[@index='1']/PVStateShard/PVStateValue/[@key='positionCurrent']/SubindexedValues/[@index='ZAxis']/SubindexedValue"
         )
 
-        # If more than one Z-axis controllers are found, 
+        # If more than one Z-axis controllers are found,
         # check which controller is changing z_field depth. Only 1 controller
         # must change depths.
         if len(z_controllers) > 1:

From 512e5afee28be7ddb9d0baf9b86480a916f62b40 Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@vathes.com>
Date: Tue, 22 Aug 2023 14:42:29 -0500
Subject: [PATCH 06/16] run_caimain - more robust

---
 element_interface/run_caiman.py | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/element_interface/run_caiman.py b/element_interface/run_caiman.py
index eb480a9..6cb1559 100644
--- a/element_interface/run_caiman.py
+++ b/element_interface/run_caiman.py
@@ -41,12 +41,19 @@ def run_caiman(
         backend="local", n_processes=None, single_thread=False
     )
 
-    cnm = CNMF(n_processes, params=opts, dview=dview)
-    cnmf_output, mc_output = cnm.fit_file(
-        motion_correct=True, include_eval=True, output_dir=output_dir, return_mc=True
-    )
-
-    cm.stop_server(dview=dview)
+    try:
+        cnm = CNMF(n_processes, params=opts, dview=dview)
+        cnmf_output, mc_output = cnm.fit_file(
+            motion_correct=True,
+            include_eval=True,
+            output_dir=output_dir,
+            return_mc=True,
+        )
+    except Exception as e:
+        dview.terminate()
+        raise e
+    else:
+        cm.stop_server(dview=dview)
 
     cnmf_output_file = pathlib.Path(cnmf_output.mmap_file[:-4] + "hdf5")
     assert cnmf_output_file.exists()

From 7d6ea0e4bc8af336394c47e9703a8dc0d71550d9 Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@vathes.com>
Date: Tue, 22 Aug 2023 14:42:40 -0500
Subject: [PATCH 07/16] add `caiman_compatible` mode

---
 element_interface/prairie_view_loader.py | 60 +++++++++++++++++-------
 1 file changed, 44 insertions(+), 16 deletions(-)

diff --git a/element_interface/prairie_view_loader.py b/element_interface/prairie_view_loader.py
index 746829b..5c4f06e 100644
--- a/element_interface/prairie_view_loader.py
+++ b/element_interface/prairie_view_loader.py
@@ -68,7 +68,7 @@ def get_prairieview_filenames(
                     f"Please specify 'channel' - Channels: {self.meta['channels']}"
                 )
             else:
-                plane_idx = self.meta["channels"][0]
+                channel = self.meta["channels"][0]
         else:
             assert (
                 channel in self.meta["channels"]
@@ -82,28 +82,56 @@ def get_prairieview_filenames(
         return fnames if not return_pln_chn else (fnames, plane_idx, channel)
 
     def write_single_tiff(
-        self, plane_idx=None, channel=None, output_prefix=None, output_dir="./"
+        self,
+        plane_idx=None,
+        channel=None,
+        output_prefix=None,
+        output_dir="./",
+        caiman_compatible=False,  # if True, save the movie as a single page (frame x height x width)
+        overwrite=False,
     ):
         tiff_names, plane_idx, channel = self.get_prairieview_filenames(
             plane_idx=plane_idx, channel=channel, return_pln_chn=True
         )
-        combined_data = []
-        for input_file in tiff_names:
-            with tifffile.TiffFile(self.prairieview_dir / input_file) as tffl:
-                assert len(tffl.pages) == 1
-                combined_data.append(tffl.asarray())
-        combined_data = np.dstack(combined_data).transpose(
-            2, 0, 1
-        )  # (frame x height x width)
-
         if output_prefix is None:
             output_prefix = os.path.commonprefix(tiff_names)
-
-        tifffile.imwrite(
-            Path(output_dir) / f"{output_prefix}_pln{plane_idx}_chn{channel}",
-            combined_data,
-            metadata={"axes": "TXY", "'fps'": self.meta["frame_rate"]},
+        output_tiff_fullpath = (
+            Path(output_dir)
+            / f"{output_prefix}_pln{plane_idx}_chn{channel}{'.ome' if not caiman_compatible else ''}.tif"
         )
+        if output_tiff_fullpath.exists() and not overwrite:
+            return output_tiff_fullpath
+
+        if not caiman_compatible:
+            with tifffile.TiffWriter(
+                output_tiff_fullpath,
+                bigtiff=True,
+            ) as tiff_writer:
+                for input_file in tiff_names:
+                    with tifffile.TiffFile(self.prairieview_dir / input_file) as tffl:
+                        assert len(tffl.pages) == 1
+                        tiff_writer.write(
+                            tffl.pages[0].asarray(),
+                            metadata={"axes": "YX", "'fps'": self.meta["frame_rate"]},
+                        )
+        else:
+            combined_data = []
+            for input_file in tiff_names:
+                with tifffile.TiffFile(self.prairieview_dir / input_file) as tffl:
+                    assert len(tffl.pages) == 1
+                    combined_data.append(tffl.pages[0].asarray())
+            combined_data = np.dstack(combined_data).transpose(
+                2, 0, 1
+            )  # (frame x height x width)
+
+            tifffile.imwrite(
+                output_tiff_fullpath,
+                combined_data,
+                metadata={"axes": "TYX", "'fps'": self.meta["frame_rate"]},
+                bigtiff=True,
+            )
+
+        return output_tiff_fullpath
 
 
 def _extract_prairieview_metadata(xml_filepath: str):

From 55e5a5c3fa50ebcdedc1f660e64f9a9b7d98a09e Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@vathes.com>
Date: Thu, 24 Aug 2023 12:13:05 -0500
Subject: [PATCH 08/16] minor cleanup, version bump

---
 element_interface/prairie_view_loader.py | 2 +-
 element_interface/version.py             | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/element_interface/prairie_view_loader.py b/element_interface/prairie_view_loader.py
index 5c4f06e..1a05139 100644
--- a/element_interface/prairie_view_loader.py
+++ b/element_interface/prairie_view_loader.py
@@ -81,7 +81,7 @@ def get_prairieview_filenames(
         fnames = [f.attrib["filename"] for f in frames]
         return fnames if not return_pln_chn else (fnames, plane_idx, channel)
 
-    def write_single_tiff(
+    def write_single_bigtiff(
         self,
         plane_idx=None,
         channel=None,
diff --git a/element_interface/version.py b/element_interface/version.py
index 0da8726..70aab85 100644
--- a/element_interface/version.py
+++ b/element_interface/version.py
@@ -1,3 +1,3 @@
 """Package metadata"""
 
-__version__ = "0.6.1"
+__version__ = "0.7.0"

From cb08034edfa1e94f509513010e802b01bb431c31 Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@vathes.com>
Date: Fri, 25 Aug 2023 09:56:32 -0500
Subject: [PATCH 09/16] Update caiman_loader.py

---
 element_interface/caiman_loader.py | 75 ++++++++++++++----------------
 1 file changed, 35 insertions(+), 40 deletions(-)

diff --git a/element_interface/caiman_loader.py b/element_interface/caiman_loader.py
index efdfc05..580e194 100644
--- a/element_interface/caiman_loader.py
+++ b/element_interface/caiman_loader.py
@@ -18,42 +18,10 @@
 
 
 class CaImAn:
-    """Parse the CaImAn output file
-
-    [CaImAn results doc](https://caiman.readthedocs.io/en/master/Getting_Started.html#result-variables-for-2p-batch-analysis)
-
-    Expecting the following objects:
-        - dims:
-        - dview:
-        - estimates:              Segmentations and traces
-        - mmap_file:
-        - params:                 Input parameters
-        - remove_very_bad_comps:
-        - skip_refinement:
-        - motion_correction:      Motion correction shifts and summary images
-
-    Example:
-        > output_dir = '<imaging_root_data_dir>/subject1/session0/caiman'
-
-        > loaded_dataset = caiman_loader.CaImAn(output_dir)
-
-    Attributes:
-        alignment_channel: hard-coded to 0
-        caiman_fp: file path with all required files:
-            "/motion_correction/reference_image",
-            "/motion_correction/correlation_image",
-            "/motion_correction/average_image",
-            "/motion_correction/max_image",
-            "/estimates/A",
-        cnmf: loaded caiman object; cm.source_extraction.cnmf.cnmf.load_CNMF(caiman_fp)
-        creation_time: file creation time
-        curation_time: file creation time
-        extract_masks: function to extract masks
-        h5f: caiman_fp read as h5py file
-        masks: dict result of extract_masks
-        motion_correction: h5f "motion_correction" property
-        params: cnmf.params
-        segmentation_channel: hard-coded to 0
+    """
+    Loader class for CaImAn analysis results
+    A top level aggregator of multiple set of CaImAn results (e.g. multi-plane analysis)
+    Calling _CaImAn (see below) under the hood
     """
 
     def __init__(self, caiman_dir: str):
@@ -85,13 +53,16 @@ def __init__(self, caiman_dir: str):
                 )
             )
 
-        self.planes = {}
+        # Extract CaImAn results from all planes, sorted by plane index
+        _planes_caiman = {}
         for idx, caiman_subdir in enumerate(sorted(caiman_subdirs)):
             pln_cm = _CaImAn(caiman_subdir.as_posix())
             pln_idx_match = re.search(r"pln(\d+)_.*")
             pln_idx = pln_idx_match.groups()[0] if pln_idx_match else idx
             pln_cm.plane_idx = pln_idx
-            self.planes[pln_idx] = pln_cm
+            _planes_caiman[pln_idx] = pln_cm
+        sorted_pln_ind = sorted(list(_planes_caiman.keys()))
+        self.planes = {k: _planes_caiman[k] for k in sorted_pln_ind}
 
         self.creation_time = min(
             [p.creation_time for p in self.planes.values()]
@@ -123,10 +94,22 @@ def __init__(self, caiman_dir: str):
         self._max_proj_image = None
         self._correlation_map = None
 
+    @property
+    def is_pw_rigid(self):
+        pw_rigid = set(p.params.motion["pw_rigid"] for p in self.planes.values())
+        assert (
+            len(pw_rigid) == 1
+        ), f"Unable to load CaImAn results mixed between rigid and pw_rigid motion correction"
+        return pw_rigid.pop()
+
     @property
     def motion_correction(self):
         if self._motion_correction is None:
-            pass
+            self._motion_correction = (
+                self.extract_pw_rigid_mc()
+                if self.is_pw_rigid
+                else self.extract_rigid_mc()
+            )
         return self._motion_correction
 
     def extract_rigid_mc(self):
@@ -272,7 +255,19 @@ def masks(self):
             for pln_idx, pln_cm in sorted(self.planes.items()):
                 mask_count = len(all_masks)  # increment mask id from all "plane"
                 all_masks.extend(
-                    [{**m, "mask_id": m["mask_id"] + mask_count} for m in pln_cm.masks]
+                    [
+                        {
+                            **m,
+                            "mask_id": m["mask_id"] + mask_count,
+                            "orig_mask_id": m["mask_id"],
+                            "accepted": (
+                                m["mask_id"] in pln_cm.cnmf.estimates.idx_components
+                                if pln_cm.cnmf.estimates.idx_components is not None
+                                else False
+                            ),
+                        }
+                        for m in pln_cm.masks
+                    ]
                 )
 
             self._masks = all_masks

From 85078fdbe7dc71b1885e94ee93a460898094ffd6 Mon Sep 17 00:00:00 2001
From: kushalbakshi <kushalbakshi@gmail.com>
Date: Mon, 28 Aug 2023 11:22:58 -0500
Subject: [PATCH 10/16] Fix regex error

---
 element_interface/caiman_loader.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/element_interface/caiman_loader.py b/element_interface/caiman_loader.py
index 580e194..1743aa7 100644
--- a/element_interface/caiman_loader.py
+++ b/element_interface/caiman_loader.py
@@ -57,7 +57,7 @@ def __init__(self, caiman_dir: str):
         _planes_caiman = {}
         for idx, caiman_subdir in enumerate(sorted(caiman_subdirs)):
             pln_cm = _CaImAn(caiman_subdir.as_posix())
-            pln_idx_match = re.search(r"pln(\d+)_.*")
+            pln_idx_match = re.search(r"pln(\d+)_.*", caiman_subdir)
             pln_idx = pln_idx_match.groups()[0] if pln_idx_match else idx
             pln_cm.plane_idx = pln_idx
             _planes_caiman[pln_idx] = pln_cm

From 7b815ff4f8a8d710c02556aee5581f579956d96e Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@vathes.com>
Date: Mon, 28 Aug 2023 12:30:46 -0500
Subject: [PATCH 11/16] bugfix

---
 element_interface/caiman_loader.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/element_interface/caiman_loader.py b/element_interface/caiman_loader.py
index 580e194..667df5a 100644
--- a/element_interface/caiman_loader.py
+++ b/element_interface/caiman_loader.py
@@ -57,7 +57,7 @@ def __init__(self, caiman_dir: str):
         _planes_caiman = {}
         for idx, caiman_subdir in enumerate(sorted(caiman_subdirs)):
             pln_cm = _CaImAn(caiman_subdir.as_posix())
-            pln_idx_match = re.search(r"pln(\d+)_.*")
+            pln_idx_match = re.search(r"pln(\d+)_.*", caiman_subdir.stem)
             pln_idx = pln_idx_match.groups()[0] if pln_idx_match else idx
             pln_cm.plane_idx = pln_idx
             _planes_caiman[pln_idx] = pln_cm

From 0bf78ed542bc8a9dba572adcf3b8d54eca985a8f Mon Sep 17 00:00:00 2001
From: kushalbakshi <kushalbakshi@gmail.com>
Date: Mon, 28 Aug 2023 12:58:40 -0500
Subject: [PATCH 12/16] Resolve TypeError: Add `.stem()`

---
 element_interface/caiman_loader.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/element_interface/caiman_loader.py b/element_interface/caiman_loader.py
index 1743aa7..667df5a 100644
--- a/element_interface/caiman_loader.py
+++ b/element_interface/caiman_loader.py
@@ -57,7 +57,7 @@ def __init__(self, caiman_dir: str):
         _planes_caiman = {}
         for idx, caiman_subdir in enumerate(sorted(caiman_subdirs)):
             pln_cm = _CaImAn(caiman_subdir.as_posix())
-            pln_idx_match = re.search(r"pln(\d+)_.*", caiman_subdir)
+            pln_idx_match = re.search(r"pln(\d+)_.*", caiman_subdir.stem)
             pln_idx = pln_idx_match.groups()[0] if pln_idx_match else idx
             pln_cm.plane_idx = pln_idx
             _planes_caiman[pln_idx] = pln_cm

From 706df0484d53dbc2aca1a29622090d1d66879e96 Mon Sep 17 00:00:00 2001
From: kushalbakshi <kushalbakshi@gmail.com>
Date: Mon, 28 Aug 2023 17:20:43 -0500
Subject: [PATCH 13/16] Debug `FileNotFoundError` at output_dir

---
 element_interface/caiman_loader.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/element_interface/caiman_loader.py b/element_interface/caiman_loader.py
index 667df5a..a87d9eb 100644
--- a/element_interface/caiman_loader.py
+++ b/element_interface/caiman_loader.py
@@ -40,7 +40,7 @@ def __init__(self, caiman_dir: str):
             raise FileNotFoundError("CaImAn directory not found: {}".format(caiman_dir))
 
         caiman_subdirs = []
-        for fp in caiman_dir.glob("*.hdf5"):
+        for fp in caiman_dir.rglob("*.hdf5"):
             with h5py.File(fp, "r") as h5f:
                 if all(s in h5f for s in _required_hdf5_fields):
                     caiman_subdirs.append(fp.parent)

From 701a5a443f3565b6c72737706007f61552b8c714 Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@vathes.com>
Date: Wed, 30 Aug 2023 19:00:16 -0500
Subject: [PATCH 14/16] subprocess `shell` as input argument

---
 element_interface/dandi.py | 27 +++++++++++++++++++--------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/element_interface/dandi.py b/element_interface/dandi.py
index 01553f4..dbaeee9 100644
--- a/element_interface/dandi.py
+++ b/element_interface/dandi.py
@@ -13,6 +13,7 @@ def upload_to_dandi(
     api_key: str = None,
     sync: bool = False,
     existing: str = "refresh",
+    shell=True,  # without this param, subprocess interprets first arg as file/dir
 ):
     """Upload NWB files to DANDI Archive
 
@@ -38,25 +39,35 @@ def upload_to_dandi(
         working_directory, str(dandiset_id)
     )  # enforce str
 
-    dandiset_url = f"https://gui-staging.dandiarchive.org/#/dandiset/{dandiset_id}" if staging else f"https://dandiarchive.org/dandiset/{dandiset_id}/draft"
+    dandiset_url = (
+        f"https://gui-staging.dandiarchive.org/#/dandiset/{dandiset_id}"
+        if staging
+        else f"https://dandiarchive.org/dandiset/{dandiset_id}/draft"
+    )
 
     subprocess.run(
-        ["dandi", "download", "--download", "dandiset.yaml", "-o", working_directory, dandiset_url],
-        shell=True, 
+        [
+            "dandi",
+            "download",
+            "--download",
+            "dandiset.yaml",
+            "-o",
+            working_directory,
+            dandiset_url,
+        ],
+        shell=shell,
     )
 
     subprocess.run(
         ["dandi", "organize", "-d", dandiset_directory, data_directory, "-f", "dry"],
-        shell=True,  # without this param, subprocess interprets first arg as file/dir
+        shell=shell,  # without this param, subprocess interprets first arg as file/dir
     )
 
     subprocess.run(
-        ["dandi", "organize", "-d", dandiset_directory, data_directory], shell=True
+        ["dandi", "organize", "-d", dandiset_directory, data_directory], shell=shell
     )
 
-    subprocess.run(
-        ["dandi", "validate", dandiset_directory], shell=True
-    )
+    subprocess.run(["dandi", "validate", dandiset_directory], shell=shell)
 
     upload(
         paths=[dandiset_directory],

From 29bc09b6d46a7cd40bae17e6962b6d074dccb2d6 Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@vathes.com>
Date: Wed, 30 Aug 2023 19:29:20 -0500
Subject: [PATCH 15/16] version pin dandi

---
 requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/requirements.txt b/requirements.txt
index 75d95e8..b18b774 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,2 @@
-dandi
+dandi>=0.56.0
 numpy

From 26c38b6c43cdb282df48637890e148a715a00a6a Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@vathes.com>
Date: Sun, 3 Sep 2023 11:21:31 -0500
Subject: [PATCH 16/16] Update dandi.py

---
 element_interface/dandi.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/element_interface/dandi.py b/element_interface/dandi.py
index dbaeee9..078e58a 100644
--- a/element_interface/dandi.py
+++ b/element_interface/dandi.py
@@ -1,7 +1,6 @@
 import os
 import subprocess
 
-from dandi.download import download
 from dandi.upload import upload
 
 
@@ -64,7 +63,18 @@ def upload_to_dandi(
     )
 
     subprocess.run(
-        ["dandi", "organize", "-d", dandiset_directory, data_directory], shell=shell
+        [
+            "dandi",
+            "organize",
+            "-d",
+            dandiset_directory,
+            data_directory,
+            "--required-field",
+            "subject_id",
+            "--required-field",
+            "session_id",
+        ],
+        shell=shell,
     )
 
     subprocess.run(["dandi", "validate", dandiset_directory], shell=shell)