add forces test

ORNL · Sep 25, 2024 · 3fe96df · 3fe96df
1 parent db705ca
commit 3fe96df
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 7 deletions.
diff --git a/examples/LennardJones/LJ.json b/examples/LennardJones/LJ.json
@@ -19,7 +19,7 @@
   "NeuralNetwork": {
     "Architecture": {
       "periodic_boundary_conditions": true,
-      "model_type": "DimeNet",
+      "model_type": "MACE",
       "radius": 5.0,
       "max_neighbours": 5,
       "int_emb_size": 32,
@@ -32,6 +32,8 @@
       "envelope_exponent": 5,
       "num_radial": 5,
       "num_spherical": 2,
+      "max_ell": 2,
+      "node_max_ell": 2,
       "hidden_dim": 20,
       "num_conv_layers": 4,
       "output_heads": {

diff --git a/hydragnn/models/MACEStack.py b/hydragnn/models/MACEStack.py
@@ -9,7 +9,7 @@
 # SPDX-License-Identifier: BSD-3-Clause                                      #
 ##############################################################################
 
-# Adapted From:
+# Adapted From (September 1/2024):
 # GitHub: https://github.com/ACEsuit/mace
 # ArXiV: https://arxiv.org/pdf/2206.07697
 ###########################################################################################
@@ -190,9 +190,10 @@ def _init_conv(self):
         self.graph_convs.append(
             self.get_conv(self.input_dim, self.hidden_dim, first_layer=True)
         )
+        irreps = hidden_irreps if not last_layer else final_hidden_irreps
         self.multihead_decoders.append(
             MultiheadDecoderBlock(
-                hidden_irreps,
+                irreps,
                 self.node_max_ell,
                 self.config_heads,
                 self.head_dims,
@@ -209,9 +210,10 @@ def _init_conv(self):
                 self.hidden_dim, self.hidden_dim, last_layer=last_layer
             )
             self.graph_convs.append(conv)
+            irreps = hidden_irreps if not last_layer else final_hidden_irreps
             self.multihead_decoders.append(
                 MultiheadDecoderBlock(
-                    final_hidden_irreps,
+                    irreps,
                     self.node_max_ell,
                     self.config_heads,
                     self.head_dims,
@@ -256,7 +258,7 @@ def get_conv(self, input_dim, output_dim, first_layer=False, last_layer=False):
                 node_feats_irreps=node_feats_irreps,
                 edge_attrs_irreps=self.sh_irreps,
                 edge_feats_irreps=self.edge_feats_irreps,
-                target_irreps=interaction_irreps,  # Replace with output?
+                target_irreps=interaction_irreps,
                 hidden_irreps=hidden_irreps_out,
                 avg_num_neighbors=self.avg_num_neighbors,
                 radial_MLP=radial_MLP,
@@ -323,7 +325,6 @@ def get_conv(self, input_dim, output_dim, first_layer=False, last_layer=False):
             )  # Change sizing to output_irreps
 
         input_args = "node_attributes, pos, node_features, edge_attributes, edge_features, edge_index"
-        # readout_args = "node_energies"
         conv_args = "node_attributes, edge_attributes, edge_features, edge_index"  # node_features is not used here because it's passed through in the forward
 
         if self.use_edge_attr:

diff --git a/tests/test_forces_equivariant.py b/tests/test_forces_equivariant.py
@@ -16,7 +16,7 @@
 
 
 @pytest.mark.parametrize("example", ["LennardJones"])
-@pytest.mark.parametrize("model_type", ["SchNet", "EGNN", "DimeNet", "PNAPlus"])
+@pytest.mark.parametrize("model_type", ["SchNet", "EGNN", "DimeNet", "PNAPlus", "MACE"])
 @pytest.mark.mpi_skip()
 def pytest_examples(example, model_type):
     path = os.path.join(os.path.dirname(__file__), "..", "examples", example)