diff --git a/ash/constants.py b/ash/constants.py index 94a2a8f18..8a9edef5c 100644 --- a/ash/constants.py +++ b/ash/constants.py @@ -9,6 +9,7 @@ hartokj = 2625.4996394798254 hartree2j = 4.3597438e-18 hartoeV = 27.211386245988 +evtohar = 0.03674932217565499 # speed of light in cm/s c = 2.99792458e10 diff --git a/ash/interfaces/interface_torch.py b/ash/interfaces/interface_torch.py index cd283cede..5128f8c80 100644 --- a/ash/interfaces/interface_torch.py +++ b/ash/interfaces/interface_torch.py @@ -4,6 +4,7 @@ from ash.modules.module_coords import elemstonuccharges from ash.functions.functions_general import ashexit, BC,print_time_rel from ash.functions.functions_general import print_line_with_mainheader +import ash.constants # TODO: Make sure energy is a general thing in PyTorch model @@ -60,13 +61,20 @@ def __init__(self, filename="torch.pt", model_name=None, model_object=None, if model_file is not None: print("model_file:", model_file) - self.load_model(model_file) + if 'aimnet2' in str(model_file).lower(): + print("AIMNet2 model selected") + self.load_aimnet2_model(model_file=model_file) + else: + self.load_model(model_file) print("Model:", self.model) elif model_name is not None: print("model_name:", model_name) if 'ani' in str(model_name).lower(): print("ANI type model selected") self.load_ani_model(model_name) + elif 'aimnet2' in str(model_name).lower(): + print("AIMNet2 model selected") + self.load_aimnet2_model(model_name=model_name) else: print("Error: Unknown model_name") ashexit() @@ -113,14 +121,38 @@ def save_model(self,filename=None, index=None): torch.jit.save(compiled_model, filename) print("Torch saved model to file:", filename) + def load_aimnet2_model(self,model_name=None, model_file=None): + print("Aimnet2-type model requested") + print("Models available: aimnet2") + try: + from aimnet2calc import AIMNet2Calculator + except ImportError as e: + print("Import error message:", e) + print("Problem importing AIMNet2Calculator or torch libraries. Make sure you have installed AIMNet2 correctly") + ashexit() + + # Model selection + print("Model:", model_name) + print("File:", model_file) + if model_name is not None: + if model_name == 'aimnet2': + self.model = AIMNet2Calculator('aimnet2') + else: + print("Error: Unknown model and no model_file selected") + ashexit() + elif model_file is not None: + print("Loading file:", model_file) + self.model = AIMNet2Calculator(model_file) + def load_ani_model(self,model): print("ANI-type model requested") print("Models available: ANI1ccx, ANI1x and ANI2x") try: import torchani import torch - except ImportError: + except ImportError as e: print("Problem importing torchani libraries. Make sure you have installed torchani correctly") + print("Import error message:", e) ashexit() # Model selection @@ -174,35 +206,36 @@ def run(self, current_coords=None, current_MM_coords=None, MMcharges=None, qm_el else: qm_elems = elems - # Converting coordinates and element information to Torch tensors - # dtype has to be float32 to get compatible Tensor for coordinate and nuc-charges - coords_torch = torch.tensor(np.array([current_coords], dtype='float32'), requires_grad=True, device=self.device) - nuc_charges_torch = torch.tensor(np.array([elemstonuccharges(qm_elems)]), device=self.device) - # Call model to get energy - print("self.model:", self.model) - energy_tensor = self.model((nuc_charges_torch, coords_torch)).energies - #result = self.model(nuc_charges_torch, coords_torch) - #print("result ", result) - #print("result 0", result[0]) - #print("result 1", result[1]) - #print("result 0 0", result[0][0]) - #print("result 0 0", result[0].item()) - #print("stuff 0 energies", stuff[0].energies) - #print("result.energies", result.energies) - #exit() - self.energy = energy_tensor.item() - #print("self.energy:", self.energy) - #exit() - # Call Grad - if Grad: - gradient_tensor = torch.autograd.grad(energy_tensor.sum(), coords_torch)[0] - self.gradient = gradient_tensor.detach().cpu().numpy()[0] - - if Hessian: - print("Calculating Hessian (requires torchani)") - import torchani - self.hessian = torchani.utils.hessian(coords_torch , energies=energy_tensor) + + # AIMNet2 + if 'aimnet2' in str(self.model).lower(): + input_data = {'coord':current_coords, 'numbers':elemstonuccharges(qm_elems), 'charge':charge, 'mult':mult} + results = self.model(input_data, forces=Grad, stress=False, hessian=False) + # print("results:", results) + self.energy = results["energy"].item() / ash.constants.hartoeV + if Grad: + self.gradient = -1*(0.03674932217565499/1.88972612546)*results["forces"].detach().cpu().numpy() + if Hessian: + self.hessian = (0.03674932217565499/1.88972612546/1.88972612546)*results["hessian"].detach().cpu().numpy() + # TorchANI + else: + # Converting coordinates and element information to Torch tensors + # dtype has to be float32 to get compatible Tensor for coordinate and nuc-charges + coords_torch = torch.tensor(np.array([current_coords], dtype='float32'), requires_grad=True, device=self.device) + nuc_charges_torch = torch.tensor(np.array([elemstonuccharges(qm_elems)]), device=self.device) + energy_tensor = self.model((nuc_charges_torch, coords_torch)).energies + self.energy = energy_tensor.item() + + # Call Grad + if Grad: + gradient_tensor = torch.autograd.grad(energy_tensor.sum(), coords_torch)[0] + self.gradient = gradient_tensor.detach().cpu().numpy()[0] + + if Hessian: + print("Calculating Hessian (requires torchani)") + import torchani + self.hessian = torchani.utils.hessian(coords_torch , energies=energy_tensor) if PC is True: print("PC not supported yet") diff --git a/ash/knarr/KNARRjobs/neb.py b/ash/knarr/KNARRjobs/neb.py index fa7a9c296..7e5916d0b 100755 --- a/ash/knarr/KNARRjobs/neb.py +++ b/ash/knarr/KNARRjobs/neb.py @@ -487,11 +487,15 @@ def DoNEB(path, calculator, neb, optimizer, second_run=False): else: print(f"Highest energy image is {checkmax}. Not turning on CI") if startci: + print("CI is active as number:", ci) if (it % 5) == 0: + print("Checking if CI should be changed (every 5th iteration)") newci = np.argmax(path.GetEnergy()) # cant allow it to be at the end points - if newci != ci and newci != 0 and newci < path.GetNim() - 1: - ci = newci + if newci != ci: + if newci != 0 and newci < path.GetNim() - 1: + print(f"Changing CI from {ci} to {newci}") + ci = newci if restart_on_ci: reset_opt = True if startci != True: diff --git a/ash/modules/module_coords.py b/ash/modules/module_coords.py index 29b5a0df6..10f688675 100644 --- a/ash/modules/module_coords.py +++ b/ash/modules/module_coords.py @@ -3141,8 +3141,7 @@ def get_linkatom_positions(qm_mm_boundary_dict, qmatoms, coords, elems, linkatom #Determine linkatom distance if linkatom_method == 'ratio': - print("Linkatom method: ratio") - + #print("Linkatom method: ratio") if linkatom_ratio == 'Auto': print("Automatic ratio. Determining ratio based on dict of equilibrium distances") #TODO @@ -3154,7 +3153,6 @@ def get_linkatom_positions(qm_mm_boundary_dict, qmatoms, coords, elems, linkatom print("Determined ratio:", linkatom_ratio) print("not yet ready") ashexit() - print("Ratio used:", linkatom_ratio) r_QM1_MM1 = distance(qmatom_coords, mmatom_coords) # See https://www.ncbi.nlm.nih.gov/pmc/articles/PMC9314059/ @@ -3163,7 +3161,7 @@ def get_linkatom_positions(qm_mm_boundary_dict, qmatoms, coords, elems, linkatom linkatom_distance = distance(qmatom_coords, linkatom_coords) print("Linkatom distance (QM1-L) determined to be:", linkatom_distance) elif linkatom_method == 'simple': - print("Linkatom method: simple") + #print("Linkatom method: simple") if linkatom_simple_distance is None: #print("linkatom_simple_distance not set. Getting standard distance from dictionary for element:", elems[dict_item[0]]) #Getting from dict