• MACE
  • Table of contents
  • About MACE
  • Documentation
  • Installation
    • pip installation
    • pip installation from source
  • Usage
    • Training
    • Evaluation
  • Tutorials
  • CUDA acceleration with cuEquivariance
  • Weights and Biases for experiment tracking
  • Pretrained Foundation Models
    • MACE-MP: Materials Project Force Fields
      • Example usage in ASE
    • MACE-OFF: Transferable Organic Force Fields
      • Example usage in ASE
    • Finetuning foundation models
    • Latest recommended foundation models
    • MACE-freeze
  • Caching
  • Development
  • References
  • Contact
  • License

MACE provides fast and accurate machine learning interatomic potentials with higher order equivariant message passing.

This repository contains the MACE reference implementation developed by Ilyes Batatia, Gregor Simm, David Kovacs, and the group of Gabor Csanyi, and friends (see Contributors).

Also available:

• MACE in JAX, currently about 2x times faster at evaluation, but training is recommended in Pytorch for optimal performances.
• MACE layers for constructing higher order equivariant graph neural networks for arbitrary 3D point clouds.

A partial documentation is available at: https://mace-docs.readthedocs.io

• Python >= 3.7 (for openMM, use Python = 3.9)
• PyTorch >= 1.12 (training with float64 is not supported with PyTorch 2.1 but is supported with 2.2 and later, Pytorch 2.4.1 is not supported)

Make sure to install PyTorch. Please refer to the official PyTorch installation for the installation instructions. Select the appropriate options for your system.

This is the recommended way to install MACE.

pip install --upgrade pip
pip install mace-torch

Note: The homonymous package on PyPI has nothing to do with this one.

git clone https://github.com/ACEsuit/mace.git
pip install ./mace

To train a MACE model, you can use the mace_run_train script, which should be in the usual place that pip places binaries (or you can explicitly run python3 <path_to_cloned_dir>/mace/cli/run_train.py)

mace_run_train \
    --name="MACE_model" \
    --train_file="train.xyz" \
    --valid_fraction=0.05 \
    --test_file="test.xyz" \
    --config_type_weights='{"Default":1.0}' \
    --E0s='{1:-13.663181292231226, 6:-1029.2809654211628, 7:-1484.1187695035828, 8:-2042.0330099956639}' \
    --model="MACE" \
    --hidden_irreps='128x0e + 128x1o' \
    --r_max=5.0 \
    --batch_size=10 \
    --max_num_epochs=1500 \
    --stage_two \
    --start_stage_two=1200 \
    --ema \
    --ema_decay=0.99 \
    --amsgrad \
    --restart_latest \
    --device=cuda \

To give a specific validation set, use the argument --valid_file. To set a larger batch size for evaluating the validation set, specify --valid_batch_size.

To control the model's size, you need to change --hidden_irreps. For most applications, the recommended default model size is --hidden_irreps='256x0e' (meaning 256 invariant messages) or --hidden_irreps='128x0e + 128x1o'. If the model is not accurate enough, you can include higher order features, e.g., 128x0e + 128x1o + 128x2e, or increase the number of channels to 256. It is also possible to specify the model using the --num_channels=128 and --max_L=1keys.

It is usually preferred to add the isolated atoms to the training set, rather than reading in their energies through the command line like in the example above. To label them in the training set, set config_type=IsolatedAtom in their info fields. If you prefer not to use or do not know the energies of the isolated atoms, you can use the option --E0s="average" which estimates the atomic energies using least squares regression. Note that using fitted E0s corresponds to fitting the deviations of the atomic energies from the average, rather than fitting the atomization energy (which is the case when using isolated-atom E0s), and this will most likely result in less stable potentials for molecular dynamics applications.

If the keyword --stage_two (previously called swa) is enabled, the energy weight of the loss is increased for the last ~20% of the training epochs (from --start_stage_two epochs). This setting usually helps lower the energy errors.

The precision can be changed using the keyword --default_dtype, the default is float64 but float32 gives a significant speed-up (usually a factor of x2 in training).

