From d43edd618d35b882d24e8bd2ba1c29c47dca05bd Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Fri, 13 Dec 2024 15:30:51 +0000 Subject: [PATCH] Make sure order of backends is reversible for delta learning. --- emle/calculator.py | 68 +++++++++++++++++++------------- tests/conftest.py | 9 +++++ tests/test_delta_learning.py | 76 ++++++++++++++++++++++++++++++++++-- 3 files changed, 122 insertions(+), 31 deletions(-) diff --git a/emle/calculator.py b/emle/calculator.py index 7ab7de0..458b07f 100644 --- a/emle/calculator.py +++ b/emle/calculator.py @@ -1000,6 +1000,10 @@ def run(self, path=None): base_model = None delta_model = None + # Zero the in vacuo energy and gradients. + E_vac = 0 + grad_vac = _np.zeros_like(xyz_qm) + # Internal backends. if not self._is_external_backend: if self._backend is not None: @@ -1016,19 +1020,17 @@ def run(self, path=None): # This is a non-Torch backend. else: try: - energy, forces = backend(atomic_numbers, xyz_qm) # Add the in vacuo energy and gradients to the total. - try: - E_vac += energy[0] - grad_vac += -forces[0] - except: - E_vac = energy[0] - grad_vac = -forces[0] + energy, forces = backend(atomic_numbers, xyz_qm) + E_vac += energy[0] + grad_vac += -forces[0] except Exception as e: msg = f"Failed to calculate in vacuo energies using {self._backend[i]} backend: {e}" _logger.error(msg) raise RuntimeError(msg) + _logger.error(f"{i}: Energy in vacuo: {E_vac}") + # No backend. else: E_vac, grad_vac = 0.0, _np.zeros_like(xyz_qm.detatch().cpu()) @@ -1044,6 +1046,8 @@ def run(self, path=None): _logger.error(msg) raise RuntimeError(msg) + _logger.error(f"Energy in vacuo: {E_vac}") + # Store a copy of the atomic numbers and QM coordinates as NumPy arrays. atomic_numbers_np = atomic_numbers xyz_qm_np = xyz_qm @@ -1064,14 +1068,14 @@ def run(self, path=None): # Apply delta-learning corrections using an EMLE model. if delta_model is not None: - model = delta_model.__class__.__name__ + model = delta_model.original_name try: # Create null MM inputs. null_charges_mm = _torch.zeros_like(charges_mm) null_xyz_mm = _torch.zeros_like(xyz_mm) # Compute the energy. - E = self._backends[1]( + E = delta_model( atomic_numbers, null_charges_mm, xyz_qm, null_xyz_mm, charge ) @@ -1079,8 +1083,8 @@ def run(self, path=None): dE_dxyz_qm = _torch.autograd.grad(E.sum(), xyz_qm) # Compute the delta correction. - delta_E = dE_dxyz_qm.cpu().numpy() * _BOHR_TO_ANGSTROM - delta_grad = dE_dxyz_qm.cpu().numpy() * _BOHR_TO_ANGSTROM + delta_E = E.sum().detach().cpu().numpy() + delta_grad = dE_dxyz_qm[0].cpu().numpy() * _BOHR_TO_ANGSTROM # Apply the correction. E_vac += delta_E @@ -1091,6 +1095,8 @@ def run(self, path=None): _logger.error(msg) raise RuntimeError(msg) + _logger.error(f"Delta-learning correction: {delta_E}") + # Compute embedding energy and gradients. if base_model is None: try: @@ -1109,10 +1115,13 @@ def run(self, path=None): _logger.error(msg) raise RuntimeError(msg) + _logger.error(f"EMLE energy: {E.sum().detach().cpu().numpy()}") + _logger.error(f"Total energy: {E_tot}") + # Compute in vacuo and embedding energies and gradients in one go using # the EMLE Torch models else: - model = base_model.__class__.__name__ + model = base_model.original_name try: with _torch.jit.optimized_execution(False): E = base_model(atomic_numbers, charges_mm, xyz_qm, xyz_mm, charge) @@ -1120,15 +1129,18 @@ def run(self, path=None): E.sum(), (xyz_qm, xyz_mm) ) - grad_qm = dE_dxyz_qm.cpu().numpy() * _BOHR_TO_ANGSTROM + grad_qm = grad_vac + dE_dxyz_qm.cpu().numpy() * _BOHR_TO_ANGSTROM grad_mm = dE_dxyz_mm.cpu().numpy() * _BOHR_TO_ANGSTROM - E_tot = E.sum().detach().cpu().numpy() + E_tot = E_vac + E.sum().detach().cpu().numpy() except Exception as e: msg = f"Failed to compute {model} energies and gradients: {e}" _logger.error(msg) raise RuntimeError(msg) + _logger.error(f"EMLE energy: {E.sum().detach().cpu().numpy()}") + _logger.error(f"Total energy: {E_tot}") + if self._is_interpolate: # Compute the in vacuo MM energy and gradients for the QM region. if self._backend != None: @@ -1384,6 +1396,10 @@ def _sire_callback(self, atomic_numbers, charges_mm, xyz_qm, xyz_mm, idx_mm=None base_model = None delta_model = None + # Zero the in vacuo energy and gradients. + E_vac = 0.0 + grad_vac = _np.zeros_like(xyz_qm) + # Internal backends. if not self._is_external_backend: if self._backend is not None: @@ -1400,14 +1416,10 @@ def _sire_callback(self, atomic_numbers, charges_mm, xyz_qm, xyz_mm, idx_mm=None # This is a non-Torch backend. else: try: - energy, forces = backend(atomic_numbers, xyz_qm) # Add the in vacuo energy and gradients to the total. - try: - E_vac += energy[0] - grad_vac += -forces[0] - except: - E_vac = energy[0] - grad_vac = -forces[0] + energy, forces = backend(atomic_numbers, xyz_qm) + E_vac += energy[0] + grad_vac += -forces[0] except Exception as e: msg = f"Failed to calculate in vacuo energies using {self._backend[i]} backend: {e}" _logger.error(msg) @@ -1448,14 +1460,14 @@ def _sire_callback(self, atomic_numbers, charges_mm, xyz_qm, xyz_mm, idx_mm=None # Apply delta-learning corrections using an EMLE model. if delta_model is not None: - model = delta_model.__class__.__name__ + model = delta_model.original_name try: # Create null MM inputs. null_charges_mm = _torch.zeros_like(charges_mm) null_xyz_mm = _torch.zeros_like(xyz_mm) # Compute the energy. - E = self._backends[1]( + E = delta_model( atomic_numbers, null_charges_mm, xyz_qm, null_xyz_mm, charge ) @@ -1463,8 +1475,8 @@ def _sire_callback(self, atomic_numbers, charges_mm, xyz_qm, xyz_mm, idx_mm=None dE_dxyz_qm = _torch.autograd.grad(E.sum(), xyz_qm) # Compute the delta correction. - delta_E = dE_dxyz_qm.cpu().numpy() * _BOHR_TO_ANGSTROM - delta_grad = dE_dxyz_qm.cpu().numpy() * _BOHR_TO_ANGSTROM + delta_E = E.sum().detach().cpu().numpy() + delta_grad = dE_dxyz_qm[0].cpu().numpy() * _BOHR_TO_ANGSTROM # Apply the correction. E_vac += delta_E @@ -1504,7 +1516,7 @@ def _sire_callback(self, atomic_numbers, charges_mm, xyz_qm, xyz_mm, idx_mm=None # Compute in vacuo and embedding energies and gradients in one go using # the EMLE Torch models else: - model = base_model.__class__.__name__ + model = base_model.original_name try: with _torch.jit.optimized_execution(False): E = base_model(atomic_numbers, charges_mm, xyz_qm, xyz_mm, charge) @@ -1512,9 +1524,9 @@ def _sire_callback(self, atomic_numbers, charges_mm, xyz_qm, xyz_mm, idx_mm=None E.sum(), (xyz_qm, xyz_mm) ) - grad_qm = dE_dxyz_qm.cpu().numpy() * _BOHR_TO_ANGSTROM + grad_qm = grad_vac + dE_dxyz_qm.cpu().numpy() * _BOHR_TO_ANGSTROM grad_mm = dE_dxyz_mm.cpu().numpy() * _BOHR_TO_ANGSTROM - E_tot = E.sum().detach().cpu().numpy() + E_tot = E_vac + E.sum().detach().cpu().numpy() except Exception as e: msg = f"Failed to compute {model} energies and gradients: {e}" diff --git a/tests/conftest.py b/tests/conftest.py index 4dd140c..ff0c565 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -13,6 +13,15 @@ def wrapper(): yield + # Kill the EMLE server. + kill_server() + + +def kill_server(): + """ + Helper function to kill the EMLE server. + """ + # Kill the EMLE server. We do this manually rather than using emle-stop # because there is sometimes a delay in the termination of the server, # which causes the next test to fail. This only seems to happen when diff --git a/tests/test_delta_learning.py b/tests/test_delta_learning.py index a129763..209a605 100644 --- a/tests/test_delta_learning.py +++ b/tests/test_delta_learning.py @@ -1,17 +1,22 @@ +import math import os import shlex import shutil import subprocess import tempfile +import time +import yaml -def test_delta_learning(): +def test_delta_learning(backend="torchani,xtb"): """ Make sure that the server can run using two backends for the in vacuo vacuo calculation. The first is the "reference" backend, the second applies delta learning corrections. """ + from conftest import kill_server + with tempfile.TemporaryDirectory() as tmpdir: # Copy files to temporary directory. shutil.copyfile("tests/input/adp.parm7", tmpdir + "/adp.parm7") @@ -39,5 +44,70 @@ def test_delta_learning(): # Make sure that the process exited successfully. assert process.returncode == 0 - # Make sure that an energy file is written. - assert os.path.isfile(tmpdir + "/emle_energy.txt") + # Load the YAML file. + with open(f"{tmpdir}/emle_settings.yaml", "r") as f: + data = yaml.safe_load(f) + + # Make sure the backend is set correctly. + assert data["backend"] == ["torchani", "xtb"] + + # Read the energy and first and last gradient from the ORCA engrad file. + with open(f"{tmpdir}/orc_job.engrad", "r") as f: + lines = f.readlines() + result_ab = ( + float(lines[2].strip()) + + float(lines[5].strip()) + + float(lines[-1].strip()) + ) + + # Kill the server. (Try twice, since there is sometimes a delay.) + kill_server() + time.sleep(1) + kill_server() + + # Now swap the order of the backends. + + with tempfile.TemporaryDirectory() as tmpdir: + # Copy files to temporary directory. + shutil.copyfile("tests/input/adp.parm7", tmpdir + "/adp.parm7") + shutil.copyfile("tests/input/adp.rst7", tmpdir + "/adp.rst7") + shutil.copyfile("tests/input/emle_sp.in", tmpdir + "/emle_sp.in") + + # Copy the current environment to a new dictionary. + env = os.environ.copy() + + # Set environment variables. + env["EMLE_BACKEND"] = "xtb,torchani" + env["EMLE_ENERGY_FREQUENCY"] = "1" + + # Create the sander command. + command = "sander -O -i emle_sp.in -p adp.parm7 -c adp.rst7 -o emle.out" + + process = subprocess.run( + shlex.split(command), + cwd=tmpdir, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + ) + + # Make sure that the process exited successfully. + assert process.returncode == 0 + + # Load the YAML file. + with open(f"{tmpdir}/emle_settings.yaml", "r") as f: + data = yaml.safe_load(f) + + # Make sure the backend is set correctly. + assert data["backend"] == ["xtb", "torchani"] + + # Read the energy and first and last gradient from the ORCA engrad file. + with open(f"{tmpdir}/orc_job.engrad", "r") as f: + result_ba = ( + float(lines[2].strip()) + + float(lines[5].strip()) + + float(lines[-1].strip()) + ) + + # Make sure that the results are the same. + assert math.isclose(result_ab, result_ba, rel_tol=1e-6)