MACE interface

MACE is an equivariant message passing neural network potential and software that looks very promising.

ASH features an interface to MACE via the PyTorch implementation that allows both the direct use of pre-trained MACE models to be used in ASH (just like any other theory) as well as direct training of a MACE-potential.

Energies and gradients can be requested (just like a regular QM or MM theory) and so a valid MACETheory object can be used for single-point energies, geometry optimizations, numerical frequencies, surface scans, NEB, molecular dynamics etc. within ASH. Even hybrid ONIOM and QM/MM calculations are possible (with some limitations).

Periodic boundary conditions are now also supported for running pretrained models.

MACETheory class:

def __init__(self, config_filename="config.yml",
             model_name=None, model_name_subtype=None, model_name_head=None,
             model_file=None, printlevel=2, mace_load_dispersion=False,
             label="MACETheory", numcores=1, platform="cpu", device=None, return_zero_gradient=False,
             polarmace=False, default_dtype="float64",
             energy_weight=None, forces_weight=None, max_num_epochs=None, valid_fraction=None,
             periodic=False, periodic_cell_vectors=None, periodic_cell_dimensions=None):

Keyword	Type	Default value	Details
config_filename	string	'config.yml'	Filename used for for MACE training configuration file.
model_file	string	None	Name of file used to load model (when using) or file created when training new model.
model_name	string	None	Name of foundational model used to load.
model_name_subtype	string	None	Name of foundational model subtype to choose (when appropriate).
model_name_head	string	None	For multihead models, which head to choose.
printlevel	integer	2	Printlevel
label	string	'MACETheory'	Label used for object
platform	string	'cpu'	Platform (device) used for training or running of MACE model within PyTorch. Options are: 'cuda', 'opencl', 'mps', 'cpu'.
default_dtype	string	'float64'	Whether to use float64 or float32 in Pytorch.
numcores	integer	1	Number of CPU cores used (if device is CPU)
energy_weight	float	None	Weight on energy for training
forces_weight	float	None	Weight on forces for training
max_num_epochs	integer	None	Max number of epochs to use for training
periodic	Boolean	False	Whether to use PBCs or not
periodic_cell_vectors	numpy array	None	Cell vectors as 3x3 numpy array in Angstrom.
periodic_cell_vectors	list	None	Cell dimensions as list of cell lengths and angles in Angstrom and degrees.

MACETheory train method:

def train(self, config_file="config.yml", name="model",model="MACE", device='cpu',
                  valid_fraction=0.1, train_file="train_data_mace.xyz",E0s=None,
                  energy_key='energy_REF', forces_key='forces_REF',
                  energy_weight=1, forces_weight=100,
                  max_num_epochs=500, swa=True, batch_size=10,
                  max_L = 0, r_max = 5.0, num_channels=128,
                  results_dir= "MACE_models", checkpoints_dir = "MACE_models",
                  log_dir ="MACE_models", model_dir="MACE_models"):

MACE installation

See MACE installation instructions: https://mace-docs.readthedocs.io/en/latest/guide/installation.html#installations Most likely the installation can be simply carried out using pip:

pip install mace-torch

MACE training

Training a MACE model is performed by first creating a MACETheory object and then calling the train method of that object.

The train method features a lot of options with the following defaults.

def train(self, config_file="config.yml", name="model",model="MACE", device='cpu',
                  valid_fraction=0.1, train_file="train_data_mace.xyz",E0s=None,
                  energy_key='energy_REF', forces_key='forces_REF',
                  energy_weight=1, forces_weight=100,
                  max_num_epochs=500, swa=True, batch_size=10,
                  max_L = 0, r_max = 5.0, num_channels=128,
                  results_dir= "MACE_models", checkpoints_dir = "MACE_models",
                  log_dir ="MACE_models", model_dir="MACE_models"):

To train a MACE model one needs a file containing the training data (specify by train_file keyword). The file is in a modified XYZ format that contains coordinates, energies and forces. It is best to use the create_ML_training_data function (see Machine learning in ASH) to generate this file (named train_data_mace.xyz). One can then train a model in the simple manner below, which uses default training parameters (see above).

theory = MACETheory()
theory.train(train_file="train_data_mace.xyz")

Calling the train method calls MACE and begins training which can take a while depending on the size of the training set, size of molecules, number of epochs requested (500 by default) etc. The training uses the CPU by default (device='cpu') but can run much faster on the GPU using e.g. device='cuda'.

Loading pre-trained or foundational MACE models

Pre-trained MACE models can either be previously trained models by yourself or a general foundational model. If you simply want to load a MACE model in MACETheory, for the purpose of using it as a Theory object to run calculations, you can either load the model via file (model_file keyword) or by name in case of a foundational model (model_name keyword). For option model_name the model is looked up inside the mace-torch library and automatically downloaded. Sometimes it may be necessary also to specify a subtype (model_name_head) or a specific head in the case of multi-head models (model_name_head).

For information on available foundational models, see:

• MACE-docs Foundational models

• MACE-foundations repository

• MACE modeles on Hugging Face

Examples on different ways to load models into MACETheory:

frag = Fragment(xyzfile="acetone.xyz", charge=0, mult=1)

