Skip to content

Commit a348fcd

Browse files
vbharadwaj-bkvbharadwaj-bk
authored
Shared Weight Problem Configuration Bugfix + UVW Benchmark Relocation (#108)
* Encapsulated the lists in the problem dictionary. * Renamed benchmark_problems.py -> prblems.py, removed standalone uvw benchmark file. * Linted.
1 parent d1f20a1 commit a348fcd

7 files changed

Lines changed: 237 additions & 252 deletions

File tree

‎openequivariance/benchmark/benchmark_configs.py‎

Lines changed: 0 additions & 141 deletions
This file was deleted.

‎openequivariance/benchmark/benchmark_routines/paper_benchmark_uvw.py‎

Lines changed: 0 additions & 75 deletions
This file was deleted.

‎openequivariance/benchmark/problems.py‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
from openequivariance.benchmark.tpp_creation_utils import (
2+
FullyConnectedTPProblem as FCTPP,
3+
)
4+
from openequivariance.benchmark.tpp_creation_utils import ChannelwiseTPP as CTPP
5+
6+
# source: https://github.com/e3nn/e3nn/blob/main/examples/tetris.py
7+
# running tetris will output the layers. I've only extracted the fully connected layers here.
8+
_e3nn_torch_tetris = [
9+
# 0th Layer
10+
FCTPP("1x0e", "1x0e", "150x0e + 50x1o + 50x2e"), # sc
11+
FCTPP("1x0e", "1x0e", "1x0e"), # lin1
12+
FCTPP("1x0e + 1x1o + 1x2e", "1x0e", "150x0e + 50x1o + 50x2e"), # lin2
13+
FCTPP("1x0e + 1x1o + 1x2e", "1x0e", "1x0e"), # alpha
14+
# 1st Layer
15+
FCTPP(
16+
"50x0e + 50x1o + 50x2e", "1x0e", "250x0e + 50x1o + 50x1e + 50x2o + 50x2e"
17+
), # sc
18+
FCTPP("50x0e + 50x1o + 50x2e", "1x0e", "50x0e + 50x1o + 50x2e"), # lin1
19+
# FCTPP("50x0e + 50x1o + 50x2e", "1x0e + 1x1o + 1x2e", "150x0e + 200x1o + 100x1e + 100x2o + 200x2e"), #tp
20+
FCTPP(
21+
"150x0e + 200x1o + 100x1e + 100x2o + 200x2e",
22+
"1x0e",
23+
"250x0e + 50x1o + 50x1e + 50x2o + 50x2e",
24+
), # lin2
25+
FCTPP("150x0e + 200x1o + 100x1e + 100x2o + 200x2e", "1x0e", "1x0e"), # alpha
26+
# 2nd Layer
27+
FCTPP(
28+
"50x0e + 50x1o + 50x1e + 50x2o + 50x2e",
29+
"1x0e",
30+
"50x0o + 250x0e + 50x1o + 50x1e + 50x2o + 50x2e",
31+
), # sc
32+
FCTPP(
33+
"50x0e + 50x1o + 50x1e + 50x2o + 50x2e",
34+
"1x0e",
35+
"50x0e + 50x1o + 50x1e + 50x2o + 50x2e",
36+
), # lin1
37+
FCTPP(
38+
"100x0o + 150x0e + 300x1o + 250x1e + 250x2o + 300x2e",
39+
"1x0e",
40+
"50x0o + 250x0e + 50x1o + 50x1e + 50x2o + 50x2e",
41+
), # lin2
42+
FCTPP(
43+
"100x0o + 150x0e + 300x1o + 250x1e + 250x2o + 300x2e", "1x0e", "1x0e"
44+
), # alpha
45+
# 3rd Layer
46+
FCTPP("50x0o + 50x0e + 50x1o + 50x1e + 50x2o + 50x2e", "1x0e", "1x0o + 6x0e"), # sc
47+
FCTPP(
48+
"50x0o + 50x0e + 50x1o + 50x1e + 50x2o + 50x2e",
49+
"1x0e",
50+
"50x0o + 50x0e + 50x1o + 50x1e + 50x2o + 50x2e",
51+
), # lin1
52+
FCTPP("150x0o + 150x0e", "1x0e", "1x0o + 6x0e"), # lin2
53+
FCTPP("150x0o + 150x0e", "1x0e", "1x0e"), # alpha
54+
]
55+
56+
57+
def e3nn_torch_tetris_poly_problems():
58+
# source: https://github.com/e3nn/e3nn/blob/f95297952303347a8a3cfe971efe449c710c43b2/examples/tetris_polynomial.py#L66-L68
59+
return [
60+
FCTPP(
61+
"1x0e + 1x1o + 1x2e + 1x3o",
62+
"1x0e + 1x1o + 1x2e + 1x3o",
63+
"64x0e + 24x1e + 24x1o + 16x2e + 16x2o",
64+
label="tetris-poly-1",
65+
), # tp1
66+
FCTPP(
67+
"64x0e + 24x1e + 24x1o + 16x2e + 16x2o",
68+
"1x0e + 1x1o + 1x2e",
69+
"0o + 6x0e",
70+
label="tetris-poly-2",
71+
), # tp2
72+
]
73+
74+
75+
# https://github.com/gcorso/DiffDock/blob/b4704d94de74d8cb2acbe7ec84ad234c09e78009/models/tensor_layers.py#L299
76+
# specific irreps come from Vivek's communication with DiffDock team
77+
def diffdock_problems():
78+
return [
79+
FCTPP(
80+
"10x1o + 10x1e + 48x0e + 48x0o",
81+
"1x0e + 1x1o",
82+
"10x1o + 10x1e + 48x0e + 48x0o",
83+
shared_weights=False,
84+
label="DiffDock-L=1",
85+
),
86+
FCTPP(
87+
"10x1o + 10x1e + 48x0e + 48x0o",
88+
"1x0e + 1x1o + 1x2e",
89+
"10x1o + 10x1e + 48x0e + 48x0o",
90+
shared_weights=False,
91+
label="DiffDock-L=2",
92+
),
93+
]
94+
95+
96+
def mace_problems():
97+
return [
98+
CTPP(*config)
99+
for config in [
100+
(
101+
"128x0e+128x1o+128x2e",
102+
"1x0e+1x1o+1x2e+1x3o",
103+
"128x0e+128x1o+128x2e+128x3o",
104+
"mace-large",
105+
),
106+
(
107+
"128x0e+128x1o",
108+
"1x0e+1x1o+1x2e+1x3o",
109+
"128x0e+128x1o+128x2e",
110+
"mace-medium",
111+
),
112+
]
113+
]
114+
115+
116+
def nequip_problems():
117+
return [
118+
CTPP(*config)
119+
for config in [
120+
(
121+
"32x0o + 32x0e + 32x1o + 32x1e + 32x2o + 32x2e",
122+
"0e + 1o + 2e",
123+
"32x0o + 32x0e + 32x1o + 32x1e + 32x2o + 32x2e",
124+
"nequip-lips",
125+
),
126+
(
127+
"64x0o + 64x0e + 64x1o + 64x1e",
128+
"0e + 1o",
129+
"64x0o + 64x0e + 64x1o + 64x1e",
130+
"nequip-revmd17-aspirin",
131+
),
132+
(
133+
"64x0o + 64x0e + 64x1o + 64x1e + 64x2o + 64x2e",
134+
"0e + 1o + 2e",
135+
"64x0o + 64x0e + 64x1o + 64x1e + 64x2o + 64x2e",
136+
"nequip-revmd17-toluene",
137+
),
138+
(
139+
"64x0o + 64x0e + 64x1o + 64x1e + 64x2o + 64x2e + 64x3o + 64x3e",
140+
"0e + 1o + 2e + 3o",
141+
"64x0o + 64x0e + 64x1o + 64x1e + 64x2o + 64x2e + 64x3o + 64x3e",
142+
"nequip-revmd17-benzene",
143+
),
144+
(
145+
"32x0o + 32x0e + 32x1o + 32x1e",
146+
"0e + 1o",
147+
"32x0o + 32x0e + 32x1o + 32x1e",
148+
"nequip-water",
149+
),
150+
]
151+
]

0 commit comments

Comments
 (0)