Merge pull request #142 from ACEsuit/develop

MACE 0.3.0 release

.gitignore

@@ -16,5 +16,5 @@ build/
 # IDE
 .idea/
 .vscode/
-logs/MACE_run-5.log
 *.txt
+*.log

README.md

@@ -90,10 +90,10 @@ pip install ./mace
 
 ### Training
 
-To train a MACE model, you can use the `run_train.py` script:
+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`)
 
 ```sh
-python ./mace/scripts/run_train.py \
+mace_run_train \
     --name="MACE_model" \
     --train_file="train.xyz" \
     --valid_fraction=0.05 \
@@ -132,18 +132,18 @@ To use Apple Silicon GPU acceleration make sure to install the latest PyTorch ve
 
 ### Evaluation
 
-To evaluate your MACE model on an XYZ file, run the `eval_configs.py`:
+To evaluate your MACE model on an XYZ file, run the `mace_eval_configs`:
 
 ```sh
-python3 ./mace/scripts/eval_configs.py \
+mace_eval_configs \
     --configs="your_configs.xyz" \
     --model="your_model.model" \
     --output="./your_output.xyz"
 ```
 
 ## 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
 

mace/calculators/__init__.py

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

mace/calculators/foundations_models.py

@@ -0,0 +1,47 @@
+import os
+
+from .mace import MACECalculator
+
+path = os.path.dirname(__file__)
+
+
+def mace_mp(
+    device="cuda",
+    model_path=None,
+) -> MACECalculator:
+    """
+    Constructs a MACECalculator with a pretrained model based on the Materials Project (89 elements).
+    The model is released under the MIT license.
+    Note:
+        If you are using this function, please cite the relevant paper for the Materials Project,
+        any paper associated with the MACE model, and also the following:
+        - "MACE-Universal by Yuan Chiang, 2023, Hugging Face, Revision e5ebd9b, DOI: 10.57967/hf/1202, URL: https://huggingface.co/cyrusyc/mace-universal"
+        - "Matbench Discovery by Janosh Riebesell, Rhys EA Goodall, Anubhav Jain, Philipp Benner, Kristin A Persson, Alpha A Lee, 2023, arXiv:2308.14920"
+    """
+    if model_path is None:
+        model_path = os.path.join(
+            path, "foundations_models/2023-08-14-mace-universal.model"
+        )
+        print(
+            "Using Materials Project model for MACECalculator, see https://huggingface.co/cyrusyc/mace-universal"
+        )
+    return MACECalculator(model_path, device=device, default_dtype="float32")
+
+
+def mace_anicc(
+    device="cuda",
+    model_path=None,
+) -> MACECalculator:
+    """
+    Constructs a MACECalculator with a pretrained model based on the ANI (H, C, N, O).
+    The model is released under the MIT license.
+    Note:
+        If you are using this function, please cite the relevant paper associated with the MACE model, ANI dataset, and also the following:
+        - "Evaluation of the MACE Force Field Architecture by Dávid Péter Kovács, Ilyes Batatia, Eszter Sára Arany, and Gábor Csányi, The Journal of Chemical Physics, 2023, URL: https://doi.org/10.1063/5.0155322
+    """
+    if model_path is None:
+        model_path = os.path.join(path, "foundations_models/ani500k_large_CC.model")
+        print(
+            "Using ANI couple cluster model for MACECalculator, see https://doi.org/10.1063/5.0155322"
+        )
+    return MACECalculator(model_path, device=device, default_dtype="float64")

mace/calculators/lammps_mace.py

@@ -3,7 +3,6 @@
 import torch
 from e3nn.util.jit import compile_mode
 
-from mace.modules.utils import get_outputs
 from mace.tools.scatter import scatter_sum
 
 
@@ -15,20 +14,19 @@ def __init__(self, model):
         self.register_buffer("atomic_numbers", model.atomic_numbers)
         self.register_buffer("r_max", model.r_max)
         self.register_buffer("num_interactions", model.num_interactions)
