lollcat / mace-jax

Equivariant machine learning interatomic potentials in JAX.

MACE                       🚀 JAX

This repository contains a porting of MACE in jax developed by Mario Geiger and Ilyes Batatia.

⚠️ This repository is still in developement. Everything can change anytime.

Test without installing

pip install nox
nox

Installation

From github:

pip install git+https://github.com/ACEsuit/mace-jax

Or locally:

python setup.py develop

Usage

Training

To train a MACE model, you can use the run_train.py script:

python -m mace_jax.run_train config.gin

An example of configuration file is located in the directory configs.

Configuration

Links to the files containing the functions configured by the gin config file.

• flags
• logs
• datasets
• model and here
• loss if loss.energy_weight, loss.forces_weight or loss.stress_weight is nonzero the loss will be applied
• optimizer
• train

Contributions

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

References

If you use this code, please cite our papers:

@misc{Batatia2022MACE,
  title = {MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields},
  author = {Batatia, Ilyes and Kov{\'a}cs, D{\'a}vid P{\'e}ter and Simm, Gregor N. C. and Ortner, Christoph and Cs{\'a}nyi, G{\'a}bor},
  year = {2022},
  number = {arXiv:2206.07697},
  eprint = {2206.07697},
  eprinttype = {arxiv},
  doi = {10.48550/ARXIV.2206.07697},
  archiveprefix = {arXiv}
}
@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}
 }

Contact

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

License

MACE is published and distributed under the MIT.