From 07070eca4586a14af8a9f84335cf9734838f67b6 Mon Sep 17 00:00:00 2001
From: MathiasMahn <53939819+MathiasMahn@users.noreply.github.com>
Date: Mon, 9 Feb 2026 17:45:43 +0100
Subject: [PATCH 1/4] Refactor motion correction grid computation logic

Refactor grid computation logic for motion correction to ensure it only attempts to build a patch grid when both strides and overlaps are provided. Added fallback for rigid correction to define grid based on full FOV.
---
 element_interface/caiman_loader.py | 52 +++++++++++++++---------------
 1 file changed, 26 insertions(+), 26 deletions(-)

diff --git a/element_interface/caiman_loader.py b/element_interface/caiman_loader.py
index be3dd8a..abf08c6 100644
--- a/element_interface/caiman_loader.py
+++ b/element_interface/caiman_loader.py
@@ -558,34 +558,34 @@ def _save_mc(
     # Load the first frame of the movie
     mc_image = np.reshape(Yr[: np.product(dims), :1], [1] + list(dims), order="F")
 
-    # Compute mc.coord_shifts_els
+    # Safe compute of mc.coord_shifts_els or mc.coord_shifts_rig
     grid = []
-    if is3D:
-        for _, _, _, x, y, z, _ in cm.motion_correction.sliding_window_3d(
-            mc_image[0, :, :, :], mc.overlaps, mc.strides
-        ):
-            grid.append(
-                [
-                    x,
-                    x + mc.overlaps[0] + mc.strides[0],
-                    y,
-                    y + mc.overlaps[1] + mc.strides[1],
-                    z,
-                    z + mc.overlaps[2] + mc.strides[2],
-                ]
-            )
+    # Only attempt to build a patch grid if both strides and overlaps are provided
+    if mc.strides is not None and mc.overlaps is not None:
+        if is3D:
+            for _, _, _, x, y, z, _ in cm.motion_correction.sliding_window_3d(
+                mc_image[0, :, :, :], mc.overlaps, mc.strides
+            ):
+                grid.append([
+                    x, x + mc.overlaps[0] + mc.strides[0],
+                    y, y + mc.overlaps[1] + mc.strides[1],
+                    z, z + mc.overlaps[2] + mc.strides[2],
+                ])
+        else:
+            for _, _, x, y, _ in cm.motion_correction.sliding_window(
+                mc_image[0, :, :], mc.overlaps, mc.strides
+            ):
+                grid.append([
+                    x, x + mc.overlaps[0] + mc.strides[0],
+                    y, y + mc.overlaps[1] + mc.strides[1],
+                ])
     else:
-        for _, _, x, y, _ in cm.motion_correction.sliding_window(
-            mc_image[0, :, :], mc.overlaps, mc.strides
-        ):
-            grid.append(
-                [
-                    x,
-                    x + mc.overlaps[0] + mc.strides[0],
-                    y,
-                    y + mc.overlaps[1] + mc.strides[1],
-                ]
-            )
+        # Fallback for Rigid correction: The grid is simply the full FOV
+        if is3D:
+            grid = [[0, dims[0], 0, dims[1], 0, dims[2]]]
+        else:
+            grid = [[0, dims[0], 0, dims[1]]]
+
 
     # Open hdf5 file and create 'motion_correction' group
     caiman_fp = pathlib.Path(caiman_fp)

From db9eb157c88295cbdb633c29c787c940288a50fc Mon Sep 17 00:00:00 2001
From: MathiasMahn <53939819+MathiasMahn@users.noreply.github.com>
Date: Mon, 9 Feb 2026 20:27:42 +0100
Subject: [PATCH 2/4] pw_rigid grid calculation only when flag == true

---
 element_interface/caiman_loader.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/element_interface/caiman_loader.py b/element_interface/caiman_loader.py
index abf08c6..c5f4d1f 100644
--- a/element_interface/caiman_loader.py
+++ b/element_interface/caiman_loader.py
@@ -594,6 +594,35 @@ def _save_mc(
 
     # Write motion correction shifts and motion corrected summary images to hdf5 file
     if mc.pw_rigid:
+        # Compute mc.coord_shifts_els
+        grid = []
+        if is3D:
+            for _, _, _, x, y, z, _ in cm.motion_correction.sliding_window_3d(
+                mc_image[0, :, :, :], mc.overlaps, mc.strides
+            ):
+                grid.append(
+                    [
+                        x,
+                        x + mc.overlaps[0] + mc.strides[0],
+                        y,
+                        y + mc.overlaps[1] + mc.strides[1],
+                        z,
+                        z + mc.overlaps[2] + mc.strides[2],
+                    ]
+                )
+        else:
+            for _, _, x, y, _ in cm.motion_correction.sliding_window(
+                mc_image[0, :, :], mc.overlaps, mc.strides
+            ):
+                grid.append(
+                    [
+                        x,
+                        x + mc.overlaps[0] + mc.strides[0],
+                        y,
+                        y + mc.overlaps[1] + mc.strides[1],
+                    ]
+                )
+        
         h5g.require_dataset(
             "x_shifts_els",
             shape=np.shape(mc.x_shifts_els),

From cec3887e7103bbcb456140477374e2f9b517a9de Mon Sep 17 00:00:00 2001
From: MathiasMahn <53939819+MathiasMahn@users.noreply.github.com>
Date: Mon, 9 Feb 2026 20:29:54 +0100
Subject: [PATCH 3/4] moved the patch grid computation in caiman_loader under
 if pw_rigid == true

moved the patch grid computation for motion correction based on strides and overlaps under if pw_rigid == true
---
 element_interface/caiman_loader.py | 29 -----------------------------
 1 file changed, 29 deletions(-)

diff --git a/element_interface/caiman_loader.py b/element_interface/caiman_loader.py
index c5f4d1f..f15b665 100644
--- a/element_interface/caiman_loader.py
+++ b/element_interface/caiman_loader.py
@@ -558,35 +558,6 @@ def _save_mc(
     # Load the first frame of the movie
     mc_image = np.reshape(Yr[: np.product(dims), :1], [1] + list(dims), order="F")
 
-    # Safe compute of mc.coord_shifts_els or mc.coord_shifts_rig
-    grid = []
-    # Only attempt to build a patch grid if both strides and overlaps are provided
-    if mc.strides is not None and mc.overlaps is not None:
-        if is3D:
-            for _, _, _, x, y, z, _ in cm.motion_correction.sliding_window_3d(
-                mc_image[0, :, :, :], mc.overlaps, mc.strides
-            ):
-                grid.append([
-                    x, x + mc.overlaps[0] + mc.strides[0],
-                    y, y + mc.overlaps[1] + mc.strides[1],
-                    z, z + mc.overlaps[2] + mc.strides[2],
-                ])
-        else:
-            for _, _, x, y, _ in cm.motion_correction.sliding_window(
-                mc_image[0, :, :], mc.overlaps, mc.strides
-            ):
-                grid.append([
-                    x, x + mc.overlaps[0] + mc.strides[0],
-                    y, y + mc.overlaps[1] + mc.strides[1],
-                ])
-    else:
-        # Fallback for Rigid correction: The grid is simply the full FOV
-        if is3D:
-            grid = [[0, dims[0], 0, dims[1], 0, dims[2]]]
-        else:
-            grid = [[0, dims[0], 0, dims[1]]]
-
-
     # Open hdf5 file and create 'motion_correction' group
     caiman_fp = pathlib.Path(caiman_fp)
     h5f = h5py.File(caiman_fp.as_posix(), "r+" if caiman_fp.exists() else "w")

From 9f84dfad1c29f3a771450f999fec027d567c54a7 Mon Sep 17 00:00:00 2001
From: MathiasMahn <53939819+MathiasMahn@users.noreply.github.com>
Date: Tue, 10 Feb 2026 13:31:54 +0100
Subject: [PATCH 4/4] Remove dataset requirement for coord_shifts_rig

Comment out unnecessary dataset requirement for global single rigid shift.
---
 element_interface/caiman_loader.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/element_interface/caiman_loader.py b/element_interface/caiman_loader.py
index f15b665..7bbcea0 100644
--- a/element_interface/caiman_loader.py
+++ b/element_interface/caiman_loader.py
@@ -632,9 +632,12 @@ def _save_mc(
             data=mc.shifts_rig,
             dtype=mc.shifts_rig[0][0].dtype,
         )
-        h5g.require_dataset(
-            "coord_shifts_rig", shape=np.shape(grid), data=grid, dtype=type(grid[0][0])
-        )
+        
+        # Not needed for global single rigid shift - there is no grid!!!
+        # h5g.require_dataset(
+        #    "coord_shifts_rig", shape=np.shape(grid), data=grid, dtype=type(grid[0][0])
+        # )
+        
         reference_image = (
             np.tile(mc.total_template_rig, (1, 1, dims[-1]))
             if is3D