def test_arrayexpr_array_expr_zero_array():
za1 = ZeroArray(k, l, m, n)
zm1 = ZeroMatrix(m, n)
za2 = ZeroArray(k, m, m, n)
zm2 = ZeroMatrix(m, m)
zm3 = ZeroMatrix(k, k)
assert _array_tensor_product(M, N, za1) == ZeroArray(k, k, k, k, k, l, m, n)
assert _array_tensor_product(M, N, zm1) == ZeroArray(k, k, k, k, m, n)
assert _array_contraction(za1, (3,)) == ZeroArray(k, l, m)
assert _array_contraction(zm1, (1,)) == ZeroArray(m)
assert _array_contraction(za2, (1, 2)) == ZeroArray(k, n)
assert _array_contraction(zm2, (0, 1)) == 0
assert _array_diagonal(za2, (1, 2)) == ZeroArray(k, n, m)
assert _array_diagonal(zm2, (0, 1)) == ZeroArray(m)
assert _permute_dims(za1, [2, 1, 3, 0]) == ZeroArray(m, l, n, k)
assert _permute_dims(zm1, [1, 0]) == ZeroArray(n, m)
assert _array_add(za1) == za1
assert _array_add(zm1) == ZeroArray(m, n)
tp1 = _array_tensor_product(MatrixSymbol("A", k, l), MatrixSymbol("B", m, n))
assert _array_add(tp1, za1) == tp1
tp2 = _array_tensor_product(MatrixSymbol("C", k, l), MatrixSymbol("D", m, n))
assert _array_add(tp1, za1, tp2) == _array_add(tp1, tp2)
assert _array_add(M, zm3) == M
assert _array_add(M, N, zm3) == _array_add(M, N)
def _(expr: PermuteDims):
if expr.permutation.array_form == [1, 0]:
return _a2m_transpose(_array2matrix(expr.expr))
elif isinstance(expr.expr, ArrayTensorProduct):
ranks = expr.expr.subranks
inv_permutation = expr.permutation**(-1)
newrange = [inv_permutation(i) for i in range(sum(ranks))]
newpos = []
counter = 0
for rank in ranks:
newpos.append(newrange[counter:counter + rank])
counter += rank
newargs = []
newperm = []
scalars = []
for pos, arg in zip(newpos, expr.expr.args):
if len(pos) == 0:
scalars.append(_array2matrix(arg))
elif pos == sorted(pos):
newargs.append((_array2matrix(arg), pos[0]))
newperm.extend(pos)
elif len(pos) == 2:
newargs.append((_a2m_transpose(_array2matrix(arg)), pos[0]))
newperm.extend(reversed(pos))
else:
raise NotImplementedError()
newargs = [i[0] for i in newargs]
return _permute_dims(_a2m_tensor_product(*scalars, *newargs),
_af_invert(newperm))
elif isinstance(expr.expr, ArrayContraction):
mat_mul_lines = _array2matrix(expr.expr)
if not isinstance(mat_mul_lines, ArrayTensorProduct):
flat_cyclic_form = [
j for i in expr.permutation.cyclic_form for j in i
]
expr_shape = get_shape(expr)
if all(expr_shape[i] == 1 for i in flat_cyclic_form):
return mat_mul_lines
return mat_mul_lines
# TODO: this assumes that all arguments are matrices, it may not be the case:
permutation = Permutation(2 * len(mat_mul_lines.args) -
1) * expr.permutation
permuted = [permutation(i) for i in range(2 * len(mat_mul_lines.args))]
args_array = [None for i in mat_mul_lines.args]
for i in range(len(mat_mul_lines.args)):
p1 = permuted[2 * i]
p2 = permuted[2 * i + 1]
if p1 // 2 != p2 // 2:
return _permute_dims(mat_mul_lines, permutation)
pos = p1 // 2
if p1 > p2:
args_array[i] = _a2m_transpose(
mat_mul_lines.args[pos]) # type: ignore
else:
args_array[i] = mat_mul_lines.args[pos] # type: ignore
return _a2m_tensor_product(*args_array)
else:
return expr
def test_arrayexpr_canonicalize_diagonal__permute_dims():
tp = _array_tensor_product(M, Q, N, P)
expr = _array_diagonal(
_permute_dims(tp, [0, 1, 2, 4, 7, 6, 3, 5]), (2, 4, 5), (6, 7),
(0, 3))
result = _array_diagonal(tp, (2, 6, 7), (3, 5), (0, 4))
assert expr == result
tp = _array_tensor_product(M, N, P, Q)
expr = _array_diagonal(_permute_dims(tp, [0, 5, 2, 4, 1, 6, 3, 7]), (1, 2, 6), (3, 4))
result = _array_diagonal(_array_tensor_product(M, P, N, Q), (3, 4, 5), (1, 2))
assert expr == result
def test_arrayexpr_convert_index_to_array_support_function():
expr = M[i, j]
assert _convert_indexed_to_array(expr) == (M, (i, j))
expr = M[i, j] * N[k, l]
assert _convert_indexed_to_array(expr) == (ArrayTensorProduct(M, N),
(i, j, k, l))
expr = M[i, j] * N[j, k]
assert _convert_indexed_to_array(expr) == (ArrayDiagonal(
ArrayTensorProduct(M, N), (1, 2)), (i, k, j))
expr = Sum(M[i, j] * N[j, k], (j, 0, k - 1))
assert _convert_indexed_to_array(expr) == (ArrayContraction(
ArrayTensorProduct(M, N), (1, 2)), (i, k))
expr = M[i, j] + N[i, j]
assert _convert_indexed_to_array(expr) == (ArrayAdd(M, N), (i, j))
expr = M[i, j] + N[j, i]
assert _convert_indexed_to_array(expr) == (ArrayAdd(
M, PermuteDims(N, Permutation([1, 0]))), (i, j))
expr = M[i, j] + M[j, i]
assert _convert_indexed_to_array(expr) == (ArrayAdd(
M, PermuteDims(M, Permutation([1, 0]))), (i, j))
expr = (M * N * P)[i, j]
assert _convert_indexed_to_array(expr) == (_array_contraction(
ArrayTensorProduct(M, N, P), (1, 2), (3, 4)), (i, j))
expr = expr.function # Disregard summation in previous expression
ret1, ret2 = _convert_indexed_to_array(expr)
assert ret1 == ArrayDiagonal(ArrayTensorProduct(M, N, P), (1, 2), (3, 4))
assert str(ret2) == "(i, j, _i_1, _i_2)"
expr = KroneckerDelta(i, j) * M[i, k]
assert _convert_indexed_to_array(expr) == (M, ({i, j}, k))
expr = KroneckerDelta(i, j) * KroneckerDelta(j, k) * M[i, l]
assert _convert_indexed_to_array(expr) == (M, ({i, j, k}, l))
expr = KroneckerDelta(j, k) * (M[i, j] * N[k, l] + N[i, j] * M[k, l])
assert _convert_indexed_to_array(expr) == (_array_diagonal(
_array_add(
ArrayTensorProduct(M, N),
_permute_dims(ArrayTensorProduct(M, N),
Permutation(0, 2)(1, 3))),
(1, 2)), (i, l, frozenset({j, k})))
expr = KroneckerDelta(j, m) * KroneckerDelta(
m, k) * (M[i, j] * N[k, l] + N[i, j] * M[k, l])
assert _convert_indexed_to_array(expr) == (_array_diagonal(
_array_add(
ArrayTensorProduct(M, N),
_permute_dims(ArrayTensorProduct(M, N),
Permutation(0, 2)(1, 3))),
(1, 2)), (i, l, frozenset({j, m, k})))
expr = KroneckerDelta(i, j) * KroneckerDelta(j, k) * KroneckerDelta(
k, m) * M[i, 0] * KroneckerDelta(m, n)
assert _convert_indexed_to_array(expr) == (M, ({i, j, k, m, n}, 0))
expr = M[i, i]
assert _convert_indexed_to_array(expr) == (ArrayDiagonal(M, (0, 1)), (i, ))
def test_arrayexpr_contraction_permutation_mix():
Me = M.subs(k, 3).as_explicit()
Ne = N.subs(k, 3).as_explicit()
cg1 = _array_contraction(PermuteDims(_array_tensor_product(M, N), Permutation([0, 2, 1, 3])), (2, 3))
cg2 = _array_contraction(_array_tensor_product(M, N), (1, 3))
assert cg1 == cg2
cge1 = tensorcontraction(permutedims(tensorproduct(Me, Ne), Permutation([0, 2, 1, 3])), (2, 3))
cge2 = tensorcontraction(tensorproduct(Me, Ne), (1, 3))
assert cge1 == cge2
cg1 = _permute_dims(_array_tensor_product(M, N), Permutation([0, 1, 3, 2]))
cg2 = _array_tensor_product(M, _permute_dims(N, Permutation([1, 0])))
assert cg1 == cg2
cg1 = _array_contraction(
_permute_dims(
_array_tensor_product(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])),
(1, 2), (3, 5)
)
cg2 = _array_contraction(
_array_tensor_product(M, N, P, _permute_dims(Q, Permutation([1, 0]))),
(1, 5), (2, 3)
)
assert cg1 == cg2
cg1 = _array_contraction(
_permute_dims(
_array_tensor_product(M, N, P, Q), Permutation([1, 0, 4, 6, 2, 7, 5, 3])),
(0, 1), (2, 6), (3, 7)
)
cg2 = _permute_dims(
_array_contraction(
_array_tensor_product(M, P, Q, N),
(0, 1), (2, 3), (4, 7)),
[1, 0]
)
assert cg1 == cg2
cg1 = _array_contraction(
_permute_dims(
_array_tensor_product(M, N, P, Q), Permutation([1, 0, 4, 6, 7, 2, 5, 3])),
(0, 1), (2, 6), (3, 7)
)
cg2 = _permute_dims(
_array_contraction(
_array_tensor_product(_permute_dims(M, [1, 0]), N, P, Q),
(0, 1), (3, 6), (4, 5)
),
Permutation([1, 0])
)
assert cg1 == cg2
def _(expr: Inverse, x: Expr):
mat = expr.I
dexpr = array_derive(mat, x)
tp = _array_tensor_product(-expr, dexpr, expr)
mp = _array_contraction(tp, (1, 4), (5, 6))
pp = _permute_dims(mp, [1, 2, 0, 3])
return pp
def _(expr: ArraySymbol, x: _ArrayExpr):
if expr == x:
return _permute_dims(
ArrayTensorProduct.fromiter(Identity(i) for i in expr.shape),
[2 * i for i in range(len(expr.shape))] +
[2 * i + 1 for i in range(len(expr.shape))])
return ZeroArray(*(x.shape + expr.shape))
def _(expr: ArrayTensorProduct, x: Expr):
args = expr.args
addend_list = []
for i, arg in enumerate(expr.args):
darg = array_derive(arg, x)
if darg == 0:
continue
args_prev = args[:i]
args_succ = args[i + 1:]
shape_prev = reduce(operator.add, map(get_shape, args_prev), ())
shape_succ = reduce(operator.add, map(get_shape, args_succ), ())
addend = _array_tensor_product(*args_prev, darg, *args_succ)
tot1 = len(get_shape(x))
tot2 = tot1 + len(shape_prev)
tot3 = tot2 + len(get_shape(arg))
tot4 = tot3 + len(shape_succ)
perm = [i for i in range(tot1, tot2)] + \
[i for i in range(tot1)] + [i for i in range(tot2, tot3)] + \
[i for i in range(tot3, tot4)]
addend = _permute_dims(addend, _af_invert(perm))
addend_list.append(addend)
if len(addend_list) == 1:
return addend_list[0]
elif len(addend_list) == 0:
return S.Zero
else:
return _array_add(*addend_list)
def test_array_symbol_and_element():
A = ArraySymbol("A", (2, ))
A0 = ArrayElement(A, (0, ))
A1 = ArrayElement(A, (1, ))
assert A[0] == A0
assert A[1] != A0
assert A.as_explicit() == ImmutableDenseNDimArray([A0, A1])
A2 = tensorproduct(A, A)
assert A2.shape == (2, 2)
# TODO: not yet supported:
# assert A2.as_explicit() == Array([[A[0]*A[0], A[1]*A[0]], [A[0]*A[1], A[1]*A[1]]])
A3 = tensorcontraction(A2, (0, 1))
assert A3.shape == ()
# TODO: not yet supported:
# assert A3.as_explicit() == Array([])
A = ArraySymbol("A", (2, 3, 4))
Ae = A.as_explicit()
assert Ae == ImmutableDenseNDimArray(
[[[ArrayElement(A, (i, j, k)) for k in range(4)] for j in range(3)]
for i in range(2)])
p = _permute_dims(A, Permutation(0, 2, 1))
assert isinstance(p, PermuteDims)
def _(expr: MatrixSymbol, x: Expr):
m, n = expr.shape
if expr == x:
return _permute_dims(
_array_tensor_product(Identity(m), Identity(n)),
[0, 2, 1, 3]
)
return ZeroArray(*(x.shape + expr.shape)) # type: ignore
def test_arrayexpr_array_diagonal():
cg = _array_diagonal(M, (1, 0))
assert cg == _array_diagonal(M, (0, 1))
cg = _array_diagonal(_array_tensor_product(M, N, P), (4, 1), (2, 0))
assert cg == _array_diagonal(_array_tensor_product(M, N, P), (1, 4), (0, 2))
cg = _array_diagonal(_array_tensor_product(M, N), (1, 2), (3,), allow_trivial_diags=True)
assert cg == _permute_dims(_array_diagonal(_array_tensor_product(M, N), (1, 2)), [0, 2, 1])
Ax = ArraySymbol("Ax", shape=(1, 2, 3, 4, 3, 5, 6, 2, 7))
cg = _array_diagonal(Ax, (1, 7), (3,), (2, 4), (6,), allow_trivial_diags=True)
assert cg == _permute_dims(_array_diagonal(Ax, (1, 7), (2, 4)), [0, 2, 4, 5, 1, 6, 3])
cg = _array_diagonal(M, (0,), allow_trivial_diags=True)
assert cg == _permute_dims(M, [1, 0])
raises(ValueError, lambda: _array_diagonal(M, (0, 0)))
def test_arrayexpr_nested_permutations():
cg = _permute_dims(_permute_dims(M, (1, 0)), (1, 0))
assert cg == M
times = 3
plist1 = [list(range(6)) for i in range(times)]
plist2 = [list(range(6)) for i in range(times)]
for i in range(times):
random.shuffle(plist1[i])
random.shuffle(plist2[i])
plist1.append([2, 5, 4, 1, 0, 3])
plist2.append([3, 5, 0, 4, 1, 2])
plist1.append([2, 5, 4, 0, 3, 1])
plist2.append([3, 0, 5, 1, 2, 4])
plist1.append([5, 4, 2, 0, 3, 1])
plist2.append([4, 5, 0, 2, 3, 1])
Me = M.subs(k, 3).as_explicit()
Ne = N.subs(k, 3).as_explicit()
Pe = P.subs(k, 3).as_explicit()
cge = tensorproduct(Me, Ne, Pe)
for permutation_array1, permutation_array2 in zip(plist1, plist2):
p1 = Permutation(permutation_array1)
p2 = Permutation(permutation_array2)
cg = _permute_dims(
_permute_dims(
_array_tensor_product(M, N, P),
p1),
p2
)
result = _permute_dims(
_array_tensor_product(M, N, P),
p2*p1
)
assert cg == result
# Check that `permutedims` behaves the same way with explicit-component arrays:
result1 = _permute_dims(_permute_dims(cge, p1), p2)
result2 = _permute_dims(cge, p2*p1)
assert result1 == result2
def test_arrayexpr_derivatives1():
res = array_derive(X, X)
assert res == PermuteDims(ArrayTensorProduct(I, I), [0, 2, 1, 3])
cg = ArrayTensorProduct(A, X, B)
res = array_derive(cg, X)
assert res == _permute_dims(ArrayTensorProduct(I, A, I, B),
[0, 4, 2, 3, 1, 5, 6, 7])
cg = ArrayContraction(X, (0, 1))
res = array_derive(cg, X)
assert res == ArrayContraction(ArrayTensorProduct(I, I), (1, 3))
cg = ArrayDiagonal(X, (0, 1))
res = array_derive(cg, X)
assert res == ArrayDiagonal(ArrayTensorProduct(I, I), (1, 3))
cg = ElementwiseApplyFunction(sin, X)
res = array_derive(cg, X)
assert res.dummy_eq(
ArrayDiagonal(
ArrayTensorProduct(ElementwiseApplyFunction(cos, X), I, I), (0, 3),
(1, 5)))
cg = ArrayElementwiseApplyFunc(sin, X)
res = array_derive(cg, X)
assert res.dummy_eq(
ArrayDiagonal(
ArrayTensorProduct(I, I, ArrayElementwiseApplyFunc(cos, X)),
(1, 4), (3, 5)))
res = array_derive(A1, A1)
assert res == PermuteDims(
ArrayTensorProduct(Identity(3), Identity(2), Identity(k)),
[0, 2, 4, 1, 3, 5])
cg = ArrayElementwiseApplyFunc(sin, A1)
res = array_derive(cg, A1)
assert res.dummy_eq(
ArrayDiagonal(
ArrayTensorProduct(Identity(3), Identity(2), Identity(k),
ArrayElementwiseApplyFunc(cos, A1)), (1, 6),
(3, 7), (5, 8)))
cg = Reshape(A, (k**2, ))
res = array_derive(cg, A)
assert res == Reshape(PermuteDims(ArrayTensorProduct(I, I), [0, 2, 1, 3]),
(k, k, k**2))
def _(expr: PermuteDims):
subexpr, subremoved = _remove_trivial_dims(expr.expr)
p = expr.permutation.array_form
pinv = _af_invert(expr.permutation.array_form)
shift = list(
accumulate([1 if i in subremoved else 0 for i in range(len(p))]))
premoved = [pinv[i] for i in subremoved]
p2 = [e - shift[e] for i, e in enumerate(p) if e not in subremoved]
# TODO: check if subremoved should be permuted as well...
newexpr = _permute_dims(subexpr, p2)
premoved = sorted(premoved)
if newexpr != expr:
newexpr, removed2 = _remove_trivial_dims(_array2matrix(newexpr))
premoved = _combine_removed(-1, premoved, removed2)
return newexpr, premoved
def test_array_symbol_and_element():
A = ArraySymbol("A", (2,))
A0 = ArrayElement(A, (0,))
A1 = ArrayElement(A, (1,))
assert A[0] == A0
assert A[1] != A0
assert A.as_explicit() == ImmutableDenseNDimArray([A0, A1])
A2 = tensorproduct(A, A)
assert A2.shape == (2, 2)
# TODO: not yet supported:
# assert A2.as_explicit() == Array([[A[0]*A[0], A[1]*A[0]], [A[0]*A[1], A[1]*A[1]]])
A3 = tensorcontraction(A2, (0, 1))
assert A3.shape == ()
# TODO: not yet supported:
# assert A3.as_explicit() == Array([])
A = ArraySymbol("A", (2, 3, 4))
Ae = A.as_explicit()
assert Ae == ImmutableDenseNDimArray(
[[[ArrayElement(A, (i, j, k)) for k in range(4)] for j in range(3)] for i in range(2)])
p = _permute_dims(A, Permutation(0, 2, 1))
assert isinstance(p, PermuteDims)
A = ArraySymbol("A", (2,))
raises(IndexError, lambda: A[()])
raises(IndexError, lambda: A[0, 1])
raises(ValueError, lambda: A[-1])
raises(ValueError, lambda: A[2])
O = OneArray(3, 4)
Z = ZeroArray(m, n)
raises(IndexError, lambda: O[()])
raises(IndexError, lambda: O[1, 2, 3])
raises(ValueError, lambda: O[3, 0])
raises(ValueError, lambda: O[0, 4])
assert O[1, 2] == 1
assert Z[1, 2] == 0
def test_arrayexpr_array_shape():
expr = _array_tensor_product(M, N, P, Q)
assert expr.shape == (k, k, k, k, k, k, k, k)
Z = MatrixSymbol("Z", m, n)
expr = _array_tensor_product(M, Z)
assert expr.shape == (k, k, m, n)
expr2 = _array_contraction(expr, (0, 1))
assert expr2.shape == (m, n)
expr2 = _array_diagonal(expr, (0, 1))
assert expr2.shape == (m, n, k)
exprp = _permute_dims(expr, [2, 1, 3, 0])
assert exprp.shape == (m, k, n, k)
expr3 = _array_tensor_product(N, Z)
expr2 = _array_add(expr, expr3)
assert expr2.shape == (k, k, m, n)
# Contraction along axes with discordant dimensions:
raises(ValueError, lambda: _array_contraction(expr, (1, 2)))
# Also diagonal needs the same dimensions:
raises(ValueError, lambda: _array_diagonal(expr, (1, 2)))
# Diagonal requires at least to axes to compute the diagonal:
raises(ValueError, lambda: _array_diagonal(expr, (1,)))
def test_arrayexpr_permute_tensor_product():
cg1 = _permute_dims(_array_tensor_product(M, N, P, Q), Permutation([2, 3, 1, 0, 5, 4, 6, 7]))
cg2 = _array_tensor_product(N, _permute_dims(M, [1, 0]),
_permute_dims(P, [1, 0]), Q)
assert cg1 == cg2
# TODO: reverse operation starting with `PermuteDims` and getting down to `bb`...
cg1 = _permute_dims(_array_tensor_product(M, N, P, Q), Permutation([2, 3, 4, 5, 0, 1, 6, 7]))
cg2 = _array_tensor_product(N, P, M, Q)
assert cg1 == cg2
cg1 = _permute_dims(_array_tensor_product(M, N, P, Q), Permutation([2, 3, 4, 6, 5, 7, 0, 1]))
assert cg1.expr == _array_tensor_product(N, P, Q, M)
assert cg1.permutation == Permutation([0, 1, 2, 4, 3, 5, 6, 7])
cg1 = _array_contraction(
_permute_dims(
_array_tensor_product(N, Q, Q, M),
[2, 1, 5, 4, 0, 3, 6, 7]),
[1, 2, 6])
cg2 = _permute_dims(_array_contraction(_array_tensor_product(Q, Q, N, M), (3, 5, 6)), [0, 2, 3, 1, 4])
assert cg1 == cg2
cg1 = _array_contraction(
_array_contraction(
_array_contraction(
_array_contraction(
_permute_dims(
_array_tensor_product(N, Q, Q, M),
[2, 1, 5, 4, 0, 3, 6, 7]),
[1, 2, 6]),
[1, 3, 4]),
[1]),
[0])
cg2 = _array_contraction(_array_tensor_product(M, N, Q, Q), (0, 3, 5), (1, 4, 7), (2,), (6,))
assert cg1 == cg2
def convert_matrix_to_array(expr: Basic) -> Basic:
if isinstance(expr, MatMul):
args_nonmat = []
args = []
for arg in expr.args:
if isinstance(arg, MatrixExpr):
args.append(arg)
else:
args_nonmat.append(convert_matrix_to_array(arg))
contractions = [(2*i+1, 2*i+2) for i in range(len(args)-1)]
scalar = _array_tensor_product(*args_nonmat) if args_nonmat else S.One
if scalar == 1:
tprod = _array_tensor_product(
*[convert_matrix_to_array(arg) for arg in args])
else:
tprod = _array_tensor_product(
scalar,
*[convert_matrix_to_array(arg) for arg in args])
return _array_contraction(
tprod,
*contractions
)
elif isinstance(expr, MatAdd):
return _array_add(
*[convert_matrix_to_array(arg) for arg in expr.args]
)
elif isinstance(expr, Transpose):
return _permute_dims(
convert_matrix_to_array(expr.args[0]), [1, 0]
)
elif isinstance(expr, Trace):
inner_expr: MatrixExpr = convert_matrix_to_array(expr.arg) # type: ignore
return _array_contraction(inner_expr, (0, len(inner_expr.shape) - 1))
elif isinstance(expr, Mul):
return _array_tensor_product(*[convert_matrix_to_array(i) for i in expr.args])
elif isinstance(expr, Pow):
base = convert_matrix_to_array(expr.base)
if (expr.exp > 0) == True:
return _array_tensor_product(*[base for i in range(expr.exp)])
else:
return expr
elif isinstance(expr, MatPow):
base = convert_matrix_to_array(expr.base)
if expr.exp.is_Integer != True:
b = symbols("b", cls=Dummy)
return ArrayElementwiseApplyFunc(Lambda(b, b**expr.exp), convert_matrix_to_array(base))
elif (expr.exp > 0) == True:
return convert_matrix_to_array(MatMul.fromiter(base for i in range(expr.exp)))
else:
return expr
elif isinstance(expr, HadamardProduct):
tp = _array_tensor_product(*[convert_matrix_to_array(arg) for arg in expr.args])
diag = [[2*i for i in range(len(expr.args))], [2*i+1 for i in range(len(expr.args))]]
return _array_diagonal(tp, *diag)
elif isinstance(expr, HadamardPower):
base, exp = expr.args
if isinstance(exp, Integer) and exp > 0:
return convert_matrix_to_array(HadamardProduct.fromiter(base for i in range(exp)))
else:
d = Dummy("d")
return ArrayElementwiseApplyFunc(Lambda(d, d**exp), base)
elif isinstance(expr, KroneckerProduct):
kp_args = [convert_matrix_to_array(arg) for arg in expr.args]
permutation = [2*i for i in range(len(kp_args))] + [2*i + 1 for i in range(len(kp_args))]
return Reshape(_permute_dims(_array_tensor_product(*kp_args), permutation), expr.shape)
else:
return expr
def test_arrayexpr_array_wrong_permutation_size():
cg = _array_tensor_product(M, N)
raises(ValueError, lambda: _permute_dims(cg, [1, 0]))
raises(ValueError, lambda: _permute_dims(cg, [1, 0, 2, 3, 5, 4]))
def _a2m_transpose(arg):
if isinstance(arg, _CodegenArrayAbstract):
return _permute_dims(arg, [1, 0])
else:
from sympy.matrices.expressions.transpose import Transpose
return Transpose(arg).doit()
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 _(expr: PermuteDims, x: Expr):
de = array_derive(expr.expr, x)
perm = [0, 1] + [i + 2 for i in expr.permutation.array_form]
return _permute_dims(de, perm)
def test_arrayexpr_permutedims_sink():
cg = _permute_dims(_array_tensor_product(M, N), [0, 1, 3, 2], nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == _array_tensor_product(M, _permute_dims(N, [1, 0]))
cg = _permute_dims(_array_tensor_product(M, N), [1, 0, 3, 2], nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == _array_tensor_product(_permute_dims(M, [1, 0]), _permute_dims(N, [1, 0]))
cg = _permute_dims(_array_tensor_product(M, N), [3, 2, 1, 0], nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == _array_tensor_product(_permute_dims(N, [1, 0]), _permute_dims(M, [1, 0]))
cg = _permute_dims(_array_contraction(_array_tensor_product(M, N), (1, 2)), [1, 0], nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == _array_contraction(_permute_dims(_array_tensor_product(M, N), [[0, 3]]), (1, 2))
cg = _permute_dims(_array_tensor_product(M, N), [1, 0, 3, 2], nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == _array_tensor_product(_permute_dims(M, [1, 0]), _permute_dims(N, [1, 0]))
cg = _permute_dims(_array_contraction(_array_tensor_product(M, N, P), (1, 2), (3, 4)), [1, 0], nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == _array_contraction(_permute_dims(_array_tensor_product(M, N, P), [[0, 5]]), (1, 2), (3, 4))
def _(expr: Transpose, x: Expr):
# D(A.T, A) ==> (m,n,i,j) ==> D(A_ji, A_mn) = d_mj d_ni
# D(B.T, A) ==> (m,n,i,j) ==> D(B_ji, A_mn)
fd = array_derive(expr.arg, x)
return _permute_dims(fd, [0, 1, 3, 2])
def test_arrayexpr_array_flatten():
# Flatten nested ArrayTensorProduct objects:
expr1 = _array_tensor_product(M, N)
expr2 = _array_tensor_product(P, Q)
expr = _array_tensor_product(expr1, expr2)
assert expr == _array_tensor_product(M, N, P, Q)
assert expr.args == (M, N, P, Q)
# Flatten mixed ArrayTensorProduct and ArrayContraction objects:
cg1 = _array_contraction(expr1, (1, 2))
cg2 = _array_contraction(expr2, (0, 3))
expr = _array_tensor_product(cg1, cg2)
assert expr == _array_contraction(_array_tensor_product(M, N, P, Q), (1, 2), (4, 7))
expr = _array_tensor_product(M, cg1)
assert expr == _array_contraction(_array_tensor_product(M, M, N), (3, 4))
# Flatten nested ArrayContraction objects:
cgnested = _array_contraction(cg1, (0, 1))
assert cgnested == _array_contraction(_array_tensor_product(M, N), (0, 3), (1, 2))
cgnested = _array_contraction(_array_tensor_product(cg1, cg2), (0, 3))
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 6), (1, 2), (4, 7))
cg3 = _array_contraction(_array_tensor_product(M, N, P, Q), (1, 3), (2, 4))
cgnested = _array_contraction(cg3, (0, 1))
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 5), (1, 3), (2, 4))
cgnested = _array_contraction(cg3, (0, 3), (1, 2))
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 7), (1, 3), (2, 4), (5, 6))
cg4 = _array_contraction(_array_tensor_product(M, N, P, Q), (1, 5), (3, 7))
cgnested = _array_contraction(cg4, (0, 1))
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 2), (1, 5), (3, 7))
cgnested = _array_contraction(cg4, (0, 1), (2, 3))
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 2), (1, 5), (3, 7), (4, 6))
cg = _array_diagonal(cg4)
assert cg == cg4
assert isinstance(cg, type(cg4))
# Flatten nested ArrayDiagonal objects:
cg1 = _array_diagonal(expr1, (1, 2))
cg2 = _array_diagonal(expr2, (0, 3))
cg3 = _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 3), (2, 4))
cg4 = _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 5), (3, 7))
cgnested = _array_diagonal(cg1, (0, 1))
assert cgnested == _array_diagonal(_array_tensor_product(M, N), (1, 2), (0, 3))
cgnested = _array_diagonal(cg3, (1, 2))
assert cgnested == _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 3), (2, 4), (5, 6))
cgnested = _array_diagonal(cg4, (1, 2))
assert cgnested == _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 5), (3, 7), (2, 4))
cg = _array_add(M, N)
cg2 = _array_add(cg, P)
assert isinstance(cg2, ArrayAdd)
assert cg2.args == (M, N, P)
assert cg2.shape == (k, k)
expr = _array_tensor_product(_array_diagonal(X, (0, 1)), _array_diagonal(A, (0, 1)))
assert expr == _array_diagonal(_array_tensor_product(X, A), (0, 1), (2, 3))
expr1 = _array_diagonal(_array_tensor_product(X, A), (1, 2))
expr2 = _array_tensor_product(expr1, a)
assert expr2 == _permute_dims(_array_diagonal(_array_tensor_product(X, A, a), (1, 2)), [0, 1, 4, 2, 3])
expr1 = _array_contraction(_array_tensor_product(X, A), (1, 2))
expr2 = _array_tensor_product(expr1, a)
assert isinstance(expr2, ArrayContraction)
assert isinstance(expr2.expr, ArrayTensorProduct)
cg = _array_tensor_product(_array_diagonal(_array_tensor_product(A, X, Y), (0, 3), (1, 5)), a, b)
assert cg == _permute_dims(_array_diagonal(_array_tensor_product(A, X, Y, a, b), (0, 3), (1, 5)), [0, 1, 6, 7, 2, 3, 4, 5])
def test_arrayexpr_convert_array_to_diagonalized_vector():
# Check matrix recognition over trivial dimensions:
cg = _array_tensor_product(a, b)
assert convert_array_to_matrix(cg) == a * b.T
cg = _array_tensor_product(I1, a, b)
assert convert_array_to_matrix(cg) == a * b.T
# Recognize trace inside a tensor product:
cg = _array_contraction(_array_tensor_product(A, B, C), (0, 3), (1, 2))
assert convert_array_to_matrix(cg) == Trace(A * B) * C
# Transform diagonal operator to contraction:
cg = _array_diagonal(_array_tensor_product(A, a), (1, 2))
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(A, OneArray(1), DiagMatrix(a)), (1, 3))
assert convert_array_to_matrix(cg) == A * DiagMatrix(a)
cg = _array_diagonal(_array_tensor_product(a, b), (0, 2))
assert _array_diag2contr_diagmatrix(cg) == _permute_dims(
_array_contraction(_array_tensor_product(DiagMatrix(a), OneArray(1), b), (0, 3)), [1, 2, 0]
)
assert convert_array_to_matrix(cg) == b.T * DiagMatrix(a)
cg = _array_diagonal(_array_tensor_product(A, a), (0, 2))
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(A, OneArray(1), DiagMatrix(a)), (0, 3))
assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a)
cg = _array_diagonal(_array_tensor_product(I, x, I1), (0, 2), (3, 5))
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(I, OneArray(1), I1, DiagMatrix(x)), (0, 5))
assert convert_array_to_matrix(cg) == DiagMatrix(x)
cg = _array_diagonal(_array_tensor_product(I, x, A, B), (1, 2), (5, 6))
assert _array_diag2contr_diagmatrix(cg) == _array_diagonal(_array_contraction(_array_tensor_product(I, OneArray(1), A, B, DiagMatrix(x)), (1, 7)), (5, 6))
# TODO: this is returning a wrong result:
# convert_array_to_matrix(cg)
cg = _array_diagonal(_array_tensor_product(I1, a, b), (1, 3, 5))
assert convert_array_to_matrix(cg) == a*b.T
cg = _array_diagonal(_array_tensor_product(I1, a, b), (1, 3))
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(OneArray(1), a, b, I1), (2, 6))
assert convert_array_to_matrix(cg) == a*b.T
cg = _array_diagonal(_array_tensor_product(x, I1), (1, 2))
assert isinstance(cg, ArrayDiagonal)
assert cg.diagonal_indices == ((1, 2),)
assert convert_array_to_matrix(cg) == x
cg = _array_diagonal(_array_tensor_product(x, I), (0, 2))
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(OneArray(1), I, DiagMatrix(x)), (1, 3))
assert convert_array_to_matrix(cg).doit() == DiagMatrix(x)
raises(ValueError, lambda: _array_diagonal(x, (1,)))
# Ignore identity matrices with contractions:
cg = _array_contraction(_array_tensor_product(I, A, I, I), (0, 2), (1, 3), (5, 7))
assert cg.split_multiple_contractions() == cg
assert convert_array_to_matrix(cg) == Trace(A) * I
cg = _array_contraction(_array_tensor_product(Trace(A) * I, I, I), (1, 5), (3, 4))
assert cg.split_multiple_contractions() == cg
assert convert_array_to_matrix(cg).doit() == Trace(A) * I
# Add DiagMatrix when required:
cg = _array_contraction(_array_tensor_product(A, a), (1, 2))
assert cg.split_multiple_contractions() == cg
assert convert_array_to_matrix(cg) == A * a
cg = _array_contraction(_array_tensor_product(A, a, B), (1, 2, 4))
assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1), B), (1, 2), (3, 5))
assert convert_array_to_matrix(cg) == A * DiagMatrix(a) * B
cg = _array_contraction(_array_tensor_product(A, a, B), (0, 2, 4))
assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1), B), (0, 2), (3, 5))
assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) * B
cg = _array_contraction(_array_tensor_product(A, a, b, a.T, B), (0, 2, 4, 7, 9))
assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1),
DiagMatrix(b), OneArray(1), DiagMatrix(a), OneArray(1), B),
(0, 2), (3, 5), (6, 9), (8, 12))
assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) * DiagMatrix(b) * DiagMatrix(a) * B.T
cg = _array_contraction(_array_tensor_product(I1, I1, I1), (1, 2, 4))
assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(I1, I1, OneArray(1), I1), (1, 2), (3, 5))
assert convert_array_to_matrix(cg) == 1
cg = _array_contraction(_array_tensor_product(I, I, I, I, A), (1, 2, 8), (5, 6, 9))
assert convert_array_to_matrix(cg.split_multiple_contractions()).doit() == A
cg = _array_contraction(_array_tensor_product(A, a, C, a, B), (1, 2, 4), (5, 6, 8))
expected = _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1), C, DiagMatrix(a), OneArray(1), B), (1, 3), (2, 5), (6, 7), (8, 10))
assert cg.split_multiple_contractions() == expected
assert convert_array_to_matrix(cg) == A * DiagMatrix(a) * C * DiagMatrix(a) * B
cg = _array_contraction(_array_tensor_product(a, I1, b, I1, (a.T*b).applyfunc(cos)), (1, 2, 8), (5, 6, 9))
expected = _array_contraction(_array_tensor_product(a, I1, OneArray(1), b, I1, OneArray(1), (a.T*b).applyfunc(cos)),
(1, 3), (2, 10), (6, 8), (7, 11))
assert cg.split_multiple_contractions().dummy_eq(expected)
assert convert_array_to_matrix(cg).doit().dummy_eq(MatMul(a, (a.T * b).applyfunc(cos), b.T))
def convert_indexed_to_array(expr, first_indices=None):
r"""
Parse indexed expression into a form useful for code generation.
Examples
========
>>> from sympy.tensor.array.expressions.conv_indexed_to_array import convert_indexed_to_array
>>> from sympy import MatrixSymbol, Sum, symbols
>>> i, j, k, d = symbols("i j k d")
>>> M = MatrixSymbol("M", d, d)
>>> N = MatrixSymbol("N", d, d)
Recognize the trace in summation form:
>>> expr = Sum(M[i, i], (i, 0, d-1))
>>> convert_indexed_to_array(expr)
ArrayContraction(M, (0, 1))
Recognize the extraction of the diagonal by using the same index `i` on
both axes of the matrix:
>>> expr = M[i, i]
>>> convert_indexed_to_array(expr)
ArrayDiagonal(M, (0, 1))
This function can help perform the transformation expressed in two
different mathematical notations as:
`\sum_{j=0}^{N-1} A_{i,j} B_{j,k} \Longrightarrow \mathbf{A}\cdot \mathbf{B}`
Recognize the matrix multiplication in summation form:
>>> expr = Sum(M[i, j]*N[j, k], (j, 0, d-1))
>>> convert_indexed_to_array(expr)
ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
Specify that ``k`` has to be the starting index:
>>> convert_indexed_to_array(expr, first_indices=[k])
ArrayContraction(ArrayTensorProduct(N, M), (0, 3))
"""
result, indices = _convert_indexed_to_array(expr)
if any(isinstance(i, (int, Integer)) for i in indices):
result = ArrayElement(result, indices)
indices = []
if not first_indices:
return result
def _check_is_in(elem, indices):
if elem in indices:
return True
if any(elem in i for i in indices if isinstance(i, frozenset)):
return True
return False
repl = {j: i for i in indices if isinstance(i, frozenset) for j in i}
first_indices = [repl.get(i, i) for i in first_indices]
for i in first_indices:
if not _check_is_in(i, indices):
first_indices.remove(i)
first_indices.extend(
[i for i in indices if not _check_is_in(i, first_indices)])
def _get_pos(elem, indices):
if elem in indices:
return indices.index(elem)
for i, e in enumerate(indices):
if not isinstance(e, frozenset):
continue
if elem in e:
return i
raise ValueError("not found")
permutation = _af_invert([_get_pos(i, first_indices) for i in indices])
if isinstance(result, ArrayAdd):
return _array_add(
*[_permute_dims(arg, permutation) for arg in result.args])
else:
return _permute_dims(result, permutation)
def _convert_indexed_to_array(expr):
if isinstance(expr, Sum):
function = expr.function
summation_indices = expr.variables
subexpr, subindices = _convert_indexed_to_array(function)
subindicessets = {
j: i
for i in subindices if isinstance(i, frozenset) for j in i
}
summation_indices = sorted(set(
[subindicessets.get(i, i) for i in summation_indices]),
key=default_sort_key)
# TODO: check that Kronecker delta is only contracted to one other element:
kronecker_indices = set([])
if isinstance(function, Mul):
for arg in function.args:
if not isinstance(arg, KroneckerDelta):
continue
arg_indices = sorted(set(arg.indices), key=default_sort_key)
if len(arg_indices) == 2:
kronecker_indices.update(arg_indices)
kronecker_indices = sorted(kronecker_indices, key=default_sort_key)
# Check dimensional consistency:
shape = get_shape(subexpr)
if shape:
for ind, istart, iend in expr.limits:
i = _get_argindex(subindices, ind)
if istart != 0 or iend + 1 != shape[i]:
raise ValueError(
"summation index and array dimension mismatch: %s" %
ind)
contraction_indices = []
subindices = list(subindices)
if isinstance(subexpr, ArrayDiagonal):
diagonal_indices = list(subexpr.diagonal_indices)
dindices = subindices[-len(diagonal_indices):]
subindices = subindices[:-len(diagonal_indices)]
for index in summation_indices:
if index in dindices:
position = dindices.index(index)
contraction_indices.append(diagonal_indices[position])
diagonal_indices[position] = None
diagonal_indices = [i for i in diagonal_indices if i is not None]
for i, ind in enumerate(subindices):
if ind in summation_indices:
pass
if diagonal_indices:
subexpr = _array_diagonal(subexpr.expr, *diagonal_indices)
else:
subexpr = subexpr.expr
axes_contraction = defaultdict(list)
for i, ind in enumerate(subindices):
include = all(j not in kronecker_indices
for j in ind) if isinstance(
ind, frozenset) else ind not in kronecker_indices
if ind in summation_indices and include:
axes_contraction[ind].append(i)
subindices[i] = None
for k, v in axes_contraction.items():
if any(i in kronecker_indices for i in k) if isinstance(
k, frozenset) else k in kronecker_indices:
continue
contraction_indices.append(tuple(v))
free_indices = [i for i in subindices if i is not None]
indices_ret = list(free_indices)
indices_ret.sort(key=lambda x: free_indices.index(x))
return _array_contraction(
subexpr, *contraction_indices,
free_indices=free_indices), tuple(indices_ret)
if isinstance(expr, Mul):
args, indices = zip(
*[_convert_indexed_to_array(arg) for arg in expr.args])
# Check if there are KroneckerDelta objects:
kronecker_delta_repl = {}
for arg in args:
if not isinstance(arg, KroneckerDelta):
continue
# Diagonalize two indices:
i, j = arg.indices
kindices = set(arg.indices)
if i in kronecker_delta_repl:
kindices.update(kronecker_delta_repl[i])
if j in kronecker_delta_repl:
kindices.update(kronecker_delta_repl[j])
kindices = frozenset(kindices)
for index in kindices:
kronecker_delta_repl[index] = kindices
# Remove KroneckerDelta objects, their relations should be handled by
# ArrayDiagonal:
newargs = []
newindices = []
for arg, loc_indices in zip(args, indices):
if isinstance(arg, KroneckerDelta):
continue
newargs.append(arg)
newindices.append(loc_indices)
flattened_indices = [
kronecker_delta_repl.get(j, j) for i in newindices for j in i
]
diagonal_indices, ret_indices = _get_diagonal_indices(
flattened_indices)
tp = _array_tensor_product(*newargs)
if diagonal_indices:
return _array_diagonal(tp, *diagonal_indices), ret_indices
else:
return tp, ret_indices
if isinstance(expr, MatrixElement):
indices = expr.args[1:]
diagonal_indices, ret_indices = _get_diagonal_indices(indices)
if diagonal_indices:
return _array_diagonal(expr.args[0],
*diagonal_indices), ret_indices
else:
return expr.args[0], ret_indices
if isinstance(expr, ArrayElement):
indices = expr.indices
diagonal_indices, ret_indices = _get_diagonal_indices(indices)
if diagonal_indices:
return _array_diagonal(expr.name, *diagonal_indices), ret_indices
else:
return expr.name, ret_indices
if isinstance(expr, Indexed):
indices = expr.indices
diagonal_indices, ret_indices = _get_diagonal_indices(indices)
if diagonal_indices:
return _array_diagonal(expr.base, *diagonal_indices), ret_indices
else:
return expr.args[0], ret_indices
if isinstance(expr, IndexedBase):
raise NotImplementedError
if isinstance(expr, KroneckerDelta):
return expr, expr.indices
if isinstance(expr, Add):
args, indices = zip(
*[_convert_indexed_to_array(arg) for arg in expr.args])
args = list(args)
# Check if all indices are compatible. Otherwise expand the dimensions:
index0 = []
shape0 = []
for arg, arg_indices in zip(args, indices):
arg_indices_set = set(arg_indices)
arg_indices_missing = arg_indices_set.difference(index0)
index0.extend([i for i in arg_indices if i in arg_indices_missing])
arg_shape = get_shape(arg)
shape0.extend([
arg_shape[i] for i, e in enumerate(arg_indices)
if e in arg_indices_missing
])
for i, (arg, arg_indices) in enumerate(zip(args, indices)):
if len(arg_indices) < len(index0):
missing_indices_pos = [
i for i, e in enumerate(index0) if e not in arg_indices
]
missing_shape = [shape0[i] for i in missing_indices_pos]
arg_indices = tuple(index0[j]
for j in missing_indices_pos) + arg_indices
args[i] = _array_tensor_product(OneArray(*missing_shape),
args[i])
permutation = Permutation([arg_indices.index(j) for j in index0])
# Perform index permutations:
args[i] = _permute_dims(args[i], permutation)
return _array_add(*args), tuple(index0)
if isinstance(expr, Pow):
subexpr, subindices = _convert_indexed_to_array(expr.base)
if isinstance(expr.exp, (int, Integer)):
diags = zip(*[(2 * i, 2 * i + 1) for i in range(expr.exp)])
arr = _array_diagonal(
_array_tensor_product(*[subexpr for i in range(expr.exp)]),
*diags)
return arr, subindices
if isinstance(expr, Function):
subexpr, subindices = _convert_indexed_to_array(expr.args[0])
return ArrayElementwiseApplyFunc(type(expr), subexpr), subindices
return expr, ()
def permutedims(expr, perm):
"""
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]]]
``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
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)
