add a test for CG

e3nn · Oct 2, 2023 · 38b46bb · 38b46bb
1 parent f65018f
commit 38b46bb
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/e3nn_jax/_src/so3.py b/e3nn_jax/_src/so3.py
@@ -13,9 +13,9 @@ def change_basis_real_to_complex(l: int) -> np.ndarray:
     for m in range(1, l + 1):
         q[l + m, l + abs(m)] = (-1) ** m / np.sqrt(2)
         q[l + m, l - abs(m)] = 1j * (-1) ** m / np.sqrt(2)
-    return (
-        -1j
-    ) ** l * q  # Added factor of 1j**l to make the Clebsch-Gordan coefficients real
+
+    # Added factor of 1j**l to make the Clebsch-Gordan coefficients real
+    return (-1j) ** l * q
 
 
 def clebsch_gordan(l1: int, l2: int, l3: int) -> np.ndarray:

diff --git a/e3nn_jax/_src/so3_test.py b/e3nn_jax/_src/so3_test.py
@@ -24,6 +24,12 @@ def test_clebsch_gordan_symmetry():
         clebsch_gordan(1, 2, 3),
         jnp.swapaxes(jnp.swapaxes(clebsch_gordan(2, 3, 1), 0, 2), 1, 2),
     )
+    assert jnp.allclose(
+        clebsch_gordan(3, 2, 4), -jnp.swapaxes(clebsch_gordan(3, 4, 2), 1, 2)
+    )
+    assert jnp.allclose(
+        clebsch_gordan(2, 3, 4), -jnp.swapaxes(clebsch_gordan(4, 3, 2), 0, 2)
+    )
 
 
 def unique_triplets(lmax):