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 parse_matrix_expression(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 parse_matrix_expression(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_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_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_codegen_einsum():
if not np:
skip("NumPy not installed")
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
cg = parse_matrix_expression(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_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 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 = parse_matrix_expression(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_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 _(expr: MatrixExpr, x: Expr):
cg = parse_matrix_expression(expr)
return array_derive(cg, x)
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()
