Add support for passing cell kwarg to EMLE forward method.

Author: lohedges

emle/calculator.py

@@ -1131,6 +1131,7 @@ def _calculate_energy_and_gradients(
         xyz_qm,
         xyz_mm,
         atoms=None,
+        cell=None,
         charge=0,
     ):
         """
@@ -1154,6 +1155,9 @@ def _calculate_energy_and_gradients(
         atoms: ase.Atoms
             The atoms object for the QM region.
 
+        cell: numpy.ndarray, (3, 3)
+            The simulation cell vectors.
+
         charge: int
             The total charge of the QM region.
 
@@ -1249,6 +1253,8 @@ def _calculate_energy_and_gradients(
         xyz_mm = _torch.tensor(
             xyz_mm, dtype=_torch.float32, device=self._device, requires_grad=True
         )
+        if cell is not None:
+            cell = _torch.tensor(cell, dtype=_torch.float32, device=self._device)
 
         # Are there any MM atoms?
         allow_unused = len(charges_mm) == 0
@@ -1263,7 +1269,7 @@ def _calculate_energy_and_gradients(
 
                 # Compute the energy.
                 E = delta_model(
-                    atomic_numbers, null_charges_mm, xyz_qm, null_xyz_mm, charge
+                    atomic_numbers, null_charges_mm, xyz_qm, null_xyz_mm, cell, charge
                 )
 
                 # Compute the gradients.
@@ -1286,7 +1292,9 @@ def _calculate_energy_and_gradients(
         if base_model is None:
             try:
                 if len(xyz_mm) > 0:
-                    E = self._emle(atomic_numbers, charges_mm, xyz_qm, xyz_mm, charge)
+                    E = self._emle(
+                        atomic_numbers, charges_mm, xyz_qm, xyz_mm, cell, charge
+                    )
                     dE_dxyz_qm, dE_dxyz_mm = _torch.autograd.grad(
                         E.sum(), (xyz_qm, xyz_mm), allow_unused=allow_unused
                     )
@@ -1313,7 +1321,9 @@ def _calculate_energy_and_gradients(
             model = base_model.original_name
             try:
                 with _torch.jit.optimized_execution(False):
-                    E = base_model(atomic_numbers, charges_mm, xyz_qm, xyz_mm, charge)
+                    E = base_model(
+                        atomic_numbers, charges_mm, xyz_qm, xyz_mm, cell, charge
+                    )
                     dE_dxyz_qm, dE_dxyz_mm = _torch.autograd.grad(
                         E.sum(), (xyz_qm, xyz_mm), allow_unused=allow_unused
                     )
@@ -1357,7 +1367,7 @@ def _calculate_energy_and_gradients(
                 E_mm_qm_vac, grad_mm_qm_vac = 0.0, _np.zeros_like(xyz_qm_np)
 
             # Compute the embedding contributions.
-            E = self._emle_mm(atomic_numbers, charges_mm, xyz_qm, xyz_mm, charge)
+            E = self._emle_mm(atomic_numbers, charges_mm, xyz_qm, xyz_mm, cell, charge)
             dE_dxyz_qm, dE_dxyz_mm = _torch.autograd.grad(
                 E.sum(), (xyz_qm, xyz_mm), allow_unused=allow_unused
             )
@@ -1448,7 +1458,9 @@ def set_lambda_interpolate(self, lambda_interpolate):
         # Reset the first step flag.
         self._is_first_step = not self._restart
 
-    def _sire_callback(self, atomic_numbers, charges_mm, xyz_qm, xyz_mm, idx_mm=None):
+    def _sire_callback(
+        self, atomic_numbers, charges_mm, xyz_qm, xyz_mm, cell=None, idx_mm=None
+    ):
         """
         A callback function to be used with Sire.
 
