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