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