Refactor job submission and environment setup for CREST integration

ReactionMechanismGenerator · Dec 22, 2024 · 881aa87 · 881aa87
1 parent 22ffb04
commit 881aa87
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 36 deletions.
diff --git a/arc/job/adapters/ts/heuristics.py b/arc/job/adapters/ts/heuristics.py
@@ -1080,12 +1080,12 @@ def h_abstraction(arc_reaction: 'ARCReaction',
             except (ValueError, KeyError) as e:
                 logger.error(f"Could not determine the H abstraction atoms, got:\n{e}")
 
-        if crest_paths:
-            crest_jobs = submit_crest_jobs(crest_paths)
-            monitor_crest_jobs(crest_jobs)  # Keep checking job statuses until complete
-            process_completed_jobs(crest_jobs, xyz_guesses)
-        else:
-            logger.error("No CREST paths found")
+    if crest_paths:
+        crest_jobs = submit_crest_jobs(crest_paths)
+        monitor_crest_jobs(crest_jobs)  # Keep checking job statuses until complete
+        process_completed_jobs(crest_jobs, xyz_guesses)
+    else:
+        logger.error("No CREST paths found")
 
 
     return xyz_guesses
@@ -1147,12 +1147,16 @@ def crest_ts_conformer_search(xyz_guess: dict, a_atom: int, h_atom: int, b_atom:
     commands = [
         f'{CREST_PATH}',
         f' -T {SERVERS["local"].get("cpus", 8)}',
-        f'{path}/coords.ref',
-        f'--cinp {path}/constraints.inp',
+        'coords.ref',
+        '--cinp constraints.inp',
         '--noreftopo'
     ]
     command = ' '.join(commands)
-    command = f"{CREST_ENV_PATH} && {command}" if CREST_ENV_PATH else command
+
+    if CREST_ENV_PATH:
+        activation_line = CREST_ENV_PATH
+    else:
+        activation_line = ''
 
     if SERVERS.get('local') is not None:
         if SERVERS['local']['cluster_soft'].lower() in ['condor', 'htcondor']:
@@ -1172,6 +1176,7 @@ def crest_ts_conformer_search(xyz_guess: dict, a_atom: int, h_atom: int, b_atom:
             # Write the crest job
             crest_job = submit_scripts['local']['crest_job']
             format_params = {
+                "activation_line": activation_line,
                 "commands": command,
             }
             crest_job = crest_job.format(**format_params)
@@ -1186,6 +1191,7 @@ def crest_ts_conformer_search(xyz_guess: dict, a_atom: int, h_atom: int, b_atom:
                 "name": f"crest_{xyz_crest_int}",
                 "cpus": SERVERS['local'].get('cpus', 8),
                 "memory": SERVERS['local'].get('memory', 32.0) * 1024,
+                "activation_line": activation_line,
                 "commands": command,
             }
             sub_job = sub_job.format(**format_params)

diff --git a/arc/settings/submit.py b/arc/settings/submit.py
@@ -311,14 +311,17 @@
 """,
         'crest_job': """#!/bin/bash
 
+source ~/.bashrc
+
 touch initial_time
 
+{activation_line}
+
 {commands}
 
 touch final_time
 
-# Remove all files except for crest_best.xyz, coords.ref constraints.inp
-rm -vfr !(crest_best.xyz|coords.ref|constraints.inp)
+find . ! -name 'crest_best.xyz' ! -name 'coords.ref' ! -name 'constraints.inp' -type f -exec rm -v {} +
 
 """,
         # Atlas uses HTCondor, see docs here: https://htcondor.readthedocs.io/en/latest/
@@ -893,13 +896,17 @@
 #PBS -o out.txt
 #PBS -e err.txt
 
+source ~/.bashrc
+
 touch initial_time
 
+{activation_line}
+
 {commands}
 
 touch final_time
 
-rm -vrf !(crest_best.xyz|coords.ref|constraints.inp)
+find . ! -name 'crest_best.xyz' ! -name 'coords.ref' ! -name 'constraints.inp' -type f -exec rm -v {} +
 
         """,
     },

diff --git a/devtools/install_crest.sh b/devtools/install_crest.sh
@@ -1,43 +1,43 @@
 #!/bin/bash -l
 
 # Check if Micromamba is installed
-if [ -x "$(command -v micromamba)" ]; then
-    echo "Micromamba is installed."
-    COMMAND_PKG=micromamba
+if [[ -x "$(command -v micromamba)" ]]; then
+	echo "Micromamba is installed."
+	COMMAND_PKG=micromamba
 # Check if Mamba is installed
-elif [ -x "$(command -v mamba)" ]; then
-    echo "Mamba is installed."
-    COMMAND_PKG=mamba
+elif [[ -x "$(command -v mamba)" ]]; then
+	echo "Mamba is installed."
+	COMMAND_PKG=mamba
 # Check if Conda is installed
-elif [ -x "$(command -v conda)" ]; then
-    echo "Conda is installed."
-    COMMAND_PKG=conda
+elif [[ -x "$(command -v conda)" ]]; then
+	echo "Conda is installed."
+	COMMAND_PKG=conda
 else
-    echo "Micromamba, Mamba, and Conda are not installed. Please download and install one of them - we strongly recommend Micromamba or Mamba."
-    exit 1
+	echo "Micromamba, Mamba, and Conda are not installed. Please download and install one of them - we strongly recommend Micromamba or Mamba."
+	exit 1
 fi
 
 # Set up Conda/Micromamba environment
-if [ "$COMMAND_PKG" = "micromamba" ]; then
-    eval "$(micromamba shell hook --shell=bash)"
-    micromamba activate base
-    BASE=$MAMBA_ROOT_PREFIX
-    # shellcheck source=/dev/null
-    . "$BASE/etc/profile.d/micromamba.sh"
+if [[ ${COMMAND_PKG} == "micromamba" ]]; then
+	eval "$(micromamba shell hook --shell=bash)"
+	micromamba activate base
+	BASE=${MAMBA_ROOT_PREFIX}
+	# shellcheck source=/dev/null
+	. "${BASE}/etc/profile.d/micromamba.sh"
 else
-    BASE=$(conda info --base)
-    # shellcheck source=/dev/null
-   . "$BASE/etc/profile.d/conda.sh"
+	BASE=$(conda info --base)
+	# shellcheck source=/dev/null
+	. "${BASE}/etc/profile.d/conda.sh"
 fi
 
 # create the environment
 echo "Creating the Crest environment..."
-$COMMAND_PKG create -n crest_env -c conda-forge python=3.10 crest=2.12 -y
+${COMMAND_PKG} create -n crest_env -c conda-forge python=3.10 crest=2.12 -y
 # Activate the environment
-if [ "$COMMAND_PKG" == "micromamba" ]; then
-    micromamba activate crest_env
+if [[ ${COMMAND_PKG} == "micromamba" ]]; then
+	micromamba activate crest_env
 else
-    conda activate crest_env
+	conda activate crest_env
 fi
 
 echo "Done installing Crest environment."