Fixed affine transformation of LazyTractogram

Author: MarcCote

nibabel/streamlines/tests/test_tractogram.py

@@ -571,18 +571,21 @@ def test_lazy_tractogram_apply_affine(self):
         tractogram = DATA['lazy_tractogram'].copy()
 
         transformed_tractogram = tractogram.apply_affine(affine)
-        assert_true(transformed_tractogram is tractogram)
-        assert_array_equal(tractogram._affine_to_apply, affine)
-        assert_array_equal(tractogram.get_affine_to_rasmm(),
+        assert_true(transformed_tractogram is not tractogram)
+        assert_array_equal(tractogram._affine_to_apply, np.eye(4))
+        assert_array_equal(tractogram.get_affine_to_rasmm(), np.eye(4))
+        assert_array_equal(transformed_tractogram._affine_to_apply, affine)
+        assert_array_equal(transformed_tractogram.get_affine_to_rasmm(),
                            np.dot(np.eye(4), np.linalg.inv(affine)))
-        check_tractogram(tractogram,
+        check_tractogram(transformed_tractogram,
                          streamlines=[s*scaling for s in DATA['streamlines']],
                          data_per_streamline=DATA['data_per_streamline'],
                          data_per_point=DATA['data_per_point'])
 
         # Apply affine again and check the affine_to_rasmm.
-        transformed_tractogram = tractogram.apply_affine(affine)
-        assert_array_equal(tractogram._affine_to_apply, np.dot(affine, affine))
+        transformed_tractogram = transformed_tractogram.apply_affine(affine)
+        assert_array_equal(transformed_tractogram._affine_to_apply,
+                           np.dot(affine, affine))
         assert_array_equal(transformed_tractogram.get_affine_to_rasmm(),
                            np.dot(np.eye(4), np.dot(np.linalg.inv(affine),
                                                     np.linalg.inv(affine))))
@@ -594,21 +597,21 @@ def test_tractogram_to_world(self):
 
         # Apply the affine to the streamlines, then bring them back
         # to world space in a lazy manner.
-        tractogram.apply_affine(affine)
-        assert_array_equal(tractogram.get_affine_to_rasmm(),
+        transformed_tractogram = tractogram.apply_affine(affine)
+        assert_array_equal(transformed_tractogram.get_affine_to_rasmm(),
                            np.linalg.inv(affine))
 
-        tractogram_world = tractogram.to_world()
-        assert_true(tractogram_world is tractogram)
-        assert_array_almost_equal(tractogram.get_affine_to_rasmm(),
+        tractogram_world = transformed_tractogram.to_world()
+        assert_true(tractogram_world is not transformed_tractogram)
+        assert_array_almost_equal(tractogram_world.get_affine_to_rasmm(),
                                   np.eye(4))
-        for s1, s2 in zip(tractogram.streamlines, DATA['streamlines']):
+        for s1, s2 in zip(tractogram_world.streamlines, DATA['streamlines']):
             assert_array_almost_equal(s1, s2)
 
         # Calling to_world twice should do nothing.
-        tractogram.to_world()
-        assert_array_almost_equal(tractogram.get_affine_to_rasmm(), np.eye(4))
-        for s1, s2 in zip(tractogram.streamlines, DATA['streamlines']):
+        tractogram_world = tractogram_world.to_world()
+        assert_array_almost_equal(tractogram_world.get_affine_to_rasmm(), np.eye(4))
+        for s1, s2 in zip(tractogram_world.streamlines, DATA['streamlines']):
             assert_array_almost_equal(s1, s2)
 
     def test_lazy_tractogram_copy(self):

nibabel/streamlines/tractogram.py

@@ -320,7 +320,7 @@ def apply_affine(self, affine, lazy=False):
         ----------
         affine : ndarray of shape (4, 4)
             Transformation that will be applied to every streamline.
-        lazy_load : {False, True}, optional
+        lazy : {False, True}, optional
             If True, streamlines are *not* transformed in-place and a
             :class:`LazyTractogram` object is returned. Otherwise, streamlines
             are modified in-place.
@@ -336,8 +336,7 @@ def apply_affine(self, affine, lazy=False):
         """
         if lazy:
             lazy_tractogram = LazyTractogram.from_tractogram(self)
-            lazy_tractogram.apply_affine(affine)
-            return lazy_tractogram
+            return lazy_tractogram.apply_affine(affine)
 
         if len(self.streamlines) == 0:
             return self
@@ -361,7 +360,7 @@ def to_world(self, lazy=False):
 
         Parameters
         ----------
-        lazy_load : {False, True}, optional
+        lazy : {False, True}, optional
             If True, streamlines are *not* transformed in-place and a
             :class:`LazyTractogram` object is returned. Otherwise, streamlines
             are modified in-place.
@@ -641,7 +640,7 @@ def copy(self):
         tractogram._affine_to_apply = self._affine_to_apply.copy()
         return tractogram
 
-    def apply_affine(self, affine):
+    def apply_affine(self, affine, lazy=True):
         """ Applies an affine transformation to the streamlines.
 
         The transformation given by the `affine` matrix is applied after any
@@ -651,29 +650,46 @@ def apply_affine(self, affine):
         ----------
         affine : 2D array (4,4)
             Transformation matrix that will be applied on each streamline.
+        lazy : True, optional
+            Should always be True for :class:`LazyTractogram` object. Doing
+            otherwise will raise a ValueError.
 
         Returns
         -------
         lazy_tractogram : :class:`LazyTractogram` object
-            Reference to this instance of :class:`LazyTractogram`.
+            A copy of this :class:`LazyTractogram` instance but with a
+            transformation to be applied on the streamlines.
         """
+        if not lazy:
+            msg = "LazyTractogram only supports lazy transformations."
+            raise ValueError(msg)
+
+        tractogram = self.copy()  # New instance.
+
         # Update the affine that will be applied when returning streamlines.
-        self._affine_to_apply = np.dot(affine, self._affine_to_apply)
+        tractogram._affine_to_apply = np.dot(affine, self._affine_to_apply)
 
         # Update the affine that brings back the streamlines to RASmm.
-        self._affine_to_rasmm = np.dot(self._affine_to_rasmm,
-                                       np.linalg.inv(affine))
-        return self
+        tractogram._affine_to_rasmm = np.dot(self._affine_to_rasmm,
+                                             np.linalg.inv(affine))
+        return tractogram
 
-    def to_world(self):
+    def to_world(self, lazy=True):
         """ Brings the streamlines to world space (i.e. RAS+ and mm).
 
-        The transformation will be applied just before returning the
-        streamlines.
+        The transformation is applied after any other pending transformations
+        to the streamline points.
+
+        Parameters
+        ----------
+        lazy : True, optional
+            Should always be True for :class:`LazyTractogram` object. Doing
+            otherwise will raise a ValueError.
 
         Returns
         -------
         lazy_tractogram : :class:`LazyTractogram` object
-            Reference to this instance of :class:`LazyTractogram`.
+            A copy of this :class:`LazyTractogram` instance but with a
+            transformation to be applied on the streamlines.
         """
-        return self.apply_affine(self._affine_to_rasmm)
+        return self.apply_affine(self._affine_to_rasmm, lazy=lazy)

nibabel/streamlines/trk.py

@@ -270,6 +270,43 @@ def __init__(self, fileobj, header):
 
         self.file.write(self.header.tostring())
 
+        # `Tractogram` streamlines are in RAS+ and mm space, we will compute
+        # the affine matrix that will bring them back to 'voxelmm' as required
+        # by the TRK format.
+        affine = np.eye(4)
+
+        # Applied the inverse of the affine found in the TRK header.
+        # rasmm -> voxel
+        affine = np.dot(np.linalg.inv(self.header[Field.VOXEL_TO_RASMM]),
+                        affine)
+
+        # If the voxel order implied by the affine does not match the voxel
+        # order in the TRK header, change the orientation.
+        # voxel (affine) -> voxel (header)
+        header_ornt = asstr(self.header[Field.VOXEL_ORDER])
+        affine_ornt = "".join(aff2axcodes(self.header[Field.VOXEL_TO_RASMM]))
+        header_ornt = axcodes2ornt(header_ornt)
+        affine_ornt = axcodes2ornt(affine_ornt)
+        ornt = nib.orientations.ornt_transform(affine_ornt, header_ornt)
+        M = nib.orientations.inv_ornt_aff(ornt, self.header[Field.DIMENSIONS])
+        affine = np.dot(M, affine)
+
+        # TrackVis considers coordinate (0,0,0) to be the corner of the
+        # voxel whereas `Tractogram` streamlines assume (0,0,0) is the
+        # center of the voxel. Thus, streamlines are shifted of half a voxel.
+        offset = np.eye(4)
+        offset[:-1, -1] += 0.5
+        affine = np.dot(offset, affine)
+
+        # Finally send the streamlines in mm space.
+        # voxel -> voxelmm
+        scale = np.eye(4)
+        scale[range(3), range(3)] *= self.header[Field.VOXEL_SIZES]
+        affine = np.dot(scale, affine)
+
+        # The TRK format uses float32 as the data type for points.
+        self._affine_rasmm_to_voxmm = affine.astype(np.float32)
+
     def write(self, tractogram):
         i4_dtype = np.dtype("<i4")  # Always save in little-endian.
         f4_dtype = np.dtype("<f4")  # Always save in little-endian.
@@ -350,50 +387,17 @@ def write(self, tractogram):
 
             self.header['scalar_name'][i] = scalar_name
 
-        # `Tractogram` streamlines are in RAS+ and mm space, we will compute
-        # the affine matrix that will bring them back to 'voxelmm' as required
-        # by the TRK format.
-        affine = np.eye(4)
-
-        # Applied the inverse of the affine found in the TRK header.
-        # rasmm -> voxel
-        affine = np.dot(np.linalg.inv(self.header[Field.VOXEL_TO_RASMM]),
-                        affine)
-
-        # If the voxel order implied by the affine does not match the voxel
-        # order in the TRK header, change the orientation.
-        # voxel (affine) -> voxel (header)
-        header_ornt = asstr(self.header[Field.VOXEL_ORDER])
-        affine_ornt = "".join(aff2axcodes(self.header[Field.VOXEL_TO_RASMM]))
-        header_ornt = axcodes2ornt(header_ornt)
-        affine_ornt = axcodes2ornt(affine_ornt)
-        ornt = nib.orientations.ornt_transform(affine_ornt, header_ornt)
-        M = nib.orientations.inv_ornt_aff(ornt, self.header[Field.DIMENSIONS])
-        affine = np.dot(M, affine)
-
-        # TrackVis considers coordinate (0,0,0) to be the corner of the
-        # voxel whereas `Tractogram` streamlines assume (0,0,0) is the
-        # center of the voxel. Thus, streamlines are shifted of half a voxel.
-        offset = np.eye(4)
-        offset[:-1, -1] += 0.5
-        affine = np.dot(offset, affine)
-
-        # Finally send the streamlines in mm space.
-        # voxel -> voxelmm
-        scale = np.eye(4)
-        scale[range(3), range(3)] *= self.header[Field.VOXEL_SIZES]
-        affine = np.dot(scale, affine)
-
-        # The TRK format uses float32 as the data type for points.
-        affine = affine.astype(np.float32)
+        # Make sure streamlines are in rasmm then send them to voxmm.
+        tractogram = tractogram.to_world(lazy=True)
+        tractogram = tractogram.apply_affine(self._affine_rasmm_to_voxmm,
+                                             lazy=True)
 
         for t in tractogram:
             if any((len(d) != len(t.streamline)
                     for d in t.data_for_points.values())):
                 raise DataError("Missing scalars for some points!")
 
-            points = apply_affine(affine,
-                                  np.asarray(t.streamline, dtype=f4_dtype))
+            points = np.asarray(t.streamline, dtype=f4_dtype)
             scalars = [np.asarray(t.data_for_points[k], dtype=f4_dtype)
                        for k in data_for_points_keys]
             scalars = np.concatenate([np.ndarray((len(points), 0),
@@ -747,8 +751,9 @@ def _read():
             for name, slice_ in data_per_streamline_slice.items():
                 tractogram.data_per_streamline[name] = properties[:, slice_]
 
-        # Bring tractogram to RAS+ and mm space
-        tractogram.apply_affine(affine.astype(np.float32))
+        # Bring tractogram to RAS+ and mm space.
+        tractogram = tractogram.apply_affine(affine.astype(np.float32))
+        tractogram._affine_to_rasmm = np.eye(4)
 
         return cls(tractogram, header=hdr)