def _(expr: ArrayContraction):
expr = expr.flatten_contraction_of_diagonal()
expr = expr.split_multiple_contractions()
expr = identify_hadamard_products(expr)
subexpr = expr.expr
contraction_indices: Tuple[Tuple[int]] = expr.contraction_indices
if isinstance(subexpr, ArrayTensorProduct):
newexpr = ArrayContraction(_array2matrix(subexpr),
*contraction_indices)
contraction_indices = newexpr.contraction_indices
if any(i > 2 for i in newexpr.subranks):
addends = ArrayAdd(*[
_a2m_tensor_product(*j) for j in itertools.product(*[
i.args if isinstance(i, ArrayAdd) else [i]
for i in expr.expr.args
])
])
newexpr = ArrayContraction(addends, *contraction_indices)
if isinstance(newexpr, ArrayAdd):
ret = _array2matrix(newexpr)
return ret
assert isinstance(newexpr, ArrayContraction)
ret = _support_function_tp1_recognize(contraction_indices,
list(newexpr.expr.args))
return ret
elif not isinstance(subexpr, _CodegenArrayAbstract):
ret = _array2matrix(subexpr)
if isinstance(ret, MatrixExpr):
assert expr.contraction_indices == ((0, 1), )
return _a2m_trace(ret)
else:
return ArrayContraction(ret, *expr.contraction_indices)
def test_arrayexpr_split_multiple_contractions():
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)
cg = ArrayContraction(
ArrayTensorProduct(A.T, a, b, b.T, (A * X * b).applyfunc(cos)),
(1, 2, 8), (5, 6, 9))
assert cg.split_multiple_contractions().dummy_eq(
ArrayContraction(
ArrayTensorProduct(DiagMatrix(a), (A * X * b).applyfunc(cos), A.T,
b, b.T), (0, 2), (1, 5), (3, 7, 8)))
# assert recognize_matrix_expression(cg)
# Check no overlap of lines:
cg = ArrayContraction(ArrayTensorProduct(A, a, C, a, B), (1, 2, 4),
(5, 6, 8), (3, 7))
assert cg.split_multiple_contractions() == cg
cg = ArrayContraction(ArrayTensorProduct(a, b, A), (0, 2, 4), (1, 3))
assert cg.split_multiple_contractions() == cg
def _(expr: ArrayContraction):
rank1 = get_rank(expr)
expr, removed1 = remove_identity_matrices(expr)
if not isinstance(expr, ArrayContraction):
expr2, removed2 = _remove_trivial_dims(expr)
return expr2, _combine_removed(rank1, removed1, removed2)
newexpr, removed2 = _remove_trivial_dims(expr.expr)
shifts = list(
accumulate(
[1 if i in removed2 else 0 for i in range(get_rank(expr.expr))]))
new_contraction_indices = [
tuple(j for j in i if j not in removed2)
for i in expr.contraction_indices
]
# Remove possible empty tuples "()":
new_contraction_indices = [
i for i in new_contraction_indices if len(i) > 0
]
contraction_indices_flat = [j for i in expr.contraction_indices for j in i]
removed2 = [i for i in removed2 if i not in contraction_indices_flat]
new_contraction_indices = [
tuple(j - shifts[j] for j in i) for i in new_contraction_indices
]
# Shift removed2:
removed2 = ArrayContraction._push_indices_up(expr.contraction_indices,
removed2)
removed = _combine_removed(rank1, removed1, removed2)
return ArrayContraction(newexpr, *new_contraction_indices), list(removed)
def test_array_expr_construction_with_functions():
tp = tensorproduct(M, N)
assert tp == ArrayTensorProduct(M, N)
expr = tensorproduct(A, eye(2))
assert expr == ArrayTensorProduct(A, eye(2))
# Contraction:
expr = tensorcontraction(M, (0, 1))
assert expr == ArrayContraction(M, (0, 1))
expr = tensorcontraction(tp, (1, 2))
assert expr == ArrayContraction(tp, (1, 2))
expr = tensorcontraction(tensorcontraction(tp, (1, 2)), (0, 1))
assert expr == ArrayContraction(tp, (0, 3), (1, 2))
# Diagonalization:
expr = tensordiagonal(M, (0, 1))
assert expr == ArrayDiagonal(M, (0, 1))
expr = tensordiagonal(tensordiagonal(tp, (0, 1)), (0, 1))
assert expr == ArrayDiagonal(tp, (0, 1), (2, 3))
# Permutation of dimensions:
expr = permutedims(M, [1, 0])
assert expr == PermuteDims(M, [1, 0])
expr = permutedims(PermuteDims(tp, [1, 0, 2, 3]), [0, 1, 3, 2])
assert expr == PermuteDims(tp, [1, 0, 3, 2])
def test_arrayexpr_convert_indexed_to_array_expression():
s = Sum(A[i] * B[i], (i, 0, 3))
cg = convert_indexed_to_array(s)
assert cg == ArrayContraction(ArrayTensorProduct(A, B), (0, 1))
expr = M * N
result = ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
elem = expr[i, j]
assert convert_indexed_to_array(elem) == result
expr = M * N * M
elem = expr[i, j]
result = _array_contraction(_array_tensor_product(M, M, N), (1, 4), (2, 5))
cg = convert_indexed_to_array(elem)
assert cg == result
cg = convert_indexed_to_array((M * N * P)[i, j])
assert cg == _array_contraction(ArrayTensorProduct(M, N, P), (1, 2),
(3, 4))
cg = convert_indexed_to_array((M * N.T * P)[i, j])
assert cg == _array_contraction(ArrayTensorProduct(M, N, P), (1, 3),
(2, 4))
expr = -2 * M * N
elem = expr[i, j]
cg = convert_indexed_to_array(elem)
assert cg == ArrayContraction(ArrayTensorProduct(-2, M, N), (1, 2))
def test_arrayexpr_permutedims_sink():
cg = PermuteDims(ArrayTensorProduct(M, N), [0, 1, 3, 2], nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == ArrayTensorProduct(M, PermuteDims(N, [1, 0]))
cg = PermuteDims(ArrayTensorProduct(M, N), [1, 0, 3, 2], nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == ArrayTensorProduct(PermuteDims(M, [1, 0]), PermuteDims(N, [1, 0]))
cg = PermuteDims(ArrayTensorProduct(M, N), [3, 2, 1, 0], nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == ArrayTensorProduct(PermuteDims(N, [1, 0]), PermuteDims(M, [1, 0]))
cg = PermuteDims(ArrayContraction(ArrayTensorProduct(M, N), (1, 2)), [1, 0], nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == ArrayContraction(PermuteDims(ArrayTensorProduct(M, N), [[0, 3]]), (1, 2))
cg = PermuteDims(ArrayTensorProduct(M, N), [1, 0, 3, 2], nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == ArrayTensorProduct(PermuteDims(M, [1, 0]), PermuteDims(N, [1, 0]))
cg = PermuteDims(ArrayContraction(ArrayTensorProduct(M, N, P), (1, 2), (3, 4)), [1, 0], nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == ArrayContraction(PermuteDims(ArrayTensorProduct(M, N, P), [[0, 5]]), (1, 2), (3, 4))
def test_arrayexpr_convert_array_to_matrix2():
cg = ArrayContraction(ArrayTensorProduct(M, N), (1, 3))
assert convert_array_to_matrix(cg) == M * N.T
cg = PermuteDims(ArrayTensorProduct(M, N), Permutation([0, 1, 3, 2]))
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M, N.T)
cg = ArrayTensorProduct(M, PermuteDims(N, Permutation([1, 0])))
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M, N.T)
cg = ArrayContraction(
PermuteDims(
ArrayTensorProduct(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])),
(1, 2), (3, 5)
)
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M * P.T * Trace(N), Q.T)
cg = ArrayContraction(
ArrayTensorProduct(M, N, P, PermuteDims(Q, Permutation([1, 0]))),
(1, 5), (2, 3)
)
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M * P.T * Trace(N), Q.T)
cg = ArrayTensorProduct(M, PermuteDims(N, [1, 0]))
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M, N.T)
cg = ArrayTensorProduct(PermuteDims(M, [1, 0]), PermuteDims(N, [1, 0]))
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M.T, N.T)
cg = ArrayTensorProduct(PermuteDims(N, [1, 0]), PermuteDims(M, [1, 0]))
assert convert_array_to_matrix(cg) == ArrayTensorProduct(N.T, M.T)
def test_arrayexpr_array_expr_zero_array():
za1 = ZeroArray(k, l, m, n)
zm1 = ZeroMatrix(m, n)
za2 = ZeroArray(k, m, m, n)
zm2 = ZeroMatrix(m, m)
zm3 = ZeroMatrix(k, k)
assert ArrayTensorProduct(M, N, za1) == ZeroArray(k, k, k, k, k, l, m, n)
assert ArrayTensorProduct(M, N, zm1) == ZeroArray(k, k, k, k, m, n)
assert ArrayContraction(za1, (3, )) == ZeroArray(k, l, m)
assert ArrayContraction(zm1, (1, )) == ZeroArray(m)
assert ArrayContraction(za2, (1, 2)) == ZeroArray(k, n)
assert ArrayContraction(zm2, (0, 1)) == 0
assert ArrayDiagonal(za2, (1, 2)) == ZeroArray(k, n, m)
assert ArrayDiagonal(zm2, (0, 1)) == ZeroArray(m)
assert PermuteDims(za1, [2, 1, 3, 0]) == ZeroArray(m, l, n, k)
assert PermuteDims(zm1, [1, 0]) == ZeroArray(n, m)
assert ArrayAdd(za1) == za1
assert ArrayAdd(zm1) == ZeroArray(m, n)
tp1 = ArrayTensorProduct(MatrixSymbol("A", k, l), MatrixSymbol("B", m, n))
assert ArrayAdd(tp1, za1) == tp1
tp2 = ArrayTensorProduct(MatrixSymbol("C", k, l), MatrixSymbol("D", m, n))
assert ArrayAdd(tp1, za1, tp2) == ArrayAdd(tp1, tp2)
assert ArrayAdd(M, zm3) == M
assert ArrayAdd(M, N, zm3) == ArrayAdd(M, N)
def test_array_printer():
A = ArraySymbol('A', (4, 4, 6, 6, 6))
I = IndexedBase('I')
i, j, k = Idx('i', (0, 1)), Idx('j', (2, 3)), Idx('k', (4, 5))
prntr = NumPyPrinter()
assert prntr.doprint(ZeroArray(5)) == 'numpy.zeros((5,))'
assert prntr.doprint(OneArray(5)) == 'numpy.ones((5,))'
assert prntr.doprint(ArrayContraction(
A, [2, 3])) == 'numpy.einsum("abccd->abd", A)'
assert prntr.doprint(I) == 'I'
assert prntr.doprint(ArrayDiagonal(
A, [2, 3, 4])) == 'numpy.einsum("abccc->abc", A)'
assert prntr.doprint(ArrayDiagonal(
A, [0, 1], [2, 3])) == 'numpy.einsum("aabbc->cab", A)'
assert prntr.doprint(ArrayContraction(
A, [2], [3])) == 'numpy.einsum("abcde->abe", A)'
assert prntr.doprint(Assignment(I[i, j, k], I[i, j, k])) == 'I = I'
prntr = TensorflowPrinter()
assert prntr.doprint(ZeroArray(5)) == 'tensorflow.zeros((5,))'
assert prntr.doprint(OneArray(5)) == 'tensorflow.ones((5,))'
assert prntr.doprint(ArrayContraction(
A, [2, 3])) == 'tensorflow.linalg.einsum("abccd->abd", A)'
assert prntr.doprint(I) == 'I'
assert prntr.doprint(ArrayDiagonal(
A, [2, 3, 4])) == 'tensorflow.linalg.einsum("abccc->abc", A)'
assert prntr.doprint(ArrayDiagonal(
A, [0, 1], [2, 3])) == 'tensorflow.linalg.einsum("aabbc->cab", A)'
assert prntr.doprint(ArrayContraction(
A, [2], [3])) == 'tensorflow.linalg.einsum("abcde->abe", A)'
assert prntr.doprint(Assignment(I[i, j, k], I[i, j, k])) == 'I = I'
def test_arrayexpr_convert_array_to_matrix():
cg = ArrayContraction(ArrayTensorProduct(M), (0, 1))
assert convert_array_to_matrix(cg) == Trace(M)
cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 1), (2, 3))
assert convert_array_to_matrix(cg) == Trace(M) * Trace(N)
cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 3), (1, 2))
assert convert_array_to_matrix(cg) == Trace(M * N)
cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 2), (1, 3))
assert convert_array_to_matrix(cg) == Trace(M * N.T)
cg = convert_matrix_to_array(M * N * P)
assert convert_array_to_matrix(cg) == M * N * P
cg = convert_matrix_to_array(M * N.T * P)
assert convert_array_to_matrix(cg) == M * N.T * P
cg = ArrayContraction(ArrayTensorProduct(M,N,P,Q), (1, 2), (5, 6))
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M * N, P * Q)
cg = ArrayContraction(ArrayTensorProduct(-2, M, N), (1, 2))
assert convert_array_to_matrix(cg) == -2 * M * N
a = MatrixSymbol("a", k, 1)
b = MatrixSymbol("b", k, 1)
c = MatrixSymbol("c", k, 1)
cg = PermuteDims(
ArrayContraction(
ArrayTensorProduct(
a,
ArrayAdd(
ArrayTensorProduct(b, c),
ArrayTensorProduct(c, b),
)
), (2, 4)), [0, 1, 3, 2])
assert convert_array_to_matrix(cg) == a * (b.T * c + c.T * b)
za = ZeroArray(m, n)
assert convert_array_to_matrix(za) == ZeroMatrix(m, n)
cg = ArrayTensorProduct(3, M)
assert convert_array_to_matrix(cg) == 3 * M
# Partial conversion to matrix multiplication:
expr = ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 2), (1, 4, 6))
assert convert_array_to_matrix(expr) == ArrayContraction(ArrayTensorProduct(M.T*N, P, Q), (0, 2, 4))
x = MatrixSymbol("x", k, 1)
cg = PermuteDims(
ArrayContraction(ArrayTensorProduct(OneArray(1), x, OneArray(1), DiagMatrix(Identity(1))),
(0, 5)), Permutation(1, 2, 3))
assert convert_array_to_matrix(cg) == x
expr = ArrayAdd(M, PermuteDims(M, [1, 0]))
assert convert_array_to_matrix(expr) == M + Transpose(M)
def convert_matrix_to_array(expr: MatrixExpr) -> Basic:
if isinstance(expr, MatMul):
args_nonmat = []
args = []
for arg in expr.args:
if isinstance(arg, MatrixExpr):
args.append(arg)
else:
args_nonmat.append(convert_matrix_to_array(arg))
contractions = [(2 * i + 1, 2 * i + 2) for i in range(len(args) - 1)]
scalar = ArrayTensorProduct.fromiter(
args_nonmat) if args_nonmat else S.One
if scalar == 1:
tprod = ArrayTensorProduct(
*[convert_matrix_to_array(arg) for arg in args])
else:
tprod = ArrayTensorProduct(
scalar, *[convert_matrix_to_array(arg) for arg in args])
return ArrayContraction(tprod, *contractions)
elif isinstance(expr, MatAdd):
return ArrayAdd(*[convert_matrix_to_array(arg) for arg in expr.args])
elif isinstance(expr, Transpose):
return PermuteDims(convert_matrix_to_array(expr.args[0]), [1, 0])
elif isinstance(expr, Trace):
inner_expr = convert_matrix_to_array(expr.arg)
return ArrayContraction(inner_expr, (0, len(inner_expr.shape) - 1))
elif isinstance(expr, Mul):
return ArrayTensorProduct.fromiter(
convert_matrix_to_array(i) for i in expr.args)
elif isinstance(expr, Pow):
base = convert_matrix_to_array(expr.base)
if (expr.exp > 0) == True:
return ArrayTensorProduct.fromiter(base for i in range(expr.exp))
else:
return expr
elif isinstance(expr, MatPow):
base = convert_matrix_to_array(expr.base)
if expr.exp.is_Integer != True:
b = symbols("b", cls=Dummy)
return ArrayElementwiseApplyFunc(Lambda(b, b**expr.exp),
convert_matrix_to_array(base))
elif (expr.exp > 0) == True:
return convert_matrix_to_array(
MatMul.fromiter(base for i in range(expr.exp)))
else:
return expr
elif isinstance(expr, HadamardProduct):
tp = ArrayTensorProduct.fromiter(expr.args)
diag = [[2 * i for i in range(len(expr.args))],
[2 * i + 1 for i in range(len(expr.args))]]
return ArrayDiagonal(tp, *diag)
elif isinstance(expr, HadamardPower):
base, exp = expr.args
return convert_matrix_to_array(
HadamardProduct.fromiter(base for i in range(exp)))
else:
return expr
def test_arrayexpr_nested_array_elementwise_add():
cg = ArrayContraction(
ArrayAdd(ArrayTensorProduct(M, N), ArrayTensorProduct(N, M)), (1, 2))
result = ArrayAdd(ArrayContraction(ArrayTensorProduct(M, N), (1, 2)),
ArrayContraction(ArrayTensorProduct(N, M), (1, 2)))
assert cg == result
cg = ArrayDiagonal(
ArrayAdd(ArrayTensorProduct(M, N), ArrayTensorProduct(N, M)), (1, 2))
result = ArrayAdd(ArrayDiagonal(ArrayTensorProduct(M, N), (1, 2)),
ArrayDiagonal(ArrayTensorProduct(N, M), (1, 2)))
assert cg == result
def test_arrayexpr_derivatives1():
res = array_derive(X, X)
assert res == PermuteDims(ArrayTensorProduct(I, I), [0, 2, 1, 3])
cg = ArrayTensorProduct(A, X, B)
res = array_derive(cg, X)
assert res == PermuteDims(
ArrayTensorProduct(I, A, I, B),
[0, 4, 2, 3, 1, 5, 6, 7])
cg = ArrayContraction(X, (0, 1))
res = array_derive(cg, X)
assert res == ArrayContraction(ArrayTensorProduct(I, I), (1, 3))
cg = ArrayDiagonal(X, (0, 1))
res = array_derive(cg, X)
assert res == ArrayDiagonal(ArrayTensorProduct(I, I), (1, 3))
cg = ElementwiseApplyFunction(sin, X)
res = array_derive(cg, X)
assert res.dummy_eq(ArrayDiagonal(
ArrayTensorProduct(
ElementwiseApplyFunction(cos, X),
I,
I
), (0, 3), (1, 5)))
cg = ArrayElementwiseApplyFunc(sin, X)
res = array_derive(cg, X)
assert res.dummy_eq(ArrayDiagonal(
ArrayTensorProduct(
I,
I,
ArrayElementwiseApplyFunc(cos, X)
), (1, 4), (3, 5)))
res = array_derive(A1, A1)
assert res == PermuteDims(
ArrayTensorProduct(Identity(3), Identity(2), Identity(k)),
[0, 2, 4, 1, 3, 5]
)
cg = ArrayElementwiseApplyFunc(sin, A1)
res = array_derive(cg, A1)
assert res.dummy_eq(ArrayDiagonal(
ArrayTensorProduct(
Identity(3), Identity(2), Identity(k),
ArrayElementwiseApplyFunc(cos, A1)
), (1, 6), (3, 7), (5, 8)
))
def test_arrayexpr_convert_index_to_array_support_function():
expr = M[i, j]
assert _convert_indexed_to_array(expr) == (M, (i, j))
expr = M[i, j] * N[k, l]
assert _convert_indexed_to_array(expr) == (ArrayTensorProduct(M, N),
(i, j, k, l))
expr = M[i, j] * N[j, k]
assert _convert_indexed_to_array(expr) == (ArrayDiagonal(
ArrayTensorProduct(M, N), (1, 2)), (i, k, j))
expr = Sum(M[i, j] * N[j, k], (j, 0, k - 1))
assert _convert_indexed_to_array(expr) == (ArrayContraction(
ArrayTensorProduct(M, N), (1, 2)), (i, k))
expr = M[i, j] + N[i, j]
assert _convert_indexed_to_array(expr) == (ArrayAdd(M, N), (i, j))
expr = M[i, j] + N[j, i]
assert _convert_indexed_to_array(expr) == (ArrayAdd(
M, PermuteDims(N, Permutation([1, 0]))), (i, j))
expr = M[i, j] + M[j, i]
assert _convert_indexed_to_array(expr) == (ArrayAdd(
M, PermuteDims(M, Permutation([1, 0]))), (i, j))
expr = (M * N * P)[i, j]
assert _convert_indexed_to_array(expr) == (ArrayContraction(
ArrayTensorProduct(M, N, P), (1, 2), (3, 4)), (i, j))
expr = expr.function # Disregard summation in previous expression
ret1, ret2 = _convert_indexed_to_array(expr)
assert ret1 == ArrayDiagonal(ArrayTensorProduct(M, N, P), (1, 2), (3, 4))
assert str(ret2) == "(i, j, _i_1, _i_2)"
expr = KroneckerDelta(i, j) * M[i, k]
assert _convert_indexed_to_array(expr) == (M, ({i, j}, k))
expr = KroneckerDelta(i, j) * KroneckerDelta(j, k) * M[i, l]
assert _convert_indexed_to_array(expr) == (M, ({i, j, k}, l))
expr = KroneckerDelta(j, k) * (M[i, j] * N[k, l] + N[i, j] * M[k, l])
assert _convert_indexed_to_array(expr) == (ArrayDiagonal(
ArrayAdd(
ArrayTensorProduct(M, N),
PermuteDims(ArrayTensorProduct(M, N),
Permutation(0, 2)(1, 3))),
(1, 2)), (i, l, frozenset({j, k})))
expr = KroneckerDelta(j, m) * KroneckerDelta(
m, k) * (M[i, j] * N[k, l] + N[i, j] * M[k, l])
assert _convert_indexed_to_array(expr) == (ArrayDiagonal(
ArrayAdd(
ArrayTensorProduct(M, N),
PermuteDims(ArrayTensorProduct(M, N),
Permutation(0, 2)(1, 3))),
(1, 2)), (i, l, frozenset({j, m, k})))
expr = KroneckerDelta(i, j) * KroneckerDelta(j, k) * KroneckerDelta(
k, m) * M[i, 0] * KroneckerDelta(m, n)
assert _convert_indexed_to_array(expr) == (M, ({i, j, k, m, n}, 0))
expr = M[i, i]
assert _convert_indexed_to_array(expr) == (ArrayDiagonal(M, (0, 1)), (i, ))
def _(expr: Inverse, x: Expr):
mat = expr.I
dexpr = array_derive(mat, x)
tp = ArrayTensorProduct(-expr, dexpr, expr)
mp = ArrayContraction(tp, (1, 4), (5, 6))
pp = PermuteDims(mp, [1, 2, 0, 3])
return pp
def test_arrayexpr_convert_indexed_to_array_and_back_to_matrix():
expr = a.T * b
elem = expr[0, 0]
cg = convert_indexed_to_array(elem)
assert cg == ArrayElement(
ArrayContraction(ArrayTensorProduct(a, b), (0, 2)), [0, 0])
expr = M[i, j] + N[i, j]
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(p1) == M + N
expr = M[i, j] + N[j, i]
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(p1) == M + N.T
expr = M[i, j] * N[k, l] + N[i, j] * M[k, l]
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(p1) == ArrayAdd(ArrayTensorProduct(M, N),
ArrayTensorProduct(N, M))
expr = (M * N * P)[i, j]
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(p1) == M * N * P
expr = Sum(M[i, j] * (N * P)[j, m], (j, 0, k - 1))
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(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 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(
p1) == M * P * N + M * P.T * N + N * P * M + N * P.T * M
def test_array2matrix():
# See issue https://github.com/sympy/sympy/pull/22877
expr = PermuteDims(
ArrayContraction(ArrayTensorProduct(x, I, I1, x), (0, 3), (1, 7)),
Permutation(2, 3))
expected = PermuteDims(ArrayTensorProduct(x * x.T, I1),
Permutation(3)(1, 2))
assert _array2matrix(expr) == expected
def _(expr: ArrayContraction, x: Expr):
fd = array_derive(expr.expr, x)
rank_x = len(get_shape(x))
contraction_indices = expr.contraction_indices
new_contraction_indices = [
tuple(j + rank_x for j in i) for i in contraction_indices
]
return ArrayContraction(fd, *new_contraction_indices)
def test_arrayexpr_push_indices_up_and_down():
indices = list(range(12))
contr_diag_indices = [(0, 6), (2, 8)]
assert ArrayContraction._push_indices_down(contr_diag_indices, indices) == (1, 3, 4, 5, 7, 9, 10, 11, 12, 13, 14, 15)
assert ArrayContraction._push_indices_up(contr_diag_indices, indices) == (None, 0, None, 1, 2, 3, None, 4, None, 5, 6, 7)
assert ArrayDiagonal._push_indices_down(contr_diag_indices, indices, 10) == (1, 3, 4, 5, 7, 9, (0, 6), (2, 8), None, None, None, None)
assert ArrayDiagonal._push_indices_up(contr_diag_indices, indices, 10) == (6, 0, 7, 1, 2, 3, 6, 4, 7, 5, None, None)
contr_diag_indices = [(1, 2), (7, 8)]
assert ArrayContraction._push_indices_down(contr_diag_indices, indices) == (0, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15)
assert ArrayContraction._push_indices_up(contr_diag_indices, indices) == (0, None, None, 1, 2, 3, 4, None, None, 5, 6, 7)
assert ArrayDiagonal._push_indices_down(contr_diag_indices, indices, 10) == (0, 3, 4, 5, 6, 9, (1, 2), (7, 8), None, None, None, None)
assert ArrayDiagonal._push_indices_up(contr_diag_indices, indices, 10) == (0, 6, 6, 1, 2, 3, 4, 7, 7, 5, None, None)
def _a2m_mul(*args):
if all(not isinstance(i, _CodegenArrayAbstract) for i in args):
from sympy import MatMul
return MatMul(*args).doit()
else:
return ArrayContraction(
ArrayTensorProduct(*args),
*[(2 * i - 1, 2 * i) for i in range(1, len(args))])
def _support_function_tp1_recognize(contraction_indices, args):
if len(contraction_indices) == 0:
return _a2m_tensor_product(*args)
ac = ArrayContraction(ArrayTensorProduct(*args), *contraction_indices)
editor = _EditArrayContraction(ac)
editor.track_permutation_start()
while True:
flag_stop: bool = True
for i, arg_with_ind in enumerate(editor.args_with_ind):
if not isinstance(arg_with_ind.element, MatrixExpr):
continue
first_index = arg_with_ind.indices[0]
second_index = arg_with_ind.indices[1]
first_frequency = editor.count_args_with_index(first_index)
second_frequency = editor.count_args_with_index(second_index)
if first_index is not None and first_frequency == 1 and first_index == second_index:
flag_stop = False
arg_with_ind.element = Trace(arg_with_ind.element)._normalize()
arg_with_ind.indices = []
break
scan_indices = []
if first_frequency == 2:
scan_indices.append(first_index)
if second_frequency == 2:
scan_indices.append(second_index)
candidate, transpose, found_index = _get_candidate_for_matmul_from_contraction(
scan_indices, editor.args_with_ind[i + 1:])
if candidate is not None:
flag_stop = False
editor.track_permutation_merge(arg_with_ind, candidate)
transpose1 = found_index == first_index
new_arge, other_index = _insert_candidate_into_editor(
editor, arg_with_ind, candidate, transpose1, transpose)
if found_index == first_index:
new_arge.indices = [second_index, other_index]
else:
new_arge.indices = [first_index, other_index]
set_indices = set(new_arge.indices)
if len(set_indices) == 1 and set_indices != {None}:
# This is a trace:
new_arge.element = Trace(new_arge.element)._normalize()
new_arge.indices = []
editor.args_with_ind[i] = new_arge
# TODO: is this break necessary?
break
if flag_stop:
break
editor.refresh_indices()
return editor.to_array_contraction()
def test_arrayexpr_convert_array_contraction_tp_additions():
a = ArrayAdd(
ArrayTensorProduct(M, N),
ArrayTensorProduct(N, M)
)
tp = ArrayTensorProduct(P, a, Q)
expr = ArrayContraction(tp, (3, 4))
expected = ArrayTensorProduct(
P,
ArrayAdd(
ArrayContraction(ArrayTensorProduct(M, N), (1, 2)),
ArrayContraction(ArrayTensorProduct(N, M), (1, 2)),
),
Q
)
assert expr == expected
assert convert_array_to_matrix(expr) == ArrayTensorProduct(P, M * N + N * M, Q)
expr = ArrayContraction(tp, (1, 2), (3, 4), (5, 6))
result = ArrayContraction(
ArrayTensorProduct(
P,
ArrayAdd(
ArrayContraction(ArrayTensorProduct(M, N), (1, 2)),
ArrayContraction(ArrayTensorProduct(N, M), (1, 2)),
),
Q
), (1, 2), (3, 4))
assert expr == result
assert convert_array_to_matrix(expr) == P * (M * N + N * M) * Q
def _(expr: ArrayContraction):
newexpr, removed = _remove_trivial_dims(expr.expr)
new_contraction_indices = [
tuple(j for j in i if j not in removed)
for i in expr.contraction_indices
]
# Remove possible empty tuples "()":
new_contraction_indices = [i for i in new_contraction_indices if i]
return ArrayContraction(newexpr, *new_contraction_indices), removed
def _array_diag2contr_diagmatrix(expr: ArrayDiagonal):
if isinstance(expr.expr, ArrayTensorProduct):
args = list(expr.expr.args)
diag_indices = list(expr.diagonal_indices)
mapping = _get_mapping_from_subranks(
[_get_subrank(arg) for arg in args])
tuple_links = [[mapping[j] for j in i] for i in diag_indices]
contr_indices = []
total_rank = get_rank(expr)
replaced = [False for arg in args]
for i, (abs_pos, rel_pos) in enumerate(zip(diag_indices, tuple_links)):
if len(abs_pos) != 2:
continue
(pos1_outer, pos1_inner), (pos2_outer, pos2_inner) = rel_pos
arg1 = args[pos1_outer]
arg2 = args[pos2_outer]
if get_rank(arg1) != 2 or get_rank(arg2) != 2:
if replaced[pos1_outer]:
diag_indices[i] = None
if replaced[pos2_outer]:
diag_indices[i] = None
continue
pos1_in2 = 1 - pos1_inner
pos2_in2 = 1 - pos2_inner
if arg1.shape[pos1_in2] == 1:
if arg1.shape[pos1_inner] != 1:
darg1 = DiagMatrix(arg1)
else:
darg1 = arg1
args.append(darg1)
contr_indices.append(
((pos2_outer, pos2_inner), (len(args) - 1, pos1_inner)))
total_rank += 1
diag_indices[i] = None
args[pos1_outer] = OneArray(arg1.shape[pos1_in2])
replaced[pos1_outer] = True
elif arg2.shape[pos2_in2] == 1:
if arg2.shape[pos2_inner] != 1:
darg2 = DiagMatrix(arg2)
else:
darg2 = arg2
args.append(darg2)
contr_indices.append(
((pos1_outer, pos1_inner), (len(args) - 1, pos2_inner)))
total_rank += 1
diag_indices[i] = None
args[pos2_outer] = OneArray(arg2.shape[pos2_in2])
replaced[pos2_outer] = True
diag_indices_new = [i for i in diag_indices if i is not None]
cumul = list(accumulate([0] + [get_rank(arg) for arg in args]))
contr_indices2 = [
tuple(cumul[a] + b for a, b in i) for i in contr_indices
]
tc = ArrayContraction(ArrayTensorProduct(*args), *contr_indices2)
td = ArrayDiagonal(tc, *diag_indices_new)
return td
return expr
def test_arrayexpr_convert_array_to_matrix_remove_trivial_dims():
# Tensor Product:
assert _remove_trivial_dims(ArrayTensorProduct(a, b)) == (a * b.T, [1, 3])
assert _remove_trivial_dims(ArrayTensorProduct(a.T,
b)) == (a * b.T, [0, 3])
assert _remove_trivial_dims(ArrayTensorProduct(a,
b.T)) == (a * b.T, [1, 2])
assert _remove_trivial_dims(ArrayTensorProduct(a.T,
b.T)) == (a * b.T, [0, 2])
assert _remove_trivial_dims(ArrayTensorProduct(I, a.T,
b.T)) == (a * b.T,
[0, 1, 2, 4])
assert _remove_trivial_dims(ArrayTensorProduct(a.T, I,
b.T)) == (a * b.T,
[0, 2, 3, 4])
assert _remove_trivial_dims(ArrayTensorProduct(a, I)) == (a, [2, 3])
assert _remove_trivial_dims(ArrayTensorProduct(I, a)) == (a, [0, 1])
assert _remove_trivial_dims(ArrayTensorProduct(a.T, b.T, c,
d)) == (ArrayTensorProduct(
a * b.T,
c * d.T), [0, 2, 5, 7])
assert _remove_trivial_dims(ArrayTensorProduct(a.T, I, b.T, c, d,
I)) == (ArrayTensorProduct(
a * b.T, c * d.T,
I), [0, 2, 3, 4, 7, 9])
# Addition:
cg = ArrayAdd(ArrayTensorProduct(a, b), ArrayTensorProduct(c, d))
assert _remove_trivial_dims(cg) == (a * b.T + c * d.T, [1, 3])
# Permute Dims:
cg = PermuteDims(ArrayTensorProduct(a, b), Permutation(3)(1, 2))
assert _remove_trivial_dims(cg) == (a * b.T, [2, 3])
cg = PermuteDims(ArrayTensorProduct(a, I, b), Permutation(5)(1, 2, 3, 4))
assert _remove_trivial_dims(cg) == (a * b.T, [1, 2, 4, 5])
cg = PermuteDims(ArrayTensorProduct(I, b, a),
Permutation(5)(1, 2, 4, 5, 3))
assert _remove_trivial_dims(cg) == (b * a.T, [0, 3, 4, 5])
# Diagonal:
cg = ArrayDiagonal(ArrayTensorProduct(M, a), (1, 2))
assert _remove_trivial_dims(cg) == (cg, [])
# Contraction:
cg = ArrayContraction(ArrayTensorProduct(M, a), (1, 2))
assert _remove_trivial_dims(cg) == (cg, [])
def test_arrayexpr_split_multiple_contractions():
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)
cg = ArrayContraction(ArrayTensorProduct(A.T, a, b, b.T, (A*X*b).applyfunc(cos)), (1, 2, 8), (5, 6, 9))
expected = ArrayContraction(ArrayTensorProduct(A.T, DiagMatrix(a), OneArray(1), b, b.T, (A*X*b).applyfunc(cos)), (1, 3), (2, 9), (6, 7, 10))
assert cg.split_multiple_contractions().dummy_eq(expected)
# Check no overlap of lines:
cg = ArrayContraction(ArrayTensorProduct(A, a, C, a, B), (1, 2, 4), (5, 6, 8), (3, 7))
assert cg.split_multiple_contractions() == cg
cg = ArrayContraction(ArrayTensorProduct(a, b, A), (0, 2, 4), (1, 3))
assert cg.split_multiple_contractions() == cg
def test_arrayexpr_contraction_permutation_mix():
Me = M.subs(k, 3).as_explicit()
Ne = N.subs(k, 3).as_explicit()
cg1 = ArrayContraction(PermuteDims(ArrayTensorProduct(M, N), Permutation([0, 2, 1, 3])), (2, 3))
cg2 = ArrayContraction(ArrayTensorProduct(M, N), (1, 3))
assert cg1 == cg2
cge1 = tensorcontraction(permutedims(tensorproduct(Me, Ne), Permutation([0, 2, 1, 3])), (2, 3))
cge2 = tensorcontraction(tensorproduct(Me, Ne), (1, 3))
assert cge1 == cge2
cg1 = PermuteDims(ArrayTensorProduct(M, N), Permutation([0, 1, 3, 2]))
cg2 = ArrayTensorProduct(M, PermuteDims(N, Permutation([1, 0])))
assert cg1 == cg2
cg1 = ArrayContraction(
PermuteDims(
ArrayTensorProduct(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])),
(1, 2), (3, 5)
)
cg2 = ArrayContraction(
ArrayTensorProduct(M, N, P, PermuteDims(Q, Permutation([1, 0]))),
(1, 5), (2, 3)
)
assert cg1 == cg2
cg1 = ArrayContraction(
PermuteDims(
ArrayTensorProduct(M, N, P, Q), Permutation([1, 0, 4, 6, 2, 7, 5, 3])),
(0, 1), (2, 6), (3, 7)
)
cg2 = PermuteDims(
ArrayContraction(
ArrayTensorProduct(M, P, Q, N),
(0, 1), (2, 3), (4, 7)),
[1, 0]
)
assert cg1 == cg2
cg1 = ArrayContraction(
PermuteDims(
ArrayTensorProduct(M, N, P, Q), Permutation([1, 0, 4, 6, 7, 2, 5, 3])),
(0, 1), (2, 6), (3, 7)
)
cg2 = PermuteDims(
ArrayContraction(
ArrayTensorProduct(PermuteDims(M, [1, 0]), N, P, Q),
(0, 1), (3, 6), (4, 5)
),
Permutation([1, 0])
)
assert cg1 == cg2
def test_array_expr_construction_with_functions():
tp = tensorproduct(M, N)
assert tp == ArrayTensorProduct(M, N)
expr = tensorproduct(A, eye(2))
assert expr == ArrayTensorProduct(A, eye(2))
# Contraction:
expr = tensorcontraction(M, (0, 1))
assert expr == ArrayContraction(M, (0, 1))
expr = tensorcontraction(tp, (1, 2))
assert expr == ArrayContraction(tp, (1, 2))
expr = tensorcontraction(tensorcontraction(tp, (1, 2)), (0, 1))
assert expr == ArrayContraction(tp, (0, 3), (1, 2))
# Diagonalization:
expr = tensordiagonal(M, (0, 1))
assert expr == ArrayDiagonal(M, (0, 1))
expr = tensordiagonal(tensordiagonal(tp, (0, 1)), (0, 1))
assert expr == ArrayDiagonal(tp, (0, 1), (2, 3))
# Permutation of dimensions:
expr = permutedims(M, [1, 0])
assert expr == PermuteDims(M, [1, 0])
expr = permutedims(PermuteDims(tp, [1, 0, 2, 3]), [0, 1, 3, 2])
assert expr == PermuteDims(tp, [1, 0, 3, 2])
expr = PermuteDims(tp, index_order_new=["a", "b", "c", "d"], index_order_old=["d", "c", "b", "a"])
assert expr == PermuteDims(tp, [3, 2, 1, 0])
arr = Array(range(32)).reshape(2, 2, 2, 2, 2)
expr = PermuteDims(arr, index_order_new=["a", "b", "c", "d", "e"], index_order_old=['b', 'e', 'a', 'd', 'c'])
assert expr == PermuteDims(arr, [2, 0, 4, 3, 1])
assert expr.as_explicit() == permutedims(arr, index_order_new=["a", "b", "c", "d", "e"], index_order_old=['b', 'e', 'a', 'd', 'c'])
def test_arrayexpr_convert_array_to_matrix_diag2contraction_diagmatrix():
cg = ArrayDiagonal(ArrayTensorProduct(M, a), (1, 2))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == ArrayContraction(
ArrayTensorProduct(M, OneArray(1), DiagMatrix(a)), (1, 3))
raises(ValueError, lambda: ArrayDiagonal(ArrayTensorProduct(a, M), (1, 2)))
cg = ArrayDiagonal(ArrayTensorProduct(a.T, M), (1, 2))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == ArrayContraction(
ArrayTensorProduct(OneArray(1), M, DiagMatrix(a.T)), (1, 4))
cg = ArrayDiagonal(ArrayTensorProduct(a.T, M, N, b.T), (1, 2), (4, 7))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == ArrayContraction(
ArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a.T),
DiagMatrix(b.T)), (1, 7), (3, 9))
cg = ArrayDiagonal(ArrayTensorProduct(a, M, N, b.T), (0, 2), (4, 7))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == ArrayContraction(
ArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a),
DiagMatrix(b.T)), (1, 6), (3, 9))
cg = ArrayDiagonal(ArrayTensorProduct(a, M, N, b.T), (0, 4), (3, 7))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == ArrayContraction(
ArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a),
DiagMatrix(b.T)), (3, 6), (2, 9))
I1 = Identity(1)
x = MatrixSymbol("x", k, 1)
A = MatrixSymbol("A", k, k)
cg = ArrayDiagonal(ArrayTensorProduct(x, A.T, I1), (0, 2))
assert _array_diag2contr_diagmatrix(cg).shape == cg.shape
assert _array2matrix(cg).shape == cg.shape
def test_array_expr_as_explicit_with_explicit_component_arrays():
# Test if .as_explicit() works with explicit-component arrays
# nested in array expressions:
from sympy.abc import x, y, z, t
A = Array([[x, y], [z, t]])
assert ArrayTensorProduct(A, A).as_explicit() == tensorproduct(A, A)
assert ArrayDiagonal(A, (0, 1)).as_explicit() == tensordiagonal(A, (0, 1))
assert ArrayContraction(A, (0, 1)).as_explicit() == tensorcontraction(A, (0, 1))
assert ArrayAdd(A, A).as_explicit() == A + A
assert ArrayElementwiseApplyFunc(sin, A).as_explicit() == A.applyfunc(sin)
assert PermuteDims(A, [1, 0]).as_explicit() == permutedims(A, [1, 0])
assert Reshape(A, [4]).as_explicit() == A.reshape(4)
