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)