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 ArrayTensorProduct(M, N, za1) == ZeroArray(k, k, k, k, k, l, m, n) assert ArrayTensorProduct(M, N, zm1) == ZeroArray(k, k, k, k, m, n) assert ArrayContraction(za1, (3, )) == ZeroArray(k, l, m) assert ArrayContraction(zm1, (1, )) == ZeroArray(m) assert ArrayContraction(za2, (1, 2)) == ZeroArray(k, n) assert ArrayContraction(zm2, (0, 1)) == 0 assert ArrayDiagonal(za2, (1, 2)) == ZeroArray(k, n, m) assert ArrayDiagonal(zm2, (0, 1)) == ZeroArray(m) assert PermuteDims(za1, [2, 1, 3, 0]) == ZeroArray(m, l, n, k) assert PermuteDims(zm1, [1, 0]) == ZeroArray(n, m) assert ArrayAdd(za1) == za1 assert ArrayAdd(zm1) == ZeroArray(m, n) tp1 = ArrayTensorProduct(MatrixSymbol("A", k, l), MatrixSymbol("B", m, n)) assert ArrayAdd(tp1, za1) == tp1 tp2 = ArrayTensorProduct(MatrixSymbol("C", k, l), MatrixSymbol("D", m, n)) assert ArrayAdd(tp1, za1, tp2) == ArrayAdd(tp1, tp2) assert ArrayAdd(M, zm3) == M assert ArrayAdd(M, N, zm3) == ArrayAdd(M, N)
def test_arrayexpr_convert_indexed_to_array_broadcast(): A = ArraySymbol("A", (3, 3)) B = ArraySymbol("B", (3, 3)) expr = A[i, j] + B[k, l] O2 = OneArray(3, 3) expected = ArrayAdd(ArrayTensorProduct(A, O2), ArrayTensorProduct(O2, B)) assert convert_indexed_to_array(expr) == expected assert convert_indexed_to_array(expr, [i, j, k, l]) == expected assert convert_indexed_to_array(expr, [l, k, i, j]) == ArrayAdd( PermuteDims(ArrayTensorProduct(O2, A), [1, 0, 2, 3]), PermuteDims(ArrayTensorProduct(B, O2), [1, 0, 2, 3])) expr = A[i, j] + B[j, k] O1 = OneArray(3) assert convert_indexed_to_array(expr, [i, j, k]) == ArrayAdd( ArrayTensorProduct(A, O1), ArrayTensorProduct(O1, B)) C = ArraySymbol("C", (d0, d1)) D = ArraySymbol("D", (d3, d1)) expr = C[i, j] + D[k, j] assert convert_indexed_to_array(expr, [i, j, k]) == ArrayAdd( ArrayTensorProduct(C, OneArray(d3)), PermuteDims(ArrayTensorProduct(OneArray(d0), D), [0, 2, 1])) X = ArraySymbol("X", (5, 3)) expr = X[i, n] - X[j, n] assert convert_indexed_to_array(expr, [i, j, n]) == ArrayAdd( ArrayTensorProduct(-1, OneArray(5), X), PermuteDims(ArrayTensorProduct(X, OneArray(5)), [0, 2, 1])) raises(ValueError, lambda: convert_indexed_to_array(C[i, j] + D[i, j]))
def test_array_expr_construction_with_functions(): tp = tensorproduct(M, N) assert tp == ArrayTensorProduct(M, N) expr = tensorproduct(A, eye(2)) assert expr == ArrayTensorProduct(A, eye(2)) # Contraction: expr = tensorcontraction(M, (0, 1)) assert expr == ArrayContraction(M, (0, 1)) expr = tensorcontraction(tp, (1, 2)) assert expr == ArrayContraction(tp, (1, 2)) expr = tensorcontraction(tensorcontraction(tp, (1, 2)), (0, 1)) assert expr == ArrayContraction(tp, (0, 3), (1, 2)) # Diagonalization: expr = tensordiagonal(M, (0, 1)) assert expr == ArrayDiagonal(M, (0, 1)) expr = tensordiagonal(tensordiagonal(tp, (0, 1)), (0, 1)) assert expr == ArrayDiagonal(tp, (0, 1), (2, 3)) # Permutation of dimensions: expr = permutedims(M, [1, 0]) assert expr == PermuteDims(M, [1, 0]) expr = permutedims(PermuteDims(tp, [1, 0, 2, 3]), [0, 1, 3, 2]) assert expr == PermuteDims(tp, [1, 0, 3, 2])
def test_array2matrix(): # See issue https://github.com/sympy/sympy/pull/22877 expr = PermuteDims( ArrayContraction(ArrayTensorProduct(x, I, I1, x), (0, 3), (1, 7)), Permutation(2, 3)) expected = PermuteDims(ArrayTensorProduct(x * x.T, I1), Permutation(3)(1, 2)) assert _array2matrix(expr) == expected
def test_arrayexpr_convert_array_to_matrix(): cg = ArrayContraction(ArrayTensorProduct(M), (0, 1)) assert convert_array_to_matrix(cg) == Trace(M) cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 1), (2, 3)) assert convert_array_to_matrix(cg) == Trace(M) * Trace(N) cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 3), (1, 2)) assert convert_array_to_matrix(cg) == Trace(M * N) cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 2), (1, 3)) assert convert_array_to_matrix(cg) == Trace(M * N.T) cg = convert_matrix_to_array(M * N * P) assert convert_array_to_matrix(cg) == M * N * P cg = convert_matrix_to_array(M * N.T * P) assert convert_array_to_matrix(cg) == M * N.T * P cg = ArrayContraction(ArrayTensorProduct(M,N,P,Q), (1, 2), (5, 6)) assert convert_array_to_matrix(cg) == ArrayTensorProduct(M * N, P * Q) cg = ArrayContraction(ArrayTensorProduct(-2, M, N), (1, 2)) assert convert_array_to_matrix(cg) == -2 * M * N a = MatrixSymbol("a", k, 1) b = MatrixSymbol("b", k, 1) c = MatrixSymbol("c", k, 1) cg = PermuteDims( ArrayContraction( ArrayTensorProduct( a, ArrayAdd( ArrayTensorProduct(b, c), ArrayTensorProduct(c, b), ) ), (2, 4)), [0, 1, 3, 2]) assert convert_array_to_matrix(cg) == a * (b.T * c + c.T * b) za = ZeroArray(m, n) assert convert_array_to_matrix(za) == ZeroMatrix(m, n) cg = ArrayTensorProduct(3, M) assert convert_array_to_matrix(cg) == 3 * M # Partial conversion to matrix multiplication: expr = ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 2), (1, 4, 6)) assert convert_array_to_matrix(expr) == ArrayContraction(ArrayTensorProduct(M.T*N, P, Q), (0, 2, 4)) x = MatrixSymbol("x", k, 1) cg = PermuteDims( ArrayContraction(ArrayTensorProduct(OneArray(1), x, OneArray(1), DiagMatrix(Identity(1))), (0, 5)), Permutation(1, 2, 3)) assert convert_array_to_matrix(cg) == x expr = ArrayAdd(M, PermuteDims(M, [1, 0])) assert convert_array_to_matrix(expr) == M + Transpose(M)
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 PermuteDims(_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 PermuteDims(mat_mul_lines, permutation) pos = p1 // 2 if p1 > p2: args_array[i] = _a2m_transpose(mat_mul_lines.args[pos]) else: args_array[i] = mat_mul_lines.args[pos] return _a2m_tensor_product(*args_array) else: return expr
def test_matrix_derivative_non_matrix_result(): # This is a 4-dimensional array: I = Identity(k) AdA = PermuteDims(ArrayTensorProduct(I, I), Permutation(3)(1, 2)) assert A.diff(A) == AdA assert A.T.diff(A) == PermuteDims(ArrayTensorProduct(I, I), Permutation(3)(1, 2, 3)) assert (2*A).diff(A) == PermuteDims(ArrayTensorProduct(2*I, I), Permutation(3)(1, 2)) assert MatAdd(A, A).diff(A) == ArrayAdd(AdA, AdA) assert (A + B).diff(A) == AdA
def test_arrayexpr_convert_array_to_matrix_remove_trivial_dims(): # Tensor Product: assert _remove_trivial_dims(ArrayTensorProduct(a, b)) == (a * b.T, [1, 3]) assert _remove_trivial_dims(ArrayTensorProduct(a.T, b)) == (a * b.T, [0, 3]) assert _remove_trivial_dims(ArrayTensorProduct(a, b.T)) == (a * b.T, [1, 2]) assert _remove_trivial_dims(ArrayTensorProduct(a.T, b.T)) == (a * b.T, [0, 2]) assert _remove_trivial_dims(ArrayTensorProduct(I, a.T, b.T)) == (a * b.T, [0, 1, 2, 4]) assert _remove_trivial_dims(ArrayTensorProduct(a.T, I, b.T)) == (a * b.T, [0, 2, 3, 4]) assert _remove_trivial_dims(ArrayTensorProduct(a, I)) == (a, [2, 3]) assert _remove_trivial_dims(ArrayTensorProduct(I, a)) == (a, [0, 1]) assert _remove_trivial_dims(ArrayTensorProduct(a.T, b.T, c, d)) == (ArrayTensorProduct( a * b.T, c * d.T), [0, 2, 5, 7]) assert _remove_trivial_dims(ArrayTensorProduct(a.T, I, b.T, c, d, I)) == (ArrayTensorProduct( a * b.T, c * d.T, I), [0, 2, 3, 4, 7, 9]) # Addition: cg = ArrayAdd(ArrayTensorProduct(a, b), ArrayTensorProduct(c, d)) assert _remove_trivial_dims(cg) == (a * b.T + c * d.T, [1, 3]) # Permute Dims: cg = PermuteDims(ArrayTensorProduct(a, b), Permutation(3)(1, 2)) assert _remove_trivial_dims(cg) == (a * b.T, [2, 3]) cg = PermuteDims(ArrayTensorProduct(a, I, b), Permutation(5)(1, 2, 3, 4)) assert _remove_trivial_dims(cg) == (a * b.T, [1, 2, 4, 5]) cg = PermuteDims(ArrayTensorProduct(I, b, a), Permutation(5)(1, 2, 4, 5, 3)) assert _remove_trivial_dims(cg) == (b * a.T, [0, 3, 4, 5]) # Diagonal: cg = ArrayDiagonal(ArrayTensorProduct(M, a), (1, 2)) assert _remove_trivial_dims(cg) == (cg, []) # Contraction: cg = ArrayContraction(ArrayTensorProduct(M, a), (1, 2)) assert _remove_trivial_dims(cg) == (cg, [])
def test_arrayexpr_convert_array_to_matrix2(): cg = ArrayContraction(ArrayTensorProduct(M, N), (1, 3)) assert convert_array_to_matrix(cg) == M * N.T cg = PermuteDims(ArrayTensorProduct(M, N), Permutation([0, 1, 3, 2])) assert convert_array_to_matrix(cg) == ArrayTensorProduct(M, N.T) cg = ArrayTensorProduct(M, PermuteDims(N, Permutation([1, 0]))) assert convert_array_to_matrix(cg) == ArrayTensorProduct(M, N.T) cg = ArrayContraction( PermuteDims( ArrayTensorProduct(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])), (1, 2), (3, 5) ) assert convert_array_to_matrix(cg) == ArrayTensorProduct(M * P.T * Trace(N), Q.T) cg = ArrayContraction( ArrayTensorProduct(M, N, P, PermuteDims(Q, Permutation([1, 0]))), (1, 5), (2, 3) ) assert convert_array_to_matrix(cg) == ArrayTensorProduct(M * P.T * Trace(N), Q.T) cg = ArrayTensorProduct(M, PermuteDims(N, [1, 0])) assert convert_array_to_matrix(cg) == ArrayTensorProduct(M, N.T) cg = ArrayTensorProduct(PermuteDims(M, [1, 0]), PermuteDims(N, [1, 0])) assert convert_array_to_matrix(cg) == ArrayTensorProduct(M.T, N.T) cg = ArrayTensorProduct(PermuteDims(N, [1, 0]), PermuteDims(M, [1, 0])) assert convert_array_to_matrix(cg) == ArrayTensorProduct(N.T, M.T)
def test_arrayexpr_contraction_permutation_mix(): Me = M.subs(k, 3).as_explicit() Ne = N.subs(k, 3).as_explicit() cg1 = ArrayContraction(PermuteDims(ArrayTensorProduct(M, N), Permutation([0, 2, 1, 3])), (2, 3)) cg2 = ArrayContraction(ArrayTensorProduct(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 = PermuteDims(ArrayTensorProduct(M, N), Permutation([0, 1, 3, 2])) cg2 = ArrayTensorProduct(M, PermuteDims(N, Permutation([1, 0]))) assert cg1 == cg2 cg1 = ArrayContraction( PermuteDims( ArrayTensorProduct(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])), (1, 2), (3, 5) ) cg2 = ArrayContraction( ArrayTensorProduct(M, N, P, PermuteDims(Q, Permutation([1, 0]))), (1, 5), (2, 3) ) assert cg1 == cg2 cg1 = ArrayContraction( PermuteDims( ArrayTensorProduct(M, N, P, Q), Permutation([1, 0, 4, 6, 2, 7, 5, 3])), (0, 1), (2, 6), (3, 7) ) cg2 = PermuteDims( ArrayContraction( ArrayTensorProduct(M, P, Q, N), (0, 1), (2, 3), (4, 7)), [1, 0] ) assert cg1 == cg2 cg1 = ArrayContraction( PermuteDims( ArrayTensorProduct(M, N, P, Q), Permutation([1, 0, 4, 6, 7, 2, 5, 3])), (0, 1), (2, 6), (3, 7) ) cg2 = PermuteDims( ArrayContraction( ArrayTensorProduct(PermuteDims(M, [1, 0]), N, P, Q), (0, 1), (3, 6), (4, 5) ), Permutation([1, 0]) ) assert cg1 == cg2
def test_arrayexpr_normalize_diagonal_permutedims(): tp = ArrayTensorProduct(M, Q, N, P) expr = ArrayDiagonal(PermuteDims(tp, [0, 1, 2, 4, 7, 6, 3, 5]), (2, 4, 5), (6, 7), (0, 3)) result = ArrayDiagonal(tp, (2, 6, 7), (3, 5), (0, 4)) assert expr == result tp = ArrayTensorProduct(M, N, P, Q) expr = ArrayDiagonal(PermuteDims(tp, [0, 5, 2, 4, 1, 6, 3, 7]), (1, 2, 6), (3, 4)) result = ArrayDiagonal(ArrayTensorProduct(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_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 == PermuteDims( 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) ))
def _(expr: Inverse, x: Expr): mat = expr.I dexpr = array_derive(mat, x) tp = ArrayTensorProduct(-expr, dexpr, expr) mp = ArrayContraction(tp, (1, 4), (5, 6)) pp = PermuteDims(mp, [1, 2, 0, 3]) return pp
def _(expr: ArraySymbol, x: Expr): if expr == x: return PermuteDims( 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 test_mixed_deriv_mixed_expressions(): expr = 3 * Trace(A) assert expr.diff(A) == 3 * Identity(k) expr = k deriv = expr.diff(A) assert isinstance(deriv, ZeroMatrix) assert deriv == ZeroMatrix(k, k) expr = Trace(A)**2 assert expr.diff(A) == (2 * Trace(A)) * Identity(k) expr = Trace(A) * A I = Identity(k) assert expr.diff(A) == ArrayAdd( ArrayTensorProduct(I, A), PermuteDims(ArrayTensorProduct(Trace(A) * I, I), Permutation(3)(1, 2))) expr = Trace(Trace(A) * A) assert expr.diff(A) == (2 * Trace(A)) * Identity(k) expr = Trace(Trace(Trace(A) * A) * A) assert expr.diff(A) == (3 * Trace(A)**2) * Identity(k)
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 = ArrayTensorProduct(*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 = PermuteDims(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 ArrayAdd(*addend_list)
def test_arrayexpr_nested_permutations(): cg = PermuteDims(PermuteDims(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 = PermuteDims( PermuteDims( ArrayTensorProduct(M, N, P), p1), p2 ) result = PermuteDims( ArrayTensorProduct(M, N, P), p2*p1 ) assert cg == result # Check that `permutedims` behaves the same way with explicit-component arrays: result1 = permutedims(permutedims(cge, p1), p2) result2 = permutedims(cge, p2*p1) assert result1 == result2
def test_arrayexpr_convert_array_to_matrix_support_function(): assert _support_function_tp1_recognize([], [2 * k]) == 2 * k assert _support_function_tp1_recognize([(1, 2)], [A, 2 * k, B, 3]) == 6 * k * A * B assert _support_function_tp1_recognize([(0, 3), (1, 2)], [A, B]) == Trace(A * B) assert _support_function_tp1_recognize([(1, 2)], [A, B]) == A * B assert _support_function_tp1_recognize([(0, 2)], [A, B]) == A.T * B assert _support_function_tp1_recognize([(1, 3)], [A, B]) == A * B.T assert _support_function_tp1_recognize([(0, 3)], [A, B]) == A.T * B.T assert _support_function_tp1_recognize([(1, 2), (5, 6)], [A, B, C, D]) == ArrayTensorProduct(A * B, C * D) assert _support_function_tp1_recognize([(1, 4), (3, 6)], [A, B, C, D]) == PermuteDims( ArrayTensorProduct(A * C, B * D), [0, 2, 1, 3]) assert _support_function_tp1_recognize([(0, 3), (1, 4)], [A, B, C]) == B * A * C assert _support_function_tp1_recognize([(9, 10), (1, 2), (5, 6), (3, 4), (7, 8)], [X, Y, A, B, C, D]) == X * Y * A * B * C * D assert _support_function_tp1_recognize([(9, 10), (1, 2), (5, 6), (3, 4)], [X, Y, A, B, C, D]) == ArrayTensorProduct(X * Y * A * B, C * D) assert _support_function_tp1_recognize([(1, 7), (3, 8), (4, 11)], [X, Y, A, B, C, D]) == PermuteDims( ArrayTensorProduct(X * B.T, Y * C, A.T * D.T), [0, 2, 4, 1, 3, 5] ) assert _support_function_tp1_recognize([(0, 1), (3, 6), (5, 8)], [X, A, B, C, D]) == PermuteDims( ArrayTensorProduct(Trace(X) * A * C, B * D), [0, 2, 1, 3]) assert _support_function_tp1_recognize([(1, 2), (3, 4), (5, 6), (7, 8)], [A, A, B, C, D]) == A ** 2 * B * C * D assert _support_function_tp1_recognize([(1, 2), (3, 4), (5, 6), (7, 8)], [X, A, B, C, D]) == X * A * B * C * D assert _support_function_tp1_recognize([(1, 6), (3, 8), (5, 10)], [X, Y, A, B, C, D]) == PermuteDims( ArrayTensorProduct(X * B, Y * C, A * D), [0, 2, 4, 1, 3, 5] ) assert _support_function_tp1_recognize([(1, 4), (3, 6)], [A, B, C, D]) == PermuteDims( ArrayTensorProduct(A * C, B * D), [0, 2, 1, 3]) assert _support_function_tp1_recognize([(0, 4), (1, 7), (2, 5), (3, 8)], [X, A, B, C, D]) == C*X.T*B*A*D assert _support_function_tp1_recognize([(0, 4), (1, 7), (2, 5), (3, 8)], [X, A, B, C, D]) == C*X.T*B*A*D
def convert_matrix_to_array(expr: MatrixExpr) -> 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 = ArrayTensorProduct.fromiter( args_nonmat) if args_nonmat else S.One if scalar == 1: tprod = ArrayTensorProduct( *[convert_matrix_to_array(arg) for arg in args]) else: tprod = ArrayTensorProduct( scalar, *[convert_matrix_to_array(arg) for arg in args]) return ArrayContraction(tprod, *contractions) elif isinstance(expr, MatAdd): return ArrayAdd(*[convert_matrix_to_array(arg) for arg in expr.args]) elif isinstance(expr, Transpose): return PermuteDims(convert_matrix_to_array(expr.args[0]), [1, 0]) elif isinstance(expr, Trace): inner_expr = convert_matrix_to_array(expr.arg) return ArrayContraction(inner_expr, (0, len(inner_expr.shape) - 1)) elif isinstance(expr, Mul): return ArrayTensorProduct.fromiter( 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 ArrayTensorProduct.fromiter(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 = ArrayTensorProduct.fromiter(expr.args) diag = [[2 * i for i in range(len(expr.args))], [2 * i + 1 for i in range(len(expr.args))]] return ArrayDiagonal(tp, *diag) elif isinstance(expr, HadamardPower): base, exp = expr.args return convert_matrix_to_array( HadamardProduct.fromiter(base for i in range(exp))) else: return expr
def test_codegen_extra(): if not np: skip("NumPy not installed") M = MatrixSymbol("M", 2, 2) N = MatrixSymbol("N", 2, 2) P = MatrixSymbol("P", 2, 2) Q = MatrixSymbol("Q", 2, 2) ma = np.matrix([[1, 2], [3, 4]]) mb = np.matrix([[1, -2], [-1, 3]]) mc = np.matrix([[2, 0], [1, 2]]) md = np.matrix([[1, -1], [4, 7]]) cg = ArrayTensorProduct(M, N) f = lambdify((M, N), cg, 'numpy') assert (f(ma, mb) == np.einsum(ma, [0, 1], mb, [2, 3])).all() cg = ArrayAdd(M, N) f = lambdify((M, N), cg, 'numpy') assert (f(ma, mb) == ma + mb).all() cg = ArrayAdd(M, N, P) f = lambdify((M, N, P), cg, 'numpy') assert (f(ma, mb, mc) == ma + mb + mc).all() cg = ArrayAdd(M, N, P, Q) f = lambdify((M, N, P, Q), cg, 'numpy') assert (f(ma, mb, mc, md) == ma + mb + mc + md).all() cg = PermuteDims(M, [1, 0]) f = lambdify((M, ), cg, 'numpy') assert (f(ma) == ma.T).all() cg = PermuteDims(ArrayTensorProduct(M, N), [1, 2, 3, 0]) f = lambdify((M, N), cg, 'numpy') assert (f(ma, mb) == np.transpose(np.einsum(ma, [0, 1], mb, [2, 3]), (1, 2, 3, 0))).all() cg = ArrayDiagonal(ArrayTensorProduct(M, N), (1, 2)) f = lambdify((M, N), cg, 'numpy') assert (f(ma, mb) == np.diagonal(np.einsum(ma, [0, 1], mb, [2, 3]), axis1=1, axis2=2)).all()
def test_array_expr_construction_with_functions(): tp = tensorproduct(M, N) assert tp == ArrayTensorProduct(M, N) expr = tensorproduct(A, eye(2)) assert expr == ArrayTensorProduct(A, eye(2)) # Contraction: expr = tensorcontraction(M, (0, 1)) assert expr == ArrayContraction(M, (0, 1)) expr = tensorcontraction(tp, (1, 2)) assert expr == ArrayContraction(tp, (1, 2)) expr = tensorcontraction(tensorcontraction(tp, (1, 2)), (0, 1)) assert expr == ArrayContraction(tp, (0, 3), (1, 2)) # Diagonalization: expr = tensordiagonal(M, (0, 1)) assert expr == ArrayDiagonal(M, (0, 1)) expr = tensordiagonal(tensordiagonal(tp, (0, 1)), (0, 1)) assert expr == ArrayDiagonal(tp, (0, 1), (2, 3)) # Permutation of dimensions: expr = permutedims(M, [1, 0]) assert expr == PermuteDims(M, [1, 0]) expr = permutedims(PermuteDims(tp, [1, 0, 2, 3]), [0, 1, 3, 2]) assert expr == PermuteDims(tp, [1, 0, 3, 2]) expr = PermuteDims(tp, index_order_new=["a", "b", "c", "d"], index_order_old=["d", "c", "b", "a"]) assert expr == PermuteDims(tp, [3, 2, 1, 0]) arr = Array(range(32)).reshape(2, 2, 2, 2, 2) expr = PermuteDims(arr, index_order_new=["a", "b", "c", "d", "e"], index_order_old=['b', 'e', 'a', 'd', 'c']) assert expr == PermuteDims(arr, [2, 0, 4, 3, 1]) assert expr.as_explicit() == permutedims(arr, index_order_new=["a", "b", "c", "d", "e"], index_order_old=['b', 'e', 'a', 'd', 'c'])
def test_array_expr_as_explicit_with_explicit_component_arrays(): # Test if .as_explicit() works with explicit-component arrays # nested in array expressions: from sympy.abc import x, y, z, t A = Array([[x, y], [z, t]]) assert ArrayTensorProduct(A, A).as_explicit() == tensorproduct(A, A) assert ArrayDiagonal(A, (0, 1)).as_explicit() == tensordiagonal(A, (0, 1)) assert ArrayContraction(A, (0, 1)).as_explicit() == tensorcontraction(A, (0, 1)) assert ArrayAdd(A, A).as_explicit() == A + A assert ArrayElementwiseApplyFunc(sin, A).as_explicit() == A.applyfunc(sin) assert PermuteDims(A, [1, 0]).as_explicit() == permutedims(A, [1, 0]) assert Reshape(A, [4]).as_explicit() == A.reshape(4)
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 not first_indices: return result for i in first_indices: if i not in indices: first_indices.remove(i) first_indices.extend([i for i in indices if i not in first_indices]) permutation = [first_indices.index(i) for i in indices] return PermuteDims(result, permutation)
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 = PermuteDims(subexpr, p2) if newexpr != expr: newexpr = _array2matrix(newexpr) return newexpr, sorted(premoved)
def test_arrayexpr_convert_array_to_implicit_matmul(): # Trivial dimensions are suppressed, so the result can be expressed in matrix form: cg = ArrayTensorProduct(a, b) assert convert_array_to_matrix(cg) == a * b.T cg = ArrayTensorProduct(a, I, b) assert convert_array_to_matrix(cg) == a * b.T cg = ArrayContraction(ArrayTensorProduct(I, I), (1, 2)) assert convert_array_to_matrix(cg) == I cg = PermuteDims(ArrayTensorProduct(I, Identity(1)), [0, 2, 1, 3]) assert convert_array_to_matrix(cg) == I
def test_arrayexpr_convert_array_to_matrix2(): cg = _array_contraction(_array_tensor_product(M, N), (1, 3)) assert convert_array_to_matrix(cg) == M * N.T cg = PermuteDims(_array_tensor_product(M, N), Permutation([0, 1, 3, 2])) assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T) cg = _array_tensor_product(M, PermuteDims(N, Permutation([1, 0]))) assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T) cg = _array_contraction( PermuteDims( _array_tensor_product(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])), (1, 2), (3, 5) ) assert convert_array_to_matrix(cg) == _array_tensor_product(M * P.T * Trace(N), Q.T) cg = _array_contraction( _array_tensor_product(M, N, P, PermuteDims(Q, Permutation([1, 0]))), (1, 5), (2, 3) ) assert convert_array_to_matrix(cg) == _array_tensor_product(M * P.T * Trace(N), Q.T) cg = _array_tensor_product(M, PermuteDims(N, [1, 0])) assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T) cg = _array_tensor_product(PermuteDims(M, [1, 0]), PermuteDims(N, [1, 0])) assert convert_array_to_matrix(cg) == _array_tensor_product(M.T, N.T) cg = _array_tensor_product(PermuteDims(N, [1, 0]), PermuteDims(M, [1, 0])) assert convert_array_to_matrix(cg) == _array_tensor_product(N.T, M.T) cg = _array_contraction(M, (0,), (1,)) assert convert_array_to_matrix(cg) == OneMatrix(1, k)*M*OneMatrix(k, 1) cg = _array_contraction(x, (0,), (1,)) assert convert_array_to_matrix(cg) == OneMatrix(1, k)*x Xm = MatrixSymbol("Xm", m, n) cg = _array_contraction(Xm, (0,), (1,)) assert convert_array_to_matrix(cg) == OneMatrix(1, m)*Xm*OneMatrix(n, 1)
def test_convert_array_to_hadamard_products(): expr = HadamardProduct(M, N) cg = convert_matrix_to_array(expr) ret = convert_array_to_matrix(cg) assert ret == expr expr = HadamardProduct(M, N) * P cg = convert_matrix_to_array(expr) ret = convert_array_to_matrix(cg) assert ret == expr expr = Q * HadamardProduct(M, N) * P cg = convert_matrix_to_array(expr) ret = convert_array_to_matrix(cg) assert ret == expr expr = Q * HadamardProduct(M, N.T) * P cg = convert_matrix_to_array(expr) ret = convert_array_to_matrix(cg) assert ret == expr expr = HadamardProduct(M, N) * HadamardProduct(Q, P) cg = convert_matrix_to_array(expr) ret = convert_array_to_matrix(cg) assert expr == ret expr = P.T * HadamardProduct(M, N) * HadamardProduct(Q, P) cg = convert_matrix_to_array(expr) ret = convert_array_to_matrix(cg) assert expr == ret # ArrayDiagonal should be converted cg = ArrayDiagonal(ArrayTensorProduct(M, N, Q), (1, 3), (0, 2, 4)) ret = convert_array_to_matrix(cg) expected = PermuteDims( ArrayDiagonal(ArrayTensorProduct(HadamardProduct(M.T, N.T), Q), (1, 2)), [1, 0, 2]) assert expected == ret # Special case that should return the same expression: cg = ArrayDiagonal(ArrayTensorProduct(HadamardProduct(M, N), Q), (0, 2)) ret = convert_array_to_matrix(cg) assert ret == cg
def test_arrayexpr_array_shape(): expr = ArrayTensorProduct(M, N, P, Q) assert expr.shape == (k, k, k, k, k, k, k, k) Z = MatrixSymbol("Z", m, n) expr = ArrayTensorProduct(M, Z) assert expr.shape == (k, k, m, n) expr2 = ArrayContraction(expr, (0, 1)) assert expr2.shape == (m, n) expr2 = ArrayDiagonal(expr, (0, 1)) assert expr2.shape == (m, n, k) exprp = PermuteDims(expr, [2, 1, 3, 0]) assert exprp.shape == (m, k, n, k) expr3 = ArrayTensorProduct(N, Z) expr2 = ArrayAdd(expr, expr3) assert expr2.shape == (k, k, m, n) # Contraction along axes with discordant dimensions: raises(ValueError, lambda: ArrayContraction(expr, (1, 2))) # Also diagonal needs the same dimensions: raises(ValueError, lambda: ArrayDiagonal(expr, (1, 2))) # Diagonal requires at least to axes to compute the diagonal: raises(ValueError, lambda: ArrayDiagonal(expr, (1, )))
def test_arrayexpr_permute_tensor_product(): cg1 = PermuteDims(ArrayTensorProduct(M, N, P, Q), Permutation([2, 3, 1, 0, 5, 4, 6, 7])) cg2 = ArrayTensorProduct(N, PermuteDims(M, [1, 0]), PermuteDims(P, [1, 0]), Q) assert cg1 == cg2 # TODO: reverse operation starting with `PermuteDims` and getting down to `bb`... cg1 = PermuteDims(ArrayTensorProduct(M, N, P, Q), Permutation([2, 3, 4, 5, 0, 1, 6, 7])) cg2 = ArrayTensorProduct(N, P, M, Q) assert cg1 == cg2 cg1 = PermuteDims(ArrayTensorProduct(M, N, P, Q), Permutation([2, 3, 4, 6, 5, 7, 0, 1])) assert cg1.expr == ArrayTensorProduct(N, P, Q, M) assert cg1.permutation == Permutation([0, 1, 2, 4, 3, 5, 6, 7]) cg1 = ArrayContraction( PermuteDims(ArrayTensorProduct(N, Q, Q, M), [2, 1, 5, 4, 0, 3, 6, 7]), [1, 2, 6]) cg2 = PermuteDims( ArrayContraction(ArrayTensorProduct(Q, Q, N, M), (3, 5, 6)), [0, 2, 3, 1, 4]) assert cg1 == cg2 cg1 = ArrayContraction( ArrayContraction( ArrayContraction( ArrayContraction( PermuteDims(ArrayTensorProduct(N, Q, Q, M), [2, 1, 5, 4, 0, 3, 6, 7]), [1, 2, 6]), [1, 3, 4]), [1]), [0]) cg2 = ArrayContraction(ArrayTensorProduct(M, N, Q, Q), (0, 3, 5), (1, 4, 7), (2, ), (6, )) assert cg1 == cg2