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_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 = CodegenArrayTensorProduct(M, N)
f = lambdify((M, N), cg, 'numpy')
assert (f(ma, mb) == np.einsum(ma, [0, 1], mb, [2, 3])).all()
cg = CodegenArrayElementwiseAdd(M, N)
f = lambdify((M, N), cg, 'numpy')
assert (f(ma, mb) == ma + mb).all()
cg = CodegenArrayElementwiseAdd(M, N, P)
f = lambdify((M, N, P), cg, 'numpy')
assert (f(ma, mb, mc) == ma + mb + mc).all()
cg = CodegenArrayElementwiseAdd(M, N, P, Q)
f = lambdify((M, N, P, Q), cg, 'numpy')
assert (f(ma, mb, mc, md) == ma + mb + mc + md).all()
cg = CodegenArrayPermuteDims(M, [1, 0])
f = lambdify((M, ), cg, 'numpy')
assert (f(ma) == ma.T).all()
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(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 = CodegenArrayDiagonal(CodegenArrayTensorProduct(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_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_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_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_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 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_codegen_array_diagonal():
cg = CodegenArrayDiagonal(M, (1, 0))
assert cg == CodegenArrayDiagonal(M, (0, 1))
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P), (4, 1), (2, 0))
assert cg == CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P), (1, 4), (0, 2))
def test_codegen_extra():
if not tf:
skip("TensorFlow not installed")
graph = tf.Graph()
with graph.as_default():
session = tf.compat.v1.Session()
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
P = MatrixSymbol("P", 2, 2)
Q = MatrixSymbol("Q", 2, 2)
ma = tf.constant([[1, 2], [3, 4]])
mb = tf.constant([[1, -2], [-1, 3]])
mc = tf.constant([[2, 0], [1, 2]])
md = tf.constant([[1, -1], [4, 7]])
cg = CodegenArrayTensorProduct(M, N)
assert tensorflow_code(cg) == \
'tensorflow.linalg.einsum("ab,cd", M, N)'
f = lambdify((M, N), cg, 'tensorflow')
y = session.run(f(ma, mb))
c = session.run(tf.einsum("ij,kl", ma, mb))
assert (y == c).all()
cg = CodegenArrayElementwiseAdd(M, N)
assert tensorflow_code(cg) == 'tensorflow.math.add(M, N)'
f = lambdify((M, N), cg, 'tensorflow')
y = session.run(f(ma, mb))
c = session.run(ma + mb)
assert (y == c).all()
cg = CodegenArrayElementwiseAdd(M, N, P)
assert tensorflow_code(cg) == \
'tensorflow.math.add(tensorflow.math.add(M, N), P)'
f = lambdify((M, N, P), cg, 'tensorflow')
y = session.run(f(ma, mb, mc))
c = session.run(ma + mb + mc)
assert (y == c).all()
cg = CodegenArrayElementwiseAdd(M, N, P, Q)
assert tensorflow_code(cg) == \
'tensorflow.math.add(' \
'tensorflow.math.add(tensorflow.math.add(M, N), P), Q)'
f = lambdify((M, N, P, Q), cg, 'tensorflow')
y = session.run(f(ma, mb, mc, md))
c = session.run(ma + mb + mc + md)
assert (y == c).all()
cg = CodegenArrayPermuteDims(M, [1, 0])
assert tensorflow_code(cg) == 'tensorflow.transpose(M, [1, 0])'
f = lambdify((M, ), cg, 'tensorflow')
y = session.run(f(ma))
c = session.run(tf.transpose(ma))
assert (y == c).all()
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N),
[1, 2, 3, 0])
assert tensorflow_code(cg) == \
'tensorflow.transpose(' \
'tensorflow.linalg.einsum("ab,cd", M, N), [1, 2, 3, 0])'
f = lambdify((M, N), cg, 'tensorflow')
y = session.run(f(ma, mb))
c = session.run(tf.transpose(tf.einsum("ab,cd", ma, mb), [1, 2, 3, 0]))
assert (y == c).all()
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (1, 2))
assert tensorflow_code(cg) == \
'tensorflow.linalg.einsum("ab,bc->acb", M, N)'
f = lambdify((M, N), cg, 'tensorflow')
y = session.run(f(ma, mb))
c = session.run(tf.einsum("ab,bc->acb", ma, mb))
assert (y == c).all()
def tensordiagonal(array, *diagonal_axes):
"""
Diagonalization of an array-like object on the specified axes.
This is equivalent to multiplying the expression by Kronecker deltas
uniting the axes.
The diagonal indices are put at the end of the axes.
Examples
========
``tensordiagonal`` acting on a 2-dimensional array by axes 0 and 1 is
equivalent to the diagonal of the matrix:
>>> from sympy import Array, tensordiagonal
>>> from sympy import Matrix, eye
>>> tensordiagonal(eye(3), (0, 1))
[1, 1, 1]
>>> from sympy.abc import a,b,c,d
>>> m1 = Matrix([[a, b], [c, d]])
>>> tensordiagonal(m1, [0, 1])
[a, d]
In case of higher dimensional arrays, the diagonalized out dimensions
are appended removed and appended as a single dimension at the end:
>>> 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]]]
>>> tensordiagonal(A, (0, 2))
[[0, 7, 14], [3, 10, 17]]
>>> from sympy import permutedims
>>> tensordiagonal(A, (0, 2)) == permutedims(Array([A[0, :, 0], A[1, :, 1], A[2, :, 2]]), [1, 0])
True
"""
if any([len(i) <= 1 for i in diagonal_axes]):
raise ValueError("need at least two axes to diagonalize")
from sympy.tensor.array.expressions.array_expressions import _ArrayExpr
from sympy.codegen.array_utils import _CodegenArrayAbstract, CodegenArrayDiagonal
from sympy import MatrixSymbol
if isinstance(array, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)):
return CodegenArrayDiagonal(array, *diagonal_axes)
array, remaining_indices, remaining_shape, diagonal_deltas = _util_contraction_diagonal(
array, *diagonal_axes)
# Compute the diagonalized array:
#
# 1. external for loops on all undiagonalized indices.
# Undiagonalized indices are determined by the combinatorial product of
# the absolute positions of the remaining indices.
# 2. internal loop on all diagonal indices.
# It appends the values of the absolute diagonalized index and the absolute
# undiagonalized index for the external loop.
diagonalized_array = []
diagonal_shape = [len(i) for i in diagonal_deltas]
for icontrib in itertools.product(*remaining_indices):
index_base_position = sum(icontrib)
isum = []
for sum_to_index in itertools.product(*diagonal_deltas):
idx = array._get_tuple_index(index_base_position +
sum(sum_to_index))
isum.append(array[idx])
isum = type(array)(isum).reshape(*diagonal_shape)
diagonalized_array.append(isum)
return type(array)(diagonalized_array, remaining_shape + diagonal_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_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(i, j) * KroneckerDelta(j, k) * M[i, l]
assert _codegen_array_parse(expr) == (M, ({i, j, k}, l))
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 = (KroneckerDelta(j, m) * KroneckerDelta(m, 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, m, k})),
)
expr = (KroneckerDelta(i, j) * KroneckerDelta(j, k) *
KroneckerDelta(k, m) * M[i, 0] * KroneckerDelta(m, n))
assert _codegen_array_parse(expr) == (M, ({i, j, k, m, n}, 0))
expr = M[i, i]
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(M,
(0, 1)), (i, ))
def _(expr: CodegenArrayDiagonal, x: Expr):
dsubexpr = array_derive(expr.expr, x)
rank_x = len(get_shape(x))
diag_indices = [[j + rank_x for j in i] for i in expr.diagonal_indices]
return CodegenArrayDiagonal(dsubexpr, *diag_indices)
def test_codegen_extra():
if not torch:
skip("Torch not installed")
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
P = MatrixSymbol("P", 2, 2)
Q = MatrixSymbol("Q", 2, 2)
ma = torch.tensor([[1, 2], [3, 4]])
mb = torch.tensor([[1, -2], [-1, 3]])
mc = torch.tensor([[2, 0], [1, 2]])
md = torch.tensor([[1, -1], [4, 7]])
cg = CodegenArrayTensorProduct(M, N)
assert torch_code(cg) == 'torch.einsum("ab,cd", [M, N])'
f = lambdify((M, N), cg, 'torch')
y = f(ma, mb)
c = torch.einsum("ij,kl", ma, mb)
assert (y == c).all()
cg = CodegenArrayElementwiseAdd(M, N)
assert torch_code(cg) == 'torch.add(M, N)'
f = lambdify((M, N), cg, 'torch')
y = f(ma, mb)
c = ma + mb
assert (y == c).all()
cg = CodegenArrayElementwiseAdd(M, N, P)
assert torch_code(cg) == 'torch.add(torch.add(M, N), P)'
f = lambdify((M, N, P), cg, 'torch')
y = f(ma, mb, mc)
c = ma + mb + mc
assert (y == c).all()
cg = CodegenArrayElementwiseAdd(M, N, P, Q)
assert torch_code(cg) == 'torch.add(torch.add(torch.add(M, N), P), Q)'
f = lambdify((M, N, P, Q), cg, 'torch')
y = f(ma, mb, mc, md)
c = ma + mb + mc + md
assert (y == c).all()
cg = CodegenArrayPermuteDims(M, [1, 0])
assert torch_code(cg) == 'M.permute(1, 0)'
f = lambdify((M, ), cg, 'torch')
y = f(ma)
c = ma.T
assert (y == c).all()
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [1, 2, 3, 0])
assert torch_code(
cg) == 'torch.einsum("ab,cd", [M, N]).permute(1, 2, 3, 0)'
f = lambdify((M, N), cg, 'torch')
y = f(ma, mb)
c = torch.einsum("ab,cd", ma, mb).permute(1, 2, 3, 0)
assert (y == c).all()
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (1, 2))
assert torch_code(cg) == 'torch.einsum("ab,bc->acb", [M, N])'
f = lambdify((M, N), cg, 'torch')
y = f(ma, mb)
c = torch.einsum("ab,bc->acb", ma, mb)
assert (y == c).all()
