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
