MACE

.github/workflows/CI.yml

@@ -40,7 +40,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           python -m pip install --upgrade -r requirements.txt -r requirements-dev.txt
-          python -m pip install --upgrade -r requirements-torch.txt --index-url https://download.pytorch.org/whl/cpu
+          python -m pip install --upgrade -r requirements-torch.txt --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple
           python -m pip install --upgrade -r requirements-pyg.txt --find-links https://data.pyg.org/whl/torch-2.0.1+cpu.html
           python -m pip install --upgrade -r requirements-deepspeed.txt
       - name: Format black

examples/LennardJones/LJ.json

@@ -30,6 +30,8 @@
       "num_before_skip": 1,
       "num_after_skip": 1,
       "envelope_exponent": 5,
+      "max_ell": 1,
+      "node_max_ell": 1,
       "num_radial": 5,
       "num_spherical": 2,
       "hidden_dim": 20,

hydragnn/models/create.py

@@ -12,6 +12,7 @@
 import os
 import torch
 from torch_geometric.data import Data
+from typing import List, Union
 
 from hydragnn.models.GINStack import GINStack
 from hydragnn.models.PNAStack import PNAStack
@@ -25,6 +26,7 @@
 from hydragnn.models.EGCLStack import EGCLStack
 from hydragnn.models.PNAEqStack import PNAEqStack
 from hydragnn.models.PAINNStack import PAINNStack
+from hydragnn.models.MACEStack import MACEStack
 
 from hydragnn.utils.distributed import get_device
 from hydragnn.utils.profiling_and_tracing.time_utils import Timer
