Merge pull request #143 from logsdail/customcommand_react

Adds 'custom' HPC option to React workflow.

carmm/err_handler/logger_set.py

@@ -0,0 +1,48 @@
+import logging
+import sys
+def set_logger(logger_name, warning_level = 1):
+
+    """
+    A general purpose error handler for CARMM - sets the Logger object
+    with a sensible set of defaults
+    
+    Args:
+        logger_name: str
+            Name of logger passed to get_logger
+        warning_level: int
+            Sets minimum-severity log message printed from lowest (0, DEBUG) to highest (4, CRITICAL)
+
+    Returns:
+        carmm_logger: logger
+            Logger object to be used by other modules for error handling
+
+    """
+
+    # Create dictionary for easy access debugs
+    warning_lvls = {0: logging.DEBUG, 1: logging.INFO, 2: logging.WARNING,
+                    3: logging.ERROR, 4: logging.CRITICAL}
+
+    # create logger
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(warning_lvls[warning_level])
+
+    # create console handler and set level to debug
+    ch = logging.StreamHandler(sys.stdout)
+    ch.setLevel(warning_lvls[warning_level])
+
+    # create formatter
+    # FULL FORMATTING
+    #formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+    # REDUCED FORMATTING
+    formatter = logging.Formatter('%(levelname)s - %(message)s')
+
+    # add formatter to ch
+    ch.setFormatter(formatter)
+
+    # add ch to logger
+    logger.addHandler(ch)
+
+    return logger
+
+
+

carmm/run/aims_path.py