@@ -1467,6 +1479,9 @@ def _sire_callback(self, atomic_numbers, charges_mm, xyz_qm, xyz_mm, idx_mm=None
         xyz_mm: [[float, float, float]]
             The coordinates of the MM atoms in Angstrom.
 
+        cell: [[float, float, float], [float, float, float], [float, float, float]]
+            The simulation box vectors.
+
         idx_mm: [int]
             A list of indices of the MM atoms in the QM/MM region.
             Note that len(idx_mm) <= len(charges_mm) since it only
@@ -1493,6 +1508,8 @@ def _sire_callback(self, atomic_numbers, charges_mm, xyz_qm, xyz_mm, idx_mm=None
         charges_mm = _np.array(charges_mm)
         xyz_qm = _np.array(xyz_qm)
         xyz_mm = _np.array(xyz_mm)
+        if cell is not None:
+            cell = _np.array(cell)
 
         # Make sure that the number of QM atoms matches the number of MM charges
         # when using mm embedding.
@@ -1512,6 +1529,7 @@ def _sire_callback(self, atomic_numbers, charges_mm, xyz_qm, xyz_mm, idx_mm=None
             charges_mm,
             xyz_qm,
             xyz_mm,
+            cell=cell,
         )
 
         # Store the number of MM atoms.

emle/models/_ani.py

@@ -348,6 +348,7 @@ def forward(
         charges_mm: Tensor,
         xyz_qm: Tensor,
         xyz_mm: Tensor,
+        cell: Optional[Tensor] = None,
         qm_charge: int = 0,
     ) -> Tensor:
         """
@@ -368,6 +369,9 @@ def forward(
         xyz_mm: torch.Tensor (N_MM_ATOMS, 3) or (BATCH, N_MM_ATOMS, 3)
             Positions of MM atoms in Angstrom.
 
+        cell: torch.Tensor (3, 3) or (BATCH, 3, 3), optional
+            The simulation cell vectors in Angstrom.
+
         qm_charge: int or torch.Tensor (BATCH,)
             The charge on the QM region.
 
@@ -404,7 +408,7 @@ def forward(
         self._emle._emle_base._emle_aev_computer._aev = self._ani2x.aev_computer._aev
 
         # Get the EMLE energy components.
-        E_emle = self._emle(atomic_numbers, charges_mm, xyz_qm, xyz_mm, qm_charge)
+        E_emle = self._emle(atomic_numbers, charges_mm, xyz_qm, xyz_mm, cell, qm_charge)
 
         # Return the ANI2x and EMLE energy components.
         return _torch.stack((E_vac, E_emle[0], E_emle[1]))

emle/models/_emle.py

@@ -34,7 +34,7 @@
 import torchani as _torchani
 
 from torch import Tensor
-from typing import Union
+from typing import Optional, Union
 
 from . import _patches
 from . import EMLEBase as _EMLEBase
@@ -417,6 +417,7 @@ def forward(
         charges_mm: Tensor,
         xyz_qm: Tensor,
         xyz_mm: Tensor,
+        cell: Optional[Tensor] = None,
         qm_charge: Union[int, Tensor] = 0,
     ) -> Tensor:
         """
@@ -437,6 +438,9 @@ def forward(
         xyz_mm: torch.Tensor (N_MM_ATOMS, 3) or (BATCH, N_MM_ATOMS, 3)
             Positions of MM atoms in Angstrom.
 
+        cell: torch.Tensor (3, 3) or (BATCH, 3, 3), optional
+            The simulation cell vectors in Angstrom.
+
         qm_charge: int or torch.Tensor (BATCH,)
             The charge on the QM region.
 
@@ -461,6 +465,14 @@ def forward(
 
         batch_size = self._atomic_numbers.shape[0]
 
+        # Ensure cell is a tensor and repeat for batch size if necessary.
+        if cell is not None:
+            if isinstance(cell, _torch.Tensor):
+                if cell.ndim == 2:
+                    cell = cell.repeat(batch_size, 1, 1).to(self._device)
+            else:
+                raise TypeError("'cell' must be of type 'torch.Tensor'")
+
         # Ensure qm_charge is a tensor and repeat for batch size if necessary
         if isinstance(qm_charge, int):
             qm_charge = _torch.full(

emle/models/_mace.py

@@ -31,7 +31,7 @@
 import torch as _torch
 import numpy as _np
 
-from typing import List
+from typing import List, Optional
 
 from ._emle import EMLE as _EMLE
 from ._emle import _has_nnpops
@@ -198,12 +198,21 @@ def __init__(
         )
 
         if not isinstance(mace_model, (list, tuple)):
-            mace_model = [mace_model] if mace_model is None or isinstance(mace_model, str) else None
+            mace_model = (
+                [mace_model]
+                if mace_model is None or isinstance(mace_model, str)
+                else None
+            )
 
-        if mace_model is None or any(not isinstance(i, (str, type(None))) for i in mace_model):
-            raise TypeError("'mace_model' must be a list, tuple, or str, with elements of type str or None")
+        if mace_model is None or any(
+            not isinstance(i, (str, type(None))) for i in mace_model
+        ):
+            raise TypeError(
+                "'mace_model' must be a list, tuple, or str, with elements of type str or None"
+            )
 
         from mace.tools.scripts_utils import extract_config_mace_model
+
         self._mace_models = _torch.nn.ModuleList()
         for model in mace_model:
             source_model = self._load_mace_model(model, device)
@@ -392,9 +401,9 @@ def to(self, *args, **kwargs):
         """
         self._emle = self._emle.to(*args, **kwargs)
         self._mace = self._mace.to(*args, **kwargs)
-        self._mace_models = _torch.nn.ModuleList([
-            model.to(*args, **kwargs) for model in self._mace_models
-        ])
+        self._mace_models = _torch.nn.ModuleList(
+            [model.to(*args, **kwargs) for model in self._mace_models]
+        )
         return self
 
     def cpu(self, **kwargs):
@@ -405,9 +414,9 @@ def cpu(self, **kwargs):
         self._mace = self._mace.cpu(**kwargs)
         if self._atomic_numbers is not None:
             self._atomic_numbers = self._atomic_numbers.cpu(**kwargs)
-        self._mace_models = _torch.nn.ModuleList([
-            model.cpu(**kwargs) for model in self._mace_models
-        ])
+        self._mace_models = _torch.nn.ModuleList(
+            [model.cpu(**kwargs) for model in self._mace_models]
+        )
         return self
 
     def cuda(self, **kwargs):
@@ -418,9 +427,9 @@ def cuda(self, **kwargs):
         self._mace = self._mace.cuda(**kwargs)
         if self._atomic_numbers is not None:
             self._atomic_numbers = self._atomic_numbers.cuda(**kwargs)
-        self._mace_models = _torch.nn.ModuleList([
-            model.cuda(**kwargs) for model in self._mace_models
-        ])
+        self._mace_models = _torch.nn.ModuleList(
+            [model.cuda(**kwargs) for model in self._mace_models]
+        )
         return self
 
     def double(self):
@@ -429,9 +438,9 @@ def double(self):
         """
         self._emle = self._emle.double()
         self._mace = self._mace.double()
-        self._mace_models = _torch.nn.ModuleList([
-            model.double() for model in self._mace_models
-        ])
+        self._mace_models = _torch.nn.ModuleList(
+            [model.double() for model in self._mace_models]
+        )
         return self
 
     def float(self):
@@ -440,9 +449,9 @@ def float(self):
         """
         self._emle = self._emle.float()
         self._mace = self._mace.float()
-        self._mace_models = _torch.nn.ModuleList([
-            model.float() for model in self._mace_models
-        ])
+        self._mace_models = _torch.nn.ModuleList(
+            [model.float() for model in self._mace_models]
+        )
         return self
 
     def forward(
@@ -451,6 +460,7 @@ def forward(
         charges_mm: Tensor,
         xyz_qm: Tensor,
         xyz_mm: Tensor,
+        cell: Optional[Tensor] = None,
         qm_charge: int = 0,
     ) -> Tensor:
         """
@@ -471,6 +481,9 @@ def forward(
         xyz_mm: torch.Tensor (N_MM_ATOMS, 3) or (BATCH, N_MM_ATOMS, 3)
             Positions of MM atoms in Angstrom.
 
+        cell: torch.Tensor (3, 3) or (BATCH, 3, 3), optional
+            The simulation cell vectors in Angstrom.
+
         qm_charge: int
             The charge on the QM region.
 
@@ -489,6 +502,8 @@ def forward(
             xyz_qm = xyz_qm.unsqueeze(0)
             xyz_mm = xyz_mm.unsqueeze(0)
             charges_mm = charges_mm.unsqueeze(0)
+            if cell is not None and cell.ndim == 2:
+                cell = cell.unsqueeze(0)
 
         # Store the number of batches.
         num_batches = atomic_numbers.shape[0]
@@ -497,8 +512,17 @@ def forward(
         num_models = len(self._mace_models)
 
         # Create tensors to store the data for QbC.
-        self._E_vac_qbc = _torch.empty(num_models, num_batches, dtype=self._dtype, device=device)
-        self._grads_qbc = _torch.empty(num_models, num_batches, xyz_qm.shape[1], 3, dtype=self._dtype, device=device)
+        self._E_vac_qbc = _torch.empty(
+            num_models, num_batches, dtype=self._dtype, device=device
+        )
+        self._grads_qbc = _torch.empty(
+            num_models,
+            num_batches,
+            xyz_qm.shape[1],
+            3,
+            dtype=self._dtype,
+            device=device,
+        )
 
         # Create tensors to store the results.
         results_E_vac = _torch.empty(num_batches, dtype=self._dtype, device=device)
@@ -516,6 +540,10 @@ def forward(
                 xyz_qm[i], None, self._r_max, self._dtype, device
             )
 
+            # Get the cell for this configuration.
+            if cell is not None:
+                cell = cell.to(self._dtype).to(device)
+
             if not _torch.equal(atomic_numbers[i], self._atomic_numbers):
                 # Update the node attributes if the atomic numbers have changed.
                 self._node_attrs = (
@@ -557,19 +585,25 @@ def forward(
             results_E_vac[i] = E_vac[0] * EV_TO_HARTREE
 
             # Decouple the positions from the computation graph for the next models.
-            input_dict["positions"] = input_dict["positions"].clone().detach().requires_grad_(True)
+            input_dict["positions"] = (
+                input_dict["positions"].clone().detach().requires_grad_(True)
+            )
 
             # Do inference for the other models.
             if len(self._mace_models) > 1:
                 for j, mace in enumerate(self._mace_models):
-                    E_vac_qbc = mace(input_dict, compute_force=False)["interaction_energy"]
+                    E_vac_qbc = mace(input_dict, compute_force=False)[
+                        "interaction_energy"
+                    ]
 
                     assert (
                         E_vac_qbc is not None
                     ), "The model did not return any energy. Please check the input."
 
                     # Calculate the gradients
-                    grads_qbc = _torch.autograd.grad([E_vac_qbc], [input_dict["positions"]])[0]
+                    grads_qbc = _torch.autograd.grad(
+                        [E_vac_qbc], [input_dict["positions"]]
+                    )[0]
                     assert grads_qbc is not None, "Gradient computation failed"
 
                     # Store the results.
@@ -585,7 +619,7 @@ def forward(
             else:
                 # Get the EMLE energy components.
                 E_emle = self._emle(
-                    atomic_numbers, charges_mm, xyz_qm, xyz_mm, qm_charge
+                    atomic_numbers, charges_mm, xyz_qm, xyz_mm, cell, qm_charge
                 )
                 results_E_emle_static[i] = E_emle[0][0]
                 results_E_emle_induced[i] = E_emle[1][0]