@@ -55,6 +57,8 @@ def create_model_config(
         config["Architecture"]["num_before_skip"],
         config["Architecture"]["num_after_skip"],
         config["Architecture"]["num_radial"],
+        config["Architecture"]["radial_type"],
+        config["Architecture"]["distance_transform"],
         config["Architecture"]["basis_emb_size"],
         config["Architecture"]["int_emb_size"],
         config["Architecture"]["out_emb_size"],
@@ -64,6 +68,10 @@ def create_model_config(
         config["Architecture"]["num_filters"],
         config["Architecture"]["radius"],
         config["Architecture"]["equivariance"],
+        config["Architecture"]["correlation"],
+        config["Architecture"]["max_ell"],
+        config["Architecture"]["node_max_ell"],
+        config["Architecture"]["avg_num_neighbors"],
         config["Training"]["conv_checkpointing"],
         verbosity,
         use_gpu,
@@ -91,6 +99,8 @@ def create_model(
     num_before_skip: int = None,
     num_after_skip: int = None,
     num_radial: int = None,
+    radial_type: str = None,
+    distance_transform: str = None,
     basis_emb_size: int = None,
     int_emb_size: int = None,
     out_emb_size: int = None,
@@ -100,6 +110,10 @@ def create_model(
     num_filters: int = None,
     radius: float = None,
     equivariance: bool = False,
+    correlation: Union[int, List[int]] = None,
+    max_ell: int = None,
+    node_max_ell: int = None,
+    avg_num_neighbors: int = None,
     conv_checkpointing: bool = False,
     verbosity: int = 0,
     use_gpu: bool = True,
@@ -371,6 +385,39 @@ def create_model(
             num_conv_layers=num_conv_layers,
             num_nodes=num_nodes,
         )
+
+    elif model_type == "MACE":
+        assert radius is not None, "MACE requires radius input."
+        assert num_radial is not None, "MACE requires num_radial input."
+        assert max_ell is not None, "MACE requires max_ell input."
+        assert node_max_ell is not None, "MACE requires node_max_ell input."
+        assert max_ell >= 1, "MACE requires max_ell >= 1."
+        assert node_max_ell >= 1, "MACE requires node_max_ell >= 1."
+        model = MACEStack(
+            radius,
+            radial_type,
+            distance_transform,
+            num_radial,
+            edge_dim,
+            max_ell,
+            node_max_ell,
+            avg_num_neighbors,
+            envelope_exponent,
+            correlation,
+            input_dim,
+            hidden_dim,
+            output_dim,
+            output_type,
+            output_heads,
+            activation_function,
+            loss_function_type,
+            equivariance,
+            loss_weights=task_weights,
+            freeze_conv=freeze_conv,
+            initial_bias=initial_bias,
+            num_conv_layers=num_conv_layers,
+            num_nodes=num_nodes,
+        )
     else:
         raise ValueError("Unknown model_type: {0}".format(model_type))
 

hydragnn/utils/input_config_parsing/config_utils.py

@@ -14,6 +14,7 @@
     check_if_graph_size_variable,
     gather_deg,
 )
+from hydragnn.utils.model.model import calculate_avg_deg
 from hydragnn.utils.distributed import get_comm_size_and_rank
 from copy import deepcopy
 import json
@@ -56,8 +57,22 @@ def update_config(config, train_loader, val_loader, test_loader):
     else:
         config["NeuralNetwork"]["Architecture"]["pna_deg"] = None
 
+    if config["NeuralNetwork"]["Architecture"]["model_type"] == "MACE":
+        if hasattr(train_loader.dataset, "avg_num_neighbors"):
+            ## Use avg neighbours used in the dataset.
+            avg_num_neighbors = torch.tensor(train_loader.dataset.avg_num_neighbors)
+        else:
+            avg_num_neighbors = float(calculate_avg_deg(train_loader.dataset))
+        config["NeuralNetwork"]["Architecture"]["avg_num_neighbors"] = avg_num_neighbors
+    else:
+        config["NeuralNetwork"]["Architecture"]["avg_num_neighbors"] = None
+
     if "radius" not in config["NeuralNetwork"]["Architecture"]:
         config["NeuralNetwork"]["Architecture"]["radius"] = None
+    if "radial_type" not in config["NeuralNetwork"]["Architecture"]:
+        config["NeuralNetwork"]["Architecture"]["radial_type"] = None
+    if "distance_transform" not in config["NeuralNetwork"]["Architecture"]:
+        config["NeuralNetwork"]["Architecture"]["distance_transform"] = None
     if "num_gaussians" not in config["NeuralNetwork"]["Architecture"]:
         config["NeuralNetwork"]["Architecture"]["num_gaussians"] = None
     if "num_filters" not in config["NeuralNetwork"]["Architecture"]:
@@ -78,6 +93,14 @@ def update_config(config, train_loader, val_loader, test_loader):
         config["NeuralNetwork"]["Architecture"]["num_radial"] = None
     if "num_spherical" not in config["NeuralNetwork"]["Architecture"]:
         config["NeuralNetwork"]["Architecture"]["num_spherical"] = None
+    if "radial_type" not in config["NeuralNetwork"]["Architecture"]:
+        config["NeuralNetwork"]["Architecture"]["radial_type"] = None
+    if "correlation" not in config["NeuralNetwork"]["Architecture"]:
+        config["NeuralNetwork"]["Architecture"]["correlation"] = None
+    if "max_ell" not in config["NeuralNetwork"]["Architecture"]:
+        config["NeuralNetwork"]["Architecture"]["max_ell"] = None
+    if "node_max_ell" not in config["NeuralNetwork"]["Architecture"]:
+        config["NeuralNetwork"]["Architecture"]["node_max_ell"] = None
 
     config["NeuralNetwork"]["Architecture"] = update_config_edge_dim(
         config["NeuralNetwork"]["Architecture"]
@@ -113,23 +136,23 @@ def update_config(config, train_loader, val_loader, test_loader):
 
 
 def update_config_equivariance(config):
-    equivariant_models = ["EGNN", "SchNet", "PNAEq", "PAINN"]
+    equivariant_models = ["EGNN", "SchNet", "PNAEq", "PAINN", "MACE"]
     if "equivariance" in config and config["equivariance"]:
         assert (
             config["model_type"] in equivariant_models
-        ), "E(3) equivariance can only be ensured for EGNN and SchNet."
+        ), "E(3) equivariance can only be ensured for EGNN, SchNet, and MACE."
     elif "equivariance" not in config:
         config["equivariance"] = False
     return config
 
 
 def update_config_edge_dim(config):
     config["edge_dim"] = None
-    edge_models = ["PNAPlus", "PNA", "CGCNN", "SchNet", "EGNN", "DimeNet"]
+    edge_models = ["PNAPlus", "PNA", "CGCNN", "SchNet", "EGNN", "DimeNet", "MACE"]
     if "edge_features" in config and config["edge_features"]:
         assert (
             config["model_type"] in edge_models
-        ), "Edge features can only be used with DimeNet EGNN, SchNet, PNA, PNAPlus, and CGCNN."
+        ), "Edge features can only be used with DimeNet, MACE, EGNN, SchNet, PNA, PNAPlus, and CGCNN."
         config["edge_dim"] = len(config["edge_features"])
     elif config["model_type"] == "CGCNN":
         # CG always needs an integer edge_dim

hydragnn/utils/model/irreps_tools.py

@@ -0,0 +1,102 @@
+###########################################################################################
+# Elementary tools for handling irreducible representations
+# Authors: Ilyes Batatia, Gregor Simm
+# This program is distributed under the MIT License (see MIT.md)
+###########################################################################################
+
+from typing import List, Tuple
+
+import torch
+from e3nn import o3
+from e3nn.util.jit import compile_mode
+
+
+# Based on mir-group/nequip
+def tp_out_irreps_with_instructions(
+    irreps1: o3.Irreps, irreps2: o3.Irreps, target_irreps: o3.Irreps
+) -> Tuple[o3.Irreps, List]:
+    trainable = True
+
+    # Collect possible irreps and their instructions
+    irreps_out_list: List[Tuple[int, o3.Irreps]] = []
+    instructions = []
+    for i, (mul, ir_in) in enumerate(irreps1):
+        for j, (_, ir_edge) in enumerate(irreps2):
+            for ir_out in ir_in * ir_edge:  # | l1 - l2 | <= l <= l1 + l2
+                if ir_out in target_irreps:
+                    k = len(irreps_out_list)  # instruction index
+                    irreps_out_list.append((mul, ir_out))
+                    instructions.append((i, j, k, "uvu", trainable))
+
+    # We sort the output irreps of the tensor product so that we can simplify them
+    # when they are provided to the second o3.Linear
+    irreps_out = o3.Irreps(irreps_out_list)
+    irreps_out, permut, _ = irreps_out.sort()
+
+    # Permute the output indexes of the instructions to match the sorted irreps:
+    instructions = [
+        (i_in1, i_in2, permut[i_out], mode, train)
+        for i_in1, i_in2, i_out, mode, train in instructions
+    ]
+
+    instructions = sorted(instructions, key=lambda x: x[2])
+
+    return irreps_out, instructions
+
+
+def linear_out_irreps(irreps: o3.Irreps, target_irreps: o3.Irreps) -> o3.Irreps:
+    # Assuming simplified irreps
+    irreps_mid = []
+    for _, ir_in in irreps:
+        found = False
+
+        for mul, ir_out in target_irreps:
+            if ir_in == ir_out:
+                irreps_mid.append((mul, ir_out))
+                found = True
+                break
+
+        if not found:
+            raise RuntimeError(f"{ir_in} not in {target_irreps}")
+
+    return o3.Irreps(irreps_mid)
+
+
+@compile_mode("script")
+class reshape_irreps(torch.nn.Module):
+    def __init__(self, irreps: o3.Irreps) -> None:
+        super().__init__()
+        self.irreps = o3.Irreps(irreps)
+        self.dims = []
+        self.muls = []
+        for mul, ir in self.irreps:
+            d = ir.dim
+            self.dims.append(d)
+            self.muls.append(mul)
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        ix = 0
+        out = []
+        batch, _ = tensor.shape
+        for mul, d in zip(self.muls, self.dims):
+            field = tensor[:, ix : ix + mul * d]  # [batch, sample, mul * repr]
+            ix += mul * d
+            field = field.reshape(batch, mul, d)
+            out.append(field)
+        return torch.cat(out, dim=-1)
+
+
+def extract_invariant(x: torch.Tensor, num_layers: int, num_features: int, l_max: int):
+    out = []
+    for i in range(num_layers - 1):
+        out.append(
+            x[
+                :,
+                i
+                * (l_max + 1) ** 2
+                * num_features : (i * (l_max + 1) ** 2 + 1)
+                * num_features,
+            ]
+        )
+    out.append(x[:, -num_features:])
+    return torch.cat(out, dim=-1)

hydragnn/utils/model/mace_utils/modules/__init__.py

@@ -0,0 +1,51 @@
+###########################################################################################
+# __init__ file for Modules
+# Authors: Ilyes Batatia, Gregor Simm
+# This program is distributed under the MIT License (see MIT.md)
+###########################################################################################
+# Taken From:
+# GitHub: https://github.com/ACEsuit/mace
+# ArXiV: https://arxiv.org/pdf/2206.07697
+# Date: August 27, 2024  |  12:37 (EST)
+###########################################################################################
+
+from typing import Callable, Dict, Optional, Type
+
+import torch
+
+from .blocks import (
+    AtomicEnergiesBlock,
+    EquivariantProductBasisBlock,
+    InteractionBlock,
+    LinearNodeEmbeddingBlock,
+    LinearReadoutBlock,
+    NonLinearReadoutBlock,
+    RadialEmbeddingBlock,
+    RealAgnosticAttResidualInteractionBlock,
+    ScaleShiftBlock,
+)
+
+from .radial import BesselBasis, GaussianBasis, PolynomialCutoff
+from .symmetric_contraction import SymmetricContraction
+
+interaction_classes: Dict[str, Type[InteractionBlock]] = {
+    "RealAgnosticAttResidualInteractionBlock": RealAgnosticAttResidualInteractionBlock,
+}
+
+__all__ = [
+    "AtomicEnergiesBlock",
+    "RadialEmbeddingBlock",
+    "LinearNodeEmbeddingBlock",
+    "LinearReadoutBlock",
+    "EquivariantProductBasisBlock",
+    "ScaleShiftBlock",
+    "LinearDipoleReadoutBlock",
+    "NonLinearDipoleReadoutBlock",
+    "InteractionBlock",
+    "NonLinearReadoutBlock",
+    "PolynomialCutoff",
+    "BesselBasis",
+    "GaussianBasis",
+    "SymmetricContraction",
+    "interaction_classes",
+]