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