updates

Author: calvinp0

arc/job/adapter.py

@@ -97,6 +97,7 @@ class JobEnum(str, Enum):
     # TS search methods
     autotst = 'autotst'  # AutoTST, 10.1021/acs.jpca.7b07361, 10.26434/chemrxiv.13277870.v2
     heuristics = 'heuristics'  # ARC's heuristics
+    crest = 'crest'  # CREST conformer/TS search
     kinbot = 'kinbot'  # KinBot, 10.1016/j.cpc.2019.106947
     gcn = 'gcn'  # Graph neural network for isomerization, https://doi.org/10.1021/acs.jpclett.0c00500
     user = 'user'  # user guesses

arc/job/adapters/common.py

@@ -41,7 +41,7 @@
                              'Cyclic_Ether_Formation': ['kinbot'],
                              'Cyclopentadiene_scission': ['gcn', 'xtb_gsm'],
                              'Diels_alder_addition': ['kinbot'],
-                             'H_Abstraction': ['heuristics', 'autotst'],
+                             'H_Abstraction': ['heuristics', 'autotst', 'crest'],
                              'carbonyl_based_hydrolysis': ['heuristics'],
                              'ether_hydrolysis': ['heuristics'],
                              'nitrile_hydrolysis': ['heuristics'],
@@ -77,7 +77,8 @@
 adapters_that_do_not_require_a_level_arg = ['xtb', 'torchani']
 
 # Default is "queue", "pipe" will be called whenever needed. So just list 'incore'.