The keywords --batch_size and --max_num_epochs should be adapted based on the size of the training set. The batch size should be increased when the number of training data increases, and the number of epochs should be decreased. An heuristic for initial settings, is to consider the number of gradient update constant to 200 000, which can be computed as $\text{max-num-epochs}*\frac{\text{num-configs-training}}{\text{batch-size}}$.

The code can handle training set with heterogeneous labels, for example containing both bulk structures with stress and isolated molecules. In this example, to make the code ignore stress on molecules, append to your molecules configuration a config_stress_weight = 0.0.

By default, a figure displaying the progression of loss and RMSEs during training, along with a scatter plot of the model's inferences on the train, validation, and test sets, will be generated in the results folder at the end of training. This can be disabled using --plot False. To track these metrics throughout training (excluding inference on the test set), you can enable periodic plotting for the train and validation sets by specifying --plot_frequency N, which updates the plots every Nth epoch.

To use Apple Silicon GPU acceleration make sure to install the latest PyTorch version and specify --device=mps.

For multi-GPU training, use the --distributed flag. This will use PyTorch's DistributedDataParallel module to train the model on multiple GPUs. Combine with on-line data loading for large datasets (see below). An example slurm script can be found in mace/scripts/distributed_example.sbatch.

Option to parse all or some arguments using a YAML is available. For example, to train a model using the arguments above, you can create a YAML file your_configs.yaml with the following content:

name: nacl
seed: 2024
train_file: train.xyz
stage_two: yes
start_stage_two: 1200
max_num_epochs: 1500
device: cpu
test_file: test.xyz
E0s:
  41: -1029.2809654211628
  38: -1484.1187695035828
  8: -2042.0330099956639
config_type_weights:
  Default: 1.0

And append to the command line --config="your_configs.yaml". Any argument specified in the command line will overwrite the one in the YAML file.

To evaluate your MACE model on an XYZ file, run the mace_eval_configs:

mace_eval_configs \
    --configs="your_configs.xyz" \
    --model="your_model.model" \
    --output="./your_output.xyz"

You can run our Colab tutorial to quickly get started with MACE.

We also have a more detailed Colab tutorials on:

• Introduction to MACE training and evaluation
• Introduction to MACE active learning and fine-tuning
• MACE theory and code (advanced)

MACE supports CUDA acceleration with the cuEquivariance library. To install the library and use the acceleration, see our documentation at https://mace-docs.readthedocs.io/en/latest/guide/cuda_acceleration.html.

If you have a large dataset that might not fit into the GPU memory it is recommended to preprocess the data on a CPU and use on-line dataloading for training the model. To preprocess your dataset specified as an xyz file run the preprocess_data.py script. An example is given here:

mkdir processed_data
python ./mace/scripts/preprocess_data.py \
    --train_file="/path/to/train_large.xyz" \
    --valid_fraction=0.05 \
    --test_file="/path/to/test_large.xyz" \
    --atomic_numbers="[1, 6, 7, 8, 9, 15, 16, 17, 35, 53]" \
    --r_max=4.5 \
    --h5_prefix="processed_data/" \
    --compute_statistics \
    --E0s="average" \
    --seed=123 \

To see all options and a little description of them run python ./mace/scripts/preprocess_data.py --help . The script will create a number of HDF5 files in the processed_data folder which can be used for training. There will be one folder for training, one for validation and a separate one for each config_type in the test set. To train the model use the run_train.py script as follows:

python ./mace/scripts/run_train.py \
    --name="MACE_on_big_data" \
    --num_workers=16 \
    --train_file="./processed_data/train.h5" \
    --valid_file="./processed_data/valid.h5" \
    --test_dir="./processed_data" \
    --statistics_file="./processed_data/statistics.json" \
    --model="ScaleShiftMACE" \
    --num_interactions=2 \
    --num_channels=128 \
    --max_L=1 \
    --correlation=3 \
    --batch_size=32 \
    --valid_batch_size=32 \
    --max_num_epochs=100 \
    --stage_two \
    --start_stage_two=60 \
    --ema \
    --ema_decay=0.99 \
    --amsgrad \
    --error_table='PerAtomMAE' \
    --device=cuda \
    --seed=123 \

