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
