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_arrayexpr_convert_indexed_to_array_broadcast():
A = ArraySymbol("A", (3, 3))
B = ArraySymbol("B", (3, 3))
expr = A[i, j] + B[k, l]
O2 = OneArray(3, 3)
expected = ArrayAdd(ArrayTensorProduct(A, O2), ArrayTensorProduct(O2, B))
assert convert_indexed_to_array(expr) == expected
assert convert_indexed_to_array(expr, [i, j, k, l]) == expected
assert convert_indexed_to_array(expr, [l, k, i, j]) == ArrayAdd(
PermuteDims(ArrayTensorProduct(O2, A), [1, 0, 2, 3]),
PermuteDims(ArrayTensorProduct(B, O2), [1, 0, 2, 3]))
expr = A[i, j] + B[j, k]
O1 = OneArray(3)
assert convert_indexed_to_array(expr, [i, j, k]) == ArrayAdd(
ArrayTensorProduct(A, O1), ArrayTensorProduct(O1, B))
C = ArraySymbol("C", (d0, d1))
D = ArraySymbol("D", (d3, d1))
expr = C[i, j] + D[k, j]
assert convert_indexed_to_array(expr, [i, j, k]) == ArrayAdd(
ArrayTensorProduct(C, OneArray(d3)),
PermuteDims(ArrayTensorProduct(OneArray(d0), D), [0, 2, 1]))
X = ArraySymbol("X", (5, 3))
expr = X[i, n] - X[j, n]
assert convert_indexed_to_array(expr, [i, j, n]) == ArrayAdd(
ArrayTensorProduct(-1, OneArray(5), X),
PermuteDims(ArrayTensorProduct(X, OneArray(5)), [0, 2, 1]))
raises(ValueError, lambda: convert_indexed_to_array(C[i, j] + D[i, j]))
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_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 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 _(expr: PermuteDims):
if expr.permutation.array_form == [1, 0]:
return _a2m_transpose(_array2matrix(expr.expr))
elif isinstance(expr.expr, ArrayTensorProduct):
ranks = expr.expr.subranks
inv_permutation = expr.permutation**(-1)
newrange = [inv_permutation(i) for i in range(sum(ranks))]
newpos = []
counter = 0
for rank in ranks:
newpos.append(newrange[counter:counter + rank])
counter += rank
newargs = []
newperm = []
scalars = []
for pos, arg in zip(newpos, expr.expr.args):
if len(pos) == 0:
scalars.append(_array2matrix(arg))
elif pos == sorted(pos):
newargs.append((_array2matrix(arg), pos[0]))
newperm.extend(pos)
elif len(pos) == 2:
newargs.append((_a2m_transpose(_array2matrix(arg)), pos[0]))
newperm.extend(reversed(pos))
else:
raise NotImplementedError()
newargs = [i[0] for i in newargs]
return PermuteDims(_a2m_tensor_product(*scalars, *newargs),
_af_invert(newperm))
elif isinstance(expr.expr, ArrayContraction):
mat_mul_lines = _array2matrix(expr.expr)
if not isinstance(mat_mul_lines, ArrayTensorProduct):
flat_cyclic_form = [
j for i in expr.permutation.cyclic_form for j in i
]
expr_shape = get_shape(expr)
if all(expr_shape[i] == 1 for i in flat_cyclic_form):
return mat_mul_lines
return mat_mul_lines
# TODO: this assumes that all arguments are matrices, it may not be the case:
permutation = Permutation(2 * len(mat_mul_lines.args) -
1) * expr.permutation
permuted = [permutation(i) for i in range(2 * len(mat_mul_lines.args))]
args_array = [None for i in mat_mul_lines.args]
for i in range(len(mat_mul_lines.args)):
p1 = permuted[2 * i]
p2 = permuted[2 * i + 1]
if p1 // 2 != p2 // 2:
return PermuteDims(mat_mul_lines, permutation)
pos = p1 // 2
if p1 > p2:
args_array[i] = _a2m_transpose(mat_mul_lines.args[pos])
else:
args_array[i] = mat_mul_lines.args[pos]
return _a2m_tensor_product(*args_array)
else:
return expr
def test_matrix_derivative_non_matrix_result():
# This is a 4-dimensional array:
I = Identity(k)
AdA = PermuteDims(ArrayTensorProduct(I, I), Permutation(3)(1, 2))
assert A.diff(A) == AdA
assert A.T.diff(A) == PermuteDims(ArrayTensorProduct(I, I), Permutation(3)(1, 2, 3))
assert (2*A).diff(A) == PermuteDims(ArrayTensorProduct(2*I, I), Permutation(3)(1, 2))
assert MatAdd(A, A).diff(A) == ArrayAdd(AdA, AdA)
assert (A + B).diff(A) == AdA
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_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_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_arrayexpr_normalize_diagonal_permutedims():
tp = ArrayTensorProduct(M, Q, N, P)
expr = ArrayDiagonal(PermuteDims(tp, [0, 1, 2, 4, 7, 6, 3, 5]), (2, 4, 5),
(6, 7), (0, 3))
result = ArrayDiagonal(tp, (2, 6, 7), (3, 5), (0, 4))
assert expr == result
tp = ArrayTensorProduct(M, N, P, Q)
expr = ArrayDiagonal(PermuteDims(tp, [0, 5, 2, 4, 1, 6, 3, 7]), (1, 2, 6),
(3, 4))
result = ArrayDiagonal(ArrayTensorProduct(M, P, N, Q), (3, 4, 5), (1, 2))
assert expr == result
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) == (_array_contraction(
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) == (_array_diagonal(
_array_add(
ArrayTensorProduct(M, N),
_permute_dims(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) == (_array_diagonal(
_array_add(
ArrayTensorProduct(M, N),
_permute_dims(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 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 _(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 _(expr: ArraySymbol, x: Expr):
if expr == x:
return PermuteDims(
ArrayTensorProduct.fromiter(Identity(i) for i in expr.shape),
[2 * i for i in range(len(expr.shape))] +
[2 * i + 1 for i in range(len(expr.shape))])
return ZeroArray(*(x.shape + expr.shape))
def test_mixed_deriv_mixed_expressions():
expr = 3 * Trace(A)
assert expr.diff(A) == 3 * Identity(k)
expr = k
deriv = expr.diff(A)
assert isinstance(deriv, ZeroMatrix)
assert deriv == ZeroMatrix(k, k)
expr = Trace(A)**2
assert expr.diff(A) == (2 * Trace(A)) * Identity(k)
expr = Trace(A) * A
I = Identity(k)
assert expr.diff(A) == ArrayAdd(
ArrayTensorProduct(I, A),
PermuteDims(ArrayTensorProduct(Trace(A) * I, I),
Permutation(3)(1, 2)))
expr = Trace(Trace(A) * A)
assert expr.diff(A) == (2 * Trace(A)) * Identity(k)
expr = Trace(Trace(Trace(A) * A) * A)
assert expr.diff(A) == (3 * Trace(A)**2) * Identity(k)
def _(expr: ArrayTensorProduct, x: Expr):
args = expr.args
addend_list = []
for i, arg in enumerate(expr.args):
darg = array_derive(arg, x)
if darg == 0:
continue
args_prev = args[:i]
args_succ = args[i + 1:]
shape_prev = reduce(operator.add, map(get_shape, args_prev), ())
shape_succ = reduce(operator.add, map(get_shape, args_succ), ())
addend = ArrayTensorProduct(*args_prev, darg, *args_succ)
tot1 = len(get_shape(x))
tot2 = tot1 + len(shape_prev)
tot3 = tot2 + len(get_shape(arg))
tot4 = tot3 + len(shape_succ)
perm = [i for i in range(tot1, tot2)] + \
[i for i in range(tot1)] + [i for i in range(tot2, tot3)] + \
[i for i in range(tot3, tot4)]
addend = PermuteDims(addend, _af_invert(perm))
addend_list.append(addend)
if len(addend_list) == 1:
return addend_list[0]
elif len(addend_list) == 0:
return S.Zero
else:
return ArrayAdd(*addend_list)
def test_arrayexpr_nested_permutations():
cg = PermuteDims(PermuteDims(M, (1, 0)), (1, 0))
assert cg == M
times = 3
plist1 = [list(range(6)) for i in range(times)]
plist2 = [list(range(6)) for i in range(times)]
for i in range(times):
random.shuffle(plist1[i])
random.shuffle(plist2[i])
plist1.append([2, 5, 4, 1, 0, 3])
plist2.append([3, 5, 0, 4, 1, 2])
plist1.append([2, 5, 4, 0, 3, 1])
plist2.append([3, 0, 5, 1, 2, 4])
plist1.append([5, 4, 2, 0, 3, 1])
plist2.append([4, 5, 0, 2, 3, 1])
Me = M.subs(k, 3).as_explicit()
Ne = N.subs(k, 3).as_explicit()
Pe = P.subs(k, 3).as_explicit()
cge = tensorproduct(Me, Ne, Pe)
for permutation_array1, permutation_array2 in zip(plist1, plist2):
p1 = Permutation(permutation_array1)
p2 = Permutation(permutation_array2)
cg = PermuteDims(
PermuteDims(
ArrayTensorProduct(M, N, P),
p1),
p2
)
result = PermuteDims(
ArrayTensorProduct(M, N, P),
p2*p1
)
assert cg == result
# Check that `permutedims` behaves the same way with explicit-component arrays:
result1 = permutedims(permutedims(cge, p1), p2)
result2 = permutedims(cge, p2*p1)
assert result1 == result2
def test_arrayexpr_convert_array_to_matrix_support_function():
assert _support_function_tp1_recognize([], [2 * k]) == 2 * k
assert _support_function_tp1_recognize([(1, 2)], [A, 2 * k, B, 3]) == 6 * k * A * B
assert _support_function_tp1_recognize([(0, 3), (1, 2)], [A, B]) == Trace(A * B)
assert _support_function_tp1_recognize([(1, 2)], [A, B]) == A * B
assert _support_function_tp1_recognize([(0, 2)], [A, B]) == A.T * B
assert _support_function_tp1_recognize([(1, 3)], [A, B]) == A * B.T
assert _support_function_tp1_recognize([(0, 3)], [A, B]) == A.T * B.T
assert _support_function_tp1_recognize([(1, 2), (5, 6)], [A, B, C, D]) == ArrayTensorProduct(A * B, C * D)
assert _support_function_tp1_recognize([(1, 4), (3, 6)], [A, B, C, D]) == PermuteDims(
ArrayTensorProduct(A * C, B * D), [0, 2, 1, 3])
assert _support_function_tp1_recognize([(0, 3), (1, 4)], [A, B, C]) == B * A * C
assert _support_function_tp1_recognize([(9, 10), (1, 2), (5, 6), (3, 4), (7, 8)],
[X, Y, A, B, C, D]) == X * Y * A * B * C * D
assert _support_function_tp1_recognize([(9, 10), (1, 2), (5, 6), (3, 4)],
[X, Y, A, B, C, D]) == ArrayTensorProduct(X * Y * A * B, C * D)
assert _support_function_tp1_recognize([(1, 7), (3, 8), (4, 11)], [X, Y, A, B, C, D]) == PermuteDims(
ArrayTensorProduct(X * B.T, Y * C, A.T * D.T), [0, 2, 4, 1, 3, 5]
)
assert _support_function_tp1_recognize([(0, 1), (3, 6), (5, 8)], [X, A, B, C, D]) == PermuteDims(
ArrayTensorProduct(Trace(X) * A * C, B * D), [0, 2, 1, 3])
assert _support_function_tp1_recognize([(1, 2), (3, 4), (5, 6), (7, 8)], [A, A, B, C, D]) == A ** 2 * B * C * D
assert _support_function_tp1_recognize([(1, 2), (3, 4), (5, 6), (7, 8)], [X, A, B, C, D]) == X * A * B * C * D
assert _support_function_tp1_recognize([(1, 6), (3, 8), (5, 10)], [X, Y, A, B, C, D]) == PermuteDims(
ArrayTensorProduct(X * B, Y * C, A * D), [0, 2, 4, 1, 3, 5]
)
assert _support_function_tp1_recognize([(1, 4), (3, 6)], [A, B, C, D]) == PermuteDims(
ArrayTensorProduct(A * C, B * D), [0, 2, 1, 3])
assert _support_function_tp1_recognize([(0, 4), (1, 7), (2, 5), (3, 8)], [X, A, B, C, D]) == C*X.T*B*A*D
assert _support_function_tp1_recognize([(0, 4), (1, 7), (2, 5), (3, 8)], [X, A, B, C, D]) == C*X.T*B*A*D
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_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 = ArrayTensorProduct(M, N)
f = lambdify((M, N), cg, 'numpy')
assert (f(ma, mb) == np.einsum(ma, [0, 1], mb, [2, 3])).all()
cg = ArrayAdd(M, N)
f = lambdify((M, N), cg, 'numpy')
assert (f(ma, mb) == ma + mb).all()
cg = ArrayAdd(M, N, P)
f = lambdify((M, N, P), cg, 'numpy')
assert (f(ma, mb, mc) == ma + mb + mc).all()
cg = ArrayAdd(M, N, P, Q)
f = lambdify((M, N, P, Q), cg, 'numpy')
assert (f(ma, mb, mc, md) == ma + mb + mc + md).all()
cg = PermuteDims(M, [1, 0])
f = lambdify((M, ), cg, 'numpy')
assert (f(ma) == ma.T).all()
cg = PermuteDims(ArrayTensorProduct(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 = ArrayDiagonal(ArrayTensorProduct(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_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_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)
def convert_indexed_to_array(expr, first_indices=None):
r"""
Parse indexed expression into a form useful for code generation.
Examples
========
>>> from sympy.tensor.array.expressions.conv_indexed_to_array import convert_indexed_to_array
>>> from sympy import MatrixSymbol, Sum, symbols
>>> i, j, k, d = symbols("i j k d")
>>> M = MatrixSymbol("M", d, d)
>>> N = MatrixSymbol("N", d, d)
Recognize the trace in summation form:
>>> expr = Sum(M[i, i], (i, 0, d-1))
>>> convert_indexed_to_array(expr)
ArrayContraction(M, (0, 1))
Recognize the extraction of the diagonal by using the same index `i` on
both axes of the matrix:
>>> expr = M[i, i]
>>> convert_indexed_to_array(expr)
ArrayDiagonal(M, (0, 1))
This function can help perform the transformation expressed in two
different mathematical notations as:
`\sum_{j=0}^{N-1} A_{i,j} B_{j,k} \Longrightarrow \mathbf{A}\cdot \mathbf{B}`
Recognize the matrix multiplication in summation form:
>>> expr = Sum(M[i, j]*N[j, k], (j, 0, d-1))
>>> convert_indexed_to_array(expr)
ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
Specify that ``k`` has to be the starting index:
>>> convert_indexed_to_array(expr, first_indices=[k])
ArrayContraction(ArrayTensorProduct(N, M), (0, 3))
"""
result, indices = _convert_indexed_to_array(expr)
if not first_indices:
return result
for i in first_indices:
if i not in indices:
first_indices.remove(i)
first_indices.extend([i for i in indices if i not in first_indices])
permutation = [first_indices.index(i) for i in indices]
return PermuteDims(result, permutation)
def _(expr: PermuteDims):
subexpr, subremoved = _remove_trivial_dims(expr.expr)
p = expr.permutation.array_form
pinv = _af_invert(expr.permutation.array_form)
shift = list(
accumulate([1 if i in subremoved else 0 for i in range(len(p))]))
premoved = [pinv[i] for i in subremoved]
p2 = [e - shift[e] for i, e in enumerate(p) if e not in subremoved]
# TODO: check if subremoved should be permuted as well...
newexpr = PermuteDims(subexpr, p2)
if newexpr != expr:
newexpr = _array2matrix(newexpr)
return newexpr, sorted(premoved)
def test_arrayexpr_convert_array_to_implicit_matmul():
# Trivial dimensions are suppressed, so the result can be expressed in matrix form:
cg = ArrayTensorProduct(a, b)
assert convert_array_to_matrix(cg) == a * b.T
cg = ArrayTensorProduct(a, I, b)
assert convert_array_to_matrix(cg) == a * b.T
cg = ArrayContraction(ArrayTensorProduct(I, I), (1, 2))
assert convert_array_to_matrix(cg) == I
cg = PermuteDims(ArrayTensorProduct(I, Identity(1)), [0, 2, 1, 3])
assert convert_array_to_matrix(cg) == I
def test_arrayexpr_convert_array_to_matrix2():
cg = _array_contraction(_array_tensor_product(M, N), (1, 3))
assert convert_array_to_matrix(cg) == M * N.T
cg = PermuteDims(_array_tensor_product(M, N), Permutation([0, 1, 3, 2]))
assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T)
cg = _array_tensor_product(M, PermuteDims(N, Permutation([1, 0])))
assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T)
cg = _array_contraction(
PermuteDims(
_array_tensor_product(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])),
(1, 2), (3, 5)
)
assert convert_array_to_matrix(cg) == _array_tensor_product(M * P.T * Trace(N), Q.T)
cg = _array_contraction(
_array_tensor_product(M, N, P, PermuteDims(Q, Permutation([1, 0]))),
(1, 5), (2, 3)
)
assert convert_array_to_matrix(cg) == _array_tensor_product(M * P.T * Trace(N), Q.T)
cg = _array_tensor_product(M, PermuteDims(N, [1, 0]))
assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T)
cg = _array_tensor_product(PermuteDims(M, [1, 0]), PermuteDims(N, [1, 0]))
assert convert_array_to_matrix(cg) == _array_tensor_product(M.T, N.T)
cg = _array_tensor_product(PermuteDims(N, [1, 0]), PermuteDims(M, [1, 0]))
assert convert_array_to_matrix(cg) == _array_tensor_product(N.T, M.T)
cg = _array_contraction(M, (0,), (1,))
assert convert_array_to_matrix(cg) == OneMatrix(1, k)*M*OneMatrix(k, 1)
cg = _array_contraction(x, (0,), (1,))
assert convert_array_to_matrix(cg) == OneMatrix(1, k)*x
Xm = MatrixSymbol("Xm", m, n)
cg = _array_contraction(Xm, (0,), (1,))
assert convert_array_to_matrix(cg) == OneMatrix(1, m)*Xm*OneMatrix(n, 1)
def test_convert_array_to_hadamard_products():
expr = HadamardProduct(M, N)
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert ret == expr
expr = HadamardProduct(M, N) * P
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert ret == expr
expr = Q * HadamardProduct(M, N) * P
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert ret == expr
expr = Q * HadamardProduct(M, N.T) * P
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert ret == expr
expr = HadamardProduct(M, N) * HadamardProduct(Q, P)
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert expr == ret
expr = P.T * HadamardProduct(M, N) * HadamardProduct(Q, P)
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert expr == ret
# ArrayDiagonal should be converted
cg = ArrayDiagonal(ArrayTensorProduct(M, N, Q), (1, 3), (0, 2, 4))
ret = convert_array_to_matrix(cg)
expected = PermuteDims(
ArrayDiagonal(ArrayTensorProduct(HadamardProduct(M.T, N.T), Q),
(1, 2)), [1, 0, 2])
assert expected == ret
# Special case that should return the same expression:
cg = ArrayDiagonal(ArrayTensorProduct(HadamardProduct(M, N), Q), (0, 2))
ret = convert_array_to_matrix(cg)
assert ret == cg
def test_arrayexpr_array_shape():
expr = ArrayTensorProduct(M, N, P, Q)
assert expr.shape == (k, k, k, k, k, k, k, k)
Z = MatrixSymbol("Z", m, n)
expr = ArrayTensorProduct(M, Z)
assert expr.shape == (k, k, m, n)
expr2 = ArrayContraction(expr, (0, 1))
assert expr2.shape == (m, n)
expr2 = ArrayDiagonal(expr, (0, 1))
assert expr2.shape == (m, n, k)
exprp = PermuteDims(expr, [2, 1, 3, 0])
assert exprp.shape == (m, k, n, k)
expr3 = ArrayTensorProduct(N, Z)
expr2 = ArrayAdd(expr, expr3)
assert expr2.shape == (k, k, m, n)
# Contraction along axes with discordant dimensions:
raises(ValueError, lambda: ArrayContraction(expr, (1, 2)))
# Also diagonal needs the same dimensions:
raises(ValueError, lambda: ArrayDiagonal(expr, (1, 2)))
# Diagonal requires at least to axes to compute the diagonal:
raises(ValueError, lambda: ArrayDiagonal(expr, (1, )))
def test_arrayexpr_permute_tensor_product():
cg1 = PermuteDims(ArrayTensorProduct(M, N, P, Q),
Permutation([2, 3, 1, 0, 5, 4, 6, 7]))
cg2 = ArrayTensorProduct(N, PermuteDims(M, [1, 0]), PermuteDims(P, [1, 0]),
Q)
assert cg1 == cg2
# TODO: reverse operation starting with `PermuteDims` and getting down to `bb`...
cg1 = PermuteDims(ArrayTensorProduct(M, N, P, Q),
Permutation([2, 3, 4, 5, 0, 1, 6, 7]))
cg2 = ArrayTensorProduct(N, P, M, Q)
assert cg1 == cg2
cg1 = PermuteDims(ArrayTensorProduct(M, N, P, Q),
Permutation([2, 3, 4, 6, 5, 7, 0, 1]))
assert cg1.expr == ArrayTensorProduct(N, P, Q, M)
assert cg1.permutation == Permutation([0, 1, 2, 4, 3, 5, 6, 7])
cg1 = ArrayContraction(
PermuteDims(ArrayTensorProduct(N, Q, Q, M), [2, 1, 5, 4, 0, 3, 6, 7]),
[1, 2, 6])
cg2 = PermuteDims(
ArrayContraction(ArrayTensorProduct(Q, Q, N, M), (3, 5, 6)),
[0, 2, 3, 1, 4])
assert cg1 == cg2
cg1 = ArrayContraction(
ArrayContraction(
ArrayContraction(
ArrayContraction(
PermuteDims(ArrayTensorProduct(N, Q, Q, M),
[2, 1, 5, 4, 0, 3, 6, 7]), [1, 2, 6]),
[1, 3, 4]), [1]), [0])
cg2 = ArrayContraction(ArrayTensorProduct(M, N, Q, Q), (0, 3, 5),
(1, 4, 7), (2, ), (6, ))
assert cg1 == cg2