If you would like to use MACE with Weights and Biases to log your experiments simply install with

pip install ./mace[wandb]

And specify the necessary keyword arguments (--wandb, --wandb_project, --wandb_entity, --wandb_name, --wandb_log_hypers)

We provide a series of pretrained foundation models for various applications. These models can be used directly for inference, or as a starting point for fine-tuning on a new dataset. Foundation models are a rapidly evolving field. Please look at the MACE-MP GitHub repository <https://github.com/ACEsuit/mace-mp/releases>_ and the MACE-OFF23 GitHub repository <https://github.com/ACEsuit/mace-off/releases>_ for the latest releases.

We have collaborated with the Materials Project (MP) to train a universal MACE potential covering 89 elements on 1.6 M bulk crystals in the MPTrj dataset selected from MP relaxation trajectories. The models are releaed on GitHub at https://github.com/ACEsuit/mace-mp. If you use them please cite our paper which also contains an large range of example applications and benchmarks.

from mace.calculators import mace_mp
from ase import build

atoms = build.molecule('H2O')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float32", device='cuda')
atoms.calc = calc
print(atoms.get_potential_energy())

There is a series (small, medium, large) transferable organic force fields. These can be used for the simulation of organic molecules, crystals and molecular liquids, or as a starting point for fine-tuning on a new dataset. The models are released under the ASL license. The models are releaed on GitHub at https://github.com/ACEsuit/mace-off. If you use them please cite our paper which also contains detailed benchmarks and example applications.

from mace.calculators import mace_off
from ase import build

atoms = build.molecule('H2O')
calc = mace_off(model="medium", device='cuda')
atoms.calc = calc
print(atoms.get_potential_energy())

To finetune one of the mace-mp-0 foundation model, you can use the mace_run_train script with the extra argument --foundation_model=model_type. For example to finetune the small model on a new dataset, you can use:

mace_run_train \
  --name="MACE" \
  --foundation_model="small" \
  --train_file="train.xyz" \
  --valid_fraction=0.05 \
  --test_file="test.xyz" \
  --energy_weight=1.0 \
  --forces_weight=1.0 \
  --E0s="average" \
  --lr=0.01 \
  --scaling="rms_forces_scaling" \
  --batch_size=2 \
  --max_num_epochs=6 \
  --ema \
  --ema_decay=0.99 \
  --amsgrad \
  --default_dtype="float32" \
  --device=cuda \
  --seed=3

Other options are "medium" and "large", or the path to a foundation model. If you want to finetune another model, the model will be loaded from the path provided --foundation_model=$path_model, all the hypers will be extracted automatically.

Installation

To install the MACE-freeze version of MACE, clone the mace-freeze branch into your work folder:

git clone -b mace-freeze https://github.com/7radians/mace-freeze.git
pip install ./mace-freeze

This functionality allows to freeze neural network layers/parameters for transfer learning or other applications.

Usage

Use the --freeze=<N> to freeze the layers from the first one to N inclusive. Freezing a layer prevents its parameters from being updated during training. For example, to freeze the first 5 layers, provide a positive integer:

--freeze=5

To freeze the last N layers, provide a negative integer:

--freeze=-5

For more fine-grained control, the --freeze_par argument allows freezing by parameter tensors or vectors that form the model layers. The log file will contain the hierarchical list of the layers and named parameters of the model within them, which can be used to identify specific components to freeze. For example, to freeze the first 6 parameters of the model, use:

--freeze_par=6 

--freeze_par takes integer values and works in the same manner as --freeze, freezing from 1 to N (if positive) or the last N (if negative).

Note

• By default, MACE-freeze assumes all layers and parameters are trainable. This is equivalent to: --freeze=0 or --freeze_par=0.
• The --freeze and --freeze_par are mutually exclusive. If both are provided, only --freeze will take effect.
• If you intend to use freezing as the sole fine-tuning strategy, ensure that --multiheads_finetuning=False.
• Conversely, if you are using the multiheads fine-tuning method independently, either omit the --freeze argument or explicitly set --freeze=0 in your training script.

