Skip to content

wip improve equivariance test #29

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
curtischong wants to merge 1 commit into
base: main
Choose a base branch
from
Open

wip improve equivariance test #29

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions e3simple_examples/test_equivariance.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from utils.model_utils import plot_3d_coords, seed_everything
import pytest

from utils.rot_utils import random_rotate_data
from utils.rot_utils import get_random_rotation_matrix_3d


@pytest.mark.skip
Expand All @@ -19,7 +19,8 @@ def test_tetris_equivariance():
model = TetrisModel(num_classes=y.shape[1])
for positions in x:
out = model(positions)
out2 = model(random_rotate_data(positions))
rot_mat = get_random_rotation_matrix_3d()
out2 = model(positions @ rot_mat.T)
assert torch.allclose(out, out2, atol=1e-6), "model is not equivariant"
print("the model is equivariant!")

Expand All @@ -34,7 +35,7 @@ def test_tetris_simple_equivariance():
num_equivariance_tests = 5
for _step in range(num_equivariance_tests):
max_equivariance_err = 0.0
model = SimpleModel(
model = TetrisModel(
num_classes=y.shape[1]
) # init a new model so it's weights are random
for positions in x:
Expand All @@ -45,7 +46,8 @@ def test_tetris_simple_equivariance():
# out = torch.mean(model(positions))
out = model(positions)

rotated_pos = random_rotate_data(positions)
rot_mat = get_random_rotation_matrix_3d()
rotated_pos = positions @ rot_mat.T
# plot_3d_coords(rotated_pos.numpy())
# print("pos", positions.tolist())
# print("rotated_pos", rotated_pos.tolist())
Expand All @@ -58,7 +60,7 @@ def test_tetris_simple_equivariance():
data1 = out[i]
data2 = out2[i]
max_equivariance_err = max(max_equivariance_err, abs(data1 - data2))
assert torch.allclose(out, out2, atol=1e-3), "model is not equivariant"
# assert torch.allclose(out, out2, atol=1e-3), "model is not equivariant"
print("max_equivariance_err", max_equivariance_err)
print("the model is equivariant!")

Expand Down
13 changes: 0 additions & 13 deletions utils/rot_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,19 +26,6 @@ def get_random_rotation_matrix_3d() -> torch.Tensor:
return I + torch.sin(angle) * K + (1 - torch.cos(angle)) * (K @ K)


def random_rotate_data(vector: torch.Tensor) -> torch.Tensor:
if vector.shape[-1] != 3:
raise ValueError(
"Input tensor must have the last dimension of size 3 (representing 3D vectors)."
)
rotation_matrix = get_random_rotation_matrix_3d()

# Apply the rotation
rotated_vector = torch.einsum("ij,...j->...i", rotation_matrix, vector)

return rotated_vector


# from e3nn
def D_from_matrix(R: torch.Tensor, l: int, parity: int) -> torch.Tensor:
r"""Matrix of the representation
Expand Down