From f897e298153ce77ae55981c110f6734c7e2c6a84 Mon Sep 17 00:00:00 2001
From: Eduardo Hirata-Miyasaki <edhiratam@gmail.com>
Date: Thu, 21 Sep 2023 22:30:39 -0700
Subject: [PATCH] slurmkit example for phase3D recon.

---
 examples/slurmkit/slurmkit_phase3D.py | 101 ++++++++++++++++++++++++++
 1 file changed, 101 insertions(+)
 create mode 100644 examples/slurmkit/slurmkit_phase3D.py

diff --git a/examples/slurmkit/slurmkit_phase3D.py b/examples/slurmkit/slurmkit_phase3D.py
new file mode 100644
index 00000000..fd81a71f
--- /dev/null
+++ b/examples/slurmkit/slurmkit_phase3D.py
@@ -0,0 +1,101 @@
+import datetime
+import os
+import glob
+from slurmkit import SlurmParams, slurm_function, submit_function
+from natsort import natsorted
+import click
+from iohub import open_ome_zarr
+from recOrder.cli.settings import ReconstructionSettings
+import torch
+from recOrder.io import utils
+from recOrder.cli.utils import (
+    apply_inverse_to_zyx_and_save,
+    create_empty_hcs_zarr,
+)
+from recOrder.cli.compute_transfer_function import (
+    compute_transfer_function_cli,
+)
+
+from recOrder.cli.apply_inverse_transfer_function import (
+    get_reconstruction_output_metadata,
+    apply_inverse_transfer_function_single_position,
+)
+from pathlib import Path
+
+
+# Reconstruction parameters
+config_path = (
+    "/home/eduardo.hirata/repos/recOrder/recOrder/tests/cli_tests/phase.yml"
+)
+transfer_function_path = "/home/eduardo.hirata/repos/recOrder/recOrder/tests/cli_tests/TF_phase3D.zarr"
+
+# io parameters
+input_position_dirpaths = "/home/eduardo.hirata/repos/recOrder/recOrder/tests/cli_tests/data_temp/2022_08_04_20x_04NA_BF.zarr/*/*/*"
+output_dirpath = "/home/eduardo.hirata/repos/recOrder/recOrder/tests/cli_tests/test_delete_phase.zarr"
+# sbatch and resource parameters
+cpus_per_task = 16
+mem_per_cpu = "16G"
+time = 40  # minutes
+
+# Path handling
+input_position_dirpaths = [
+    Path(path) for path in natsorted(glob.glob(input_position_dirpaths))
+]
+output_dirpath = Path(output_dirpath)
+slurm_out_path = output_dirpath.parent / "slurm_output/recon-%j.out"
+
+transfer_function_path = Path(transfer_function_path)
+config_path = Path(config_path)
+
+## First compute-tf
+# Handle transfer function path
+# Compute transfer function
+compute_transfer_function_cli(
+    input_position_dirpath=input_position_dirpaths[0],
+    config_filepath=config_path,
+    output_dirpath=transfer_function_path,
+)
+
+## Second apply-inv-tf
+output_metadata = get_reconstruction_output_metadata(
+    input_position_dirpaths[0], config_path
+)
+
+create_empty_hcs_zarr(
+    store_path=output_dirpath,
+    position_keys=[p.parts[-3:] for p in input_position_dirpaths],
+    **output_metadata,
+)
+
+click.echo(f"in: {input_position_dirpaths}, out: {output_dirpath}")
+
+# prepare slurm parameters
+params = SlurmParams(
+    partition="cpu",
+    cpus_per_task=cpus_per_task,
+    mem_per_cpu=mem_per_cpu,
+    time=datetime.timedelta(minutes=time),
+    output=slurm_out_path,
+)
+
+# wrap our process_single_position() function with slurmkit
+slurm_reconstruct_single_position = slurm_function(
+    apply_inverse_transfer_function_single_position
+)
+reconstruct_func = slurm_reconstruct_single_position(
+    transfer_function_dirpath=transfer_function_path,
+    config_filepath=config_path,
+    num_processes=cpus_per_task,
+)
+
+# generate an array of jobs by passing the in_path and out_path to slurm wrapped function
+recon_jobs = [
+    submit_function(
+        reconstruct_func,
+        slurm_params=params,
+        input_position_dirpath=input_position_dirpath,
+        output_position_dirpath=output_dirpath
+        / Path(*input_position_dirpath.parts[-3:]),
+    )
+    for input_position_dirpath in input_position_dirpaths
+]