Merge pull request #195 from ttngu207/dev_per_field_processing

fix: skip ingesting for rigid motion correction - suite2p no longer output the shifts
datajoint · Apr 30, 2024 · 3187b04 · 3187b04
2 parents 27da38d + ca291cd
commit 3187b04
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 34 deletions.
diff --git a/element_calcium_imaging/field_processing.py b/element_calcium_imaging/field_processing.py
@@ -1,3 +1,4 @@
+import gc
 import importlib
 import inspect
 import pathlib
@@ -299,6 +300,10 @@ def make(self, key):
                 f"Field processing for {acq_software} scans with {method} is not yet supported in this table."
             )
 
+        # explicitly garbage collect after the processing step (caiman or suite2p)
+        # as these may have large memory footprint and may not be cleared fast enough
+        gc.collect()
+
         exec_dur = (datetime.utcnow() - execution_time).total_seconds() / 3600
         self.insert1(
             {

diff --git a/element_calcium_imaging/imaging_no_curation.py b/element_calcium_imaging/imaging_no_curation.py
@@ -755,42 +755,42 @@ def make(self, key):
             rigid_correction, nonrigid_correction, nonrigid_blocks = {}, {}, {}
             summary_images = []
             for idx, (plane, s2p) in enumerate(suite2p_dataset.planes.items()):
-                # -- rigid motion correction --
-                if idx == 0:
-                    rigid_correction = {
-                        **key,
-                        "y_shifts": s2p.ops["yoff"],
-                        "x_shifts": s2p.ops["xoff"],
-                        "z_shifts": np.full_like(s2p.ops["xoff"], 0),
-                        "y_std": np.nanstd(s2p.ops["yoff"]),
-                        "x_std": np.nanstd(s2p.ops["xoff"]),
-                        "z_std": np.nan,
-                        "outlier_frames": s2p.ops["badframes"],
-                    }
-                else:
-                    rigid_correction["y_shifts"] = np.vstack(
-                        [rigid_correction["y_shifts"], s2p.ops["yoff"]]
-                    )
-                    rigid_correction["y_std"] = np.nanstd(
-                        rigid_correction["y_shifts"].flatten()
+                if not all(k in s2p.ops for k in ["xblock", "yblock", "nblocks"]):
+                    logger.warning(
+                        f"Unable to load/ingest nonrigid motion correction for plane {plane}. "
+                        f"Nonrigid motion correction output is no longer saved by Suite2p for version above 0.10.*."
                     )
-                    rigid_correction["x_shifts"] = np.vstack(
-                        [rigid_correction["x_shifts"], s2p.ops["xoff"]]
-                    )
-                    rigid_correction["x_std"] = np.nanstd(
-                        rigid_correction["x_shifts"].flatten()
-                    )
-                    rigid_correction["outlier_frames"] = np.logical_or(
-                        rigid_correction["outlier_frames"], s2p.ops["badframes"]
-                    )
-                # -- non-rigid motion correction --
-                if s2p.ops["nonrigid"]:
-                    if not all(k in s2p.ops for k in ["xblock", "yblock", "nblocks"]):
-                        logger.warning(
-                            f"Unable to load/ingest nonrigid motion correction for plane {plane}. "
-                            f"Nonrigid motion correction output is no longer saved by Suite2p for version above 0.10.*."
-                        )
+                else:
+                    # -- rigid motion correction --
+                    if idx == 0:
+                        rigid_correction = {
+                            **key,
+                            "y_shifts": s2p.ops["yoff"],
+                            "x_shifts": s2p.ops["xoff"],
+                            "z_shifts": np.full_like(s2p.ops["xoff"], 0),
+                            "y_std": np.nanstd(s2p.ops["yoff"]),
+                            "x_std": np.nanstd(s2p.ops["xoff"]),
+                            "z_std": np.nan,
+                            "outlier_frames": s2p.ops["badframes"],
+                        }
                     else:
+                        rigid_correction["y_shifts"] = np.vstack(
+                            [rigid_correction["y_shifts"], s2p.ops["yoff"]]
+                        )
+                        rigid_correction["y_std"] = np.nanstd(
+                            rigid_correction["y_shifts"].flatten()
+                        )
+                        rigid_correction["x_shifts"] = np.vstack(
+                            [rigid_correction["x_shifts"], s2p.ops["xoff"]]
+                        )
+                        rigid_correction["x_std"] = np.nanstd(
+                            rigid_correction["x_shifts"].flatten()
+                        )
+                        rigid_correction["outlier_frames"] = np.logical_or(
+                            rigid_correction["outlier_frames"], s2p.ops["badframes"]
+                        )
+                    # -- non-rigid motion correction --
+                    if s2p.ops["nonrigid"]:
                         if idx == 0:
                             nonrigid_correction = {
                                 **key,