Setitem method for ArraySequence

Author: MarcCote

nibabel/streamlines/array_sequence.py

@@ -6,6 +6,9 @@
 
 import numpy as np
 
+from nibabel.deprecated import deprecate_with_version
+
+
 MEGABYTE = 1024 * 1024
 
 
@@ -148,16 +151,15 @@ def total_nb_rows(self):
         return np.sum(self._lengths)
 
     @property
+    @deprecate_with_version("'ArraySequence.data' property is deprecated.\n"
+                            "Please use the 'ArraySequence.get_data()' method instead",
+                            '3.0', '4.0')
     def data(self):
         """ Elements in this array sequence. """
-        warnings.warn("The 'ArraySequence.data' property has been deprecated"
-                      " in favor of 'ArraySequence.get_data()'.",
-                      DeprecationWarning,
-                      stacklevel=2)
         return self.get_data()
 
     def get_data(self):
-        """ Returns a copy of the elements in this array sequence.
+        """ Returns a *copy* of the elements in this array sequence.
 
         Notes
         -----
@@ -384,7 +386,7 @@ def __getitem__(self, idx):
             seq._lengths = self._lengths[off_idx]
             return seq
 
-        if isinstance(off_idx, list) or is_ndarray_of_int_or_bool(off_idx):
+        if isinstance(off_idx, (list, range)) or is_ndarray_of_int_or_bool(off_idx):
             # Fancy indexing
             seq._offsets = self._offsets[off_idx]
             seq._lengths = self._lengths[off_idx]
@@ -425,7 +427,7 @@ def __setitem__(self, idx, elements):
             offsets = self._offsets[off_idx]
             lengths = self._lengths[off_idx]
 
-        elif isinstance(off_idx, list) or is_ndarray_of_int_or_bool(off_idx):
+        elif isinstance(off_idx, (list, range)) or is_ndarray_of_int_or_bool(off_idx):
             # Fancy indexing
             offsets = self._offsets[off_idx]
             lengths = self._lengths[off_idx]
@@ -434,12 +436,25 @@ def __setitem__(self, idx, elements):
             raise TypeError("Index must be either an int, a slice, a list of int"
                             " or a ndarray of bool! Not " + str(type(idx)))
 
-        if len(lengths) != elements.total_nb_rows:
-            msg = "Trying to set {} sequences with {} sequences."
-            raise TypeError(msg.format(len(lengths), elements.total_nb_rows))
+        if is_array_sequence(elements):
+            if len(lengths) != len(elements):
+                msg = "Trying to set {} sequences with {} sequences."
+                raise ValueError(msg.format(len(lengths), len(elements)))
+
+            if sum(lengths) != elements.total_nb_rows:
+                msg = "Trying to set {} points with {} points."
+                raise ValueError(msg.format(sum(lengths), elements.total_nb_rows))
+
+            for o1, l1, o2, l2 in zip(offsets, lengths, elements._offsets, elements._lengths):
+                data[o1:o1 + l1] = elements._data[o2:o2 + l2]
+
+        elif isinstance(elements, numbers.Number):
+            for o1, l1 in zip(offsets, lengths):
+                data[o1:o1 + l1] = elements
 
