From 7cade60719eb972594b2ea0644a9667cb8c1ad2f Mon Sep 17 00:00:00 2001
From: lauraporta <29216006+lauraporta@users.noreply.github.com>
Date: Mon, 11 Dec 2023 18:11:58 +0000
Subject: [PATCH] Add script to run derotation in the cluster

---
 examples/derotation_slurm_job.py | 85 ++++++++++++++++++++++++++++++++
 1 file changed, 85 insertions(+)
 create mode 100644 examples/derotation_slurm_job.py

diff --git a/examples/derotation_slurm_job.py b/examples/derotation_slurm_job.py
new file mode 100644
index 0000000..bd781ce
--- /dev/null
+++ b/examples/derotation_slurm_job.py
@@ -0,0 +1,85 @@
+import os
+import sys
+from pathlib import Path
+
+import yaml
+
+from derotation.analysis.full_rotation_pipeline import FullPipeline
+from derotation.analysis.incremental_rotation_pipeline import (
+    IncrementalPipeline,
+)
+
+job_id = int(sys.argv[1:][0])
+dataset_path = sys.argv[1:][1]
+datasets = [path for path in os.listdir(dataset_path) if path.startswith("23")]
+dataset = datasets[job_id]
+
+bin_files = [
+    file
+    for file in os.listdir(f"{dataset_path}/{dataset}/aux_stim/")
+    if file.endswith(".bin")
+]
+full_rotation_bin = [file for file in bin_files if "_rotation" in file][0]
+incremental_bin = [file for file in bin_files if "increment" in file][0]
+
+image_files = [
+    file
+    for file in os.listdir(f"{dataset_path}/{dataset}/imaging/")
+    if file.endswith(".tif")
+]
+full_rotation_image = [file for file in image_files if "rotation_0" in file][0]
+incremental_image = [file for file in image_files if "increment_0" in file][0]
+
+#  make debug_plots and logs and derotated folder in dataset
+Path(f"{dataset_path}/{dataset}/debug_plots_incremental/").mkdir(
+    parents=True, exist_ok=True
+)
+Path(f"{dataset_path}/{dataset}/debug_plots_full/").mkdir(
+    parents=True, exist_ok=True
+)
+Path(f"{dataset_path}/{dataset}/logs/").mkdir(parents=True, exist_ok=True)
+Path(f"{dataset_path}/{dataset}/derotated/").mkdir(parents=True, exist_ok=True)
+
+for config_name in ["incremental_rotation", "full_rotation"]:
+    with open(f"derotation/config/{config_name}.yml") as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+
+    config["paths_read"][
+        "path_to_randperm"
+    ] = f"{dataset_path}/stimlus_randperm.mat"
+    bin_name = (
+        incremental_bin
+        if config_name == "incremental_rotation"
+        else full_rotation_bin
+    )
+    config["paths_read"][
+        "path_to_aux"
+    ] = f"{dataset_path}/{dataset}/aux_stim/{bin_name}"
+    image_name = (
+        incremental_image
+        if config_name == "incremental_rotation"
+        else full_rotation_image
+    )
+    config["paths_read"][
+        "path_to_tif"
+    ] = f"{dataset_path}/{dataset}/imaging/{image_name}"
+    config["paths_write"][
+        "debug_plots_folder"
+    ] = f"{dataset_path}/{dataset}/debug_plots_{config_name.split('_')[0]}"
+    config["paths_write"]["logs_folder"] = f"{dataset_path}/{dataset}/logs/"
+    config["paths_write"][
+        "derotated_tiff_folder"
+    ] = f"{dataset_path}/{dataset}/derotated/"
+    config["paths_write"][
+        "saving_name"
+    ] = f"derotated_image_stack_{config_name.split('_')[0]}"
+
+    with open(f"derotation/config/{config_name}_{job_id}.yml", "w") as f:
+        yaml.dump(config, f)
+
+
+derotate = IncrementalPipeline(f"incremental_rotation_{job_id}")
+derotate()
+
+derotate_full = FullPipeline(f"full_rotation_{job_id}")
+derotate_full()