Skip to content

Commit

Permalink
Make sure order of backends is reversible for delta learning.
Browse files Browse the repository at this point in the history
  • Loading branch information
lohedges committed Dec 13, 2024
1 parent a304697 commit d43edd6
Show file tree
Hide file tree
Showing 3 changed files with 122 additions and 31 deletions.
68 changes: 40 additions & 28 deletions emle/calculator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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())
Expand All @@ -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
Expand All @@ -1064,23 +1068,23 @@ 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
)

# Compute the gradients.
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
Expand All @@ -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:
Expand All @@ -1109,26 +1115,32 @@ 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)
dE_dxyz_qm, dE_dxyz_mm = _torch.autograd.grad(
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:
Expand Down Expand Up @@ -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:
Expand All @@ -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)
Expand Down Expand Up @@ -1448,23 +1460,23 @@ 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
)

# Compute the gradients.
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
Expand Down Expand Up @@ -1504,17 +1516,17 @@ 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)
dE_dxyz_qm, dE_dxyz_mm = _torch.autograd.grad(
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}"
Expand Down
9 changes: 9 additions & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
76 changes: 73 additions & 3 deletions tests/test_delta_learning.py
Original file line number Diff line number Diff line change
@@ -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")
Expand Down Expand Up @@ -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)

0 comments on commit d43edd6

Please sign in to comment.