def cg_product_tau(tau1, tau2, maxdim=inf): """ Calulate output multiplicity of the CG Product of two G Vectors given the multiplicty of two input G Vectors. Parameters ---------- tau1 : :class:`list` of :class:`int`, :class:`GTau`. Multiplicity of first representation. tau2 : :class:`list` of :class:`int`, :class:`GTau`. Multiplicity of second representation. maxdim : :class:`int` Largest weight to include in CG Product. Return ------ tau : :class:`GTau` Multiplicity of output representation. """ tau1 = GTau(tau1) tau2 = GTau(tau2) tau = {} for (k1, n1) in tau1.keys(): for (k2, n2) in tau2.keys(): if max(k1, n1, k2, n2) >= maxdim: continue kmin, kmax = abs(k1 - k2), min(k1 + k2, maxdim - 1) nmin, nmax = abs(n1 - n2), min(n1 + n2, maxdim - 1) for k in range(kmin, kmax + 1, 2): for n in range(nmin, nmax + 1, 2): tau.setdefault((k, n), 0) tau[(k, n)] += tau1[(k1, n1)] * tau2[(k2, n2)] return GTau(tau)
nbatch = 1 # number of batches (will be equal for both representations) natoms = 1 # number of atoms in each batch (will be equal for both representations) nchan = 1 # number of channels in each irrep of each representation (must be uniform because of our restrictions) tau1 = GTau({(1, 1): nchan }) # the representation types of the two vectors to be multiplied tau2 = GTau({(1, 1): nchan}) aggregate = True # whether you want to aggregate. This creates rep1 of batch dimension (nbatch, natoms, natoms) and rep2 of (nbatch, natoms) and then sums over one pair of atom indices. (alpha, beta, gamma) = (1 + 2j, 1 + 3j, 1 + 1j ) # Complex Euler angles to rotate by accuracy = 1e-4 # absolute error up to which answers should match ############################################################## TEST maxk1 = max({key[0] for key in tau1.keys()}) maxn1 = max({key[1] for key in tau1.keys()}) maxk2 = max({key[0] for key in tau2.keys()}) maxn2 = max({key[1] for key in tau2.keys()}) maxdim = max(maxk1 + maxk2, maxn1 + maxn2) + 1 cg_dict = CGDict(maxdim=maxdim) print("Running covariance test with parameters:") print("nbatch=", nbatch) print("natoms=", natoms) print("tau1=", tau1) print("tau2=", tau2) print("(alpha,beta,gamma)=", (alpha, beta, gamma)) print("accuracy=", accuracy, "\n--------------------------------") cartesian4 = torch.tensor([[[1, 0, 0, 0], [0, 1 / sqrt(2.), 0, 0],
class CatReps(Module): """ Module to concanteate a list of reps. Specify input type for error checking and to allow network to fit into main architecture. Parameters ---------- taus_in : :obj:`list` of :obj:`GTau` or compatible. List of taus of input reps. maxdim : :obj:`bool`, optional Maximum weight to include in concatenation. """ def __init__(self, taus_in, maxdim=None): super().__init__() self.taus_in = taus_in = [GTau(tau) for tau in taus_in if tau] if maxdim is None: maxdim = max(sum(dict(i for tau in taus_in for i in tau.items()), ())) + 1 self.maxdim = maxdim self.taus_in = taus_in self.tau_out = {} for tau in taus_in: for key, val in tau.items(): if val > 0: if max(key) <= maxdim - 1: self.tau_out.setdefault(key, 0) self.tau_out[key] += val self.tau_out = GTau(self.tau_out) self.all_keys = list(self.tau_out.keys()) def forward(self, reps): """ Concatenate a list of reps Parameters ---------- reps : :obj:`list` of :obj:`GTensor` subclasses List of representations to concatenate. Returns ------- reps_cat : :obj:`list` of :obj:`torch.Tensor` """ # Drop Nones reps = [rep for rep in reps if rep is not None] # Error checking reps_taus_in = [rep.tau for rep in reps] if reps_taus_in != self.taus_in: raise ValueError('Tau of input reps does not match predefined version!' 'got: {} expected: {}'.format(reps_taus_in, self.taus_in)) if self.maxdim is not None: reps = [rep.truncate(self.maxdim) for rep in reps] return g_torch.cat(reps) @property def tau(self): return self.tau_out