From 60fc32928073a38ba799e5693c851fc942f885f2 Mon Sep 17 00:00:00 2001
From: Ji Hoon Jeong <knowblesse@gmail.com>
Date: Mon, 1 Apr 2024 16:37:44 -0400
Subject: [PATCH] fixed further index problem with 5D data array

---
 suite2p/io/h5.py | 42 +++++++++++++++++++++++++++++++++---------
 1 file changed, 33 insertions(+), 9 deletions(-)

diff --git a/suite2p/io/h5.py b/suite2p/io/h5.py
index 079c9654..9fd653da 100644
--- a/suite2p/io/h5.py
+++ b/suite2p/io/h5.py
@@ -52,19 +52,16 @@ def h5py_to_binary(ops):
 
     for ih5, h5 in enumerate(h5list):
         with h5py.File(h5, "r") as f:
-            # if h5py data is 5D or 4D instead of 3D, assume that
-            # data = (nchan x) (nframes x) nplanes x pixels x pixels
             # 5D/4D data is flattened to process the same way as interleaved data
             for key in keys:
                 hdims = f[key].ndim
                 # keep track of the plane identity of the first frame (channel identity is assumed always 0)
                 ncp = nplanes * nchannels
                 nbatch = ncp * math.ceil(ops1[0]["batch_size"] / ncp)
-                nframes_all = f[key].shape[
-                    0] if hdims == 3 else f[key].shape[0] * f[key].shape[1]
+                nframes_all = np.prod(f[key].shape[:-2]) # product except for the last two dim
                 nbatch = min(nbatch, nframes_all)
                 nfunc = ops["functional_chan"] - 1 if nchannels > 1 else 0
-                # loop over all tiffs
+                # loop over all frames
                 ik = 0
                 while 1:
                     if hdims == 3:
@@ -78,23 +75,50 @@ def h5py_to_binary(ops):
                             1)
                         if irange.size == 0:
                             break
-                        im = f[key][irange, ...]
-                        if im.ndim == 5 and im.shape[0] == nchannels:
+                        if hdims == 4:
+                            # for 4D, expect
+                            # data = nframes x nplanes x pixels x pixels or
+                            # data = nframes x nchans x pixels x pixels
+                            # Data sanity check
+                            if not (
+                                    ((f[key].shape[1] == nchannels) and (nplanes == 1)) or
+                                    ((f[key].shape[1] == nplanes) and (nchannels == 1))):
+                                raise ValueError("Channel/Plane dimension mismatch.\n"
+                                                 "h5 file should follow the following dimension\n"
+                                                 "nframes x nplanes x pixels x pixels or"
+                                                 "nframes x nchans x pixels x pixels")
+                            im = f[key][irange, ...]
+                        elif hdims == 5:
+                            # for 5D, expect
+                            # data = nchans x nframes x nplanes x pixels x pixels
+                            # Data sanity check
+                            if not (
+                                    (f[key].shape[0] == nchannels) and
+                                    (f[key].shape[2] == nplanes)):
+                                raise ValueError("Channel/Plane dimension mismatch.\n"
+                                                 "h5 file should follow the following dimension\n"
+                                                 "nchans x nframes x nplanes x pixels x pixels")
+                            im = f[key][:, irange, :, :, :]
                             im = im.transpose((1, 0, 2, 3, 4))
                         # flatten to frames x pixels x pixels
+                        # for 5D data with 2 channels, 2 planes, this results
+                        # [fr0 ch0 pl0], [fr0 ch0 pl1], [fr0 ch1 pl0],
+                        # [fr0 ch1 pl1] [fr1 ch0 pl0], ...
                         im = np.reshape(im, (-1, im.shape[-2], im.shape[-1]))
                     nframes = im.shape[0]
-                    if type(im[0, 0, 0]) == np.uint16:
+                    if im.dtype == np.uint16:
                         im = im / 2
                     for j in range(0, nplanes):
                         if iall == 0:
                             ops1[j]["meanImg"] = np.zeros((im.shape[1], im.shape[2]),
                                                           np.float32)
                             if nchannels > 1:
+                                # TODO: either show error when nchannel > 2 or
+                                # implement generalized func. for nchannel > 2
                                 ops1[j]["meanImg_chan2"] = np.zeros(
                                     (im.shape[1], im.shape[2]), np.float32)
                             ops1[j]["nframes"] = 0
-                        i0 = nchannels * ((j) % nplanes)
+                        i0 = j
                         im2write = im[np.arange(int(i0) +
                                                 nfunc, nframes, ncp), :, :].astype(
                                                     np.int16)