Mantis deskewing with multiprocessing and slurm

examples/slurm_option1/batch.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+INPUT_DATA="/hpc/instruments/cm.mantis/2023_05_10_PCNA_RAC1/timelapse_2_3/timelapse_2_lightsheet_1.zarr/*/*/*"
+OUTPUT_PATH="./timelapse_2_lightsheet_1_deskewed.zarr"
+DESKEW_PARAMS="./deskew_settings.yml"
+
+# Make an array of positions
+POSPATHS=($(natsort -p $INPUT_DATA))
+POSITIONS=${#POSPATHS[@]}
+
+# Get the Zarrstore name
+IFS='/' read -ra path <<< "${POSPATHS[0]}"
+ZARR_STORE=$(IFS='/'; echo "${path[*]:0:${#path[@]}-3}")
+echo "Zarr Store: $ZARR_STORE"
+
+rm ./output/deskew/*.out
+
+# Create an empty array to pre-initialize positions taking pos 0 as sample for the shape
+ZARR_JOB_ID=$(sbatch --parsable empty_zarr.sh "$INPUT_DATA" $DESKEW_PARAMS $OUTPUT_PATH)
+echo "DONE ZARR JOB: $ZARR_JOB_ID"
+DESKEW_JOB_ID=$(sbatch --parsable --array=0-$((POSITIONS-1)) -d after:$ZARR_JOB_ID deskew.sh "$INPUT_DATA" $OUTPUT_PATH)

examples/slurm_option1/deskew.sh

@@ -0,0 +1,52 @@
+#!/bin/bash
+#SBATCH --job-name=deskew
+#SBATCH --time=4:00:00
+#SBATCH --partition=cpu
+#SBATCH --cpus-per-task=16
+#SBATCH --mem-per-cpu=64G
+#SBATCH --output=./output/deskew/deskew-%A-%a.out
+env | grep "^SLURM" | sort
+
+module load anaconda
+module load comp_micro
+conda activate pyplay
+
+now=$(date '+%y-%m-%d')
+logpath=./logs/$now/deskew
+mkdir -p $logpath
+logfile="$logpath/deskew_$SLURM_ARRAY_TASK_ID.out"
+
+INPUT_PATH=$1
+OUTPUT_PATH=$2
+NUM_CORES=128
+
+echo  "in:  $INPUT_PATH " >> ${logfile}
+echo "out: $OUTPUT_PATH " >> ${logfile}
+echo "idx: $SLURM_ARRAY_TASK_ID " >> ${logfile}
+
+#Get the array of positions
+POSPATHS=($(natsort -p $INPUT_PATH)) >> ${logfile}
+echo "pospaths: ${POSPATHS[$SLURM_ARRAY_TASK_ID]}" >> ${logfile}
+
+# Create an array to store the last three directories for each path
+OUTPUT_FOVS=()
+
+# Extract the last three directories for each path
+for path in "${POSPATHS[@]}"; do
+  IFS='/' read -ra last_three_dirs <<< "$path"
+  OUTPUT_FOVS+=("${last_three_dirs[-3]}/${last_three_dirs[-2]}/${last_three_dirs[-1]}")
+done
+
+# Start measuring the execution time
+start_time=$(date +%s.%N) 
+echo "Starting the deskew" >> ${logfile}
+
+# Key Code
+mantis deskew ${POSPATHS[${SLURM_ARRAY_TASK_ID}]} ./deskew_settings.yml -o "$OUTPUT_PATH/${OUT_FOV_DIR[SLURM_ARRAY_TASK_ID]}" -j $NUM_CORES --slurm >> ${logfile}
+
+# End measuring the execution time
+end_time=$(date +%s.%N)
+# Calculate the elapsed time
+elapsed_time=$(echo "$end_time - $start_time" | bc)
+
+echo "Script execution time: $elapsed_time seconds" >> ${logfile}

examples/slurm_option1/deskew_settings.yml

@@ -0,0 +1,4 @@
+pixel_size_um: 0.116
+ls_angle_deg: 36.17  # can be calibrated using estimate_deskew.py
+px_to_scan_ratio: 0.371  # Optional, can be calibrated using estimate_deskew.py
+scan_step_um: 0.313  # Optional, corresponds to 10 mV

examples/slurm_option1/empty_zarr.py

@@ -0,0 +1,108 @@
+# %%%
+from iohub import open_ome_zarr
+
+from mantis.cli.parsing import (
+    input_data_path_argument,
+    deskew_param_argument,
+    output_dataset_options,
+)
+from mantis.analysis.deskew import get_deskewed_data_shape
+import yaml
+from mantis.analysis.AnalysisSettings import DeskewSettings
+from dataclasses import asdict
+
+# debugging
+from tqdm import tqdm
+import click
+import numpy as np
+from datetime import datetime
+import os
+from iohub.ngff_meta import TransformationMeta
+from natsort import natsorted
+
+def _get_deskew_params(deskew_params_path):
+    # Load params
+    with open(deskew_params_path) as file:
+        raw_settings = yaml.safe_load(file)
+    settings = DeskewSettings(**raw_settings)
+    print(f"Deskewing parameters: {asdict(settings)}")
+    return settings
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.command()
+@input_data_path_argument()
+@deskew_param_argument()
+@output_dataset_options(default="./deskewed.zarr")
+@click.option(
+    "--keep-overhang",
+    "-ko",
+    default=False,
+    is_flag=True,
+    help="Keep the overhanging region.",
+)
+@click.help_option("-h", "--help")
+def create_empty_zarr(input_data_path, deskew_param_path, output_path, keep_overhang):
+    # Load the "0" position to infer dataset information
+    input_data_path = natsorted(input_data_path)
+    click.echo(f"Input data folders {input_data_path})")
+
+    input_dataset = open_ome_zarr(str(input_data_path[0]), mode="r")
+    T, C, Z, Y, X = input_dataset.data.shape
+
+    # Get the deskewing parameters
+    settings = _get_deskew_params(deskew_param_path)
+    deskewed_shape, voxel_size = get_deskewed_data_shape(
+        (Z, Y, X),
+        settings.pixel_size_um,
+        settings.ls_angle_deg,
+        settings.px_to_scan_ratio,
+        keep_overhang,
+    )
+
+    click.echo("Creating empty array...")
+
+    # Handle transforms and metadata
+    transform = TransformationMeta(
+        type="scale",
+        scale=2 * (1,) + voxel_size,
+    )
+
+    # Prepare output dataset
+    channel_names = input_dataset.channel_names
+
+    # Output shape based on the type of reconstruction
+    output_shape = (T, len(channel_names)) + deskewed_shape
+    click.echo(f"Number of positions: {len(input_data_path)}")
+    click.echo(f"Output shape: {output_shape}")
+    chunk_size = (1, 1, 64) + deskewed_shape[-2:]
+    click.echo(f"Chunk size {chunk_size}")
+
+    # Create output dataset
+    output_dataset = open_ome_zarr(
+        output_path, layout="hcs", mode="w", channel_names=channel_names
+    )
+    # This takes care of the logic for single position or multiple position by wildcards
+    for filepath in input_data_path:
+        path_strings = filepath.split(os.path.sep)[-3:]
+        pos = output_dataset.create_position(
+            str(path_strings[0]), str(path_strings[1]), str(path_strings[2])
+        )
+        _ = pos.create_zeros(
+            name="0",
+            shape=(
+                T,
+                C,
+            )
+            + deskewed_shape,
+            chunks=chunk_size,
+            dtype=np.uint16,
+            transform=[transform],
+        )
+
+if __name__ == "__main__":
+    create_empty_zarr()

examples/slurm_option1/empty_zarr.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+
+#SBATCH --job-name=ZARR_INIT
+#SBATCH --time=0:10:00
+#SBATCH --partition=cpu
+#SBATCH --cpus-per-task=1
+#SBATCH --mem-per-cpu=5G
+#SBATCH --output=./output/deskew/empty_zarr_%j.out
+env | grep "^SLURM" | sort
+
+module load anaconda
+module load comp_micro
+conda activate pyplay
+
+IN_DATA=$1
+DESKEW_PARAMS=$2
+OUT_DATA=$3
+
+# Logging parameters
+rm -r $OUT_DATA
+now=$(date '+%y-%m-%d')
+logpath=./logs/$now/deskew
+mkdir -p $logpath
+rm $logpath/*.out
+logfile="$logpath/empty_zarr.out"
+
+echo  "raw: $IN_DATA" >> ${logfile}
+echo "out: $OUT_DATA " >> ${logfile}
+echo "deskew params: $DESKEW_PARAMS " >> ${logfile}
+
+python -u ./empty_zarr.py $IN_DATA $DESKEW_PARAMS -o $OUT_DATA >> ${logfile}
+