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()