Change of license

README.md

@@ -149,4 +149,4 @@ For bugs or feature requests, please use [GitHub Issues](https://github.com/ACEs
 
 ## License
 
-MACE is published and distributed under the [Academic Software License v1.0 ](ASL.md).
+MACE is published and distributed under the [MIT](MIT.md).

mace/calculators/mace.py

@@ -1,7 +1,7 @@
 ###########################################################################################
 # The ASE Calculator for MACE (based on https://github.com/mir-group/nequip)
 # Authors: Ilyes Batatia, David Kovacs
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 
@@ -87,7 +87,7 @@ def calculate(self, atoms=None, properties=None, system_changes=all_changes):
             stress = out["stress"].detach().cpu().numpy()
             # stress has units eng / len^3:
             self.results["stress"] = (
-                stress * (self.energy_units_to_eV / self.length_units_to_A**3)
+                stress * (self.energy_units_to_eV / self.length_units_to_A ** 3)
             )[0]
             self.results["stress"] = full_3x3_to_voigt_6_stress(self.results["stress"])
 
@@ -248,5 +248,5 @@ def calculate(self, atoms=None, properties=None, system_changes=all_changes):
         if out["stress"] is not None:
             stress = out["stress"].detach().cpu().numpy()
             self.results["stress"] = (
-                stress * (self.energy_units_to_eV / self.length_units_to_A**3)
+                stress * (self.energy_units_to_eV / self.length_units_to_A ** 3)
             )[0]

mace/data/atomic_data.py

@@ -1,7 +1,7 @@
 ###########################################################################################
 # Atomic Data Class for handling molecules as graphs
 # Authors: Ilyes Batatia, Gregor Simm
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 from typing import Optional, Sequence
@@ -210,14 +210,8 @@ def from_config(
 
 
 def get_data_loader(
-    dataset: Sequence[AtomicData],
-    batch_size: int,
-    shuffle=True,
-    drop_last=False,
+    dataset: Sequence[AtomicData], batch_size: int, shuffle=True, drop_last=False,
 ) -> torch.utils.data.DataLoader:
     return torch_geometric.dataloader.DataLoader(
-        dataset=dataset,
-        batch_size=batch_size,
-        shuffle=shuffle,
-        drop_last=drop_last,
+        dataset=dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last,
     )

mace/data/neighborhood.py

@@ -1,7 +1,7 @@
 ###########################################################################################
 # Neighborhood construction
 # Authors: Ilyes Batatia, Gregor Simm
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 from typing import Optional, Tuple

mace/data/utils.py

@@ -1,7 +1,7 @@
 ###########################################################################################
 # Data parsing utilities
 # Authors: Ilyes Batatia, Gregor Simm and David Kovacs
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 import logging

mace/modules/blocks.py

@@ -1,7 +1,7 @@
 ###########################################################################################
 # Elementary Block for Building O(3) Equivariant Higher Order Message Passing Neural Network
 # Authors: Ilyes Batatia, Gregor Simm
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 from abc import ABC, abstractmethod
@@ -28,8 +28,7 @@ def __init__(self, irreps_in: o3.Irreps, irreps_out: o3.Irreps):
         self.linear = o3.Linear(irreps_in=irreps_in, irreps_out=irreps_out)
 
     def forward(
-        self,
-        node_attrs: torch.Tensor,  # [n_nodes, irreps]
+        self, node_attrs: torch.Tensor,  # [n_nodes, irreps]
     ):
         return self.linear(node_attrs)
 
@@ -142,8 +141,7 @@ def __init__(self, r_max: float, num_bessel: int, num_polynomial_cutoff: int):
         self.out_dim = num_bessel
 
     def forward(
-        self,
-        edge_lengths: torch.Tensor,  # [n_edges, 1]
+        self, edge_lengths: torch.Tensor,  # [n_edges, 1]
     ):
         bessel = self.bessel_fn(edge_lengths)  # [n_edges, n_basis]
         cutoff = self.cutoff_fn(edge_lengths)  # [n_edges, 1]
@@ -172,10 +170,7 @@ def __init__(
         )
         # Update linear
         self.linear = o3.Linear(
-            target_irreps,
-            target_irreps,
-            internal_weights=True,
-            shared_weights=True,
+            target_irreps, target_irreps, internal_weights=True, shared_weights=True,
         )
 
     def forward(
@@ -342,8 +337,7 @@ def _setup(self) -> None:
         # Convolution weights
         input_dim = self.edge_feats_irreps.num_irreps
         self.conv_tp_weights = nn.FullyConnectedNet(
-            [input_dim] + 3 * [64] + [self.conv_tp.weight_numel],
-            torch.nn.SiLU(),
+            [input_dim] + 3 * [64] + [self.conv_tp.weight_numel], torch.nn.SiLU(),
         )
 
         # Linear
@@ -407,8 +401,7 @@ def _setup(self) -> None:
         # Convolution weights
         input_dim = self.edge_feats_irreps.num_irreps
         self.conv_tp_weights = nn.FullyConnectedNet(
-            [input_dim] + 3 * [64] + [self.conv_tp.weight_numel],
-            torch.nn.SiLU(),
+            [input_dim] + 3 * [64] + [self.conv_tp.weight_numel], torch.nn.SiLU(),
         )
 
         # Linear
@@ -459,9 +452,7 @@ def _setup(self) -> None:
         )
         # TensorProduct
         irreps_mid, instructions = tp_out_irreps_with_instructions(
-            self.node_feats_irreps,
-            self.edge_attrs_irreps,
-            self.target_irreps,
+            self.node_feats_irreps, self.edge_attrs_irreps, self.target_irreps,
         )
         self.conv_tp = o3.TensorProduct(
             self.node_feats_irreps,
@@ -475,8 +466,7 @@ def _setup(self) -> None:
         # Convolution weights
         input_dim = self.edge_feats_irreps.num_irreps
         self.conv_tp_weights = nn.FullyConnectedNet(
-            [input_dim] + 3 * [64] + [self.conv_tp.weight_numel],
-            torch.nn.SiLU(),
+            [input_dim] + 3 * [64] + [self.conv_tp.weight_numel], torch.nn.SiLU(),
         )
 
         # Linear
@@ -531,9 +521,7 @@ def _setup(self) -> None:
         )
         # TensorProduct
         irreps_mid, instructions = tp_out_irreps_with_instructions(
-            self.node_feats_irreps,
-            self.edge_attrs_irreps,
-            self.target_irreps,
+            self.node_feats_irreps, self.edge_attrs_irreps, self.target_irreps,
         )
         self.conv_tp = o3.TensorProduct(
             self.node_feats_irreps,
@@ -547,8 +535,7 @@ def _setup(self) -> None:
         # Convolution weights
         input_dim = self.edge_feats_irreps.num_irreps
         self.conv_tp_weights = nn.FullyConnectedNet(
-            [input_dim] + 3 * [64] + [self.conv_tp.weight_numel],
-            torch.nn.SiLU(),
+            [input_dim] + 3 * [64] + [self.conv_tp.weight_numel], torch.nn.SiLU(),
         )
 
         # Linear

mace/modules/irreps_tools.py

@@ -1,7 +1,7 @@
 ###########################################################################################
 # Elementary tools for handling irreducible representations
 # Authors: Ilyes Batatia, Gregor Simm
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 from typing import List, Tuple

mace/modules/loss.py

@@ -1,7 +1,7 @@
 ###########################################################################################
 # Implementation of different loss functions
 # Authors: Ilyes Batatia, Gregor Simm
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 import torch

mace/modules/models.py

@@ -1,7 +1,7 @@
 ###########################################################################################
 # Implementation of MACE models and other models based E(3)-Equivariant MPNNs
 # Authors: Ilyes Batatia, Gregor Simm
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 from typing import Any, Callable, Dict, List, Optional, Type
@@ -224,10 +224,7 @@ def forward(
 
 class ScaleShiftMACE(MACE):
     def __init__(
-        self,
-        atomic_inter_scale: float,
-        atomic_inter_shift: float,
-        **kwargs,
+        self, atomic_inter_scale: float, atomic_inter_shift: float, **kwargs,
     ):
         super().__init__(**kwargs)
         self.scale_shift = ScaleShiftBlock(
@@ -445,10 +442,7 @@ def forward(self, data: AtomicData, training=False) -> Dict[str, Any]:
 
 class ScaleShiftBOTNet(BOTNet):
     def __init__(
-        self,
-        atomic_inter_scale: float,
-        atomic_inter_shift: float,
-        **kwargs,
+        self, atomic_inter_scale: float, atomic_inter_shift: float, **kwargs,
     ):
         super().__init__(**kwargs)
         self.scale_shift = ScaleShiftBlock(

mace/modules/radial.py

@@ -1,7 +1,7 @@
 ###########################################################################################
 # Radial basis and cutoff
 # Authors: Ilyes Batatia, Gregor Simm
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 import numpy as np
@@ -40,10 +40,7 @@ def __init__(self, r_max: float, num_basis=8, trainable=False):
             torch.tensor(np.sqrt(2.0 / r_max), dtype=torch.get_default_dtype()),
         )
 
-    def forward(
-        self,
-        x: torch.Tensor,
-    ) -> torch.Tensor:  # [..., 1]
+    def forward(self, x: torch.Tensor,) -> torch.Tensor:  # [..., 1]
         numerator = torch.sin(self.bessel_weights * x)  # [..., num_basis]
         return self.prefactor * (numerator / x)
 

mace/modules/symmetric_contraction.py

@@ -2,7 +2,7 @@
 # Implementation of the symmetric contraction algorithm presented in the MACE paper
 # (Batatia et al, MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields , Eq.10 and 11)
 # Authors: Ilyes Batatia
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 from typing import Dict, Optional, Union

mace/modules/utils.py

@@ -1,7 +1,7 @@
 ###########################################################################################
 # Utilities
 # Authors: Ilyes Batatia, Gregor Simm and David Kovacs
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 import logging
@@ -60,9 +60,7 @@ def compute_forces_virials(
     if compute_stress and virials is not None:
         cell = cell.view(-1, 3, 3)
         volume = torch.einsum(
-            "zi,zi->z",
-            cell[:, 0, :],
-            torch.cross(cell[:, 1, :], cell[:, 2, :], dim=1),
+            "zi,zi->z", cell[:, 0, :], torch.cross(cell[:, 1, :], cell[:, 2, :], dim=1),
         ).unsqueeze(-1)
         stress = virials / volume.view(-1, 1, 1)
     if forces is None:
@@ -86,10 +84,7 @@ def get_symmetric_displacement(
     if cell is None:
         logging.info("Virial required but no cell provided")
         cell = torch.zeros(
-            num_graphs * 3,
-            3,
-            dtype=positions.dtype,
-            device=positions.device,
+            num_graphs * 3, 3, dtype=positions.dtype, device=positions.device,
         )
     sender = edge_index[0]
     displacement = torch.zeros(
@@ -106,11 +101,7 @@ def get_symmetric_displacement(
     )
     cell = cell.view(-1, 3, 3)
     cell = cell + torch.matmul(cell, symmetric_displacement)
-    shifts = torch.einsum(
-        "be,bec->bc",
-        unit_shifts,
-        cell[batch[sender]],
-    )
+    shifts = torch.einsum("be,bec->bc", unit_shifts, cell[batch[sender]],)
     return positions, shifts, displacement
 
 
@@ -166,8 +157,7 @@ def get_edge_vectors_and_lengths(
 
 
 def compute_mean_std_atomic_inter_energy(
-    data_loader: torch.utils.data.DataLoader,
-    atomic_energies: np.ndarray,
+    data_loader: torch.utils.data.DataLoader, atomic_energies: np.ndarray,
 ) -> Tuple[float, float]:
     atomic_energies_fn = AtomicEnergiesBlock(atomic_energies=atomic_energies)
 
@@ -191,8 +181,7 @@ def compute_mean_std_atomic_inter_energy(
 
 
 def compute_mean_rms_energy_forces(
-    data_loader: torch.utils.data.DataLoader,
-    atomic_energies: np.ndarray,
+    data_loader: torch.utils.data.DataLoader, atomic_energies: np.ndarray,
 ) -> Tuple[float, float]:
     atomic_energies_fn = AtomicEnergiesBlock(atomic_energies=atomic_energies)
 

mace/tools/arg_parser.py

@@ -1,7 +1,7 @@
 ###########################################################################################
 # Parsing functionalities
 # Authors: Ilyes Batatia, Gregor Simm, David Kovacs
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 import argparse
@@ -197,9 +197,7 @@ def build_default_arg_parser() -> argparse.ArgumentParser:
         required=False,
     )
     parser.add_argument(
-        "--test_file",
-        help="Test set xyz file",
-        type=str,
+        "--test_file", help="Test set xyz file", type=str,
     )
     parser.add_argument(
         "--E0s",

mace/tools/cg.py

@@ -1,7 +1,7 @@
 ###########################################################################################
 # Higher Order Real Clebsch Gordan (based on e3nn by Mario Geiger)
 # Authors: Ilyes Batatia
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 import collections
@@ -54,9 +54,9 @@ def _wigner_nj(
 
                 C = o3.wigner_3j(ir_out.l, ir_left.l, ir.l, dtype=dtype)
                 if normalization == "component":
-                    C *= ir_out.dim**0.5
+                    C *= ir_out.dim ** 0.5
                 if normalization == "norm":
-                    C *= ir_left.dim**0.5 * ir.dim**0.5
+                    C *= ir_left.dim ** 0.5 * ir.dim ** 0.5
 
                 C = torch.einsum("jk,ijl->ikl", C_left.flatten(1), C)
                 C = C.reshape(

mace/tools/checkpoint.py

@@ -1,7 +1,7 @@
 ###########################################################################################
 # Checkpointing
 # Authors: Gregor Simm
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 import dataclasses
@@ -86,9 +86,7 @@ def _parse_checkpoint_path(self, path: str) -> Optional[CheckpointPathInfo]:
             return None
 
         return CheckpointPathInfo(
-            path=path,
-            tag=match.group("tag"),
-            epochs=int(match.group("epochs")),
+            path=path, tag=match.group("tag"), epochs=int(match.group("epochs")),
         )
 
     def _get_latest_checkpoint_path(self) -> Optional[str]:

mace/tools/scripts_utils.py

@@ -1,7 +1,7 @@
 ###########################################################################################
 # Training utils
 # Authors: David Kovacs, Ilyes Batatia
-# This program is distributed under the ASL License (see ASL.md)
+# This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
 import dataclasses