Repository: sympy/sympy
File: sympy/tensor/array/expressions/tests/test_array_expressions.py (Python)

Code before:
    assert PermuteDims(A, [1, 0]).as_explicit() == permutedims(A, [1, 0])
    assert Reshape(A, [4]).as_explicit() == A.reshape(4)


def test_array_sum():
    expr = ArraySum(X, (i, 1, j))
    assert isinstance(expr, Sum)
    assert expr.doit() == j*X

    expr = ArraySum(X, (i, 1, 5))
    assert expr.doit() == 5*X

    expr = ArrayTensorProduct(ArraySum(X, (i, 1, 5)), Y)
    assert expr.doit() == ArrayTensorProduct(5*X, Y)

    expr = ArrayTensorProduct(A, ArraySum(X, (i, 1, j)), B, ArraySum(Y, (i, 1, m)), C)
    assert expr.doit() == ArrayTensorProduct(A, j*X, B, m*Y, C)

    expr = ArrayTensorProduct(A, ArraySum(X*sin(i), (i, 1, j)), B, ArraySum(Y*sin(i), (i, 1, m)), C)
    assert expr.doit().dummy_eq(ArraySum(ArrayTensorProduct(A, sin(i)*X, B, sin(n)*Y, C), (i, 1, j), (n, 1, m)))

    expr = ArrayTensorProduct(ArraySum(X**sin(i), (i, 1, j)), Y)
    assert expr.doit().dummy_eq(ArraySum(ArrayTensorProduct(X**sin(i), Y), (i, 1, j)))

    expr = ArrayTensorProduct(ArraySum(X**sin(i), (i, 1, j)), i*Y)
    assert expr.doit().dummy_eq(ArraySum(X**sin(m)*i*Y, (m, 1, j)))

    expr = ArrayContraction(ArraySum(X, (i, 0, j)), (0, 1))
    # assert expr.doit() == ArrayContraction(j*T, (0, 1))  # TODO: Not working!

    T = MatrixSymbol("T", 3, 3)
    expr = ArrayContraction(ArraySum(T, (i, 1, j)), (0, 1))
    assert expr.doit() == ArrayContraction(j*T, (0, 1))

    expr = ArrayContraction(ArraySum(T*sin(i), (i, 1, j)), (0, 1))
    assert expr.doit() == ArraySum(ArrayContraction(T*sin(i), (0, 1)), (i, 1, j))


Review comment (bug):
This is wrong, `X**sin(i)` creates a `Pow` object, which has shape `()`. This should not be allowed.

Code after:
    assert PermuteDims(A, [1, 0]).as_explicit() == permutedims(A, [1, 0])
    assert Reshape(A, [4]).as_explicit() == A.reshape(4)


def test_array_sum():
    X = MatrixSymbol("X", k, k)
    Y = MatrixSymbol("Y", k, k)

    expr = ArraySum(X, (i, 1, j))
    assert isinstance(expr, Sum)
    assert expr.doit() == j*X

    expr = ArraySum(X, (i, 1, 5))
    assert expr.doit() == 5*X

    expr = ArrayTensorProduct(ArraySum(X, (i, 1, 5)), Y)
    assert expr.doit() == ArrayTensorProduct(5*X, Y)

    expr = ArrayTensorProduct(A, ArraySum(X, (i, 1, j)), B, ArraySum(Y, (i, 1, m)), C)
    assert expr.doit() == ArrayTensorProduct(A, j*X, B, m*Y, C)

    expr = ArrayTensorProduct(A, ArraySum(X*sin(i), (i, 1, j)), B, ArraySum(Y*sin(i), (i, 1, m)), C)
    assert expr.doit().dummy_eq(ArraySum(ArrayTensorProduct(A, sin(i)*X, B, sin(n)*Y, C), (i, 1, j), (n, 1, m)))

    expr = ArrayTensorProduct(ArraySum(X**sin(i), (i, 1, j)), Y)
    assert expr.doit().dummy_eq(ArraySum(ArrayTensorProduct(X**sin(i), Y), (i, 1, j)))

    expr = ArrayTensorProduct(ArraySum(X**sin(i), (i, 1, j)), i*Y)
    assert expr.doit().dummy_eq(ArraySum(ArrayTensorProduct(X**sin(m), i*Y), (m, 1, j)))

    expr = ArrayContraction(ArraySum(X, (i, 0, j)), (0, 1))
    assert expr.doit() == ArrayContraction((j + 1)*X, (0, 1))

    T = MatrixSymbol("T", 3, 3)
    expr = ArrayContraction(ArraySum(T, (i, 1, j)), (0, 1))
    assert expr.doit() == ArrayContraction(j*T, (0, 1))

    expr = ArrayContraction(ArraySum(T*sin(i), (i, 1, j)), (0, 1))
    assert expr.doit() == ArraySum(ArrayContraction(T*sin(i), (0, 1)), (i, 1, j))