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_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_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 __matmul__(self, other):
# Tensor product
# Better handle mismatch of features indices
Rs_out, C = rs.tensor_product(self.Rs, other.Rs, o3.selection_rule)
new_signal = torch.einsum('kij,...i,...j->...k',
(C, self.signal, other.signal))
return IrrepTensor(new_signal, Rs_out)
def kernel_geometric(Rs_in,
Rs_out,
selection_rule=o3.selection_rule_in_out_sh,
normalization='component'):
# Compute Clebsh-Gordan coefficients
Rs_f, Q = rs.tensor_product(Rs_in, selection_rule, Rs_out,
normalization) # [out, in, Y]
# Sort filters representation
Rs_f, perm = rs.sort(Rs_f)
Rs_f = rs.simplify(Rs_f)
Q = torch.einsum('ijk,lk->ijl', Q, perm)
del perm
# Normalize the spherical harmonics
if normalization == 'component':
diag = torch.ones(rs.irrep_dim(Rs_f))
if normalization == 'norm':
diag = torch.cat(
[torch.ones(2 * l + 1) / math.sqrt(2 * l + 1) for _, l, _ in Rs_f])
norm_Y = math.sqrt(4 * math.pi) * torch.diag(diag) # [Y, Y]
# Matrix to dispatch the spherical harmonics
mat_Y = rs.map_irrep_to_Rs(Rs_f) # [Rs_f, Y]
mat_Y = mat_Y @ norm_Y
# Create the radial model: R+ -> R^n_path
mat_R = rs.map_mul_to_Rs(Rs_f) # [Rs_f, R]
mixing_matrix = torch.einsum('ijk,ky,kw->ijyw', Q, mat_Y,
mat_R) # [out, in, Y, R]
return Rs_f, mixing_matrix
def __init__(self, Rs_1, Rs_2, selection_rule=o3.selection_rule):
super().__init__()
self.Rs_1 = rs.simplify(Rs_1)
self.Rs_2 = rs.simplify(Rs_2)
Rs_out, mixing_matrix = rs.tensor_product(Rs_1, Rs_2, selection_rule)
self.Rs_out = rs.simplify(Rs_out)
self.register_buffer('mixing_matrix', mixing_matrix)
def kernel_linear(Rs_in, Rs_out):
# Compute Clebsh-Gordan coefficients
def selection_rule(l_in, p_in, l_out, p_out):
if l_in == l_out and p_out in [0, p_in]:
return [0]
return []
Rs_f, Q = rs.tensor_product(Rs_in, selection_rule, Rs_out) # [out, in, w]
Rs_f = rs.simplify(Rs_f)
[(_n_path, l, p)] = Rs_f
assert l == 0 and p in [0, 1]
return Q
def test_tensor_product(self):
torch.set_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)]
Rs_out, m = rs.tensor_product(Rs_1, Rs_2, o3.selection_rule)
mul = TensorProduct(Rs_1, Rs_2)
x1 = torch.randn(1, rs.dim(Rs_1))
x2 = torch.randn(1, rs.dim(Rs_2))
y1 = mul(x1, x2)
y2 = torch.einsum('kij,zi,zj->zk', m, x1, x2)
self.assertEqual(rs.dim(Rs_out), y1.shape[1])
self.assertLess((y1 - y2).abs().max(), 1e-7 * y1.abs().max())
def __init__(self, Rs_in1, Rs_in2, Rs_out, allow_change_output=False):
super().__init__()
self.Rs_in1 = rs.simplify(Rs_in1)
self.Rs_in2 = rs.simplify(Rs_in2)
self.Rs_out = rs.simplify(Rs_out)
ls = [l for _, l, _ in self.Rs_out]
selection_rule = partial(o3.selection_rule, lfilter=lambda l: l in ls)
Rs_ts, T = rs.tensor_product(self.Rs_in1, self.Rs_in2, selection_rule)
register_sparse_buffer(self, 'T', T) # [out, in1 * in2]
ls = [l for _, l, _ in Rs_ts]
if allow_change_output:
self.Rs_out = [(mul, l, p) for mul, l, p in self.Rs_out if l in ls]
else:
assert all(l in ls for _, l, _ in self.Rs_out)
self.kernel = KernelLinear(Rs_ts, self.Rs_out) # [out, in, w]
def test_tensor_product_equal_TensorProduct():
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)]
Rs_out, m = rs.tensor_product(Rs_1,
Rs_2,
o3.selection_rule,
sorted=True)
mul = rs.TensorProduct(Rs_1, Rs_2, o3.selection_rule)
x1 = rs.randn(1, Rs_1)
x2 = rs.randn(1, Rs_2)
y1 = mul(x1, x2)
y2 = torch.einsum('zi,zj->ijz', x1, x2)
y2 = (m @ y2.reshape(rs.dim(Rs_1) * rs.dim(Rs_2), -1)).T
assert rs.dim(Rs_out) == y1.shape[1]
assert (y1 - y2).abs().max() < 1e-10 * y1.abs().max()
