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)
def test_convert_array_element_to_matrix():
expr = ArrayElement(M, (i, j))
assert convert_array_to_matrix(expr) == MatrixElement(M, i, j)
expr = ArrayElement(_array_contraction(_array_tensor_product(M, N), (1, 3)), (i, j))
assert convert_array_to_matrix(expr) == MatrixElement(M*N.T, i, j)
expr = ArrayElement(_array_tensor_product(M, N), (i, j, m, n))
assert convert_array_to_matrix(expr) == expr
def test_arrayexpr_convert_indexed_to_array_and_back_to_matrix():
expr = a.T * b
elem = expr[0, 0]
cg = convert_indexed_to_array(elem)
assert cg == ArrayElement(
ArrayContraction(ArrayTensorProduct(a, b), (0, 2)), [0, 0])
expr = M[i, j] + N[i, j]
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(p1) == M + N
expr = M[i, j] + N[j, i]
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(p1) == M + N.T
expr = M[i, j] * N[k, l] + N[i, j] * M[k, l]
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(p1) == ArrayAdd(ArrayTensorProduct(M, N),
ArrayTensorProduct(N, M))
expr = (M * N * P)[i, j]
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(p1) == M * N * P
expr = Sum(M[i, j] * (N * P)[j, m], (j, 0, k - 1))
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(p1) == M * N * P
expr = Sum((P[j, m] + P[m, j]) * (M[i, j] * N[m, n] + N[i, j] * M[m, n]),
(j, 0, k - 1), (m, 0, k - 1))
p1, p2 = _convert_indexed_to_array(expr)
assert convert_array_to_matrix(
p1) == M * P * N + M * P.T * N + N * P * M + N * P.T * M
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_printing_str_array_expressions():
assert sstr(ArraySymbol("A", 2, 3, 4)) == "A"
assert sstr(ArrayElement("A", (2, 1/(1-x), 0))) == "A[2, 1/(1 - x), 0]"
def convert_indexed_to_array(expr, first_indices=None):
r"""
Parse indexed expression into a form useful for code generation.
Examples
========
>>> from sympy.tensor.array.expressions.conv_indexed_to_array import convert_indexed_to_array
>>> from sympy import MatrixSymbol, Sum, symbols
>>> i, j, k, d = symbols("i j k d")
>>> M = MatrixSymbol("M", d, d)
>>> N = MatrixSymbol("N", d, d)
Recognize the trace in summation form:
>>> expr = Sum(M[i, i], (i, 0, d-1))
>>> convert_indexed_to_array(expr)
ArrayContraction(M, (0, 1))
Recognize the extraction of the diagonal by using the same index `i` on
both axes of the matrix:
>>> expr = M[i, i]
>>> convert_indexed_to_array(expr)
ArrayDiagonal(M, (0, 1))
This function can help perform the transformation expressed in two
different mathematical notations as:
`\sum_{j=0}^{N-1} A_{i,j} B_{j,k} \Longrightarrow \mathbf{A}\cdot \mathbf{B}`
Recognize the matrix multiplication in summation form:
>>> expr = Sum(M[i, j]*N[j, k], (j, 0, d-1))
>>> convert_indexed_to_array(expr)
ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
Specify that ``k`` has to be the starting index:
>>> convert_indexed_to_array(expr, first_indices=[k])
ArrayContraction(ArrayTensorProduct(N, M), (0, 3))
"""
result, indices = _convert_indexed_to_array(expr)
if any(isinstance(i, (int, Integer)) for i in indices):
result = ArrayElement(result, indices)
indices = []
if not first_indices:
return result
def _check_is_in(elem, indices):
if elem in indices:
return True
if any(elem in i for i in indices if isinstance(i, frozenset)):
return True
return False
repl = {j: i for i in indices if isinstance(i, frozenset) for j in i}
first_indices = [repl.get(i, i) for i in first_indices]
for i in first_indices:
if not _check_is_in(i, indices):
first_indices.remove(i)
first_indices.extend(
[i for i in indices if not _check_is_in(i, first_indices)])
def _get_pos(elem, indices):
if elem in indices:
return indices.index(elem)
for i, e in enumerate(indices):
if not isinstance(e, frozenset):
continue
if elem in e:
return i
raise ValueError("not found")
permutation = _af_invert([_get_pos(i, first_indices) for i in indices])
if isinstance(result, ArrayAdd):
return _array_add(
*[_permute_dims(arg, permutation) for arg in result.args])
else:
return _permute_dims(result, permutation)
def test_printing_str_array_expressions():
assert sstr(ArraySymbol("A", (2, 3, 4))) == "A"
assert sstr(ArrayElement("A", (2, 1/(1-x), 0))) == "A[2, 1/(1 - x), 0]"
M = MatrixSymbol("M", 3, 3)
N = MatrixSymbol("N", 3, 3)
assert sstr(ArrayElement(M*N, [x, 0])) == "(M*N)[x, 0]"
def _(expr: ArrayElement):
ret = _array2matrix(expr.name)
if isinstance(ret, MatrixExpr):
return MatrixElement(ret, *expr.indices)
return ArrayElement(ret, expr.indices)