Skip to content

Commit e723ea7

Browse files
asgloverAustin Glover
andauthored
add Seven-net and nequix configs (#169)
* add nequix and seven-net problems --------- Co-authored-by: Austin Glover <[email protected]>
1 parent bcb7fdb commit e723ea7

File tree

1 file changed

+48
-0
lines changed

1 file changed

+48
-0
lines changed

openequivariance/benchmark/problems.py

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,54 @@ def nequip_problems():
151151
]
152152

153153

154+
# https://github.com/atomicarchitects/nequix/blob/main/configs/nequix-mp-1.yml
155+
def nequix_problems():
156+
return [
157+
CTPP(
158+
"89x0e",
159+
"1x0e+1x1o+1x2e+1x3o",
160+
"89x0e+89x1o+89x2e+89x3o",
161+
"nequix-mp-1-first_layer",
162+
),
163+
CTPP(
164+
"128x0e+64x1o+32x2e+32x3o",
165+
"1x0e+1x1o+1x2e+1x3o",
166+
"128x0e+128x1o+128x2e+128x3o+64x1o+64x0e+64x2e+64x1o+64x3o+64x2e+32x2e+32x1o+32x3o+32x0e+32x2e+32x1o+32x3o+32x3o+32x2e+32x1o+32x3o+32x0e+32x2e",
167+
"nequix-mp-1-main_layers",
168+
),
169+
CTPP(
170+
"128x0e+64x1o+32x2e+32x3o",
171+
"1x0e+1x1o+1x2e+1x3o",
172+
"128x0e+64x0e+32x0e+32x0e",
173+
"nequix-mp-1-last_layer",
174+
),
175+
]
176+
177+
178+
# https://github.com/MDIL-SNU/SevenNet/tree/main/sevenn/pretrained_potentials/SevenNet_l3i5
179+
def seven_net_problems():
180+
return [
181+
CTPP(
182+
"128x0e",
183+
"1x0e+1x1e+1x2e+1x3e",
184+
"128x0e+128x1e+128x2e+128x3e",
185+
"SevenNet_l3i5-first-layer",
186+
),
187+
CTPP(
188+
"128x0e+64x1e+32x2e+32x3e",
189+
"1x0e+1x1e+1x2e+1x3e",
190+
"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",
191+
"SevenNet_l3i5-main-layers",
192+
),
193+
CTPP(
194+
"128x0e+64x1e+32x2e+32x3e",
195+
"1x0e+1x1e+1x2e+1x3e",
196+
"128x0e+64x0e+32x0e+32x0e",
197+
"SevenNet_l3i5-last-layer",
198+
),
199+
]
200+
201+
154202
def e3tools_problems():
155203
return [
156204
FCTPP(in1, in2, out, label=label, shared_weights=sw, internal_weights=iw)

0 commit comments

Comments
 (0)