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_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_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_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_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_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: ArrayDiagonal):
newexpr, removed = _remove_trivial_dims(expr.expr)
shifts = list(
accumulate(
[0] +
[1 if i in removed else 0 for i in range(get_rank(expr.expr))]))
new_diag_indices_map = {
i: tuple(j for j in i if j not in removed)
for i in expr.diagonal_indices
}
for old_diag_tuple, new_diag_tuple in new_diag_indices_map.items():
if len(new_diag_tuple) == 1:
removed = [i for i in removed if i not in old_diag_tuple]
new_diag_indices = [
tuple(j - shifts[j] for j in i) for i in new_diag_indices_map.values()
]
rank = get_rank(expr.expr)
removed = ArrayDiagonal._push_indices_up(expr.diagonal_indices, removed,
rank)
removed = sorted({i for i in removed})
# If there are single axes to diagonalize remaining, it means that their
# corresponding dimension has been removed, they no longer need diagonalization:
new_diag_indices = [i for i in new_diag_indices if len(i) > 0]
if len(new_diag_indices) > 0:
newexpr2 = _array_diagonal(newexpr,
*new_diag_indices,
allow_trivial_diags=True)
else:
newexpr2 = newexpr
if isinstance(newexpr2, ArrayDiagonal):
newexpr3, removed2 = _remove_diagonalized_identity_matrices(newexpr2)
removed = _combine_removed(-1, removed, removed2)
return newexpr3, removed
else:
return newexpr2, removed
def _(expr: ArrayDiagonal):
expr2 = _array2matrix(expr.expr)
pexpr = _array_diag2contr_diagmatrix(
ArrayDiagonal(expr2, *expr.diagonal_indices))
if expr == pexpr:
return expr
return _array2matrix(pexpr)
def _(expr: ElementwiseApplyFunction, x: Expr):
assert get_rank(expr) == 2
assert get_rank(x) == 2
fdiff = expr._get_function_fdiff()
dexpr = array_derive(expr.expr, x)
tp = ArrayTensorProduct(ElementwiseApplyFunction(fdiff, expr.expr), dexpr)
td = ArrayDiagonal(tp, (0, 4), (1, 5))
return td
def _(expr: ArrayDiagonal):
pexpr = ArrayDiagonal(_array2matrix(expr.expr), *expr.diagonal_indices)
pexpr = identify_hadamard_products(pexpr)
if isinstance(pexpr, ArrayDiagonal):
pexpr = _array_diag2contr_diagmatrix(pexpr)
if expr == pexpr:
return expr
return _array2matrix(pexpr)
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 _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 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_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_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 _(expr: ArrayElementwiseApplyFunc, x: Expr):
fdiff = expr._get_function_fdiff()
subexpr = expr.expr
dsubexpr = array_derive(subexpr, x)
tp = ArrayTensorProduct(dsubexpr,
ArrayElementwiseApplyFunc(fdiff, subexpr))
b = get_rank(x)
c = get_rank(expr)
diag_indices = [(b + i, b + c + i) for i in range(c)]
return ArrayDiagonal(tp, *diag_indices)
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 _(expr: ArrayDiagonal):
newexpr, removed = _remove_trivial_dims(expr.expr)
shifts = list(
accumulate(
[0] +
[1 if i in removed else 0 for i in range(get_rank(expr.expr))]))
new_diag_indices = [
tuple(j for j in i if j not in removed) for i in expr.diagonal_indices
]
new_diag_indices = [
tuple(j - shifts[j] for j in i) for i in new_diag_indices
]
rank = get_rank(expr.expr)
removed = ArrayDiagonal._push_indices_up(expr.diagonal_indices, removed,
rank)
removed = sorted({i for i in removed})
# If there are single axes to diagonalize remaining, it means that their
# corresponding dimension has been removed, they no longer need diagonalization:
new_diag_indices = [i for i in new_diag_indices if len(i) > 1]
return ArrayDiagonal(newexpr, *new_diag_indices), removed
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 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_convert_array_element_to_array_expression():
A = ArraySymbol("A", (k, ))
B = ArraySymbol("B", (k, ))
s = Sum(A[i] * B[i], (i, 0, k - 1))
cg = convert_indexed_to_array(s)
assert cg == ArrayContraction(ArrayTensorProduct(A, B), (0, 1))
s = A[i] * B[i]
cg = convert_indexed_to_array(s)
assert cg == ArrayDiagonal(ArrayTensorProduct(A, B), (0, 1))
s = A[i] * B[j]
cg = convert_indexed_to_array(s, [i, j])
assert cg == ArrayTensorProduct(A, B)
cg = convert_indexed_to_array(s, [j, i])
assert cg == ArrayTensorProduct(B, A)
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_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_arrayexpr_normalize_diagonal_contraction():
tp = ArrayTensorProduct(M, N, P, Q)
expr = ArrayContraction(ArrayDiagonal(tp, (1, 3, 4)), (0, 3))
result = ArrayDiagonal(ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 6)), (0, 2, 3))
assert expr == result
expr = ArrayContraction(ArrayDiagonal(tp, (0, 1, 2, 3, 7)), (1, 2, 3))
result = ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 1, 2, 3, 5, 6, 7))
assert expr == result
expr = ArrayContraction(ArrayDiagonal(tp, (0, 2, 6, 7)), (1, 2, 3))
result = ArrayDiagonal(ArrayContraction(tp, (3, 4, 5)), (0, 2, 3, 4))
assert expr == result
td = ArrayDiagonal(ArrayTensorProduct(M, N, P, Q), (0, 3))
expr = ArrayContraction(td, (2, 1), (0, 4, 6, 5, 3))
result = ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 1, 3, 5, 6, 7), (2, 4))
assert expr == result
def _remove_diagonalized_identity_matrices(expr: ArrayDiagonal):
assert isinstance(expr, ArrayDiagonal)
editor = _EditArrayContraction(expr)
mapping = {
i: {j
for j in editor.args_with_ind if i in j.indices}
for i in range(-1, -1 - editor.number_of_diagonal_indices, -1)
}
removed = []
counter: int = 0
for i, arg_with_ind in enumerate(editor.args_with_ind):
counter += len(arg_with_ind.indices)
if isinstance(arg_with_ind.element, Identity):
if None in arg_with_ind.indices and any(
i is not None and (i < 0) == True
for i in arg_with_ind.indices):
diag_ind = [j for j in arg_with_ind.indices
if j is not None][0]
other = [j for j in mapping[diag_ind] if j != arg_with_ind][0]
if not isinstance(other.element, MatrixExpr):
continue
if 1 not in other.element.shape:
continue
if None not in other.indices:
continue
editor.args_with_ind[i].element = None
none_index = other.indices.index(None)
other.element = DiagMatrix(other.element)
other_range = editor.get_absolute_range(other)
removed.extend([other_range[0] + none_index])
editor.args_with_ind = [
i for i in editor.args_with_ind if i.element is not None
]
removed = ArrayDiagonal._push_indices_up(expr.diagonal_indices, removed,
get_rank(expr.expr))
return editor.to_array_contraction(), removed
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
# Hadamard products with traces:
expr = Trace(HadamardProduct(M, N))
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert ret == Trace(HadamardProduct(M.T, N.T))
expr = Trace(A*HadamardProduct(M, N))
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert ret == Trace(HadamardProduct(M, N)*A)
expr = Trace(HadamardProduct(A, M)*N)
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert ret == Trace(HadamardProduct(M.T, N)*A)
# These should not be converted into Hadamard products:
cg = ArrayDiagonal(ArrayTensorProduct(M, N), (0, 1, 2, 3))
ret = convert_array_to_matrix(cg)
assert ret == cg
cg = ArrayDiagonal(ArrayTensorProduct(A), (0, 1))
ret = convert_array_to_matrix(cg)
assert ret == cg
cg = ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 2, 4), (1, 3, 5))
assert convert_array_to_matrix(cg) == HadamardProduct(M, N, P)
cg = ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 3, 4), (1, 2, 5))
assert convert_array_to_matrix(cg) == HadamardProduct(M, P, N.T)
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, [])
# A few more cases to test the removal and shift of nested removed axes
# with array contractions and array diagonals:
tp = ArrayTensorProduct(
OneMatrix(1, 1),
M,
x,
OneMatrix(1, 1),
Identity(1),
)
expr = ArrayContraction(tp, (1, 8))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [0, 5, 6, 7]
expr = ArrayContraction(tp, (1, 8), (3, 4))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [0, 3, 4, 5]
expr = ArrayDiagonal(tp, (1, 8))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [0, 5, 6, 7, 8]
expr = ArrayDiagonal(tp, (1, 8), (3, 4))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [0, 3, 4, 5, 6]
expr = ArrayDiagonal(ArrayContraction(ArrayTensorProduct(A, x, I, I1), (1, 2, 5)), (1, 4))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [1, 2]
cg = ArrayDiagonal(ArrayTensorProduct(PermuteDims(ArrayTensorProduct(x, I1), Permutation(1, 2, 3)), (x.T*x).applyfunc(sqrt)), (2, 4), (3, 5))
rexpr, removed = _remove_trivial_dims(cg)
assert removed == [1, 2]
# Contractions with identity matrices need to be followed by a permutation
# in order
cg = ArrayContraction(ArrayTensorProduct(A, B, C, M, I), (1, 8))
ret, removed = _remove_trivial_dims(cg)
assert ret == PermuteDims(ArrayTensorProduct(A, B, C, M), [0, 2, 3, 4, 5, 6, 7, 1])
assert removed == []
cg = ArrayContraction(ArrayTensorProduct(A, B, C, M, I), (1, 8), (3, 4))
ret, removed = _remove_trivial_dims(cg)
assert ret == PermuteDims(ArrayContraction(ArrayTensorProduct(A, B, C, M), (3, 4)), [0, 2, 3, 4, 5, 1])
assert removed == []
def test_arrayexpr_convert_array_to_diagonalized_vector():
# Check matrix recognition over trivial dimensions:
cg = ArrayTensorProduct(a, b)
assert convert_array_to_matrix(cg) == a * b.T
cg = ArrayTensorProduct(I1, a, b)
assert convert_array_to_matrix(cg) == a * b.T
# Recognize trace inside a tensor product:
cg = ArrayContraction(ArrayTensorProduct(A, B, C), (0, 3), (1, 2))
assert convert_array_to_matrix(cg) == Trace(A * B) * C
# Transform diagonal operator to contraction:
cg = ArrayDiagonal(ArrayTensorProduct(A, a), (1, 2))
assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(A, OneArray(1), DiagMatrix(a)), (1, 3))
assert convert_array_to_matrix(cg) == A * DiagMatrix(a)
cg = ArrayDiagonal(ArrayTensorProduct(a, b), (0, 2))
assert _array_diag2contr_diagmatrix(cg) == PermuteDims(
ArrayContraction(ArrayTensorProduct(DiagMatrix(a), OneArray(1), b), (0, 3)), [1, 2, 0]
)
assert convert_array_to_matrix(cg) == b.T * DiagMatrix(a)
cg = ArrayDiagonal(ArrayTensorProduct(A, a), (0, 2))
assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(A, OneArray(1), DiagMatrix(a)), (0, 3))
assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a)
cg = ArrayDiagonal(ArrayTensorProduct(I, x, I1), (0, 2), (3, 5))
assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(I, OneArray(1), I1, DiagMatrix(x)), (0, 5))
assert convert_array_to_matrix(cg) == DiagMatrix(x)
cg = ArrayDiagonal(ArrayTensorProduct(I, x, A, B), (1, 2), (5, 6))
assert _array_diag2contr_diagmatrix(cg) == ArrayDiagonal(ArrayContraction(ArrayTensorProduct(I, OneArray(1), A, B, DiagMatrix(x)), (1, 7)), (5, 6))
# TODO: this is returning a wrong result:
# convert_array_to_matrix(cg)
cg = ArrayDiagonal(ArrayTensorProduct(I1, a, b), (1, 3, 5))
assert convert_array_to_matrix(cg) == a*b.T
cg = ArrayDiagonal(ArrayTensorProduct(I1, a, b), (1, 3))
assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(OneArray(1), a, b, I1), (2, 6))
assert convert_array_to_matrix(cg) == a*b.T
cg = ArrayDiagonal(ArrayTensorProduct(x, I1), (1, 2))
assert isinstance(cg, ArrayDiagonal)
assert cg.diagonal_indices == ((1, 2),)
assert convert_array_to_matrix(cg) == x
cg = ArrayDiagonal(ArrayTensorProduct(x, I), (0, 2))
assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(OneArray(1), I, DiagMatrix(x)), (1, 3))
assert convert_array_to_matrix(cg).doit() == DiagMatrix(x)
raises(ValueError, lambda: ArrayDiagonal(x, (1,)))
# Ignore identity matrices with contractions:
cg = ArrayContraction(ArrayTensorProduct(I, A, I, I), (0, 2), (1, 3), (5, 7))
assert cg.split_multiple_contractions() == cg
assert convert_array_to_matrix(cg) == Trace(A) * I
cg = ArrayContraction(ArrayTensorProduct(Trace(A) * I, I, I), (1, 5), (3, 4))
assert cg.split_multiple_contractions() == cg
assert convert_array_to_matrix(cg).doit() == Trace(A) * I
# Add DiagMatrix when required:
cg = ArrayContraction(ArrayTensorProduct(A, a), (1, 2))
assert cg.split_multiple_contractions() == cg
assert convert_array_to_matrix(cg) == A * a
cg = ArrayContraction(ArrayTensorProduct(A, a, B), (1, 2, 4))
assert cg.split_multiple_contractions() == ArrayContraction(ArrayTensorProduct(A, DiagMatrix(a), OneArray(1), B), (1, 2), (3, 5))
assert convert_array_to_matrix(cg) == A * DiagMatrix(a) * B
cg = ArrayContraction(ArrayTensorProduct(A, a, B), (0, 2, 4))
assert cg.split_multiple_contractions() == ArrayContraction(ArrayTensorProduct(A, DiagMatrix(a), OneArray(1), B), (0, 2), (3, 5))
assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) * B
cg = ArrayContraction(ArrayTensorProduct(A, a, b, a.T, B), (0, 2, 4, 7, 9))
assert cg.split_multiple_contractions() == ArrayContraction(ArrayTensorProduct(A, DiagMatrix(a), OneArray(1),
DiagMatrix(b), OneArray(1), DiagMatrix(a), OneArray(1), B),
(0, 2), (3, 5), (6, 9), (8, 12))
assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) * DiagMatrix(b) * DiagMatrix(a) * B.T
cg = ArrayContraction(ArrayTensorProduct(I1, I1, I1), (1, 2, 4))
assert cg.split_multiple_contractions() == ArrayContraction(ArrayTensorProduct(I1, I1, OneArray(1), I1), (1, 2), (3, 5))
assert convert_array_to_matrix(cg) == 1
cg = ArrayContraction(ArrayTensorProduct(I, I, I, I, A), (1, 2, 8), (5, 6, 9))
assert convert_array_to_matrix(cg.split_multiple_contractions()).doit() == A
cg = ArrayContraction(ArrayTensorProduct(A, a, C, a, B), (1, 2, 4), (5, 6, 8))
expected = ArrayContraction(ArrayTensorProduct(A, DiagMatrix(a), OneArray(1), C, DiagMatrix(a), OneArray(1), B), (1, 3), (2, 5), (6, 7), (8, 10))
assert cg.split_multiple_contractions() == expected
assert convert_array_to_matrix(cg) == A * DiagMatrix(a) * C * DiagMatrix(a) * B
cg = ArrayContraction(ArrayTensorProduct(a, I1, b, I1, (a.T*b).applyfunc(cos)), (1, 2, 8), (5, 6, 9))
expected = ArrayContraction(ArrayTensorProduct(a, I1, OneArray(1), b, I1, OneArray(1), (a.T*b).applyfunc(cos)),
(1, 3), (2, 10), (6, 8), (7, 11))
assert cg.split_multiple_contractions().dummy_eq(expected)
assert convert_array_to_matrix(cg).doit().dummy_eq(MatMul(a, (a.T * b).applyfunc(cos), b.T))
def test_array_expressions_no_canonicalization():
tp = _array_tensor_product(M, N, P)
# ArrayTensorProduct:
expr = ArrayTensorProduct(tp, N)
assert str(expr) == "ArrayTensorProduct(ArrayTensorProduct(M, N, P), N)"
assert expr.doit() == ArrayTensorProduct(M, N, P, N)
expr = ArrayTensorProduct(ArrayContraction(M, (0, 1)), N)
assert str(expr) == "ArrayTensorProduct(ArrayContraction(M, (0, 1)), N)"
assert expr.doit() == ArrayContraction(ArrayTensorProduct(M, N), (0, 1))
expr = ArrayTensorProduct(ArrayDiagonal(M, (0, 1)), N)
assert str(expr) == "ArrayTensorProduct(ArrayDiagonal(M, (0, 1)), N)"
assert expr.doit() == PermuteDims(ArrayDiagonal(ArrayTensorProduct(M, N), (0, 1)), [2, 0, 1])
expr = ArrayTensorProduct(PermuteDims(M, [1, 0]), N)
assert str(expr) == "ArrayTensorProduct(PermuteDims(M, (0 1)), N)"
assert expr.doit() == PermuteDims(ArrayTensorProduct(M, N), [1, 0, 2, 3])
# ArrayContraction:
expr = ArrayContraction(_array_contraction(tp, (0, 2)), (0, 1))
assert isinstance(expr, ArrayContraction)
assert isinstance(expr.expr, ArrayContraction)
assert str(expr) == "ArrayContraction(ArrayContraction(ArrayTensorProduct(M, N, P), (0, 2)), (0, 1))"
assert expr.doit() == ArrayContraction(tp, (0, 2), (1, 3))
expr = ArrayContraction(ArrayContraction(ArrayContraction(tp, (0, 1)), (0, 1)), (0, 1))
assert expr.doit() == ArrayContraction(tp, (0, 1), (2, 3), (4, 5))
# assert expr._canonicalize() == ArrayContraction(ArrayContraction(tp, (0, 1)), (0, 1), (2, 3))
expr = ArrayContraction(ArrayDiagonal(tp, (0, 1)), (0, 1))
assert str(expr) == "ArrayContraction(ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 1)), (0, 1))"
assert expr.doit() == ArrayDiagonal(ArrayContraction(ArrayTensorProduct(N, M, P), (0, 1)), (0, 1))
expr = ArrayContraction(PermuteDims(M, [1, 0]), (0, 1))
assert str(expr) == "ArrayContraction(PermuteDims(M, (0 1)), (0, 1))"
assert expr.doit() == ArrayContraction(M, (0, 1))
# ArrayDiagonal:
expr = ArrayDiagonal(ArrayDiagonal(tp, (0, 2)), (0, 1))
assert str(expr) == "ArrayDiagonal(ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 2)), (0, 1))"
assert expr.doit() == ArrayDiagonal(tp, (0, 2), (1, 3))
expr = ArrayDiagonal(ArrayDiagonal(ArrayDiagonal(tp, (0, 1)), (0, 1)), (0, 1))
assert expr.doit() == ArrayDiagonal(tp, (0, 1), (2, 3), (4, 5))
assert expr._canonicalize() == expr.doit()
expr = ArrayDiagonal(ArrayContraction(tp, (0, 1)), (0, 1))
assert str(expr) == "ArrayDiagonal(ArrayContraction(ArrayTensorProduct(M, N, P), (0, 1)), (0, 1))"
assert expr.doit() == expr
expr = ArrayDiagonal(PermuteDims(M, [1, 0]), (0, 1))
assert str(expr) == "ArrayDiagonal(PermuteDims(M, (0 1)), (0, 1))"
assert expr.doit() == ArrayDiagonal(M, (0, 1))
# ArrayAdd:
expr = ArrayAdd(M)
assert isinstance(expr, ArrayAdd)
assert expr.doit() == M
expr = ArrayAdd(ArrayAdd(M, N), P)
assert str(expr) == "ArrayAdd(ArrayAdd(M, N), P)"
assert expr.doit() == ArrayAdd(M, N, P)
expr = ArrayAdd(M, ArrayAdd(N, ArrayAdd(P, M)))
assert expr.doit() == ArrayAdd(M, N, P, M)
assert expr._canonicalize() == ArrayAdd(M, N, ArrayAdd(P, M))
expr = ArrayAdd(M, ZeroArray(k, k), N)
assert str(expr) == "ArrayAdd(M, ZeroArray(k, k), N)"
assert expr.doit() == ArrayAdd(M, N)
# PermuteDims:
expr = PermuteDims(PermuteDims(M, [1, 0]), [1, 0])
assert str(expr) == "PermuteDims(PermuteDims(M, (0 1)), (0 1))"
assert expr.doit() == M
expr = PermuteDims(PermuteDims(PermuteDims(M, [1, 0]), [1, 0]), [1, 0])
assert expr.doit() == PermuteDims(M, [1, 0])
assert expr._canonicalize() == expr.doit()