from copy import copy from sympy.tensor.array.dense_ndim_array import ImmutableDenseNDimArray from sympy.core.containers import Dict from sympy.core.function import diff from sympy.core.numbers import Rational from sympy.core.singleton import S from sympy.core.symbol import (Symbol, symbols) from sympy.matrices import SparseMatrix from sympy.tensor.indexed import (Indexed, IndexedBase) from sympy.matrices import Matrix from sympy.tensor.array.sparse_ndim_array import ImmutableSparseNDimArray from sympy.testing.pytest import raises def test_ndim_array_initiation(): arr_with_no_elements = ImmutableDenseNDimArray([], shape=(0,)) assert len(arr_with_no_elements) == 0 assert arr_with_no_elements.rank() == 1 raises(ValueError, lambda: ImmutableDenseNDimArray([0], shape=(0,))) raises(ValueError, lambda: ImmutableDenseNDimArray([1, 2, 3], shape=(0,))) raises(ValueError, lambda: ImmutableDenseNDimArray([], shape=())) raises(ValueError, lambda: ImmutableSparseNDimArray([0], shape=(0,))) raises(ValueError, lambda: ImmutableSparseNDimArray([1, 2, 3], shape=(0,))) raises(ValueError, lambda: ImmutableSparseNDimArray([], shape=())) arr_with_one_element = ImmutableDenseNDimArray([23]) assert len(arr_with_one_element) == 1 assert arr_with_one_element[0] == 23 assert arr_with_one_element[:] == ImmutableDenseNDimArray([23]) assert arr_with_one_element.rank() == 1 arr_with_symbol_element = ImmutableDenseNDimArray([Symbol('x')]) assert len(arr_with_symbol_element) == 1 assert arr_with_symbol_element[0] == Symbol('x') assert arr_with_symbol_element[:] == ImmutableDenseNDimArray([Symbol('x')]) assert arr_with_symbol_element.rank() == 1 number5 = 5 vector = ImmutableDenseNDimArray.zeros(number5) assert len(vector) == number5 assert vector.shape == (number5,) assert vector.rank() == 1 vector = ImmutableSparseNDimArray.zeros(number5) assert len(vector) == number5 assert vector.shape == (number5,) assert vector._sparse_array == Dict() assert vector.rank() == 1 n_dim_array = ImmutableDenseNDimArray(range(3**4), (3, 3, 3, 3,)) assert len(n_dim_array) == 3 * 3 * 3 * 3 assert n_dim_array.shape == (3, 3, 3, 3) assert n_dim_array.rank() == 4 array_shape = (3, 3, 3, 3) sparse_array = ImmutableSparseNDimArray.zeros(*array_shape) assert len(sparse_array._sparse_array) == 0 assert len(sparse_array) == 3 * 3 * 3 * 3 assert n_dim_array.shape == array_shape assert n_dim_array.rank() == 4 one_dim_array = ImmutableDenseNDimArray([2, 3, 1]) assert len(one_dim_array) == 3 assert one_dim_array.shape == (3,) assert one_dim_array.rank() == 1 assert one_dim_array.tolist() == [2, 3, 1] shape = (3, 3) array_with_many_args = ImmutableSparseNDimArray.zeros(*shape) assert len(array_with_many_args) == 3 * 3 assert array_with_many_args.shape == shape assert array_with_many_args[0, 0] == 0 assert array_with_many_args.rank() == 2 shape = (int(3), int(3)) array_with_long_shape = ImmutableSparseNDimArray.zeros(*shape) assert len(array_with_long_shape) == 3 * 3 assert array_with_long_shape.shape == shape assert array_with_long_shape[int(0), int(0)] == 0 assert array_with_long_shape.rank() == 2 vector_with_long_shape = ImmutableDenseNDimArray(range(5), int(5)) assert len(vector_with_long_shape) == 5 assert vector_with_long_shape.shape == (int(5),) assert vector_with_long_shape.rank() == 1 raises(ValueError, lambda: vector_with_long_shape[int(5)]) from sympy.abc import x for ArrayType in [ImmutableDenseNDimArray, ImmutableSparseNDimArray]: rank_zero_array = ArrayType(x) assert len(rank_zero_array) == 1 assert rank_zero_array.shape == () assert rank_zero_array.rank() == 0 assert rank_zero_array[()] == x raises(ValueError, lambda: rank_zero_array[0]) def test_reshape(): array = ImmutableDenseNDimArray(range(50), 50) assert array.shape == (50,) assert array.rank() == 1 array = array.reshape(5, 5, 2) assert array.shape == (5, 5, 2) assert array.rank() == 3 assert len(array) == 50 def test_getitem(): for ArrayType in [ImmutableDenseNDimArray, ImmutableSparseNDimArray]: array = ArrayType(range(24)).reshape(2, 3, 4) assert array.tolist() == [[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]] assert array[0] == ArrayType([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]) assert array[0, 0] == ArrayType([0, 1, 2, 3]) value = 0 for i in range(2): for j in range(3): for k in range(4): assert array[i, j, k] == value value += 1 raises(ValueError, lambda: array[3, 4, 5]) raises(ValueError, lambda: array[3, 4, 5, 6]) raises(ValueError, lambda: array[3, 4, 5, 3:4]) def test_iterator(): array = ImmutableDenseNDimArray(range(4), (2, 2)) assert array[0] == ImmutableDenseNDimArray([0, 1]) assert array[1] == ImmutableDenseNDimArray([2, 3]) array = array.reshape(4) j = 0 for i in array: assert i == j j += 1 def test_sparse(): sparse_array = ImmutableSparseNDimArray([0, 0, 0, 1], (2, 2)) assert len(sparse_array) == 2 * 2 # dictionary where all data is, only non-zero entries are actually stored: assert len(sparse_array._sparse_array) == 1 assert sparse_array.tolist() == [[0, 0], [0, 1]] for i, j in zip(sparse_array, [[0, 0], [0, 1]]): assert i == ImmutableSparseNDimArray(j) def sparse_assignment(): sparse_array[0, 0] = 123 assert len(sparse_array._sparse_array) == 1 raises(TypeError, sparse_assignment) assert len(sparse_array._sparse_array) == 1 assert sparse_array[0, 0] == 0 assert sparse_array/0 == ImmutableSparseNDimArray([[S.NaN, S.NaN], [S.NaN, S.ComplexInfinity]], (2, 2)) # test for large scale sparse array # equality test assert ImmutableSparseNDimArray.zeros(100000, 200000) == ImmutableSparseNDimArray.zeros(100000, 200000) # __mul__ and __rmul__ a = ImmutableSparseNDimArray({200001: 1}, (100000, 200000)) assert a * 3 == ImmutableSparseNDimArray({200001: 3}, (100000, 200000)) assert 3 * a == ImmutableSparseNDimArray({200001: 3}, (100000, 200000)) assert a * 0 == ImmutableSparseNDimArray({}, (100000, 200000)) assert 0 * a == ImmutableSparseNDimArray({}, (100000, 200000)) # __truediv__ assert a/3 == ImmutableSparseNDimArray({200001: Rational(1, 3)}, (100000, 200000)) # __neg__ assert -a == ImmutableSparseNDimArray({200001: -1}, (100000, 200000)) def test_calculation(): a = ImmutableDenseNDimArray([1]*9, (3, 3)) b = ImmutableDenseNDimArray([9]*9, (3, 3)) c = a + b for i in c: assert i == ImmutableDenseNDimArray([10, 10, 10]) assert c == ImmutableDenseNDimArray([10]*9, (3, 3)) assert c == ImmutableSparseNDimArray([10]*9, (3, 3)) c = b - a for i in c: assert i == ImmutableDenseNDimArray([8, 8, 8]) assert c == ImmutableDenseNDimArray([8]*9, (3, 3)) assert c == ImmutableSparseNDimArray([8]*9, (3, 3)) def test_ndim_array_converting(): dense_array = ImmutableDenseNDimArray([1, 2, 3, 4], (2, 2)) alist = dense_array.tolist() assert alist == [[1, 2], [3, 4]] matrix = dense_array.tomatrix() assert (isinstance(matrix, Matrix)) for i in range(len(dense_array)): assert dense_array[dense_array._get_tuple_index(i)] == matrix[i] assert matrix.shape == dense_array.shape assert ImmutableDenseNDimArray(matrix) == dense_array assert ImmutableDenseNDimArray(matrix.as_immutable()) == dense_array assert ImmutableDenseNDimArray(matrix.as_mutable()) == dense_array sparse_array = ImmutableSparseNDimArray([1, 2, 3, 4], (2, 2)) alist = sparse_array.tolist() assert alist == [[1, 2], [3, 4]] matrix = sparse_array.tomatrix() assert(isinstance(matrix, SparseMatrix)) for i in range(len(sparse_array)): assert sparse_array[sparse_array._get_tuple_index(i)] == matrix[i] assert matrix.shape == sparse_array.shape assert ImmutableSparseNDimArray(matrix) == sparse_array assert ImmutableSparseNDimArray(matrix.as_immutable()) == sparse_array assert ImmutableSparseNDimArray(matrix.as_mutable()) == sparse_array def test_converting_functions(): arr_list = [1, 2, 3, 4] arr_matrix = Matrix(((1, 2), (3, 4))) # list arr_ndim_array = ImmutableDenseNDimArray(arr_list, (2, 2)) assert (isinstance(arr_ndim_array, ImmutableDenseNDimArray)) assert arr_matrix.tolist() == arr_ndim_array.tolist() # Matrix arr_ndim_array = ImmutableDenseNDimArray(arr_matrix) assert (isinstance(arr_ndim_array, ImmutableDenseNDimArray)) assert arr_matrix.tolist() == arr_ndim_array.tolist() assert arr_matrix.shape == arr_ndim_array.shape def test_equality(): first_list = [1, 2, 3, 4] second_list = [1, 2, 3, 4] third_list = [4, 3, 2, 1] assert first_list == second_list assert first_list != third_list first_ndim_array = ImmutableDenseNDimArray(first_list, (2, 2)) second_ndim_array = ImmutableDenseNDimArray(second_list, (2, 2)) fourth_ndim_array = ImmutableDenseNDimArray(first_list, (2, 2)) assert first_ndim_array == second_ndim_array def assignment_attempt(a): a[0, 0] = 0 raises(TypeError, lambda: assignment_attempt(second_ndim_array)) assert first_ndim_array == second_ndim_array assert first_ndim_array == fourth_ndim_array def test_arithmetic(): a = ImmutableDenseNDimArray([3 for i in range(9)], (3, 3)) b = ImmutableDenseNDimArray([7 for i in range(9)], (3, 3)) c1 = a + b c2 = b + a assert c1 == c2 d1 = a - b d2 = b - a assert d1 == d2 * (-1) e1 = a * 5 e2 = 5 * a e3 = copy(a) e3 *= 5 assert e1 == e2 == e3 f1 = a / 5 f2 = copy(a) f2 /= 5 assert f1 == f2 assert f1[0, 0] == f1[0, 1] == f1[0, 2] == f1[1, 0] == f1[1, 1] == \ f1[1, 2] == f1[2, 0] == f1[2, 1] == f1[2, 2] == Rational(3, 5) assert type(a) == type(b) == type(c1) == type(c2) == type(d1) == type(d2) \ == type(e1) == type(e2) == type(e3) == type(f1) z0 = -a assert z0 == ImmutableDenseNDimArray([-3 for i in range(9)], (3, 3)) def test_higher_dimenions(): m3 = ImmutableDenseNDimArray(range(10, 34), (2, 3, 4)) assert m3.tolist() == [[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]] assert m3._get_tuple_index(0) == (0, 0, 0) assert m3._get_tuple_index(1) == (0, 0, 1) assert m3._get_tuple_index(4) == (0, 1, 0) assert m3._get_tuple_index(12) == (1, 0, 0) assert str(m3) == '[[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]]' m3_rebuilt = ImmutableDenseNDimArray([[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]]) assert m3 == m3_rebuilt m3_other = ImmutableDenseNDimArray([[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]], (2, 3, 4)) assert m3 == m3_other def test_rebuild_immutable_arrays(): sparr = ImmutableSparseNDimArray(range(10, 34), (2, 3, 4)) densarr = ImmutableDenseNDimArray(range(10, 34), (2, 3, 4)) assert sparr == sparr.func(*sparr.args) assert densarr == densarr.func(*densarr.args) def test_slices(): md = ImmutableDenseNDimArray(range(10, 34), (2, 3, 4)) assert md[:] == ImmutableDenseNDimArray(range(10, 34), (2, 3, 4)) assert md[:, :, 0].tomatrix() == Matrix([[10, 14, 18], [22, 26, 30]]) assert md[0, 1:2, :].tomatrix() == Matrix([[14, 15, 16, 17]]) assert md[0, 1:3, :].tomatrix() == Matrix([[14, 15, 16, 17], [18, 19, 20, 21]]) assert md[:, :, :] == md sd = ImmutableSparseNDimArray(range(10, 34), (2, 3, 4)) assert sd == ImmutableSparseNDimArray(md) assert sd[:] == ImmutableSparseNDimArray(range(10, 34), (2, 3, 4)) assert sd[:, :, 0].tomatrix() == Matrix([[10, 14, 18], [22, 26, 30]]) assert sd[0, 1:2, :].tomatrix() == Matrix([[14, 15, 16, 17]]) assert sd[0, 1:3, :].tomatrix() == Matrix([[14, 15, 16, 17], [18, 19, 20, 21]]) assert sd[:, :, :] == sd def test_diff_and_applyfunc(): from sympy.abc import x, y, z md = ImmutableDenseNDimArray([[x, y], [x*z, x*y*z]]) assert md.diff(x) == ImmutableDenseNDimArray([[1, 0], [z, y*z]]) assert diff(md, x) == ImmutableDenseNDimArray([[1, 0], [z, y*z]]) sd = ImmutableSparseNDimArray(md) assert sd == ImmutableSparseNDimArray([x, y, x*z, x*y*z], (2, 2)) assert sd.diff(x) == ImmutableSparseNDimArray([[1, 0], [z, y*z]]) assert diff(sd, x) == ImmutableSparseNDimArray([[1, 0], [z, y*z]]) mdn = md.applyfunc(lambda x: x*3) assert mdn == ImmutableDenseNDimArray([[3*x, 3*y], [3*x*z, 3*x*y*z]]) assert md != mdn sdn = sd.applyfunc(lambda x: x/2) assert sdn == ImmutableSparseNDimArray([[x/2, y/2], [x*z/2, x*y*z/2]]) assert sd != sdn sdp = sd.applyfunc(lambda x: x+1) assert sdp == ImmutableSparseNDimArray([[x + 1, y + 1], [x*z + 1, x*y*z + 1]]) assert sd != sdp def test_op_priority(): from sympy.abc import x md = ImmutableDenseNDimArray([1, 2, 3]) e1 = (1+x)*md e2 = md*(1+x) assert e1 == ImmutableDenseNDimArray([1+x, 2+2*x, 3+3*x]) assert e1 == e2 sd = ImmutableSparseNDimArray([1, 2, 3]) e3 = (1+x)*sd e4 = sd*(1+x) assert e3 == ImmutableDenseNDimArray([1+x, 2+2*x, 3+3*x]) assert e3 == e4 def test_symbolic_indexing(): x, y, z, w = symbols("x y z w") M = ImmutableDenseNDimArray([[x, y], [z, w]]) i, j = symbols("i, j") Mij = M[i, j] assert isinstance(Mij, Indexed) Ms = ImmutableSparseNDimArray([[2, 3*x], [4, 5]]) msij = Ms[i, j] assert isinstance(msij, Indexed) for oi, oj in [(0, 0), (0, 1), (1, 0), (1, 1)]: assert Mij.subs({i: oi, j: oj}) == M[oi, oj] assert msij.subs({i: oi, j: oj}) == Ms[oi, oj] A = IndexedBase("A", (0, 2)) assert A[0, 0].subs(A, M) == x assert A[i, j].subs(A, M) == M[i, j] assert M[i, j].subs(M, A) == A[i, j] assert isinstance(M[3 * i - 2, j], Indexed) assert M[3 * i - 2, j].subs({i: 1, j: 0}) == M[1, 0] assert isinstance(M[i, 0], Indexed) assert M[i, 0].subs(i, 0) == M[0, 0] assert M[0, i].subs(i, 1) == M[0, 1] assert M[i, j].diff(x) == ImmutableDenseNDimArray([[1, 0], [0, 0]])[i, j] assert Ms[i, j].diff(x) == ImmutableSparseNDimArray([[0, 3], [0, 0]])[i, j] Mo = ImmutableDenseNDimArray([1, 2, 3]) assert Mo[i].subs(i, 1) == 2 Mos = ImmutableSparseNDimArray([1, 2, 3]) assert Mos[i].subs(i, 1) == 2 raises(ValueError, lambda: M[i, 2]) raises(ValueError, lambda: M[i, -1]) raises(ValueError, lambda: M[2, i]) raises(ValueError, lambda: M[-1, i]) raises(ValueError, lambda: Ms[i, 2]) raises(ValueError, lambda: Ms[i, -1]) raises(ValueError, lambda: Ms[2, i]) raises(ValueError, lambda: Ms[-1, i]) def test_issue_12665(): # Testing Python 3 hash of immutable arrays: arr = ImmutableDenseNDimArray([1, 2, 3]) # This should NOT raise an exception: hash(arr) def test_zeros_without_shape(): arr = ImmutableDenseNDimArray.zeros() assert arr == ImmutableDenseNDimArray(0) def test_issue_21870(): a0 = ImmutableDenseNDimArray(0) assert a0.rank() == 0 a1 = ImmutableDenseNDimArray(a0) assert a1.rank() == 0