diff --git a/bin/emle-server b/bin/emle-server index 13b47d8..c36113a 100755 --- a/bin/emle-server +++ b/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, diff --git a/emle/__init__.py b/emle/__init__.py index 389e302..cdb813d 100644 --- a/emle/__init__.py +++ b/emle/__init__.py @@ -36,6 +36,7 @@ _supported_backends = [ "torchani", "mace", + "ace", "deepmd", "orca", "rascal", diff --git a/emle/_backends/__init__.py b/emle/_backends/__init__.py index 5be18e6..c66de39 100644 --- a/emle/_backends/__init__.py +++ b/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 * diff --git a/emle/_backends/_ace.py b/emle/_backends/_ace.py new file mode 100644 index 0000000..589345e --- /dev/null +++ b/emle/_backends/_ace.py @@ -0,0 +1,152 @@ +####################################################################### +# EMLE-Engine: https://github.com/chemle/emle-engine +# +# Copyright: 2023-2024 +# +# Authors: Lester Hedges +# Kirill Zinovjev +# +# 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 . +##################################################################### + +"""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 diff --git a/emle/_backends/_rascal.py b/emle/_backends/_rascal.py index b986735..3c0abf9 100644 --- a/emle/_backends/_rascal.py +++ b/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) diff --git a/emle/calculator.py b/emle/calculator.py index e00ad40..d4cdecd 100644 --- a/emle/calculator.py +++ b/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