Add conditional ACE backend.

bin/emle-server

@@ -97,6 +97,10 @@ try:
     mace_model = os.getenv("EMLE_MACE_MODEL")
 except:
     mace_model = None
+try:
+    ace_model = os.getenv("EMLE_ACE_MODEL")
+except:
+    ace_model = None
 rascal_model = os.getenv("EMLE_RASCAL_MODEL")
 parm7 = os.getenv("EMLE_PARM7")
 try:
@@ -152,7 +156,8 @@ env = {
     "pc_xyz_file": pc_xyz_file,
     "qm_xyz_frequency": qm_xyz_frequency,
     "ani2x_model_index": ani2x_model_index,
-    "macemodel": mace_model,
+    "mace_model": mace_model,
+    "ace_model": ace_model,
     "rascal_model": rascal_model,
     "lambda_interpolate": lambda_interpolate,
     "interpolate_steps": interpolate_steps,
@@ -267,6 +272,12 @@ parser.add_argument(
     help="name of the MACE-OFF23 model, or path to the MACE model file",
     required=False,
 )
+parser.add_argument(
+    "--ace-model",
+    type=str,
+    help="path to the ACE model file",
+    required=False,
+)
 parser.add_argument(
     "--deepmd-model",
     type=str,

emle/__init__.py

@@ -36,6 +36,7 @@
 _supported_backends = [
     "torchani",
     "mace",
+    "ace",
     "deepmd",
     "orca",
     "rascal",

emle/_backends/__init__.py

@@ -24,6 +24,7 @@
 Backends for in-vacuo calculations.
 """
 
+from ._ace import *
 from ._deepmd import *
 from ._orca import *
 from ._rascal import *

emle/_backends/_ace.py

@@ -0,0 +1,152 @@
+#######################################################################
+# EMLE-Engine: https://github.com/chemle/emle-engine
+#
+# Copyright: 2023-2024
+#
+# Authors: Lester Hedges   <[email protected]>
+#          Kirill Zinovjev <[email protected]>
+#
+# EMLE-Engine is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 2 of the License, or
+# (at your option) any later version.
+#
+# EMLE-Engine is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with EMLE-Engine. If not, see <http://www.gnu.org/licenses/>.
+#####################################################################
+
+"""ACE in-vacuo backend implementation."""
+
+try:
+    import pyjulip as _pyjulip
+
+    __all__ = ["ACE"]
+except:
+    __all__ = []
+
+import ase as _ase
+import os as _os
+import numpy as _np
+
+
+from .._units import _EV_TO_HARTREE, _BOHR_TO_ANGSTROM
+
+from ._backend import Backend as _Backend
+
+
+class ACE(_Backend):
+    """
+    ACE in-vacuo backend implementation.
+    """
+
+    def __init__(self, model):
+        """
+        Initialize the ACE in-vacuo backend.
+
+        Parameters
+        ----------
+
+        model: str
+            The path to the ACE model.
+        """
+
+        # Validate the model path.
+        if not isinstance(model, str):
+            raise TypeError("'model' must be of type 'str'")
+        if not _os.path.exists(model):
+            raise FileNotFoundError(f"Model file '{model}' not found")
+
+        # Try to create the ACE calculator.
+
+        try:
+            self._calculator = _pyjulip.ACE(model)
+        except Exception as e:
+            raise RuntimeError(f"Failed to create ACE calculator: {e}")
+
+    def calculate(self, atomic_numbers, xyz, forces=True):
+        """
+        Compute the energy and forces.
+
+        Parameters
+        ----------
+
+        atomic_numbers: numpy.ndarray, (N_BATCH, N_QM_ATOMS,)
+            The atomic numbers of the atoms in the QM region.
+
+        xyz: numpy.ndarray, (N_BATCH, N_QM_ATOMS, 3)
+            The coordinates of the atoms in the QM region in Angstrom.
+
+        forces: bool
+            Whether to calculate and return forces.
+
+        Returns
+        -------
+
+        energy: float
+            The in-vacuo energy in Eh.
+
+        forces: numpy.ndarray
+            The in-vacuo forces in Eh/Bohr.
+        """
+
+        if not isinstance(atomic_numbers, _np.ndarray):
+            raise TypeError("'atomic_numbers' must be of type 'numpy.ndarray'")
+        if not isinstance(xyz, _np.ndarray):
+            raise TypeError("'xyz' must be of type 'numpy.ndarray'")
+
+        if len(atomic_numbers) != len(xyz):
+            raise ValueError(
+                f"Length of 'atomic_numbers' ({len(atomic_numbers)}) does not "
+                f"match length of 'xyz' ({len(xyz)})"
+            )
+
+        # Convert to batched NumPy arrays.
+        if len(atomic_numbers.shape) == 1:
+            atomic_numbers = _np.expand_dims(atomic_numbers, axis=0)
+            xyz = _np.expand_dims(xyz, axis=0)
+
+        # Lists to store results.
+        results_energy = []
+        results_forces = []
+
+        # Loop over batches.
+        for i, (atomic_numbers_i, xyz_i) in enumerate(zip(atomic_numbers, xyz)):
+            if len(atomic_numbers_i) != len(xyz_i):
+                raise ValueError(
+                    f"Length of 'atomic_numbers' ({len(atomic_numbers_i)}) does not "
+                    f"match length of 'xyz' ({len(xyz_i)}) for index {i}"
+                )
+
+            # Create an ASE atoms object.
+            atoms = _ase.Atoms(
+                numbers=atomic_numbers_i,
+                positions=xyz_i,
+            )
+
+            # The calculator requires periodic box information so we translate the atoms
+            # so that the lowest (x, y, z) position is zero, then set the cell to the
+            # maximum position.
+            atoms.positions -= _np.min(atoms.positions, axis=0)
+            atoms.cell = _np.max(atoms.positions, axis=0)
+
+            # Set the calculator.
+            atoms.calc = self._calculator
+
+            # Get the energy.
+            results_energy.append(atoms.get_potential_energy() * _EV_TO_HARTREE)
+
+            if forces:
+                results_forces.append(
+                    atoms.get_forces() * _EV_TO_HARTREE * _BOHR_TO_ANGSTROM
+                )
+
+        # Convert to NumPy arrays.
+        results_energy = _np.array(results_energy)
+        results_forces = _np.array(results_forces)
+
+        return results_energy, results_forces if forces else results_energy

emle/_backends/_rascal.py

@@ -140,9 +140,9 @@ def calculate(self, atomic_numbers, xyz, forces=True):
                 positions=xyz_i,
             )
 
-            # Rascal requires periodic box information so we translate the atoms so that
-            # the lowest (x, y, z) position is zero, then set the cell to the maximum
-            # position.
+            # The calculator requires periodic box information so we translate the atoms
+            # so that the lowest (x, y, z) position is zero, then set the cell to the
+            # maximum position.
             atoms.positions -= _np.min(atoms.positions, axis=0)
             atoms.cell = _np.max(atoms.positions, axis=0)
 

emle/calculator.py

@@ -63,6 +63,7 @@ class EMLECalculator:
     _supported_backends = [
         "torchani",
         "mace",
+        "ace",
         "deepmd",
         "orca",
         "rascal",
@@ -100,6 +101,7 @@ def __init__(
         qm_xyz_frequency=0,
         ani2x_model_index=None,
         mace_model=None,
+        ace_model=None,
         rascal_model=None,
         parm7=None,
         qm_indices=None,
@@ -216,6 +218,10 @@ def __init__(
             To use a locally trained MACE model, provide the path to the model file.
             If None, the MACE-OFF23(S) model will be used by default.
 
+        ace_model: str
+            Path to the ACE model file to use for in vacuo calculations. This
+            must be specified if 'ace' is the selected backend.
+
         rascal_model: str
             Path to the Rascal model file to use for in vacuo calculations. This
             must be specified if "rascal" is the selected backend.
@@ -638,6 +644,16 @@ def __init__(
                 # Store the MACE model.
                 self._mace_model = mace_model
 
+            elif backend == "ace":
+                try:
+                    from ._backends import ACE
+
+                    b = ACE(ace_model)
+                except:
+                    msg = "Unable to create ACE backend. Please ensure PyJulip is installed."
+                    _logger.error(msg)
+                    raise RuntimeError(msg)
+
             elif backend == "orca":
                 try:
                     from ._backends import ORCA