Package: mace-torch
MACETable of contentsAbout MACEDocumentationInstallationUsageTrainingEvaluationTutorialWeights and BiasesDevelopmentPretrained foundation modelsMACE-MP: Materials Project Force FieldsMACE-OFF: Transferable Organic Force FieldsReferencesContactLicenseAbout MACEMACE 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, and David Kovacs.Also available:MACE in JAX, currently about 2x times faster at evaluation, but training is recommended in Pytorch for optimal performances.MACE layersfor constructing higher order equivariant graph neural networks for arbitrary 3D point clouds.DocumentationA partial documentation is available at:https://mace-docs.readthedocs.ioInstallationRequirements:Python >= 3.7PyTorch>= 1.12(for openMM, use Python = 3.9)pip installationTo install viapip, follow the steps below:pipinstall--upgradepip
pipinstalltorchtorchvisiontorchaudio--index-urlhttps://download.pytorch.org/whl/cu118
pipinstallmace-torchFor CPU or MPS (Apple Silicon) installation, usepip install torch torchvision torchaudioinstead.conda installationIf you do not have CUDA pre-installed, it isrecommendedto follow the conda installation process:# Create a virtual environment and activate itcondacreate--namemace_env
condaactivatemace_env# Install PyTorchcondainstallpytorchtorchvisiontorchaudiopytorch-cuda=11.6-cpytorch-cnvidia# (optional) Install MACE's dependencies from Conda as wellcondainstallnumpyscipymatplotlibaseopt_einsumprettytablepandase3nn# Clone and install MACE (and all required packages)gitclonehttps://github.com/ACEsuit/mace.git
pipinstall./macepip installation from sourceTo install viapip, follow the steps below:# Create a virtual environment and activate itpython-mvenvmace-venvsourcemace-venv/bin/activate# Install PyTorch (for example, for CUDA 11.6 [cu116])pip3installtorchtorchvisiontorchaudio--extra-index-urlhttps://download.pytorch.org/whl/cu116# Clone and install MACE (and all required packages)gitclonehttps://github.com/ACEsuit/mace.git
pipinstall./maceNote:The homonymous package onPyPIhas nothing to do with this one.UsageTrainingTo train a MACE model, you can use themace_run_trainscript, which should be in the usual place that pip places binaries (or you can explicitly runpython3 <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\--swa\--start_swa=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 to256.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, setconfig_type=IsolatedAtomin 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.If the keyword--swais enabled, the energy weight of the loss is increased for the last ~20% of the training epochs (from--start_swaepochs). This setting usually helps lower the energy errors.The precision can be changed using the keyword--default_dtype, the default isfloat64butfloat32gives a significant speed-up (usually a factor of x2 in training).The keywords--batch_sizeand--max_num_epochsshould 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 aconfig_stress_weight = 0.0.To use Apple Silicon GPU acceleration make sure to install the latest PyTorch version and specify--device=mps.EvaluationTo evaluate your MACE model on an XYZ file, run themace_eval_configs:mace_eval_configs\--configs="your_configs.xyz"\--model="your_model.model"\--output="./your_output.xyz"TutorialYou can run ourColab tutorialto quickly get started with MACE.We also have a more detailed user and developer tutorial athttps://github.com/ilyes319/mace-tutorialsWeights and Biases for experiment trackingIf you would like to use MACE with Weights and Biases to log your experiments simply install withpipinstall./mace[wandb]And specify the necessary keyword arguments (--wandb,--wandb_project,--wandb_entity,--wandb_name,--wandb_log_hypers)Pretrained Foundation ModelsMACE-MP: Materials Project Force FieldsWe have collaborated with the Materials Project (MP) to train a universal MACE potential covering 89 elements on 1.6 M bulk crystals in theMPTrj datasetselected from MP relaxation trajectories.
The models are releaed on GitHub athttps://github.com/ACEsuit/mace-mp.
If you use them please citeour paperwhich also contains an large range of example applications and benchmarks.Example usage in ASEfrommace.calculatorsimportmace_mpfromaseimportbuildatoms=build.molecule('H2O')calc=mace_mp(model="medium",dispersion=False,default_dtype="float32",device='cuda')atoms.calc=calcprint(atoms.get_potential_energy())MACE-OFF: Transferable Organic Force FieldsThere 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 theASL license.
The models are releaed on GitHub athttps://github.com/ACEsuit/mace-off.
If you use them please citeour paperwhich also contains detailed benchmarks and example applications.Example usage in ASEfrommace.calculatorsimportmace_offfromaseimportbuildatoms=build.molecule('H2O')calc=mace_off(model="medium",device='cuda')atoms.calc=calcprint(atoms.get_potential_energy())DevelopmentWe useblack,isort,pylint, andmypy.
Run the following to format and check your code:bash./scripts/run_checks.shWe have CI set up to check this, but wehighlyrecommend that you run those commands
before you commit (and push) to avoid accidentally committing bad code.We are happy to accept pull requests under anMIT license. Please copy/paste the license text as a comment into your pull request.ReferencesIf 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}
 }ContactIf you have any questions, please contact us atilyes.batatia@ens-paris-saclay.fr.For bugs or feature requests, please useGitHub Issues.LicenseMACE is published and distributed under theMIT License.