def test_arrayexpr_convert_indexed_to_array_broadcast():
A = ArraySymbol("A", (3, 3))
B = ArraySymbol("B", (3, 3))
expr = A[i, j] + B[k, l]
O2 = OneArray(3, 3)
expected = ArrayAdd(ArrayTensorProduct(A, O2), ArrayTensorProduct(O2, B))
assert convert_indexed_to_array(expr) == expected
assert convert_indexed_to_array(expr, [i, j, k, l]) == expected
assert convert_indexed_to_array(expr, [l, k, i, j]) == ArrayAdd(
PermuteDims(ArrayTensorProduct(O2, A), [1, 0, 2, 3]),
PermuteDims(ArrayTensorProduct(B, O2), [1, 0, 2, 3]))
expr = A[i, j] + B[j, k]
O1 = OneArray(3)
assert convert_indexed_to_array(expr, [i, j, k]) == ArrayAdd(
ArrayTensorProduct(A, O1), ArrayTensorProduct(O1, B))
C = ArraySymbol("C", (d0, d1))
D = ArraySymbol("D", (d3, d1))
expr = C[i, j] + D[k, j]
assert convert_indexed_to_array(expr, [i, j, k]) == ArrayAdd(
ArrayTensorProduct(C, OneArray(d3)),
PermuteDims(ArrayTensorProduct(OneArray(d0), D), [0, 2, 1]))
X = ArraySymbol("X", (5, 3))
expr = X[i, n] - X[j, n]
assert convert_indexed_to_array(expr, [i, j, n]) == ArrayAdd(
ArrayTensorProduct(-1, OneArray(5), X),
PermuteDims(ArrayTensorProduct(X, OneArray(5)), [0, 2, 1]))
raises(ValueError, lambda: convert_indexed_to_array(C[i, j] + D[i, j]))
def test_arrayexpr_convert_array_contraction_tp_additions():
a = ArrayAdd(
ArrayTensorProduct(M, N),
ArrayTensorProduct(N, M)
)
tp = ArrayTensorProduct(P, a, Q)
expr = ArrayContraction(tp, (3, 4))
expected = ArrayTensorProduct(
P,
ArrayAdd(
ArrayContraction(ArrayTensorProduct(M, N), (1, 2)),
ArrayContraction(ArrayTensorProduct(N, M), (1, 2)),
),
Q
)
assert expr == expected
assert convert_array_to_matrix(expr) == ArrayTensorProduct(P, M * N + N * M, Q)
expr = ArrayContraction(tp, (1, 2), (3, 4), (5, 6))
result = ArrayContraction(
ArrayTensorProduct(
P,
ArrayAdd(
ArrayContraction(ArrayTensorProduct(M, N), (1, 2)),
ArrayContraction(ArrayTensorProduct(N, M), (1, 2)),
),
Q
), (1, 2), (3, 4))
assert expr == result
assert convert_array_to_matrix(expr) == P * (M * N + N * M) * Q
def test_array_as_explicit_matrix_symbol():
A = MatrixSymbol("A", 3, 3)
B = MatrixSymbol("B", 3, 3)
texpr = tensorproduct(A, B)
assert isinstance(texpr, ArrayTensorProduct)
assert texpr.as_explicit() == tensorproduct(A.as_explicit(),
B.as_explicit())
texpr = tensorcontraction(A, (0, 1))
assert isinstance(texpr, ArrayContraction)
assert texpr.as_explicit() == A[0, 0] + A[1, 1] + A[2, 2]
texpr = tensordiagonal(A, (0, 1))
assert isinstance(texpr, ArrayDiagonal)
assert texpr.as_explicit() == ImmutableDenseNDimArray(
[A[0, 0], A[1, 1], A[2, 2]])
texpr = permutedims(A, [1, 0])
assert isinstance(texpr, PermuteDims)
assert texpr.as_explicit() == permutedims(A.as_explicit(), [1, 0])
expr = ArrayAdd(ArrayTensorProduct(A, B), ArrayTensorProduct(B, A))
assert expr.as_explicit(
) == expr.args[0].as_explicit() + expr.args[1].as_explicit()
def test_arrayexpr_convert_array_to_matrix():
cg = ArrayContraction(ArrayTensorProduct(M), (0, 1))
assert convert_array_to_matrix(cg) == Trace(M)
cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 1), (2, 3))
assert convert_array_to_matrix(cg) == Trace(M) * Trace(N)
cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 3), (1, 2))
assert convert_array_to_matrix(cg) == Trace(M * N)
cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 2), (1, 3))
assert convert_array_to_matrix(cg) == Trace(M * N.T)
cg = convert_matrix_to_array(M * N * P)
assert convert_array_to_matrix(cg) == M * N * P
cg = convert_matrix_to_array(M * N.T * P)
assert convert_array_to_matrix(cg) == M * N.T * P
cg = ArrayContraction(ArrayTensorProduct(M,N,P,Q), (1, 2), (5, 6))
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M * N, P * Q)
cg = ArrayContraction(ArrayTensorProduct(-2, M, N), (1, 2))
assert convert_array_to_matrix(cg) == -2 * M * N
a = MatrixSymbol("a", k, 1)
b = MatrixSymbol("b", k, 1)
c = MatrixSymbol("c", k, 1)
cg = PermuteDims(
ArrayContraction(
ArrayTensorProduct(
a,
ArrayAdd(
ArrayTensorProduct(b, c),
ArrayTensorProduct(c, b),
)
), (2, 4)), [0, 1, 3, 2])
assert convert_array_to_matrix(cg) == a * (b.T * c + c.T * b)
za = ZeroArray(m, n)
assert convert_array_to_matrix(za) == ZeroMatrix(m, n)
cg = ArrayTensorProduct(3, M)
assert convert_array_to_matrix(cg) == 3 * M
# Partial conversion to matrix multiplication:
expr = ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 2), (1, 4, 6))
assert convert_array_to_matrix(expr) == ArrayContraction(ArrayTensorProduct(M.T*N, P, Q), (0, 2, 4))
x = MatrixSymbol("x", k, 1)
cg = PermuteDims(
ArrayContraction(ArrayTensorProduct(OneArray(1), x, OneArray(1), DiagMatrix(Identity(1))),
(0, 5)), Permutation(1, 2, 3))
assert convert_array_to_matrix(cg) == x
expr = ArrayAdd(M, PermuteDims(M, [1, 0]))
assert convert_array_to_matrix(expr) == M + Transpose(M)
def 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_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 _(expr: ArrayContraction):
expr = expr.flatten_contraction_of_diagonal()
expr = expr.split_multiple_contractions()
expr = identify_hadamard_products(expr)
subexpr = expr.expr
contraction_indices: Tuple[Tuple[int]] = expr.contraction_indices
if isinstance(subexpr, ArrayTensorProduct):
newexpr = ArrayContraction(_array2matrix(subexpr),
*contraction_indices)
contraction_indices = newexpr.contraction_indices
if any(i > 2 for i in newexpr.subranks):
addends = ArrayAdd(*[
_a2m_tensor_product(*j) for j in itertools.product(*[
i.args if isinstance(i, ArrayAdd) else [i]
for i in expr.expr.args
])
])
newexpr = ArrayContraction(addends, *contraction_indices)
if isinstance(newexpr, ArrayAdd):
ret = _array2matrix(newexpr)
return ret
assert isinstance(newexpr, ArrayContraction)
ret = _support_function_tp1_recognize(contraction_indices,
list(newexpr.expr.args))
return ret
elif not isinstance(subexpr, _CodegenArrayAbstract):
ret = _array2matrix(subexpr)
if isinstance(ret, MatrixExpr):
assert expr.contraction_indices == ((0, 1), )
return _a2m_trace(ret)
else:
return ArrayContraction(ret, *expr.contraction_indices)
def test_mixed_deriv_mixed_expressions():
expr = 3 * Trace(A)
assert expr.diff(A) == 3 * Identity(k)
expr = k
deriv = expr.diff(A)
assert isinstance(deriv, ZeroMatrix)
assert deriv == ZeroMatrix(k, k)
expr = Trace(A)**2
assert expr.diff(A) == (2 * Trace(A)) * Identity(k)
expr = Trace(A) * A
I = Identity(k)
assert expr.diff(A) == ArrayAdd(
ArrayTensorProduct(I, A),
PermuteDims(ArrayTensorProduct(Trace(A) * I, I),
Permutation(3)(1, 2)))
expr = Trace(Trace(A) * A)
assert expr.diff(A) == (2 * Trace(A)) * Identity(k)
expr = Trace(Trace(Trace(A) * A) * A)
assert expr.diff(A) == (3 * Trace(A)**2) * Identity(k)
def test_arrayexpr_convert_indexed_to_array_and_back_to_matrix():
expr = a.T * b
elem = expr[0, 0]
cg = convert_indexed_to_array(elem)
assert cg == ArrayElement(
ArrayContraction(ArrayTensorProduct(a, b), (0, 2)), [0, 0])
expr = M[i, j] + N[i, j]
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(p1) == M + N
expr = M[i, j] + N[j, i]
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(p1) == M + N.T
expr = M[i, j] * N[k, l] + N[i, j] * M[k, l]
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(p1) == ArrayAdd(ArrayTensorProduct(M, N),
ArrayTensorProduct(N, M))
expr = (M * N * P)[i, j]
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(p1) == M * N * P
expr = Sum(M[i, j] * (N * P)[j, m], (j, 0, k - 1))
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(p1) == M * N * P
expr = Sum((P[j, m] + P[m, j]) * (M[i, j] * N[m, n] + N[i, j] * M[m, n]),
(j, 0, k - 1), (m, 0, k - 1))
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(
p1) == M * P * N + M * P.T * N + N * P * M + N * P.T * M
def _(expr: ArrayTensorProduct, x: Expr):
args = expr.args
addend_list = []
for i, arg in enumerate(expr.args):
darg = array_derive(arg, x)
if darg == 0:
continue
args_prev = args[:i]
args_succ = args[i + 1:]
shape_prev = reduce(operator.add, map(get_shape, args_prev), ())
shape_succ = reduce(operator.add, map(get_shape, args_succ), ())
addend = ArrayTensorProduct(*args_prev, darg, *args_succ)
tot1 = len(get_shape(x))
tot2 = tot1 + len(shape_prev)
tot3 = tot2 + len(get_shape(arg))
tot4 = tot3 + len(shape_succ)
perm = [i for i in range(tot1, tot2)] + \
[i for i in range(tot1)] + [i for i in range(tot2, tot3)] + \
[i for i in range(tot3, tot4)]
addend = PermuteDims(addend, _af_invert(perm))
addend_list.append(addend)
if len(addend_list) == 1:
return addend_list[0]
elif len(addend_list) == 0:
return S.Zero
else:
return ArrayAdd(*addend_list)
def test_matrix_derivative_non_matrix_result():
# This is a 4-dimensional array:
I = Identity(k)
AdA = PermuteDims(ArrayTensorProduct(I, I), Permutation(3)(1, 2))
assert A.diff(A) == AdA
assert A.T.diff(A) == PermuteDims(ArrayTensorProduct(I, I), Permutation(3)(1, 2, 3))
assert (2*A).diff(A) == PermuteDims(ArrayTensorProduct(2*I, I), Permutation(3)(1, 2))
assert MatAdd(A, A).diff(A) == ArrayAdd(AdA, AdA)
assert (A + B).diff(A) == AdA
def convert_matrix_to_array(expr: MatrixExpr) -> Basic:
if isinstance(expr, MatMul):
args_nonmat = []
args = []
for arg in expr.args:
if isinstance(arg, MatrixExpr):
args.append(arg)
else:
args_nonmat.append(convert_matrix_to_array(arg))
contractions = [(2 * i + 1, 2 * i + 2) for i in range(len(args) - 1)]
scalar = ArrayTensorProduct.fromiter(
args_nonmat) if args_nonmat else S.One
if scalar == 1:
tprod = ArrayTensorProduct(
*[convert_matrix_to_array(arg) for arg in args])
else:
tprod = ArrayTensorProduct(
scalar, *[convert_matrix_to_array(arg) for arg in args])
return ArrayContraction(tprod, *contractions)
elif isinstance(expr, MatAdd):
return ArrayAdd(*[convert_matrix_to_array(arg) for arg in expr.args])
elif isinstance(expr, Transpose):
return PermuteDims(convert_matrix_to_array(expr.args[0]), [1, 0])
elif isinstance(expr, Trace):
inner_expr = convert_matrix_to_array(expr.arg)
return ArrayContraction(inner_expr, (0, len(inner_expr.shape) - 1))
elif isinstance(expr, Mul):
return ArrayTensorProduct.fromiter(
convert_matrix_to_array(i) for i in expr.args)
elif isinstance(expr, Pow):
base = convert_matrix_to_array(expr.base)
if (expr.exp > 0) == True:
return ArrayTensorProduct.fromiter(base for i in range(expr.exp))
else:
return expr
elif isinstance(expr, MatPow):
base = convert_matrix_to_array(expr.base)
if expr.exp.is_Integer != True:
b = symbols("b", cls=Dummy)
return ArrayElementwiseApplyFunc(Lambda(b, b**expr.exp),
convert_matrix_to_array(base))
elif (expr.exp > 0) == True:
return convert_matrix_to_array(
MatMul.fromiter(base for i in range(expr.exp)))
else:
return expr
elif isinstance(expr, HadamardProduct):
tp = ArrayTensorProduct.fromiter(expr.args)
diag = [[2 * i for i in range(len(expr.args))],
[2 * i + 1 for i in range(len(expr.args))]]
return ArrayDiagonal(tp, *diag)
elif isinstance(expr, HadamardPower):
base, exp = expr.args
return convert_matrix_to_array(
HadamardProduct.fromiter(base for i in range(exp)))
else:
return expr
def test_arrayexpr_convert_array_to_matrix_remove_trivial_dims():
# Tensor Product:
assert _remove_trivial_dims(ArrayTensorProduct(a, b)) == (a * b.T, [1, 3])
assert _remove_trivial_dims(ArrayTensorProduct(a.T,
b)) == (a * b.T, [0, 3])
assert _remove_trivial_dims(ArrayTensorProduct(a,
b.T)) == (a * b.T, [1, 2])
assert _remove_trivial_dims(ArrayTensorProduct(a.T,
b.T)) == (a * b.T, [0, 2])
assert _remove_trivial_dims(ArrayTensorProduct(I, a.T,
b.T)) == (a * b.T,
[0, 1, 2, 4])
assert _remove_trivial_dims(ArrayTensorProduct(a.T, I,
b.T)) == (a * b.T,
[0, 2, 3, 4])
assert _remove_trivial_dims(ArrayTensorProduct(a, I)) == (a, [2, 3])
assert _remove_trivial_dims(ArrayTensorProduct(I, a)) == (a, [0, 1])
assert _remove_trivial_dims(ArrayTensorProduct(a.T, b.T, c,
d)) == (ArrayTensorProduct(
a * b.T,
c * d.T), [0, 2, 5, 7])
assert _remove_trivial_dims(ArrayTensorProduct(a.T, I, b.T, c, d,
I)) == (ArrayTensorProduct(
a * b.T, c * d.T,
I), [0, 2, 3, 4, 7, 9])
# Addition:
cg = ArrayAdd(ArrayTensorProduct(a, b), ArrayTensorProduct(c, d))
assert _remove_trivial_dims(cg) == (a * b.T + c * d.T, [1, 3])
# Permute Dims:
cg = PermuteDims(ArrayTensorProduct(a, b), Permutation(3)(1, 2))
assert _remove_trivial_dims(cg) == (a * b.T, [2, 3])
cg = PermuteDims(ArrayTensorProduct(a, I, b), Permutation(5)(1, 2, 3, 4))
assert _remove_trivial_dims(cg) == (a * b.T, [1, 2, 4, 5])
cg = PermuteDims(ArrayTensorProduct(I, b, a),
Permutation(5)(1, 2, 4, 5, 3))
assert _remove_trivial_dims(cg) == (b * a.T, [0, 3, 4, 5])
# Diagonal:
cg = ArrayDiagonal(ArrayTensorProduct(M, a), (1, 2))
assert _remove_trivial_dims(cg) == (cg, [])
# Contraction:
cg = ArrayContraction(ArrayTensorProduct(M, a), (1, 2))
assert _remove_trivial_dims(cg) == (cg, [])
def test_codegen_extra():
if not np:
skip("NumPy not installed")
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
P = MatrixSymbol("P", 2, 2)
Q = MatrixSymbol("Q", 2, 2)
ma = np.matrix([[1, 2], [3, 4]])
mb = np.matrix([[1, -2], [-1, 3]])
mc = np.matrix([[2, 0], [1, 2]])
md = np.matrix([[1, -1], [4, 7]])
cg = ArrayTensorProduct(M, N)
f = lambdify((M, N), cg, 'numpy')
assert (f(ma, mb) == np.einsum(ma, [0, 1], mb, [2, 3])).all()
cg = ArrayAdd(M, N)
f = lambdify((M, N), cg, 'numpy')
assert (f(ma, mb) == ma + mb).all()
cg = ArrayAdd(M, N, P)
f = lambdify((M, N, P), cg, 'numpy')
assert (f(ma, mb, mc) == ma + mb + mc).all()
cg = ArrayAdd(M, N, P, Q)
f = lambdify((M, N, P, Q), cg, 'numpy')
assert (f(ma, mb, mc, md) == ma + mb + mc + md).all()
cg = PermuteDims(M, [1, 0])
f = lambdify((M, ), cg, 'numpy')
assert (f(ma) == ma.T).all()
cg = PermuteDims(ArrayTensorProduct(M, N), [1, 2, 3, 0])
f = lambdify((M, N), cg, 'numpy')
assert (f(ma, mb) == np.transpose(np.einsum(ma, [0, 1], mb, [2, 3]),
(1, 2, 3, 0))).all()
cg = ArrayDiagonal(ArrayTensorProduct(M, N), (1, 2))
f = lambdify((M, N), cg, 'numpy')
assert (f(ma, mb) == np.diagonal(np.einsum(ma, [0, 1], mb, [2, 3]),
axis1=1,
axis2=2)).all()
def test_array_expr_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_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 _convert_indexed_to_array(expr):
if isinstance(expr, Sum):
function = expr.function
summation_indices = expr.variables
subexpr, subindices = _convert_indexed_to_array(function)
# Check dimensional consistency:
shape = subexpr.shape
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 = ArrayDiagonal(subexpr.expr, *diagonal_indices)
else:
subexpr = subexpr.expr
axes_contraction = defaultdict(list)
for i, ind in enumerate(subindices):
if ind in summation_indices:
axes_contraction[ind].append(i)
subindices[i] = None
for k, v in axes_contraction.items():
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 ArrayContraction(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 = ArrayTensorProduct(*newargs)
if diagonal_indices:
return (ArrayDiagonal(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 (ArrayDiagonal(expr.args[0],
*diagonal_indices), ret_indices)
else:
return expr.args[0], ret_indices
if isinstance(expr, Indexed):
indices = expr.indices
diagonal_indices, ret_indices = _get_diagonal_indices(indices)
if diagonal_indices:
return (ArrayDiagonal(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:
index0set = set(indices[0])
index0 = indices[0]
for i in range(1, len(args)):
if set(indices[i]) != index0set:
raise NotImplementedError("indices must be the same")
permutation = Permutation([index0.index(j) for j in indices[i]])
# Perform index permutations:
args[i] = PermuteDims(args[i], permutation)
return ArrayAdd(*args), index0
return expr, ()
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()
def test_codegen_extra():
if not tf:
skip("TensorFlow not installed")
graph = tf.Graph()
with graph.as_default():
session = tf.compat.v1.Session()
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
P = MatrixSymbol("P", 2, 2)
Q = MatrixSymbol("Q", 2, 2)
ma = tf.constant([[1, 2], [3, 4]])
mb = tf.constant([[1, -2], [-1, 3]])
mc = tf.constant([[2, 0], [1, 2]])
md = tf.constant([[1, -1], [4, 7]])
cg = ArrayTensorProduct(M, N)
assert tensorflow_code(cg) == \
'tensorflow.linalg.einsum("ab,cd", M, N)'
f = lambdify((M, N), cg, 'tensorflow')
y = session.run(f(ma, mb))
c = session.run(tf.einsum("ij,kl", ma, mb))
assert (y == c).all()
cg = ArrayAdd(M, N)
assert tensorflow_code(cg) == 'tensorflow.math.add(M, N)'
f = lambdify((M, N), cg, 'tensorflow')
y = session.run(f(ma, mb))
c = session.run(ma + mb)
assert (y == c).all()
cg = ArrayAdd(M, N, P)
assert tensorflow_code(cg) == \
'tensorflow.math.add(tensorflow.math.add(M, N), P)'
f = lambdify((M, N, P), cg, 'tensorflow')
y = session.run(f(ma, mb, mc))
c = session.run(ma + mb + mc)
assert (y == c).all()
cg = ArrayAdd(M, N, P, Q)
assert tensorflow_code(cg) == \
'tensorflow.math.add(' \
'tensorflow.math.add(tensorflow.math.add(M, N), P), Q)'
f = lambdify((M, N, P, Q), cg, 'tensorflow')
y = session.run(f(ma, mb, mc, md))
c = session.run(ma + mb + mc + md)
assert (y == c).all()
cg = PermuteDims(M, [1, 0])
assert tensorflow_code(cg) == 'tensorflow.transpose(M, [1, 0])'
f = lambdify((M, ), cg, 'tensorflow')
y = session.run(f(ma))
c = session.run(tf.transpose(ma))
assert (y == c).all()
cg = PermuteDims(ArrayTensorProduct(M, N), [1, 2, 3, 0])
assert tensorflow_code(cg) == \
'tensorflow.transpose(' \
'tensorflow.linalg.einsum("ab,cd", M, N), [1, 2, 3, 0])'
f = lambdify((M, N), cg, 'tensorflow')
y = session.run(f(ma, mb))
c = session.run(tf.transpose(tf.einsum("ab,cd", ma, mb), [1, 2, 3, 0]))
assert (y == c).all()
cg = ArrayDiagonal(ArrayTensorProduct(M, N), (1, 2))
assert tensorflow_code(cg) == \
'tensorflow.linalg.einsum("ab,bc->acb", M, N)'
f = lambdify((M, N), cg, 'tensorflow')
y = session.run(f(ma, mb))
c = session.run(tf.einsum("ab,bc->acb", ma, mb))
assert (y == c).all()
def test_matrix_derivatives_of_traces():
expr = Trace(A) * A
I = Identity(k)
assert expr.diff(A) == ArrayAdd(
ArrayTensorProduct(I, A),
PermuteDims(ArrayTensorProduct(Trace(A) * I, I),
Permutation(3)(1, 2)))
assert expr[i, j].diff(
A[m, n]).doit() == (KDelta(i, m) * KDelta(j, n) * Trace(A) +
KDelta(m, n) * A[i, j])
## First order:
# Cookbook example 99:
expr = Trace(X)
assert expr.diff(X) == Identity(k)
assert expr.rewrite(Sum).diff(X[m, n]).doit() == KDelta(m, n)
# Cookbook example 100:
expr = Trace(X * A)
assert expr.diff(X) == A.T
assert expr.rewrite(Sum).diff(X[m, n]).doit() == A[n, m]
# Cookbook example 101:
expr = Trace(A * X * B)
assert expr.diff(X) == A.T * B.T
assert expr.rewrite(Sum).diff(X[m, n]).doit().dummy_eq((A.T * B.T)[m, n])
# Cookbook example 102:
expr = Trace(A * X.T * B)
assert expr.diff(X) == B * A
# Cookbook example 103:
expr = Trace(X.T * A)
assert expr.diff(X) == A
# Cookbook example 104:
expr = Trace(A * X.T)
assert expr.diff(X) == A
# Cookbook example 105:
# TODO: TensorProduct is not supported
#expr = Trace(TensorProduct(A, X))
#assert expr.diff(X) == Trace(A)*Identity(k)
## Second order:
# Cookbook example 106:
expr = Trace(X**2)
assert expr.diff(X) == 2 * X.T
# Cookbook example 107:
expr = Trace(X**2 * B)
assert expr.diff(X) == (X * B + B * X).T
expr = Trace(MatMul(X, X, B))
assert expr.diff(X) == (X * B + B * X).T
# Cookbook example 108:
expr = Trace(X.T * B * X)
assert expr.diff(X) == B * X + B.T * X
# Cookbook example 109:
expr = Trace(B * X * X.T)
assert expr.diff(X) == B * X + B.T * X
# Cookbook example 110:
expr = Trace(X * X.T * B)
assert expr.diff(X) == B * X + B.T * X
# Cookbook example 111:
expr = Trace(X * B * X.T)
assert expr.diff(X) == X * B.T + X * B
# Cookbook example 112:
expr = Trace(B * X.T * X)
assert expr.diff(X) == X * B.T + X * B
# Cookbook example 113:
expr = Trace(X.T * X * B)
assert expr.diff(X) == X * B.T + X * B
# Cookbook example 114:
expr = Trace(A * X * B * X)
assert expr.diff(X) == A.T * X.T * B.T + B.T * X.T * A.T
# Cookbook example 115:
expr = Trace(X.T * X)
assert expr.diff(X) == 2 * X
expr = Trace(X * X.T)
assert expr.diff(X) == 2 * X
# Cookbook example 116:
expr = Trace(B.T * X.T * C * X * B)
assert expr.diff(X) == C.T * X * B * B.T + C * X * B * B.T
# Cookbook example 117:
expr = Trace(X.T * B * X * C)
assert expr.diff(X) == B * X * C + B.T * X * C.T
# Cookbook example 118:
expr = Trace(A * X * B * X.T * C)
assert expr.diff(X) == A.T * C.T * X * B.T + C * A * X * B
# Cookbook example 119:
expr = Trace((A * X * B + C) * (A * X * B + C).T)
assert expr.diff(X) == 2 * A.T * (A * X * B + C) * B.T
# Cookbook example 120:
# TODO: no support for TensorProduct.
# expr = Trace(TensorProduct(X, X))
# expr = Trace(X)*Trace(X)
# expr.diff(X) == 2*Trace(X)*Identity(k)
# Higher Order
# Cookbook example 121:
expr = Trace(X**k)
#assert expr.diff(X) == k*(X**(k-1)).T
# Cookbook example 122:
expr = Trace(A * X**k)
#assert expr.diff(X) == # Needs indices
# Cookbook example 123:
expr = Trace(B.T * X.T * C * X * X.T * C * X * B)
assert expr.diff(
X
) == C * X * X.T * C * X * B * B.T + C.T * X * B * B.T * X.T * C.T * X + C * X * B * B.T * X.T * C * X + C.T * X * X.T * C.T * X * B * B.T
# Other
# Cookbook example 124:
expr = Trace(A * X**(-1) * B)
assert expr.diff(X) == -Inverse(X).T * A.T * B.T * Inverse(X).T
# Cookbook example 125:
expr = Trace(Inverse(X.T * C * X) * A)
# Warning: result in the cookbook is equivalent if B and C are symmetric:
assert expr.diff(X) == -X.inv().T * A.T * X.inv() * C.inv().T * X.inv(
).T - X.inv().T * A * X.inv() * C.inv() * X.inv().T
# Cookbook example 126:
expr = Trace((X.T * C * X).inv() * (X.T * B * X))
assert expr.diff(X) == -2 * C * X * (X.T * C * X).inv() * X.T * B * X * (
X.T * C * X).inv() + 2 * B * X * (X.T * C * X).inv()
# Cookbook example 127:
expr = Trace((A + X.T * C * X).inv() * (X.T * B * X))
# Warning: result in the cookbook is equivalent if B and C are symmetric:
assert expr.diff(X) == B * X * Inverse(A + X.T * C * X) - C * X * Inverse(
A + X.T * C *
X) * X.T * B * X * Inverse(A + X.T * C * X) - C.T * X * Inverse(
A.T + (C * X).T * X) * X.T * B.T * X * Inverse(
A.T + (C * X).T * X) + B.T * X * Inverse(A.T + (C * X).T * 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 _a2m_add(*args):
if all(not isinstance(i, _CodegenArrayAbstract) for i in args):
from sympy import MatAdd
return MatAdd(*args).doit()
else:
return ArrayAdd(*args)
def test_arrayexpr_array_flatten():
# Flatten nested ArrayTensorProduct objects:
expr1 = ArrayTensorProduct(M, N)
expr2 = ArrayTensorProduct(P, Q)
expr = ArrayTensorProduct(expr1, expr2)
assert expr == ArrayTensorProduct(M, N, P, Q)
assert expr.args == (M, N, P, Q)
# Flatten mixed ArrayTensorProduct and ArrayContraction objects:
cg1 = ArrayContraction(expr1, (1, 2))
cg2 = ArrayContraction(expr2, (0, 3))
expr = ArrayTensorProduct(cg1, cg2)
assert expr == ArrayContraction(ArrayTensorProduct(M, N, P, Q), (1, 2),
(4, 7))
expr = ArrayTensorProduct(M, cg1)
assert expr == ArrayContraction(ArrayTensorProduct(M, M, N), (3, 4))
# Flatten nested ArrayContraction objects:
cgnested = ArrayContraction(cg1, (0, 1))
assert cgnested == ArrayContraction(ArrayTensorProduct(M, N), (0, 3),
(1, 2))
cgnested = ArrayContraction(ArrayTensorProduct(cg1, cg2), (0, 3))
assert cgnested == ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 6),
(1, 2), (4, 7))
cg3 = ArrayContraction(ArrayTensorProduct(M, N, P, Q), (1, 3), (2, 4))
cgnested = ArrayContraction(cg3, (0, 1))
assert cgnested == ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 5),
(1, 3), (2, 4))
cgnested = ArrayContraction(cg3, (0, 3), (1, 2))
assert cgnested == ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 7),
(1, 3), (2, 4), (5, 6))
cg4 = ArrayContraction(ArrayTensorProduct(M, N, P, Q), (1, 5), (3, 7))
cgnested = ArrayContraction(cg4, (0, 1))
assert cgnested == ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 2),
(1, 5), (3, 7))
cgnested = ArrayContraction(cg4, (0, 1), (2, 3))
assert cgnested == ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 2),
(1, 5), (3, 7), (4, 6))
cg = ArrayDiagonal(cg4)
assert cg == cg4
assert isinstance(cg, type(cg4))
# Flatten nested ArrayDiagonal objects:
cg1 = ArrayDiagonal(expr1, (1, 2))
cg2 = ArrayDiagonal(expr2, (0, 3))
cg3 = ArrayDiagonal(ArrayTensorProduct(M, N, P, Q), (1, 3), (2, 4))
cg4 = ArrayDiagonal(ArrayTensorProduct(M, N, P, Q), (1, 5), (3, 7))
cgnested = ArrayDiagonal(cg1, (0, 1))
assert cgnested == ArrayDiagonal(ArrayTensorProduct(M, N), (1, 2), (0, 3))
cgnested = ArrayDiagonal(cg3, (1, 2))
assert cgnested == ArrayDiagonal(ArrayTensorProduct(M, N, P, Q), (1, 3),
(2, 4), (5, 6))
cgnested = ArrayDiagonal(cg4, (1, 2))
assert cgnested == ArrayDiagonal(ArrayTensorProduct(M, N, P, Q), (1, 5),
(3, 7), (2, 4))
cg = ArrayAdd(M, N)
cg2 = ArrayAdd(cg, P)
assert isinstance(cg2, ArrayAdd)
assert cg2.args == (M, N, P)
assert cg2.shape == (k, k)
expr = ArrayTensorProduct(ArrayDiagonal(X, (0, 1)),
ArrayDiagonal(A, (0, 1)))
assert expr == ArrayDiagonal(ArrayTensorProduct(X, A), (0, 1), (2, 3))
expr1 = ArrayDiagonal(ArrayTensorProduct(X, A), (1, 2))
expr2 = ArrayTensorProduct(expr1, a)
assert expr2 == PermuteDims(
ArrayDiagonal(ArrayTensorProduct(X, A, a), (1, 2)), [0, 1, 3, 4, 2])
expr1 = ArrayContraction(ArrayTensorProduct(X, A), (1, 2))
expr2 = ArrayTensorProduct(expr1, a)
assert isinstance(expr2, ArrayContraction)
assert isinstance(expr2.expr, ArrayTensorProduct)
def test_arrayexpr_convert_array_to_matrix():
cg = ArrayContraction(ArrayTensorProduct(M), (0, 1))
assert convert_array_to_matrix(cg) == Trace(M)
cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 1), (2, 3))
assert convert_array_to_matrix(cg) == Trace(M) * Trace(N)
cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 3), (1, 2))
assert convert_array_to_matrix(cg) == Trace(M * N)
cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 2), (1, 3))
assert convert_array_to_matrix(cg) == Trace(M * N.T)
cg = convert_matrix_to_array(M * N * P)
assert convert_array_to_matrix(cg) == M * N * P
cg = convert_matrix_to_array(M * N.T * P)
assert convert_array_to_matrix(cg) == M * N.T * P
cg = ArrayContraction(ArrayTensorProduct(M, N, P, Q), (1, 2), (5, 6))
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M * N, P * Q)
cg = ArrayContraction(ArrayTensorProduct(-2, M, N), (1, 2))
assert convert_array_to_matrix(cg) == -2 * M * N
a = MatrixSymbol("a", k, 1)
b = MatrixSymbol("b", k, 1)
c = MatrixSymbol("c", k, 1)
cg = PermuteDims(
ArrayContraction(
ArrayTensorProduct(
a,
ArrayAdd(
ArrayTensorProduct(b, c),
ArrayTensorProduct(c, b),
)), (2, 4)), [0, 1, 3, 2])
assert convert_array_to_matrix(cg) == a * (b.T * c + c.T * b)
za = ZeroArray(m, n)
assert convert_array_to_matrix(za) == ZeroMatrix(m, n)
cg = ArrayTensorProduct(3, M)
assert convert_array_to_matrix(cg) == 3 * M
# TODO: not yet supported:
# cg = ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 2, 4), (1, 3, 5))
# assert recognize_matrix_expression(cg) == HadamardProduct(M, N, P)
# cg = ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 3, 4), (1, 2, 5))
# assert recognize_matrix_expression(cg) == HadamardProduct(M, N.T, P)
x = MatrixSymbol("x", k, 1)
cg = PermuteDims(
ArrayContraction(
ArrayTensorProduct(OneArray(1), x, OneArray(1),
DiagMatrix(Identity(1))), (0, 5)),
Permutation(1, 2, 3))
assert convert_array_to_matrix(cg) == x
expr = ArrayAdd(M, PermuteDims(M, [1, 0]))
assert convert_array_to_matrix(expr) == M + Transpose(M)
def _(expr: ArrayAdd, x: Expr):
return ArrayAdd(*[array_derive(arg, x) for arg in expr.args])