From e30672e8fb8063023a196a3a2dbc470f1ec37448 Mon Sep 17 00:00:00 2001
From: Ali Khan <akhanf@users.noreply.github.com>
Date: Tue, 24 Dec 2024 09:26:30 -0500
Subject: [PATCH] Behaviour and location of workdir, and removing hippunfold/
 toplevel folder (#338)

* removes --keep-work, adds --workdir

- WIP
- sets workdir to a temp directory by default
- the files won't be deleted by the workflow, but we could make use of
temp() to do so..

* updates
---
 hippunfold/config/snakebids.yml      | 14 ++++---
 hippunfold/workflow/Snakefile        |  7 +++-
 hippunfold/workflow/rules/common.smk | 60 +---------------------------
 3 files changed, 15 insertions(+), 66 deletions(-)

diff --git a/hippunfold/config/snakebids.yml b/hippunfold/config/snakebids.yml
index 6af8d752..5ac42e30 100644
--- a/hippunfold/config/snakebids.yml
+++ b/hippunfold/config/snakebids.yml
@@ -264,11 +264,6 @@ parse_args:
     nargs: '+'
     help: 'Run hipp (CA + subiculum) alone or include dentate (default: %(default)s)'
 
-  --keep_work:
-    help: 'Keep work folder intact instead of archiving it for each subject (default: %(default)s)'
-    default: False
-    action: 'store_true'
-
   --force_nnunet_model:
     help: 'Force nnunet model to use (expert option). (default: %(default)s)'
     default: False
@@ -294,6 +289,14 @@ parse_args:
     action: "store_false"
     default: True
 
+  --workdir:
+    help: |
+      Folder for storing working files. If not specified, a temporary folder
+      will be made in your system's temp directory (e.g. /tmp). You can also
+      use environment variables when setting the workdir, e.g. --workdir '$SLURM_TMPDIR'.
+    default: null
+    type: str
+
 
   --version:
     help: 'Print the version of HippUnfold'
@@ -688,6 +691,7 @@ t1_reg_template: False
 generate_myelin_map: False
 no_unfolded_reg: False
 root: results
+workdir: null
 use_template_seg: False
 template_seg_smoothing_factor: 2
 native_surf: False
diff --git a/hippunfold/workflow/Snakefile b/hippunfold/workflow/Snakefile
index c9952da9..5cbfc74d 100644
--- a/hippunfold/workflow/Snakefile
+++ b/hippunfold/workflow/Snakefile
@@ -2,6 +2,7 @@
 
 import snakebids
 from snakebids import bids, set_bids_spec
+from tempfile import mkdtemp
 
 
 configfile: "config/snakebids.yml"
@@ -94,8 +95,10 @@ wildcard_constraints:
     template="[a-zA-Z0-9]+",
 
 
-root = os.path.normpath(os.path.join(config["root"], "hippunfold"))
-work = os.path.normpath(os.path.join(config["root"], "work"))
+if config["workdir"] == None:
+    config["workdir"] = mkdtemp()
+work = os.path.expandvars(config["workdir"])
+root = os.path.expandvars(config["root"])
 
 
 include: "rules/common.smk"
diff --git a/hippunfold/workflow/rules/common.smk b/hippunfold/workflow/rules/common.smk
index 91fd1c23..5aed626c 100644
--- a/hippunfold/workflow/rules/common.smk
+++ b/hippunfold/workflow/rules/common.smk
@@ -332,7 +332,7 @@ def get_final_qc():
     return qc
 
 
-def get_final_subj_output():
+def get_final_output():
     subj_output = []
     subj_output.extend(get_final_spec())
     subj_output.extend(get_final_subfields())
@@ -343,33 +343,6 @@ def get_final_subj_output():
     return subj_output
 
 
-def get_final_output():
-    if config["keep_work"]:
-        subj_output = get_final_subj_output()
-    else:
-        subj_output = get_final_work_tar()
-
-    final_output = []
-
-    modality_suffix = get_modality_suffix(config["modality"])
-    modality_key = config["modality"]
-
-    # use a zip list for subject/session:
-    zip_list = inputs[modality_key].zip_lists
-    if "session" in zip_list:
-        zip_list = snakebids.filter_list(
-            zip_list,
-            {"session": inputs[config["modality"]].zip_lists["session"]},
-        )
-    final_output.extend(
-        expand(
-            expand(subj_output, zip, allow_missing=True, **zip_list),
-            modality_suffix=modality_suffix,
-        )
-    )
-    return final_output
-
-
 if "corobl" in ref_spaces:
 
     rule copy_coords_to_results:
@@ -403,19 +376,6 @@ if "corobl" in ref_spaces:
             "cp {input} {output}"
 
 
-def get_final_work_tar():
-    bids = bids_factory(specs.v0_0_0(subject_dir=False, session_dir=False))
-    return bids(root=work, suffix="work.tar.gz", **inputs.subj_wildcards)
-
-
-def get_work_dir(wildcards):
-    folder_with_file = inputs[config["modality"]].expand(
-        bids(root=work, **inputs.subj_wildcards), **wildcards
-    )
-    folder_without_file = os.path.dirname(folder_with_file[0])
-    return folder_without_file
-
-
 def get_download_dir():
     if "HIPPUNFOLD_CACHE_DIR" in os.environ.keys():
         download_dir = os.environ["HIPPUNFOLD_CACHE_DIR"]
@@ -424,21 +384,3 @@ def get_download_dir():
         dirs = AppDirs("hippunfold", "khanlab")
         download_dir = dirs.user_cache_dir
     return download_dir
-
-
-rule archive_work_after_final:
-    input:
-        get_final_subj_output(),
-    params:
-        work_dir=get_work_dir,
-    output:
-        get_final_work_tar(),
-    group:
-        "subj"
-    shell:
-        #exit code 0 or 1 is acceptable (2 is fatal)
-        "tar -czf {output} {params.work_dir}; "
-        "if [ $? -le 1 ]; then "
-        "  rm -rf {params.work_dir}; "
-        "else exit 1; "
-        "fi"