def test_array_as_explicit_call(): assert ZeroArray(3, 2, 4).as_explicit() == ImmutableDenseNDimArray.zeros( 3, 2, 4) assert OneArray(3, 2, 4).as_explicit() == ImmutableDenseNDimArray( [1 for i in range(3 * 2 * 4)]).reshape(3, 2, 4) k = Symbol("k") X = ArraySymbol("X", k, 3, 2) raises(ValueError, lambda: X.as_explicit()) raises(ValueError, lambda: ZeroArray(k, 2, 3).as_explicit()) raises(ValueError, lambda: OneArray(2, k, 2).as_explicit()) A = ArraySymbol("A", 3, 3) B = ArraySymbol("B", 3, 3) texpr = tensorproduct(A, B) assert isinstance(texpr, ArrayTensorProduct) assert texpr.as_explicit() == tensorproduct(A.as_explicit(), B.as_explicit()) texpr = tensorcontraction(A, (0, 1)) assert isinstance(texpr, ArrayContraction) assert texpr.as_explicit() == A[0, 0] + A[1, 1] + A[2, 2] texpr = tensordiagonal(A, (0, 1)) assert isinstance(texpr, ArrayDiagonal) assert texpr.as_explicit() == ImmutableDenseNDimArray( [A[0, 0], A[1, 1], A[2, 2]]) texpr = permutedims(A, [1, 0]) assert isinstance(texpr, PermuteDims) assert texpr.as_explicit() == permutedims(A.as_explicit(), [1, 0])
def test_array_printer(): A = ArraySymbol('A', (4, 4, 6, 6, 6)) I = IndexedBase('I') i, j, k = Idx('i', (0, 1)), Idx('j', (2, 3)), Idx('k', (4, 5)) prntr = NumPyPrinter() assert prntr.doprint(ZeroArray(5)) == 'numpy.zeros((5,))' assert prntr.doprint(OneArray(5)) == 'numpy.ones((5,))' assert prntr.doprint(ArrayContraction( A, [2, 3])) == 'numpy.einsum("abccd->abd", A)' assert prntr.doprint(I) == 'I' assert prntr.doprint(ArrayDiagonal( A, [2, 3, 4])) == 'numpy.einsum("abccc->abc", A)' assert prntr.doprint(ArrayDiagonal( A, [0, 1], [2, 3])) == 'numpy.einsum("aabbc->cab", A)' assert prntr.doprint(ArrayContraction( A, [2], [3])) == 'numpy.einsum("abcde->abe", A)' assert prntr.doprint(Assignment(I[i, j, k], I[i, j, k])) == 'I = I' prntr = TensorflowPrinter() assert prntr.doprint(ZeroArray(5)) == 'tensorflow.zeros((5,))' assert prntr.doprint(OneArray(5)) == 'tensorflow.ones((5,))' assert prntr.doprint(ArrayContraction( A, [2, 3])) == 'tensorflow.linalg.einsum("abccd->abd", A)' assert prntr.doprint(I) == 'I' assert prntr.doprint(ArrayDiagonal( A, [2, 3, 4])) == 'tensorflow.linalg.einsum("abccc->abc", A)' assert prntr.doprint(ArrayDiagonal( A, [0, 1], [2, 3])) == 'tensorflow.linalg.einsum("aabbc->cab", A)' assert prntr.doprint(ArrayContraction( A, [2], [3])) == 'tensorflow.linalg.einsum("abcde->abe", A)' assert prntr.doprint(Assignment(I[i, j, k], I[i, j, k])) == 'I = I'
def test_arrayexpr_convert_array_to_matrix(): cg = ArrayContraction(ArrayTensorProduct(M), (0, 1)) assert convert_array_to_matrix(cg) == Trace(M) cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 1), (2, 3)) assert convert_array_to_matrix(cg) == Trace(M) * Trace(N) cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 3), (1, 2)) assert convert_array_to_matrix(cg) == Trace(M * N) cg = ArrayContraction(ArrayTensorProduct(M, N), (0, 2), (1, 3)) assert convert_array_to_matrix(cg) == Trace(M * N.T) cg = convert_matrix_to_array(M * N * P) assert convert_array_to_matrix(cg) == M * N * P cg = convert_matrix_to_array(M * N.T * P) assert convert_array_to_matrix(cg) == M * N.T * P cg = ArrayContraction(ArrayTensorProduct(M,N,P,Q), (1, 2), (5, 6)) assert convert_array_to_matrix(cg) == ArrayTensorProduct(M * N, P * Q) cg = ArrayContraction(ArrayTensorProduct(-2, M, N), (1, 2)) assert convert_array_to_matrix(cg) == -2 * M * N a = MatrixSymbol("a", k, 1) b = MatrixSymbol("b", k, 1) c = MatrixSymbol("c", k, 1) cg = PermuteDims( ArrayContraction( ArrayTensorProduct( a, ArrayAdd( ArrayTensorProduct(b, c), ArrayTensorProduct(c, b), ) ), (2, 4)), [0, 1, 3, 2]) assert convert_array_to_matrix(cg) == a * (b.T * c + c.T * b) za = ZeroArray(m, n) assert convert_array_to_matrix(za) == ZeroMatrix(m, n) cg = ArrayTensorProduct(3, M) assert convert_array_to_matrix(cg) == 3 * M # Partial conversion to matrix multiplication: expr = ArrayContraction(ArrayTensorProduct(M, N, P, Q), (0, 2), (1, 4, 6)) assert convert_array_to_matrix(expr) == ArrayContraction(ArrayTensorProduct(M.T*N, P, Q), (0, 2, 4)) x = MatrixSymbol("x", k, 1) cg = PermuteDims( ArrayContraction(ArrayTensorProduct(OneArray(1), x, OneArray(1), DiagMatrix(Identity(1))), (0, 5)), Permutation(1, 2, 3)) assert convert_array_to_matrix(cg) == x expr = ArrayAdd(M, PermuteDims(M, [1, 0])) assert convert_array_to_matrix(expr) == M + Transpose(M)
def _array_diag2contr_diagmatrix(expr: ArrayDiagonal): if isinstance(expr.expr, ArrayTensorProduct): args = list(expr.expr.args) diag_indices = list(expr.diagonal_indices) mapping = _get_mapping_from_subranks( [_get_subrank(arg) for arg in args]) tuple_links = [[mapping[j] for j in i] for i in diag_indices] contr_indices = [] total_rank = get_rank(expr) replaced = [False for arg in args] for i, (abs_pos, rel_pos) in enumerate(zip(diag_indices, tuple_links)): if len(abs_pos) != 2: continue (pos1_outer, pos1_inner), (pos2_outer, pos2_inner) = rel_pos arg1 = args[pos1_outer] arg2 = args[pos2_outer] if get_rank(arg1) != 2 or get_rank(arg2) != 2: if replaced[pos1_outer]: diag_indices[i] = None if replaced[pos2_outer]: diag_indices[i] = None continue pos1_in2 = 1 - pos1_inner pos2_in2 = 1 - pos2_inner if arg1.shape[pos1_in2] == 1: if arg1.shape[pos1_inner] != 1: darg1 = DiagMatrix(arg1) else: darg1 = arg1 args.append(darg1) contr_indices.append( ((pos2_outer, pos2_inner), (len(args) - 1, pos1_inner))) total_rank += 1 diag_indices[i] = None args[pos1_outer] = OneArray(arg1.shape[pos1_in2]) replaced[pos1_outer] = True elif arg2.shape[pos2_in2] == 1: if arg2.shape[pos2_inner] != 1: darg2 = DiagMatrix(arg2) else: darg2 = arg2 args.append(darg2) contr_indices.append( ((pos1_outer, pos1_inner), (len(args) - 1, pos2_inner))) total_rank += 1 diag_indices[i] = None args[pos2_outer] = OneArray(arg2.shape[pos2_in2]) replaced[pos2_outer] = True diag_indices_new = [i for i in diag_indices if i is not None] cumul = list(accumulate([0] + [get_rank(arg) for arg in args])) contr_indices2 = [ tuple(cumul[a] + b for a, b in i) for i in contr_indices ] tc = _array_contraction(_array_tensor_product(*args), *contr_indices2) td = _array_diagonal(tc, *diag_indices_new) return td return expr
def test_one_array(): assert OneArray() == 1 assert OneArray().is_Integer oa = OneArray(3, 2, 4) assert oa.shape == (3, 2, 4) oa_e = oa.as_explicit() assert oa_e.shape == (3, 2, 4) m, n, k = symbols("m n k") oa = OneArray(m, n, k, 2) assert oa.shape == (m, n, k, 2) raises(ValueError, lambda: oa.as_explicit())
def test_diag2contraction_diagmatrix(): cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, a), (1, 2)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == CodegenArrayContraction(CodegenArrayTensorProduct(M, OneArray(1), DiagMatrix(a)), (1, 3)) raises(ValueError, lambda: CodegenArrayDiagonal(CodegenArrayTensorProduct(a, M), (1, 2))) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a.T, M), (1, 2)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == CodegenArrayContraction(CodegenArrayTensorProduct(OneArray(1), M, DiagMatrix(a.T)), (1, 4)) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a.T, M, N, b.T), (1, 2), (4, 7)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == CodegenArrayContraction( CodegenArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a.T), DiagMatrix(b.T)), (1, 7), (3, 9)) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a, M, N, b.T), (0, 2), (4, 7)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == CodegenArrayContraction( CodegenArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a), DiagMatrix(b.T)), (1, 6), (3, 9)) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a, M, N, b.T), (0, 4), (3, 7)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == CodegenArrayContraction( CodegenArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a), DiagMatrix(b.T)), (3, 6), (2, 9))
def test_arrayexpr_split_multiple_contractions(): a = MatrixSymbol("a", k, 1) b = MatrixSymbol("b", k, 1) A = MatrixSymbol("A", k, k) B = MatrixSymbol("B", k, k) C = MatrixSymbol("C", k, k) X = MatrixSymbol("X", k, k) cg = _array_contraction( _array_tensor_product(A.T, a, b, b.T, (A * X * b).applyfunc(cos)), (1, 2, 8), (5, 6, 9)) expected = _array_contraction( _array_tensor_product(A.T, DiagMatrix(a), OneArray(1), b, b.T, (A * X * b).applyfunc(cos)), (1, 3), (2, 9), (6, 7, 10)) assert cg.split_multiple_contractions().dummy_eq(expected) # Check no overlap of lines: cg = _array_contraction(_array_tensor_product(A, a, C, a, B), (1, 2, 4), (5, 6, 8), (3, 7)) assert cg.split_multiple_contractions() == cg cg = _array_contraction(_array_tensor_product(a, b, A), (0, 2, 4), (1, 3)) assert cg.split_multiple_contractions() == cg
def test_diag2contraction_diagmatrix(): cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, a), (1, 2)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == CodegenArrayContraction( CodegenArrayTensorProduct(M, OneArray(1), DiagMatrix(a)), (1, 3)) raises( ValueError, lambda: CodegenArrayDiagonal(CodegenArrayTensorProduct(a, M), (1, 2))) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a.T, M), (1, 2)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == CodegenArrayContraction( CodegenArrayTensorProduct(OneArray(1), M, DiagMatrix(a.T)), (1, 4)) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a.T, M, N, b.T), (1, 2), (4, 7)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == CodegenArrayContraction( CodegenArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a.T), DiagMatrix(b.T)), (1, 7), (3, 9)) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a, M, N, b.T), (0, 2), (4, 7)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == CodegenArrayContraction( CodegenArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a), DiagMatrix(b.T)), (1, 6), (3, 9)) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a, M, N, b.T), (0, 4), (3, 7)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == CodegenArrayContraction( CodegenArrayTensorProduct(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 = CodegenArrayDiagonal(CodegenArrayTensorProduct(x, A.T, I1), (0, 2)) assert _array_diag2contr_diagmatrix(cg).shape == cg.shape assert array2matrix(cg).shape == cg.shape
def test_arrayexpr_convert_indexed_to_array_broadcast(): A = ArraySymbol("A", (3, 3)) B = ArraySymbol("B", (3, 3)) expr = A[i, j] + B[k, l] O2 = OneArray(3, 3) expected = ArrayAdd(ArrayTensorProduct(A, O2), ArrayTensorProduct(O2, B)) assert convert_indexed_to_array(expr) == expected assert convert_indexed_to_array(expr, [i, j, k, l]) == expected assert convert_indexed_to_array(expr, [l, k, i, j]) == ArrayAdd( PermuteDims(ArrayTensorProduct(O2, A), [1, 0, 2, 3]), PermuteDims(ArrayTensorProduct(B, O2), [1, 0, 2, 3])) expr = A[i, j] + B[j, k] O1 = OneArray(3) assert convert_indexed_to_array(expr, [i, j, k]) == ArrayAdd( ArrayTensorProduct(A, O1), ArrayTensorProduct(O1, B)) C = ArraySymbol("C", (d0, d1)) D = ArraySymbol("D", (d3, d1)) expr = C[i, j] + D[k, j] assert convert_indexed_to_array(expr, [i, j, k]) == ArrayAdd( ArrayTensorProduct(C, OneArray(d3)), PermuteDims(ArrayTensorProduct(OneArray(d0), D), [0, 2, 1])) X = ArraySymbol("X", (5, 3)) expr = X[i, n] - X[j, n] assert convert_indexed_to_array(expr, [i, j, n]) == ArrayAdd( ArrayTensorProduct(-1, OneArray(5), X), PermuteDims(ArrayTensorProduct(X, OneArray(5)), [0, 2, 1])) raises(ValueError, lambda: convert_indexed_to_array(C[i, j] + D[i, j]))
def test_array_symbol_and_element(): A = ArraySymbol("A", (2,)) A0 = ArrayElement(A, (0,)) A1 = ArrayElement(A, (1,)) assert A[0] == A0 assert A[1] != A0 assert A.as_explicit() == ImmutableDenseNDimArray([A0, A1]) A2 = tensorproduct(A, A) assert A2.shape == (2, 2) # TODO: not yet supported: # assert A2.as_explicit() == Array([[A[0]*A[0], A[1]*A[0]], [A[0]*A[1], A[1]*A[1]]]) A3 = tensorcontraction(A2, (0, 1)) assert A3.shape == () # TODO: not yet supported: # assert A3.as_explicit() == Array([]) A = ArraySymbol("A", (2, 3, 4)) Ae = A.as_explicit() assert Ae == ImmutableDenseNDimArray( [[[ArrayElement(A, (i, j, k)) for k in range(4)] for j in range(3)] for i in range(2)]) p = _permute_dims(A, Permutation(0, 2, 1)) assert isinstance(p, PermuteDims) A = ArraySymbol("A", (2,)) raises(IndexError, lambda: A[()]) raises(IndexError, lambda: A[0, 1]) raises(ValueError, lambda: A[-1]) raises(ValueError, lambda: A[2]) O = OneArray(3, 4) Z = ZeroArray(m, n) raises(IndexError, lambda: O[()]) raises(IndexError, lambda: O[1, 2, 3]) raises(ValueError, lambda: O[3, 0]) raises(ValueError, lambda: O[0, 4]) assert O[1, 2] == 1 assert Z[1, 2] == 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_recognize_diagonalized_vectors(): a = MatrixSymbol("a", k, 1) b = MatrixSymbol("b", k, 1) A = MatrixSymbol("A", k, k) B = MatrixSymbol("B", k, k) C = MatrixSymbol("C", k, k) X = MatrixSymbol("X", k, k) x = MatrixSymbol("x", k, 1) I1 = Identity(1) I = Identity(k) # Check matrix recognition over trivial dimensions: cg = CodegenArrayTensorProduct(a, b) assert recognize_matrix_expression(cg) == a * b.T cg = CodegenArrayTensorProduct(I1, a, b) assert recognize_matrix_expression(cg) == a * b.T # Recognize trace inside a tensor product: cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, B, C), (0, 3), (1, 2)) assert recognize_matrix_expression(cg) == Trace(A * B) * C # Transform diagonal operator to contraction: cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(A, a), (1, 2)) assert _array_diag2contr_diagmatrix(cg) == CodegenArrayContraction( CodegenArrayTensorProduct(A, OneArray(1), DiagMatrix(a)), (1, 3)) assert recognize_matrix_expression(cg) == A * DiagMatrix(a) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a, b), (0, 2)) assert _array_diag2contr_diagmatrix(cg) == CodegenArrayPermuteDims( CodegenArrayContraction( CodegenArrayTensorProduct(DiagMatrix(a), OneArray(1), b), (0, 3)), [1, 2, 0]) assert recognize_matrix_expression(cg) == b.T * DiagMatrix(a) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(A, a), (0, 2)) assert _array_diag2contr_diagmatrix(cg) == CodegenArrayContraction( CodegenArrayTensorProduct(A, OneArray(1), DiagMatrix(a)), (0, 3)) assert recognize_matrix_expression(cg) == A.T * DiagMatrix(a) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(I, x, I1), (0, 2), (3, 5)) assert _array_diag2contr_diagmatrix(cg) == CodegenArrayContraction( CodegenArrayTensorProduct(I, OneArray(1), I1, DiagMatrix(x)), (0, 5)) assert recognize_matrix_expression(cg) == DiagMatrix(x) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(I, x, A, B), (1, 2), (5, 6)) assert _array_diag2contr_diagmatrix(cg) == CodegenArrayDiagonal( CodegenArrayContraction( CodegenArrayTensorProduct(I, OneArray(1), A, B, DiagMatrix(x)), (1, 7)), (5, 6)) # TODO: not yet working # assert recognize_matrix_expression(cg) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(x, I1), (1, 2)) assert isinstance(cg, CodegenArrayDiagonal) assert cg.diagonal_indices == ((1, 2), ) assert recognize_matrix_expression(cg) == x cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(x, I), (0, 2)) assert _array_diag2contr_diagmatrix(cg) == CodegenArrayContraction( CodegenArrayTensorProduct(OneArray(1), I, DiagMatrix(x)), (1, 3)) assert recognize_matrix_expression(cg).doit() == DiagMatrix(x) raises(ValueError, lambda: CodegenArrayDiagonal(x, (1, ))) # Ignore identity matrices with contractions: cg = CodegenArrayContraction(CodegenArrayTensorProduct(I, A, I, I), (0, 2), (1, 3), (5, 7)) assert cg.split_multiple_contractions() == cg assert recognize_matrix_expression(cg) == Trace(A) * I cg = CodegenArrayContraction(CodegenArrayTensorProduct(Trace(A) * I, I, I), (1, 5), (3, 4)) assert cg.split_multiple_contractions() == cg assert recognize_matrix_expression(cg).doit() == Trace(A) * I # Add DiagMatrix when required: cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a), (1, 2)) assert cg.split_multiple_contractions() == cg assert recognize_matrix_expression(cg) == A * a cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, B), (1, 2, 4)) assert cg.split_multiple_contractions() == CodegenArrayContraction( CodegenArrayTensorProduct(A, DiagMatrix(a), B), (1, 2), (3, 4)) assert recognize_matrix_expression(cg) == A * DiagMatrix(a) * B cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, B), (0, 2, 4)) assert cg.split_multiple_contractions() == CodegenArrayContraction( CodegenArrayTensorProduct(A, DiagMatrix(a), B), (0, 2), (3, 4)) assert recognize_matrix_expression(cg) == A.T * DiagMatrix(a) * B cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, b, a.T, B), (0, 2, 4, 7, 9)) assert cg.split_multiple_contractions() == CodegenArrayContraction( CodegenArrayTensorProduct(A, DiagMatrix(a), DiagMatrix(b), DiagMatrix(a), B), (0, 2), (3, 4), (5, 7), (6, 9)) assert recognize_matrix_expression( cg).doit() == A.T * DiagMatrix(a) * DiagMatrix(b) * DiagMatrix(a) * B.T cg = CodegenArrayContraction(CodegenArrayTensorProduct(I1, I1, I1), (1, 2, 4)) assert cg.split_multiple_contractions() == CodegenArrayContraction( CodegenArrayTensorProduct(I1, I1, I1), (1, 2), (3, 4)) assert recognize_matrix_expression(cg).doit() == Identity(1) cg = CodegenArrayContraction(CodegenArrayTensorProduct(I, I, I, I, A), (1, 2, 8), (5, 6, 9)) assert recognize_matrix_expression( cg.split_multiple_contractions()).doit() == A cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, C, a, B), (1, 2, 4), (5, 6, 8)) expected = CodegenArrayContraction( CodegenArrayTensorProduct(DiagMatrix(a), DiagMatrix(a), C, A, B), (0, 4), (1, 7), (2, 5), (3, 8)) assert cg.split_multiple_contractions() == expected assert recognize_matrix_expression( cg) == A * DiagMatrix(a) * C * DiagMatrix(a) * B cg = CodegenArrayContraction( CodegenArrayTensorProduct(a, I1, b, I1, (a.T * b).applyfunc(cos)), (1, 2, 8), (5, 6, 9)) assert cg.split_multiple_contractions().dummy_eq( CodegenArrayContraction( CodegenArrayTensorProduct((a.T * b).applyfunc(cos), I1, I1, a, b), (0, 2), (1, 4), (3, 7), (5, 9))) assert recognize_matrix_expression(cg).doit().dummy_eq( MatMul(a, (a.T * b).applyfunc(cos), b.T)) cg = CodegenArrayContraction( CodegenArrayTensorProduct(A.T, a, b, b.T, (A * X * b).applyfunc(cos)), (1, 2, 8), (5, 6, 9)) assert cg.split_multiple_contractions().dummy_eq( CodegenArrayContraction( CodegenArrayTensorProduct(DiagMatrix(a), (A * X * b).applyfunc(cos), A.T, b, b.T), (0, 2), (1, 5), (3, 7, 8))) # assert recognize_matrix_expression(cg) # Check no overlap of lines: cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, C, a, B), (1, 2, 4), (5, 6, 8), (3, 7)) assert cg.split_multiple_contractions() == cg cg = CodegenArrayContraction(CodegenArrayTensorProduct(a, b, A), (0, 2, 4), (1, 3)) assert cg.split_multiple_contractions() == cg
def _convert_indexed_to_array(expr): if isinstance(expr, Sum): function = expr.function summation_indices = expr.variables subexpr, subindices = _convert_indexed_to_array(function) subindicessets = { j: i for i in subindices if isinstance(i, frozenset) for j in i } summation_indices = sorted(set( [subindicessets.get(i, i) for i in summation_indices]), key=default_sort_key) # TODO: check that Kronecker delta is only contracted to one other element: kronecker_indices = set([]) if isinstance(function, Mul): for arg in function.args: if not isinstance(arg, KroneckerDelta): continue arg_indices = sorted(set(arg.indices), key=default_sort_key) if len(arg_indices) == 2: kronecker_indices.update(arg_indices) kronecker_indices = sorted(kronecker_indices, key=default_sort_key) # Check dimensional consistency: shape = get_shape(subexpr) if shape: for ind, istart, iend in expr.limits: i = _get_argindex(subindices, ind) if istart != 0 or iend + 1 != shape[i]: raise ValueError( "summation index and array dimension mismatch: %s" % ind) contraction_indices = [] subindices = list(subindices) if isinstance(subexpr, ArrayDiagonal): diagonal_indices = list(subexpr.diagonal_indices) dindices = subindices[-len(diagonal_indices):] subindices = subindices[:-len(diagonal_indices)] for index in summation_indices: if index in dindices: position = dindices.index(index) contraction_indices.append(diagonal_indices[position]) diagonal_indices[position] = None diagonal_indices = [i for i in diagonal_indices if i is not None] for i, ind in enumerate(subindices): if ind in summation_indices: pass if diagonal_indices: subexpr = _array_diagonal(subexpr.expr, *diagonal_indices) else: subexpr = subexpr.expr axes_contraction = defaultdict(list) for i, ind in enumerate(subindices): include = all(j not in kronecker_indices for j in ind) if isinstance( ind, frozenset) else ind not in kronecker_indices if ind in summation_indices and include: axes_contraction[ind].append(i) subindices[i] = None for k, v in axes_contraction.items(): if any(i in kronecker_indices for i in k) if isinstance( k, frozenset) else k in kronecker_indices: continue contraction_indices.append(tuple(v)) free_indices = [i for i in subindices if i is not None] indices_ret = list(free_indices) indices_ret.sort(key=lambda x: free_indices.index(x)) return _array_contraction( subexpr, *contraction_indices, free_indices=free_indices), tuple(indices_ret) if isinstance(expr, Mul): args, indices = zip( *[_convert_indexed_to_array(arg) for arg in expr.args]) # Check if there are KroneckerDelta objects: kronecker_delta_repl = {} for arg in args: if not isinstance(arg, KroneckerDelta): continue # Diagonalize two indices: i, j = arg.indices kindices = set(arg.indices) if i in kronecker_delta_repl: kindices.update(kronecker_delta_repl[i]) if j in kronecker_delta_repl: kindices.update(kronecker_delta_repl[j]) kindices = frozenset(kindices) for index in kindices: kronecker_delta_repl[index] = kindices # Remove KroneckerDelta objects, their relations should be handled by # ArrayDiagonal: newargs = [] newindices = [] for arg, loc_indices in zip(args, indices): if isinstance(arg, KroneckerDelta): continue newargs.append(arg) newindices.append(loc_indices) flattened_indices = [ kronecker_delta_repl.get(j, j) for i in newindices for j in i ] diagonal_indices, ret_indices = _get_diagonal_indices( flattened_indices) tp = _array_tensor_product(*newargs) if diagonal_indices: return _array_diagonal(tp, *diagonal_indices), ret_indices else: return tp, ret_indices if isinstance(expr, MatrixElement): indices = expr.args[1:] diagonal_indices, ret_indices = _get_diagonal_indices(indices) if diagonal_indices: return _array_diagonal(expr.args[0], *diagonal_indices), ret_indices else: return expr.args[0], ret_indices if isinstance(expr, ArrayElement): indices = expr.indices diagonal_indices, ret_indices = _get_diagonal_indices(indices) if diagonal_indices: return _array_diagonal(expr.name, *diagonal_indices), ret_indices else: return expr.name, ret_indices if isinstance(expr, Indexed): indices = expr.indices diagonal_indices, ret_indices = _get_diagonal_indices(indices) if diagonal_indices: return _array_diagonal(expr.base, *diagonal_indices), ret_indices else: return expr.args[0], ret_indices if isinstance(expr, IndexedBase): raise NotImplementedError if isinstance(expr, KroneckerDelta): return expr, expr.indices if isinstance(expr, Add): args, indices = zip( *[_convert_indexed_to_array(arg) for arg in expr.args]) args = list(args) # Check if all indices are compatible. Otherwise expand the dimensions: index0 = [] shape0 = [] for arg, arg_indices in zip(args, indices): arg_indices_set = set(arg_indices) arg_indices_missing = arg_indices_set.difference(index0) index0.extend([i for i in arg_indices if i in arg_indices_missing]) arg_shape = get_shape(arg) shape0.extend([ arg_shape[i] for i, e in enumerate(arg_indices) if e in arg_indices_missing ]) for i, (arg, arg_indices) in enumerate(zip(args, indices)): if len(arg_indices) < len(index0): missing_indices_pos = [ i for i, e in enumerate(index0) if e not in arg_indices ] missing_shape = [shape0[i] for i in missing_indices_pos] arg_indices = tuple(index0[j] for j in missing_indices_pos) + arg_indices args[i] = _array_tensor_product(OneArray(*missing_shape), args[i]) permutation = Permutation([arg_indices.index(j) for j in index0]) # Perform index permutations: args[i] = _permute_dims(args[i], permutation) return _array_add(*args), tuple(index0) if isinstance(expr, Pow): subexpr, subindices = _convert_indexed_to_array(expr.base) if isinstance(expr.exp, (int, Integer)): diags = zip(*[(2 * i, 2 * i + 1) for i in range(expr.exp)]) arr = _array_diagonal( _array_tensor_product(*[subexpr for i in range(expr.exp)]), *diags) return arr, subindices if isinstance(expr, Function): subexpr, subindices = _convert_indexed_to_array(expr.args[0]) return ArrayElementwiseApplyFunc(type(expr), subexpr), subindices return expr, ()
def identify_removable_identity_matrices(expr): editor = _EditArrayContraction(expr) flag: bool = True while flag: flag = False for arg_with_ind in editor.args_with_ind: if isinstance(arg_with_ind.element, Identity): k = arg_with_ind.element.shape[0] # Candidate for removal: if arg_with_ind.indices == [None, None]: # Free identity matrix, will be cleared by _remove_trivial_dims: continue elif None in arg_with_ind.indices: ind = [j for j in arg_with_ind.indices if j is not None][0] counted = editor.count_args_with_index(ind) if counted == 1: # Identity matrix contracted only on one index with itself, # transform to a OneArray(k) element: editor.insert_after(arg_with_ind, OneArray(k)) editor.args_with_ind.remove(arg_with_ind) flag = True break elif counted > 2: # Case counted = 2 is a matrix multiplication by identity matrix, skip it. # Case counted > 2 is a multiple contraction, # this is a case where the contraction becomes a diagonalization if the # identity matrix is dropped. continue elif arg_with_ind.indices[0] == arg_with_ind.indices[1]: ind = arg_with_ind.indices[0] counted = editor.count_args_with_index(ind) if counted > 1: editor.args_with_ind.remove(arg_with_ind) flag = True break else: # This is a trace, skip it as it will be recognized somewhere else: pass elif ask(Q.diagonal(arg_with_ind.element)): if arg_with_ind.indices == [None, None]: continue elif None in arg_with_ind.indices: pass elif arg_with_ind.indices[0] == arg_with_ind.indices[1]: ind = arg_with_ind.indices[0] counted = editor.count_args_with_index(ind) if counted == 3: # A_ai B_bi D_ii ==> A_ai D_ij B_bj ind_new = editor.get_new_contraction_index() other_args = [ j for j in editor.args_with_ind if j != arg_with_ind ] other_args[1].indices = [ ind_new if j == ind else j for j in other_args[1].indices ] arg_with_ind.indices = [ind, ind_new] flag = True break return editor.to_array_contraction()
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)