def test_codegen_array_recognize_matrix_mul_lines():
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M), (0, 1))
assert recognize_matrix_expression(cg) == Trace(M)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 1), (2, 3))
assert recognize_matrix_expression(cg) == [Trace(M), Trace(N)]
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 3), (1, 2))
assert recognize_matrix_expression(cg) == Trace(M*N)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 2), (1, 3))
assert recognize_matrix_expression(cg) == Trace(M*N.T)
cg = parse_indexed_expression((M*N*P)[i,j])
assert recognize_matrix_expression(cg) == M*N*P
cg = CodegenArrayContraction.from_MatMul(M*N*P)
assert recognize_matrix_expression(cg) == M*N*P
cg = parse_indexed_expression((M*N.T*P)[i,j])
assert recognize_matrix_expression(cg) == M*N.T*P
cg = CodegenArrayContraction.from_MatMul(M*N.T*P)
assert recognize_matrix_expression(cg) == M*N.T*P
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M,N,P,Q), (1, 2), (5, 6))
assert recognize_matrix_expression(cg) == [M*N, P*Q]
expr = -2*M*N
elem = expr[i, j]
cg = parse_indexed_expression(elem)
assert recognize_matrix_expression(cg) == -2*M*N
def test_codegen_array_contraction_construction():
cg = CodegenArrayContraction(A)
assert cg == A
s = Sum(A[i]*B[i], (i, 0, 3))
cg = parse_indexed_expression(s)
assert cg == CodegenArrayContraction(CodegenArrayTensorProduct(A, B), (0, 1))
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, B), (1, 0))
assert cg == CodegenArrayContraction(CodegenArrayTensorProduct(A, B), (0, 1))
expr = M*N
result = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2))
assert CodegenArrayContraction.from_MatMul(expr) == result
elem = expr[i, j]
assert parse_indexed_expression(elem) == result
expr = M*N*M
result = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, M), (1, 2), (3, 4))
assert CodegenArrayContraction.from_MatMul(expr) == result
elem = expr[i, j]
result = CodegenArrayContraction(CodegenArrayTensorProduct(M, M, N), (1, 4), (2, 5))
cg = parse_indexed_expression(elem)
cg = cg.sort_args_by_name()
assert cg == result
def test_push_indices_up_and_down():
indices = list(range(10))
contraction_indices = [(0, 6), (2, 8)]
assert CodegenArrayContraction._push_indices_down(contraction_indices, indices) == (1, 3, 4, 5, 7, 9, 10, 11, 12, 13)
assert CodegenArrayContraction._push_indices_up(contraction_indices, indices) == (None, 0, None, 1, 2, 3, None, 4, None, 5)
assert CodegenArrayDiagonal._push_indices_down(contraction_indices, indices) == (0, 1, 2, 3, 4, 5, 7, 9, 10, 11)
assert CodegenArrayDiagonal._push_indices_up(contraction_indices, indices) == (0, 1, 2, 3, 4, 5, None, 6, None, 7)
contraction_indices = [(1, 2), (7, 8)]
assert CodegenArrayContraction._push_indices_down(contraction_indices, indices) == (0, 3, 4, 5, 6, 9, 10, 11, 12, 13)
assert CodegenArrayContraction._push_indices_up(contraction_indices, indices) == (0, None, None, 1, 2, 3, 4, None, None, 5)
assert CodegenArrayContraction._push_indices_down(contraction_indices, indices) == (0, 3, 4, 5, 6, 9, 10, 11, 12, 13)
assert CodegenArrayDiagonal._push_indices_up(contraction_indices, indices) == (0, 1, None, 2, 3, 4, 5, 6, None, 7)
def test_contraction_tp_additions():
a = CodegenArrayElementwiseAdd(CodegenArrayTensorProduct(M, N),
CodegenArrayTensorProduct(N, M))
tp = CodegenArrayTensorProduct(P, a, Q)
expr = CodegenArrayContraction(tp, (3, 4))
expected = CodegenArrayTensorProduct(
P,
CodegenArrayElementwiseAdd(
CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2)),
CodegenArrayContraction(CodegenArrayTensorProduct(N, M), (1, 2)),
), Q)
assert expr == expected
assert recognize_matrix_expression(expr) == CodegenArrayTensorProduct(
P, M * N + N * M, Q)
expr = CodegenArrayContraction(tp, (1, 2), (3, 4), (5, 6))
result = CodegenArrayContraction(
CodegenArrayTensorProduct(
P,
CodegenArrayElementwiseAdd(
CodegenArrayContraction(CodegenArrayTensorProduct(M, N),
(1, 2)),
CodegenArrayContraction(CodegenArrayTensorProduct(N, M),
(1, 2)),
), Q), (1, 2), (3, 4))
assert expr == result
assert recognize_matrix_expression(expr) == P * (M * N + N * M) * Q
def test_contraction_permutation_mix():
Me = M.subs(k, 3).as_explicit()
Ne = N.subs(k, 3).as_explicit()
cg1 = CodegenArrayContraction(
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N),
Permutation([0, 2, 1, 3])), (2, 3))
cg2 = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 3))
assert cg1 == cg2
assert recognize_matrix_expression(cg2) == M * N.T
cge1 = tensorcontraction(
permutedims(tensorproduct(Me, Ne), Permutation([0, 2, 1, 3])), (2, 3))
cge2 = tensorcontraction(tensorproduct(Me, Ne), (1, 3))
assert cge1 == cge2
cg1 = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N),
Permutation([0, 1, 3, 2]))
cg2 = CodegenArrayTensorProduct(
M, CodegenArrayPermuteDims(N, Permutation([1, 0])))
assert cg1 == cg2
assert recognize_matrix_expression(cg1) == [M, N.T]
assert recognize_matrix_expression(cg2) == [M, N.T]
cg1 = CodegenArrayContraction(
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q),
Permutation([0, 2, 3, 1, 4, 5, 7, 6])), (1, 2),
(3, 5))
cg2 = CodegenArrayContraction(
CodegenArrayTensorProduct(
M, N, P, CodegenArrayPermuteDims(Q, Permutation([1, 0]))), (1, 5),
(2, 3))
assert cg1 == cg2
assert recognize_matrix_expression(cg1) == [M * P.T * Trace(N), Q.T]
assert recognize_matrix_expression(cg2) == [M * P.T * Trace(N), Q.T]
cg1 = CodegenArrayContraction(
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q),
Permutation([1, 0, 4, 6, 2, 7, 5, 3])), (0, 1),
(2, 6), (3, 7))
cg2 = CodegenArrayPermuteDims(
CodegenArrayContraction(CodegenArrayTensorProduct(M, P, Q, N), (0, 1),
(2, 3), (4, 7)), [1, 0])
assert cg1 == cg2
cg1 = CodegenArrayContraction(
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q),
Permutation([1, 0, 4, 6, 7, 2, 5, 3])), (0, 1),
(2, 6), (3, 7))
cg2 = CodegenArrayPermuteDims(
CodegenArrayContraction(
CodegenArrayTensorProduct(CodegenArrayPermuteDims(M, [1, 0]), N, P,
Q), (0, 1), (3, 6), (4, 5)),
Permutation([1, 0]))
assert cg1 == cg2
def test_codegen_array_recognize_matrix_mul_lines():
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M), (0, 1))
assert recognize_matrix_expression(cg) == Trace(M)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 1),
(2, 3))
assert recognize_matrix_expression(cg) == [Trace(M), Trace(N)]
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 3),
(1, 2))
assert recognize_matrix_expression(cg) == Trace(M * N)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 2),
(1, 3))
assert recognize_matrix_expression(cg) == Trace(M * N.T)
cg = parse_indexed_expression((M * N * P)[i, j])
assert recognize_matrix_expression(cg) == M * N * P
cg = CodegenArrayContraction.from_MatMul(M * N * P)
assert recognize_matrix_expression(cg) == M * N * P
cg = parse_indexed_expression((M * N.T * P)[i, j])
assert recognize_matrix_expression(cg) == M * N.T * P
cg = CodegenArrayContraction.from_MatMul(M * N.T * P)
assert recognize_matrix_expression(cg) == M * N.T * P
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (1, 2),
(5, 6))
assert recognize_matrix_expression(cg) == [M * N, P * Q]
expr = -2 * M * N
elem = expr[i, j]
cg = parse_indexed_expression(elem)
assert recognize_matrix_expression(cg) == -2 * M * N
def test_special_matrices():
a = MatrixSymbol("a", k, 1)
b = MatrixSymbol("b", k, 1)
expr = a.T*b
elem = expr[0, 0]
cg = parse_indexed_expression(elem)
assert cg == CodegenArrayContraction(CodegenArrayTensorProduct(a, b), (0, 2))
assert recognize_matrix_expression(cg) == a.T*b
def test_recognize_diagonalized_vectors():
a = MatrixSymbol("a", k, 1)
b = MatrixSymbol("b", k, 1)
A = MatrixSymbol("A", k, k)
B = MatrixSymbol("B", k, k)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a), (1, 2))
assert recognize_matrix_expression(cg) == A*a
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, B), (1, 2, 4))
assert recognize_matrix_expression(cg) == A*DiagonalizeVector(a)*B
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, B), (0, 2, 4))
assert recognize_matrix_expression(cg) == A.T*DiagonalizeVector(a)*B
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, b, a.T, B), (0, 2, 4, 7, 9))
assert recognize_matrix_expression(cg).doit() == A.T*DiagonalizeVector(a)*DiagonalizeVector(b)*DiagonalizeVector(a)*B.T
def test_codegen_permutedims_sink():
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [0, 1, 3, 2],
nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == CodegenArrayTensorProduct(
M, CodegenArrayPermuteDims(N, [1, 0]))
assert recognize_matrix_expression(sunk) == [M, N.T]
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [1, 0, 3, 2],
nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == CodegenArrayTensorProduct(
CodegenArrayPermuteDims(M, [1, 0]), CodegenArrayPermuteDims(N, [1, 0]))
assert recognize_matrix_expression(sunk) == [M.T, N.T]
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [3, 2, 1, 0],
nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == CodegenArrayTensorProduct(
CodegenArrayPermuteDims(N, [1, 0]), CodegenArrayPermuteDims(M, [1, 0]))
assert recognize_matrix_expression(sunk) == [N.T, M.T]
cg = CodegenArrayPermuteDims(CodegenArrayContraction(
CodegenArrayTensorProduct(M, N), (1, 2)), [1, 0],
nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == CodegenArrayContraction(
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [[0, 3]]),
(1, 2))
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [1, 0, 3, 2],
nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == CodegenArrayTensorProduct(
CodegenArrayPermuteDims(M, [1, 0]), CodegenArrayPermuteDims(N, [1, 0]))
cg = CodegenArrayPermuteDims(CodegenArrayContraction(
CodegenArrayTensorProduct(M, N, P), (1, 2), (3, 4)), [1, 0],
nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == CodegenArrayContraction(
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P), [[0, 5]]),
(1, 2), (3, 4))
def test_codegen_array_contraction_indices_types():
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 1))
indtup = cg._get_contraction_tuples()
assert indtup == [[(0, 0), (0, 1)]]
assert cg._contraction_tuples_to_contraction_indices(cg.expr, indtup) == [(0, 1)]
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2))
indtup = cg._get_contraction_tuples()
assert indtup == [[(0, 1), (1, 0)]]
assert cg._contraction_tuples_to_contraction_indices(cg.expr, indtup) == [(1, 2)]
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, M, N), (1, 4), (2, 5))
indtup = cg._get_contraction_tuples()
assert indtup == [[(0, 1), (2, 0)], [(1, 0), (2, 1)]]
assert cg._contraction_tuples_to_contraction_indices(cg.expr, indtup) == [(1, 4), (2, 5)]
def test_nested_array_elementwise_add():
cg = CodegenArrayContraction(CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(M, N),
CodegenArrayTensorProduct(N, M)
), (1, 2))
result = CodegenArrayElementwiseAdd(
CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2)),
CodegenArrayContraction(CodegenArrayTensorProduct(N, M), (1, 2))
)
assert cg == result
cg = CodegenArrayDiagonal(CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(M, N),
CodegenArrayTensorProduct(N, M)
), (1, 2))
result = CodegenArrayElementwiseAdd(
CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (1, 2)),
CodegenArrayDiagonal(CodegenArrayTensorProduct(N, M), (1, 2))
)
assert cg == result
def test_codegen_array_shape():
expr = CodegenArrayTensorProduct(M, N, P, Q)
assert expr.shape == (k, k, k, k, k, k, k, k)
Z = MatrixSymbol("Z", m, n)
expr = CodegenArrayTensorProduct(M, Z)
assert expr.shape == (k, k, m, n)
expr2 = CodegenArrayContraction(expr, (0, 1))
assert expr2.shape == (m, n)
expr2 = CodegenArrayDiagonal(expr, (0, 1))
assert expr2.shape == (m, n, k)
exprp = CodegenArrayPermuteDims(expr, [2, 1, 3, 0])
assert exprp.shape == (m, k, n, k)
expr3 = CodegenArrayTensorProduct(N, Z)
expr2 = CodegenArrayElementwiseAdd(expr, expr3)
assert expr2.shape == (k, k, m, n)
# Contraction along axes with discordant dimensions:
raises(ValueError, lambda: CodegenArrayContraction(expr, (1, 2)))
# Also diagonal needs the same dimensions:
raises(ValueError, lambda: CodegenArrayDiagonal(expr, (1, 2)))
def test_codegen_array_parse():
expr = M[i, j]
assert _codegen_array_parse(expr) == (M, (i, j))
expr = M[i, j] * N[k, l]
assert _codegen_array_parse(expr) == (CodegenArrayTensorProduct(M, N),
(i, j, k, l))
expr = M[i, j] * N[j, k]
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(
CodegenArrayTensorProduct(M, N), (1, 2)), (i, k, j))
expr = Sum(M[i, j] * N[j, k], (j, 0, k - 1))
assert _codegen_array_parse(expr) == (CodegenArrayContraction(
CodegenArrayTensorProduct(M, N), (1, 2)), (i, k))
expr = M[i, j] + N[i, j]
assert _codegen_array_parse(expr) == (CodegenArrayElementwiseAdd(M, N),
(i, j))
expr = M[i, j] + N[j, i]
assert _codegen_array_parse(expr) == (CodegenArrayElementwiseAdd(
M, CodegenArrayPermuteDims(N, Permutation([1, 0]))), (i, j))
expr = M[i, j] + M[j, i]
assert _codegen_array_parse(expr) == (CodegenArrayElementwiseAdd(
M, CodegenArrayPermuteDims(M, Permutation([1, 0]))), (i, j))
expr = (M * N * P)[i, j]
assert _codegen_array_parse(expr) == (CodegenArrayContraction(
CodegenArrayTensorProduct(M, N, P), (1, 2), (3, 4)), (i, j))
expr = expr.function # Disregard summation in previous expression
ret1, ret2 = _codegen_array_parse(expr)
assert ret1 == CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P),
(1, 2), (3, 4))
assert str(ret2) == "(i, j, _i_1, _i_2)"
expr = KroneckerDelta(i, j) * M[i, k]
assert _codegen_array_parse(expr) == (M, ({i, j}, k))
expr = KroneckerDelta(j, k) * (M[i, j] * N[k, l] + N[i, j] * M[k, l])
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(
CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(M, N),
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N),
Permutation(0, 2)(1, 3))),
(1, 2)), (i, l, frozenset({j, k})))
expr = M[i, i]
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(M,
(0, 1)), (i, ))
def test_parsing_of_matrix_expressions():
expr = M * N
assert parse_matrix_expression(expr) == CodegenArrayContraction(
CodegenArrayTensorProduct(M, N), (1, 2))
expr = Transpose(M)
assert parse_matrix_expression(expr) == CodegenArrayPermuteDims(M, [1, 0])
expr = M * Transpose(N)
assert parse_matrix_expression(expr) == CodegenArrayContraction(
CodegenArrayTensorProduct(M, CodegenArrayPermuteDims(N, [1, 0])),
(1, 2))
expr = 3 * M * N
res = parse_matrix_expression(expr)
rexpr = recognize_matrix_expression(res)
assert expr == rexpr
expr = 3 * M + N * M.T * M + 4 * k * N
res = parse_matrix_expression(expr)
rexpr = recognize_matrix_expression(res)
assert expr == rexpr
expr = Inverse(M) * N
rexpr = recognize_matrix_expression(parse_matrix_expression(expr))
assert expr == rexpr
expr = M**2
rexpr = recognize_matrix_expression(parse_matrix_expression(expr))
assert expr == rexpr
expr = M * (2 * N + 3 * M)
res = parse_matrix_expression(expr)
rexpr = recognize_matrix_expression(res)
assert expr.expand() == rexpr.doit()
expr = Trace(M)
result = CodegenArrayContraction(M, (0, 1))
assert parse_matrix_expression(expr) == result
def test_codegen_einsum():
if not np:
skip("NumPy not installed")
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
cg = CodegenArrayContraction.from_MatMul(M*N)
f = lambdify((M, N), cg, 'numpy')
ma = np.matrix([[1, 2], [3, 4]])
mb = np.matrix([[1,-2], [-1, 3]])
assert (f(ma, mb) == ma*mb).all()
def test_codegen_array_doit():
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
P = MatrixSymbol("P", 2, 2)
Q = MatrixSymbol("Q", 2, 2)
M = M.as_explicit()
N = N.as_explicit()
P = P.as_explicit()
Q = Q.as_explicit()
expr = CodegenArrayTensorProduct(M, N, P, Q)
assert expr.doit() == tensorproduct(M, N, P, Q)
expr2 = CodegenArrayContraction(expr, (0, 1))
assert expr2.doit() == tensorcontraction(tensorproduct(M, N, P, Q), (0, 1))
expr2 = CodegenArrayDiagonal(expr, (0, 1))
#assert expr2 = ... # TODO: not implemented
expr = CodegenArrayTensorProduct(M, N)
exprp = CodegenArrayPermuteDims(expr, [2, 1, 3, 0])
assert exprp.doit() == permutedims(tensorproduct(M, N), [2, 1, 3, 0])
expr = CodegenArrayElementwiseAdd(M, N)
assert expr.doit() == M + N
def test_recognize_expression_implicit_mul():
cg = CodegenArrayTensorProduct(a, b)
assert recognize_matrix_expression(cg) == a*b.T
cg = CodegenArrayTensorProduct(a, I, b)
assert recognize_matrix_expression(cg) == a*b.T
cg = CodegenArrayContraction(CodegenArrayTensorProduct(I, I), (1, 2))
assert recognize_matrix_expression(cg) == I
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(I, Identity(1)), [0, 2, 1, 3])
assert recognize_matrix_expression(cg) == I
def test_codegen_einsum():
if not np:
skip("NumPy not installed")
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
cg = CodegenArrayContraction.from_MatMul(M*N)
f = lambdify((M, N), cg, 'numpy')
ma = np.matrix([[1, 2], [3, 4]])
mb = np.matrix([[1,-2], [-1, 3]])
assert (f(ma, mb) == ma*mb).all()
def test_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 CodegenArrayTensorProduct(M, N,
za1) == ZeroArray(k, k, k, k, k, l, m, n)
assert CodegenArrayTensorProduct(M, N, zm1) == ZeroArray(k, k, k, k, m, n)
assert CodegenArrayContraction(za1, (3, )) == ZeroArray(k, l, m)
assert CodegenArrayContraction(zm1, (1, )) == ZeroArray(m)
assert CodegenArrayContraction(za2, (1, 2)) == ZeroArray(k, n)
assert CodegenArrayContraction(zm2, (0, 1)) == 0
assert CodegenArrayDiagonal(za2, (1, 2)) == ZeroArray(k, n, m)
assert CodegenArrayDiagonal(zm2, (0, 1)) == ZeroArray(m)
assert CodegenArrayPermuteDims(za1, [2, 1, 3, 0]) == ZeroArray(m, l, n, k)
assert CodegenArrayPermuteDims(zm1, [1, 0]) == ZeroArray(n, m)
assert CodegenArrayElementwiseAdd(za1) == za1
assert CodegenArrayElementwiseAdd(zm1) == ZeroArray(m, n)
tp1 = CodegenArrayTensorProduct(MatrixSymbol("A", k, l),
MatrixSymbol("B", m, n))
assert CodegenArrayElementwiseAdd(tp1, za1) == tp1
tp2 = CodegenArrayTensorProduct(MatrixSymbol("C", k, l),
MatrixSymbol("D", m, n))
assert CodegenArrayElementwiseAdd(tp1, za1,
tp2) == CodegenArrayElementwiseAdd(
tp1, tp2)
assert CodegenArrayElementwiseAdd(M, zm3) == M
assert CodegenArrayElementwiseAdd(M, N,
zm3) == CodegenArrayElementwiseAdd(M, N)
def test_codegen_array_recognize_matrix_mul_lines():
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M), (0, 1))
assert recognize_matrix_expression(cg) == Trace(M)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 1),
(2, 3))
assert recognize_matrix_expression(cg) == Trace(M) * Trace(N)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 3),
(1, 2))
assert recognize_matrix_expression(cg) == Trace(M * N)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 2),
(1, 3))
assert recognize_matrix_expression(cg) == Trace(M * N.T)
cg = parse_indexed_expression((M * N * P)[i, j])
assert recognize_matrix_expression(cg) == M * N * P
cg = parse_matrix_expression(M * N * P)
assert recognize_matrix_expression(cg) == M * N * P
cg = parse_indexed_expression((M * N.T * P)[i, j])
assert recognize_matrix_expression(cg) == M * N.T * P
cg = parse_matrix_expression(M * N.T * P)
assert recognize_matrix_expression(cg) == M * N.T * P
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (1, 2),
(5, 6))
assert recognize_matrix_expression(cg) == CodegenArrayTensorProduct(
M * N, P * Q)
expr = -2 * M * N
elem = expr[i, j]
cg = parse_indexed_expression(elem)
assert recognize_matrix_expression(cg) == -2 * M * N
a = MatrixSymbol("a", k, 1)
b = MatrixSymbol("b", k, 1)
c = MatrixSymbol("c", k, 1)
cg = CodegenArrayPermuteDims(
CodegenArrayContraction(
CodegenArrayTensorProduct(
a,
CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(b, c),
CodegenArrayTensorProduct(c, b),
)), (2, 4)), [0, 1, 3, 2])
assert recognize_matrix_expression(cg) == a * (b.T * c + c.T * b)
za = ZeroArray(m, n)
assert recognize_matrix_expression(za) == ZeroMatrix(m, n)
cg = CodegenArrayTensorProduct(3, M)
assert recognize_matrix_expression(cg) == 3 * M
def test_normalize_diagonal_contraction():
tp = CodegenArrayTensorProduct(M, N, P, Q)
expr = CodegenArrayContraction(CodegenArrayDiagonal(tp, (1, 3, 4)), (0, 3))
result = CodegenArrayDiagonal(CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (0, 6)), (0, 2, 3))
assert expr == result
expr = CodegenArrayContraction(CodegenArrayDiagonal(tp, (0, 1, 2, 3, 7)), (1, 2, 3))
result = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (0, 1, 2, 3, 5, 6, 7))
assert expr == result
expr = CodegenArrayContraction(CodegenArrayDiagonal(tp, (0, 2, 6, 7)), (1, 2, 3))
result = CodegenArrayDiagonal(CodegenArrayContraction(tp, (3, 4, 5)), (0, 2, 3, 4))
assert expr == result
td = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P, Q), (0, 3))
expr = CodegenArrayContraction(td, (2, 1), (0, 4, 6, 5, 3))
result = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (0, 1, 3, 5, 6, 7), (2, 4))
assert expr == result
def test_codegen_einsum():
if not tf:
skip("TensorFlow not installed")
session = tf.Session()
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
cg = CodegenArrayContraction.from_MatMul(M*N)
f = lambdify((M, N), cg, 'tensorflow')
ma = tf.constant([[1, 2], [3, 4]])
mb = tf.constant([[1,-2], [-1, 3]])
y = session.run(f(ma, mb))
c = session.run(tf.matmul(ma, mb))
assert (y == c).all()
def test_codegen_einsum():
if not tf:
skip("TensorFlow not installed")
session = tf.Session()
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
cg = CodegenArrayContraction.from_MatMul(M * N)
f = lambdify((M, N), cg, 'tensorflow')
ma = tf.constant([[1, 2], [3, 4]])
mb = tf.constant([[1, -2], [-1, 3]])
y = session.run(f(ma, mb))
c = session.run(tf.matmul(ma, mb))
assert (y == c).all()
def test_remove_trivial_dims():
# Tensor Product:
assert _remove_trivial_dims(CodegenArrayTensorProduct(a, b)) == (a * b.T, [1, 3])
assert _remove_trivial_dims(CodegenArrayTensorProduct(a.T, b)) == (a * b.T, [0, 3])
assert _remove_trivial_dims(CodegenArrayTensorProduct(a, b.T)) == (a * b.T, [1, 2])
assert _remove_trivial_dims(CodegenArrayTensorProduct(a.T, b.T)) == (a * b.T, [0, 2])
assert _remove_trivial_dims(CodegenArrayTensorProduct(I, a.T, b.T)) == (a * b.T, [0, 1, 2, 4])
assert _remove_trivial_dims(CodegenArrayTensorProduct(a.T, I, b.T)) == (a * b.T, [0, 2, 3, 4])
assert _remove_trivial_dims(CodegenArrayTensorProduct(a, I)) == (a, [2, 3])
assert _remove_trivial_dims(CodegenArrayTensorProduct(I, a)) == (a, [0, 1])
assert _remove_trivial_dims(CodegenArrayTensorProduct(a.T, b.T, c, d)) == (
CodegenArrayTensorProduct(a * b.T, c * d.T), [0, 2, 5, 7])
assert _remove_trivial_dims(CodegenArrayTensorProduct(a.T, I, b.T, c, d, I)) == (
CodegenArrayTensorProduct(a * b.T, c * d.T, I), [0, 2, 3, 4, 7, 9])
# Addition:
cg = CodegenArrayElementwiseAdd(CodegenArrayTensorProduct(a, b), CodegenArrayTensorProduct(c, d))
assert _remove_trivial_dims(cg) == (a * b.T + c * d.T, [1, 3])
# Permute Dims:
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(a, b), Permutation(3)(1, 2))
assert _remove_trivial_dims(cg) == (a * b.T, [2, 3])
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(a, I, b), Permutation(5)(1, 2, 3, 4))
assert _remove_trivial_dims(cg) == (a * b.T, [2, 3, 4, 5])
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(I, b, a), Permutation(5)(1, 2, 4, 5, 3))
assert _remove_trivial_dims(cg) == (b * a.T, [0, 1, 2, 3])
# Diagonal:
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, a), (1, 2))
assert _remove_trivial_dims(cg) == (cg, [])
# Contraction:
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, a), (1, 2))
assert _remove_trivial_dims(cg) == (cg, [])
def test_codegen_einsum():
if not tf:
skip("TensorFlow not installed")
graph = tf.Graph()
with graph.as_default():
session = tf.compat.v1.Session(graph=graph)
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
cg = CodegenArrayContraction.from_MatMul(M * N)
f = lambdify((M, N), cg, "tensorflow")
ma = tf.constant([[1, 2], [3, 4]])
mb = tf.constant([[1, -2], [-1, 3]])
y = session.run(f(ma, mb))
c = session.run(tf.matmul(ma, mb))
assert (y == c).all()
def test_permute_tensor_product():
cg1 = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q),
Permutation([2, 3, 1, 0, 5, 4, 6, 7]))
cg2 = CodegenArrayTensorProduct(N, CodegenArrayPermuteDims(M, [1, 0]),
CodegenArrayPermuteDims(P, [1, 0]), Q)
assert cg1 == cg2
# TODO: reverse operation starting with `CodegenArrayPermuteDims` and getting down to `bb`...
cg1 = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q),
Permutation([2, 3, 4, 5, 0, 1, 6, 7]))
cg2 = CodegenArrayTensorProduct(N, P, M, Q)
assert cg1 == cg2
cg1 = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q),
Permutation([2, 3, 4, 6, 5, 7, 0, 1]))
assert cg1.expr == CodegenArrayTensorProduct(N, P, Q, M)
assert cg1.permutation == Permutation([0, 1, 2, 4, 3, 5, 6, 7])
cg1 = CodegenArrayContraction(
CodegenArrayPermuteDims(CodegenArrayTensorProduct(N, Q, Q, M),
[2, 1, 5, 4, 0, 3, 6, 7]), [1, 2, 6])
cg2 = CodegenArrayPermuteDims(
CodegenArrayContraction(CodegenArrayTensorProduct(Q, Q, N, M),
(3, 5, 6)), [0, 2, 3, 1, 4])
assert cg1 == cg2
cg1 = CodegenArrayContraction(
CodegenArrayContraction(
CodegenArrayContraction(
CodegenArrayContraction(
CodegenArrayPermuteDims(
CodegenArrayTensorProduct(N, Q, Q, M),
[2, 1, 5, 4, 0, 3, 6, 7]), [1, 2, 6]), [1, 3, 4]),
[1]), [0])
cg2 = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, Q, Q),
(0, 3, 5), (1, 4, 7), (2, ), (6, ))
assert cg1 == cg2
def test_recognize_diagonalized_vectors():
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)
x = MatrixSymbol("x", k, 1)
I1 = Identity(1)
I = Identity(k)
# Check matrix recognition over trivial dimensions:
cg = CodegenArrayTensorProduct(a, b)
assert recognize_matrix_expression(cg) == a * b.T
cg = CodegenArrayTensorProduct(I1, a, b)
assert recognize_matrix_expression(cg) == a * I1 * b.T
# Recognize trace inside a tensor product:
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, B, C), (0, 3),
(1, 2))
assert recognize_matrix_expression(cg) == Trace(A * B) * C
# Transform diagonal operator to contraction:
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(A, a), (1, 2))
assert cg.transform_to_product() == CodegenArrayContraction(
CodegenArrayTensorProduct(A, DiagMatrix(a)), (1, 2))
assert recognize_matrix_expression(cg) == A * DiagMatrix(a)
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a, b), (0, 2))
assert cg.transform_to_product() == CodegenArrayContraction(
CodegenArrayTensorProduct(DiagMatrix(a), b), (0, 2))
assert recognize_matrix_expression(cg).doit() == DiagMatrix(a) * b
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(A, a), (0, 2))
assert cg.transform_to_product() == CodegenArrayContraction(
CodegenArrayTensorProduct(A, DiagMatrix(a)), (0, 2))
assert recognize_matrix_expression(cg) == A.T * DiagMatrix(a)
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(I, x, I1), (0, 2),
(3, 5))
assert cg.transform_to_product() == CodegenArrayContraction(
CodegenArrayTensorProduct(I, DiagMatrix(x), I1), (0, 2))
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(I, x, A, B), (1, 2),
(5, 6))
assert cg.transform_to_product() == CodegenArrayDiagonal(
CodegenArrayContraction(
CodegenArrayTensorProduct(I, DiagMatrix(x), A, B), (1, 2)), (3, 4))
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(x, I1), (1, 2))
assert isinstance(cg, CodegenArrayDiagonal)
assert cg.diagonal_indices == ((1, 2), )
assert recognize_matrix_expression(cg) == x
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(x, I), (0, 2))
assert cg.transform_to_product() == CodegenArrayContraction(
CodegenArrayTensorProduct(DiagMatrix(x), I), (0, 2))
assert recognize_matrix_expression(cg).doit() == DiagMatrix(x)
cg = CodegenArrayDiagonal(x, (1, ))
assert cg == x
# Ignore identity matrices with contractions:
cg = CodegenArrayContraction(CodegenArrayTensorProduct(I, A, I, I), (0, 2),
(1, 3), (5, 7))
assert cg.split_multiple_contractions() == cg
assert recognize_matrix_expression(cg) == Trace(A) * I
cg = CodegenArrayContraction(CodegenArrayTensorProduct(Trace(A) * I, I, I),
(1, 5), (3, 4))
assert cg.split_multiple_contractions() == cg
assert recognize_matrix_expression(cg).doit() == Trace(A) * I
# Add DiagMatrix when required:
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a), (1, 2))
assert cg.split_multiple_contractions() == cg
assert recognize_matrix_expression(cg) == A * a
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, B), (1, 2, 4))
assert cg.split_multiple_contractions() == CodegenArrayContraction(
CodegenArrayTensorProduct(A, DiagMatrix(a), B), (1, 2), (3, 4))
assert recognize_matrix_expression(cg) == A * DiagMatrix(a) * B
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, B), (0, 2, 4))
assert cg.split_multiple_contractions() == CodegenArrayContraction(
CodegenArrayTensorProduct(A, DiagMatrix(a), B), (0, 2), (3, 4))
assert recognize_matrix_expression(cg) == A.T * DiagMatrix(a) * B
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, b, a.T, B),
(0, 2, 4, 7, 9))
assert cg.split_multiple_contractions() == CodegenArrayContraction(
CodegenArrayTensorProduct(A, DiagMatrix(a), DiagMatrix(b),
DiagMatrix(a), B), (0, 2), (3, 4), (5, 7),
(6, 9))
assert recognize_matrix_expression(
cg).doit() == A.T * DiagMatrix(a) * DiagMatrix(b) * DiagMatrix(a) * B.T
cg = CodegenArrayContraction(CodegenArrayTensorProduct(I1, I1, I1),
(1, 2, 4))
assert cg.split_multiple_contractions() == CodegenArrayContraction(
CodegenArrayTensorProduct(I1, I1, I1), (1, 2), (3, 4))
assert recognize_matrix_expression(cg).doit() == Identity(1)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(I, I, I, I, A),
(1, 2, 8), (5, 6, 9))
assert recognize_matrix_expression(
cg.split_multiple_contractions()).doit() == A
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, C, a, B),
(1, 2, 4), (5, 6, 8))
assert cg.split_multiple_contractions() == CodegenArrayContraction(
CodegenArrayTensorProduct(A, DiagMatrix(a), C, DiagMatrix(a), B),
(1, 2), (3, 4), (5, 6), (7, 8))
assert recognize_matrix_expression(
cg) == A * DiagMatrix(a) * C * DiagMatrix(a) * B
cg = CodegenArrayContraction(
CodegenArrayTensorProduct(a, I1, b, I1, (a.T * b).applyfunc(cos)),
(1, 2, 8), (5, 6, 9))
assert cg.split_multiple_contractions().dummy_eq(
CodegenArrayContraction(
CodegenArrayTensorProduct(a, I1, b, I1, (a.T * b).applyfunc(cos)),
(1, 2), (3, 8), (5, 6), (7, 9)))
assert recognize_matrix_expression(cg).dummy_eq(
MatMul(a, I1, (a.T * b).applyfunc(cos), Transpose(I1), b.T))
cg = CodegenArrayContraction(
CodegenArrayTensorProduct(A.T, a, b, b.T, (A * X * b).applyfunc(cos)),
(1, 2, 8), (5, 6, 9))
assert cg.split_multiple_contractions().dummy_eq(
CodegenArrayContraction(
CodegenArrayTensorProduct(A.T, DiagMatrix(a), b, b.T,
(A * X * b).applyfunc(cos)), (1, 2),
(3, 8), (5, 6, 9)))
# assert recognize_matrix_expression(cg)
# Check no overlap of lines:
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, C, a, B),
(1, 2, 4), (5, 6, 8), (3, 7))
assert cg.split_multiple_contractions() == cg
cg = CodegenArrayContraction(CodegenArrayTensorProduct(a, b, A), (0, 2, 4),
(1, 3))
assert cg.split_multiple_contractions() == cg
def test_codegen_array_contraction_construction():
cg = CodegenArrayContraction(A)
assert cg == A
s = Sum(A[i] * B[i], (i, 0, 3))
cg = parse_indexed_expression(s)
assert cg == CodegenArrayContraction(CodegenArrayTensorProduct(A, B),
(0, 1))
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, B), (1, 0))
assert cg == CodegenArrayContraction(CodegenArrayTensorProduct(A, B),
(0, 1))
expr = M * N
result = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2))
assert CodegenArrayContraction.from_MatMul(expr) == result
elem = expr[i, j]
assert parse_indexed_expression(elem) == result
expr = M * N * M
result = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, M),
(1, 2), (3, 4))
assert CodegenArrayContraction.from_MatMul(expr) == result
elem = expr[i, j]
result = CodegenArrayContraction(CodegenArrayTensorProduct(M, M, N),
(1, 4), (2, 5))
cg = parse_indexed_expression(elem)
cg = cg.sort_args_by_name()
assert cg == result
def test_codegen_array_flatten():
# Flatten nested CodegenArrayTensorProduct objects:
expr1 = CodegenArrayTensorProduct(M, N)
expr2 = CodegenArrayTensorProduct(P, Q)
expr = CodegenArrayTensorProduct(expr1, expr2)
assert expr == CodegenArrayTensorProduct(M, N, P, Q)
assert expr.args == (M, N, P, Q)
# Flatten mixed CodegenArrayTensorProduct and CodegenArrayContraction objects:
cg1 = CodegenArrayContraction(expr1, (1, 2))
cg2 = CodegenArrayContraction(expr2, (0, 3))
expr = CodegenArrayTensorProduct(cg1, cg2)
assert expr == CodegenArrayContraction(
CodegenArrayTensorProduct(M, N, P, Q), (1, 2), (4, 7))
expr = CodegenArrayTensorProduct(M, cg1)
assert expr == CodegenArrayContraction(CodegenArrayTensorProduct(M, M, N),
(3, 4))
# Flatten nested CodegenArrayContraction objects:
cgnested = CodegenArrayContraction(cg1, (0, 1))
assert cgnested == CodegenArrayContraction(CodegenArrayTensorProduct(M, N),
(0, 3), (1, 2))
cgnested = CodegenArrayContraction(CodegenArrayTensorProduct(cg1, cg2),
(0, 3))
assert cgnested == CodegenArrayContraction(
CodegenArrayTensorProduct(M, N, P, Q), (0, 6), (1, 2), (4, 7))
cg3 = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q),
(1, 3), (2, 4))
cgnested = CodegenArrayContraction(cg3, (0, 1))
assert cgnested == CodegenArrayContraction(
CodegenArrayTensorProduct(M, N, P, Q), (0, 5), (1, 3), (2, 4))
cgnested = CodegenArrayContraction(cg3, (0, 3), (1, 2))
assert cgnested == CodegenArrayContraction(
CodegenArrayTensorProduct(M, N, P, Q), (0, 7), (1, 3), (2, 4), (5, 6))
cg4 = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q),
(1, 5), (3, 7))
cgnested = CodegenArrayContraction(cg4, (0, 1))
assert cgnested == CodegenArrayContraction(
CodegenArrayTensorProduct(M, N, P, Q), (0, 2), (1, 5), (3, 7))
cgnested = CodegenArrayContraction(cg4, (0, 1), (2, 3))
assert cgnested == CodegenArrayContraction(
CodegenArrayTensorProduct(M, N, P, Q), (0, 2), (1, 5), (3, 7), (4, 6))
cg = CodegenArrayDiagonal(cg4)
assert cg == cg4
assert isinstance(cg, type(cg4))
# Flatten nested CodegenArrayDiagonal objects:
cg1 = CodegenArrayDiagonal(expr1, (1, 2))
cg2 = CodegenArrayDiagonal(expr2, (0, 3))
cg3 = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P, Q), (1, 3),
(2, 4))
cg4 = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P, Q), (1, 5),
(3, 7))
cgnested = CodegenArrayDiagonal(cg1, (0, 1))
assert cgnested == CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N),
(1, 2), (0, 3))
cgnested = CodegenArrayDiagonal(cg3, (1, 2))
assert cgnested == CodegenArrayDiagonal(
CodegenArrayTensorProduct(M, N, P, Q), (1, 3), (2, 4), (5, 6))
cgnested = CodegenArrayDiagonal(cg4, (1, 2))
assert cgnested == CodegenArrayDiagonal(
CodegenArrayTensorProduct(M, N, P, Q), (1, 5), (3, 7), (2, 4))
def tensorcontraction(array, *contraction_axes):
"""
Contraction of an array-like object on the specified axes.
Examples
========
>>> from sympy import Array, tensorcontraction
>>> from sympy import Matrix, eye
>>> tensorcontraction(eye(3), (0, 1))
3
>>> A = Array(range(18), (3, 2, 3))
>>> A
[[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]], [[12, 13, 14], [15, 16, 17]]]
>>> tensorcontraction(A, (0, 2))
[21, 30]
Matrix multiplication may be emulated with a proper combination of
``tensorcontraction`` and ``tensorproduct``
>>> from sympy import tensorproduct
>>> from sympy.abc import a,b,c,d,e,f,g,h
>>> m1 = Matrix([[a, b], [c, d]])
>>> m2 = Matrix([[e, f], [g, h]])
>>> p = tensorproduct(m1, m2)
>>> p
[[[[a*e, a*f], [a*g, a*h]], [[b*e, b*f], [b*g, b*h]]], [[[c*e, c*f], [c*g, c*h]], [[d*e, d*f], [d*g, d*h]]]]
>>> tensorcontraction(p, (1, 2))
[[a*e + b*g, a*f + b*h], [c*e + d*g, c*f + d*h]]
>>> m1*m2
Matrix([
[a*e + b*g, a*f + b*h],
[c*e + d*g, c*f + d*h]])
"""
from sympy.codegen.array_utils import _CodegenArrayAbstract, CodegenArrayContraction
from sympy.tensor.array.expressions.array_expressions import _ArrayExpr
if isinstance(array, (_ArrayExpr, _CodegenArrayAbstract)):
return CodegenArrayContraction(array, *contraction_axes)
array, remaining_indices, remaining_shape, summed_deltas = _util_contraction_diagonal(array, *contraction_axes)
# Compute the contracted array:
#
# 1. external for loops on all uncontracted indices.
# Uncontracted indices are determined by the combinatorial product of
# the absolute positions of the remaining indices.
# 2. internal loop on all contracted indices.
# It sums the values of the absolute contracted index and the absolute
# uncontracted index for the external loop.
contracted_array = []
for icontrib in itertools.product(*remaining_indices):
index_base_position = sum(icontrib)
isum = S.Zero
for sum_to_index in itertools.product(*summed_deltas):
idx = array._get_tuple_index(index_base_position + sum(sum_to_index))
isum += array[idx]
contracted_array.append(isum)
if len(remaining_indices) == 0:
assert len(contracted_array) == 1
return contracted_array[0]
return type(array)(contracted_array, remaining_shape)
def test_push_indices_up_and_down():
indices = list(range(10))
contraction_indices = [(0, 6), (2, 8)]
assert CodegenArrayContraction._push_indices_down(contraction_indices,
indices) == (
1,
3,
4,
5,
7,
9,
10,
11,
12,
13,
)
assert CodegenArrayContraction._push_indices_up(contraction_indices,
indices) == (
None,
0,
None,
1,
2,
3,
None,
4,
None,
5,
)
assert CodegenArrayDiagonal._push_indices_down(contraction_indices,
indices) == (
0,
1,
2,
3,
4,
5,
7,
9,
10,
11,
)
assert CodegenArrayDiagonal._push_indices_up(contraction_indices,
indices) == (
0,
1,
2,
3,
4,
5,
None,
6,
None,
7,
)
contraction_indices = [(1, 2), (7, 8)]
assert CodegenArrayContraction._push_indices_down(contraction_indices,
indices) == (
0,
3,
4,
5,
6,
9,
10,
11,
12,
13,
)
assert CodegenArrayContraction._push_indices_up(contraction_indices,
indices) == (
0,
None,
None,
1,
2,
3,
4,
None,
None,
5,
)
assert CodegenArrayContraction._push_indices_down(contraction_indices,
indices) == (
0,
3,
4,
5,
6,
9,
10,
11,
12,
13,
)
assert CodegenArrayDiagonal._push_indices_up(contraction_indices,
indices) == (
0,
1,
None,
2,
3,
4,
5,
6,
None,
7,
)
def test_codegen_array_contraction_indices_types():
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 1))
indtup = cg._get_contraction_tuples()
assert indtup == [[(0, 0), (0, 1)]]
assert cg._contraction_tuples_to_contraction_indices(cg.expr,
indtup) == [(0, 1)]
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2))
indtup = cg._get_contraction_tuples()
assert indtup == [[(0, 1), (1, 0)]]
assert cg._contraction_tuples_to_contraction_indices(cg.expr,
indtup) == [(1, 2)]
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, M, N), (1, 4),
(2, 5))
indtup = cg._get_contraction_tuples()
assert indtup == [[(0, 1), (2, 0)], [(1, 0), (2, 1)]]
assert cg._contraction_tuples_to_contraction_indices(cg.expr,
indtup) == [(1, 4),
(2, 5)]
def test_parsing_of_matrix_expressions():
expr = M*N
assert parse_matrix_expression(expr) == CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2))
expr = Transpose(M)
assert parse_matrix_expression(expr) == CodegenArrayPermuteDims(M, [1, 0])
expr = M*Transpose(N)
assert parse_matrix_expression(expr) == CodegenArrayContraction(CodegenArrayTensorProduct(M, CodegenArrayPermuteDims(N, [1, 0])), (1, 2))
expr = 3*M*N
res = parse_matrix_expression(expr)
rexpr = recognize_matrix_expression(res)
assert expr == rexpr
expr = 3*M + N*M.T*M + 4*k*N
res = parse_matrix_expression(expr)
rexpr = recognize_matrix_expression(res)
assert expr == rexpr
expr = Inverse(M)*N
rexpr = recognize_matrix_expression(parse_matrix_expression(expr))
assert expr == rexpr
expr = M**2
rexpr = recognize_matrix_expression(parse_matrix_expression(expr))
assert expr == rexpr
expr = M*(2*N + 3*M)
res = parse_matrix_expression(expr)
rexpr = recognize_matrix_expression(res)
assert expr == rexpr
expr = Trace(M)
result = CodegenArrayContraction(M, (0, 1))
assert parse_matrix_expression(expr) == result
expr = 3*Trace(M)
result = CodegenArrayContraction(CodegenArrayTensorProduct(3, M), (0, 1))
assert parse_matrix_expression(expr) == result
expr = 3*Trace(Trace(M) * M)
result = CodegenArrayContraction(CodegenArrayTensorProduct(3, M, M), (0, 1), (2, 3))
assert parse_matrix_expression(expr) == result
expr = 3*Trace(M)**2
result = CodegenArrayContraction(CodegenArrayTensorProduct(3, M, M), (0, 1), (2, 3))
assert parse_matrix_expression(expr) == result
expr = HadamardProduct(M, N)
result = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (0, 2), (1, 3))
assert parse_matrix_expression(expr) == result
expr = HadamardPower(M, 2)
result = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, M), (0, 2), (1, 3))
assert parse_matrix_expression(expr) == result
expr = M**2
assert isinstance(expr, MatPow)
assert parse_matrix_expression(expr) == CodegenArrayContraction(CodegenArrayTensorProduct(M, M), (1, 2))
def test_recognize_diagonalized_vectors():
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)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a), (1, 2))
assert cg.split_multiple_contractions() == cg
assert recognize_matrix_expression(cg) == A * a
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, B), (1, 2, 4))
assert cg.split_multiple_contractions() == CodegenArrayContraction(
CodegenArrayTensorProduct(A, DiagonalizeVector(a), B), (1, 2), (3, 4))
assert recognize_matrix_expression(cg) == A * DiagonalizeVector(a) * B
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, B), (0, 2, 4))
assert cg.split_multiple_contractions() == CodegenArrayContraction(
CodegenArrayTensorProduct(A, DiagonalizeVector(a), B), (0, 2), (3, 4))
assert recognize_matrix_expression(cg) == A.T * DiagonalizeVector(a) * B
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, b, a.T, B),
(0, 2, 4, 7, 9))
assert cg.split_multiple_contractions() == CodegenArrayContraction(
CodegenArrayTensorProduct(A,
DiagonalizeVector(a), DiagonalizeVector(b),
DiagonalizeVector(a), B), (0, 2), (3, 4),
(5, 7), (6, 9))
assert recognize_matrix_expression(cg).doit() == A.T * DiagonalizeVector(
a) * DiagonalizeVector(b) * DiagonalizeVector(a) * B.T
I1 = Identity(1)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(I1, I1, I1),
(1, 2, 4))
assert cg.split_multiple_contractions() == CodegenArrayContraction(
CodegenArrayTensorProduct(I1, I1, I1), (1, 2), (3, 4))
assert recognize_matrix_expression(cg).doit() == Identity(1)
I = Identity(k)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(I, I, I, I, A),
(1, 2, 8), (5, 6, 9))
assert recognize_matrix_expression(
cg.split_multiple_contractions()).doit() == A
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, C, a, B),
(1, 2, 4), (5, 6, 8))
assert cg.split_multiple_contractions() == CodegenArrayContraction(
CodegenArrayTensorProduct(A, DiagonalizeVector(a), C,
DiagonalizeVector(a), B), (1, 2), (3, 4),
(5, 6), (7, 8))
assert recognize_matrix_expression(
cg) == A * DiagonalizeVector(a) * C * DiagonalizeVector(a) * B
cg = CodegenArrayContraction(
CodegenArrayTensorProduct(a, I1, b, I1, (a.T * b).applyfunc(cos)),
(1, 2, 8), (5, 6, 9))
assert cg.split_multiple_contractions() == CodegenArrayContraction(
CodegenArrayTensorProduct(a, I1, b, I1, (a.T * b).applyfunc(cos)),
(1, 2), (3, 8), (5, 6), (7, 9))
assert recognize_matrix_expression(cg) == MatMul(a, I1,
(a.T * b).applyfunc(cos),
Transpose(I1), b.T)
# Check no overlap of lines:
cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, C, a, B),
(1, 2, 4), (5, 6, 8), (3, 7))
raises(ValueError, lambda: cg.split_multiple_contractions())
cg = CodegenArrayContraction(CodegenArrayTensorProduct(a, b, A), (0, 2, 4),
(1, 3))
raises(ValueError, lambda: cg.split_multiple_contractions())
