EDGE_CUTOFF_KEY

mir-group · Apr 7, 2023 · cc76ba8 · cc76ba8
1 parent 9e20b15
commit cc76ba8
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/allegro/nn/_radial_basis.py b/allegro/nn/_radial_basis.py
@@ -122,7 +122,10 @@ def __init__(
         self._init_irreps(
             irreps_in=irreps_in,
             irreps_out={
-                AtomicDataDict.EDGE_EMBEDDING_KEY: o3.Irreps([(self.num_basis, (0, 1))])
+                AtomicDataDict.EDGE_EMBEDDING_KEY: o3.Irreps(
+                    [(self.num_basis, (0, 1))]
+                ),
+                AtomicDataDict.EDGE_CUTOFF_KEY: "0e",
             },
         )
 
@@ -163,4 +166,5 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         data[AtomicDataDict.EDGE_EMBEDDING_KEY] = (
             bessel.view(-1, self.num_basis) * cutoff
         )
+        data[AtomicDataDict.EDGE_CUTOFF_KEY] = cutoff
         return data
diff --git a/tests/model/test_allegro.py b/tests/model/test_allegro.py
@@ -8,13 +8,15 @@
     "r_max": 4,
     "num_bessels_per_basis": 4,
     "num_bases": 4,
-    "num_types": 3,
-    "type_names": ["H", "C", "O"],
+    "chemical_symbol_to_type": {"H": 0, "C": 1, "O": 2},
     "num_tensor_features": 4,
     "two_body_latent_mlp_latent_dimensions": [32],
     "latent_mlp_latent_dimensions": [32, 32],
     "env_embed_mlp_latent_dimensions": [],
     "edge_eng_mlp_latent_dimensions": [8],
+    # Just in case for when that builder exists:
+    "pair_style": "ZBL",
+    "units": "metal",
 }
 # TODO: test so3 mode when can pass down option to assert equivariance to ignore parity
 minimal_config1 = dict(
@@ -100,6 +102,16 @@ def base_config(self, request):
                     AtomicDataDict.VIRIAL_KEY,
                 ],
             ),
+            (
+                ["allegro.model.Allegro", "PairPotentialTerm", "StressForceOutput"],
+                [
+                    AtomicDataDict.TOTAL_ENERGY_KEY,
+                    AtomicDataDict.PER_ATOM_ENERGY_KEY,
+                    AtomicDataDict.FORCE_KEY,
+                    AtomicDataDict.STRESS_KEY,
+                    AtomicDataDict.VIRIAL_KEY,
+                ],
+            ),
         ],
         scope="class",
     )