XY and Z drift stabilization; replaces #108

examples/slurmkit_example/slurmkit_stabilize.py

@@ -0,0 +1,88 @@
+import datetime
+import os
+import glob
+from slurmkit import SlurmParams, slurm_function, submit_function
+from natsort import natsorted
+from iohub import open_ome_zarr
+from pathlib import Path
+from mantis.cli.utils import (
+    yaml_to_model,
+    create_empty_hcs_zarr,
+    process_single_position_v2,
+)
+from mantis.cli.stabilize import apply_stabilization_transform
+import numpy as np
+from mantis.analysis.AnalysisSettings import StabilizationSettings
+
+# NOTE: this pipeline uses the focus found on well one for all. Perhaps this should be done per FOV(?)
+
+# io parameters
+input_position_dirpaths = "/hpc/projects/comp.micro/mantis/2023_11_08_Opencell_infection_redo/1-recon/OC43_infection_timelapse_3.zarr/*/*/*"
+output_dirpath = "/hpc/projects/comp.micro/mantis/2023_11_08_Opencell_infection_redo/3-stabilization/OC43_infection_timelapse_3.zarr"
+config_filepath = "/hpc/projects/comp.micro/mantis/2023_11_08_Opencell_infection_redo/3-stabilization/stabilization.yml"
+
+# batch and resource parameters
+partition = 'cpu'
+cpus_per_task = 5
+mem_per_cpu = "4G"
+time = 60  # minutes
+simultaneous_processes_per_node = 5
+
+# convert to Path
+input_position_dirpaths = [Path(p) for p in natsorted(glob.glob(input_position_dirpaths))]
+output_dirpath = Path(output_dirpath)
+config_filepath = Path(config_filepath)
+
+settings = yaml_to_model(config_filepath, StabilizationSettings)
+combined_mats = settings.affine_transform_zyx_list
+combined_mats = np.array(combined_mats)
+
+with open_ome_zarr(input_position_dirpaths[0]) as dataset:
+    T, C, Z, Y, X = dataset.data.shape
+    channel_names = dataset.channel_names
+
+output_metadata = {
+    "shape": (T, C, Z, Y, X),
+    "scale": dataset.scale,
+    "channel_names": channel_names,
+    "dtype": np.float32,
+}
+
+# Create the output zarr mirroring input_position_dirpaths
+create_empty_hcs_zarr(
+    store_path=output_dirpath,
+    position_keys=[p.parts[-3:] for p in input_position_dirpaths],
+    **output_metadata,
+)
+
+slurm_out_path = str(os.path.join(output_dirpath.parent, "slurm_output/stabilization-%j.out"))
+
+# prepare slurm parameters
+params = SlurmParams(
+    partition=partition,
+    cpus_per_task=cpus_per_task,
+    mem_per_cpu=mem_per_cpu,
+    time=datetime.timedelta(minutes=time),
+    output=slurm_out_path,
+)
+
+# wrap our utils.process_single_position() function with slurmkit
+slurm_process_single_position = slurm_function(process_single_position_v2)
+stabilization_function = slurm_process_single_position(
+    func=apply_stabilization_transform,
+    list_of_shifts=combined_mats,
+    time_indices=list(range(T)),
+    input_channel_idx=None,
+    output_channel_idx=None,
+    num_processes=simultaneous_processes_per_node,
+)
+
+stabilization_jobs = [
+    submit_function(
+        stabilization_function,
+        slurm_params=params,
+        input_data_path=in_path,
+        output_path=output_dirpath,
+    )
+    for in_path in input_position_dirpaths
+]

mantis/analysis/AnalysisSettings.py

@@ -67,3 +67,23 @@ def check_affine_transform(cls, v):
             raise ValueError("The array must contain valid numerical values.")
 
         return v
+
+
+class StabilizationSettings(MyBaseModel):
+    stabilization_estimation_channel: str
+    stabilization_type: Literal["z", "xy", "xyz"]
+    stabilization_channels: list
+    affine_transform_zyx_list: list
+    time_indices: Union[NonNegativeInt, list[NonNegativeInt], Literal["all"]] = "all"
+
+    @validator("affine_transform_zyx_list")
+    def check_affine_transform_list(cls, v):
+        if not isinstance(v, list):
+            raise ValueError("affine_transform_list must be a list")
+
+        for arr in v:
+            arr = np.array(arr)
+            if arr.shape != (4, 4):
+                raise ValueError("Each element in affine_transform_list must be a 4x4 ndarray")
+
+        return v

mantis/analysis/settings/example_stabilize_timelapse_settings.yml

@@ -0,0 +1,5 @@
+stabilization_estimation_channel: "GFP"
+stabilization_type: "xyz"
+stabilization_channels: ['GFP']
+affine_transform_zyx_list: [[[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0],[0.0, 0.0, 0.0,1.0]],[[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0],[0.0, 0.0, 0.0,1.0]],[[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0],[0.0, 0.0, 0.0,1.0]]]
+time_indices: 'all'