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_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_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_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_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_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: ArrayDiagonal): newexpr, removed = _remove_trivial_dims(expr.expr) shifts = list( accumulate( [0] + [1 if i in removed else 0 for i in range(get_rank(expr.expr))])) new_diag_indices_map = { i: tuple(j for j in i if j not in removed) for i in expr.diagonal_indices } for old_diag_tuple, new_diag_tuple in new_diag_indices_map.items(): if len(new_diag_tuple) == 1: removed = [i for i in removed if i not in old_diag_tuple] new_diag_indices = [ tuple(j - shifts[j] for j in i) for i in new_diag_indices_map.values() ] rank = get_rank(expr.expr) removed = ArrayDiagonal._push_indices_up(expr.diagonal_indices, removed, rank) removed = sorted({i for i in removed}) # If there are single axes to diagonalize remaining, it means that their # corresponding dimension has been removed, they no longer need diagonalization: new_diag_indices = [i for i in new_diag_indices if len(i) > 0] if len(new_diag_indices) > 0: newexpr2 = _array_diagonal(newexpr, *new_diag_indices, allow_trivial_diags=True) else: newexpr2 = newexpr if isinstance(newexpr2, ArrayDiagonal): newexpr3, removed2 = _remove_diagonalized_identity_matrices(newexpr2) removed = _combine_removed(-1, removed, removed2) return newexpr3, removed else: return newexpr2, removed
def _(expr: ArrayDiagonal): expr2 = _array2matrix(expr.expr) pexpr = _array_diag2contr_diagmatrix( ArrayDiagonal(expr2, *expr.diagonal_indices)) if expr == pexpr: return expr return _array2matrix(pexpr)
def _(expr: ElementwiseApplyFunction, x: Expr): assert get_rank(expr) == 2 assert get_rank(x) == 2 fdiff = expr._get_function_fdiff() dexpr = array_derive(expr.expr, x) tp = ArrayTensorProduct(ElementwiseApplyFunction(fdiff, expr.expr), dexpr) td = ArrayDiagonal(tp, (0, 4), (1, 5)) return td
def _(expr: ArrayDiagonal): pexpr = ArrayDiagonal(_array2matrix(expr.expr), *expr.diagonal_indices) pexpr = identify_hadamard_products(pexpr) if isinstance(pexpr, ArrayDiagonal): pexpr = _array_diag2contr_diagmatrix(pexpr) if expr == pexpr: return expr return _array2matrix(pexpr)
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 _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 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_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_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 _(expr: ArrayElementwiseApplyFunc, x: Expr): fdiff = expr._get_function_fdiff() subexpr = expr.expr dsubexpr = array_derive(subexpr, x) tp = ArrayTensorProduct(dsubexpr, ArrayElementwiseApplyFunc(fdiff, subexpr)) b = get_rank(x) c = get_rank(expr) diag_indices = [(b + i, b + c + i) for i in range(c)] return ArrayDiagonal(tp, *diag_indices)
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 _(expr: ArrayDiagonal): newexpr, removed = _remove_trivial_dims(expr.expr) shifts = list( accumulate( [0] + [1 if i in removed else 0 for i in range(get_rank(expr.expr))])) new_diag_indices = [ tuple(j for j in i if j not in removed) for i in expr.diagonal_indices ] new_diag_indices = [ tuple(j - shifts[j] for j in i) for i in new_diag_indices ] rank = get_rank(expr.expr) removed = ArrayDiagonal._push_indices_up(expr.diagonal_indices, removed, rank) removed = sorted({i for i in removed}) # If there are single axes to diagonalize remaining, it means that their # corresponding dimension has been removed, they no longer need diagonalization: new_diag_indices = [i for i in new_diag_indices if len(i) > 1] return ArrayDiagonal(newexpr, *new_diag_indices), removed
def 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 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_convert_array_element_to_array_expression(): A = ArraySymbol("A", (k, )) B = ArraySymbol("B", (k, )) s = Sum(A[i] * B[i], (i, 0, k - 1)) cg = convert_indexed_to_array(s) assert cg == ArrayContraction(ArrayTensorProduct(A, B), (0, 1)) s = A[i] * B[i] cg = convert_indexed_to_array(s) assert cg == ArrayDiagonal(ArrayTensorProduct(A, B), (0, 1)) s = A[i] * B[j] cg = convert_indexed_to_array(s, [i, j]) assert cg == ArrayTensorProduct(A, B) cg = convert_indexed_to_array(s, [j, i]) assert cg == ArrayTensorProduct(B, A)
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_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_arrayexpr_normalize_diagonal_contraction(): tp = ArrayTensorProduct(M, N, P, Q) expr = ArrayContraction(ArrayDiagonal(tp, (1, 3, 4)), (0, 3)) result = ArrayDiagonal(ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 6)), (0, 2, 3)) assert expr == result expr = ArrayContraction(ArrayDiagonal(tp, (0, 1, 2, 3, 7)), (1, 2, 3)) result = ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 1, 2, 3, 5, 6, 7)) assert expr == result expr = ArrayContraction(ArrayDiagonal(tp, (0, 2, 6, 7)), (1, 2, 3)) result = ArrayDiagonal(ArrayContraction(tp, (3, 4, 5)), (0, 2, 3, 4)) assert expr == result td = ArrayDiagonal(ArrayTensorProduct(M, N, P, Q), (0, 3)) expr = ArrayContraction(td, (2, 1), (0, 4, 6, 5, 3)) result = ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 1, 3, 5, 6, 7), (2, 4)) assert expr == result
def _remove_diagonalized_identity_matrices(expr: ArrayDiagonal): assert isinstance(expr, ArrayDiagonal) editor = _EditArrayContraction(expr) mapping = { i: {j for j in editor.args_with_ind if i in j.indices} for i in range(-1, -1 - editor.number_of_diagonal_indices, -1) } removed = [] counter: int = 0 for i, arg_with_ind in enumerate(editor.args_with_ind): counter += len(arg_with_ind.indices) if isinstance(arg_with_ind.element, Identity): if None in arg_with_ind.indices and any( i is not None and (i < 0) == True for i in arg_with_ind.indices): diag_ind = [j for j in arg_with_ind.indices if j is not None][0] other = [j for j in mapping[diag_ind] if j != arg_with_ind][0] if not isinstance(other.element, MatrixExpr): continue if 1 not in other.element.shape: continue if None not in other.indices: continue editor.args_with_ind[i].element = None none_index = other.indices.index(None) other.element = DiagMatrix(other.element) other_range = editor.get_absolute_range(other) removed.extend([other_range[0] + none_index]) editor.args_with_ind = [ i for i in editor.args_with_ind if i.element is not None ] removed = ArrayDiagonal._push_indices_up(expr.diagonal_indices, removed, get_rank(expr.expr)) return editor.to_array_contraction(), removed
def 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 # Hadamard products with traces: expr = Trace(HadamardProduct(M, N)) cg = convert_matrix_to_array(expr) ret = convert_array_to_matrix(cg) assert ret == Trace(HadamardProduct(M.T, N.T)) expr = Trace(A*HadamardProduct(M, N)) cg = convert_matrix_to_array(expr) ret = convert_array_to_matrix(cg) assert ret == Trace(HadamardProduct(M, N)*A) expr = Trace(HadamardProduct(A, M)*N) cg = convert_matrix_to_array(expr) ret = convert_array_to_matrix(cg) assert ret == Trace(HadamardProduct(M.T, N)*A) # These should not be converted into Hadamard products: cg = ArrayDiagonal(ArrayTensorProduct(M, N), (0, 1, 2, 3)) ret = convert_array_to_matrix(cg) assert ret == cg cg = ArrayDiagonal(ArrayTensorProduct(A), (0, 1)) ret = convert_array_to_matrix(cg) assert ret == cg cg = ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 2, 4), (1, 3, 5)) assert convert_array_to_matrix(cg) == HadamardProduct(M, N, P) cg = ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 3, 4), (1, 2, 5)) assert convert_array_to_matrix(cg) == HadamardProduct(M, P, N.T)
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, []) # A few more cases to test the removal and shift of nested removed axes # with array contractions and array diagonals: tp = ArrayTensorProduct( OneMatrix(1, 1), M, x, OneMatrix(1, 1), Identity(1), ) expr = ArrayContraction(tp, (1, 8)) rexpr, removed = _remove_trivial_dims(expr) assert removed == [0, 5, 6, 7] expr = ArrayContraction(tp, (1, 8), (3, 4)) rexpr, removed = _remove_trivial_dims(expr) assert removed == [0, 3, 4, 5] expr = ArrayDiagonal(tp, (1, 8)) rexpr, removed = _remove_trivial_dims(expr) assert removed == [0, 5, 6, 7, 8] expr = ArrayDiagonal(tp, (1, 8), (3, 4)) rexpr, removed = _remove_trivial_dims(expr) assert removed == [0, 3, 4, 5, 6] expr = ArrayDiagonal(ArrayContraction(ArrayTensorProduct(A, x, I, I1), (1, 2, 5)), (1, 4)) rexpr, removed = _remove_trivial_dims(expr) assert removed == [1, 2] cg = ArrayDiagonal(ArrayTensorProduct(PermuteDims(ArrayTensorProduct(x, I1), Permutation(1, 2, 3)), (x.T*x).applyfunc(sqrt)), (2, 4), (3, 5)) rexpr, removed = _remove_trivial_dims(cg) assert removed == [1, 2] # Contractions with identity matrices need to be followed by a permutation # in order cg = ArrayContraction(ArrayTensorProduct(A, B, C, M, I), (1, 8)) ret, removed = _remove_trivial_dims(cg) assert ret == PermuteDims(ArrayTensorProduct(A, B, C, M), [0, 2, 3, 4, 5, 6, 7, 1]) assert removed == [] cg = ArrayContraction(ArrayTensorProduct(A, B, C, M, I), (1, 8), (3, 4)) ret, removed = _remove_trivial_dims(cg) assert ret == PermuteDims(ArrayContraction(ArrayTensorProduct(A, B, C, M), (3, 4)), [0, 2, 3, 4, 5, 1]) assert removed == []
def test_arrayexpr_convert_array_to_diagonalized_vector(): # Check matrix recognition over trivial dimensions: cg = ArrayTensorProduct(a, b) assert convert_array_to_matrix(cg) == a * b.T cg = ArrayTensorProduct(I1, a, b) assert convert_array_to_matrix(cg) == a * b.T # Recognize trace inside a tensor product: cg = ArrayContraction(ArrayTensorProduct(A, B, C), (0, 3), (1, 2)) assert convert_array_to_matrix(cg) == Trace(A * B) * C # Transform diagonal operator to contraction: cg = ArrayDiagonal(ArrayTensorProduct(A, a), (1, 2)) assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(A, OneArray(1), DiagMatrix(a)), (1, 3)) assert convert_array_to_matrix(cg) == A * DiagMatrix(a) cg = ArrayDiagonal(ArrayTensorProduct(a, b), (0, 2)) assert _array_diag2contr_diagmatrix(cg) == PermuteDims( ArrayContraction(ArrayTensorProduct(DiagMatrix(a), OneArray(1), b), (0, 3)), [1, 2, 0] ) assert convert_array_to_matrix(cg) == b.T * DiagMatrix(a) cg = ArrayDiagonal(ArrayTensorProduct(A, a), (0, 2)) assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(A, OneArray(1), DiagMatrix(a)), (0, 3)) assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) cg = ArrayDiagonal(ArrayTensorProduct(I, x, I1), (0, 2), (3, 5)) assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(I, OneArray(1), I1, DiagMatrix(x)), (0, 5)) assert convert_array_to_matrix(cg) == DiagMatrix(x) cg = ArrayDiagonal(ArrayTensorProduct(I, x, A, B), (1, 2), (5, 6)) assert _array_diag2contr_diagmatrix(cg) == ArrayDiagonal(ArrayContraction(ArrayTensorProduct(I, OneArray(1), A, B, DiagMatrix(x)), (1, 7)), (5, 6)) # TODO: this is returning a wrong result: # convert_array_to_matrix(cg) cg = ArrayDiagonal(ArrayTensorProduct(I1, a, b), (1, 3, 5)) assert convert_array_to_matrix(cg) == a*b.T cg = ArrayDiagonal(ArrayTensorProduct(I1, a, b), (1, 3)) assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(OneArray(1), a, b, I1), (2, 6)) assert convert_array_to_matrix(cg) == a*b.T cg = ArrayDiagonal(ArrayTensorProduct(x, I1), (1, 2)) assert isinstance(cg, ArrayDiagonal) assert cg.diagonal_indices == ((1, 2),) assert convert_array_to_matrix(cg) == x cg = ArrayDiagonal(ArrayTensorProduct(x, I), (0, 2)) assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(OneArray(1), I, DiagMatrix(x)), (1, 3)) assert convert_array_to_matrix(cg).doit() == DiagMatrix(x) raises(ValueError, lambda: ArrayDiagonal(x, (1,))) # Ignore identity matrices with contractions: cg = ArrayContraction(ArrayTensorProduct(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 = ArrayContraction(ArrayTensorProduct(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 = ArrayContraction(ArrayTensorProduct(A, a), (1, 2)) assert cg.split_multiple_contractions() == cg assert convert_array_to_matrix(cg) == A * a cg = ArrayContraction(ArrayTensorProduct(A, a, B), (1, 2, 4)) assert cg.split_multiple_contractions() == ArrayContraction(ArrayTensorProduct(A, DiagMatrix(a), OneArray(1), B), (1, 2), (3, 5)) assert convert_array_to_matrix(cg) == A * DiagMatrix(a) * B cg = ArrayContraction(ArrayTensorProduct(A, a, B), (0, 2, 4)) assert cg.split_multiple_contractions() == ArrayContraction(ArrayTensorProduct(A, DiagMatrix(a), OneArray(1), B), (0, 2), (3, 5)) assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) * B cg = ArrayContraction(ArrayTensorProduct(A, a, b, a.T, B), (0, 2, 4, 7, 9)) assert cg.split_multiple_contractions() == ArrayContraction(ArrayTensorProduct(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 = ArrayContraction(ArrayTensorProduct(I1, I1, I1), (1, 2, 4)) assert cg.split_multiple_contractions() == ArrayContraction(ArrayTensorProduct(I1, I1, OneArray(1), I1), (1, 2), (3, 5)) assert convert_array_to_matrix(cg) == 1 cg = ArrayContraction(ArrayTensorProduct(I, I, I, I, A), (1, 2, 8), (5, 6, 9)) assert convert_array_to_matrix(cg.split_multiple_contractions()).doit() == A cg = ArrayContraction(ArrayTensorProduct(A, a, C, a, B), (1, 2, 4), (5, 6, 8)) expected = ArrayContraction(ArrayTensorProduct(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 = ArrayContraction(ArrayTensorProduct(a, I1, b, I1, (a.T*b).applyfunc(cos)), (1, 2, 8), (5, 6, 9)) expected = ArrayContraction(ArrayTensorProduct(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 test_array_expressions_no_canonicalization(): tp = _array_tensor_product(M, N, P) # ArrayTensorProduct: expr = ArrayTensorProduct(tp, N) assert str(expr) == "ArrayTensorProduct(ArrayTensorProduct(M, N, P), N)" assert expr.doit() == ArrayTensorProduct(M, N, P, N) expr = ArrayTensorProduct(ArrayContraction(M, (0, 1)), N) assert str(expr) == "ArrayTensorProduct(ArrayContraction(M, (0, 1)), N)" assert expr.doit() == ArrayContraction(ArrayTensorProduct(M, N), (0, 1)) expr = ArrayTensorProduct(ArrayDiagonal(M, (0, 1)), N) assert str(expr) == "ArrayTensorProduct(ArrayDiagonal(M, (0, 1)), N)" assert expr.doit() == PermuteDims(ArrayDiagonal(ArrayTensorProduct(M, N), (0, 1)), [2, 0, 1]) expr = ArrayTensorProduct(PermuteDims(M, [1, 0]), N) assert str(expr) == "ArrayTensorProduct(PermuteDims(M, (0 1)), N)" assert expr.doit() == PermuteDims(ArrayTensorProduct(M, N), [1, 0, 2, 3]) # ArrayContraction: expr = ArrayContraction(_array_contraction(tp, (0, 2)), (0, 1)) assert isinstance(expr, ArrayContraction) assert isinstance(expr.expr, ArrayContraction) assert str(expr) == "ArrayContraction(ArrayContraction(ArrayTensorProduct(M, N, P), (0, 2)), (0, 1))" assert expr.doit() == ArrayContraction(tp, (0, 2), (1, 3)) expr = ArrayContraction(ArrayContraction(ArrayContraction(tp, (0, 1)), (0, 1)), (0, 1)) assert expr.doit() == ArrayContraction(tp, (0, 1), (2, 3), (4, 5)) # assert expr._canonicalize() == ArrayContraction(ArrayContraction(tp, (0, 1)), (0, 1), (2, 3)) expr = ArrayContraction(ArrayDiagonal(tp, (0, 1)), (0, 1)) assert str(expr) == "ArrayContraction(ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 1)), (0, 1))" assert expr.doit() == ArrayDiagonal(ArrayContraction(ArrayTensorProduct(N, M, P), (0, 1)), (0, 1)) expr = ArrayContraction(PermuteDims(M, [1, 0]), (0, 1)) assert str(expr) == "ArrayContraction(PermuteDims(M, (0 1)), (0, 1))" assert expr.doit() == ArrayContraction(M, (0, 1)) # ArrayDiagonal: expr = ArrayDiagonal(ArrayDiagonal(tp, (0, 2)), (0, 1)) assert str(expr) == "ArrayDiagonal(ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 2)), (0, 1))" assert expr.doit() == ArrayDiagonal(tp, (0, 2), (1, 3)) expr = ArrayDiagonal(ArrayDiagonal(ArrayDiagonal(tp, (0, 1)), (0, 1)), (0, 1)) assert expr.doit() == ArrayDiagonal(tp, (0, 1), (2, 3), (4, 5)) assert expr._canonicalize() == expr.doit() expr = ArrayDiagonal(ArrayContraction(tp, (0, 1)), (0, 1)) assert str(expr) == "ArrayDiagonal(ArrayContraction(ArrayTensorProduct(M, N, P), (0, 1)), (0, 1))" assert expr.doit() == expr expr = ArrayDiagonal(PermuteDims(M, [1, 0]), (0, 1)) assert str(expr) == "ArrayDiagonal(PermuteDims(M, (0 1)), (0, 1))" assert expr.doit() == ArrayDiagonal(M, (0, 1)) # ArrayAdd: expr = ArrayAdd(M) assert isinstance(expr, ArrayAdd) assert expr.doit() == M expr = ArrayAdd(ArrayAdd(M, N), P) assert str(expr) == "ArrayAdd(ArrayAdd(M, N), P)" assert expr.doit() == ArrayAdd(M, N, P) expr = ArrayAdd(M, ArrayAdd(N, ArrayAdd(P, M))) assert expr.doit() == ArrayAdd(M, N, P, M) assert expr._canonicalize() == ArrayAdd(M, N, ArrayAdd(P, M)) expr = ArrayAdd(M, ZeroArray(k, k), N) assert str(expr) == "ArrayAdd(M, ZeroArray(k, k), N)" assert expr.doit() == ArrayAdd(M, N) # PermuteDims: expr = PermuteDims(PermuteDims(M, [1, 0]), [1, 0]) assert str(expr) == "PermuteDims(PermuteDims(M, (0 1)), (0 1))" assert expr.doit() == M expr = PermuteDims(PermuteDims(PermuteDims(M, [1, 0]), [1, 0]), [1, 0]) assert expr.doit() == PermuteDims(M, [1, 0]) assert expr._canonicalize() == expr.doit()