Soft-freezing provides a flexible approach to fine-tuning models by assigning a reduced learning rate to selected layers, rather than fully freezing them. This can be useful when adapting pretrained models, enabling a smooth transition between frozen and actively trained parameters.

Overview

Soft-freezing scales down the learning rate in a subset of layers using a multiplicative factor. This allows certain layers to update their weights more slowly than others, preserving learned representations while still allowing gradual adaptation. It can be used:

• On its own.
• In conjunction with layer freezing.
• Alongside multi-head fine-tuning.

Configuration

• --soft_freeze=<N>: Specifies the number of layers to soft-freeze.
• --soft_freeze_factor=<float>: A scaling factor between 0 and 1 applied to the base learning rate (--lr) for soft-frozen parameters.
• --soft_freeze_swa will also apply the soft-freeze logic to the Stage Two training by applying --soft_freeze_factor scaling to --lr_swa. Omit this flag to make the Stage Two learning rate uniform across parameters.

Behavior with freezing

When combined with standard layer freezing (--freeze):

• The first --freeze layers are fully frozen (i.e., excluded from optimization).
• The soft-freezing begins immediately after the frozen layers. For example: --freeze=5 and --soft_freeze=1 will freeze layers 1–5, soft-freeze layer 6, and leave all subsequent layers fully trainable.

When combined with parameter-level freezing (--freeze_par):

• If a layer contains a mix of frozen and active parameters, the soft-freezing logic begins at the first layer with any active parameters.
• A single layer can contain both soft-frozen and fully frozen parameters, depending on the granularity of the freeze.

Logging

During training, the structure of the model, including its layers and parameters, is printed to the logs in a hierarchical format. Each parameter is annotated with its learning rate, showing:

• Fully frozen parameters (either showing the baseline --lr or n/a, as the learning rate is irrelevant for parameters that are not updated)
• Soft-frozen parameters (lr = --lr × --soft_freeze_factor)
• Actively trained parameters (lr = --lr) This detailed breakdown allows for informed decision-making about the fine-tuning strategy applied to each part of the model.

If you use the freezing functionality, please cite this paper:

By default automatically downloaded models, like mace_mp, mace_off and data for fine tuning, end up in ~/.cache/mace. The path can be changed by using the environment variable XDG_CACHE_HOME. When set, the new cache path expands to $XDG_CACHE_HOME/.cache/mace

This project uses pre-commit to execute code formatting and linting on commit. We also use black, isort, pylint, and mypy. We recommend setting up your development environment by installing the dev packages into your python environment:

pip install -e ".[dev]"
pre-commit install

The second line will initialise pre-commit to automaticaly run code checks on commit. We have CI set up to check this, but we highly recommend that you run those commands before you commit (and push) to avoid accidentally committing bad code.

We are happy to accept pull requests under an MIT license. Please copy/paste the license text as a comment into your pull request.

If you use this code, please cite our papers:

@inproceedings{Batatia2022mace,
  title={{MACE}: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields},
  author={Ilyes Batatia and David Peter Kovacs and Gregor N. C. Simm and Christoph Ortner and Gabor Csanyi},
  booktitle={Advances in Neural Information Processing Systems},
  editor={Alice H. Oh and Alekh Agarwal and Danielle Belgrave and Kyunghyun Cho},
  year={2022},
  url={https://openreview.net/forum?id=YPpSngE-ZU}
}

@misc{Batatia2022Design,
  title = {The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials},
  author = {Batatia, Ilyes and Batzner, Simon and Kov{\'a}cs, D{\'a}vid P{\'e}ter and Musaelian, Albert and Simm, Gregor N. C. and Drautz, Ralf and Ortner, Christoph and Kozinsky, Boris and Cs{\'a}nyi, G{\'a}bor},
  year = {2022},
  number = {arXiv:2205.06643},
  eprint = {2205.06643},
  eprinttype = {arxiv},
  doi = {10.48550/arXiv.2205.06643},
  archiveprefix = {arXiv}
 }

If you have any questions, please contact us at [email protected].

For bugs or feature requests, please use GitHub Issues.

The MACE code is published and distributed under the MIT License. (Note that some of the models linked above come with different licenses).