Skip to content

Commit

Permalink
Fix issues from git stash application.
Browse files Browse the repository at this point in the history
  • Loading branch information
lohedges committed Jun 26, 2024
1 parent aeeef24 commit 73d4415
Showing 1 changed file with 13 additions and 15 deletions.
28 changes: 13 additions & 15 deletions emle/calculator.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ def __init__(self, ref_values, ref_features, n_ref, sigma, device):
device: torch device
The PyTorch device to use for calculations.
"""
self._device = device
self.device = device
self.ref_features = ref_features
Kinv = _torch.tensor(
self.get_Kinv(ref_features, sigma), dtype=_torch.float32, device=self.device
Expand Down Expand Up @@ -126,19 +126,19 @@ def __call__(self, mol_features, zid):
The values of the predicted property for each atom.
"""

result = _torch.zeros(
len(zid), dtype=_torch.float32, device=self.device, requires_grad=True
)
tmp = _torch.zeros(len(zid), dtype=_torch.float32, device=self.device)
for i in range(self.n_z):
n_ref = self.n_ref[i]
ref_features_z = _torch.tensor(
self.ref_features[i, :n_ref], dtype=_torch.float32, device=self.device
)
mol_soap_z = mol_soap[zid == i, :, None]
mol_features_z = mol_features[zid == i, :, None]

K_mol_ref2 = (ref_features_z @ mol_features_z) ** 2
K_mol_ref2 = K_mol_ref2.reshape(K_mol_ref2.shape[:-1])
result[zid == i] = K_mol_ref2 @ self.c[i, :n_ref] + self.ref_mean[i]
tmp[zid == i] = K_mol_ref2 @ self.c[i, :n_ref] + self.ref_mean[i]

result = _torch.tensor(tmp, dtype=_torch.float32, device=self.device, requires_grad=True)

return result

Expand Down Expand Up @@ -225,7 +225,7 @@ def __call__(self, z, xyz):
device=self._device,
)

aev = self._aev_computer.forward((atomic_numbers, coords)).aevs[0]
aev = self._aev_computer.forward((atomic_numbers, xyz)).aevs[0]
return aev / _torch.linalg.norm(aev, axis=1, keepdims=True)


Expand Down Expand Up @@ -1165,7 +1165,6 @@ def __init__(
# Store the settings as a dictionary.
self._settings = {
"model": None if model is None else self._model,
"features": self._features,
"method": self._method,
"backend": self._backend,
"external_backend": None if external_backend is None else external_backend,
Expand Down Expand Up @@ -1379,7 +1378,7 @@ def run(self, path=None):
)

# Get the features.
mol_features = self._get_features(atomic_numbers, xyz_qm_bohr)
mol_features = self._get_features(atomic_numbers, xyz_qm)
s = self._get_s(mol_features, self._species_id)
chi = self._get_chi(mol_features, self._species_id)

Expand All @@ -1391,13 +1390,13 @@ def run(self, path=None):
charges_mm, dtype=_torch.float32, device=self._device
)

# Convert units.
xyz_qm_bohr = xyz_qm * _ANGSTROM_TO_BOHR
xyz_qm_bohr2 *= _ANGSTROM_TO_BOHR

# Compute energy and gradients.
E = self._get_E(charges_mm, xyz_qm_bohr2, xyz_mm_bohr, s, chi)
E = self._get_E(charges_mm, xyz_qm_bohr, xyz_mm_bohr, s, chi)
dE_dxyz_qm_bohr, dE_dxyz_mm_bohr = _torch.autograd.grad(
E, (xyz_qm_bohr2, xyz_mm_bohr)
E, (xyz_qm_bohr, xyz_mm_bohr)
)
dE_dxyz_qm_bohr = dE_dxyz_qm_bohr.cpu().numpy()
dE_dxyz_mm_bohr = dE_dxyz_mm_bohr.cpu().numpy()
Expand Down Expand Up @@ -2689,9 +2688,8 @@ def _run_torchani(self, xyz, atomic_numbers):
# Flag that NNPOps is active.
self._nnpops_active = True

# Applyt the optimised AEV symmetry functions.
if self._features == "aev":
self._aev_computer = self._torchani_model.aev_computer
# Apply the optimised AEV symmetry functions.
self._aev_computer = self._torchani_model.aev_computer

# Compute the energy and gradient.
energy = self._torchani_model((atomic_numbers, coords)).energies
Expand Down

0 comments on commit 73d4415

Please sign in to comment.