update foundations with develop

.gitignore

@@ -18,3 +18,6 @@ build/
 .vscode/
 *.txt
 *.log
+
+# Distribution
+dist/

README.md

@@ -1,51 +1,70 @@
-<span style="font-size:larger;">MACE</span>
-========
+# <span style="font-size:larger;">MACE</span>
+
 [![GitHub release](https://img.shields.io/github/release/ACEsuit/mace.svg)](https://GitHub.com/ACEsuit/mace/releases/)
 [![Paper](https://img.shields.io/badge/Paper-NeurIPs2022-blue)](https://openreview.net/forum?id=YPpSngE-ZU)
 [![License](https://img.shields.io/badge/License-MIT%202.0-blue.svg)](https://opensource.org/licenses/mit)
 [![GitHub issues](https://img.shields.io/github/issues/ACEsuit/mace.svg)](https://GitHub.com/ACEsuit/mace/issues/)
 [![Documentation Status](https://readthedocs.org/projects/mace/badge/)](https://mace-docs.readthedocs.io/en/latest/)
 
-# Table of contents
+## Table of contents
+
 - [About MACE](#about-mace)
 - [Documentation](#documentation)
 - [Installation](#installation)
 - [Usage](#usage)
-    - [Training](#training)
-    - [Evaluation](#evaluation)
+  - [Training](#training)
+  - [Evaluation](#evaluation)
 - [Tutorial](#tutorial)
 - [Weights and Biases](#weights-and-biases-for-experiment-tracking)
 - [Development](#development)
+- [Pretrained foundation models](#pretrained-foundation-models)
+  - [MACE-MP: Materials Project Force Fields](#mace-mp-materials-project-force-fields)
+  - [MACE-OFF: Transferable Organic Force Fields](#mace-off-transferable-organic-force-fields)
 - [References](#references)
 - [Contact](#contact)
 - [License](#license)
 
+## About MACE
 
-##  About MACE
 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, and David Kovacs.
 
-Also available: 
-* [MACE in JAX](https://github.com/ACEsuit/mace-jax), currently about 2x times faster at evaluation, but training is recommended in Pytorch for optimal performances.
-* [MACE layers](https://github.com/ACEsuit/mace-layer) for constructing higher order equivariant graph neural networks for arbitrary 3D point clouds.
+Also available:
+
+- [MACE in JAX](https://github.com/ACEsuit/mace-jax), currently about 2x times faster at evaluation, but training is recommended in Pytorch for optimal performances.
+- [MACE layers](https://github.com/ACEsuit/mace-layer) for constructing higher order equivariant graph neural networks for arbitrary 3D point clouds.
 
 ## Documentation
 
-A partial documentation is available at: https://mace-docs.readthedocs.io/en/latest/
+A partial documentation is available at: https://mace-docs.readthedocs.io
 
 ## Installation
 
 Requirements:
-* Python >= 3.7
-* [PyTorch](https://pytorch.org/) >= 1.12
+
+- Python >= 3.7
+- [PyTorch](https://pytorch.org/) >= 1.12
 
 (for openMM, use Python = 3.9)
 
+### pip installation
+
+To install via `pip`, follow the steps below:
+
+```sh
+pip install --upgrade pip
+pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
+pip install mace-torch
+```
+
+For CPU or MPS (Apple Silicon) installation, use `pip install torch torchvision torchaudio` instead.
+
 ### conda installation
 
 If you do not have CUDA pre-installed, it is **recommended** to follow the conda installation process:
+
 ```sh
 # Create a virtual environment and activate it
 conda create --name mace_env
@@ -58,13 +77,14 @@ conda install pytorch torchvision torchaudio pytorch-cuda=11.6 -c pytorch -c nvi
 conda install numpy scipy matplotlib ase opt_einsum prettytable pandas e3nn
 
 # Clone and install MACE (and all required packages)
-git clone git@github.com:ACEsuit/mace.git 
+git clone https://github.com/ACEsuit/mace.git
 pip install ./mace
 ```
 
-### pip installation
+### pip installation from source
 
 To install via `pip`, follow the steps below:
+
 ```sh
 # Create a virtual environment and activate it
 python -m venv mace-venv
@@ -74,15 +94,15 @@ source mace-venv/bin/activate
 pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116
 
 # Clone and install MACE (and all required packages)
-git clone git@github.com:ACEsuit/mace.git
+git clone https://github.com/ACEsuit/mace.git
 pip install ./mace
 ```
 
 **Note:** The homonymous package on [PyPI](https://pypi.org/project/MACE/) has nothing to do with this one.
 
 ## Usage
 
-### Training 
+### Training
 
 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`)
 
@@ -108,21 +128,21 @@ mace_run_train \
     --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 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`. 
+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 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. 
+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.
 
-If the keyword `--swa` is enabled, the energy weight of the loss is increased for the last ~20% of the training epochs (from `--start_swa` epochs). This setting usually helps lower the energy errors. 
+If the keyword `--swa` is enabled, the energy weight of the loss is increased for the last ~20% of the training epochs (from `--start_swa` 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 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 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``.
+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`.
 
-To use Apple Silicon GPU acceleration make sure to install the latest PyTorch version and specify ``--device=mps``. 
+To use Apple Silicon GPU acceleration make sure to install the latest PyTorch version and specify `--device=mps`.
 
 ### Evaluation
 
@@ -137,22 +157,62 @@ mace_eval_configs \
 
 ## Tutorial
 
-You can run our [Colab tutorial](https://colab.research.google.com/drive/1D6EtMUjQPey_GkuxUAbPgld6_9ibIa-V?authuser=1#scrollTo=Z10787RE1N8T) to quickly get started with MACE.
+You can run our [Colab tutorial](https://colab.research.google.com/drive/1D6EtMUjQPey_GkuxUAbPgld6_9ibIa-V?authuser=1#scrollTo=Z10787RE1N8T) to quickly get started with MACE. 
+
+We also have a more detailed user and developer tutorial at https://github.com/ilyes319/mace-tutorials
 
 ## Weights and Biases for experiment tracking
 
-If you would like to use MACE with Weights and Biases to log your experiments simply install with 
+If you would like to use MACE with Weights and Biases to log your experiments simply install with
 
 ```sh
 pip install ./mace[wandb]
 ```
 
 And specify the necessary keyword arguments (`--wandb`, `--wandb_project`, `--wandb_entity`, `--wandb_name`, `--wandb_log_hypers`)
 
+
+## Pretrained Foundation Models
+
+### MACE-MP: Materials Project Force Fields
+
+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](https://figshare.com/articles/dataset/23713842) 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](https://arxiv.org/abs/2401.00096) which also contains an large range of example applications and benchmarks.
+
+#### Example usage in ASE
+```py
+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())
+```
+
+### MACE-OFF: Transferable Organic Force Fields
+
+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](https://github.com/gabor1/ASL). 
+The models are releaed on GitHub at https://github.com/ACEsuit/mace-off.
+If you use them please cite [our paper](https://arxiv.org/abs/2312.15211) which also contains detailed benchmarks and example applications.
+
+#### Example usage in ASE
+```py
+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())
+```
+
 ## Development
 
 We use `black`, `isort`, `pylint`, and `mypy`.
 Run the following to format and check your code:
+
 ```sh
 bash ./scripts/run_checks.sh
 ```
@@ -165,15 +225,15 @@ We are happy to accept pull requests under an [MIT license](https://choosealicen
 ## References
 
 If you use this code, please cite our papers:
+
 ```text
-@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}
+@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,

mace/__version__.py

@@ -1 +1 @@
-__version__ = "0.2.0"
+__version__ = "0.3.4"

mace/calculators/__init__.py

@@ -1,10 +1,11 @@
+from .foundations_models import mace_anicc, mace_mp, mace_off
 from .lammps_mace import LAMMPS_MACE
 from .mace import MACECalculator
-from .foundations_models import mace_mp, mace_anicc
 
 __all__ = [
     "MACECalculator",
     "LAMMPS_MACE",
     "mace_mp",
+    "mace_off",
     "mace_anicc",
 ]