Modified behavior of __getitem__ for arrays

sympy/tensor/array/arrayop.py

@@ -160,7 +160,8 @@ def tensorcontraction(array, *contraction_axes):
         index_base_position = sum(icontrib)
         isum = S.Zero
         for sum_to_index in itertools.product(*summed_deltas):
-            isum += array[index_base_position + sum(sum_to_index)]
+            idx = array._get_tuple_index(index_base_position + sum(sum_to_index))
+            isum += array[idx]
 
         contracted_array.append(isum)
 

sympy/tensor/array/dense_ndim_array.py

@@ -5,6 +5,9 @@
 from sympy.core.sympify import _sympify
 from sympy.tensor.array.mutable_ndim_array import MutableNDimArray
 from sympy.tensor.array.ndim_array import NDimArray, ImmutableNDimArray
+from sympy.core.compatibility import SYMPY_INTS
+from sympy.core.numbers import Integer
+
 
 
 class DenseNDimArray(NDimArray):
@@ -27,6 +30,11 @@ def __getitem__(self, index):
         0
         >>> a[1, 1]
         3
+        >>> a[0]
+        [0, 1]
+        >>> a[1]
+        [2, 3]
+
 
         Symbolic index:
 
@@ -44,6 +52,12 @@ def __getitem__(self, index):
         if syindex is not None:
             return syindex
 
+        if isinstance(index, (SYMPY_INTS, Integer)):
+            index = (index, )
+        if not isinstance(index, slice) and len(index) < self.rank():
+            index = tuple([i for i in index] + \
+                          [slice(None) for i in range(len(index), self.rank())])
+
         if isinstance(index, tuple) and any([isinstance(i, slice) for i in index]):
             sl_factors, eindices = self._get_slice_data_for_array_access(index)
             array = [self._array[self._parse_index(i)] for i in eindices]
@@ -53,6 +67,8 @@ def __getitem__(self, index):
             if isinstance(index, slice):
                 return self._array[index]
             else:
+                if self.shape == ():
+                    index = ()
                 index = self._parse_index(index)
                 return self._array[index]
 
@@ -86,7 +102,10 @@ def tomatrix(self):
         return Matrix(self.shape[0], self.shape[1], self._array)
 
     def __iter__(self):
-        return self._array.__iter__()
+        def iterator():
+            for i in range(self._loop_size):
+                yield self[self._get_tuple_index(i)]
+        return iterator()
 
     def reshape(self, *newshape):
         """
@@ -164,7 +183,7 @@ def _new(cls, iterable, shape, **kwargs):
         self._shape = shape
         self._array = list(flat_list)
         self._rank = len(shape)
-        self._loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else 0
+        self._loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else len(flat_list)
         return self
 
     def __setitem__(self, index, value):

sympy/tensor/array/ndim_array.py

@@ -69,11 +69,11 @@ def __new__(cls, iterable, shape=None, **kwargs):
         return ImmutableDenseNDimArray(iterable, shape, **kwargs)
 
     def _parse_index(self, index):
-
         if isinstance(index, (SYMPY_INTS, Integer)):
-            if index >= self._loop_size:
-                raise ValueError("index out of range")
-            return index
+            raise ValueError("Only a tuple index is accepted")
+
+        if self._loop_size == 0:
+            raise ValueError("Index not valide with an empty array")
 
         if len(index) != self._rank:
             raise ValueError('Wrong number of array axes')
@@ -306,7 +306,7 @@ def __str__(self):
         """
         def f(sh, shape_left, i, j):
             if len(shape_left) == 1:
-                return "["+", ".join([str(self[e]) for e in range(i, j)])+"]"
+                return "["+", ".join([str(self[self._get_tuple_index(e)]) for e in range(i, j)])+"]"
 
             sh //= shape_left[0]
             return "[" + ", ".join([f(sh, shape_left[1:], i+e*sh, i+(e+1)*sh) for e in range(shape_left[0])]) + "]" # + "\n"*len(shape_left)
@@ -356,7 +356,7 @@ def tolist(self):
 
         def f(sh, shape_left, i, j):
             if len(shape_left) == 1:
-                return [self[e] for e in range(i, j)]
+                return [self[self._get_tuple_index(e)] for e in range(i, j)]
             result = []
             sh //= shape_left[0]
             for e in range(shape_left[0]):

sympy/tensor/array/sparse_ndim_array.py

@@ -4,6 +4,8 @@
 from sympy.core.sympify import _sympify
 from sympy.tensor.array.mutable_ndim_array import MutableNDimArray
 from sympy.tensor.array.ndim_array import NDimArray, ImmutableNDimArray