# Load a MACE foundational model via previously downloaded file (here MACE-OFF23)
theory = MACETheory(model_file="/path/to/MACE-foundational-models/MACE-OFF23_medium.model")
# Load a MACE foundational model via name (here MACE-OFF23)
theory = MACETheory(model_name="mace-off23")
# Load a MACE foundational model via name (here MACE-OFF23) and specify subtype (size)
theory = MACETheory(model_name="mace-off23", model_name_subtype="large")
# Load a MACE MH-1 foundational model via name and specify a specific head
theory = MACETheory(model_file="/path/to/MACE-foundational-models/mace-mh-1.model", model_name_head="omat_pbe")

Recognized model_name options by MACETheory in ASH:

• mace-off23 (model_name_subtype options: 'small', 'medium', 'large')

• mace-mp or mace-mh (model_name_subtype options: 'small', 'medium', 'large' or 'medium-mpa-0', 'mh-1', 'mh-0')

• mace-polar (model_name_subtype options: 'polar-1-s', 'polar-1-m', 'polar-1-l')

• mace-ani-cc

• mace_omol

Note that for mace-mp / mace-mh models, D3 dispersion can also be activated via the mace_load_dispersion =True keyword and the mace_dispersion_xc keyword can be used to specify the name of the DFT functional to use for evaluating the dispersion formula (e.g. *mace_dispersion_xc*="PBE" for PBE-D3 correction). Dispersion is often recommended for use with those models (see literature). Note that the torch-dftd library is needed for this (pip install torch-dftd) Code block below shows how to perform a mace-mh-1-omat-D3 model calculation (see preprint: https://arxiv.org/pdf/2510.25380):

# Choosing MACE-MH-1 model by name, subtype and head. And activating D3 dispersion
qm = MACETheory(model_name="mace-mh-1", model_name_subtype="mh-1", model_name_head="omat_pbe",
                mace_load_dispersion=True, mace_dispersion_xc="pbe")
# Choosing MACE-MH-1 model by file and head.  And activating D3 dispersion
qm = MACETheory(model_file="/path/to/MACE-foundational-models/mace-mh-1.model", model_name_head="omat_pbe",
                mace_load_dispersion=True, mace_dispersion_xc="pbe")

Note that for other MACE foundational models it is easiest to download them separately and select via model_file keyword.

MACE Examples

Train a MACE model

Here we train a MACE potential from scratch. We first need to create some basic training data for a small molecule in the form of geometries and energies and forces. Here we generate a high-temperature Wigner ensemble from a Hessian calculated at our target level of theory. This is an easy way of generating some relevant geometries (though unknown how good this choice is for any given molecule).

We then call the create_ML_training_data function to generate energies and forces for these geometries or possibly a subset (this is our training data).

After that we can create a MACETheory object and then call the train method of the object using the train_data_mace.xyz file, a multigeometry XYZ-file that contains coordinates, energies and forces (generated by create_ML_training_data). We can choose to either train on the CPU (device="cpu") or alternatively using a GPU (use device="cuda" or device="opencl"). Here we modify the weight on forces vs. energy to be 100:1 (forces_weight=100). The code below will create multiple model files but the file named model_stagetwo_compiled.model would typically be the one to use.

from ash import *

# Define initial fragment
frag = Fragment(xyzfile="3fgaba.xyz", charge=0, mult=1)

# Define a QM theory
xtb = xTBTheory(xtbmethod="GFN2-xTB")

###########################################
# Generate training data (Wigner ensemble)
###########################################
#Run Hessian
result_freq = NumFreq(theory=xtb, fragment=frag)
# Generate Wigner ensemble of 100 samples at elevated temperature
wigner_frags = wigner_distribution(fragment=frag, hessian=result_freq.hessian,
    temperature=500, num_samples=100)
# This creates file Wigner_traj.xyz

# Read the Wigner ensemle snapshots and generate energies and forces training data
create_ML_training_data(xyz_trajectory="Wigner_traj.xyz", theory_1=xtb, Grad=True)
# This creates file train_data_mace.xyz

###################################
# TRAIN MACE model
###################################
#Create MACE model object
macetheory = MACETheory()
#Train model
macetheory.train(train_file="train_data_mace.xyz", device="cpu", max_num_epochs="3", forces_weight=100)

Use a previously trained model

Here it is assumed that you have already either trained a model yourself previously (see above) or that you want to use a foundational MACE model from the literature.

from ash import *

frag = Fragment(xyzfile="3fgaba.xyz", charge=0, mult=1)

#Create MACE model object and point to modelfile
macetheory = MACETheory(model_file="model_stagetwo_compiled.model")

#Run a job with Theory
res = Optimize(theory=macetheory, fragment=frag)

Run PBC calculations with MACE foundational models

MACE foundational models can be used for periodic calculations. Example below activates PBCs in MACETheory object and then uses the geomeTRICOptimizer to optimize atom and cell positions.

from ash import *

numcores=1
frag = Fragment(xyzfile="ammonia.xyz", charge=0, mult=1)

#Cell
cell_vectors = np.array([[5.01336,0.0,0.0],[0.0,5.01336,0.0],[0.0,0.0,5.01336]])

#Periodic CP2KTheory definition with specified cell dimensions
theory = MACETheory(model_file="/path/to/MACE-foundational-models/MACE-POLAR-1-M.model",
                periodic=True, periodic_cell_vectors=cell_vectors)

# Here using the geomeTRIC Optimizer that recognizes that PBCs are active in MACETheory and optimizes atoms and cell positions
Optimizer(theory=theory, fragment=frag, coordsystem="hdlc")