Merge branch 'develop' of https://github.com/abinit/abipy into develop

dev_scripts/abiml_md.py

@@ -1,14 +1,16 @@
 #!/usr/bin/env python
-
+import sys
 import os
+
 from abipy.flowtk.qutils import SlurmJobArray
 
-def main():
- # IMPORTANT: You need customize the slurm options according to your machine.
- conda_env = os.environ['CONDA_DEFAULT_ENV']
- print(f"Slurm script will be executed in {conda_env=}")
 
- header = f"""\
+def get_slurm_header(conda_env)
+ # IMPORTANT: You need customize the slurm options below according to your machine.
+ #conda_env = os.environ['CONDA_DEFAULT_ENV']
+ #print(f"Slurm script will be executed in {conda_env=}")
+
+ slurm_header = f"""\
 #!/bin/bash
 
 #SBATCH --account=battab
@@ -24,32 +26,39 @@ def main():
 ulimit -s unlimited
 
 # Important: This is the conda environment providing the NN potential(s) nn_names
+source $HOME/.bashrc
 conda activate {conda_env}
 """
+ return slurm_header
+
+
+def main():
  filepath = "LLZO_cubic.vasp"
- nn_names = ["m3gnet", "chgnet"]
- temp_list = [600, 800, 1000, 1200]
+ nn_names_env = {"mace_mp": "mace_mp"}
+ temperature_list = [600, 800, 1000, 1200]
  steps = 4000
+ steps = 100
 
  print(f"""\
 Performing MD calculations with the following parameters:
 
 {filepath=}
 {nn_names=}
-{temp_list=}
+{temperature_list=}
 {steps=}
 """)
 
  arr_options = []
- for nn_name in nn_names:
- for temperature in temp_list:
+ for nn_name, conda_env in nn_names_env.items():
+ slurm_header = get_slurm_header(conda_env)
+ for temperature in temperature_list:
  workdir = f"nn-{nn_name}_T-{temperature}"
  opts = f"{filepath} --nn-name {nn_name} --temperature {temperature} --timestep 1 \
  --loginterval 100 --steps {steps} --ensemble nvt -w {workdir}"
  arr_options.append(opts)
 
- job_array = SlurmJobArray(header, "abiml.py md", arr_options)
- job_array.sbatch("job_array.sh")
+  job_array = SlurmJobArray(slurm_header, "abiml.py md", arr_options)
+  job_array.sbatch("job_array.sh")
 
 
 if __name__ == "__main__":