+from sympy.core.numbers import Integer
+from sympy.core.compatibility import SYMPY_INTS
 
 import functools
 
@@ -28,9 +30,9 @@ def __getitem__(self, index):
         >>> a[1, 1]
         3
         >>> a[0]
-        0
-        >>> a[2]
-        2
+        [0, 1]
+        >>> a[1]
+        [2, 3]
 
         Symbolic indexing:
 
@@ -48,6 +50,12 @@ def __getitem__(self, index):
         if syindex is not None:
             return syindex
 
+        if isinstance(index, (SYMPY_INTS, Integer)):
+            index = (index, )
+        if not isinstance(index, slice) and len(index) < self.rank():
+            index = tuple([i for i in index] + \
+                          [slice(None) for i in range(len(index), self.rank())])
+
         # `index` is a tuple with one or more slices:
         if isinstance(index, tuple) and any([isinstance(i, slice) for i in index]):
             sl_factors, eindices = self._get_slice_data_for_array_access(index)
@@ -101,7 +109,7 @@ def tomatrix(self):
     def __iter__(self):
         def iterator():
             for i in range(self._loop_size):
-                yield self[i]
+                yield self[self._get_tuple_index(i)]
         return iterator()
 
     def reshape(self, *newshape):
@@ -119,7 +127,7 @@ def __new__(cls, iterable=None, shape=None, **kwargs):
         shape, flat_list = cls._handle_ndarray_creation_inputs(iterable, shape, **kwargs)
         shape = Tuple(*map(_sympify, shape))
         cls._check_special_bounds(flat_list, shape)
-        loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else 0
+        loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else len(flat_list)
 
         # Sparse array:
         if isinstance(flat_list, (dict, Dict)):
@@ -156,7 +164,7 @@ def __new__(cls, iterable=None, shape=None, **kwargs):
         self = object.__new__(cls)
         self._shape = shape
         self._rank = len(shape)
-        self._loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else 0
+        self._loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else len(flat_list)
 
         # Sparse array:
         if isinstance(flat_list, (dict, Dict)):

sympy/tensor/array/tests/test_immutable_ndim_array.py

@@ -103,6 +103,23 @@ def test_reshape():
     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])
+
+
 def test_iterator():
     array = ImmutableDenseNDimArray(range(4), (2, 2))
     j = 0
@@ -185,7 +202,7 @@ def test_ndim_array_converting():
     assert (isinstance(matrix, Matrix))
 
     for i in range(len(dense_array)):
-        assert dense_array[i] == matrix[i]
+        assert dense_array[dense_array._get_tuple_index(i)] == matrix[i]
     assert matrix.shape == dense_array.shape
 
     assert ImmutableDenseNDimArray(matrix) == dense_array
@@ -201,7 +218,7 @@ def test_ndim_array_converting():
     assert(isinstance(matrix, SparseMatrix))
 
     for i in range(len(sparse_array)):
-        assert sparse_array[i] == matrix[i]
+        assert sparse_array[sparse_array._get_tuple_index(i)] == matrix[i]
     assert matrix.shape == sparse_array.shape
 
     assert ImmutableSparseNDimArray(matrix) == sparse_array

sympy/tensor/array/tests/test_mutable_ndim_array.py

@@ -108,6 +108,23 @@ def test_iterator():
         j += 1
 
 
+def test_getitem():
+    for ArrayType in [MutableDenseNDimArray, MutableSparseNDimArray]:
+        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])
+
+
 def test_sparse():
     sparse_array = MutableSparseNDimArray([0, 0, 0, 1], (2, 2))
     assert len(sparse_array) == 2 * 2
@@ -184,7 +201,7 @@ def test_ndim_array_converting():
     assert (isinstance(matrix, Matrix))
 
     for i in range(len(dense_array)):
-        assert dense_array[i] == matrix[i]
+        assert dense_array[dense_array._get_tuple_index(i)] == matrix[i]
     assert matrix.shape == dense_array.shape
 
     assert MutableDenseNDimArray(matrix) == dense_array
@@ -200,7 +217,7 @@ def test_ndim_array_converting():
     assert(isinstance(matrix, SparseMatrix))
 
     for i in range(len(sparse_array)):
-        assert sparse_array[i] == matrix[i]
+        assert sparse_array[sparse_array._get_tuple_index(i)] == matrix[i]
     assert matrix.shape == sparse_array.shape
 
     assert MutableSparseNDimArray(matrix) == sparse_array