def test_sh_is_in_irrep():
torch.set_default_dtype(torch.float64)
for l in range(4 + 1):
a, b, _ = o3.rand_angles()
Y = o3.spherical_harmonics_alpha_beta([l], a, b) * math.sqrt(
4 * math.pi) / math.sqrt(2 * l + 1) * (-1)**l
D = o3.wigner_D(l, a, b, 0)
assert (Y - D[:, l]).abs().max() < 1e-10
def test_reduce_tensor_antisymmetric_L2():
torch.set_default_dtype(torch.float64)
irreps, Q = o3.reduce_tensor('ijk=-ikj=-jik', i='2e')
assert irreps[0] == (1, (1, 1))
q = Q[:3].reshape(3, 5, 5, 5)
r = o3.rand_angles()
D1 = o3.wigner_D(1, *r)
D2 = o3.wigner_D(2, *r)
Q1 = torch.einsum('il,jm,kn,zijk->zlmn', D2, D2, D2, q)
Q2 = torch.einsum('yz,zijk->yijk', D1, q)
assert (Q1 - Q2).abs().max() < 1e-10
assert (q + q.transpose(1, 2)).abs().max() < 1e-10
assert (q + q.transpose(1, 3)).abs().max() < 1e-10
assert (q + q.transpose(3, 2)).abs().max() < 1e-10
def test_reduce_tensor_Levi_Civita_symbol():
torch.set_default_dtype(torch.float64)
irreps, Q = o3.reduce_tensor('ijk=-ikj=-jik', i='1e')
assert irreps == ((1, (0, 1)),)
r = o3.rand_angles()
D = o3.wigner_D(1, *r)
Q = Q.reshape(3, 3, 3)
Q1 = torch.einsum('li,mj,nk,ijk', D, D, D, Q)
assert (Q1 - Q).abs().max() < 1e-10
def test_sh_equivariance1():
"""test
- compose
- spherical_harmonics_alpha_beta
- irrep
"""
torch.set_default_dtype(torch.float64)
for l in range(7 + 1):
a, b, _ = o3.rand_angles()
alpha, beta, gamma = o3.rand_angles()
ra, rb, _ = o3.compose_angles(alpha, beta, gamma, a, b,
torch.tensor(0.0))
Yrx = o3.spherical_harmonics_alpha_beta([l], ra, rb)
Y = o3.spherical_harmonics_alpha_beta([l], a, b)
DrY = o3.wigner_D(l, alpha, beta, gamma) @ Y
assert (Yrx - DrY).abs().max() < 1e-10 * Y.abs().max()
def D(self, alpha, beta, gamma, parity=0):
return o3.wigner_D(self.l, alpha, beta, gamma) * self.p**parity
def f(g):
*abc, k = g
return o3.wigner_D(l, *abc) * p**k
def f(g):
return o3.wigner_D(r, *g)