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)