BF: only apply affine to selected streamlines

Author: MarcCote

nibabel/streamlines/array_sequence.py

@@ -329,6 +329,54 @@ def __getitem__(self, idx):
         raise TypeError("Index must be either an int, a slice, a list of int"
                         " or a ndarray of bool! Not " + str(type(idx)))
 
+    def __setitem__(self, idx, elements):
+        """ Set sequence(s) through standard or advanced numpy indexing.
+
+        Parameters
+        ----------
+        idx : int or slice or list or ndarray
+            If int, index of the element to retrieve.
+            If slice, use slicing to retrieve elements.
+            If list, indices of the elements to retrieve.
+            If ndarray with dtype int, indices of the elements to retrieve.
+            If ndarray with dtype bool, only retrieve selected elements.
+        elements: ndarray or :class:`ArraySequence`
+            Data that will overwrite selected sequences.
+            If `idx` is an int, `elements` is expected to be a ndarray.
+            Otherwise, `elements` is expected a :class:`ArraySequence` object.
+        """
+        if isinstance(idx, (numbers.Integral, np.integer)):
+            start = self._offsets[idx]
+            self._data[start:start + self._lengths[idx]] = elements
+            return
+
+        if isinstance(idx, tuple):
+            off_idx = idx[0]
+            data = self._data.__getitem__((slice(None),) + idx[1:])
+        else:
+            off_idx = idx
+            data = self._data
+
+        if isinstance(off_idx, slice):  # Standard list slicing
+            offsets = self._offsets[off_idx]
+            lengths = self._lengths[off_idx]
+
+        elif isinstance(off_idx, list) or is_ndarray_of_int_or_bool(off_idx):
+            # Fancy indexing
+            offsets = self._offsets[off_idx]
+            lengths = self._lengths[off_idx]
+
+        else:
+            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))
+
+        for o1, l1, o2, l2 in zip(offsets, lengths, elements._offsets, elements._lengths):
+            data[o1:o1 + l1] = elements._data[o2:o2 + l2]
+
     def __iter__(self):
         if len(self._lengths) != len(self._offsets):
             raise ValueError("ArraySequence object corrupted:"

nibabel/streamlines/tests/test_array_sequence.py

@@ -277,6 +277,67 @@ 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'])
+
+        if sys.version_info < (3,):
+            seq = ArraySequence(SEQ_DATA['seq'] * 0)
+            for i, e in enumerate(SEQ_DATA['seq']):
+                seq[long(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'])

nibabel/streamlines/tests/test_streamlines.py

@@ -267,6 +267,19 @@ def test_save_complex_file(self):
                     tfile = nib.streamlines.load(filename, lazy_load=False)
                     assert_tractogram_equal(tfile.tractogram, tractogram)
 
+    def test_save_sliced_tractogram(self):
+        tractogram = Tractogram(DATA['streamlines'],
+                                affine_to_rasmm=np.eye(4))
+        original_tractogram = tractogram.copy()
+        for ext, cls in FORMATS.items():
+            with InTemporaryDirectory():
+                filename = 'streamlines' + ext
+                nib.streamlines.save(tractogram[::2], filename)
+                tfile = nib.streamlines.load(filename, lazy_load=False)
+                assert_tractogram_equal(tfile.tractogram, tractogram[::2])
+                # Make sure original tractogram hasn't changed.
+                assert_tractogram_equal(tractogram, original_tractogram)
+
     def test_load_unknown_format(self):
         assert_raises(ValueError, nib.streamlines.load, "")
 

nibabel/streamlines/tests/test_tractogram.py

@@ -688,6 +688,22 @@ def test_tractogram_apply_affine(self):
                            np.dot(np.eye(4), np.dot(np.linalg.inv(affine),
                                                     np.linalg.inv(affine))))
 
+        # Applying the affine to a tractogram that has been indexed or sliced
+        # shouldn't affect the remaining streamlines.
+        tractogram = DATA['tractogram'].copy()
+        transformed_tractogram = tractogram[::2].apply_affine(affine)
+        assert_true(transformed_tractogram is not tractogram)
+        check_tractogram(tractogram[::2],
+                         streamlines=[s*scaling for s in DATA['streamlines'][::2]],
+                         data_per_streamline=DATA['tractogram'].data_per_streamline[::2],
+                         data_per_point=DATA['tractogram'].data_per_point[::2])
+
+        # Remaining streamlines should match the original ones.
+        check_tractogram(tractogram[1::2],
+                         streamlines=DATA['streamlines'][1::2],
+                         data_per_streamline=DATA['tractogram'].data_per_streamline[1::2],
+                         data_per_point=DATA['tractogram'].data_per_point[1::2])
+
         # Check that applying an affine and its inverse give us back the
         # original streamlines.
         tractogram = DATA['tractogram'].copy()

nibabel/streamlines/tractogram.py

@@ -432,11 +432,8 @@ def apply_affine(self, affine, lazy=False):
         if np.all(affine == np.eye(4)):
             return self  # No transformation.
 
-        BUFFER_SIZE = 10000000  # About 128 Mb since pts shape is 3.
-        for start in range(0, len(self.streamlines.data), BUFFER_SIZE):
-            end = start + BUFFER_SIZE
-            pts = self.streamlines._data[start:end]
-            self.streamlines.data[start:end] = apply_affine(affine, pts)
+        for i in range(len(self.streamlines)):
+            self.streamlines[i] = apply_affine(affine, self.streamlines[i])
 
         if self.affine_to_rasmm is not None:
             # Update the affine that brings back the streamlines to RASmm.