def __init__(self,
input1,
input2,
output,
normalization='component',
sorted=True):
super().__init__()
Rs, mat = tensor_product(input1, input2, output, normalization, sorted)
if callable(input1):
self.Rs_in1 = Rs
self.Rs_in2 = convention(input2)
self.Rs_out = convention(output)
self._complete = 'in1'
if callable(input2):
self.Rs_in1 = convention(input1)
self.Rs_in2 = Rs
self.Rs_out = convention(output)
self._complete = 'in2'
if callable(output):
self.Rs_in1 = convention(input1)
self.Rs_in2 = convention(input2)
self.Rs_out = Rs
self._complete = 'out'
register_sparse_buffer(self, 'mixing_matrix', mat)
def __init__(self,
Rs_in,
Rs_out,
linear=True,
allow_change_output=False,
allow_zero_outputs=False):
super().__init__()
self.Rs_in = rs.simplify(Rs_in)
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)
if linear:
Rs_in = [(1, 0, 1)] + self.Rs_in
else:
Rs_in = self.Rs_in
self.linear = linear
Rs_ts, T = rs.tensor_square(Rs_in, 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]
elif not allow_zero_outputs:
assert all(l in ls for _, l, _ in self.Rs_out)
self.kernel = KernelLinear(Rs_ts, self.Rs_out) # [out, in, w]
def __init__(self, Rs_in, selection_rule=o3.selection_rule):
super().__init__()
self.Rs_in = simplify(Rs_in)
self.Rs_out, mixing_matrix = tensor_square(self.Rs_in, selection_rule, sorted=True)
register_sparse_buffer(self, 'mixing_matrix', mixing_matrix)
def __init__(self, Rs_in1, Rs_in2, selection_rule=o3.selection_rule):
super().__init__()
self.Rs_in1 = simplify(Rs_in1)
self.Rs_in2 = simplify(Rs_in2)
assert sum(mul for mul, _, _ in self.Rs_in1) == sum(mul for mul, _, _ in self.Rs_in2)
self.Rs_out, mixing_matrix = elementwise_tensor_product(self.Rs_in1, self.Rs_in2, selection_rule)
register_sparse_buffer(self, "mixing_matrix", mixing_matrix)
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 __init__(self, Rs):
super().__init__()
self.Rs_in = convention(Rs)
self.mul, self.Rs_out = transpose_mul(self.Rs_in)
register_sparse_buffer(self, 'mixing_matrix',
rearrange(self.Rs_in, self.mul * self.Rs_out))
