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))