def _array_diag2contr_diagmatrix(expr: ArrayDiagonal): if isinstance(expr.expr, ArrayTensorProduct): args = list(expr.expr.args) diag_indices = list(expr.diagonal_indices) mapping = _get_mapping_from_subranks( [_get_subrank(arg) for arg in args]) tuple_links = [[mapping[j] for j in i] for i in diag_indices] contr_indices = [] total_rank = get_rank(expr) replaced = [False for arg in args] for i, (abs_pos, rel_pos) in enumerate(zip(diag_indices, tuple_links)): if len(abs_pos) != 2: continue (pos1_outer, pos1_inner), (pos2_outer, pos2_inner) = rel_pos arg1 = args[pos1_outer] arg2 = args[pos2_outer] if get_rank(arg1) != 2 or get_rank(arg2) != 2: if replaced[pos1_outer]: diag_indices[i] = None if replaced[pos2_outer]: diag_indices[i] = None continue pos1_in2 = 1 - pos1_inner pos2_in2 = 1 - pos2_inner if arg1.shape[pos1_in2] == 1: if arg1.shape[pos1_inner] != 1: darg1 = DiagMatrix(arg1) else: darg1 = arg1 args.append(darg1) contr_indices.append( ((pos2_outer, pos2_inner), (len(args) - 1, pos1_inner))) total_rank += 1 diag_indices[i] = None args[pos1_outer] = OneArray(arg1.shape[pos1_in2]) replaced[pos1_outer] = True elif arg2.shape[pos2_in2] == 1: if arg2.shape[pos2_inner] != 1: darg2 = DiagMatrix(arg2) else: darg2 = arg2 args.append(darg2) contr_indices.append( ((pos1_outer, pos1_inner), (len(args) - 1, pos2_inner))) total_rank += 1 diag_indices[i] = None args[pos2_outer] = OneArray(arg2.shape[pos2_in2]) replaced[pos2_outer] = True diag_indices_new = [i for i in diag_indices if i is not None] cumul = list(accumulate([0] + [get_rank(arg) for arg in args])) contr_indices2 = [ tuple(cumul[a] + b for a, b in i) for i in contr_indices ] tc = _array_contraction(_array_tensor_product(*args), *contr_indices2) td = _array_diagonal(tc, *diag_indices_new) return td return expr
def __new__(cls, expr, *contraction_indices, **kwargs): contraction_indices = _sort_contraction_indices(contraction_indices) expr = _sympify(expr) if len(contraction_indices) == 0: return expr if isinstance(expr, ArrayContraction): return cls._flatten(expr, *contraction_indices) if isinstance(expr, (ZeroArray, ZeroMatrix)): contraction_indices_flat = [j for i in contraction_indices for j in i] shape = [e for i, e in enumerate(expr.shape) if i not in contraction_indices_flat] return ZeroArray(*shape) if isinstance(expr, PermuteDims): return cls._handle_nested_permute_dims(expr, *contraction_indices) if isinstance(expr, ArrayTensorProduct): expr, contraction_indices = cls._sort_fully_contracted_args(expr, contraction_indices) expr, contraction_indices = cls._lower_contraction_to_addends(expr, contraction_indices) if len(contraction_indices) == 0: return expr if isinstance(expr, ArrayDiagonal): return cls._handle_nested_diagonal(expr, *contraction_indices) if isinstance(expr, ArrayAdd): return ArrayAdd(*[ArrayContraction(i, *contraction_indices) for i in expr.args]) obj = Basic.__new__(cls, expr, *contraction_indices) obj._subranks = _get_subranks(expr) obj._mapping = _get_mapping_from_subranks(obj._subranks) free_indices_to_position = {i: i for i in range(sum(obj._subranks)) if all([i not in cind for cind in contraction_indices])} obj._free_indices_to_position = free_indices_to_position shape = expr.shape cls._validate(expr, *contraction_indices) if shape: shape = tuple(shp for i, shp in enumerate(shape) if not any(i in j for j in contraction_indices)) obj._shape = shape return obj
def identify_hadamard_products(expr: Union[ArrayContraction, ArrayDiagonal]): mapping = _get_mapping_from_subranks(expr.subranks) editor: _EditArrayContraction if isinstance(expr, ArrayContraction): editor = _EditArrayContraction(expr) elif isinstance(expr, ArrayDiagonal): if isinstance(expr.expr, ArrayContraction): editor = _EditArrayContraction(expr.expr) diagonalized = ArrayContraction._push_indices_down( expr.expr.contraction_indices, expr.diagonal_indices) elif isinstance(expr.expr, ArrayTensorProduct): editor = _EditArrayContraction(None) editor.args_with_ind = [ _ArgE(arg) for i, arg in enumerate(expr.expr.args) ] diagonalized = expr.diagonal_indices else: raise NotImplementedError("not implemented") # Trick: add diagonalized indices as negative indices into the editor object: for i, e in enumerate(diagonalized): for j in e: arg_pos, rel_pos = mapping[j] editor.args_with_ind[arg_pos].indices[rel_pos] = -1 - i map_contr_to_args: Dict[FrozenSet, List[_ArgE]] = defaultdict(list) map_ind_to_inds = defaultdict(int) for arg_with_ind in editor.args_with_ind: for ind in arg_with_ind.indices: map_ind_to_inds[ind] += 1 if None in arg_with_ind.indices: continue map_contr_to_args[frozenset(arg_with_ind.indices)].append(arg_with_ind) k: FrozenSet[int] v: List[_ArgE] for k, v in map_contr_to_args.items(): if len(k) != 2: # Hadamard product only defined for matrices: continue if len(v) == 1: # Hadamard product with a single argument makes no sense: continue for ind in k: if map_ind_to_inds[ind] <= 2: # There is no other contraction, skip: continue # Check if expression is a trace: if all([map_ind_to_inds[j] == len(v) and j >= 0 for j in k]): # This is a trace continue # This is a Hadamard product: def check_transpose(x): x = [i if i >= 0 else -1 - i for i in x] return x == sorted(x) hp = hadamard_product(*[ i.element if check_transpose(i.indices) else Transpose(i.element) for i in v ]) hp_indices = v[0].indices if not check_transpose(v[0].indices): hp_indices = list(reversed(hp_indices)) editor.insert_after(v[0], _ArgE(hp, hp_indices)) for i in v: editor.args_with_ind.remove(i) # Count the ranks of the arguments: counter = 0 # Create a collector for the new diagonal indices: diag_indices = defaultdict(list) count_index_freq = Counter() for arg_with_ind in editor.args_with_ind: count_index_freq.update(Counter(arg_with_ind.indices)) free_index_count = count_index_freq[None] # Construct the inverse permutation: inv_perm1 = [] inv_perm2 = [] # Keep track of which diagonal indices have already been processed: done = set([]) # Counter for the diagonal indices: counter4 = 0 for arg_with_ind in editor.args_with_ind: # If some diagonalization axes have been removed, they should be # permuted in order to keep the permutation. # Add permutation here counter2 = 0 # counter for the indices for i in arg_with_ind.indices: if i is None: inv_perm1.append(counter4) counter2 += 1 counter4 += 1 continue if i >= 0: continue # Reconstruct the diagonal indices: diag_indices[-1 - i].append(counter + counter2) if count_index_freq[i] == 1 and i not in done: inv_perm1.append(free_index_count - 1 - i) done.add(i) elif i not in done: inv_perm2.append(free_index_count - 1 - i) done.add(i) counter2 += 1 # Remove negative indices to restore a proper editor object: arg_with_ind.indices = [ i if i is not None and i >= 0 else None for i in arg_with_ind.indices ] counter += len([i for i in arg_with_ind.indices if i is None or i < 0]) inverse_permutation = inv_perm1 + inv_perm2 permutation = _af_invert(inverse_permutation) if isinstance(expr, ArrayContraction): return editor.to_array_contraction() else: # Get the diagonal indices after the detection of HadamardProduct in the expression: diag_indices_filtered = [ tuple(v) for v in diag_indices.values() if len(v) > 1 ] expr1 = editor.to_array_contraction() expr2 = ArrayDiagonal(expr1, *diag_indices_filtered) expr3 = PermuteDims(expr2, permutation) return expr3
def split_multiple_contractions(self): """ Recognize multiple contractions and attempt at rewriting them as paired-contractions. """ from sympy import ask, Q contraction_indices = self.contraction_indices if isinstance(self.expr, ArrayTensorProduct): args = list(self.expr.args) else: args = [self.expr] # TODO: unify API, best location in ArrayTensorProduct subranks = [get_rank(i) for i in args] # TODO: unify API mapping = _get_mapping_from_subranks(subranks) reverse_mapping = {v:k for k, v in mapping.items()} new_contraction_indices = [] for indl, links in enumerate(contraction_indices): if len(links) <= 2: new_contraction_indices.append(links) continue # Check multiple contractions: # # Examples: # # * `A_ij b_j0 C_jk` ===> `A*DiagMatrix(b)*C` # # Care for: # - matrix being diagonalized (i.e. `A_ii`) # - vectors being diagonalized (i.e. `a_i0`) # Also consider the case of diagonal matrices being contracted: current_dimension = self.expr.shape[links[0]] tuple_links = [mapping[i] for i in links] arg_indices, arg_positions = zip(*tuple_links) args_updates = {} if len(arg_indices) != len(set(arg_indices)): # Maybe trace should be supported? raise NotImplementedError not_vectors = [] vectors = [] for arg_ind, arg_pos in tuple_links: mat = args[arg_ind] other_arg_pos = 1-arg_pos other_arg_abs = reverse_mapping[arg_ind, other_arg_pos] if (((1 not in mat.shape) and (not ask(Q.diagonal(mat)))) or ((current_dimension == 1) is True and mat.shape != (1, 1)) or any([other_arg_abs in l for li, l in enumerate(contraction_indices) if li != indl]) ): not_vectors.append((arg_ind, arg_pos)) continue args_updates[arg_ind] = diagonalize_vector(mat) vectors.append((arg_ind, arg_pos)) vectors.append((arg_ind, 1-arg_pos)) if len(not_vectors) > 2: new_contraction_indices.append(links) continue if len(not_vectors) == 0: new_sequence = vectors[:1] + vectors[2:] elif len(not_vectors) == 1: new_sequence = not_vectors[:1] + vectors[:-1] else: new_sequence = not_vectors[:1] + vectors + not_vectors[1:] for i in range(0, len(new_sequence) - 1, 2): arg1, pos1 = new_sequence[i] arg2, pos2 = new_sequence[i+1] if arg1 == arg2: raise NotImplementedError continue abspos1 = reverse_mapping[arg1, pos1] abspos2 = reverse_mapping[arg2, pos2] new_contraction_indices.append((abspos1, abspos2)) for ind, newarg in args_updates.items(): args[ind] = newarg return ArrayContraction( ArrayTensorProduct(*args), *new_contraction_indices )
def _support_function_tp1_recognize(contraction_indices, args): subranks = [get_rank(i) for i in args] coeff = reduce(lambda x, y: x * y, [arg for arg, srank in zip(args, subranks) if srank == 0], S.One) mapping = _get_mapping_from_subranks(subranks) new_contraction_indices = list(contraction_indices) newargs = args[:] # make a copy of the list removed = [None for i in newargs] cumul = list(accumulate([0] + [get_rank(arg) for arg in args])) new_perms = [ list(range(cumul[i], cumul[i + 1])) for i, arg in enumerate(args) ] for pi, contraction_pair in enumerate(contraction_indices): if len(contraction_pair) != 2: continue i1, i2 = contraction_pair a1, e1 = mapping[i1] a2, e2 = mapping[i2] while removed[a1] is not None: a1, e1 = removed[a1] while removed[a2] is not None: a2, e2 = removed[a2] if a1 == a2: trace_arg = newargs[a1] newargs[a1] = Trace(trace_arg)._normalize() new_contraction_indices[pi] = None continue if not isinstance(newargs[a1], MatrixExpr) or not isinstance( newargs[a2], MatrixExpr): continue arg1 = newargs[a1] arg2 = newargs[a2] if (e1 == 1 and e2 == 1) or (e1 == 0 and e2 == 0): arg2 = Transpose(arg2) if e1 == 1: argnew = arg1 * arg2 else: argnew = arg2 * arg1 removed[a2] = a1, e1 new_perms[a1][e1] = new_perms[a2][1 - e2] new_perms[a2] = None newargs[a1] = argnew newargs[a2] = None new_contraction_indices[pi] = None new_contraction_indices = [ i for i in new_contraction_indices if i is not None ] newargs2 = [arg for arg in newargs if arg is not None] if len(newargs2) == 0: return coeff tp = _a2m_tensor_product(*newargs2) tc = ArrayContraction(tp, *new_contraction_indices) new_perms2 = ArrayContraction._push_indices_up( contraction_indices, [i for i in new_perms if i is not None]) permutation = _af_invert( [j for i in new_perms2 for j in i if j is not None]) if permutation == [1, 0] and len(newargs2) == 1: return Transpose(newargs2[0]).doit() tperm = PermuteDims(tc, permutation) return tperm