Add initial tests.

Author: ameya98

tests/test_extract_irreps.py

@@ -0,0 +1,27 @@
+import pytest
+
+import e3nn
+import torch
+
+from e3tools.nn import ExtractIrreps
+
+
+def test_extract_irreps():
+    irreps_in = e3nn.o3.Irreps("0e + 1o + 2e")
+    input = torch.as_tensor([1., 2., 3., 4., 5., 6., 7., 8., 9.])
+    assert input.shape[-1] == irreps_in.dim
+
+
+    layer = ExtractIrreps(irreps_in, "0e")
+    output = layer(input)
+    assert torch.allclose(output, torch.as_tensor([1.]))
+
+
+    layer = ExtractIrreps(irreps_in, "1o")
+    output = layer(input)
+    assert torch.allclose(output, torch.as_tensor([2., 3., 4.]))
+
+    layer = ExtractIrreps(irreps_in, "2e")
+    output = layer(input)
+    assert torch.allclose(output, torch.as_tensor([5., 6., 7., 8., 9.]))
+

tests/test_pack_unpack.py

@@ -0,0 +1,66 @@
+import pytest
+
+import e3nn
+import torch
+
+from e3tools.nn import AxisToMul, MulToAxis
+
+
+
+@pytest.mark.parametrize(
+    "irreps_in, factor",
+    zip(
+        ["0e + 1o", "8x0e + 8x1o + 8x2e", "8x0e + 8x1o + 8x2e", "3x1o + 3x2o"],
+        [1, 2, 4, 3],
+    ),
+)
+def test_axis_to_mul_shape(irreps_in: str, factor: int, batch_size: int = 5):
+    irreps_in = e3nn.o3.Irreps(irreps_in)
+    layer = AxisToMul(irreps_in, factor)
+    assert layer.irreps_in == irreps_in
+
+    input = irreps_in.randn(batch_size, factor, -1)
+    output = layer(input)
+
+    assert output.shape == (batch_size, factor * irreps_in.dim)
+
+
+@pytest.mark.parametrize(
+    "irreps_in, factor",
+    zip(
+        ["0e + 1o", "8x0e + 8x1o + 8x2e", "8x0e + 8x1o + 8x2e", "3x1o + 3x2o"],
+        [1, 2, 4, 3],
+    ),
+)
+def test_mul_to_axis_shape(irreps_in: str, factor: int, batch_size: int = 5):
+    irreps_in = e3nn.o3.Irreps(irreps_in)
+    layer = MulToAxis(irreps_in, factor)
+    assert layer.irreps_in == irreps_in
+
+    input = irreps_in.randn(batch_size, -1)
+    output = layer(input)
+
+    assert output.shape == (batch_size, factor, irreps_in.dim // factor)
+
+
+
+@pytest.mark.parametrize(
+    "irreps_in, factor",
+    zip(
+        ["0e + 1o", "8x0e + 8x1o + 8x2e", "8x0e + 8x1o + 8x2e", "3x1o + 3x2o"],
+        [1, 2, 4, 3],
+    ),
+)
+def test_inverse(irreps_in: str, factor: int, batch_size: int = 5):
+    irreps_in = e3nn.o3.Irreps(irreps_in)
+    layer = MulToAxis(irreps_in, factor)
+    inv_layer = AxisToMul(layer.irreps_out, factor)
+    
+    assert layer.irreps_in == irreps_in
+    assert inv_layer.irreps_out == irreps_in
+
+    input = irreps_in.randn(batch_size, -1)
+    output = layer(input)
+    recovered = inv_layer(output)
+
+    assert torch.allclose(input, recovered)

tests/test_scaling.py

@@ -0,0 +1,36 @@
+import pytest
+
+import e3nn
+import torch
+
+from e3tools.nn import ScaleIrreps
+
+
+@pytest.mark.parametrize("irreps_in", ["0e + 1o", "0e + 1o + 2e", "3x1o + 2x2o"])
+def test_scale_irreps_by_one(irreps_in):
+    irreps_in = e3nn.o3.Irreps(irreps_in)
+    layer = ScaleIrreps(irreps_in)
+    assert layer.irreps_in == irreps_in
+    assert layer.irreps_out == irreps_in
+
+    input = irreps_in.randn(-1)
+    weight = torch.ones(irreps_in.num_irreps)
+    output = layer(input, weight)
+
+    assert torch.allclose(input, output)
+
+
+@pytest.mark.parametrize("irreps_in", ["0e + 1o", "0e + 1o + 2e", "3x1o + 2x2o"])
+def test_scale_irreps_random(irreps_in: str):
+    irreps_in = e3nn.o3.Irreps(irreps_in)
+    layer = ScaleIrreps(irreps_in)
+    assert layer.irreps_in == irreps_in
+    assert layer.irreps_out == irreps_in
+
+    input = irreps_in.randn(-1)
+    weight = torch.randn(irreps_in.num_irreps)
+    output = layer(input, weight)
+
+    norm = e3nn.o3.Norm(irreps_in)
+    factor = norm(output) / norm(input)
+    assert torch.allclose(factor, torch.abs(weight))