From 62cd1fefafd8b78dd0a2f146323b637423de4c9e Mon Sep 17 00:00:00 2001 From: kushalbakshi Date: Tue, 30 Apr 2024 09:45:58 -0500 Subject: [PATCH] Fix: skip rigid/non-rigid motion correction - shifts not in suite2p output --- .../imaging_no_curation.py | 158 +++++++++--------- 1 file changed, 82 insertions(+), 76 deletions(-) diff --git a/element_calcium_imaging/imaging_no_curation.py b/element_calcium_imaging/imaging_no_curation.py index e8ebd204..d1769f36 100644 --- a/element_calcium_imaging/imaging_no_curation.py +++ b/element_calcium_imaging/imaging_no_curation.py @@ -828,93 +828,99 @@ 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() - ) - 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"] + if not all(k in s2p.ops for k in ["xblock", "yblock", "nblocks"]): + logger.warning( + f"Unable to load/ingest non-rigid motion correction for plane {plane}." + "Non-rigid motion correction data is not saved by Suite2p for versions above 0.10.*." ) - # -- non-rigid motion correction -- - if s2p.ops["nonrigid"]: + else: + # -- rigid motion correction -- if idx == 0: - nonrigid_correction = { + rigid_correction = { **key, - "block_height": s2p.ops["block_size"][0], - "block_width": s2p.ops["block_size"][1], - "block_depth": 1, - "block_count_y": s2p.ops["nblocks"][0], - "block_count_x": s2p.ops["nblocks"][1], - "block_count_z": len(suite2p_dataset.planes), + "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: - nonrigid_correction["outlier_frames"] = np.logical_or( - nonrigid_correction["outlier_frames"], - s2p.ops["badframes"], + rigid_correction["y_shifts"] = np.vstack( + [rigid_correction["y_shifts"], s2p.ops["yoff"]] ) - for b_id, (b_y, b_x, bshift_y, bshift_x) in enumerate( - zip( - s2p.ops["xblock"], - s2p.ops["yblock"], - s2p.ops["yoff1"].T, - s2p.ops["xoff1"].T, + rigid_correction["y_std"] = np.nanstd( + rigid_correction["y_shifts"].flatten() ) - ): - if b_id in nonrigid_blocks: - nonrigid_blocks[b_id]["y_shifts"] = np.vstack( - [nonrigid_blocks[b_id]["y_shifts"], bshift_y] - ) - nonrigid_blocks[b_id]["y_std"] = np.nanstd( - nonrigid_blocks[b_id]["y_shifts"].flatten() - ) - nonrigid_blocks[b_id]["x_shifts"] = np.vstack( - [nonrigid_blocks[b_id]["x_shifts"], bshift_x] + 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, + "block_height": s2p.ops["block_size"][0], + "block_width": s2p.ops["block_size"][1], + "block_depth": 1, + "block_count_y": s2p.ops["nblocks"][0], + "block_count_x": s2p.ops["nblocks"][1], + "block_count_z": len(suite2p_dataset.planes), + "outlier_frames": s2p.ops["badframes"], + } + else: + nonrigid_correction["outlier_frames"] = np.logical_or( + nonrigid_correction["outlier_frames"], + s2p.ops["badframes"], ) - nonrigid_blocks[b_id]["x_std"] = np.nanstd( - nonrigid_blocks[b_id]["x_shifts"].flatten() + for b_id, (b_y, b_x, bshift_y, bshift_x) in enumerate( + zip( + s2p.ops["xblock"], + s2p.ops["yblock"], + s2p.ops["yoff1"].T, + s2p.ops["xoff1"].T, ) - else: - nonrigid_blocks[b_id] = { - **key, - "block_id": b_id, - "block_y": b_y, - "block_x": b_x, - "block_z": np.full_like(b_x, plane), - "y_shifts": bshift_y, - "x_shifts": bshift_x, - "z_shifts": np.full( - ( - len(suite2p_dataset.planes), - len(bshift_x), + ): + if b_id in nonrigid_blocks: + nonrigid_blocks[b_id]["y_shifts"] = np.vstack( + [nonrigid_blocks[b_id]["y_shifts"], bshift_y] + ) + nonrigid_blocks[b_id]["y_std"] = np.nanstd( + nonrigid_blocks[b_id]["y_shifts"].flatten() + ) + nonrigid_blocks[b_id]["x_shifts"] = np.vstack( + [nonrigid_blocks[b_id]["x_shifts"], bshift_x] + ) + nonrigid_blocks[b_id]["x_std"] = np.nanstd( + nonrigid_blocks[b_id]["x_shifts"].flatten() + ) + else: + nonrigid_blocks[b_id] = { + **key, + "block_id": b_id, + "block_y": b_y, + "block_x": b_x, + "block_z": np.full_like(b_x, plane), + "y_shifts": bshift_y, + "x_shifts": bshift_x, + "z_shifts": np.full( + ( + len(suite2p_dataset.planes), + len(bshift_x), + ), + 0, ), - 0, - ), - "y_std": np.nanstd(bshift_y), - "x_std": np.nanstd(bshift_x), - "z_std": np.nan, - } + "y_std": np.nanstd(bshift_y), + "x_std": np.nanstd(bshift_x), + "z_std": np.nan, + } # -- summary images -- motion_correction_key = (