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 = ArrayContraction(
ArrayTensorProduct(A.T, a, b, b.T, (A * X * b).applyfunc(cos)),
(1, 2, 8), (5, 6, 9))
assert cg.split_multiple_contractions().dummy_eq(
ArrayContraction(
ArrayTensorProduct(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 = ArrayContraction(ArrayTensorProduct(A, a, C, a, B), (1, 2, 4),
(5, 6, 8), (3, 7))
assert cg.split_multiple_contractions() == cg
cg = ArrayContraction(ArrayTensorProduct(a, b, A), (0, 2, 4), (1, 3))
assert cg.split_multiple_contractions() == cg
def _(expr: ArrayContraction):
expr = expr.flatten_contraction_of_diagonal()
expr = expr.split_multiple_contractions()
expr = identify_hadamard_products(expr)
subexpr = expr.expr
contraction_indices: Tuple[Tuple[int]] = expr.contraction_indices
if isinstance(subexpr, ArrayTensorProduct):
newexpr = ArrayContraction(_array2matrix(subexpr),
*contraction_indices)
contraction_indices = newexpr.contraction_indices
if any(i > 2 for i in newexpr.subranks):
addends = ArrayAdd(*[
_a2m_tensor_product(*j) for j in itertools.product(*[
i.args if isinstance(i, ArrayAdd) else [i]
for i in expr.expr.args
])
])
newexpr = ArrayContraction(addends, *contraction_indices)
if isinstance(newexpr, ArrayAdd):
ret = _array2matrix(newexpr)
return ret
assert isinstance(newexpr, ArrayContraction)
ret = _support_function_tp1_recognize(contraction_indices,
list(newexpr.expr.args))
return ret
elif not isinstance(subexpr, _CodegenArrayAbstract):
ret = _array2matrix(subexpr)
if isinstance(ret, MatrixExpr):
assert expr.contraction_indices == ((0, 1), )
return _a2m_trace(ret)
else:
return ArrayContraction(ret, *expr.contraction_indices)
def test_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 = ArrayContraction(ArrayTensorProduct(A.T, a, b, b.T, (A*X*b).applyfunc(cos)), (1, 2, 8), (5, 6, 9))
expected = ArrayContraction(ArrayTensorProduct(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 = ArrayContraction(ArrayTensorProduct(A, a, C, a, B), (1, 2, 4), (5, 6, 8), (3, 7))
assert cg.split_multiple_contractions() == cg
cg = ArrayContraction(ArrayTensorProduct(a, b, A), (0, 2, 4), (1, 3))
assert cg.split_multiple_contractions() == cg