-        for o1, l1, o2, l2 in zip(offsets, lengths, elements._offsets, elements._lengths):
-            data[o1:o1 + l1] = elements._data[o2:o2 + l2]
+        else:  # Try to iterate over it.
+            for o1, l1, element in zip(offsets, lengths, elements):
+                data[o1:o1 + l1] = element
 
     def _op(self, op, value=None, inplace=False):
         """ Applies some operator to this arraysequence.

nibabel/streamlines/tests/test_array_sequence.py

@@ -274,6 +274,56 @@ def test_arraysequence_getitem(self):
         check_arr_seq_view(seq_view, SEQ_DATA['seq'])
         check_arr_seq(seq_view, [d[:, 2] for d in SEQ_DATA['data'][::-2]])
 
+    def test_arraysequence_setitem(self):
+        # Set one item
+        seq = SEQ_DATA['seq'] * 0
+        for i, e in enumerate(SEQ_DATA['seq']):
+            seq[i] = e
+
+        check_arr_seq(seq, SEQ_DATA['seq'])
+
+        # Setitem with a scalar.
+        seq = SEQ_DATA['seq'].copy()
+        seq[:] = 0
+        assert_true(seq._data.sum() == 0)
+
+        # Setitem with a list of ndarray.
+        seq = SEQ_DATA['seq'] * 0
+        seq[:] = SEQ_DATA['data']
+        check_arr_seq(seq, SEQ_DATA['data'])
+
+        # Setitem using tuple indexing.
+        seq = ArraySequence(np.arange(900).reshape((50,6,3)))
+        seq[:, 0] = 0
+        assert_true(seq._data[:, 0].sum() == 0)
+
+        # Setitem using tuple indexing.
+        seq = ArraySequence(np.arange(900).reshape((50,6,3)))
+        seq[range(len(seq))] = 0
+        assert_true(seq._data.sum() == 0)
+
+        # Setitem of a slice using another slice.
+        seq = ArraySequence(np.arange(900).reshape((50,6,3)))
+        seq[0:4] = seq[5:9]
+        check_arr_seq(seq[0:4], seq[5:9])
+
+        # Setitem between array sequences with different number of sequences.
+        seq = ArraySequence(np.arange(900).reshape((50,6,3)))
+        assert_raises(ValueError, seq.__setitem__, slice(0, 4), seq[5:10])
+
+        # Setitem between array sequences with different amount of points.
+        seq1 = ArraySequence(np.arange(10).reshape(5, 2))
+        seq2 = ArraySequence(np.arange(15).reshape(5, 3))
+        assert_raises(ValueError, seq1.__setitem__, slice(0, 5), seq2)
+
+        # Setitem between array sequences with different common shape.
+        seq1 = ArraySequence(np.arange(12).reshape(2, 2, 3))
+        seq2 = ArraySequence(np.arange(8).reshape(2, 2, 2))
+        assert_raises(ValueError, seq1.__setitem__, slice(0, 2), seq2)
+
+        # Invalid index.
+        assert_raises(TypeError, seq.__setitem__, object(), None)
+
     def test_arraysequence_operators(self):
         # Disable division per zero warnings.
         flags = np.seterr(divide='ignore', invalid='ignore')
@@ -375,61 +425,6 @@ def _test_binary(op, arrseq, scalars, seqs, inplace=False):
         # Restore flags.
         np.seterr(**flags)
 
-
-    def test_arraysequence_setitem(self):
-        # Set one item
-        seq = SEQ_DATA['seq'] * 0
-        for i, e in enumerate(SEQ_DATA['seq']):
-            seq[i] = e
-
-        check_arr_seq(seq, SEQ_DATA['seq'])
-
-        # Get all items using indexing (creates a view).
-        indices = list(range(len(SEQ_DATA['seq'])))
-        seq_view = SEQ_DATA['seq'][indices]
-        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
-        # We took all elements so the view should match the original.
-        check_arr_seq(seq_view, SEQ_DATA['seq'])
-
-        # Get multiple items using ndarray of dtype integer.
-        for dtype in [np.int8, np.int16, np.int32, np.int64]:
-            seq_view = SEQ_DATA['seq'][np.array(indices, dtype=dtype)]
-            check_arr_seq_view(seq_view, SEQ_DATA['seq'])
-            # We took all elements so the view should match the original.
-            check_arr_seq(seq_view, SEQ_DATA['seq'])
-
-        # Get multiple items out of order (creates a view).
-        SEQ_DATA['rng'].shuffle(indices)
-        seq_view = SEQ_DATA['seq'][indices]
-        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
-        check_arr_seq(seq_view, [SEQ_DATA['data'][i] for i in indices])
-
-        # Get slice (this will create a view).
-        seq_view = SEQ_DATA['seq'][::2]
-        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
-        check_arr_seq(seq_view, SEQ_DATA['data'][::2])
-
-        # Use advanced indexing with ndarray of data type bool.
-        selection = np.array([False, True, True, False, True])
-        seq_view = SEQ_DATA['seq'][selection]
-        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
-        check_arr_seq(seq_view,
-                      [SEQ_DATA['data'][i]
-                       for i, keep in enumerate(selection) if keep])
-
-        # Test invalid indexing
-        assert_raises(TypeError, SEQ_DATA['seq'].__getitem__, 'abc')
-
-        # Get specific columns.
-        seq_view = SEQ_DATA['seq'][:, 2]
-        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
-        check_arr_seq(seq_view, [d[:, 2] for d in SEQ_DATA['data']])
-
-        # Combining multiple slicing and indexing operations.
-        seq_view = SEQ_DATA['seq'][::-2][:, 2]
-        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
-        check_arr_seq(seq_view, [d[:, 2] for d in SEQ_DATA['data'][::-2]])
-
     def test_arraysequence_repr(self):
         # Test that calling repr on a ArraySequence object is not falling.
         repr(SEQ_DATA['seq'])