def _(expr: ArrayContraction): expr = expr.flatten_contraction_of_diagonal() expr = expr.split_multiple_contractions() expr = identify_hadamard_products(expr) subexpr = expr.expr contraction_indices: Tuple[Tuple[int]] = expr.contraction_indices if isinstance(subexpr, ArrayTensorProduct): newexpr = ArrayContraction(_array2matrix(subexpr), *contraction_indices) contraction_indices = newexpr.contraction_indices if any(i > 2 for i in newexpr.subranks): addends = ArrayAdd(*[ _a2m_tensor_product(*j) for j in itertools.product(*[ i.args if isinstance(i, ArrayAdd) else [i] for i in expr.expr.args ]) ]) newexpr = ArrayContraction(addends, *contraction_indices) if isinstance(newexpr, ArrayAdd): ret = _array2matrix(newexpr) return ret assert isinstance(newexpr, ArrayContraction) ret = _support_function_tp1_recognize(contraction_indices, list(newexpr.expr.args)) return ret elif not isinstance(subexpr, _CodegenArrayAbstract): ret = _array2matrix(subexpr) if isinstance(ret, MatrixExpr): assert expr.contraction_indices == ((0, 1), ) return _a2m_trace(ret) else: return ArrayContraction(ret, *expr.contraction_indices)
def test_arrayexpr_split_multiple_contractions(): a = MatrixSymbol("a", k, 1) b = MatrixSymbol("b", k, 1) A = MatrixSymbol("A", k, k) B = MatrixSymbol("B", k, k) C = MatrixSymbol("C", k, k) X = MatrixSymbol("X", k, k) cg = ArrayContraction( ArrayTensorProduct(A.T, a, b, b.T, (A * X * b).applyfunc(cos)), (1, 2, 8), (5, 6, 9)) assert cg.split_multiple_contractions().dummy_eq( ArrayContraction( ArrayTensorProduct(DiagMatrix(a), (A * X * b).applyfunc(cos), A.T, b, b.T), (0, 2), (1, 5), (3, 7, 8))) # assert recognize_matrix_expression(cg) # Check no overlap of lines: cg = ArrayContraction(ArrayTensorProduct(A, a, C, a, B), (1, 2, 4), (5, 6, 8), (3, 7)) assert cg.split_multiple_contractions() == cg cg = ArrayContraction(ArrayTensorProduct(a, b, A), (0, 2, 4), (1, 3)) assert cg.split_multiple_contractions() == cg
def _(expr: ArrayContraction): rank1 = get_rank(expr) expr, removed1 = remove_identity_matrices(expr) if not isinstance(expr, ArrayContraction): expr2, removed2 = _remove_trivial_dims(expr) return expr2, _combine_removed(rank1, removed1, removed2) newexpr, removed2 = _remove_trivial_dims(expr.expr) shifts = list( accumulate( [1 if i in removed2 else 0 for i in range(get_rank(expr.expr))])) new_contraction_indices = [ tuple(j for j in i if j not in removed2) for i in expr.contraction_indices ] # Remove possible empty tuples "()": new_contraction_indices = [ i for i in new_contraction_indices if len(i) > 0 ] contraction_indices_flat = [j for i in expr.contraction_indices for j in i] removed2 = [i for i in removed2 if i not in contraction_indices_flat] new_contraction_indices = [ tuple(j - shifts[j] for j in i) for i in new_contraction_indices ] # Shift removed2: removed2 = ArrayContraction._push_indices_up(expr.contraction_indices, removed2) removed = _combine_removed(rank1, removed1, removed2) return ArrayContraction(newexpr, *new_contraction_indices), list(removed)
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_arrayexpr_convert_indexed_to_array_expression(): s = Sum(A[i] * B[i], (i, 0, 3)) cg = convert_indexed_to_array(s) assert cg == ArrayContraction(ArrayTensorProduct(A, B), (0, 1)) expr = M * N result = ArrayContraction(ArrayTensorProduct(M, N), (1, 2)) elem = expr[i, j] assert convert_indexed_to_array(elem) == result expr = M * N * M elem = expr[i, j] result = _array_contraction(_array_tensor_product(M, M, N), (1, 4), (2, 5)) cg = convert_indexed_to_array(elem) assert cg == result cg = convert_indexed_to_array((M * N * P)[i, j]) assert cg == _array_contraction(ArrayTensorProduct(M, N, P), (1, 2), (3, 4)) cg = convert_indexed_to_array((M * N.T * P)[i, j]) assert cg == _array_contraction(ArrayTensorProduct(M, N, P), (1, 3), (2, 4)) expr = -2 * M * N elem = expr[i, j] cg = convert_indexed_to_array(elem) assert cg == ArrayContraction(ArrayTensorProduct(-2, M, N), (1, 2))
def test_arrayexpr_permutedims_sink(): cg = PermuteDims(ArrayTensorProduct(M, N), [0, 1, 3, 2], nest_permutation=False) sunk = nest_permutation(cg) assert sunk == ArrayTensorProduct(M, PermuteDims(N, [1, 0])) cg = PermuteDims(ArrayTensorProduct(M, N), [1, 0, 3, 2], nest_permutation=False) sunk = nest_permutation(cg) assert sunk == ArrayTensorProduct(PermuteDims(M, [1, 0]), PermuteDims(N, [1, 0])) cg = PermuteDims(ArrayTensorProduct(M, N), [3, 2, 1, 0], nest_permutation=False) sunk = nest_permutation(cg) assert sunk == ArrayTensorProduct(PermuteDims(N, [1, 0]), PermuteDims(M, [1, 0])) cg = PermuteDims(ArrayContraction(ArrayTensorProduct(M, N), (1, 2)), [1, 0], nest_permutation=False) sunk = nest_permutation(cg) assert sunk == ArrayContraction(PermuteDims(ArrayTensorProduct(M, N), [[0, 3]]), (1, 2)) cg = PermuteDims(ArrayTensorProduct(M, N), [1, 0, 3, 2], nest_permutation=False) sunk = nest_permutation(cg) assert sunk == ArrayTensorProduct(PermuteDims(M, [1, 0]), PermuteDims(N, [1, 0])) cg = PermuteDims(ArrayContraction(ArrayTensorProduct(M, N, P), (1, 2), (3, 4)), [1, 0], nest_permutation=False) sunk = nest_permutation(cg) assert sunk == ArrayContraction(PermuteDims(ArrayTensorProduct(M, N, P), [[0, 5]]), (1, 2), (3, 4))
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_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_array_printer(): A = ArraySymbol('A', (4, 4, 6, 6, 6)) I = IndexedBase('I') i, j, k = Idx('i', (0, 1)), Idx('j', (2, 3)), Idx('k', (4, 5)) prntr = NumPyPrinter() assert prntr.doprint(ZeroArray(5)) == 'numpy.zeros((5,))' assert prntr.doprint(OneArray(5)) == 'numpy.ones((5,))' assert prntr.doprint(ArrayContraction( A, [2, 3])) == 'numpy.einsum("abccd->abd", A)' assert prntr.doprint(I) == 'I' assert prntr.doprint(ArrayDiagonal( A, [2, 3, 4])) == 'numpy.einsum("abccc->abc", A)' assert prntr.doprint(ArrayDiagonal( A, [0, 1], [2, 3])) == 'numpy.einsum("aabbc->cab", A)' assert prntr.doprint(ArrayContraction( A, [2], [3])) == 'numpy.einsum("abcde->abe", A)' assert prntr.doprint(Assignment(I[i, j, k], I[i, j, k])) == 'I = I' prntr = TensorflowPrinter() assert prntr.doprint(ZeroArray(5)) == 'tensorflow.zeros((5,))' assert prntr.doprint(OneArray(5)) == 'tensorflow.ones((5,))' assert prntr.doprint(ArrayContraction( A, [2, 3])) == 'tensorflow.linalg.einsum("abccd->abd", A)' assert prntr.doprint(I) == 'I' assert prntr.doprint(ArrayDiagonal( A, [2, 3, 4])) == 'tensorflow.linalg.einsum("abccc->abc", A)' assert prntr.doprint(ArrayDiagonal( A, [0, 1], [2, 3])) == 'tensorflow.linalg.einsum("aabbc->cab", A)' assert prntr.doprint(ArrayContraction( A, [2], [3])) == 'tensorflow.linalg.einsum("abcde->abe", A)' assert prntr.doprint(Assignment(I[i, j, k], I[i, j, k])) == 'I = I'
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 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_arrayexpr_nested_array_elementwise_add(): cg = ArrayContraction( ArrayAdd(ArrayTensorProduct(M, N), ArrayTensorProduct(N, M)), (1, 2)) result = ArrayAdd(ArrayContraction(ArrayTensorProduct(M, N), (1, 2)), ArrayContraction(ArrayTensorProduct(N, M), (1, 2))) assert cg == result cg = ArrayDiagonal( ArrayAdd(ArrayTensorProduct(M, N), ArrayTensorProduct(N, M)), (1, 2)) result = ArrayAdd(ArrayDiagonal(ArrayTensorProduct(M, N), (1, 2)), ArrayDiagonal(ArrayTensorProduct(N, M), (1, 2))) assert cg == result
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 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) == (ArrayContraction( 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) == (ArrayDiagonal( ArrayAdd( ArrayTensorProduct(M, N), PermuteDims(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) == (ArrayDiagonal( ArrayAdd( ArrayTensorProduct(M, N), PermuteDims(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 _(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 test_arrayexpr_convert_indexed_to_array_and_back_to_matrix(): expr = a.T * b elem = expr[0, 0] cg = convert_indexed_to_array(elem) assert cg == ArrayElement( ArrayContraction(ArrayTensorProduct(a, b), (0, 2)), [0, 0]) expr = M[i, j] + N[i, j] p1, p2 = _convert_indexed_to_array(expr) assert convert_array_to_matrix(p1) == M + N expr = M[i, j] + N[j, i] p1, p2 = _convert_indexed_to_array(expr) assert convert_array_to_matrix(p1) == M + N.T expr = M[i, j] * N[k, l] + N[i, j] * M[k, l] p1, p2 = _convert_indexed_to_array(expr) assert convert_array_to_matrix(p1) == ArrayAdd(ArrayTensorProduct(M, N), ArrayTensorProduct(N, M)) expr = (M * N * P)[i, j] p1, p2 = _convert_indexed_to_array(expr) assert convert_array_to_matrix(p1) == M * N * P expr = Sum(M[i, j] * (N * P)[j, m], (j, 0, k - 1)) p1, p2 = _convert_indexed_to_array(expr) assert convert_array_to_matrix(p1) == M * N * P expr = Sum((P[j, m] + P[m, j]) * (M[i, j] * N[m, n] + N[i, j] * M[m, n]), (j, 0, k - 1), (m, 0, k - 1)) p1, p2 = _convert_indexed_to_array(expr) assert convert_array_to_matrix( p1) == M * P * N + M * P.T * N + N * P * M + N * P.T * M
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 _(expr: ArrayContraction, x: Expr): fd = array_derive(expr.expr, x) rank_x = len(get_shape(x)) contraction_indices = expr.contraction_indices new_contraction_indices = [ tuple(j + rank_x for j in i) for i in contraction_indices ] return ArrayContraction(fd, *new_contraction_indices)
def test_arrayexpr_push_indices_up_and_down(): indices = list(range(12)) contr_diag_indices = [(0, 6), (2, 8)] assert ArrayContraction._push_indices_down(contr_diag_indices, indices) == (1, 3, 4, 5, 7, 9, 10, 11, 12, 13, 14, 15) assert ArrayContraction._push_indices_up(contr_diag_indices, indices) == (None, 0, None, 1, 2, 3, None, 4, None, 5, 6, 7) assert ArrayDiagonal._push_indices_down(contr_diag_indices, indices, 10) == (1, 3, 4, 5, 7, 9, (0, 6), (2, 8), None, None, None, None) assert ArrayDiagonal._push_indices_up(contr_diag_indices, indices, 10) == (6, 0, 7, 1, 2, 3, 6, 4, 7, 5, None, None) contr_diag_indices = [(1, 2), (7, 8)] assert ArrayContraction._push_indices_down(contr_diag_indices, indices) == (0, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15) assert ArrayContraction._push_indices_up(contr_diag_indices, indices) == (0, None, None, 1, 2, 3, 4, None, None, 5, 6, 7) assert ArrayDiagonal._push_indices_down(contr_diag_indices, indices, 10) == (0, 3, 4, 5, 6, 9, (1, 2), (7, 8), None, None, None, None) assert ArrayDiagonal._push_indices_up(contr_diag_indices, indices, 10) == (0, 6, 6, 1, 2, 3, 4, 7, 7, 5, None, None)
def _a2m_mul(*args): if all(not isinstance(i, _CodegenArrayAbstract) for i in args): from sympy import MatMul return MatMul(*args).doit() else: return ArrayContraction( ArrayTensorProduct(*args), *[(2 * i - 1, 2 * i) for i in range(1, len(args))])
def _support_function_tp1_recognize(contraction_indices, args): if len(contraction_indices) == 0: return _a2m_tensor_product(*args) ac = ArrayContraction(ArrayTensorProduct(*args), *contraction_indices) editor = _EditArrayContraction(ac) editor.track_permutation_start() while True: flag_stop: bool = True for i, arg_with_ind in enumerate(editor.args_with_ind): if not isinstance(arg_with_ind.element, MatrixExpr): continue first_index = arg_with_ind.indices[0] second_index = arg_with_ind.indices[1] first_frequency = editor.count_args_with_index(first_index) second_frequency = editor.count_args_with_index(second_index) if first_index is not None and first_frequency == 1 and first_index == second_index: flag_stop = False arg_with_ind.element = Trace(arg_with_ind.element)._normalize() arg_with_ind.indices = [] break scan_indices = [] if first_frequency == 2: scan_indices.append(first_index) if second_frequency == 2: scan_indices.append(second_index) candidate, transpose, found_index = _get_candidate_for_matmul_from_contraction( scan_indices, editor.args_with_ind[i + 1:]) if candidate is not None: flag_stop = False editor.track_permutation_merge(arg_with_ind, candidate) transpose1 = found_index == first_index new_arge, other_index = _insert_candidate_into_editor( editor, arg_with_ind, candidate, transpose1, transpose) if found_index == first_index: new_arge.indices = [second_index, other_index] else: new_arge.indices = [first_index, other_index] set_indices = set(new_arge.indices) if len(set_indices) == 1 and set_indices != {None}: # This is a trace: new_arge.element = Trace(new_arge.element)._normalize() new_arge.indices = [] editor.args_with_ind[i] = new_arge # TODO: is this break necessary? break if flag_stop: break editor.refresh_indices() return editor.to_array_contraction()
def test_arrayexpr_convert_array_contraction_tp_additions(): a = ArrayAdd( ArrayTensorProduct(M, N), ArrayTensorProduct(N, M) ) tp = ArrayTensorProduct(P, a, Q) expr = ArrayContraction(tp, (3, 4)) expected = ArrayTensorProduct( P, ArrayAdd( ArrayContraction(ArrayTensorProduct(M, N), (1, 2)), ArrayContraction(ArrayTensorProduct(N, M), (1, 2)), ), Q ) assert expr == expected assert convert_array_to_matrix(expr) == ArrayTensorProduct(P, M * N + N * M, Q) expr = ArrayContraction(tp, (1, 2), (3, 4), (5, 6)) result = ArrayContraction( ArrayTensorProduct( P, ArrayAdd( ArrayContraction(ArrayTensorProduct(M, N), (1, 2)), ArrayContraction(ArrayTensorProduct(N, M), (1, 2)), ), Q ), (1, 2), (3, 4)) assert expr == result assert convert_array_to_matrix(expr) == P * (M * N + N * M) * Q
def _(expr: ArrayContraction): newexpr, removed = _remove_trivial_dims(expr.expr) new_contraction_indices = [ tuple(j for j in i if j not in removed) for i in expr.contraction_indices ] # Remove possible empty tuples "()": new_contraction_indices = [i for i in new_contraction_indices if i] return ArrayContraction(newexpr, *new_contraction_indices), removed
def _array_diag2contr_diagmatrix(expr: ArrayDiagonal): if isinstance(expr.expr, ArrayTensorProduct): args = list(expr.expr.args) diag_indices = list(expr.diagonal_indices) mapping = _get_mapping_from_subranks( [_get_subrank(arg) for arg in args]) tuple_links = [[mapping[j] for j in i] for i in diag_indices] contr_indices = [] total_rank = get_rank(expr) replaced = [False for arg in args] for i, (abs_pos, rel_pos) in enumerate(zip(diag_indices, tuple_links)): if len(abs_pos) != 2: continue (pos1_outer, pos1_inner), (pos2_outer, pos2_inner) = rel_pos arg1 = args[pos1_outer] arg2 = args[pos2_outer] if get_rank(arg1) != 2 or get_rank(arg2) != 2: if replaced[pos1_outer]: diag_indices[i] = None if replaced[pos2_outer]: diag_indices[i] = None continue pos1_in2 = 1 - pos1_inner pos2_in2 = 1 - pos2_inner if arg1.shape[pos1_in2] == 1: if arg1.shape[pos1_inner] != 1: darg1 = DiagMatrix(arg1) else: darg1 = arg1 args.append(darg1) contr_indices.append( ((pos2_outer, pos2_inner), (len(args) - 1, pos1_inner))) total_rank += 1 diag_indices[i] = None args[pos1_outer] = OneArray(arg1.shape[pos1_in2]) replaced[pos1_outer] = True elif arg2.shape[pos2_in2] == 1: if arg2.shape[pos2_inner] != 1: darg2 = DiagMatrix(arg2) else: darg2 = arg2 args.append(darg2) contr_indices.append( ((pos1_outer, pos1_inner), (len(args) - 1, pos2_inner))) total_rank += 1 diag_indices[i] = None args[pos2_outer] = OneArray(arg2.shape[pos2_in2]) replaced[pos2_outer] = True diag_indices_new = [i for i in diag_indices if i is not None] cumul = list(accumulate([0] + [get_rank(arg) for arg in args])) contr_indices2 = [ tuple(cumul[a] + b for a, b in i) for i in contr_indices ] tc = ArrayContraction(ArrayTensorProduct(*args), *contr_indices2) td = ArrayDiagonal(tc, *diag_indices_new) return td return expr
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_split_multiple_contractions(): a = MatrixSymbol("a", k, 1) b = MatrixSymbol("b", k, 1) A = MatrixSymbol("A", k, k) B = MatrixSymbol("B", k, k) C = MatrixSymbol("C", k, k) X = MatrixSymbol("X", k, k) cg = ArrayContraction(ArrayTensorProduct(A.T, a, b, b.T, (A*X*b).applyfunc(cos)), (1, 2, 8), (5, 6, 9)) expected = ArrayContraction(ArrayTensorProduct(A.T, DiagMatrix(a), OneArray(1), b, b.T, (A*X*b).applyfunc(cos)), (1, 3), (2, 9), (6, 7, 10)) assert cg.split_multiple_contractions().dummy_eq(expected) # Check no overlap of lines: cg = ArrayContraction(ArrayTensorProduct(A, a, C, a, B), (1, 2, 4), (5, 6, 8), (3, 7)) assert cg.split_multiple_contractions() == cg cg = ArrayContraction(ArrayTensorProduct(a, b, A), (0, 2, 4), (1, 3)) assert cg.split_multiple_contractions() == cg
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_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_arrayexpr_convert_array_to_matrix_diag2contraction_diagmatrix(): cg = ArrayDiagonal(ArrayTensorProduct(M, a), (1, 2)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == ArrayContraction( ArrayTensorProduct(M, OneArray(1), DiagMatrix(a)), (1, 3)) raises(ValueError, lambda: ArrayDiagonal(ArrayTensorProduct(a, M), (1, 2))) cg = ArrayDiagonal(ArrayTensorProduct(a.T, M), (1, 2)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == ArrayContraction( ArrayTensorProduct(OneArray(1), M, DiagMatrix(a.T)), (1, 4)) cg = ArrayDiagonal(ArrayTensorProduct(a.T, M, N, b.T), (1, 2), (4, 7)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == ArrayContraction( ArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a.T), DiagMatrix(b.T)), (1, 7), (3, 9)) cg = ArrayDiagonal(ArrayTensorProduct(a, M, N, b.T), (0, 2), (4, 7)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == ArrayContraction( ArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a), DiagMatrix(b.T)), (1, 6), (3, 9)) cg = ArrayDiagonal(ArrayTensorProduct(a, M, N, b.T), (0, 4), (3, 7)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == ArrayContraction( ArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a), DiagMatrix(b.T)), (3, 6), (2, 9)) I1 = Identity(1) x = MatrixSymbol("x", k, 1) A = MatrixSymbol("A", k, k) cg = ArrayDiagonal(ArrayTensorProduct(x, A.T, I1), (0, 2)) assert _array_diag2contr_diagmatrix(cg).shape == cg.shape assert _array2matrix(cg).shape == cg.shape
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)