Seven net and nequix configs

openequivariance/benchmark/problems.py

@@ -151,6 +151,54 @@ def nequip_problems():
     ]
 
 
+# https://github.com/atomicarchitects/nequix/blob/main/configs/nequix-mp-1.yml
+def nequix_problems():
+    return [
+        CTPP(
+            "89x0e",
+            "1x0e+1x1o+1x2e+1x3o",
+            "89x0e+89x1o+89x2e+89x3o",
+            "nequix-mp-1-first_layer",
+        ),
+        CTPP(
+            "128x0e+64x1o+32x2e+32x3o",
+            "1x0e+1x1o+1x2e+1x3o",
+            "128x0e+128x1o+128x2e+128x3o+64x1o+64x0e+64x2e+64x1o+64x3o+64x2e+32x2e+32x1o+32x3o+32x0e+32x2e+32x1o+32x3o+32x3o+32x2e+32x1o+32x3o+32x0e+32x2e",
+            "nequix-mp-1-main_layers",
+        ),
+        CTPP(
+            "128x0e+64x1o+32x2e+32x3o",
+            "1x0e+1x1o+1x2e+1x3o",
+            "128x0e+64x0e+32x0e+32x0e",
+            "nequix-mp-1-last_layer",
+        ),
+    ]
+
+
+# https://github.com/MDIL-SNU/SevenNet/tree/main/sevenn/pretrained_potentials/SevenNet_l3i5
+def seven_net_problems():
+    return [
+        CTPP(
+            "128x0e",
+            "1x0e+1x1e+1x2e+1x3e",
+            "128x0e+128x1e+128x2e+128x3e",
+            "SevenNet_l3i5-first-layer",
+        ),
+        CTPP(
+            "128x0e+64x1e+32x2e+32x3e",
+            "1x0e+1x1e+1x2e+1x3e",
+            "128x0e+64x0e+32x0e+32x0e+128x1e+64x1e+64x1e+64x1e+32x1e+32x1e+32x1e+32x1e+32x1e+128x2e+64x2e+64x2e+64x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+128x3e+64x3e+64x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e",
+            "SevenNet_l3i5-main-layers",
+        ),
+        CTPP(
+            "128x0e+64x1e+32x2e+32x3e",
+            "1x0e+1x1e+1x2e+1x3e",
+            "128x0e+64x0e+32x0e+32x0e",
+            "SevenNet_l3i5-last-layer",
+        ),
+    ]
+
+
 def e3tools_problems():
     return [
         FCTPP(in1, in2, out, label=label, shared_weights=sw, internal_weights=iw)

tests/benchmark.py

@@ -46,6 +46,8 @@
     diffdock_problems,
     mace_problems,
     nequip_problems,
+    nequix_problems,
+    seven_net_problems,
 )
 
 from torch._functorch import config
@@ -91,11 +93,19 @@
 
 
 def benchmark_uvu(params):
-    float64_problems = mace_problems() + nequip_problems()
+    def get_problems():
+        return (
+            mace_problems()
+            + nequip_problems()
+            + nequix_problems()
+            + seven_net_problems()
+        )
+
+    float64_problems = get_problems()
     for problem in float64_problems:
         problem.irrep_dtype = np.float64
         problem.weight_dtype = np.float64
-    problems = mace_problems() + nequip_problems() + float64_problems
+    problems = get_problems() + float64_problems
 
     implementations = [implementation_map_tp[impl] for impl in params.implementations]
     directions = params.directions
@@ -280,21 +290,20 @@ def benchmark_convolution(params):
     graphs = download_graphs(params, filenames)
 
     if not params.disable_bench:
-        configs = [
-            ChannelwiseTPP(
-                "128x0e+128x1o+128x2e",
-                "1x0e+1x1o+1x2e+1x3o",
-                "128x0e+128x1o+128x2e+128x3o",
-            ),
-            ChannelwiseTPP(
-                "128x0e+128x1o+128x2e",
-                "1x0e+1x1o+1x2e+1x3o",
-                "128x0e+128x1o+128x2e+128x3o",
-            ),
-        ]  # MACE-large
-
-        configs[1].irrep_dtype = np.float64
-        configs[1].weight_dtype = np.float64
+
+        def get_problems():
+            return (
+                mace_problems()
+                + nequip_problems()
+                + nequix_problems()
+                + seven_net_problems()
+            )
+
+        float64_problems = get_problems()
+        for problem in float64_problems:
+            problem.irrep_dtype = np.float64
+            problem.weight_dtype = np.float64
+        configs = get_problems() + float64_problems
 
         bench = ConvBenchmarkSuite(configs, test_name="convolution")