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 _array_tensor_product(M, N, za1) == ZeroArray(k, k, k, k, k, l, m, n)
assert _array_tensor_product(M, N, zm1) == ZeroArray(k, k, k, k, m, n)
assert _array_contraction(za1, (3,)) == ZeroArray(k, l, m)
assert _array_contraction(zm1, (1,)) == ZeroArray(m)
assert _array_contraction(za2, (1, 2)) == ZeroArray(k, n)
assert _array_contraction(zm2, (0, 1)) == 0
assert _array_diagonal(za2, (1, 2)) == ZeroArray(k, n, m)
assert _array_diagonal(zm2, (0, 1)) == ZeroArray(m)
assert _permute_dims(za1, [2, 1, 3, 0]) == ZeroArray(m, l, n, k)
assert _permute_dims(zm1, [1, 0]) == ZeroArray(n, m)
assert _array_add(za1) == za1
assert _array_add(zm1) == ZeroArray(m, n)
tp1 = _array_tensor_product(MatrixSymbol("A", k, l), MatrixSymbol("B", m, n))
assert _array_add(tp1, za1) == tp1
tp2 = _array_tensor_product(MatrixSymbol("C", k, l), MatrixSymbol("D", m, n))
assert _array_add(tp1, za1, tp2) == _array_add(tp1, tp2)
assert _array_add(M, zm3) == M
assert _array_add(M, N, zm3) == _array_add(M, N)
def test_arrayexpr_canonicalize_diagonal__permute_dims():
tp = _array_tensor_product(M, Q, N, P)
expr = _array_diagonal(
_permute_dims(tp, [0, 1, 2, 4, 7, 6, 3, 5]), (2, 4, 5), (6, 7),
(0, 3))
result = _array_diagonal(tp, (2, 6, 7), (3, 5), (0, 4))
assert expr == result
tp = _array_tensor_product(M, N, P, Q)
expr = _array_diagonal(_permute_dims(tp, [0, 5, 2, 4, 1, 6, 3, 7]), (1, 2, 6), (3, 4))
result = _array_diagonal(_array_tensor_product(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_nested_array_elementwise_add():
cg = _array_contraction(
_array_add(_array_tensor_product(M, N), _array_tensor_product(N, M)),
(1, 2))
result = _array_add(
_array_contraction(_array_tensor_product(M, N), (1, 2)),
_array_contraction(_array_tensor_product(N, M), (1, 2)))
assert cg == result
cg = _array_diagonal(
_array_add(_array_tensor_product(M, N), _array_tensor_product(N, M)),
(1, 2))
result = _array_add(_array_diagonal(_array_tensor_product(M, N), (1, 2)),
_array_diagonal(_array_tensor_product(N, M), (1, 2)))
assert cg == result
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):
pexpr = _array_diagonal(_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 _(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 = _array_tensor_product(ElementwiseApplyFunction(fdiff, expr.expr),
dexpr)
td = _array_diagonal(tp, (0, 4), (1, 5))
return td
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 = _array_contraction(_array_tensor_product(*args), *contr_indices2)
td = _array_diagonal(tc, *diag_indices_new)
return td
return expr
def test_arrayexpr_array_diagonal():
cg = _array_diagonal(M, (1, 0))
assert cg == _array_diagonal(M, (0, 1))
cg = _array_diagonal(_array_tensor_product(M, N, P), (4, 1), (2, 0))
assert cg == _array_diagonal(_array_tensor_product(M, N, P), (1, 4), (0, 2))
cg = _array_diagonal(_array_tensor_product(M, N), (1, 2), (3,), allow_trivial_diags=True)
assert cg == _permute_dims(_array_diagonal(_array_tensor_product(M, N), (1, 2)), [0, 2, 1])
Ax = ArraySymbol("Ax", shape=(1, 2, 3, 4, 3, 5, 6, 2, 7))
cg = _array_diagonal(Ax, (1, 7), (3,), (2, 4), (6,), allow_trivial_diags=True)
assert cg == _permute_dims(_array_diagonal(Ax, (1, 7), (2, 4)), [0, 2, 4, 5, 1, 6, 3])
cg = _array_diagonal(M, (0,), allow_trivial_diags=True)
assert cg == _permute_dims(M, [1, 0])
raises(ValueError, lambda: _array_diagonal(M, (0, 0)))
def _(expr: ArrayElementwiseApplyFunc, x: Expr):
fdiff = expr._get_function_fdiff()
subexpr = expr.expr
dsubexpr = array_derive(subexpr, x)
tp = _array_tensor_product(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 _array_diagonal(tp, *diag_indices)
def test_arrayexpr_array_shape():
expr = _array_tensor_product(M, N, P, Q)
assert expr.shape == (k, k, k, k, k, k, k, k)
Z = MatrixSymbol("Z", m, n)
expr = _array_tensor_product(M, Z)
assert expr.shape == (k, k, m, n)
expr2 = _array_contraction(expr, (0, 1))
assert expr2.shape == (m, n)
expr2 = _array_diagonal(expr, (0, 1))
assert expr2.shape == (m, n, k)
exprp = _permute_dims(expr, [2, 1, 3, 0])
assert exprp.shape == (m, k, n, k)
expr3 = _array_tensor_product(N, Z)
expr2 = _array_add(expr, expr3)
assert expr2.shape == (k, k, m, n)
# Contraction along axes with discordant dimensions:
raises(ValueError, lambda: _array_contraction(expr, (1, 2)))
# Also diagonal needs the same dimensions:
raises(ValueError, lambda: _array_diagonal(expr, (1, 2)))
# Diagonal requires at least to axes to compute the diagonal:
raises(ValueError, lambda: _array_diagonal(expr, (1,)))
def test_arrayexpr_convert_array_to_matrix_diag2contraction_diagmatrix():
cg = _array_diagonal(_array_tensor_product(M, a), (1, 2))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == _array_contraction(
_array_tensor_product(M, OneArray(1), DiagMatrix(a)), (1, 3))
raises(ValueError, lambda: _array_diagonal(_array_tensor_product(a, M),
(1, 2)))
cg = _array_diagonal(_array_tensor_product(a.T, M), (1, 2))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == _array_contraction(
_array_tensor_product(OneArray(1), M, DiagMatrix(a.T)), (1, 4))
cg = _array_diagonal(_array_tensor_product(a.T, M, N, b.T), (1, 2), (4, 7))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == _array_contraction(
_array_tensor_product(OneArray(1), M, N, OneArray(1), DiagMatrix(a.T),
DiagMatrix(b.T)), (1, 7), (3, 9))
cg = _array_diagonal(_array_tensor_product(a, M, N, b.T), (0, 2), (4, 7))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == _array_contraction(
_array_tensor_product(OneArray(1), M, N, OneArray(1), DiagMatrix(a),
DiagMatrix(b.T)), (1, 6), (3, 9))
cg = _array_diagonal(_array_tensor_product(a, M, N, b.T), (0, 4), (3, 7))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == _array_contraction(
_array_tensor_product(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 = _array_diagonal(_array_tensor_product(x, A.T, I1), (0, 2))
assert _array_diag2contr_diagmatrix(cg).shape == cg.shape
assert _array2matrix(cg).shape == cg.shape
def test_arrayexpr_canonicalize_diagonal_contraction():
tp = _array_tensor_product(M, N, P, Q)
expr = _array_contraction(_array_diagonal(tp, (1, 3, 4)), (0, 3))
result = _array_diagonal(_array_contraction(_array_tensor_product(M, N, P, Q), (0, 6)), (0, 2, 3))
assert expr == result
expr = _array_contraction(_array_diagonal(tp, (0, 1, 2, 3, 7)), (1, 2, 3))
result = _array_contraction(_array_tensor_product(M, N, P, Q), (0, 1, 2, 3, 5, 6, 7))
assert expr == result
expr = _array_contraction(_array_diagonal(tp, (0, 2, 6, 7)), (1, 2, 3))
result = _array_diagonal(_array_contraction(tp, (3, 4, 5)), (0, 2, 3, 4))
assert expr == result
td = _array_diagonal(_array_tensor_product(M, N, P, Q), (0, 3))
expr = _array_contraction(td, (2, 1), (0, 4, 6, 5, 3))
result = _array_contraction(_array_tensor_product(M, N, P, Q), (0, 1, 3, 5, 6, 7), (2, 4))
assert expr == result
def convert_matrix_to_array(expr: Basic) -> 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 = _array_tensor_product(*args_nonmat) if args_nonmat else S.One
if scalar == 1:
tprod = _array_tensor_product(
*[convert_matrix_to_array(arg) for arg in args])
else:
tprod = _array_tensor_product(
scalar,
*[convert_matrix_to_array(arg) for arg in args])
return _array_contraction(
tprod,
*contractions
)
elif isinstance(expr, MatAdd):
return _array_add(
*[convert_matrix_to_array(arg) for arg in expr.args]
)
elif isinstance(expr, Transpose):
return _permute_dims(
convert_matrix_to_array(expr.args[0]), [1, 0]
)
elif isinstance(expr, Trace):
inner_expr: MatrixExpr = convert_matrix_to_array(expr.arg) # type: ignore
return _array_contraction(inner_expr, (0, len(inner_expr.shape) - 1))
elif isinstance(expr, Mul):
return _array_tensor_product(*[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 _array_tensor_product(*[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 = _array_tensor_product(*[convert_matrix_to_array(arg) for arg in expr.args])
diag = [[2*i for i in range(len(expr.args))], [2*i+1 for i in range(len(expr.args))]]
return _array_diagonal(tp, *diag)
elif isinstance(expr, HadamardPower):
base, exp = expr.args
if isinstance(exp, Integer) and exp > 0:
return convert_matrix_to_array(HadamardProduct.fromiter(base for i in range(exp)))
else:
d = Dummy("d")
return ArrayElementwiseApplyFunc(Lambda(d, d**exp), base)
elif isinstance(expr, KroneckerProduct):
kp_args = [convert_matrix_to_array(arg) for arg in expr.args]
permutation = [2*i for i in range(len(kp_args))] + [2*i + 1 for i in range(len(kp_args))]
return Reshape(_permute_dims(_array_tensor_product(*kp_args), permutation), expr.shape)
else:
return expr
def _(expr: ArrayDiagonal, x: Expr):
dsubexpr = array_derive(expr.expr, x)
rank_x = len(get_shape(x))
diag_indices = [[j + rank_x for j in i] for i in expr.diagonal_indices]
return _array_diagonal(dsubexpr, *diag_indices)
def tensordiagonal(array, *diagonal_axes):
"""
Diagonalization of an array-like object on the specified axes.
This is equivalent to multiplying the expression by Kronecker deltas
uniting the axes.
The diagonal indices are put at the end of the axes.
Examples
========
``tensordiagonal`` acting on a 2-dimensional array by axes 0 and 1 is
equivalent to the diagonal of the matrix:
>>> from sympy import Array, tensordiagonal
>>> from sympy import Matrix, eye
>>> tensordiagonal(eye(3), (0, 1))
[1, 1, 1]
>>> from sympy.abc import a,b,c,d
>>> m1 = Matrix([[a, b], [c, d]])
>>> tensordiagonal(m1, [0, 1])
[a, d]
In case of higher dimensional arrays, the diagonalized out dimensions
are appended removed and appended as a single dimension at the end:
>>> A = Array(range(18), (3, 2, 3))
>>> A
[[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]], [[12, 13, 14], [15, 16, 17]]]
>>> tensordiagonal(A, (0, 2))
[[0, 7, 14], [3, 10, 17]]
>>> from sympy import permutedims
>>> tensordiagonal(A, (0, 2)) == permutedims(Array([A[0, :, 0], A[1, :, 1], A[2, :, 2]]), [1, 0])
True
"""
if any(len(i) <= 1 for i in diagonal_axes):
raise ValueError("need at least two axes to diagonalize")
from sympy.tensor.array.expressions.array_expressions import _ArrayExpr
from sympy.tensor.array.expressions.array_expressions import _CodegenArrayAbstract
from sympy.tensor.array.expressions.array_expressions import ArrayDiagonal, _array_diagonal
from sympy.matrices.expressions.matexpr import MatrixSymbol
if isinstance(array, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)):
return _array_diagonal(array, *diagonal_axes)
ArrayDiagonal._validate(array, *diagonal_axes)
array, remaining_indices, remaining_shape, diagonal_deltas = _util_contraction_diagonal(
array, *diagonal_axes)
# Compute the diagonalized array:
#
# 1. external for loops on all undiagonalized indices.
# Undiagonalized indices are determined by the combinatorial product of
# the absolute positions of the remaining indices.
# 2. internal loop on all diagonal indices.
# It appends the values of the absolute diagonalized index and the absolute
# undiagonalized index for the external loop.
diagonalized_array = []
diagonal_shape = [len(i) for i in diagonal_deltas]
for icontrib in itertools.product(*remaining_indices):
index_base_position = sum(icontrib)
isum = []
for sum_to_index in itertools.product(*diagonal_deltas):
idx = array._get_tuple_index(index_base_position +
sum(sum_to_index))
isum.append(array[idx])
isum = type(array)(isum).reshape(*diagonal_shape)
diagonalized_array.append(isum)
return type(array)(diagonalized_array, remaining_shape + diagonal_shape)
def _convert_indexed_to_array(expr):
if isinstance(expr, Sum):
function = expr.function
summation_indices = expr.variables
subexpr, subindices = _convert_indexed_to_array(function)
subindicessets = {
j: i
for i in subindices if isinstance(i, frozenset) for j in i
}
summation_indices = sorted(set(
[subindicessets.get(i, i) for i in summation_indices]),
key=default_sort_key)
# TODO: check that Kronecker delta is only contracted to one other element:
kronecker_indices = set([])
if isinstance(function, Mul):
for arg in function.args:
if not isinstance(arg, KroneckerDelta):
continue
arg_indices = sorted(set(arg.indices), key=default_sort_key)
if len(arg_indices) == 2:
kronecker_indices.update(arg_indices)
kronecker_indices = sorted(kronecker_indices, key=default_sort_key)
# Check dimensional consistency:
shape = get_shape(subexpr)
if shape:
for ind, istart, iend in expr.limits:
i = _get_argindex(subindices, ind)
if istart != 0 or iend + 1 != shape[i]:
raise ValueError(
"summation index and array dimension mismatch: %s" %
ind)
contraction_indices = []
subindices = list(subindices)
if isinstance(subexpr, ArrayDiagonal):
diagonal_indices = list(subexpr.diagonal_indices)
dindices = subindices[-len(diagonal_indices):]
subindices = subindices[:-len(diagonal_indices)]
for index in summation_indices:
if index in dindices:
position = dindices.index(index)
contraction_indices.append(diagonal_indices[position])
diagonal_indices[position] = None
diagonal_indices = [i for i in diagonal_indices if i is not None]
for i, ind in enumerate(subindices):
if ind in summation_indices:
pass
if diagonal_indices:
subexpr = _array_diagonal(subexpr.expr, *diagonal_indices)
else:
subexpr = subexpr.expr
axes_contraction = defaultdict(list)
for i, ind in enumerate(subindices):
include = all(j not in kronecker_indices
for j in ind) if isinstance(
ind, frozenset) else ind not in kronecker_indices
if ind in summation_indices and include:
axes_contraction[ind].append(i)
subindices[i] = None
for k, v in axes_contraction.items():
if any(i in kronecker_indices for i in k) if isinstance(
k, frozenset) else k in kronecker_indices:
continue
contraction_indices.append(tuple(v))
free_indices = [i for i in subindices if i is not None]
indices_ret = list(free_indices)
indices_ret.sort(key=lambda x: free_indices.index(x))
return _array_contraction(
subexpr, *contraction_indices,
free_indices=free_indices), tuple(indices_ret)
if isinstance(expr, Mul):
args, indices = zip(
*[_convert_indexed_to_array(arg) for arg in expr.args])
# Check if there are KroneckerDelta objects:
kronecker_delta_repl = {}
for arg in args:
if not isinstance(arg, KroneckerDelta):
continue
# Diagonalize two indices:
i, j = arg.indices
kindices = set(arg.indices)
if i in kronecker_delta_repl:
kindices.update(kronecker_delta_repl[i])
if j in kronecker_delta_repl:
kindices.update(kronecker_delta_repl[j])
kindices = frozenset(kindices)
for index in kindices:
kronecker_delta_repl[index] = kindices
# Remove KroneckerDelta objects, their relations should be handled by
# ArrayDiagonal:
newargs = []
newindices = []
for arg, loc_indices in zip(args, indices):
if isinstance(arg, KroneckerDelta):
continue
newargs.append(arg)
newindices.append(loc_indices)
flattened_indices = [
kronecker_delta_repl.get(j, j) for i in newindices for j in i
]
diagonal_indices, ret_indices = _get_diagonal_indices(
flattened_indices)
tp = _array_tensor_product(*newargs)
if diagonal_indices:
return _array_diagonal(tp, *diagonal_indices), ret_indices
else:
return tp, ret_indices
if isinstance(expr, MatrixElement):
indices = expr.args[1:]
diagonal_indices, ret_indices = _get_diagonal_indices(indices)
if diagonal_indices:
return _array_diagonal(expr.args[0],
*diagonal_indices), ret_indices
else:
return expr.args[0], ret_indices
if isinstance(expr, ArrayElement):
indices = expr.indices
diagonal_indices, ret_indices = _get_diagonal_indices(indices)
if diagonal_indices:
return _array_diagonal(expr.name, *diagonal_indices), ret_indices
else:
return expr.name, ret_indices
if isinstance(expr, Indexed):
indices = expr.indices
diagonal_indices, ret_indices = _get_diagonal_indices(indices)
if diagonal_indices:
return _array_diagonal(expr.base, *diagonal_indices), ret_indices
else:
return expr.args[0], ret_indices
if isinstance(expr, IndexedBase):
raise NotImplementedError
if isinstance(expr, KroneckerDelta):
return expr, expr.indices
if isinstance(expr, Add):
args, indices = zip(
*[_convert_indexed_to_array(arg) for arg in expr.args])
args = list(args)
# Check if all indices are compatible. Otherwise expand the dimensions:
index0 = []
shape0 = []
for arg, arg_indices in zip(args, indices):
arg_indices_set = set(arg_indices)
arg_indices_missing = arg_indices_set.difference(index0)
index0.extend([i for i in arg_indices if i in arg_indices_missing])
arg_shape = get_shape(arg)
shape0.extend([
arg_shape[i] for i, e in enumerate(arg_indices)
if e in arg_indices_missing
])
for i, (arg, arg_indices) in enumerate(zip(args, indices)):
if len(arg_indices) < len(index0):
missing_indices_pos = [
i for i, e in enumerate(index0) if e not in arg_indices
]
missing_shape = [shape0[i] for i in missing_indices_pos]
arg_indices = tuple(index0[j]
for j in missing_indices_pos) + arg_indices
args[i] = _array_tensor_product(OneArray(*missing_shape),
args[i])
permutation = Permutation([arg_indices.index(j) for j in index0])
# Perform index permutations:
args[i] = _permute_dims(args[i], permutation)
return _array_add(*args), tuple(index0)
if isinstance(expr, Pow):
subexpr, subindices = _convert_indexed_to_array(expr.base)
if isinstance(expr.exp, (int, Integer)):
diags = zip(*[(2 * i, 2 * i + 1) for i in range(expr.exp)])
arr = _array_diagonal(
_array_tensor_product(*[subexpr for i in range(expr.exp)]),
*diags)
return arr, subindices
if isinstance(expr, Function):
subexpr, subindices = _convert_indexed_to_array(expr.args[0])
return ArrayElementwiseApplyFunc(type(expr), subexpr), subindices
return expr, ()
def test_arrayexpr_array_flatten():
# Flatten nested ArrayTensorProduct objects:
expr1 = _array_tensor_product(M, N)
expr2 = _array_tensor_product(P, Q)
expr = _array_tensor_product(expr1, expr2)
assert expr == _array_tensor_product(M, N, P, Q)
assert expr.args == (M, N, P, Q)
# Flatten mixed ArrayTensorProduct and ArrayContraction objects:
cg1 = _array_contraction(expr1, (1, 2))
cg2 = _array_contraction(expr2, (0, 3))
expr = _array_tensor_product(cg1, cg2)
assert expr == _array_contraction(_array_tensor_product(M, N, P, Q), (1, 2), (4, 7))
expr = _array_tensor_product(M, cg1)
assert expr == _array_contraction(_array_tensor_product(M, M, N), (3, 4))
# Flatten nested ArrayContraction objects:
cgnested = _array_contraction(cg1, (0, 1))
assert cgnested == _array_contraction(_array_tensor_product(M, N), (0, 3), (1, 2))
cgnested = _array_contraction(_array_tensor_product(cg1, cg2), (0, 3))
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 6), (1, 2), (4, 7))
cg3 = _array_contraction(_array_tensor_product(M, N, P, Q), (1, 3), (2, 4))
cgnested = _array_contraction(cg3, (0, 1))
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 5), (1, 3), (2, 4))
cgnested = _array_contraction(cg3, (0, 3), (1, 2))
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 7), (1, 3), (2, 4), (5, 6))
cg4 = _array_contraction(_array_tensor_product(M, N, P, Q), (1, 5), (3, 7))
cgnested = _array_contraction(cg4, (0, 1))
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 2), (1, 5), (3, 7))
cgnested = _array_contraction(cg4, (0, 1), (2, 3))
assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 2), (1, 5), (3, 7), (4, 6))
cg = _array_diagonal(cg4)
assert cg == cg4
assert isinstance(cg, type(cg4))
# Flatten nested ArrayDiagonal objects:
cg1 = _array_diagonal(expr1, (1, 2))
cg2 = _array_diagonal(expr2, (0, 3))
cg3 = _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 3), (2, 4))
cg4 = _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 5), (3, 7))
cgnested = _array_diagonal(cg1, (0, 1))
assert cgnested == _array_diagonal(_array_tensor_product(M, N), (1, 2), (0, 3))
cgnested = _array_diagonal(cg3, (1, 2))
assert cgnested == _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 3), (2, 4), (5, 6))
cgnested = _array_diagonal(cg4, (1, 2))
assert cgnested == _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 5), (3, 7), (2, 4))
cg = _array_add(M, N)
cg2 = _array_add(cg, P)
assert isinstance(cg2, ArrayAdd)
assert cg2.args == (M, N, P)
assert cg2.shape == (k, k)
expr = _array_tensor_product(_array_diagonal(X, (0, 1)), _array_diagonal(A, (0, 1)))
assert expr == _array_diagonal(_array_tensor_product(X, A), (0, 1), (2, 3))
expr1 = _array_diagonal(_array_tensor_product(X, A), (1, 2))
expr2 = _array_tensor_product(expr1, a)
assert expr2 == _permute_dims(_array_diagonal(_array_tensor_product(X, A, a), (1, 2)), [0, 1, 4, 2, 3])
expr1 = _array_contraction(_array_tensor_product(X, A), (1, 2))
expr2 = _array_tensor_product(expr1, a)
assert isinstance(expr2, ArrayContraction)
assert isinstance(expr2.expr, ArrayTensorProduct)
cg = _array_tensor_product(_array_diagonal(_array_tensor_product(A, X, Y), (0, 3), (1, 5)), a, b)
assert cg == _permute_dims(_array_diagonal(_array_tensor_product(A, X, Y, a, b), (0, 3), (1, 5)), [0, 1, 6, 7, 2, 3, 4, 5])
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 = _array_diagonal(_array_tensor_product(M, N, Q), (1, 3), (0, 2, 4))
ret = convert_array_to_matrix(cg)
expected = PermuteDims(_array_diagonal(_array_tensor_product(HadamardProduct(M.T, N.T), Q), (1, 2)), [1, 0, 2])
assert expected == ret
# Special case that should return the same expression:
cg = _array_diagonal(_array_tensor_product(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 = _array_diagonal(_array_tensor_product(M, N), (0, 1, 2, 3))
ret = convert_array_to_matrix(cg)
assert ret == cg
cg = _array_diagonal(_array_tensor_product(A), (0, 1))
ret = convert_array_to_matrix(cg)
assert ret == cg
cg = _array_diagonal(_array_tensor_product(M, N, P), (0, 2, 4), (1, 3, 5))
assert convert_array_to_matrix(cg) == HadamardProduct(M, N, P)
cg = _array_diagonal(_array_tensor_product(M, N, P), (0, 3, 4), (1, 2, 5))
assert convert_array_to_matrix(cg) == HadamardProduct(M, P, N.T)
cg = _array_diagonal(_array_tensor_product(I, I1, x), (1, 4), (3, 5))
assert convert_array_to_matrix(cg) == DiagMatrix(x)
def test_arrayexpr_convert_array_to_matrix_remove_trivial_dims():
# Tensor Product:
assert _remove_trivial_dims(_array_tensor_product(a, b)) == (a * b.T, [1, 3])
assert _remove_trivial_dims(_array_tensor_product(a.T, b)) == (a * b.T, [0, 3])
assert _remove_trivial_dims(_array_tensor_product(a, b.T)) == (a * b.T, [1, 2])
assert _remove_trivial_dims(_array_tensor_product(a.T, b.T)) == (a * b.T, [0, 2])
assert _remove_trivial_dims(_array_tensor_product(I, a.T, b.T)) == (_array_tensor_product(I, a * b.T), [2, 4])
assert _remove_trivial_dims(_array_tensor_product(a.T, I, b.T)) == (_array_tensor_product(a.T, I, b.T), [])
assert _remove_trivial_dims(_array_tensor_product(a, I)) == (_array_tensor_product(a, I), [])
assert _remove_trivial_dims(_array_tensor_product(I, a)) == (_array_tensor_product(I, a), [])
assert _remove_trivial_dims(_array_tensor_product(a.T, b.T, c, d)) == (
_array_tensor_product(a * b.T, c * d.T), [0, 2, 5, 7])
assert _remove_trivial_dims(_array_tensor_product(a.T, I, b.T, c, d, I)) == (
_array_tensor_product(a.T, I, b*c.T, d, I), [4, 7])
# Addition:
cg = ArrayAdd(_array_tensor_product(a, b), _array_tensor_product(c, d))
assert _remove_trivial_dims(cg) == (a * b.T + c * d.T, [1, 3])
# Permute Dims:
cg = PermuteDims(_array_tensor_product(a, b), Permutation(3)(1, 2))
assert _remove_trivial_dims(cg) == (a * b.T, [2, 3])
cg = PermuteDims(_array_tensor_product(a, I, b), Permutation(5)(1, 2, 3, 4))
assert _remove_trivial_dims(cg) == (cg, [])
cg = PermuteDims(_array_tensor_product(I, b, a), Permutation(5)(1, 2, 4, 5, 3))
assert _remove_trivial_dims(cg) == (PermuteDims(_array_tensor_product(I, b * a.T), [0, 2, 3, 1]), [4, 5])
# Diagonal:
cg = _array_diagonal(_array_tensor_product(M, a), (1, 2))
assert _remove_trivial_dims(cg) == (cg, [])
# Contraction:
cg = _array_contraction(_array_tensor_product(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 = _array_tensor_product(
OneMatrix(1, 1),
M,
x,
OneMatrix(1, 1),
Identity(1),
)
expr = _array_contraction(tp, (1, 8))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [0, 5, 6, 7]
expr = _array_contraction(tp, (1, 8), (3, 4))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [0, 3, 4, 5]
expr = _array_diagonal(tp, (1, 8))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [0, 5, 6, 7, 8]
expr = _array_diagonal(tp, (1, 8), (3, 4))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [0, 3, 4, 5, 6]
expr = _array_diagonal(_array_contraction(_array_tensor_product(A, x, I, I1), (1, 2, 5)), (1, 4))
rexpr, removed = _remove_trivial_dims(expr)
assert removed == [2, 3]
cg = _array_diagonal(_array_tensor_product(PermuteDims(_array_tensor_product(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 = _array_contraction(_array_tensor_product(A, B, C, M, I), (1, 8))
ret, removed = _remove_trivial_dims(cg)
assert ret == PermuteDims(_array_tensor_product(A, B, C, M), [0, 2, 3, 4, 5, 6, 7, 1])
assert removed == []
cg = _array_contraction(_array_tensor_product(A, B, C, M, I), (1, 8), (3, 4))
ret, removed = _remove_trivial_dims(cg)
assert ret == PermuteDims(_array_contraction(_array_tensor_product(A, B, C, M), (3, 4)), [0, 2, 3, 4, 5, 1])
assert removed == []
# Trivial matrices are sometimes inserted into MatMul expressions:
cg = _array_tensor_product(b*b.T, a.T*a)
ret, removed = _remove_trivial_dims(cg)
assert ret == b*a.T*a*b.T
assert removed == [2, 3]
Xs = ArraySymbol("X", (3, 2, k))
cg = _array_tensor_product(M, Xs, b.T*c, a*a.T, b*b.T, c.T*d)
ret, removed = _remove_trivial_dims(cg)
assert ret == _array_tensor_product(M, Xs, a*b.T*c*c.T*d*a.T, b*b.T)
assert removed == [5, 6, 11, 12]
cg = _array_diagonal(_array_tensor_product(I, I1, x), (1, 4), (3, 5))
assert _remove_trivial_dims(cg) == (PermuteDims(_array_diagonal(_array_tensor_product(I, x), (1, 2)), Permutation(1, 2)), [1])
expr = _array_diagonal(_array_tensor_product(x, I, y), (0, 2))
assert _remove_trivial_dims(expr) == (PermuteDims(_array_tensor_product(DiagMatrix(x), y), [1, 2, 3, 0]), [0])
expr = _array_diagonal(_array_tensor_product(x, I, y), (0, 2), (3, 4))
assert _remove_trivial_dims(expr) == (expr, [])
def test_arrayexpr_convert_array_to_diagonalized_vector():
# Check matrix recognition over trivial dimensions:
cg = _array_tensor_product(a, b)
assert convert_array_to_matrix(cg) == a * b.T
cg = _array_tensor_product(I1, a, b)
assert convert_array_to_matrix(cg) == a * b.T
# Recognize trace inside a tensor product:
cg = _array_contraction(_array_tensor_product(A, B, C), (0, 3), (1, 2))
assert convert_array_to_matrix(cg) == Trace(A * B) * C
# Transform diagonal operator to contraction:
cg = _array_diagonal(_array_tensor_product(A, a), (1, 2))
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(A, OneArray(1), DiagMatrix(a)), (1, 3))
assert convert_array_to_matrix(cg) == A * DiagMatrix(a)
cg = _array_diagonal(_array_tensor_product(a, b), (0, 2))
assert _array_diag2contr_diagmatrix(cg) == _permute_dims(
_array_contraction(_array_tensor_product(DiagMatrix(a), OneArray(1), b), (0, 3)), [1, 2, 0]
)
assert convert_array_to_matrix(cg) == b.T * DiagMatrix(a)
cg = _array_diagonal(_array_tensor_product(A, a), (0, 2))
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(A, OneArray(1), DiagMatrix(a)), (0, 3))
assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a)
cg = _array_diagonal(_array_tensor_product(I, x, I1), (0, 2), (3, 5))
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(I, OneArray(1), I1, DiagMatrix(x)), (0, 5))
assert convert_array_to_matrix(cg) == DiagMatrix(x)
cg = _array_diagonal(_array_tensor_product(I, x, A, B), (1, 2), (5, 6))
assert _array_diag2contr_diagmatrix(cg) == _array_diagonal(_array_contraction(_array_tensor_product(I, OneArray(1), A, B, DiagMatrix(x)), (1, 7)), (5, 6))
# TODO: this is returning a wrong result:
# convert_array_to_matrix(cg)
cg = _array_diagonal(_array_tensor_product(I1, a, b), (1, 3, 5))
assert convert_array_to_matrix(cg) == a*b.T
cg = _array_diagonal(_array_tensor_product(I1, a, b), (1, 3))
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(OneArray(1), a, b, I1), (2, 6))
assert convert_array_to_matrix(cg) == a*b.T
cg = _array_diagonal(_array_tensor_product(x, I1), (1, 2))
assert isinstance(cg, ArrayDiagonal)
assert cg.diagonal_indices == ((1, 2),)
assert convert_array_to_matrix(cg) == x
cg = _array_diagonal(_array_tensor_product(x, I), (0, 2))
assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(OneArray(1), I, DiagMatrix(x)), (1, 3))
assert convert_array_to_matrix(cg).doit() == DiagMatrix(x)
raises(ValueError, lambda: _array_diagonal(x, (1,)))
# Ignore identity matrices with contractions:
cg = _array_contraction(_array_tensor_product(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 = _array_contraction(_array_tensor_product(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 = _array_contraction(_array_tensor_product(A, a), (1, 2))
assert cg.split_multiple_contractions() == cg
assert convert_array_to_matrix(cg) == A * a
cg = _array_contraction(_array_tensor_product(A, a, B), (1, 2, 4))
assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1), B), (1, 2), (3, 5))
assert convert_array_to_matrix(cg) == A * DiagMatrix(a) * B
cg = _array_contraction(_array_tensor_product(A, a, B), (0, 2, 4))
assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1), B), (0, 2), (3, 5))
assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) * B
cg = _array_contraction(_array_tensor_product(A, a, b, a.T, B), (0, 2, 4, 7, 9))
assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(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 = _array_contraction(_array_tensor_product(I1, I1, I1), (1, 2, 4))
assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(I1, I1, OneArray(1), I1), (1, 2), (3, 5))
assert convert_array_to_matrix(cg) == 1
cg = _array_contraction(_array_tensor_product(I, I, I, I, A), (1, 2, 8), (5, 6, 9))
assert convert_array_to_matrix(cg.split_multiple_contractions()).doit() == A
cg = _array_contraction(_array_tensor_product(A, a, C, a, B), (1, 2, 4), (5, 6, 8))
expected = _array_contraction(_array_tensor_product(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 = _array_contraction(_array_tensor_product(a, I1, b, I1, (a.T*b).applyfunc(cos)), (1, 2, 8), (5, 6, 9))
expected = _array_contraction(_array_tensor_product(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))
