def _(expr: ArrayDiagonal): newexpr, removed = _remove_trivial_dims(expr.expr) shifts = list( accumulate( [0] + [1 if i in removed else 0 for i in range(get_rank(expr.expr))])) new_diag_indices_map = { i: tuple(j for j in i if j not in removed) for i in expr.diagonal_indices } for old_diag_tuple, new_diag_tuple in new_diag_indices_map.items(): if len(new_diag_tuple) == 1: removed = [i for i in removed if i not in old_diag_tuple] new_diag_indices = [ tuple(j - shifts[j] for j in i) for i in new_diag_indices_map.values() ] rank = get_rank(expr.expr) removed = ArrayDiagonal._push_indices_up(expr.diagonal_indices, removed, rank) removed = sorted({i for i in removed}) # If there are single axes to diagonalize remaining, it means that their # corresponding dimension has been removed, they no longer need diagonalization: new_diag_indices = [i for i in new_diag_indices if len(i) > 0] if len(new_diag_indices) > 0: newexpr2 = _array_diagonal(newexpr, *new_diag_indices, allow_trivial_diags=True) else: newexpr2 = newexpr if isinstance(newexpr2, ArrayDiagonal): newexpr3, removed2 = _remove_diagonalized_identity_matrices(newexpr2) removed = _combine_removed(-1, removed, removed2) return newexpr3, removed else: return newexpr2, removed
def _(expr: ArrayContraction): new_expr, removed0 = _array_contraction_to_diagonal_multiple_identity(expr) if new_expr != expr: new_expr2, removed1 = _remove_trivial_dims(_array2matrix(new_expr)) removed = _combine_removed(-1, removed0, removed1) return new_expr2, removed rank1 = get_rank(expr) expr, removed1 = remove_identity_matrices(expr) if not isinstance(expr, ArrayContraction): expr2, removed2 = _remove_trivial_dims(expr) return expr2, _combine_removed(rank1, removed1, removed2) newexpr, removed2 = _remove_trivial_dims(expr.expr) shifts = list( accumulate( [1 if i in removed2 else 0 for i in range(get_rank(expr.expr))])) new_contraction_indices = [ tuple(j for j in i if j not in removed2) for i in expr.contraction_indices ] # Remove possible empty tuples "()": new_contraction_indices = [ i for i in new_contraction_indices if len(i) > 0 ] contraction_indices_flat = [j for i in expr.contraction_indices for j in i] removed2 = [i for i in removed2 if i not in contraction_indices_flat] new_contraction_indices = [ tuple(j - shifts[j] for j in i) for i in new_contraction_indices ] # Shift removed2: removed2 = ArrayContraction._push_indices_up(expr.contraction_indices, removed2) removed = _combine_removed(rank1, removed1, removed2) return _array_contraction(newexpr, *new_contraction_indices), list(removed)
def _(expr: ElementwiseApplyFunction, x: Expr): assert get_rank(expr) == 2 assert get_rank(x) == 2 fdiff = expr._get_function_fdiff() dexpr = array_derive(expr.expr, x) tp = ArrayTensorProduct(ElementwiseApplyFunction(fdiff, expr.expr), dexpr) td = ArrayDiagonal(tp, (0, 4), (1, 5)) return td
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 _(expr: ArrayElementwiseApplyFunc, x: Expr): fdiff = expr._get_function_fdiff() subexpr = expr.expr dsubexpr = array_derive(subexpr, x) tp = _array_tensor_product(dsubexpr, ArrayElementwiseApplyFunc(fdiff, subexpr)) b = get_rank(x) c = get_rank(expr) diag_indices = [(b + i, b + c + i) for i in range(c)] return _array_diagonal(tp, *diag_indices)
def _find_trivial_kronecker_products_broadcast(expr: ArrayTensorProduct): newargs: List[Basic] = [] removed = [] count_dims = 0 for i, arg in enumerate(expr.args): count_dims += get_rank(arg) shape = get_shape(arg) current_range = [count_dims - i for i in range(len(shape), 0, -1)] if (shape == (1, 1) and len(newargs) > 0 and 1 not in get_shape(newargs[-1]) and isinstance(newargs[-1], MatrixExpr) and isinstance(arg, MatrixExpr)): # KroneckerProduct object allows the trick of broadcasting: newargs[-1] = KroneckerProduct(newargs[-1], arg) removed.extend(current_range) elif 1 not in shape and len(newargs) > 0 and get_shape( newargs[-1]) == (1, 1): # Broadcast: newargs[-1] = KroneckerProduct(newargs[-1], arg) prev_range = [ i for i in range(min(current_range)) if i not in removed ] removed.extend(prev_range[-2:]) else: newargs.append(arg) return _array_tensor_product(*newargs), removed
def _(expr: ArrayDiagonal): newexpr, removed = _remove_trivial_dims(expr.expr) shifts = list( accumulate( [0] + [1 if i in removed else 0 for i in range(get_rank(expr.expr))])) new_diag_indices = [ tuple(j for j in i if j not in removed) for i in expr.diagonal_indices ] new_diag_indices = [ tuple(j - shifts[j] for j in i) for i in new_diag_indices ] rank = get_rank(expr.expr) removed = ArrayDiagonal._push_indices_up(expr.diagonal_indices, removed, rank) removed = sorted({i for i in removed}) # If there are single axes to diagonalize remaining, it means that their # corresponding dimension has been removed, they no longer need diagonalization: new_diag_indices = [i for i in new_diag_indices if len(i) > 1] return ArrayDiagonal(newexpr, *new_diag_indices), removed
def _(expr: ArrayContraction): newexpr, removed = _remove_trivial_dims(expr.expr) shifts = list( accumulate( [1 if i in removed else 0 for i in range(get_rank(expr.expr))])) new_contraction_indices = [ tuple(j for j in i if j not in removed) for i in expr.contraction_indices ] # Remove possible empty tuples "()": new_contraction_indices = [ i for i in new_contraction_indices if len(i) > 0 ] contraction_indices_flat = [j for i in expr.contraction_indices for j in i] removed = [i for i in removed if i not in contraction_indices_flat] new_contraction_indices = [ tuple(j - shifts[j] for j in i) for i in new_contraction_indices ] # Shift removed: removed = ArrayContraction._push_indices_up(expr.contraction_indices, removed) return ArrayContraction(newexpr, *new_contraction_indices), list(removed)
def _remove_diagonalized_identity_matrices(expr: ArrayDiagonal): assert isinstance(expr, ArrayDiagonal) editor = _EditArrayContraction(expr) mapping = { i: {j for j in editor.args_with_ind if i in j.indices} for i in range(-1, -1 - editor.number_of_diagonal_indices, -1) } removed = [] counter: int = 0 for i, arg_with_ind in enumerate(editor.args_with_ind): counter += len(arg_with_ind.indices) if isinstance(arg_with_ind.element, Identity): if None in arg_with_ind.indices and any( i is not None and (i < 0) == True for i in arg_with_ind.indices): diag_ind = [j for j in arg_with_ind.indices if j is not None][0] other = [j for j in mapping[diag_ind] if j != arg_with_ind][0] if not isinstance(other.element, MatrixExpr): continue if 1 not in other.element.shape: continue if None not in other.indices: continue editor.args_with_ind[i].element = None none_index = other.indices.index(None) other.element = DiagMatrix(other.element) other_range = editor.get_absolute_range(other) removed.extend([other_range[0] + none_index]) editor.args_with_ind = [ i for i in editor.args_with_ind if i.element is not None ] removed = ArrayDiagonal._push_indices_up(expr.diagonal_indices, removed, get_rank(expr.expr)) return editor.to_array_contraction(), removed
def _find_trivial_matrices_rewrite(expr: ArrayTensorProduct): # If there are matrices of trivial shape in the tensor product (i.e. shape # (1, 1)), try to check if there is a suitable non-trivial MatMul where the # expression can be inserted. # For example, if "a" has shape (1, 1) and "b" has shape (k, 1), the # expressions "_array_tensor_product(a, b*b.T)" can be rewritten as # "b*a*b.T" trivial_matrices = [] pos: Optional[int] = None first: Optional[MatrixExpr] = None second: Optional[MatrixExpr] = None removed: List[int] = [] counter: int = 0 args: List[Optional[Basic]] = [i for i in expr.args] for i, arg in enumerate(expr.args): if isinstance(arg, MatrixExpr): if arg.shape == (1, 1): trivial_matrices.append(arg) args[i] = None removed.extend([counter, counter + 1]) elif pos is None and isinstance(arg, MatMul): margs = arg.args for j, e in enumerate(margs): if isinstance(e, MatrixExpr) and e.shape[1] == 1: pos = i first = MatMul.fromiter(margs[:j + 1]) second = MatMul.fromiter(margs[j + 1:]) break counter += get_rank(arg) if pos is None: return expr, [] args[pos] = (first * MatMul.fromiter(i for i in trivial_matrices) * second).doit() return _array_tensor_product(*[i for i in args if i is not None]), removed
def do_convert(self, expr, indices): if isinstance(expr, ArrayTensorProduct): cumul = list(accumulate([0] + [get_rank(arg) for arg in expr.args])) indices_grp = [indices[cumul[i]:cumul[i+1]] for i in range(len(expr.args))] return Mul.fromiter(self.do_convert(arg, ind) for arg, ind in zip(expr.args, indices_grp)) if isinstance(expr, ArrayContraction): new_indices = [None for i in range(get_rank(expr.expr))] limits = [] bottom_shape = get_shape(expr.expr) for contraction_index_grp in expr.contraction_indices: d = Dummy(f"d{self.count_dummies}") self.count_dummies += 1 dim = bottom_shape[contraction_index_grp[0]] limits.append((d, 0, dim-1)) for i in contraction_index_grp: new_indices[i] = d j = 0 for i in range(len(new_indices)): if new_indices[i] is None: new_indices[i] = indices[j] j += 1 newexpr = self.do_convert(expr.expr, new_indices) return Sum(newexpr, *limits) if isinstance(expr, ArrayDiagonal): new_indices = [None for i in range(get_rank(expr.expr))] ind_pos = expr._push_indices_down(expr.diagonal_indices, list(range(len(indices))), get_rank(expr)) for i, index in zip(ind_pos, indices): if isinstance(i, collections.abc.Iterable): for j in i: new_indices[j] = index else: new_indices[i] = index newexpr = self.do_convert(expr.expr, new_indices) return newexpr if isinstance(expr, PermuteDims): permuted_indices = _apply_permutation_to_list(expr.permutation, indices) return self.do_convert(expr.expr, permuted_indices) if isinstance(expr, ArrayAdd): return Add.fromiter(self.do_convert(arg, indices) for arg in expr.args) if isinstance(expr, _ArrayExpr): return expr.__getitem__(tuple(indices)) if isinstance(expr, ArrayElementwiseApplyFunc): return expr.function(self.do_convert(expr.expr, indices)) if isinstance(expr, Reshape): shape_up = expr.shape shape_down = get_shape(expr.expr) cumul = list(accumulate([1] + list(reversed(shape_up)), operator.mul)) one_index = Add.fromiter(i*s for i, s in zip(reversed(indices), cumul)) dest_indices = [None for _ in shape_down] c = 1 for i, e in enumerate(reversed(shape_down)): if c == 1: if i == len(shape_down) - 1: dest_indices[i] = one_index else: dest_indices[i] = one_index % e elif i == len(shape_down) - 1: dest_indices[i] = one_index // c else: dest_indices[i] = one_index // c % e c *= e dest_indices.reverse() return self.do_convert(expr.expr, dest_indices) return _get_array_element_or_slice(expr, indices)
def _(expr: ArrayTensorProduct): # Recognize expressions like [x, y] with shape (k, 1, k, 1) as `x*y.T`. # The matrix expression has to be equivalent to the tensor product of the # matrices, with trivial dimensions (i.e. dim=1) dropped. # That is, add contractions over trivial dimensions: removed = [] newargs = [] cumul = list(accumulate([0] + [get_rank(arg) for arg in expr.args])) pending = None prev_i = None for i, arg in enumerate(expr.args): current_range = list(range(cumul[i], cumul[i + 1])) if isinstance(arg, OneArray): removed.extend(current_range) continue if not isinstance(arg, (MatrixExpr, MatrixCommon)): rarg, rem = _remove_trivial_dims(arg) removed.extend(rem) newargs.append(rarg) continue elif getattr(arg, "is_Identity", False): if arg.shape == (1, 1): # Ignore identity matrices of shape (1, 1) - they are equivalent to scalar 1. removed.extend(current_range) continue k = arg.shape[0] if pending == k: # OK, there is already removed.extend(current_range) continue elif pending is None: newargs.append(arg) pending = k prev_i = i else: pending = k prev_i = i newargs.append(arg) elif arg.shape == (1, 1): arg, _ = _remove_trivial_dims(arg) # Matrix is equivalent to scalar: if len(newargs) == 0: newargs.append(arg) elif 1 in get_shape(newargs[-1]): if newargs[-1].shape[1] == 1: newargs[-1] = newargs[-1] * arg else: newargs[-1] = arg * newargs[-1] removed.extend(current_range) else: newargs.append(arg) elif 1 in arg.shape: k = [i for i in arg.shape if i != 1][0] if pending is None: pending = k prev_i = i newargs.append(arg) elif pending == k: prev = newargs[-1] if prev.is_Identity: removed.extend([cumul[prev_i], cumul[prev_i] + 1]) newargs[-1] = arg prev_i = i continue if prev.shape[0] == 1: d1 = cumul[prev_i] prev = _a2m_transpose(prev) else: d1 = cumul[prev_i] + 1 if arg.shape[1] == 1: d2 = cumul[i] + 1 arg = _a2m_transpose(arg) else: d2 = cumul[i] newargs[-1] = prev * arg pending = None removed.extend([d1, d2]) else: newargs.append(arg) pending = k prev_i = i else: newargs.append(arg) pending = None return _a2m_tensor_product(*newargs), sorted(removed)
def _(expr: ZeroArray): if get_rank(expr) == 2: return ZeroMatrix(*expr.shape) else: return expr
def remove_identity_matrices(expr: ArrayContraction): editor = _EditArrayContraction(expr) removed: List[int] = [] permutation_map = {} free_indices = list( accumulate([0] + [ sum([i is None for i in arg.indices]) for arg in editor.args_with_ind ])) free_map = {k: v for k, v in zip(editor.args_with_ind, free_indices[:-1])} update_pairs = {} for ind in range(editor.number_of_contraction_indices): args = editor.get_args_with_index(ind) identity_matrices = [ i for i in args if isinstance(i.element, Identity) ] number_identity_matrices = len(identity_matrices) # If the contraction involves a non-identity matrix and multiple identity matrices: if number_identity_matrices != len( args) - 1 or number_identity_matrices == 0: continue # Get the non-identity element: non_identity = [ i for i in args if not isinstance(i.element, Identity) ][0] # Check that all identity matrices have at least one free index # (otherwise they would be contractions to some other elements) if any([None not in i.indices for i in identity_matrices]): continue # Mark the identity matrices for removal: for i in identity_matrices: i.element = None removed.extend( range(free_map[i], free_map[i] + len([j for j in i.indices if j is None]))) last_removed = removed.pop(-1) update_pairs[last_removed, ind] = non_identity.indices[:] # Remove the indices from the non-identity matrix, as the contraction # no longer exists: non_identity.indices = [ None if i == ind else i for i in non_identity.indices ] removed.sort() shifts = list( accumulate([1 if i in removed else 0 for i in range(get_rank(expr))])) for (last_removed, ind), non_identity_indices in update_pairs.items(): pos = [ free_map[non_identity] + i for i, e in enumerate(non_identity_indices) if e == ind ] assert len(pos) == 1 for j in pos: permutation_map[j] = last_removed editor.args_with_ind = [ i for i in editor.args_with_ind if i.element is not None ] ret_expr = editor.to_array_contraction() permutation = [] counter = 0 counter2 = 0 for j in range(get_rank(expr)): if j in removed: continue if counter2 in permutation_map: target = permutation_map[counter2] permutation.append(target - shifts[target]) counter2 += 1 else: while counter in permutation_map.values(): counter += 1 permutation.append(counter) counter += 1 counter2 += 1 ret_expr2 = _permute_dims(ret_expr, _af_invert(permutation)) return ret_expr2, removed
def permutedims(expr, perm=None, index_order_old=None, index_order_new=None): """ Permutes the indices of an array. Parameter specifies the permutation of the indices. Examples ======== >>> from sympy.abc import x, y, z, t >>> from sympy import sin >>> from sympy import Array, permutedims >>> a = Array([[x, y, z], [t, sin(x), 0]]) >>> a [[x, y, z], [t, sin(x), 0]] >>> permutedims(a, (1, 0)) [[x, t], [y, sin(x)], [z, 0]] If the array is of second order, ``transpose`` can be used: >>> from sympy import transpose >>> transpose(a) [[x, t], [y, sin(x)], [z, 0]] Examples on higher dimensions: >>> b = Array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) >>> permutedims(b, (2, 1, 0)) [[[1, 5], [3, 7]], [[2, 6], [4, 8]]] >>> permutedims(b, (1, 2, 0)) [[[1, 5], [2, 6]], [[3, 7], [4, 8]]] An alternative way to specify the same permutations as in the previous lines involves passing the *old* and *new* indices, either as a list or as a string: >>> permutedims(b, index_order_old="cba", index_order_new="abc") [[[1, 5], [3, 7]], [[2, 6], [4, 8]]] >>> permutedims(b, index_order_old="cab", index_order_new="abc") [[[1, 5], [2, 6]], [[3, 7], [4, 8]]] ``Permutation`` objects are also allowed: >>> from sympy.combinatorics import Permutation >>> permutedims(b, Permutation([1, 2, 0])) [[[1, 5], [2, 6]], [[3, 7], [4, 8]]] """ from sympy.tensor.array import SparseNDimArray from sympy.tensor.array.expressions.array_expressions import _ArrayExpr from sympy.tensor.array.expressions.array_expressions import _CodegenArrayAbstract from sympy.tensor.array.expressions.array_expressions import _permute_dims from sympy.matrices.expressions.matexpr import MatrixSymbol from sympy.tensor.array.expressions import PermuteDims from sympy.tensor.array.expressions.array_expressions import get_rank perm = PermuteDims._get_permutation_from_arguments(perm, index_order_old, index_order_new, get_rank(expr)) if isinstance(expr, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)): return _permute_dims(expr, perm) if not isinstance(expr, NDimArray): expr = ImmutableDenseNDimArray(expr) from sympy.combinatorics import Permutation if not isinstance(perm, Permutation): perm = Permutation(list(perm)) if perm.size != expr.rank(): raise ValueError("wrong permutation size") # Get the inverse permutation: iperm = ~perm new_shape = perm(expr.shape) if isinstance(expr, SparseNDimArray): return type(expr)({ tuple(perm(expr._get_tuple_index(k))): v for k, v in expr._sparse_array.items() }, new_shape) indices_span = perm([range(i) for i in expr.shape]) new_array = [None] * len(expr) for i, idx in enumerate(itertools.product(*indices_span)): t = iperm(idx) new_array[i] = expr[t] return type(expr)(new_array, new_shape)
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