Add tests.

e3nn · Dec 2, 2024 · 7ef3f63 · 7ef3f63
1 parent 8f22a3a
commit 7ef3f63
Showing 1 changed file with 51 additions and 0 deletions.
diff --git a/tests/_src/so3grid_test.py b/tests/_src/so3grid_test.py
@@ -46,3 +46,54 @@ def test_sampling(num_seeds: int = 10):
     # Check that the samples are orthogonal.
     for R in Rs:
         assert jnp.allclose(R @ R.T, jnp.eye(3), atol=1e-6)
+
+
+def test_multiplication_scalar():
+    sig1 = SO3Signal.from_function(
+        lambda R: jnp.trace(R @ R),
+        res_beta=40,
+        res_alpha=39,
+        res_theta=40,
+        quadrature="gausslegendre",
+    )
+    integral1 = sig1.integrate()
+    sig2 = sig1 * 2.7
+    integral2 = sig2.integrate()
+    assert jnp.isclose(integral2, 2.7 * integral1)
+
+
+def test_multiplication_signal():
+    sig1 = SO3Signal.from_function(
+        lambda R: jnp.trace(R @ R),
+        res_beta=40,
+        res_alpha=39,
+        res_theta=40,
+        quadrature="gausslegendre",
+    )
+    integral1 = sig1.integrate()
+
+    ones = SO3Signal.from_function(
+        lambda R: 1.0,
+        res_beta=40,
+        res_alpha=39,
+        res_theta=40,
+        quadrature="gausslegendre",
+    )
+    sig2 = sig1 * ones
+    integral2 = sig2.integrate()
+    assert jnp.isclose(integral2, integral1)
+
+
+def test_division_scalar():
+    sig1 = SO3Signal.from_function(
+        lambda R: jnp.trace(R @ R),
+        res_beta=40,
+        res_alpha=39,
+        res_theta=40,
+        quadrature="gausslegendre",
+    )
+    integral1 = sig1.integrate()
+    sig2 = sig1 / 2.7
+    integral2 = sig2.integrate()
+    assert jnp.isclose(integral2, integral1 / 2.7)
+