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_MatrixNormal():
M = MatrixNormal('M', [[5, 6]], [4], [[2, 1], [1, 2]])
assert M.pspace.distribution.set == MatrixSet(1, 2, S.Reals)
X = MatrixSymbol('X', 1, 2)
term1 = exp(-Trace(Matrix([[ S(2)/3, -S(1)/3], [-S(1)/3, S(2)/3]])*(
Matrix([[-5], [-6]]) + X.T)*Matrix([[1/4]])*(Matrix([[-5, -6]]) + X))/2)
assert density(M)(X).doit() == term1/(24*pi)
assert density(M)([[7, 8]]).doit() == exp(-S(1)/3)/(24*pi)
d, n = symbols('d n', positive=True, integer=True)
SM2 = MatrixSymbol('SM2', d, d)
SM1 = MatrixSymbol('SM1', n, n)
LM = MatrixSymbol('LM', n, d)
Y = MatrixSymbol('Y', n, d)
M = MatrixNormal('M', LM, SM1, SM2)
exprd = 4*(2*pi)**(-d*n/2)*exp(-Trace(SM2**(-1)*(-LM.T + Y.T)*SM1**(-1)*(-LM + Y)
)/2)*Determinant(SM1)**(-d)*Determinant(SM2)**(-n)
assert density(M)(Y).doit() == exprd
raises(ValueError, lambda: density(M)(1))
raises(ValueError, lambda: MatrixNormal('M', [1, 2], [[1, 0], [0, 1]], [[1, 0], [2, 1]]))
raises(ValueError, lambda: MatrixNormal('M', [1, 2], [[1, 0], [2, 1]], [[1, 0], [0, 1]]))
raises(ValueError, lambda: MatrixNormal('M', [1, 2], [[1, 0], [0, 1]], [[1, 0], [0, 1]]))
raises(ValueError, lambda: MatrixNormal('M', [1, 2], [[1, 0], [2]], [[1, 0], [0, 1]]))
raises(ValueError, lambda: MatrixNormal('M', [1, 2], [[1, 0], [2, 1]], [[1, 0], [0]]))
raises(ValueError, lambda: MatrixNormal('M', [[1, 2]], [[1, 0], [0, 1]], [[1, 0]]))
raises(ValueError, lambda: MatrixNormal('M', [[1, 2]], [1], [[1, 0]]))
def test_codegen_array_recognize_matrix_mul_lines():
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M), (0, 1))
assert recognize_matrix_expression(cg) == Trace(M)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 1),
(2, 3))
assert recognize_matrix_expression(cg) == Trace(M) * Trace(N)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 3),
(1, 2))
assert recognize_matrix_expression(cg) == Trace(M * N)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 2),
(1, 3))
assert recognize_matrix_expression(cg) == Trace(M * N.T)
cg = parse_indexed_expression((M * N * P)[i, j])
assert recognize_matrix_expression(cg) == M * N * P
cg = parse_matrix_expression(M * N * P)
assert recognize_matrix_expression(cg) == M * N * P
cg = parse_indexed_expression((M * N.T * P)[i, j])
assert recognize_matrix_expression(cg) == M * N.T * P
cg = parse_matrix_expression(M * N.T * P)
assert recognize_matrix_expression(cg) == M * N.T * P
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (1, 2),
(5, 6))
assert recognize_matrix_expression(cg) == [M * N, P * Q]
expr = -2 * M * N
elem = expr[i, j]
cg = parse_indexed_expression(elem)
assert recognize_matrix_expression(cg) == -2 * M * N
def test_MatrixGamma():
M = MatrixGamma('M', 1, 2, [[1, 0], [0, 1]])
assert M.pspace.distribution.set == MatrixSet(2, 2, S.Reals)
assert isinstance(density(M), MatrixGammaDistribution)
X = MatrixSymbol('X', 2, 2)
num = exp(Trace(Matrix([[-S(1) / 2, 0], [0, -S(1) / 2]]) * X))
assert density(M)(X).doit() == num / (4 * pi * sqrt(Determinant(X)))
assert density(M)([[2, 1], [1, 2]]).doit() == sqrt(3) * exp(-2) / (12 * pi)
X = MatrixSymbol('X', 1, 2)
Y = MatrixSymbol('Y', 1, 2)
assert density(M)([X, Y]).doit() == exp(-X[0, 0] / 2 - Y[0, 1] / 2) / (
4 * pi * sqrt(X[0, 0] * Y[0, 1] - X[0, 1] * Y[0, 0]))
# symbolic
a, b = symbols('a b', positive=True)
d = symbols('d', positive=True, integer=True)
Y = MatrixSymbol('Y', d, d)
Z = MatrixSymbol('Z', 2, 2)
SM = MatrixSymbol('SM', d, d)
M2 = MatrixGamma('M2', a, b, SM)
M3 = MatrixGamma('M3', 2, 3, [[2, 1], [1, 2]])
k = Dummy('k')
exprd = pi**(-d * (d - 1) / 4) * b**(-a * d) * exp(
Trace((-1 / b) * SM**(-1) * Y)) * Determinant(SM)**(-a) * Determinant(
Y)**(a - d / 2 - S(1) / 2) / Product(gamma(-k / 2 + a + S(1) / 2),
(k, 1, d))
assert density(M2)(Y).dummy_eq(exprd)
raises(NotImplementedError, lambda: density(M3 + M)(Z))
raises(ValueError, lambda: density(M)(1))
raises(ValueError, lambda: MatrixGamma('M', -1, 2, [[1, 0], [0, 1]]))
raises(ValueError, lambda: MatrixGamma('M', -1, -2, [[1, 0], [0, 1]]))
raises(ValueError, lambda: MatrixGamma('M', -1, 2, [[1, 0], [2, 1]]))
raises(ValueError, lambda: MatrixGamma('M', -1, 2, [[1, 0], [0]]))
def test_contraction_permutation_mix():
Me = M.subs(k, 3).as_explicit()
Ne = N.subs(k, 3).as_explicit()
cg1 = CodegenArrayContraction(
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N),
Permutation([0, 2, 1, 3])), (2, 3))
cg2 = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 3))
assert cg1 == cg2
assert recognize_matrix_expression(cg2) == M * N.T
cge1 = tensorcontraction(
permutedims(tensorproduct(Me, Ne), Permutation([0, 2, 1, 3])), (2, 3))
cge2 = tensorcontraction(tensorproduct(Me, Ne), (1, 3))
assert cge1 == cge2
cg1 = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N),
Permutation([0, 1, 3, 2]))
cg2 = CodegenArrayTensorProduct(
M, CodegenArrayPermuteDims(N, Permutation([1, 0])))
assert cg1 == cg2
assert recognize_matrix_expression(cg1) == CodegenArrayTensorProduct(
M, N.T)
assert recognize_matrix_expression(cg2) == CodegenArrayTensorProduct(
M, N.T)
cg1 = CodegenArrayContraction(
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q),
Permutation([0, 2, 3, 1, 4, 5, 7, 6])), (1, 2),
(3, 5))
cg2 = CodegenArrayContraction(
CodegenArrayTensorProduct(
M, N, P, CodegenArrayPermuteDims(Q, Permutation([1, 0]))), (1, 5),
(2, 3))
assert cg1 == cg2
assert recognize_matrix_expression(cg1) == CodegenArrayTensorProduct(
M * P.T * Trace(N), Q.T)
assert recognize_matrix_expression(cg2) == CodegenArrayTensorProduct(
M * P.T * Trace(N), Q.T)
cg1 = CodegenArrayContraction(
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q),
Permutation([1, 0, 4, 6, 2, 7, 5, 3])), (0, 1),
(2, 6), (3, 7))
cg2 = CodegenArrayPermuteDims(
CodegenArrayContraction(CodegenArrayTensorProduct(M, P, Q, N), (0, 1),
(2, 3), (4, 7)), [1, 0])
assert cg1 == cg2
cg1 = CodegenArrayContraction(
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q),
Permutation([1, 0, 4, 6, 7, 2, 5, 3])), (0, 1),
(2, 6), (3, 7))
cg2 = CodegenArrayPermuteDims(
CodegenArrayContraction(
CodegenArrayTensorProduct(CodegenArrayPermuteDims(M, [1, 0]), N, P,
Q), (0, 1), (3, 6), (4, 5)),
Permutation([1, 0]))
assert cg1 == cg2
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 test_codegen_array_recognize_matrix_mul_lines():
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M), (0, 1))
assert recognize_matrix_expression(cg) == Trace(M)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 1),
(2, 3))
assert recognize_matrix_expression(cg) == Trace(M) * Trace(N)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 3),
(1, 2))
assert recognize_matrix_expression(cg) == Trace(M * N)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 2),
(1, 3))
assert recognize_matrix_expression(cg) == Trace(M * N.T)
cg = parse_indexed_expression((M * N * P)[i, j])
assert recognize_matrix_expression(cg) == M * N * P
cg = parse_matrix_expression(M * N * P)
assert recognize_matrix_expression(cg) == M * N * P
cg = parse_indexed_expression((M * N.T * P)[i, j])
assert recognize_matrix_expression(cg) == M * N.T * P
cg = parse_matrix_expression(M * N.T * P)
assert recognize_matrix_expression(cg) == M * N.T * P
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (1, 2),
(5, 6))
assert recognize_matrix_expression(cg) == CodegenArrayTensorProduct(
M * N, P * Q)
expr = -2 * M * N
elem = expr[i, j]
cg = parse_indexed_expression(elem)
assert recognize_matrix_expression(cg) == -2 * M * N
a = MatrixSymbol("a", k, 1)
b = MatrixSymbol("b", k, 1)
c = MatrixSymbol("c", k, 1)
cg = CodegenArrayPermuteDims(
CodegenArrayContraction(
CodegenArrayTensorProduct(
a,
CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(b, c),
CodegenArrayTensorProduct(c, b),
)), (2, 4)), [0, 1, 3, 2])
assert recognize_matrix_expression(cg) == a * (b.T * c + c.T * b)
za = ZeroArray(m, n)
assert recognize_matrix_expression(za) == ZeroMatrix(m, n)
cg = CodegenArrayTensorProduct(3, M)
assert recognize_matrix_expression(cg) == 3 * M
def test_transpose():
Sq = MatrixSymbol('Sq', n, n)
assert Transpose(A).shape == (m, n)
assert Transpose(A * B).shape == (l, n)
assert Transpose(Transpose(A)) == A
assert Transpose(eye(3)) == eye(3)
assert Transpose(S(5)) == S(5)
assert Transpose(Matrix([[1, 2], [3, 4]])) == Matrix([[1, 3], [2, 4]])
assert Transpose(Trace(Sq)) == Trace(Sq)
def test_transpose():
n, m, l = symbols('n m l', integer=True)
A = MatrixSymbol('A', n, m)
B = MatrixSymbol('B', m, l)
Sq = MatrixSymbol('Sq', n, n)
assert Transpose(A).shape == (m, n)
assert Transpose(A * B).shape == (l, n)
assert Transpose(Transpose(A)) == A
assert Transpose(eye(3)) == eye(3)
assert Transpose(S(5)) == S(5)
assert Transpose(Matrix([[1, 2], [3, 4]])) == Matrix([[1, 3], [2, 4]])
assert Transpose(Trace(Sq)) == Trace(Sq)
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, A.T * D.T), [0, 2, 4, 1, 3, 5]
)
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_matrix2():
cg = _array_contraction(_array_tensor_product(M, N), (1, 3))
assert convert_array_to_matrix(cg) == M * N.T
cg = PermuteDims(_array_tensor_product(M, N), Permutation([0, 1, 3, 2]))
assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T)
cg = _array_tensor_product(M, PermuteDims(N, Permutation([1, 0])))
assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T)
cg = _array_contraction(
PermuteDims(
_array_tensor_product(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])),
(1, 2), (3, 5)
)
assert convert_array_to_matrix(cg) == _array_tensor_product(M * P.T * Trace(N), Q.T)
cg = _array_contraction(
_array_tensor_product(M, N, P, PermuteDims(Q, Permutation([1, 0]))),
(1, 5), (2, 3)
)
assert convert_array_to_matrix(cg) == _array_tensor_product(M * P.T * Trace(N), Q.T)
cg = _array_tensor_product(M, PermuteDims(N, [1, 0]))
assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T)
cg = _array_tensor_product(PermuteDims(M, [1, 0]), PermuteDims(N, [1, 0]))
assert convert_array_to_matrix(cg) == _array_tensor_product(M.T, N.T)
cg = _array_tensor_product(PermuteDims(N, [1, 0]), PermuteDims(M, [1, 0]))
assert convert_array_to_matrix(cg) == _array_tensor_product(N.T, M.T)
cg = _array_contraction(M, (0,), (1,))
assert convert_array_to_matrix(cg) == OneMatrix(1, k)*M*OneMatrix(k, 1)
cg = _array_contraction(x, (0,), (1,))
assert convert_array_to_matrix(cg) == OneMatrix(1, k)*x
Xm = MatrixSymbol("Xm", m, n)
cg = _array_contraction(Xm, (0,), (1,))
assert convert_array_to_matrix(cg) == OneMatrix(1, m)*Xm*OneMatrix(n, 1)
def pdf(self, x):
M , U , V = self.location_matrix, self.scale_matrix_1, self.scale_matrix_2
n, p = M.shape
if isinstance(x, list):
x = ImmutableMatrix(x)
if not isinstance(x, (MatrixBase, MatrixSymbol)):
raise ValueError("%s should be an isinstance of Matrix "
"or MatrixSymbol" % str(x))
term1 = Inverse(V)*Transpose(x - M)*Inverse(U)*(x - M)
num = exp(-Trace(term1)/S(2))
den = (2*pi)**(S(n*p)/2) * Determinant(U)**S(p)/2 * Determinant(V)**S(n)/2
return num/den
def test_Wishart():
W = Wishart('W', 5, [[1, 0], [0, 1]])
assert W.pspace.distribution.set == MatrixSet(2, 2, S.Reals)
X = MatrixSymbol('X', 2, 2)
term1 = exp(Trace(Matrix([[-S(1)/2, 0], [0, -S(1)/2]])*X))
assert density(W)(X).doit() == term1 * Determinant(X)/(24*pi)
assert density(W)([[2, 1], [1, 2]]).doit() == exp(-2)/(8*pi)
n = symbols('n', positive=True)
d = symbols('d', positive=True, integer=True)
Y = MatrixSymbol('Y', d, d)
SM = MatrixSymbol('SM', d, d)
W = Wishart('W', n, SM)
k = Dummy('k')
exprd = 2**(-d*n/2)*pi**(-d*(d - 1)/4)*exp(Trace(-(S(1)/2)*SM**(-1)*Y)
)*Determinant(SM)**(-n/2)*Determinant(Y)**(
-d/2 + n/2 - S(1)/2)/Product(gamma(-k/2 + n/2 + S(1)/2), (k, 1, d))
assert density(W)(Y).dummy_eq(exprd)
raises(ValueError, lambda: density(W)(1))
raises(ValueError, lambda: Wishart('W', -1, [[1, 0], [0, 1]]))
raises(ValueError, lambda: Wishart('W', -1, [[1, 0], [2, 1]]))
raises(ValueError, lambda: Wishart('W', 2, [[1, 0], [0]]))
def pdf(self, x):
alpha , beta , scale_matrix = self.alpha, self.beta, self.scale_matrix
p = scale_matrix.shape[0]
if isinstance(x, list):
x = ImmutableMatrix(x)
if not isinstance(x, (MatrixBase, MatrixSymbol)):
raise ValueError("%s should be an isinstance of Matrix "
"or MatrixSymbol" % str(x))
sigma_inv_x = - Inverse(scale_matrix)*x / beta
term1 = exp(Trace(sigma_inv_x))/((beta**(p*alpha)) * multigamma(alpha, p))
term2 = (Determinant(scale_matrix))**(-alpha)
term3 = (Determinant(x))**(alpha - S(p + 1)/2)
return term1 * term2 * term3
def pdf(self, x):
n, scale_matrix = self.n, self.scale_matrix
p = scale_matrix.shape[0]
if isinstance(x, list):
x = ImmutableMatrix(x)
if not isinstance(x, (MatrixBase, MatrixSymbol)):
raise ValueError("%s should be an isinstance of Matrix "
"or MatrixSymbol" % str(x))
sigma_inv_x = - Inverse(scale_matrix)*x / S(2)
term1 = exp(Trace(sigma_inv_x))/((2**(p*n/S(2))) * multigamma(n/S(2), p))
term2 = (Determinant(scale_matrix))**(-n/S(2))
term3 = (Determinant(x))**(S(n - p - 1)/2)
return term1 * term2 * term3
def test_parsing_of_matrix_expressions():
expr = M * N
assert parse_matrix_expression(expr) == CodegenArrayContraction(
CodegenArrayTensorProduct(M, N), (1, 2))
expr = Transpose(M)
assert parse_matrix_expression(expr) == CodegenArrayPermuteDims(M, [1, 0])
expr = M * Transpose(N)
assert parse_matrix_expression(expr) == CodegenArrayContraction(
CodegenArrayTensorProduct(M, CodegenArrayPermuteDims(N, [1, 0])),
(1, 2))
expr = 3 * M * N
res = parse_matrix_expression(expr)
rexpr = recognize_matrix_expression(res)
assert expr == rexpr
expr = 3 * M + N * M.T * M + 4 * k * N
res = parse_matrix_expression(expr)
rexpr = recognize_matrix_expression(res)
assert expr == rexpr
expr = Inverse(M) * N
rexpr = recognize_matrix_expression(parse_matrix_expression(expr))
assert expr == rexpr
expr = M**2
rexpr = recognize_matrix_expression(parse_matrix_expression(expr))
assert expr == rexpr
expr = M * (2 * N + 3 * M)
res = parse_matrix_expression(expr)
rexpr = recognize_matrix_expression(res)
assert expr.expand() == rexpr.doit()
expr = Trace(M)
result = CodegenArrayContraction(M, (0, 1))
assert parse_matrix_expression(expr) == result
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_convert_array_to_hadamard_products():
expr = HadamardProduct(M, N)
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert ret == expr
expr = HadamardProduct(M, N)*P
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert ret == expr
expr = Q*HadamardProduct(M, N)*P
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert ret == expr
expr = Q*HadamardProduct(M, N.T)*P
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert ret == expr
expr = HadamardProduct(M, N)*HadamardProduct(Q, P)
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert expr == ret
expr = P.T*HadamardProduct(M, N)*HadamardProduct(Q, P)
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert expr == ret
# ArrayDiagonal should be converted
cg = ArrayDiagonal(ArrayTensorProduct(M, N, Q), (1, 3), (0, 2, 4))
ret = convert_array_to_matrix(cg)
expected = PermuteDims(ArrayDiagonal(ArrayTensorProduct(HadamardProduct(M.T, N.T), Q), (1, 2)), [1, 0, 2])
assert expected == ret
# Special case that should return the same expression:
cg = ArrayDiagonal(ArrayTensorProduct(HadamardProduct(M, N), Q), (0, 2))
ret = convert_array_to_matrix(cg)
assert ret == cg
# Hadamard products with traces:
expr = Trace(HadamardProduct(M, N))
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert ret == Trace(HadamardProduct(M.T, N.T))
expr = Trace(A*HadamardProduct(M, N))
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert ret == Trace(HadamardProduct(M, N)*A)
expr = Trace(HadamardProduct(A, M)*N)
cg = convert_matrix_to_array(expr)
ret = convert_array_to_matrix(cg)
assert ret == Trace(HadamardProduct(M.T, N)*A)
# These should not be converted into Hadamard products:
cg = ArrayDiagonal(ArrayTensorProduct(M, N), (0, 1, 2, 3))
ret = convert_array_to_matrix(cg)
assert ret == cg
cg = ArrayDiagonal(ArrayTensorProduct(A), (0, 1))
ret = convert_array_to_matrix(cg)
assert ret == cg
cg = ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 2, 4), (1, 3, 5))
assert convert_array_to_matrix(cg) == HadamardProduct(M, N, P)
cg = ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 3, 4), (1, 2, 5))
assert convert_array_to_matrix(cg) == HadamardProduct(M, P, N.T)
def test_BlockMatrix_Trace():
A, B, C, D = map(lambda s: MatrixSymbol(s, 3, 3), 'ABCD')
X = BlockMatrix([[A, B], [C, D]])
assert Trace(X) == Trace(A) + Trace(D)
def test_trace():
n, m, l = symbols('n m l', integer=True)
A = MatrixSymbol('A', n, n)
B = MatrixSymbol('B', n, n)
assert isinstance(Trace(A), Trace)
assert not isinstance(Trace(A), MatrixExpr)
raises(ShapeError, lambda: Trace(MatrixSymbol('B', 3, 4)))
assert Trace(eye(3)) == 3
assert Trace(Matrix(3, 3, [1, 2, 3, 4, 5, 6, 7, 8, 9])) == 15
A / Trace(A) # Make sure this is possible
# Some easy simplifications
assert Trace(Identity(5)) == 5
assert Trace(ZeroMatrix(5, 5)) == 0
assert Trace(2 * A * B) == 2 * Trace(A * B)
assert Trace(A.T) == Trace(A)
i, j = symbols('i,j')
F = FunctionMatrix(3, 3, Lambda((i, j), i + j))
assert Trace(F).doit() == (0 + 0) + (1 + 1) + (2 + 2)
raises(TypeError, lambda: Trace(S.One))
assert Trace(A).arg is A
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_products_support_function():
A = MatrixSymbol("A", k, k)
B = MatrixSymbol("B", k, k)
C = MatrixSymbol("C", k, k)
D = MatrixSymbol("D", k, k)
X = MatrixSymbol("X", k, k)
Y = MatrixSymbol("Y", k, k)
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]) == CodegenArrayTensorProduct(A * B, C * D)
assert _support_function_tp1_recognize(
[(1, 4), (3, 6)], [A, B, C, D]) == CodegenArrayPermuteDims(
CodegenArrayTensorProduct(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]) == CodegenArrayTensorProduct(X * Y * A * B, C * D)
assert _support_function_tp1_recognize(
[(1, 7), (3, 8),
(4, 11)], [X, Y, A, B, C, D]) == CodegenArrayPermuteDims(
CodegenArrayTensorProduct(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]) == CodegenArrayPermuteDims(
CodegenArrayTensorProduct(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]) == CodegenArrayPermuteDims(
CodegenArrayTensorProduct(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]) == CodegenArrayPermuteDims(
CodegenArrayTensorProduct(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_parsing_of_matrix_expressions():
expr = M*N
assert parse_matrix_expression(expr) == CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2))
expr = Transpose(M)
assert parse_matrix_expression(expr) == CodegenArrayPermuteDims(M, [1, 0])
expr = M*Transpose(N)
assert parse_matrix_expression(expr) == CodegenArrayContraction(CodegenArrayTensorProduct(M, CodegenArrayPermuteDims(N, [1, 0])), (1, 2))
expr = 3*M*N
res = parse_matrix_expression(expr)
rexpr = recognize_matrix_expression(res)
assert expr == rexpr
expr = 3*M + N*M.T*M + 4*k*N
res = parse_matrix_expression(expr)
rexpr = recognize_matrix_expression(res)
assert expr == rexpr
expr = Inverse(M)*N
rexpr = recognize_matrix_expression(parse_matrix_expression(expr))
assert expr == rexpr
expr = M**2
rexpr = recognize_matrix_expression(parse_matrix_expression(expr))
assert expr == rexpr
expr = M*(2*N + 3*M)
res = parse_matrix_expression(expr)
rexpr = recognize_matrix_expression(res)
assert expr == rexpr
expr = Trace(M)
result = CodegenArrayContraction(M, (0, 1))
assert parse_matrix_expression(expr) == result
expr = 3*Trace(M)
result = CodegenArrayContraction(CodegenArrayTensorProduct(3, M), (0, 1))
assert parse_matrix_expression(expr) == result
expr = 3*Trace(Trace(M) * M)
result = CodegenArrayContraction(CodegenArrayTensorProduct(3, M, M), (0, 1), (2, 3))
assert parse_matrix_expression(expr) == result
expr = 3*Trace(M)**2
result = CodegenArrayContraction(CodegenArrayTensorProduct(3, M, M), (0, 1), (2, 3))
assert parse_matrix_expression(expr) == result
expr = HadamardProduct(M, N)
result = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (0, 2), (1, 3))
assert parse_matrix_expression(expr) == result
expr = HadamardPower(M, 2)
result = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, M), (0, 2), (1, 3))
assert parse_matrix_expression(expr) == result
expr = M**2
assert isinstance(expr, MatPow)
assert parse_matrix_expression(expr) == CodegenArrayContraction(CodegenArrayTensorProduct(M, M), (1, 2))
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 * I1 * 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 cg.transform_to_product() == CodegenArrayContraction(
CodegenArrayTensorProduct(A, DiagMatrix(a)), (1, 2))
assert recognize_matrix_expression(cg) == A * DiagMatrix(a)
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a, b), (0, 2))
assert cg.transform_to_product() == CodegenArrayContraction(
CodegenArrayTensorProduct(DiagMatrix(a), b), (0, 2))
assert recognize_matrix_expression(cg).doit() == DiagMatrix(a) * b
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(A, a), (0, 2))
assert cg.transform_to_product() == CodegenArrayContraction(
CodegenArrayTensorProduct(A, DiagMatrix(a)), (0, 2))
assert recognize_matrix_expression(cg) == A.T * DiagMatrix(a)
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(I, x, I1), (0, 2),
(3, 5))
assert cg.transform_to_product() == CodegenArrayContraction(
CodegenArrayTensorProduct(I, DiagMatrix(x), I1), (0, 2))
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(I, x, A, B), (1, 2),
(5, 6))
assert cg.transform_to_product() == CodegenArrayDiagonal(
CodegenArrayContraction(
CodegenArrayTensorProduct(I, DiagMatrix(x), A, B), (1, 2)), (3, 4))
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 cg.transform_to_product() == CodegenArrayContraction(
CodegenArrayTensorProduct(DiagMatrix(x), I), (0, 2))
assert recognize_matrix_expression(cg).doit() == DiagMatrix(x)
cg = CodegenArrayDiagonal(x, (1, ))
assert cg == x
# 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))
assert cg.split_multiple_contractions() == CodegenArrayContraction(
CodegenArrayTensorProduct(A, DiagMatrix(a), C, DiagMatrix(a), B),
(1, 2), (3, 4), (5, 6), (7, 8))
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, I1, b, I1, (a.T * b).applyfunc(cos)),
(1, 2), (3, 8), (5, 6), (7, 9)))
assert recognize_matrix_expression(cg).dummy_eq(
MatMul(a, I1, (a.T * b).applyfunc(cos), Transpose(I1), 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(A.T, DiagMatrix(a), b, b.T,
(A * X * b).applyfunc(cos)), (1, 2),
(3, 8), (5, 6, 9)))
# 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
