Refactor MACE subclasses - reduce code duplication & clearer logic

mace/modules/__init__.py

@@ -36,6 +36,7 @@
     EnergyDipolesMACE,
     ScaleShiftBOTNet,
     ScaleShiftMACE,
+    ScaleShiftEnergyDipoleMACE,
 )
 from .radial import BesselBasis, PolynomialCutoff
 from .symmetric_contraction import SymmetricContraction

mace/modules/blocks.py

@@ -66,12 +66,20 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:  # [n_nodes, irreps]  # [...
 
 @compile_mode("script")
 class LinearDipoleReadoutBlock(torch.nn.Module):
-    def __init__(self, irreps_in: o3.Irreps, dipole_only: bool = False):
+    def __init__(self, irreps_in: o3.Irreps):
         super().__init__()
-        if dipole_only:
-            self.irreps_out = o3.Irreps("1x1o")
-        else:
-            self.irreps_out = o3.Irreps("1x0e + 1x1o")
+        self.irreps_out = o3.Irreps("1x0e + 1x1o")
+        self.linear = o3.Linear(irreps_in=irreps_in, irreps_out=self.irreps_out)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:  # [n_nodes, irreps]  # [..., ]
+        return self.linear(x)  # [n_nodes, 1]
+
+
+@compile_mode("script")
+class LinearDipoleOnlyReadoutBlock(torch.nn.Module):
+    def __init__(self, irreps_in: o3.Irreps):
+        super().__init__()
+        self.irreps_out = o3.Irreps("1x1o")
         self.linear = o3.Linear(irreps_in=irreps_in, irreps_out=self.irreps_out)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:  # [n_nodes, irreps]  # [..., ]
@@ -100,24 +108,38 @@ def __init__(
             [(mul, ir) for mul, ir in MLP_irreps if ir.l > 0 and ir in self.irreps_out]
         )
         irreps_gates = o3.Irreps([mul, "0e"] for mul, _ in irreps_gated)
-        self.equivariant_nonlin = nn.Gate(
+        self.non_linearity = nn.Gate(
             irreps_scalars=irreps_scalars,
             act_scalars=[gate for _, ir in irreps_scalars],
             irreps_gates=irreps_gates,
             act_gates=[gate] * len(irreps_gates),
             irreps_gated=irreps_gated,
         )
-        self.irreps_nonlin = self.equivariant_nonlin.irreps_in.simplify()
+        self.irreps_nonlin = self.non_linearity.irreps_in.simplify()
         self.linear_1 = o3.Linear(irreps_in=irreps_in, irreps_out=self.irreps_nonlin)
         self.linear_2 = o3.Linear(
             irreps_in=self.hidden_irreps, irreps_out=self.irreps_out
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:  # [n_nodes, irreps]  # [..., ]
-        x = self.equivariant_nonlin(self.linear_1(x))
+        x = self.non_linearity(self.linear_1(x))
         return self.linear_2(x)  # [n_nodes, 1]
 
 
+@compile_mode("script")
+class NonLinearDipoleOnlyReadoutBlock(NonLinearDipoleReadoutBlock):
+    def __init__(
+        self,
+        irreps_in: o3.Irreps,
+        MLP_irreps: o3.Irreps,
+        gate: Callable,
+    ):
+        # fixme: use more reasonable inheritance
+        super().__init__(
+            irreps_in=irreps_in, MLP_irreps=MLP_irreps, gate=gate, dipole_only=True
+        )
+
+
 @compile_mode("script")
 class AtomicEnergiesBlock(torch.nn.Module):
     atomic_energies: torch.Tensor
@@ -258,7 +280,8 @@ def __init__(self, num_elements: int, num_edge_feats: int, num_feats_out: int):
 
     def forward(
         self,
-        sender_or_receiver_node_attrs: torch.Tensor,  # assumes that the node attributes are one-hot encoded
+        sender_or_receiver_node_attrs: torch.Tensor,
+        # assumes that the node attributes are one-hot encoded
         edge_feats: torch.Tensor,
     ):
         return torch.einsum(