MACE Preoptimisation for FHI-aims

carmm/run/workflows/react.py

@@ -87,7 +87,7 @@
             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 = {}):
+                      restart: bool = True, optimiser=None, opt_kwargs: dict = {}, mace_preopt=False):
 
         """
          The function needs information about structure geometry (model), name of hpc system
@@ -132,6 +132,12 @@
         if initial is not None:
             self.initial = initial
 
+        if mace_preopt and not initial:
+            """Make sure to skip preoptimisations if AIMs restart information is available."""
+            assert self._MaceReactor is not None, "Please set MaceReact_Preoptimiser if mace_preopt is True"
+
+            self.initial = self._mace_preoptimise(self.initial, fmax, relax_unit_cell, optimiser, opt_kwargs)
+
         self._perform_optimization(subdirectory_name, out, counter, fmax, relax_unit_cell, optimiser, opt_kwargs)
         self._finalize_optimization(subdirectory_name, post_process)
 
@@ -323,11 +329,10 @@
 
         return self.initial
 
-
     def search_ts(self, initial: Atoms, final: Atoms,
                   fmax: float, unc: float, interpolation=None,
                   n=0.25, steps=40, restart=True, prev_calcs=None,
-                  input_check=0.01):
+                  input_check=0.01, mace_preopt=0, preopt_maxsteps=200):
         """
         This function allows calculation of the transition state using the CatLearn software package in an
         ASE/sockets/FHI-aims setup. The resulting converged band will be located in the MLNEB.traj file.
@@ -358,12 +363,24 @@
             input_check: float or None
                 If float the calculators of the input structures will be checked if the structures are below
                 the requested fmax and an optimisation will be performed if not.
+            mace_preopt: None or int
+                Controls whether to use a MACE preoptimised TS path, and the work flow used for preoptimisation
+                Requires a MACE calculator be attached to the React_AIMs object via the MaceReact_Preoptimiser
+                property.
+                0: Off      -     No MACE pre-optimisation
+                1: Full Preopt -  MACE preoptimises TS before FHI-aims - FHI-aims inherits all structures from MACE.
+                2: TS only     -  MACE preoptimises after FHI-aims check - MACE receives initial and reactant
+                                  structures from FHI-aims and does not optimise
+            preopt_maxsteps: int
+                Controls the maximum number of steps used in the preoptimiser before giving up and passing the 
+                calculation onto MLNEB with FHI-aims.
 
         Returns: Atoms object
             Transition state geometry structure
         """
 
         from catlearn.optimize.mlneb import MLNEB
+        from ase.io import write
 
         #TODO: calling mlneb.run() generates files in the current directory, reverting to os.chdir() necessary
         assert not self.nodes_per_instance, "ReactAims.nodes_per_instance is not None \n" \
@@ -382,10 +399,36 @@
 
         self._initialize_parameters(initial)
 
-
         """Set the environment parameters"""
         set_aims_command(hpc=hpc, basis_set=basis_set, defaults=2020, nodes_per_instance=self.nodes_per_instance)
 
+        '''Setup the TS calculation'''
+        helper = CalculationHelper(calc_type="TS",
+                                   parent_dir=os.getcwd(),
+                                   filename=self.filename,
+                                   restart=restart,
+                                   verbose=self.verbose)
+
+        counter, out, subdirectory_name, minimum_energy_path = helper.restart_setup()
+
+        """Set MACE Pre-optimisation flavor"""
+        self.mace_preopt_flavour = 0
+        if mace_preopt > 0 and not minimum_energy_path:
+            assert isinstance(n, int), "Integer number of images required for MACE TS preoptimiser"
+            assert self._MaceReactor is not None, "Please set MaceReact_Preoptimiser if mace_preopt is True."
+            assert input_check, "Mace Preoptimisation workflow requires input be set to 'float'."
+            self.mace_preopt_flavour = mace_preopt
+
+        """Run preoptimisation flavour 1"""
+        if self.mace_preopt_flavour == 1:
+
+            preopt_ts = self._mace_preoptimise_ts(initial, final, fmax, n, self.interpolation,
+                                                  input_check, max_steps=preopt_maxsteps)
+
+            initial = preopt_ts[0]
+            final = preopt_ts[-1]
+            self.mace_interpolation = preopt_ts
+
         if input_check:
             filename_copy = self.filename
             """Ensure input is converged"""
@@ -399,20 +442,21 @@
             """Set original name after input check is complete"""
             self.filename = filename_copy
 
-        '''Setup the TS calculation'''
-        helper = CalculationHelper(calc_type="TS",
-                                   parent_dir=os.getcwd(),
-                                   filename=self.filename,
-                                   restart=restart,
-                                   verbose=self.verbose)
+        """Run preotimisation flavour 2"""
+        if self.mace_preopt_flavour == 2:
+
+            preopt_ts = self._mace_preoptimise_ts(initial, final, fmax, n, self.interpolation,
+                                                  input_check, max_steps=preopt_maxsteps)
+
+            self.mace_interpolation = preopt_ts
 
-        counter, out, subdirectory_name, minimum_energy_path = helper.restart_setup()
         if not minimum_energy_path:
             minimum_energy_path = [None, None]
 
         traj_name = f"{subdirectory_name}/ML-NEB.traj"
 
         if not minimum_energy_path[0]:
+
             """Create the sockets calculator - using a with statement means the object is closed at the end."""
             with _calc_generator(params,
                                  out_fn=out,
@@ -432,11 +476,27 @@
 
                 iterations = 0
 
+                if self.mace_preopt_flavour > 0 and not os.path.exists(traj_name):
+                    """GAB: ML-NEB misbehaves if a calculator is not provided for interpolated images"""
+                    """     following function ensures correct calculators are attached with closed  """
+                    """     sockets.                                                                 """
+                    self.interpolation = [ image.copy() for image in self.mace_interpolation[1:-1] ]
+                    self.interpolation = [initial] + self.interpolation + [final]
+
+                    for idx, image in enumerate(self.interpolation):
+                        self.attach_calculator(self.interpolation[idx], params,
+                                           out_fn=out, dimensions=self.dimensions, directory=subdirectory_name,
+                                           calc_e=True)
+
+                    initial = self.interpolation[0]
+                    final   = self.interpolation[-1]
+
                 while not os.path.exists(traj_name):
                     if iterations > 0:
                         self.prev_calcs = read(f"{subdirectory_name}/last_predicted_path.traj@:")
 
                     os.chdir(subdirectory_name)
+
                     """Setup the Catlearn object for MLNEB"""
                     neb_catlearn = MLNEB(start=initial,
                                          end=final,
@@ -459,6 +519,8 @@
                         iterations += 1
                         os.chdir(parent_dir)
                     else:
+
+                        os.chdir(parent_dir)
                         return None
 
         """Find maximum energy, i.e. transition state to return it"""
@@ -470,7 +532,6 @@
 
         return self.ts
 
-
     def vibrate(self, atoms: Atoms, indices: list, read_only=False):
         """
 
