def test_arrayexpr_convert_matrix_to_array():
expr = M * N
result = ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
assert convert_matrix_to_array(expr) == result
expr = M * N * M
result = _array_contraction(ArrayTensorProduct(M, N, M), (1, 2), (3, 4))
assert convert_matrix_to_array(expr) == result
expr = Transpose(M)
assert convert_matrix_to_array(expr) == PermuteDims(M, [1, 0])
expr = M * Transpose(N)
assert convert_matrix_to_array(expr) == _array_contraction(
_array_tensor_product(M, PermuteDims(N, [1, 0])), (1, 2))
expr = 3 * M * N
res = convert_matrix_to_array(expr)
rexpr = convert_array_to_matrix(res)
assert expr == rexpr
expr = 3 * M + N * M.T * M + 4 * k * N
res = convert_matrix_to_array(expr)
rexpr = convert_array_to_matrix(res)
assert expr == rexpr
expr = Inverse(M) * N
rexpr = convert_array_to_matrix(convert_matrix_to_array(expr))
assert expr == rexpr
expr = M**2
rexpr = convert_array_to_matrix(convert_matrix_to_array(expr))
assert expr == rexpr
expr = M * (2 * N + 3 * M)
res = convert_matrix_to_array(expr)
rexpr = convert_array_to_matrix(res)
assert expr == rexpr
expr = Trace(M)
result = ArrayContraction(M, (0, 1))
assert convert_matrix_to_array(expr) == result
expr = 3 * Trace(M)
result = ArrayContraction(ArrayTensorProduct(3, M), (0, 1))
assert convert_matrix_to_array(expr) == result
expr = 3 * Trace(Trace(M) * M)
result = ArrayContraction(ArrayTensorProduct(3, M, M), (0, 1), (2, 3))
assert convert_matrix_to_array(expr) == result
expr = 3 * Trace(M)**2
result = ArrayContraction(ArrayTensorProduct(3, M, M), (0, 1), (2, 3))
assert convert_matrix_to_array(expr) == result
expr = HadamardProduct(M, N)
result = ArrayDiagonal(ArrayTensorProduct(M, N), (0, 2), (1, 3))
assert convert_matrix_to_array(expr) == result
expr = HadamardProduct(M * N, N * M)
result = ArrayDiagonal(
ArrayContraction(ArrayTensorProduct(M, N, N, M), (1, 2), (5, 6)),
(0, 2), (1, 3))
assert convert_matrix_to_array(expr) == result
expr = HadamardPower(M, 2)
result = ArrayDiagonal(ArrayTensorProduct(M, M), (0, 2), (1, 3))
assert convert_matrix_to_array(expr) == result
expr = HadamardPower(M * N, 2)
result = ArrayDiagonal(
ArrayContraction(ArrayTensorProduct(M, N, M, N), (1, 2), (5, 6)),
(0, 2), (1, 3))
assert convert_matrix_to_array(expr) == result
expr = HadamardPower(M, n)
d0 = Dummy("d0")
result = ArrayElementwiseApplyFunc(Lambda(d0, d0**n), M)
assert convert_matrix_to_array(expr).dummy_eq(result)
expr = M**2
assert isinstance(expr, MatPow)
assert convert_matrix_to_array(expr) == ArrayContraction(
ArrayTensorProduct(M, M), (1, 2))
expr = a.T * b
cg = convert_matrix_to_array(expr)
assert cg == ArrayContraction(ArrayTensorProduct(a, b), (0, 2))
r"""
Array expressions are expressions representing N-dimensional arrays, without
evaluating them. These expressions represent in a certain way abstract syntax
trees of operations on N-dimensional arrays.
Every N-dimensional array operator has a corresponding array expression object.
Table of correspondences:
=============================== =============================
Array operator Array expression operator
=============================== =============================
tensorproduct ArrayTensorProduct
tensorcontraction ArrayContraction
tensordiagonal ArrayDiagonal
permutedims PermuteDims
=============================== =============================
Examples
========
``ArraySymbol`` objects are the N-dimensional equivalent of ``MatrixSymbol``
objects in the matrix module:
>>> from sympy.tensor.array.expressions import ArraySymbol
>>> from sympy.abc import i, j, k
>>> A = ArraySymbol("A", (3, 2, 4))
>>> A.shape
(3, 2, 4)
>>> A[i, j, k]
def test_arrayexpr_convert_matrix_to_array():
expr = M * N
result = ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
assert convert_matrix_to_array(expr) == result
expr = M * N * M
result = ArrayContraction(ArrayTensorProduct(M, N, M), (1, 2), (3, 4))
assert convert_matrix_to_array(expr) == result
expr = Transpose(M)
assert convert_matrix_to_array(expr) == PermuteDims(M, [1, 0])
expr = M * Transpose(N)
assert convert_matrix_to_array(expr) == ArrayContraction(
ArrayTensorProduct(M, PermuteDims(N, [1, 0])), (1, 2))
expr = 3 * M * N
res = convert_matrix_to_array(expr)
rexpr = convert_array_to_matrix(res)
assert expr == rexpr
expr = 3 * M + N * M.T * M + 4 * k * N
res = convert_matrix_to_array(expr)
rexpr = convert_array_to_matrix(res)
assert expr == rexpr
expr = Inverse(M) * N
rexpr = convert_array_to_matrix(convert_matrix_to_array(expr))
assert expr == rexpr
expr = M**2
rexpr = convert_array_to_matrix(convert_matrix_to_array(expr))
assert expr == rexpr
expr = M * (2 * N + 3 * M)
res = convert_matrix_to_array(expr)
rexpr = convert_array_to_matrix(res)
assert expr == rexpr
expr = Trace(M)
result = ArrayContraction(M, (0, 1))
assert convert_matrix_to_array(expr) == result
expr = 3 * Trace(M)
result = ArrayContraction(ArrayTensorProduct(3, M), (0, 1))
assert convert_matrix_to_array(expr) == result
expr = 3 * Trace(Trace(M) * M)
result = ArrayContraction(ArrayTensorProduct(3, M, M), (0, 1), (2, 3))
assert convert_matrix_to_array(expr) == result
expr = 3 * Trace(M)**2
result = ArrayContraction(ArrayTensorProduct(3, M, M), (0, 1), (2, 3))
assert convert_matrix_to_array(expr) == result
expr = HadamardProduct(M, N)
result = ArrayDiagonal(ArrayTensorProduct(M, N), (0, 2), (1, 3))
assert convert_matrix_to_array(expr) == result
expr = HadamardPower(M, 2)
result = ArrayDiagonal(ArrayTensorProduct(M, M), (0, 2), (1, 3))
assert convert_matrix_to_array(expr) == result
expr = M**2
assert isinstance(expr, MatPow)
assert convert_matrix_to_array(expr) == ArrayContraction(
ArrayTensorProduct(M, M), (1, 2))
expr = a.T * b
cg = convert_matrix_to_array(expr)
assert cg == ArrayContraction(ArrayTensorProduct(a, b), (0, 2))
def test_matrix_expression_from_index_summation():
from sympy.abc import a,b,c,d
A = MatrixSymbol("A", k, k)
B = MatrixSymbol("B", k, k)
C = MatrixSymbol("C", k, k)
w1 = MatrixSymbol("w1", k, 1)
i0, i1, i2, i3, i4 = symbols("i0:5", cls=Dummy)
expr = Sum(W[a,b]*X[b,c]*Z[c,d], (b, 0, l-1), (c, 0, m-1))
assert MatrixExpr.from_index_summation(expr, a) == W*X*Z
expr = Sum(W.T[b,a]*X[b,c]*Z[c,d], (b, 0, l-1), (c, 0, m-1))
assert MatrixExpr.from_index_summation(expr, a) == W*X*Z
expr = Sum(A[b, a]*B[b, c]*C[c, d], (b, 0, k-1), (c, 0, k-1))
assert MatrixSymbol.from_index_summation(expr, a) == A.T*B*C
expr = Sum(A[b, a]*B[c, b]*C[c, d], (b, 0, k-1), (c, 0, k-1))
assert MatrixSymbol.from_index_summation(expr, a) == A.T*B.T*C
expr = Sum(C[c, d]*A[b, a]*B[c, b], (b, 0, k-1), (c, 0, k-1))
assert MatrixSymbol.from_index_summation(expr, a) == A.T*B.T*C
expr = Sum(A[a, b] + B[a, b], (a, 0, k-1), (b, 0, k-1))
assert MatrixExpr.from_index_summation(expr, a) == OneMatrix(1, k)*A*OneMatrix(k, 1) + OneMatrix(1, k)*B*OneMatrix(k, 1)
expr = Sum(A[a, b]**2, (a, 0, k - 1), (b, 0, k - 1))
assert MatrixExpr.from_index_summation(expr, a) == Trace(A * A.T)
expr = Sum(A[a, b]**3, (a, 0, k - 1), (b, 0, k - 1))
assert MatrixExpr.from_index_summation(expr, a) == Trace(HadamardPower(A.T, 2) * A)
expr = Sum((A[a, b] + B[a, b])*C[b, c], (b, 0, k-1))
assert MatrixExpr.from_index_summation(expr, a) == (A+B)*C
expr = Sum((A[a, b] + B[b, a])*C[b, c], (b, 0, k-1))
assert MatrixExpr.from_index_summation(expr, a) == (A+B.T)*C
expr = Sum(A[a, b]*A[b, c]*A[c, d], (b, 0, k-1), (c, 0, k-1))
assert MatrixExpr.from_index_summation(expr, a) == A**3
expr = Sum(A[a, b]*A[b, c]*B[c, d], (b, 0, k-1), (c, 0, k-1))
assert MatrixExpr.from_index_summation(expr, a) == A**2*B
# Parse the trace of a matrix:
expr = Sum(A[a, a], (a, 0, k-1))
assert MatrixExpr.from_index_summation(expr, None) == trace(A)
expr = Sum(A[a, a]*B[b, c]*C[c, d], (a, 0, k-1), (c, 0, k-1))
assert MatrixExpr.from_index_summation(expr, b) == trace(A)*B*C
# Check wrong sum ranges (should raise an exception):
## Case 1: 0 to m instead of 0 to m-1
expr = Sum(W[a,b]*X[b,c]*Z[c,d], (b, 0, l-1), (c, 0, m))
raises(ValueError, lambda: MatrixExpr.from_index_summation(expr, a))
## Case 2: 1 to m-1 instead of 0 to m-1
expr = Sum(W[a,b]*X[b,c]*Z[c,d], (b, 0, l-1), (c, 1, m-1))
raises(ValueError, lambda: MatrixExpr.from_index_summation(expr, a))
# Parse nested sums:
expr = Sum(A[a, b]*Sum(B[b, c]*C[c, d], (c, 0, k-1)), (b, 0, k-1))
assert MatrixExpr.from_index_summation(expr, a) == A*B*C
# Test Kronecker delta:
expr = Sum(A[a, b]*KroneckerDelta(b, c)*B[c, d], (b, 0, k-1), (c, 0, k-1))
assert MatrixExpr.from_index_summation(expr, a) == A*B
expr = Sum(KroneckerDelta(i1, m)*KroneckerDelta(i2, n)*A[i, i1]*A[j, i2], (i1, 0, k-1), (i2, 0, k-1))
assert MatrixExpr.from_index_summation(expr, m) == ArrayTensorProduct(A.T, A)
# Test numbered indices:
expr = Sum(A[i1, i2]*w1[i2, 0], (i2, 0, k-1))
assert MatrixExpr.from_index_summation(expr, i1) == MatrixElement(A*w1, i1, 0)
expr = Sum(A[i1, i2]*B[i2, 0], (i2, 0, k-1))
assert MatrixExpr.from_index_summation(expr, i1) == MatrixElement(A*B, i1, 0)
def test_arrayexpr_convert_array_to_diagonalized_vector():
# Check matrix recognition over trivial dimensions:
cg = ArrayTensorProduct(a, b)
assert convert_array_to_matrix(cg) == a * b.T
cg = ArrayTensorProduct(I1, a, b)
assert convert_array_to_matrix(cg) == a * b.T
# Recognize trace inside a tensor product:
cg = ArrayContraction(ArrayTensorProduct(A, B, C), (0, 3), (1, 2))
assert convert_array_to_matrix(cg) == Trace(A * B) * C
# Transform diagonal operator to contraction:
cg = ArrayDiagonal(ArrayTensorProduct(A, a), (1, 2))
assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(A, OneArray(1), DiagMatrix(a)), (1, 3))
assert convert_array_to_matrix(cg) == A * DiagMatrix(a)
cg = ArrayDiagonal(ArrayTensorProduct(a, b), (0, 2))
assert _array_diag2contr_diagmatrix(cg) == PermuteDims(
ArrayContraction(ArrayTensorProduct(DiagMatrix(a), OneArray(1), b), (0, 3)), [1, 2, 0]
)
assert convert_array_to_matrix(cg) == b.T * DiagMatrix(a)
cg = ArrayDiagonal(ArrayTensorProduct(A, a), (0, 2))
assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(A, OneArray(1), DiagMatrix(a)), (0, 3))
assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a)
cg = ArrayDiagonal(ArrayTensorProduct(I, x, I1), (0, 2), (3, 5))
assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(I, OneArray(1), I1, DiagMatrix(x)), (0, 5))
assert convert_array_to_matrix(cg) == DiagMatrix(x)
cg = ArrayDiagonal(ArrayTensorProduct(I, x, A, B), (1, 2), (5, 6))
assert _array_diag2contr_diagmatrix(cg) == ArrayDiagonal(ArrayContraction(ArrayTensorProduct(I, OneArray(1), A, B, DiagMatrix(x)), (1, 7)), (5, 6))
# TODO: this is returning a wrong result:
# convert_array_to_matrix(cg)
cg = ArrayDiagonal(ArrayTensorProduct(I1, a, b), (1, 3, 5))
assert convert_array_to_matrix(cg) == a*b.T
cg = ArrayDiagonal(ArrayTensorProduct(I1, a, b), (1, 3))
assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(OneArray(1), a, b, I1), (2, 6))
assert convert_array_to_matrix(cg) == a*b.T
cg = ArrayDiagonal(ArrayTensorProduct(x, I1), (1, 2))
assert isinstance(cg, ArrayDiagonal)
assert cg.diagonal_indices == ((1, 2),)
assert convert_array_to_matrix(cg) == x
cg = ArrayDiagonal(ArrayTensorProduct(x, I), (0, 2))
assert _array_diag2contr_diagmatrix(cg) == ArrayContraction(ArrayTensorProduct(OneArray(1), I, DiagMatrix(x)), (1, 3))
assert convert_array_to_matrix(cg).doit() == DiagMatrix(x)
raises(ValueError, lambda: ArrayDiagonal(x, (1,)))
# Ignore identity matrices with contractions:
cg = ArrayContraction(ArrayTensorProduct(I, A, I, I), (0, 2), (1, 3), (5, 7))
assert cg.split_multiple_contractions() == cg
assert convert_array_to_matrix(cg) == Trace(A) * I
cg = ArrayContraction(ArrayTensorProduct(Trace(A) * I, I, I), (1, 5), (3, 4))
assert cg.split_multiple_contractions() == cg
assert convert_array_to_matrix(cg).doit() == Trace(A) * I
# Add DiagMatrix when required:
cg = ArrayContraction(ArrayTensorProduct(A, a), (1, 2))
assert cg.split_multiple_contractions() == cg
assert convert_array_to_matrix(cg) == A * a
cg = ArrayContraction(ArrayTensorProduct(A, a, B), (1, 2, 4))
assert cg.split_multiple_contractions() == ArrayContraction(ArrayTensorProduct(A, DiagMatrix(a), OneArray(1), B), (1, 2), (3, 5))
assert convert_array_to_matrix(cg) == A * DiagMatrix(a) * B
cg = ArrayContraction(ArrayTensorProduct(A, a, B), (0, 2, 4))
assert cg.split_multiple_contractions() == ArrayContraction(ArrayTensorProduct(A, DiagMatrix(a), OneArray(1), B), (0, 2), (3, 5))
assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) * B
cg = ArrayContraction(ArrayTensorProduct(A, a, b, a.T, B), (0, 2, 4, 7, 9))
assert cg.split_multiple_contractions() == ArrayContraction(ArrayTensorProduct(A, DiagMatrix(a), OneArray(1),
DiagMatrix(b), OneArray(1), DiagMatrix(a), OneArray(1), B),
(0, 2), (3, 5), (6, 9), (8, 12))
assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) * DiagMatrix(b) * DiagMatrix(a) * B.T
cg = ArrayContraction(ArrayTensorProduct(I1, I1, I1), (1, 2, 4))
assert cg.split_multiple_contractions() == ArrayContraction(ArrayTensorProduct(I1, I1, OneArray(1), I1), (1, 2), (3, 5))
assert convert_array_to_matrix(cg) == 1
cg = ArrayContraction(ArrayTensorProduct(I, I, I, I, A), (1, 2, 8), (5, 6, 9))
assert convert_array_to_matrix(cg.split_multiple_contractions()).doit() == A
cg = ArrayContraction(ArrayTensorProduct(A, a, C, a, B), (1, 2, 4), (5, 6, 8))
expected = ArrayContraction(ArrayTensorProduct(A, DiagMatrix(a), OneArray(1), C, DiagMatrix(a), OneArray(1), B), (1, 3), (2, 5), (6, 7), (8, 10))
assert cg.split_multiple_contractions() == expected
assert convert_array_to_matrix(cg) == A * DiagMatrix(a) * C * DiagMatrix(a) * B
cg = ArrayContraction(ArrayTensorProduct(a, I1, b, I1, (a.T*b).applyfunc(cos)), (1, 2, 8), (5, 6, 9))
expected = ArrayContraction(ArrayTensorProduct(a, I1, OneArray(1), b, I1, OneArray(1), (a.T*b).applyfunc(cos)),
(1, 3), (2, 10), (6, 8), (7, 11))
assert cg.split_multiple_contractions().dummy_eq(expected)
assert convert_array_to_matrix(cg).doit().dummy_eq(MatMul(a, (a.T * b).applyfunc(cos), b.T))
def test_arrayexpr_convert_array_to_matrix_remove_trivial_dims():
# Tensor Product:
assert _remove_trivial_dims(ArrayTensorProduct(a, b)) == (a * b.T, [1, 3])
assert _remove_trivial_dims(ArrayTensorProduct(a.T,
b)) == (a * b.T, [0, 3])
assert _remove_trivial_dims(ArrayTensorProduct(a,
b.T)) == (a * b.T, [1, 2])
assert _remove_trivial_dims(ArrayTensorProduct(a.T,
b.T)) == (a * b.T, [0, 2])
assert _remove_trivial_dims(ArrayTensorProduct(I, a.T,
b.T)) == (a * b.T,
[0, 1, 2, 4])
assert _remove_trivial_dims(ArrayTensorProduct(a.T, I,
b.T)) == (a * b.T,
[0, 2, 3, 4])
assert _remove_trivial_dims(ArrayTensorProduct(a, I)) == (a, [2, 3])
assert _remove_trivial_dims(ArrayTensorProduct(I, a)) == (a, [0, 1])
assert _remove_trivial_dims(ArrayTensorProduct(a.T, b.T, c,
d)) == (ArrayTensorProduct(
a * b.T,
c * d.T), [0, 2, 5, 7])
assert _remove_trivial_dims(ArrayTensorProduct(a.T, I, b.T, c, d,
I)) == (ArrayTensorProduct(
a * b.T, c * d.T,
I), [0, 2, 3, 4, 7, 9])
# Addition:
cg = ArrayAdd(ArrayTensorProduct(a, b), ArrayTensorProduct(c, d))
assert _remove_trivial_dims(cg) == (a * b.T + c * d.T, [1, 3])
# Permute Dims:
cg = PermuteDims(ArrayTensorProduct(a, b), Permutation(3)(1, 2))
assert _remove_trivial_dims(cg) == (a * b.T, [2, 3])
cg = PermuteDims(ArrayTensorProduct(a, I, b), Permutation(5)(1, 2, 3, 4))
assert _remove_trivial_dims(cg) == (a * b.T, [1, 2, 4, 5])
cg = PermuteDims(ArrayTensorProduct(I, b, a),
Permutation(5)(1, 2, 4, 5, 3))
assert _remove_trivial_dims(cg) == (b * a.T, [0, 3, 4, 5])
# Diagonal:
cg = ArrayDiagonal(ArrayTensorProduct(M, a), (1, 2))
assert _remove_trivial_dims(cg) == (cg, [])
# Contraction:
cg = ArrayContraction(ArrayTensorProduct(M, a), (1, 2))
assert _remove_trivial_dims(cg) == (cg, [])
def test_arrayexpr_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_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 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(
[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 ArrayDiagonal(tp, *diag)
elif isinstance(expr, HadamardPower):
base, exp = expr.args
if exp.is_Integer and exp > 0:
return convert_matrix_to_array(
HadamardProduct.fromiter(base for i in range(exp)))
else:
raise NotImplementedError(
"conversion of Hadamard symbolic power is currently not supported"
)
else:
return expr
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_arrayexpr_convert_array_to_matrix_support_function():
assert _support_function_tp1_recognize([], [2 * k]) == 2 * k
assert _support_function_tp1_recognize([(1, 2)],
[A, 2 * k, B, 3]) == 6 * k * A * B
assert _support_function_tp1_recognize([(0, 3), (1, 2)],
[A, B]) == Trace(A * B)
assert _support_function_tp1_recognize([(1, 2)], [A, B]) == A * B
assert _support_function_tp1_recognize([(0, 2)], [A, B]) == A.T * B
assert _support_function_tp1_recognize([(1, 3)], [A, B]) == A * B.T
assert _support_function_tp1_recognize([(0, 3)], [A, B]) == A.T * B.T
assert _support_function_tp1_recognize([(1, 2), (5, 6)],
[A, B, C, D]) == ArrayTensorProduct(
A * B, C * D)
assert _support_function_tp1_recognize(
[(1, 4), (3, 6)],
[A, B, C, D]) == PermuteDims(ArrayTensorProduct(A * C, B * D),
[0, 2, 1, 3])
assert _support_function_tp1_recognize([(0, 3), (1, 4)],
[A, B, C]) == B * A * C
assert _support_function_tp1_recognize(
[(9, 10), (1, 2), (5, 6), (3, 4),
(7, 8)], [X, Y, A, B, C, D]) == X * Y * A * B * C * D
assert _support_function_tp1_recognize(
[(9, 10), (1, 2), (5, 6), (3, 4)],
[X, Y, A, B, C, D]) == ArrayTensorProduct(X * Y * A * B, C * D)
assert _support_function_tp1_recognize(
[(1, 7), (3, 8), (4, 11)], [X, Y, A, B, C, D]) == PermuteDims(
ArrayTensorProduct(X * B.T, Y * C, D * A), [0, 2, 5, 1, 3, 4])
assert _support_function_tp1_recognize(
[(0, 1), (3, 6), (5, 8)], [X, A, B, C, D]) == PermuteDims(
ArrayTensorProduct(Trace(X) * A * C, B * D), [0, 2, 1, 3])
assert _support_function_tp1_recognize([(1, 2), (3, 4), (5, 6), (7, 8)],
[A, A, B, C, D]) == A**2 * B * C * D
assert _support_function_tp1_recognize(
[(1, 2), (3, 4), (5, 6), (7, 8)], [X, A, B, C, D]) == X * A * B * C * D
assert _support_function_tp1_recognize(
[(1, 6), (3, 8), (5, 10)], [X, Y, A, B, C, D]) == PermuteDims(
ArrayTensorProduct(X * B, Y * C, A * D), [0, 2, 4, 1, 3, 5])
assert _support_function_tp1_recognize(
[(1, 4), (3, 6)],
[A, B, C, D]) == PermuteDims(ArrayTensorProduct(A * C, B * D),
[0, 2, 1, 3])
assert _support_function_tp1_recognize(
[(0, 4), (1, 7), (2, 5),
(3, 8)], [X, A, B, C, D]) == C * X.T * B * A * D
assert _support_function_tp1_recognize(
[(0, 4), (1, 7), (2, 5),
(3, 8)], [X, A, B, C, D]) == C * X.T * B * A * D
def 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 test_arrayexpr_convert_array_to_matrix_diag2contraction_diagmatrix():
cg = ArrayDiagonal(ArrayTensorProduct(M, a), (1, 2))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == ArrayContraction(
ArrayTensorProduct(M, OneArray(1), DiagMatrix(a)), (1, 3))
raises(ValueError, lambda: ArrayDiagonal(ArrayTensorProduct(a, M), (1, 2)))
cg = ArrayDiagonal(ArrayTensorProduct(a.T, M), (1, 2))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == ArrayContraction(
ArrayTensorProduct(OneArray(1), M, DiagMatrix(a.T)), (1, 4))
cg = ArrayDiagonal(ArrayTensorProduct(a.T, M, N, b.T), (1, 2), (4, 7))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == ArrayContraction(
ArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a.T),
DiagMatrix(b.T)), (1, 7), (3, 9))
cg = ArrayDiagonal(ArrayTensorProduct(a, M, N, b.T), (0, 2), (4, 7))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == ArrayContraction(
ArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a),
DiagMatrix(b.T)), (1, 6), (3, 9))
cg = ArrayDiagonal(ArrayTensorProduct(a, M, N, b.T), (0, 4), (3, 7))
res = _array_diag2contr_diagmatrix(cg)
assert res.shape == cg.shape
assert res == ArrayContraction(
ArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a),
DiagMatrix(b.T)), (3, 6), (2, 9))
I1 = Identity(1)
x = MatrixSymbol("x", k, 1)
A = MatrixSymbol("A", k, k)
cg = ArrayDiagonal(ArrayTensorProduct(x, A.T, I1), (0, 2))
assert _array_diag2contr_diagmatrix(cg).shape == cg.shape
assert _array2matrix(cg).shape == cg.shape
def test_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_array_diagonal():
cg = ArrayDiagonal(M, (1, 0))
assert cg == ArrayDiagonal(M, (0, 1))
cg = ArrayDiagonal(ArrayTensorProduct(M, N, P), (4, 1), (2, 0))
assert cg == ArrayDiagonal(ArrayTensorProduct(M, N, P), (1, 4), (0, 2))
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_arrayexpr_permutedims_sink():
cg = PermuteDims(ArrayTensorProduct(M, N), [0, 1, 3, 2],
nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == ArrayTensorProduct(M, PermuteDims(N, [1, 0]))
cg = PermuteDims(ArrayTensorProduct(M, N), [1, 0, 3, 2],
nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == ArrayTensorProduct(PermuteDims(M, [1, 0]),
PermuteDims(N, [1, 0]))
cg = PermuteDims(ArrayTensorProduct(M, N), [3, 2, 1, 0],
nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == ArrayTensorProduct(PermuteDims(N, [1, 0]),
PermuteDims(M, [1, 0]))
cg = PermuteDims(ArrayContraction(ArrayTensorProduct(M, N), (1, 2)),
[1, 0],
nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == ArrayContraction(
PermuteDims(ArrayTensorProduct(M, N), [[0, 3]]), (1, 2))
cg = PermuteDims(ArrayTensorProduct(M, N), [1, 0, 3, 2],
nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == ArrayTensorProduct(PermuteDims(M, [1, 0]),
PermuteDims(N, [1, 0]))
cg = PermuteDims(ArrayContraction(ArrayTensorProduct(M, N, P), (1, 2),
(3, 4)), [1, 0],
nest_permutation=False)
sunk = nest_permutation(cg)
assert sunk == ArrayContraction(
PermuteDims(ArrayTensorProduct(M, N, P), [[0, 5]]), (1, 2), (3, 4))
def test_arrayexpr_convert_array_to_matrix2():
cg = ArrayContraction(ArrayTensorProduct(M, N), (1, 3))
assert convert_array_to_matrix(cg) == M * N.T
cg = PermuteDims(ArrayTensorProduct(M, N), Permutation([0, 1, 3, 2]))
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M, N.T)
cg = ArrayTensorProduct(M, PermuteDims(N, Permutation([1, 0])))
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M, N.T)
cg = ArrayContraction(
PermuteDims(
ArrayTensorProduct(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])),
(1, 2), (3, 5)
)
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M * P.T * Trace(N), Q.T)
cg = ArrayContraction(
ArrayTensorProduct(M, N, P, PermuteDims(Q, Permutation([1, 0]))),
(1, 5), (2, 3)
)
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M * P.T * Trace(N), Q.T)
cg = ArrayTensorProduct(M, PermuteDims(N, [1, 0]))
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M, N.T)
cg = ArrayTensorProduct(PermuteDims(M, [1, 0]), PermuteDims(N, [1, 0]))
assert convert_array_to_matrix(cg) == ArrayTensorProduct(M.T, N.T)
cg = ArrayTensorProduct(PermuteDims(N, [1, 0]), PermuteDims(M, [1, 0]))
assert convert_array_to_matrix(cg) == ArrayTensorProduct(N.T, M.T)
def test_arrayexpr_array_wrong_permutation_size():
cg = ArrayTensorProduct(M, N)
raises(ValueError, lambda: PermuteDims(cg, [1, 0]))
raises(ValueError, lambda: PermuteDims(cg, [1, 0, 2, 3, 5, 4]))
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 _(expr: MatrixSymbol, x: Expr):
m, n = expr.shape
if expr == x:
return PermuteDims(ArrayTensorProduct(Identity(m), Identity(n)),
[0, 2, 1, 3])
return ZeroArray(*(x.shape + expr.shape))
