From 73d4415390ff3bc12949816c18f9bd7515ac723a Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 26 Jun 2024 16:45:09 +0100
Subject: [PATCH] Fix issues from git stash application.

---
 emle/calculator.py | 28 +++++++++++++---------------
 1 file changed, 13 insertions(+), 15 deletions(-)

diff --git a/emle/calculator.py b/emle/calculator.py
index 648e95b..9c1e5aa 100644
--- a/emle/calculator.py
+++ b/emle/calculator.py
@@ -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
@@ -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
 
@@ -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)
 
 
@@ -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,
@@ -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)
 
@@ -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()
@@ -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