Add Lennard-Jones implementation and training of dispersion coefficients

emle/_analyzer.py

@@ -49,7 +49,7 @@ def __init__(
         parser=None,
         q_total=None,
         start=None,
-        end=None
+        end=None,
     ):
         """
         Constructor.
@@ -168,7 +168,7 @@ def __init__(
             self.qm_xyz,
             self.q_total,
         )
-        self.atomic_alpha = 1. / _torch.diagonal(self.A_thole, dim1=1, dim2=2)[:, ::3]
+        self.atomic_alpha = 1.0 / _torch.diagonal(self.A_thole, dim1=1, dim2=2)[:, ::3]
         self.alpha = self._get_mol_alpha(self.A_thole, self.atomic_numbers)
 
         mask = (self.atomic_numbers > 0).unsqueeze(-1)

emle/_orca_parser.py

@@ -125,7 +125,6 @@ def __init__(self, filename, decompose=False, alpha=False):
 
         try:
             with _tarfile.open(filename, "r") as tar:
-
                 self._tar = tar
                 self.names = self._get_names(tar)
 

emle/calculator.py

@@ -107,6 +107,10 @@ def __init__(
         interpolate_steps=None,
         restart=False,
         device=None,
+        lj_mode=None,
+        lj_params_mm=None,
+        lj_params_qm=None,
+        lj_xyz_qm=None,
         orca_template=None,
         energy_frequency=0,
         energy_file="emle_energy.txt",
@@ -301,6 +305,26 @@ def __init__(
                 ! Angs NoUseSym
                 *xyzfile 0 1 inpfile.xyz
 
+        lj_mode: str
+            The mode to use for the Lennard-Jones parameters. Options are:
+                "fixed":
+                    Lennard-Jones parameters are fixed and provided in lj_params_qm.
+                "flexible":
+                    Lennard-Jones parameters are determined from a configuration.
+                    This requires lj_xyz_qm and atomic_numbers to be provided.
+
+        lj_params_mm: List[List[float]], Tuple[List[List[Float]]], numpy.ndarray, torch.Tensor
+            Lennard-Jones parameters for each atom in the MM region (sigma, epsilon) in units of nanometers (sigma)
+            and kJ/mol (epsilon). This is required for both "fixed" and "flexible" modes.
+
+        lj_params_qm: List[List[float]], Tuple[List[List[Float]]], numpy.ndarray, torch.Tensor
+            Lennard-Jones parameters for each atom in the QM region (sigma, epsilon) in units of nanometers (sigma)
+            and kJ/mol (epsilon). This is required if the "lj_mode" is "fixed" and lj_xyz_qm is not provided.
+            Takes precedence over lj_xyz_qm.
+
+        lj_xyz_qm: List[List[float]], Tuple[List[List[Float]]], numpy.ndarray, torch.Tensor
+            Positions of the atoms in the QM region for which the Lennard-Jones parameters are to be determined.
+
         energy_frequency: int
             The frequency of logging energies to file. If 0, then no energies are
             logged.
@@ -459,6 +483,47 @@ def __init__(
                 raise TypeError(msg)
             self._qm_charge = qm_charge
 
+        if lj_mode is not None:
+            if lj_mode.lower().replace(" ", "") not in ["fixed", "flexible"]:
+                msg = f"Unsupported Lennard-Jones mode: {lj_mode}. Options are: 'fixed', 'flexible'"
+                _logger.error(msg)
+                raise ValueError(msg)
+
+            self._lj_mode = lj_mode.lower().replace(" ", "")
+
+            if lj_params_mm is None:
+                msg = (
+                    "lj_params_mm must be provided if lj_mode is 'fixed' or 'flexible'"
+                )
+                _logger.error(msg)
+                raise ValueError(msg)
+            else:
+                if not isinstance(
+                    lj_params_mm, (list, tuple, _np.ndarray, _torch.Tensor)
+                ) or not isinstance(
+                    lj_params_mm[0], (list, tuple, _np.ndarray, _torch.Tensor)
+                ):
+                    msg = "lj_params_mm must be a list of lists, tuples, or arrays"
+                    _logger.error(msg)
+                    raise TypeError(msg)
+                self._lj_params_mm = _torch.tensor(
+                    lj_params_mm, dtype=self._device.dtype, device=self._device
+                )
+
+            if lambda_interpolate is not None:
+                if not isinstance(
+                    lj_params_mm, (list, tuple, _np.ndarray, _torch.Tensor)
+                ) or not isinstance(
+                    lj_params_mm[0], (list, tuple, _np.ndarray, _torch.Tensor)
+                ):
+                    msg = "When using interpolation, 'lj_params_qm' must be provided and must be a list of lists, tuples, or arrays."
+                    _logger.error(msg)
+                    raise TypeError(msg)
+
+                self._lj_params_qm = _torch.tensor(
+                    lj_params_qm, dtype=self._device.dtype, device=self._device
+                )
+
         # Create the EMLE model instance.
         self._emle = _EMLE(
             model=model,
@@ -468,6 +533,9 @@ def __init__(
             mm_charges=self._mm_charges,
             qm_charge=self._qm_charge,
             device=self._device,
+            lj_mode=lj_mode,
+            lj_params_qm=lj_params_qm,
+            lj_xyz_qm=lj_xyz_qm,
         )
 
         # Validate the backend(s).
@@ -868,6 +936,8 @@ def __init__(
                 method="mm",
                 mm_charges=self._mm_charges,
                 device=self._device,
+                lj_mode="fixed" if self._lj_mode is not None else None,
+                lj_params_qm=self._lj_params_qm,
             )
 
         else:
@@ -1131,6 +1201,7 @@ def _calculate_energy_and_gradients(
         xyz_qm,
         xyz_mm,
         atoms=None,
+        lj_params_mm=None,
         charge=0,
     ):
         """
