def forward(self, rep): """ Linearly mix a represention. Parameters ---------- rep : :obj:`list` of :obj:`torch.Tensor` Representation to mix. Returns ------- rep : :obj:`list` of :obj:`torch.Tensor` Mixed representation. """ if GTau.from_rep(rep) != self.tau_in: raise ValueError('Tau of input rep does not match initialized tau!' ' rep: {} tau: {}'.format(GTau.from_rep(rep), self.tau_in)) return g_torch.mix(self.weights, rep)
def cg_product(cg_dict, rep1, rep2, maxdim=inf, aggregate=False, ignore_check=False): """ Explicit function to calculate the Clebsch-Gordan product. See the documentation for CGProduct for more information. rep1 : list of :obj:`torch.Tensors` First :obj:`GVector` in the CG product rep2 : list of :obj:`torch.Tensors` First :obj:`GVector` in the CG product maxdim : :obj:`int`, optional Minimum weight to include in CG Product aggregate : :obj:`bool`, optional Apply an "aggregation" operation, or a pointwise convolution with a :obj:`GVector` as a filter. cg_dict : :obj:`CGDict`, optional Specify a Clebsch-Gordan dictionary. If not specified, one will be generated automatically at runtime based upon maxdim. ignore_check : :obj:`bool` Ignore GVec initialization check. Necessary for current implementation of :obj:`zonal_functions`. Use with caution. """ tau1 = GTau.from_rep(rep1) tau2 = GTau.from_rep(rep2) assert tau1.channels and ( tau1.channels == tau2.channels ), 'The number of fragments must be same for each part! {} {}'.format( tau1, tau2) maxk1 = max({key[0] for key in rep1.keys()}) maxn1 = max({key[1] for key in rep1.keys()}) maxk2 = max({key[0] for key in rep2.keys()}) maxn2 = max({key[1] for key in rep2.keys()}) maxDim = min(max(maxk1 + maxk2, maxn1 + maxn2) + 1, maxdim) if (cg_dict.maxdim < maxDim) or (cg_dict.maxdim < max( maxk1, maxn1, maxk2, maxn2)): raise ValueError( 'CG Dictionary maxdim ({}) not sufficiently large for (maxdim, L1, L2) = ({} {} {})' .format(cg_dict.maxdim, maxdim, max1, max2)) assert ( cg_dict.transpose), 'This operation uses transposed CG coefficients!' new_rep = {} for (k1, n1), irrep1 in rep1.items(): for (k2, n2), irrep2 in rep2.items(): if max(k1, n1, k2, n2) > maxDim - 1 or irrep1.shape[ -2] == 0 or irrep2.shape[-2] == 0: continue # cg_mat, aka H, is initially a dictionary {(k,n):rectangular matrix}, # which when flattened/stacked over keys becomes an orthogonal square matrix # we create a sorted list of keys first and then stack the rectangular matrices over keys cg_mat_keys = [ (k, n) for k in range(abs(k1 - k2), min(maxdim, k1 + k2 + 1), 2) for n in range(abs(n1 - n2), min(maxdim, n1 + n2 + 1), 2) ] cg_mat = torch.cat( [cg_dict[((k1, n1), (k2, n2))][key] for key in cg_mat_keys], -2) # Pairwise tensor multiply parts, loop over atom parts accumulating each. irrep_prod = complex_kron_product(irrep1, irrep2, aggregate=aggregate) # Multiply by the CG matrix, effectively turning the product into stacked irreps. Channels are preserved # Have to add a dummy index because matmul acts over the last two dimensions, so the vector dimension on the right needs to be -2 cg_decomp = torch.squeeze( torch.matmul(cg_mat, torch.unsqueeze(irrep_prod, -1)), -1) # Split the result into a list of separate irreps split = [(k + 1) * (n + 1) for (k, n) in cg_mat_keys] cg_decomp = torch.split(cg_decomp, split, dim=-1) # Add the irreps to the dictionary entries, first keeping the channel dimension as a list for idx, key in enumerate(cg_mat_keys): new_rep.setdefault(key, []) new_rep[key].append(cg_decomp[idx]) # at the end concatenate over the channel dimension back into torch tensors new_rep = {key: torch.cat(val, dim=-2) for key, val in new_rep.items()} # TODO: Rewrite so ignore_check not necessary return GVec(new_rep, ignore_check=ignore_check)