def test_codegen_permutedims_sink():
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [0, 1, 3, 2])
sunk = cg.nest_permutation()
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])
sunk = cg.nest_permutation()
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])
sunk = cg.nest_permutation()
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])
sunk = cg.nest_permutation()
assert sunk == CodegenArrayContraction(CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [[0, 3]]), (1, 2))
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [1, 0, 3, 2])
sunk = cg.nest_permutation()
assert sunk == CodegenArrayTensorProduct(CodegenArrayPermuteDims(M, [1, 0]), CodegenArrayPermuteDims(N, [1, 0]))
cg = CodegenArrayPermuteDims(CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P), (1, 2), (3, 4)), [1, 0])
sunk = cg.nest_permutation()
assert sunk == CodegenArrayContraction(CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P), [[0, 5]]), (1, 2), (3, 4))
sunk2 = sunk.expr.nest_permutation()
def test_codegen_permutedims_sink():
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [0, 1, 3, 2])
sunk = cg.nest_permutation()
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])
sunk = cg.nest_permutation()
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])
sunk = cg.nest_permutation()
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])
sunk = cg.nest_permutation()
assert sunk == CodegenArrayContraction(CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [[0, 3]]), (1, 2))
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [1, 0, 3, 2])
sunk = cg.nest_permutation()
assert sunk == CodegenArrayTensorProduct(CodegenArrayPermuteDims(M, [1, 0]), CodegenArrayPermuteDims(N, [1, 0]))
cg = CodegenArrayPermuteDims(CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P), (1, 2), (3, 4)), [1, 0])
sunk = cg.nest_permutation()
assert sunk == CodegenArrayContraction(CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P), [[0, 5]]), (1, 2), (3, 4))
sunk2 = sunk.expr.nest_permutation()
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) == CodegenArrayTensorProduct(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) == CodegenArrayTensorProduct(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) == CodegenArrayTensorProduct(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_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_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_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) == [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_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) == CodegenArrayTensorProduct(
M, N.T)
assert recognize_matrix_expression(cg2) == CodegenArrayTensorProduct(
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) == CodegenArrayTensorProduct(
M * P.T * Trace(N), Q.T)
assert recognize_matrix_expression(cg2) == CodegenArrayTensorProduct(
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_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_recognize_matrix_expression():
expr = CodegenArrayElementwiseAdd(M, CodegenArrayPermuteDims(M, [1, 0]))
assert recognize_matrix_expression(expr) == M + Transpose(M)
expr = M[i,j] + N[i,j]
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == M + N
expr = M[i,j] + N[j,i]
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == M + N.T
expr = M[i,j]*N[k,l] + N[i,j]*M[k,l]
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(M, N),
CodegenArrayTensorProduct(N, M))
expr = (M*N*P)[i, j]
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == M*N*P
expr = Sum(M[i,j]*(N*P)[j,m], (j, 0, k-1))
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(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 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == M*P*N + M*P.T*N + N*P*M + N*P.T*M
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_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_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_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 _matrix_derivative(expr, x):
from sympy import Derivative
lines = expr._eval_derivative_matrix_lines(x)
parts = [i.build() for i in lines]
from sympy.codegen.array_utils import recognize_matrix_expression
parts = [[recognize_matrix_expression(j).doit() for j in i] for i in parts]
def _get_shape(elem):
if isinstance(elem, MatrixExpr):
return elem.shape
return (1, 1)
def get_rank(parts):
return sum([j not in (1, None) for i in parts for j in _get_shape(i)])
ranks = [get_rank(i) for i in parts]
rank = ranks[0]
def contract_one_dims(parts):
if len(parts) == 1:
return parts[0]
else:
p1, p2 = parts[:2]
if p2.is_Matrix:
p2 = p2.T
if p1 == Identity(1):
pbase = p2
elif p2 == Identity(1):
pbase = p1
else:
pbase = p1 * p2
if len(parts) == 2:
return pbase
else: # len(parts) > 2
if pbase.is_Matrix:
raise ValueError("")
return pbase * Mul.fromiter(parts[2:])
if rank <= 2:
return Add.fromiter([contract_one_dims(i) for i in parts])
return Derivative(expr, x)
def _matrix_derivative(expr, x):
from sympy import Derivative
lines = expr._eval_derivative_matrix_lines(x)
parts = [i.build() for i in lines]
from sympy.codegen.array_utils import recognize_matrix_expression
parts = [[recognize_matrix_expression(j).doit() for j in i] for i in parts]
def _get_shape(elem):
if isinstance(elem, MatrixExpr):
return elem.shape
return (1, 1)
def get_rank(parts):
return sum([j not in (1, None) for i in parts for j in _get_shape(i)])
ranks = [get_rank(i) for i in parts]
rank = ranks[0]
def contract_one_dims(parts):
if len(parts) == 1:
return parts[0]
else:
p1, p2 = parts[:2]
if p2.is_Matrix:
p2 = p2.T
pbase = p1*p2
if len(parts) == 2:
return pbase
else: # len(parts) > 2
if pbase.is_Matrix:
raise ValueError("")
return pbase*Mul.fromiter(parts[2:])
if rank <= 2:
return Add.fromiter([contract_one_dims(i) for i in parts])
return Derivative(expr, x)
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, DiagonalizeVector(a)), (1, 2))
assert recognize_matrix_expression(cg) == A*DiagonalizeVector(a)
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a, b), (0, 2))
assert cg.transform_to_product() == CodegenArrayContraction(CodegenArrayTensorProduct(DiagonalizeVector(a), b), (0, 2))
assert recognize_matrix_expression(cg).doit() == DiagonalizeVector(a)*b
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(A, a), (0, 2))
assert cg.transform_to_product() == CodegenArrayContraction(CodegenArrayTensorProduct(A, DiagonalizeVector(a)), (0, 2))
assert recognize_matrix_expression(cg) == A.T*DiagonalizeVector(a)
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(I, x, I1), (0, 2), (3, 5))
assert cg.transform_to_product() == CodegenArrayContraction(CodegenArrayTensorProduct(I, DiagonalizeVector(x), I1), (0, 2))
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(I, x, A, B), (1, 2), (5, 6))
assert cg.transform_to_product() == CodegenArrayDiagonal(CodegenArrayContraction(CodegenArrayTensorProduct(I, DiagonalizeVector(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(DiagonalizeVector(x), I), (0, 2))
assert recognize_matrix_expression(cg).doit() == DiagonalizeVector(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 DiagonalizeVector 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, 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
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, 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)
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() == CodegenArrayContraction(
CodegenArrayTensorProduct(A.T, DiagonalizeVector(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_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_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_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 matrix_derive(expr, x):
from sympy.codegen.array_utils import parse_matrix_expression, recognize_matrix_expression
ce = parse_matrix_expression(expr)
dce = array_derive(ce, x)
return recognize_matrix_expression(dce).doit()
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())