def test_rot_to_abc(self): with o3.torch_default_dtype(torch.float64): R = o3.rand_rot() abc = o3.rot_to_abc(R) R2 = o3.rot(*abc) d = (R - R2).norm() / R.norm() self.assertTrue(d < 1e-10, d)
def test_tensor_square_norm(self): for Rs_in in [[(1, 0), (2, 1), (4, 3)]]: with o3.torch_default_dtype(torch.float64): Rs_out, Q = rs.tensor_square(Rs_in, o3.selection_rule, normalization='component', sorted=True) abc = o3.rand_angles() D_in = rs.rep(Rs_in, *abc) D_out = rs.rep(Rs_out, *abc) Q1 = torch.einsum("ijk,il->ljk", (Q, D_out)) Q2 = torch.einsum("li,mj,kij->klm", (D_in, D_in, Q)) d = (Q1 - Q2).pow(2).mean().sqrt() / Q1.pow(2).mean().sqrt() self.assertLess(d, 1e-10) n = Q.size(0) M = Q.reshape(n, -1) I = torch.eye(n) d = ((M @ M.t()) - I).pow(2).mean().sqrt() self.assertLess(d, 1e-10)
def test_reduce_tensor_product(self): for Rs_i, Rs_j in [([(1, 0)], [(2, 0)]), ([(3, 1), (2, 2)], [(2, 0), (1, 1), (1, 3)])]: with o3.torch_default_dtype(torch.float64): Rs, Q = rs.tensor_product(Rs_i, Rs_j) abc = torch.rand(3, dtype=torch.float64) D_i = o3.direct_sum(*[ o3.irr_repr(l, *abc) for mul, l in Rs_i for _ in range(mul) ]) D_j = o3.direct_sum(*[ o3.irr_repr(l, *abc) for mul, l in Rs_j for _ in range(mul) ]) D = o3.direct_sum(*[ o3.irr_repr(l, *abc) for mul, l, _ in Rs for _ in range(mul) ]) Q1 = torch.einsum("ijk,il->ljk", (Q, D)) Q2 = torch.einsum("li,mj,kij->klm", (D_i, D_j, Q)) d = (Q1 - Q2).pow(2).mean().sqrt() / Q1.pow(2).mean().sqrt() self.assertLess(d, 1e-10) n = Q.size(0) M = Q.view(n, n) I = torch.eye(n, dtype=M.dtype) d = ((M @ M.t()) - I).pow(2).mean().sqrt() self.assertLess(d, 1e-10) d = ((M.t() @ M) - I).pow(2).mean().sqrt() self.assertLess(d, 1e-10)
def test_sh_dirac(): with o3.torch_default_dtype(torch.float64): for l in range(5): r = torch.randn(3) a = spherical_harmonics_dirac(r, l) v = SphericalTensor(a).signal_xyz(r) assert v.sub(1).abs() < 1e-10
def test_sh_is_in_irrep(): with o3.torch_default_dtype(torch.float64): for l in range(4 + 1): a, b = 3.14 * torch.rand(2) # works only for beta in [0, pi] Y = rsh.spherical_harmonics_alpha_beta([l], a, b) * math.sqrt(4 * math.pi) / math.sqrt(2 * l + 1) * (-1) ** l D = o3.irr_repr(l, a, b, 0) assert (Y - D[:, l]).norm() < 1e-10
def test_sh_dirac(self): with o3.torch_default_dtype(torch.float64): for l in range(5): angles = torch.tensor(1.2), torch.tensor(2.1) a = sphten.spherical_harmonics_dirac(torch.stack(o3.angles_to_xyz(*angles), dim=-1), l) v = sphten.SphericalTensor(a, 1, l).value(*angles) self.assertAlmostEqual(v.item(), 1)
def test_tensor_product_norm(self): for Rs_in1, Rs_in2 in [([(1, 0)], [(2, 0)]), ([(3, 1), (2, 2)], [(2, 0), (1, 1), (1, 3)])]: with o3.torch_default_dtype(torch.float64): Rs_out, Q = rs.tensor_product(Rs_in1, Rs_in2, o3.selection_rule) abc = torch.rand(3, dtype=torch.float64) D_in1 = rs.rep(Rs_in1, *abc) D_in2 = rs.rep(Rs_in2, *abc) D_out = rs.rep(Rs_out, *abc) Q1 = torch.einsum("ijk,il->ljk", (Q, D_out)) Q2 = torch.einsum("li,mj,kij->klm", (D_in1, D_in2, Q)) d = (Q1 - Q2).pow(2).mean().sqrt() / Q1.pow(2).mean().sqrt() self.assertLess(d, 1e-10) n = Q.size(0) M = Q.reshape(n, n) I = torch.eye(n, dtype=M.dtype) d = ((M @ M.t()) - I).pow(2).mean().sqrt() self.assertLess(d, 1e-10) d = ((M.t() @ M) - I).pow(2).mean().sqrt() self.assertLess(d, 1e-10)
def test_scipy_spherical_harmonics(): with o3.torch_default_dtype(torch.float64): ls = [0, 1, 2, 3, 4, 5] beta = torch.linspace(1e-3, math.pi - 1e-3, 100, requires_grad=True).reshape(1, -1) alpha = torch.linspace(0, 2 * math.pi, 100, requires_grad=True).reshape(-1, 1) Y1 = rsh.spherical_harmonics_alpha_beta(ls, alpha, beta) Y2 = rsh.spherical_harmonics_alpha_beta(ls, alpha.detach(), beta.detach()) assert (Y1 - Y2).abs().max() < 1e-10
def test_tensor_product_to_dense(): with o3.torch_default_dtype(torch.float64): Rs_1 = [(3, 0), (2, 1), (5, 2)] Rs_2 = [(1, 0), (2, 1), (2, 2), (2, 0), (2, 1), (1, 2)] mul = rs.TensorProduct(Rs_1, Rs_2, o3.selection_rule) assert mul.to_dense().shape == (rs.dim(mul.Rs_out), rs.dim(Rs_1), rs.dim(Rs_2))
def test_wigner_3j_orthogonal(self): with o3.torch_default_dtype(torch.float64): for l_out in range(3 + 1): for l_in in range(l_out, 4 + 1): for l_f in range(abs(l_out - l_in), l_out + l_in + 1): Q = o3.wigner_3j(l_f, l_in, l_out).reshape(2 * l_f + 1, -1) e = (2 * l_f + 1) * Q @ Q.t() d = e - torch.eye(2 * l_f + 1) self.assertLess(d.pow(2).mean().sqrt(), 1e-10)
def test_wigner_3j_sh_norm(): with o3.torch_default_dtype(torch.float64): for l_out in range(3 + 1): for l_in in range(l_out, 4 + 1): for l_f in range(abs(l_out - l_in), l_out + l_in + 1): Q = o3.wigner_3j(l_out, l_in, l_f) Y = rsh.spherical_harmonics_xyz([l_f], torch.randn(3)) QY = math.sqrt(4 * math.pi) * Q @ Y assert abs(QY.norm() - 1) < 1e-10
def test_map_irrep_to_Rs(): with o3.torch_default_dtype(torch.float64): Rs = [(3, 0)] mapping_matrix = rs.map_irrep_to_Rs(Rs) assert torch.allclose(mapping_matrix, torch.ones(3, 1)) Rs = [(1, 0), (1, 1), (1, 2)] mapping_matrix = rs.map_irrep_to_Rs(Rs) assert torch.allclose(mapping_matrix, torch.eye(1 + 3 + 5))
def test_kron(self): with o3.torch_default_dtype(torch.float64): m1 = torch.randn(4, 4) m2 = torch.randn(3, 5) m3 = torch.randn(6, 6) x1 = o3.kron(m1, m2, m3) x2 = o3.kron(m1, o3.kron(m2, m3)) assert torch.allclose(x1, x2)
def test_xyz_to_irreducible_basis(self, ): with o3.torch_default_dtype(torch.float64): A = o3.xyz_to_irreducible_basis() a, b, c = torch.rand(3) r1 = A.t() @ o3.irr_repr(1, a, b, c) @ A r2 = o3.rot(a, b, c) assert torch.allclose(r1, r2)
def test_xyz_vector_basis_to_spherical_basis(self, ): with o3.torch_default_dtype(torch.float64): A = o3.xyz_vector_basis_to_spherical_basis() a, b, c = torch.rand(3) r1 = A.t() @ o3.irr_repr(1, a, b, c) @ A r2 = o3.rot(a, b, c) self.assertLess((r1 - r2).abs().max(), 1e-10)
def test_sh_norm(): with o3.torch_default_dtype(torch.float64): l_filter = list(range(15)) Ys = [ rsh.spherical_harmonics_xyz([l], torch.randn(10, 3)) for l in l_filter ] s = torch.stack([Y.pow(2).mean(-1) for Y in Ys]) d = s - 1 / (4 * math.pi) assert d.pow(2).mean().sqrt() < 1e-10
def test_sh_parity(self): """ (-1)^l Y(x) = Y(-x) """ with o3.torch_default_dtype(torch.float64): for l in range(7 + 1): x = torch.randn(3) Y1 = (-1)**l * o3.spherical_harmonics_xyz(l, x) Y2 = o3.spherical_harmonics_xyz(l, -x) self.assertLess((Y1 - Y2).abs().max(), 1e-10 * Y1.abs().max())
def test_sh_norm(self): with o3.torch_default_dtype(torch.float64): l_filter = list(range(15)) Ys = [ o3.spherical_harmonics_xyz(l, torch.randn(10, 3)) for l in l_filter ] s = torch.stack([Y.pow(2).mean(0) for Y in Ys]) d = s - 1 / (4 * math.pi) self.assertLess(d.pow(2).mean().sqrt(), 1e-10)
def test_clebsch_gordan_sh_norm(self): with o3.torch_default_dtype(torch.float64): for l_out in range(6): for l_in in range(6): for l_f in range(abs(l_out - l_in), l_out + l_in + 1): Q = o3.clebsch_gordan(l_out, l_in, l_f) Y = o3.spherical_harmonics_xyz(l_f, torch.randn( 1, 3)).view(2 * l_f + 1) QY = math.sqrt(4 * math.pi) * Q @ Y self.assertLess(abs(QY.norm() - 1), 1e-10)
def test_sh_parity(): """ (-1)^l Y(x) = Y(-x) """ with o3.torch_default_dtype(torch.float64): for l in range(7 + 1): x = torch.randn(3) Y1 = (-1) ** l * rsh.spherical_harmonics_xyz([l], x) Y2 = rsh.spherical_harmonics_xyz([l], -x) assert (Y1 - Y2).abs().max() < 1e-10 * Y1.abs().max()
def test_clebsch_gordan_orthogonal(self): with o3.torch_default_dtype(torch.float64): for l_out in range(6): for l_in in range(6): for l_f in range(abs(l_out - l_in), l_out + l_in + 1): Q = o3.clebsch_gordan(l_f, l_in, l_out).view(2 * l_f + 1, -1) e = (2 * l_f + 1) * Q @ Q.t() d = e - torch.eye(2 * l_f + 1) self.assertLess(d.pow(2).mean().sqrt(), 1e-10)
def test_sh_cuda_ordered_full(self): if torch.cuda.is_available(): with o3.torch_default_dtype(torch.float64): l = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] x = torch.randn(10, 3) x_cuda = x.cuda() Y1 = o3.spherical_harmonics_xyz(l, x) Y2 = o3.spherical_harmonics_xyz(l, x_cuda).cpu() self.assertLess((Y1 - Y2).abs().max(), 1e-7) else: print("Cuda is not available! test_sh_cuda_ordered_full skipped!")
def test_sh_cuda_ordered_partial(): if torch.cuda.is_available(): with o3.torch_default_dtype(torch.float64): l = [0, 2, 5, 7, 10] x = torch.randn(10, 3) x_cuda = x.cuda() Y1 = rsh.spherical_harmonics_xyz(l, x) Y2 = rsh.spherical_harmonics_xyz(l, x_cuda).cpu() assert (Y1 - Y2).abs().max() < 1e-7 else: print("Cuda is not available! test_sh_cuda_ordered_partial skipped!")
def test_sh_cuda_single(): if torch.cuda.is_available(): with o3.torch_default_dtype(torch.float64): for l in range(10 + 1): x = torch.randn(10, 3) x_cuda = x.cuda() Y1 = rsh.spherical_harmonics_xyz([l], x) Y2 = rsh.spherical_harmonics_xyz([l], x_cuda).cpu() assert (Y1 - Y2).abs().max() < 1e-7 else: print("Cuda is not available! test_sh_cuda_single skipped!")
def test_normalization(self): with o3.torch_default_dtype(torch.float64): lmax = 5 res = (20, 30) for normalization in ['component', 'norm']: to = s2grid.ToS2Grid(lmax, res, normalization=normalization) x = rs.randn(50, [(1, l) for l in range(lmax + 1)], normalization=normalization) y = to(x) self.assertAlmostEqual(y.var().item(), 1, delta=0.2)
def test_tensor_product_in_out_normalization(Rs_in1, Rs_out): with o3.torch_default_dtype(torch.float64): n = rs.dim(Rs_out) I = torch.eye(n) _, Q = rs.tensor_product(Rs_in1, o3.selection_rule, Rs_out) d = ((Q @ Q.t()).to_dense() - I).pow(2).mean().sqrt() assert d < 1e-10 _, Q = rs.tensor_product(o3.selection_rule, Rs_in1, Rs_out) d = ((Q @ Q.t()).to_dense() - I).pow(2).mean().sqrt() assert d < 1e-10
def test_derivative_irr_repr(self): with o3.torch_default_dtype(torch.float64): l = 4 angles = o3.rand_angles() da, db, dc = o3.derivative_irr_repr(l, *angles) h = 1e-7 da1 = (o3.irr_repr(l, angles[0] + h, angles[1], angles[2]) - o3.irr_repr(l, *angles)) / h db1 = (o3.irr_repr(l, angles[0], angles[1] + h, angles[2]) - o3.irr_repr(l, *angles)) / h dc1 = (o3.irr_repr(l, angles[0], angles[1], angles[2] + h) - o3.irr_repr(l, *angles)) / h self.assertLess((da1 - da).abs().max(), 1e-5) self.assertLess((db1 - db).abs().max(), 1e-5) self.assertLess((dc1 - dc).abs().max(), 1e-5)
def test_tensor_product_in_in_normalization_norm(Rs_in1, Rs_in2): with o3.torch_default_dtype(torch.float64): tp = rs.TensorProduct(Rs_in1, Rs_in2, o3.selection_rule, normalization='norm') x1 = rs.randn(10, Rs_in1, normalization='norm') x2 = rs.randn(10, Rs_in2, normalization='norm') n = Norm(tp.Rs_out, normalization='norm') x = n(tp(x1, x2)).mean(0) assert (x.log10().abs() < 1).all()
def test_map_mul_to_Rs(): with o3.torch_default_dtype(torch.float64): Rs = [(3, 0)] mapping_matrix = rs.map_mul_to_Rs(Rs) assert torch.allclose(mapping_matrix, torch.eye(3)) Rs = [(1, 0), (1, 1), (1, 2)] mapping_matrix = rs.map_mul_to_Rs(Rs) check_matrix = torch.zeros(1 + 3 + 5, 3) check_matrix[0, 0] = 1. check_matrix[1:4, 1] = 1. check_matrix[4:, 2] = 1. assert torch.allclose(mapping_matrix, check_matrix)
def test_tensor_square_norm(): for Rs_in in [[(1, 0), (1, 1)]]: with o3.torch_default_dtype(torch.float64): Rs_out, Q = rs.tensor_square(Rs_in, o3.selection_rule, normalization='component', sorted=True) I1 = (Q @ Q.t()).to_dense() I2 = torch.eye(rs.dim(Rs_out)) d = (I1 - I2).pow(2).mean().sqrt() assert d < 1e-10