@@ -7,7 +7,12 @@ def set_aims_command(hpc='hawk', basis_set='light', defaults=2010, nodes_per_ins
     Parameters:
     hpc: String
         Name of the HPC facility where the jobs are being run
-        Options: 'hawk', 'hawk-amd', 'isambard', 'archer2', 'young', 'aws'
+        Options: 'hawk', 'hawk-amd', 'isambard', 'archer2', 'young', 'aws', 'custom'
+        NOTE 1: 'custom' requires the environmental variable "CARMM_AIMS_ROOT_DIRECTORY"
+        before running to allow logic of basis set selection, while maintaining
+        free choice of basis set folders.
+        NOTE 2: 'custom' requires environmental variable ASE_AIMS_COMMAND be set, with
+        the desired number of mpi processes and WITHOUT a default output file
     basis_set: String
         Name of basis set for FHI-aims
         Options: 'light', 'intermediate', 'tight', 'really_tight' etc.
@@ -23,6 +28,17 @@ def set_aims_command(hpc='hawk', basis_set='light', defaults=2010, nodes_per_ins
 
     hpc = hpc.lower()
 
+    #TODO: Handling of AIMs ASE calculator is due to change in ASE 3.23
+    #TODO: This method will need to be futureproof/backwards compatible
+    if hpc == "custom":
+        assert "CARMM_AIMS_ROOT_DIRECTORY" in os.environ, \
+            "hpc is 'custom' but environmental variable CARMM_AIMS_ROOT_DIRECTORY not specified."
+
+        custom_root_dir = os.environ["CARMM_AIMS_ROOT_DIRECTORY"]
+    else:
+        custom_root_dir = None
+
+
     species = "species_defaults/" + "defaults_" + str(defaults) + "/" + basis_set
 
     preamble = {
@@ -31,7 +47,8 @@ def set_aims_command(hpc='hawk', basis_set='light', defaults=2010, nodes_per_ins
         "isambard": "time aprun",
         "archer2": "srun --cpu-bind=cores --distribution=block:block --hint=nomultithread",
         "young": "gerun",
-        "aws": "time srun --mpi=pmi2 --hint=nomultithread --distribution=block:block"
+        "aws": "time srun --mpi=pmi2 --hint=nomultithread --distribution=block:block",
+        "custom": ""
     }
 
     assert hpc in preamble, "Inappropriate HPC facility: " + hpc + "is not recognised."
@@ -43,6 +60,7 @@ def set_aims_command(hpc='hawk', basis_set='light', defaults=2010, nodes_per_ins
         "archer2": "/work/e05/e05-files-log/shared/software/fhi-aims/",
         "young": "/home/mmm0170/Software/fhi-aims/",
         "aws": "/shared/logsdail_group/sing/",
+        "custom": custom_root_dir
     }
 
     executable_d = {"compiled": "bin/aims.$VERSION.scalapack.mpi.x",
@@ -53,11 +71,11 @@ def set_aims_command(hpc='hawk', basis_set='light', defaults=2010, nodes_per_ins
     '''Handle compiled and containerized FHIaims versions'''
     if hpc == "aws":
         executable = executable_d["apptainer"]
-    else:
+    elif hpc != "custom":
         executable = fhi_aims_directory[hpc] + executable_d["compiled"]
 
     """Set the relevant environment variables based on HPC"""
-    os.environ["AIMS_SPECIES_DIR"] = fhi_aims_directory[hpc] + species     
+    os.environ["AIMS_SPECIES_DIR"] = fhi_aims_directory[hpc] + species
 
     # Define the executable command
     if nodes_per_instance:
@@ -69,11 +87,15 @@ def set_aims_command(hpc='hawk', basis_set='light', defaults=2010, nodes_per_ins
         if hpc == "aws":
             assert nodes_per_instance == 1, "FHI-aims does not run on more than one node on AWS at present."
 
-    # This has a helper function as we need to take different actions 
-    # if running single or task-farmed calculations
-    cpu_command = _get_cpu_command(hpc, nodes_per_instance)
+    if hpc == 'custom':
+        assert "ASE_AIMS_COMMAND" in os.environ, \
+            "set_aims_command: option hpc is 'custom', but ASE_AIMS_COMMAND not set."
+    else:
+        # This has a helper function as we need to take different actions
+        # if running single or task-farmed calculations
+        cpu_command = _get_cpu_command(hpc, nodes_per_instance)
 
-    os.environ["ASE_AIMS_COMMAND"] = f"{preamble[hpc]} {cpu_command} {executable}"
+        os.environ["ASE_AIMS_COMMAND"] = f"{preamble[hpc]} {cpu_command} {executable}"
 
 def _get_cpu_command(hpc, nodes_per_instance=None):
     """

carmm/run/workflows/react.py

@@ -9,6 +9,7 @@
 from carmm.run.aims_path import set_aims_command
 from ase.io import Trajectory
 from carmm.run.workflows.helper import CalculationHelper
+from carmm.err_handler.logger_set import set_logger
 from ase.optimize import BFGS
 
 # TODO: Enable serialization with ASE db - save locations of converged files as well as all properties
@@ -25,7 +26,9 @@ def __init__(self,
                  filename: str = None,
                  nodes_per_instance: int = None,
                  dry_run: bool = False,
-                 verbose: bool = True):
+                 verbose: bool = True,
+                 warning_lvl: int = 1
+                 ):
         """
         Args:
             params: dict
@@ -58,6 +61,8 @@ def __init__(self,
         self.nodes_per_instance = nodes_per_instance
         self.verbose = verbose
 
+        self.logger = set_logger("react_logger", warning_lvl)
+
         """Define additional parameters"""
         self.initial = None  # input for optimisation or input for NEB initial image
         self.model_optimised = None  # optimised geometry with calculator attached
@@ -71,6 +76,16 @@ def __init__(self,
         """ Set the test flag"""
         self.dry_run = dry_run
 
+        if hpc == "custom":
+            self.logger.debug(f"WARNING: You have selected 'custom' as an option for HPC.                      ")
+            self.logger.debug(f"         This requires a couple of extra steps from the user side.             ")
+            self.logger.debug(f"         1) A new environmental variable - CARMM_AIMS_ROOT_DIRECTORY           ")
+            self.logger.debug(f"            - must be set. This helps find the folder containing the           ")
+            self.logger.debug(f"             default basis function.                                           ")
+            self.logger.debug(f"         2) ASE_AIMS_COMMAND must be set, with the correct number of           ")
+            self.logger.debug(f"            mpi tasks if desired. Avoid piping output, as React has its own    ")
+            self.logger.debug(f"            output folder names.                                               ")
+
     def aims_optimise(self, atoms: Atoms, fmax: float = 0.01, post_process: str = None, relax_unit_cell: bool = False,
                       restart: bool = True, optimiser=None, opt_kwargs: dict = {}):
 

examples/err_handler_testoutput.py

@@ -0,0 +1,11 @@
+
+def test_set_logger():
+    from carmm.err_handler.logger_set import set_logger
+
+    #TODO: MAKE THIS TEST ACTUALLY TEST CAPTURED OUTPUT
+
+    logger = set_logger("test_logger", 0)
+
+    logger.debug('test_set_logger test passed! DEBUG error correctly printed')
+
+test_set_logger()

examples/run_workflows_ReactAims.py

@@ -35,12 +35,11 @@ def test_run_workflows_ReactAims():
     model_optimised, model_postprocessed = reactor.aims_optimise(atoms, fmax=0.05, restart=True, optimiser=FIRE)
 
     zero_point_energy = reactor.vibrate(atoms, indices =[atom.index for atom in atoms]).get_zero_point_energy()
-    
+
     assert is_converged(reactor.model_optimised, 0.01), \
     '''The structure saved in React_Aims is not converged'''
     assert round(zero_point_energy, 3) == 0.275
 
-
     '''Create a reaction pathway'''
     atoms1 = atoms.copy()
     atoms1[1].x += 8
@@ -54,6 +53,24 @@ def test_run_workflows_ReactAims():
     molecule_with_charges = reactor.get_mulliken_charges(model_optimised)
     assert molecule_with_charges[0].charge == 0.0
 
+
+    '''Test 'custom' HPC by replicationg the settings for HAWK'''
+    atoms = molecule("H2")
+
+    os.environ['CARMM_AIMS_ROOT_DIRECTORY'] = "/apps/local/projects/scw1057/software/fhi-aims/"
+    os.environ['ASE_AIMS_COMMAND'] = "time srun" + \
+                                     f"--nodes=$SLURM_NNODES --ntasks=$SLURM_NTASKS -d mpirun" + \
+                                     "/apps/local/projects/scw1057/software/fhi-aims/bin/aims.$VERSION.scalapack.mpi.x"
+
+    reactor = ReactAims(params, basis_set, hpc="custom", filename="H2")
+    model_optimised, model_postprocessed = reactor.aims_optimise(atoms, fmax=0.05, restart=True, optimiser=FIRE)
+    zero_point_energy = reactor.vibrate(atoms, indices =[atom.index for atom in atoms]).get_zero_point_energy()
+
+    assert is_converged(reactor.model_optimised, 0.01), \
+    '''The structure saved in React_Aims is not converged'''
+    assert round(zero_point_energy, 3) == 0.275
+
+
     '''The below uses the "dry_run" flag and uses an EMT calculator instead of FHI-aims to test code in CI'''
     reactor = ReactAims(params, basis_set, hpc, dry_run=True, filename="H")
     H_atom = Atoms("H", positions=[(0,0,0)])
@@ -70,6 +87,22 @@ def test_run_workflows_ReactAims():
     '''Calculate the optimal unit cell and post process the calculation with a larger "tight" basis set'''
     light, tight = reactor.aims_optimise(Al_bulk, 0.01, relax_unit_cell=True, post_process="tight")
 
+    '''The below uses the "dry_run" flag and uses'''
+    reactor = ReactAims(params, basis_set, hpc, dry_run=True, filename="H")
+    H_atom = Atoms("H", positions=[(0, 0, 0)])
+    reactor.aims_optimise(atoms, post_process="tight")
+    reactor.get_mulliken_charges(H_atom)
+    reactor.vibrate(H_atom, indices=[0])
+    vib = reactor.vibrate(H_atom, indices=[0], read_only=True)
+
+    '''Optimise the bulk metal using stress tensor calculations and ExpCellFilter to later cut a surface model'''
+    reactor.filename = "Al"
+    reactor.params["k_grid"] = (8, 8, 8)
+    Al_bulk = bulk("Al")
+
+    '''Calculate the optimal unit cell and post process the calculation with a larger "tight" basis set'''
+    light, tight = reactor.aims_optimise(Al_bulk, 0.01, relax_unit_cell=True, post_process="tight")
+
     """
     '''Cut a 2x2-Al(001) surface with 3 layers and an
      Au atom adsorbed in a hollow site:'''