@@ -1154,6 +1225,9 @@ def _calculate_energy_and_gradients(
         atoms: ase.Atoms
             The atoms object for the QM region.
 
+        lj_params_mm: numpy.ndarray, (N_MM_ATOMS, 2, 2)
+            The LJ parameters for the MM atoms.
+
         charge: int
             The total charge of the QM region.
 
@@ -1286,7 +1360,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, lj_params_mm, charge
+                    )
                     dE_dxyz_qm, dE_dxyz_mm = _torch.autograd.grad(
                         E.sum(), (xyz_qm, xyz_mm), allow_unused=allow_unused
                     )
@@ -1327,12 +1403,15 @@ def _calculate_energy_and_gradients(
                 _logger.error(msg)
                 raise RuntimeError(msg)
 
-            if (self._qbc_deviation):
+            if self._qbc_deviation:
                 E_std = _torch.std(base_model._E_vac_qbc).item()
                 max_f_std = _torch.max(_torch.std(base_model._grads_qbc, axis=0)).item()
                 with open(self._qbc_deviation, "a") as f:
                     f.write(f"{E_std:12.5f}{max_f_std:12.5f}\n")
-                if self._qbc_deviation_threshold and max_f_std > self._qbc_deviation_threshold:
+                if (
+                    self._qbc_deviation_threshold
+                    and max_f_std > self._qbc_deviation_threshold
+                ):
                     msg = "Force deviation threshold reached!"
                     raise ValueError(msg)
 
@@ -1371,7 +1450,7 @@ def _calculate_energy_and_gradients(
             else:
                 offset = int(not self._restart)
                 lam = self._lambda_interpolate[0] + (
-                    (self._step / (self._interpolate_steps - offset))
+                    self._step / (self._interpolate_steps - offset)
                 ) * (self._lambda_interpolate[1] - self._lambda_interpolate[0])
                 if lam < 0.0:
                     lam = 0.0
@@ -1503,12 +1582,17 @@ def _sire_callback(self, atomic_numbers, charges_mm, xyz_qm, xyz_mm, idx_mm=None
                 _logger.error(msg)
                 raise ValueError(msg)
 
+        if self._lj_mode is not None:
+            # Determine LJ parameters for the QM region.
+            lj_params_mm = self._lj_params_mm[idx_mm, :, :]
+
         # Compute the energy and gradients.
         E_vac, grad_vac, E_tot, grad_qm, grad_mm = self._calculate_energy_and_gradients(
             atomic_numbers,
             charges_mm,
             xyz_qm,
             xyz_mm,
+            lj_params_mm=lj_params_mm,
         )
 
         # Store the number of MM atoms.

emle/models/_ani.py

@@ -375,7 +375,7 @@ def forward(
         -------
 
         result: torch.Tensor (3,) or (3, BATCH)
-            The ANI2x and static and induced EMLE energy components in Hartree.
+            The ANI2x and static, induced and LJ EMLE energy components in Hartree.
         """
         # Batch the inputs if necessary.
         if atomic_numbers.ndim == 1:
@@ -407,4 +407,4 @@ def forward(
         E_emle = self._emle(atomic_numbers, charges_mm, xyz_qm, xyz_mm, qm_charge)
 
         # Return the ANI2x and EMLE energy components.
-        return _torch.stack((E_vac, E_emle[0], E_emle[1]))
+        return _torch.stack((E_vac, E_emle[0], E_emle[1], E_emle[2]))