-default_incore_adapters = ['autotst', 'gcn', 'heuristics', 'kinbot', 'psi4', 'xtb', 'xtb_gsm', 'torchani', 'openbabel']
+default_incore_adapters = ['autotst', 'crest', 'gcn', 'heuristics', 'kinbot', 'psi4', 'xtb', 'xtb_gsm', 'torchani',
+                           'openbabel']
 
 
 def _initialize_adapter(obj: 'JobAdapter',

arc/job/adapters/ts/__init__.py

@@ -1,4 +1,5 @@
 import arc.job.adapters.ts.autotst_ts
+import arc.job.adapters.ts.crest
 import arc.job.adapters.ts.gcn_ts
 import arc.job.adapters.ts.heuristics
 import arc.job.adapters.ts.kinbot_ts

arc/job/adapters/ts/crest.py

@@ -4,18 +4,29 @@
 Separated from heuristics so CREST can be conditionally imported and reused.
 """
 
+import datetime
 import os
 import re
 import time
-from typing import List
+from typing import TYPE_CHECKING, List, Optional, Union
 
 import numpy as np
 import pandas as pd
 
-from arc.common import get_logger
+from arc.common import almost_equal_coords, get_logger
 from arc.imports import settings, submit_scripts
+from arc.job.adapter import JobAdapter
+from arc.job.adapters.common import _initialize_adapter, ts_adapters_by_rmg_family
+from arc.job.factory import register_job_adapter
 from arc.job.local import check_job_status, submit_job
+from arc.plotter import save_geo
 from arc.species.converter import reorder_xyz_string, str_to_xyz, xyz_to_dmat, xyz_to_str
+from arc.species.species import ARCSpecies, TSGuess
+from arc.job.adapters.ts.heuristics import h_abstraction
+
+if TYPE_CHECKING:
+    from arc.level import Level
+    from arc.reaction import ARCReaction
 
 logger = get_logger()
 
@@ -38,6 +49,229 @@ def crest_available() -> bool:
     return bool(SERVERS.get("local")) and bool(CREST_PATH or CREST_ENV_PATH)
 
 
+class CrestAdapter(JobAdapter):
+    """
+    A class for executing CREST TS conformer searches based on heuristics-generated guesses.
+    """
+
+    def __init__(self,
+                 project: str,
+                 project_directory: str,
+                 job_type: Union[List[str], str],
+                 args: Optional[dict] = None,
+                 bath_gas: Optional[str] = None,
+                 checkfile: Optional[str] = None,
+                 conformer: Optional[int] = None,
+                 constraints: Optional[List] = None,
+                 cpu_cores: Optional[str] = None,
+                 dihedral_increment: Optional[float] = None,
+                 dihedrals: Optional[List[float]] = None,
+                 directed_scan_type: Optional[str] = None,
+                 ess_settings: Optional[dict] = None,
+                 ess_trsh_methods: Optional[List[str]] = None,
+                 execution_type: Optional[str] = None,
+                 fine: bool = False,
+                 initial_time: Optional[Union['datetime.datetime', str]] = None,
+                 irc_direction: Optional[str] = None,
+                 job_id: Optional[int] = None,
+                 job_memory_gb: float = 14.0,
+                 job_name: Optional[str] = None,
+                 job_num: Optional[int] = None,
+                 job_server_name: Optional[str] = None,
+                 job_status: Optional[List[Union[dict, str]]] = None,
+                 level: Optional['Level'] = None,
+                 max_job_time: Optional[float] = None,
+                 run_multi_species: bool = False,
+                 reactions: Optional[List['ARCReaction']] = None,
+                 rotor_index: Optional[int] = None,
+                 server: Optional[str] = None,
+                 server_nodes: Optional[list] = None,
+                 queue: Optional[str] = None,
+                 attempted_queues: Optional[List[str]] = None,
+                 species: Optional[List[ARCSpecies]] = None,
+                 testing: bool = False,
+                 times_rerun: int = 0,
+                 torsions: Optional[List[List[int]]] = None,
+                 tsg: Optional[int] = None,
+                 xyz: Optional[dict] = None,
+                 ):
+
+        self.incore_capacity = 50
+        self.job_adapter = 'crest'
+        self.command = None
+        self.execution_type = execution_type or 'incore'
+
+        if reactions is None:
+            raise ValueError('Cannot execute TS CREST without ARCReaction object(s).')
+
+        dihedral_increment = dihedral_increment or 30
+
+        _initialize_adapter(obj=self,
+                            is_ts=True,
+                            project=project,
+                            project_directory=project_directory,
+                            job_type=job_type,
+                            args=args,
+                            bath_gas=bath_gas,
+                            checkfile=checkfile,
+                            conformer=conformer,
+                            constraints=constraints,
+                            cpu_cores=cpu_cores,
+                            dihedral_increment=dihedral_increment,
+                            dihedrals=dihedrals,
+                            directed_scan_type=directed_scan_type,
+                            ess_settings=ess_settings,
+                            ess_trsh_methods=ess_trsh_methods,
+                            fine=fine,
+                            initial_time=initial_time,
+                            irc_direction=irc_direction,
+                            job_id=job_id,
+                            job_memory_gb=job_memory_gb,
+                            job_name=job_name,
+                            job_num=job_num,
+                            job_server_name=job_server_name,
+                            job_status=job_status,
+                            level=level,
+                            max_job_time=max_job_time,
+                            run_multi_species=run_multi_species,
+                            reactions=reactions,
+                            rotor_index=rotor_index,
+                            server=server,
+                            server_nodes=server_nodes,
+                            queue=queue,
+                            attempted_queues=attempted_queues,
+                            species=species,
+                            testing=testing,
+                            times_rerun=times_rerun,
+                            torsions=torsions,
+                            tsg=tsg,
+                            xyz=xyz,
+                            )
+
+    def write_input_file(self) -> None:
+        pass
+
+    def set_files(self) -> None:
+        pass
+
+    def set_additional_file_paths(self) -> None:
+        pass
+
+    def set_input_file_memory(self) -> None:
+        pass
+
+    def execute_incore(self):
+        self._log_job_execution()
+        self.initial_time = self.initial_time if self.initial_time else datetime.datetime.now()
+
+        supported_families = [key for key, val in ts_adapters_by_rmg_family.items() if 'crest' in val]
+
+        self.reactions = [self.reactions] if not isinstance(self.reactions, list) else self.reactions
+        for rxn in self.reactions:
+            if rxn.family not in supported_families:
+                logger.warning(f'The CREST TS search adapter does not support the {rxn.family} reaction family.')
+                continue
+            if any(spc.get_xyz() is None for spc in rxn.r_species + rxn.p_species):
+                logger.warning(f'The CREST TS search adapter cannot process a reaction if 3D coordinates of '
+                               f'some/all of its reactants/products are missing.\nNot processing {rxn}.')
+                continue
+            if not crest_available():
+                logger.warning('CREST is not available. Skipping CREST TS search.')
+                continue
+
+            if rxn.ts_species is None:
+                rxn.ts_species = ARCSpecies(label='TS',
+                                            is_ts=True,
+                                            charge=rxn.charge,
+                                            multiplicity=rxn.multiplicity,
+                                            )
+
+            tsg = TSGuess(method='CREST')
+            tsg.tic()
+
+            crest_job_dirs = []
+            xyz_guesses = h_abstraction(reaction=rxn,
+                                        dihedral_increment=self.dihedral_increment,
+                                        path=self.local_path)
+            if not xyz_guesses:
+                logger.warning(f'CREST TS search failed to generate any seed guesses for {rxn.label}.')
+                tsg.tok()
+                continue
+
+            for iteration, xyz_entry in enumerate(xyz_guesses):
+                if isinstance(xyz_entry, dict):
+                    xyz_guess = xyz_entry.get("xyz")
+                else:
+                    xyz_guess = xyz_entry
+                if xyz_guess is None:
+                    continue
+
+                df_dmat = convert_xyz_to_df(xyz_guess)
+                try:
+                    h_abs_atoms_dict = get_h_abs_atoms(df_dmat)
+                    crest_job_dir = crest_ts_conformer_search(
+                        xyz_guess,
+                        h_abs_atoms_dict["A"],
+                        h_abs_atoms_dict["H"],
+                        h_abs_atoms_dict["B"],
+                        path=self.local_path,
+                        xyz_crest_int=iteration,
+                    )
+                    crest_job_dirs.append(crest_job_dir)
+                except (ValueError, KeyError) as e:
+                    logger.error(f"Could not determine the H abstraction atoms, got:\n{e}")
+
+            if not crest_job_dirs:
+                logger.warning(f'CREST TS search failed to prepare any jobs for {rxn.label}.')
+                tsg.tok()
+                continue
+
+            crest_jobs = submit_crest_jobs(crest_job_dirs)
+            monitor_crest_jobs(crest_jobs)
+            xyz_guesses_crest = process_completed_jobs(crest_jobs)
+            tsg.tok()
+
+            for method_index, xyz in enumerate(xyz_guesses_crest):
+                if xyz is None:
+                    continue
+                unique = True
+                for other_tsg in rxn.ts_species.ts_guesses:
+                    if almost_equal_coords(xyz, other_tsg.initial_xyz):
+                        if 'crest' not in other_tsg.method.lower():
+                            other_tsg.method += ' and CREST'
+                        unique = False
+                        break
+                if unique:
+                    ts_guess = TSGuess(method='CREST',
+                                       index=len(rxn.ts_species.ts_guesses),
+                                       method_index=method_index,
+                                       t0=tsg.t0,
+                                       execution_time=tsg.execution_time,
+                                       success=True,
+                                       family=rxn.family,
+                                       xyz=xyz,
+                                       )
+                    rxn.ts_species.ts_guesses.append(ts_guess)
+                    save_geo(xyz=xyz,
+                             path=self.local_path,
+                             filename=f'CREST_{method_index}',
+                             format_='xyz',
+                             comment=f'CREST {method_index}, family: {rxn.family}',
+                             )
+
+            if len(self.reactions) < 5:
+                successes = [tsg for tsg in rxn.ts_species.ts_guesses if tsg.success and 'crest' in tsg.method.lower()]
+                if successes:
+                    logger.info(f'CREST successfully found {len(successes)} TS guesses for {rxn.label}.')
+                else:
+                    logger.info(f'CREST did not find any successful TS guesses for {rxn.label}.')
+
+        self.final_time = datetime.datetime.now()
+
+    def execute_queue(self):
+        self.execute_incore()
+
+
 def crest_ts_conformer_search(
     xyz_guess: dict,
     a_atom: int,
@@ -376,3 +610,6 @@ def get_h_abs_atoms(dataframe: pd.DataFrame) -> dict:
             }
         else:
             raise ValueError("No valid hydrogen atom found.")
+
+
+register_job_adapter('crest', CrestAdapter)