@@ -544,6 +605,120 @@
             vib.write_mode()
             return vib
 
+    @property
+    def MaceReact_Preoptimiser(self):
+        """
+        Get or set the MACE preoptimiser used in aims_optimise or ts_search with setting
+
+        Args:
+             MaceReact MaceReact Obj:
+               User-defined MaceReact.
+
+        Returns:
+             self._MaceReactor MaceReact Obj:
+               User-defined MaceReact.
+        """
+
+        return self._MaceReactor
+
+    @MaceReact_Preoptimiser.setter
+    def MaceReact_Preoptimiser(self, MaceReact):
+
+        self._MaceReactor = MaceReact
+
+    def _mace_preoptimise(self, atoms: Atoms, fmax, relax_unit_cell, optimiser, opt_kwargs = {}):
+        """
+        Internal function for preoptimisation of structures with a pre-set MACE calculator
+
+        Args:
+            atoms: Atoms object
+            fmax: flaot
+            relax_unit_cell: bool
+            optimiser: bool or optimiser class
+            opt_kwargs: dict
+        Return:
+           preopt_atoms: Atoms object
+              Atoms object with the calculator object removed
+        """
+
+        filname = self._MaceReactor.filename
+        self._MaceReactor.filename = "MACE_PREOPT_" + self.filename
+
+        preopt_atoms = self._MaceReactor.mace_optimise(atoms, fmax, restart=False, relax_unit_cell=relax_unit_cell,
+                                                       optimiser = optimiser, opt_kwargs = opt_kwargs)
+
+        self._MaceReactor.filename = filname
+
+        # Clean calculator to prevent future problems
+        preopt_atoms.calc = None
+
+        return preopt_atoms
+
+    def _mace_preoptimise_ts(self, initial, final, fmax, n, interpolation, input_check, max_steps=200):
+        """
+        Internal function for preoptimisation of NEB with a pre-set MACE calculator
+
+        Args:
+            initial: Atoms object
+            final: Atoms object
+            fmax: float
+            n: int
+            interpolation: string, list of n Atoms objects 
+            input_check: None or float
+
+        Returns:
+
+        """
+
+        filname = self._MaceReactor.filename
+        self._MaceReactor.filename = "MACE_PREOPT_" + self.filename
+
+        if self.mace_preopt_flavour == 1:
+            input_check = input_check
+        elif self.mace_preopt_flavour == 2:
+            input_check = None
+
+        preopt_ts = self._MaceReactor.search_ts_neb(initial, final, fmax, n, k=0.05, method="improvedtangent",
+                                        interpolation=interpolation, input_check=input_check,
+                                        max_steps=max_steps, restart=True)
+
+        self._MaceReactor.filename = filname
+
+        return self._MaceReactor.interpolation.copy()
+
+    def attach_calculator(self, atoms, params,
+                    out_fn="aims.out",
+                    forces=True,
+                    dimensions=2,
+                    relax_unit_cell=False,
+                    directory=".",
+                    calc_e=False):
+        """
+        Potentially a redundant functionality, but condenses attaching and closing a socket
+        calculator to a given Atoms object
+
+        Args:
+            calc_e: Boolean, optional
+                Calculate total energy through Atoms.get_potential_energy().
+
+        Returns:
+            atoms: Atoms
+                Input atoms object with correctly closed socket.
+
+        """
+
+        with _calc_generator(params, out_fn=out_fn, forces=forces, dimensions=dimensions,
+                             relax_unit_cell=relax_unit_cell,directory=directory)[0] as calculator:
+
+            if self.dry_run:
+                calculator = EMT()
+
+            atoms.calc = calculator
+
+            if calc_e:
+                atoms.get_potential_energy()
+
+        return atoms
 
 def _calc_generator(params,
                     out_fn="aims.out",

examples/run_workflows_ReactAims_MACEpreopt.py

@@ -0,0 +1,69 @@
+'''
+This example shows how to use the work flow functionality with CARMM to optimise
+reactants/products and find transition states in an automated manner using FHI-aims.
+@Igor: Can you update with any further information?
+'''
+
+def test_run_workflows_ReactAims_MACE_preopt():
+    from carmm.run.workflows.react import ReactAims
+    from carmm.run.workflows.react_mace import ReactMACE
+    from ase.build import molecule, bulk
+    from carmm.analyse.forces import is_converged
+    import os
+    import numpy as np
+    from ase import Atoms
+    from ase.optimize import FIRE
+    from ase.build import surface, add_adsorbate
+    from ase.constraints import FixAtoms
+
+    '''Work in a dedicated folder'''
+    parent_dir = os.getcwd()
+    os.makedirs("data/react_preopt", exist_ok=True)
+    os.chdir("data/react_preopt")
+
+    '''Determine calculation input settings'''
+    params = {"xc":"pbe"}
+    basis_set = "light"
+    hpc = "hawk"
+
+    mace_params = {'model': 'small',
+                   'dispersion': False,
+                   'default_dtype': 'float64',
+                   'device': 'cpu'}
+
+    '''Prepare the Atoms object geometry'''
+    atoms = molecule("H2")
+
+    '''Build the React_Aims object using the set of settings'''
+    reactor = ReactAims(params, basis_set, hpc, dry_run=True)
+    reactor.MaceReact_Preoptimiser = ReactMACE(mace_params)
+
+    '''Call relevant calculations'''
+    '''The below has been previously calculated and data is retrieved from saved trajectories'''
+    model_optimised, _ = reactor.aims_optimise(atoms, fmax=0.01, restart=True, optimiser=FIRE, mace_preopt=True)
+    '''I do not like this test - it tests EMT more than it does the workflow'''
+    assert is_converged(reactor.model_optimised, 0.01), \
+        "Model is not optimised after preoptimisation"
+
+    '''Create a reaction pathway'''
+    atoms1 = atoms.copy()
+    atoms1[1].x += 8
+    atoms2 = atoms.copy()
+    '''Tests each flavour of ML-NEB'''
+    transition_state = reactor.search_ts(atoms1, atoms2, 0.05, 0.03, n=6, input_check=0.01, mace_preopt=2, restart=False)
+    print(len(reactor.mace_interpolation))
+    assert is_converged(reactor.mace_interpolation[3], 0.03), \
+        "MACE pre-optimised NEB flavour 2 is not converged"
+    transition_state = reactor.search_ts(atoms1, atoms2, 0.05, 0.03, n=6, input_check=0.01, mace_preopt=1, restart=False)
+    assert is_converged(reactor.mace_interpolation[3], 0.03), \
+        "MACE pre-optimised NEB flavour 1 is not converged"
+    transition_state = reactor.search_ts(atoms1, atoms2, 0.05, 0.03, n=6, input_check=0.01, mace_preopt=0, restart=False)
+    assert isinstance(reactor.interpolation, str), \
+        "ML-NEB is somehow running even though not requested"
+
+    '''Return to parent directory'''
+    os.chdir(parent_dir)
+
+
+test_run_workflows_ReactAims_MACE_preopt()
+