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_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_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: 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 = _array_tensor_product(*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 = _permute_dims(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 _array_add(*addend_list)
def _(expr: ArrayContraction):
expr = expr.flatten_contraction_of_diagonal()
expr = identify_removable_identity_matrices(expr)
expr = expr.split_multiple_contractions()
expr = identify_hadamard_products(expr)
if not isinstance(expr, ArrayContraction):
return _array2matrix(expr)
subexpr = expr.expr
contraction_indices: tTuple[tTuple[int]] = expr.contraction_indices
if contraction_indices == ((0, ), (1, )) or (contraction_indices == (
(0, ), ) and subexpr.shape[1] == 1) or (contraction_indices == (
(1, ), ) and subexpr.shape[0] == 1):
shape = subexpr.shape
subexpr = _array2matrix(subexpr)
if isinstance(subexpr, MatrixExpr):
return OneMatrix(1, shape[0]) * subexpr * OneMatrix(shape[1], 1)
if isinstance(subexpr, ArrayTensorProduct):
newexpr = _array_contraction(_array2matrix(subexpr),
*contraction_indices)
contraction_indices = newexpr.contraction_indices
if any(i > 2 for i in newexpr.subranks):
addends = _array_add(*[
_a2m_tensor_product(*j) for j in itertools.product(*[
i.args if isinstance(i, ArrayAdd) else [i]
for i in expr.expr.args
])
])
newexpr = _array_contraction(addends, *contraction_indices)
if isinstance(newexpr, ArrayAdd):
ret = _array2matrix(newexpr)
return ret
assert isinstance(newexpr, ArrayContraction)
ret = _support_function_tp1_recognize(contraction_indices,
list(newexpr.expr.args))
return ret
elif not isinstance(subexpr, _CodegenArrayAbstract):
ret = _array2matrix(subexpr)
if isinstance(ret, MatrixExpr):
assert expr.contraction_indices == ((0, 1), )
return _a2m_trace(ret)
else:
return _array_contraction(ret, *expr.contraction_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_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 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 any(isinstance(i, (int, Integer)) for i in indices):
result = ArrayElement(result, indices)
indices = []
if not first_indices:
return result
def _check_is_in(elem, indices):
if elem in indices:
return True
if any(elem in i for i in indices if isinstance(i, frozenset)):
return True
return False
repl = {j: i for i in indices if isinstance(i, frozenset) for j in i}
first_indices = [repl.get(i, i) for i in first_indices]
for i in first_indices:
if not _check_is_in(i, indices):
first_indices.remove(i)
first_indices.extend(
[i for i in indices if not _check_is_in(i, first_indices)])
def _get_pos(elem, indices):
if elem in indices:
return indices.index(elem)
for i, e in enumerate(indices):
if not isinstance(e, frozenset):
continue
if elem in e:
return i
raise ValueError("not found")
permutation = _af_invert([_get_pos(i, first_indices) for i in indices])
if isinstance(result, ArrayAdd):
return _array_add(
*[_permute_dims(arg, permutation) for arg in result.args])
else:
return _permute_dims(result, permutation)
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 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: ArrayAdd, x: Expr):
return _array_add(*[array_derive(arg, x) for arg in expr.args])
def _a2m_add(*args):
if not any(isinstance(i, _CodegenArrayAbstract) for i in args):
from sympy.matrices.expressions.matadd import MatAdd
return MatAdd(*args).doit()
else:
return _array_add(*args)
