Fix motion correction and CNMF for 3D movies

environment.yml

@@ -1,7 +1,7 @@
 channels:
   - conda-forge
 dependencies:
-  - caiman >= 1.9.10
+  - caiman >= 1.11.2
   - pandas >= 1.5.0
   - requests
   - click

environment_rtd.yml

@@ -1,7 +1,7 @@
 channels:
   - conda-forge
 dependencies:
-  - caiman >= 1.9.10
+  - caiman >= 1.11.2
   - pandas >= 1.5.0
   - requests
   - click

mesmerize_core/algorithms/cnmf.py

@@ -99,7 +99,7 @@ def run_algo(batch_path, uuid, data_path: str = None):
 
         cnm.save(str(output_path))
 
-        Cn = cm.local_correlations(images.transpose(1, 2, 0))
+        Cn = cm.local_correlations(images, swap_dim=False)
         Cn[np.isnan(Cn)] = 0
 
         corr_img_path = output_dir.joinpath(f"{uuid}_cn.npy").resolve()

mesmerize_core/algorithms/mcorr.py

@@ -90,7 +90,7 @@ def run_algo(batch_path, uuid, data_path: str = None):
 
         print("Computing correlation image")
         Cns = local_correlations_movie_offline(
-            [str(mcorr_memmap_path)],
+            str(mcorr_memmap_path),
             remove_baseline=True,
             window=1000,
             stride=1000,
@@ -102,24 +102,27 @@ def run_algo(batch_path, uuid, data_path: str = None):
         Cn[np.isnan(Cn)] = 0
         cn_path = output_dir.joinpath(f"{uuid}_cn.npy")
         np.save(str(cn_path), Cn, allow_pickle=False)
-
-        # output dict for pandas series for dataframe row
-        d = dict()
-
+
         print("finished computing correlation image")
 
+
         # Compute shifts
         if opts.motion["pw_rigid"] == True:
             x_shifts = mc.x_shifts_els
             y_shifts = mc.y_shifts_els
             shifts = [x_shifts, y_shifts]
+            if hasattr(mc, 'z_shifts_els'):
+                shifts += mc.z_shifts_els
             shift_path = output_dir.joinpath(f"{uuid}_shifts.npy")
             np.save(str(shift_path), shifts)
         else:
             shifts = mc.shifts_rig
             shift_path = output_dir.joinpath(f"{uuid}_shifts.npy")
             np.save(str(shift_path), shifts)
 
+        # output dict for pandas series for dataframe row
+        d = dict()
+
         # save paths as relative path strings with forward slashes
         cn_path = str(PurePosixPath(cn_path.relative_to(output_dir.parent)))
         mcorr_memmap_path = str(PurePosixPath(mcorr_memmap_path.relative_to(output_dir.parent)))