+        for param in self.model.parameters():
+            param.requires_grad = False
 
     def forward(
         self,
         data: Dict[str, torch.Tensor],
-        mask_ghost: torch.Tensor,
-        compute_force: bool = True,
+        local_or_ghost: torch.Tensor,
         compute_virials: bool = False,
-        compute_stress: bool = False,
     ) -> Dict[str, Optional[torch.Tensor]]:
         num_graphs = data["ptr"].numel() - 1
         compute_displacement = False
-        if compute_virials or compute_stress:
+        if compute_virials:
             compute_displacement = True
-
         out = self.model(
             data,
             training=False,
@@ -39,59 +37,58 @@ def forward(
         )
         node_energy = out["node_energy"]
         if node_energy is None:
-            return {"energy": None, "forces": None, "virials": None, "stress": None}
+            return {
+                "total_energy_local": None,
+                "node_energy": None,
+                "forces": None,
+                "virials": None,
+            }
+        positions = data["positions"]
         displacement = out["displacement"]
+        forces: Optional[torch.Tensor] = torch.zeros_like(positions)
         virials: Optional[torch.Tensor] = torch.zeros_like(data["cell"])
-        stress: Optional[torch.Tensor] = torch.zeros_like(data["cell"])
-        if mask_ghost is not None and displacement is not None:
-            # displacement.requires_grad_(True)  # For some reason torchscript needs that.
-            node_energy_ghost = node_energy * mask_ghost
-            total_energy_ghost = scatter_sum(
-                src=node_energy_ghost, index=data["batch"], dim=-1, dim_size=num_graphs
-            )
-            grad_outputs: List[Optional[torch.Tensor]] = [
-                torch.ones_like(total_energy_ghost)
-            ]
-            virials = torch.autograd.grad(
-                outputs=[total_energy_ghost],
-                inputs=[displacement],
+        # accumulate energies of local atoms
+        node_energy_local = node_energy * local_or_ghost
+        total_energy_local = scatter_sum(
+            src=node_energy_local, index=data["batch"], dim=-1, dim_size=num_graphs
+        )
+        # compute partial forces and (possibly) partial virials
+        grad_outputs: List[Optional[torch.Tensor]] = [
+            torch.ones_like(total_energy_local)
+        ]
+        if compute_virials and displacement is not None:
+            forces, virials = torch.autograd.grad(
+                outputs=[total_energy_local],
+                inputs=[positions, displacement],
                 grad_outputs=grad_outputs,
-                retain_graph=True,
-                create_graph=True,
+                retain_graph=False,
+                create_graph=False,
                 allow_unused=True,
-            )[0]
-
+            )
+            if forces is not None:
+                forces = -1 * forces
+            else:
+                forces = torch.zeros_like(positions)
             if virials is not None:
                 virials = -1 * virials
-                cell = data["cell"].view(-1, 3, 3)
-                volume = torch.einsum(
-                    "zi,zi->z",
-                    cell[:, 0, :],
-                    torch.cross(cell[:, 1, :], cell[:, 2, :], dim=1),
-                ).unsqueeze(-1)
-                stress = virials / volume.view(-1, 1, 1)
             else:
                 virials = torch.zeros_like(displacement)
-
-        total_energy = scatter_sum(
-            src=node_energy, index=data["batch"], dim=-1, dim_size=num_graphs
-        )
-
-        forces, _, _ = get_outputs(
-            energy=total_energy,
-            positions=data["positions"],
-            displacement=displacement,
-            cell=data["cell"],
-            training=False,
-            compute_force=compute_force,
-            compute_virials=False,
-            compute_stress=False,
-        )
-
+        else:
+            forces = torch.autograd.grad(
+                outputs=[total_energy_local],
+                inputs=[positions],
+                grad_outputs=grad_outputs,
+                retain_graph=False,
+                create_graph=False,
+                allow_unused=True,
+            )[0]
+            if forces is not None:
+                forces = -1 * forces
+            else:
+                forces = torch.zeros_like(positions)
         return {
-            "energy": total_energy,
+            "total_energy_local": total_energy_local,
             "node_energy": node_energy,
             "forces": forces,
             "virials": virials,
-            "stress": stress,
         }