Added Tractogram.to_world and LazyTractogram.to_world methods

Author: MarcCote

nibabel/streamlines/tests/test_tractogram.py

@@ -389,34 +389,73 @@ def test_tractogram_apply_affine(self):
         # Apply the affine to the streamline in a lazy manner.
         transformed_tractogram = tractogram.apply_affine(affine, lazy=True)
         assert_true(type(transformed_tractogram) is LazyTractogram)
-        assert_true(check_iteration(transformed_tractogram))
-        assert_equal(len(transformed_tractogram), len(DATA['streamlines']))
-        for s1, s2 in zip(transformed_tractogram.streamlines,
-                          DATA['streamlines']):
-            assert_array_almost_equal(s1, s2*scaling)
-
-        for s1, s2 in zip(transformed_tractogram.streamlines,
-                          tractogram.streamlines):
-            assert_array_almost_equal(s1, s2*scaling)
-
+        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'])
         assert_array_equal(transformed_tractogram.get_affine_to_rasmm(),
                            np.dot(np.eye(4), np.linalg.inv(affine)))
+        # Make sure streamlines of the original tractogram have not been modified.
+        assert_arrays_equal(tractogram.streamlines, DATA['streamlines'])
 
         # Apply the affine to the streamlines in-place.
         transformed_tractogram = tractogram.apply_affine(affine)
         assert_true(transformed_tractogram is tractogram)
-        assert_true(check_iteration(transformed_tractogram))
-        assert_equal(len(transformed_tractogram), len(DATA['streamlines']))
-        for s1, s2 in zip(transformed_tractogram.streamlines,
-                          DATA['streamlines']):
-            assert_array_almost_equal(s1, s2*scaling)
+        check_tractogram(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(transformed_tractogram.get_affine_to_rasmm(),
                            np.dot(np.eye(4), np.dot(np.linalg.inv(affine),
                                                     np.linalg.inv(affine))))
 
+        # Check that applying an affine and its inverse give us back the
+        # original streamlines.
+        tractogram = DATA['tractogram'].copy()
+        affine = np.random.RandomState(1234).randn(4, 4)
+        affine[-1] = [0, 0, 0, 1]  # Remove perspective projection.
+
+        tractogram.apply_affine(affine)
+        tractogram.apply_affine(np.linalg.inv(affine))
+        assert_array_almost_equal(tractogram.get_affine_to_rasmm(), np.eye(4))
+        for s1, s2 in zip(tractogram.streamlines, DATA['streamlines']):
+            assert_array_almost_equal(s1, s2)
+
+    def test_tractogram_to_world(self):
+        tractogram = DATA['tractogram'].copy()
+        affine = np.random.RandomState(1234).randn(4, 4)
+        affine[-1] = [0, 0, 0, 1]  # Remove perspective projection.
+
+        # Apply the affine to the streamlines, then bring them back
+        # to world space in a lazy manner.
+        transformed_tractogram = tractogram.apply_affine(affine)
+        assert_array_equal(transformed_tractogram.get_affine_to_rasmm(),
+                           np.linalg.inv(affine))
+
+        tractogram_world = transformed_tractogram.to_world(lazy=True)
+        assert_true(type(tractogram_world) is LazyTractogram)
+        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)
+
+        # Bring them back streamlines to world space in a in-place manner.
+        tractogram_world = transformed_tractogram.to_world()
+        assert_true(tractogram_world is tractogram)
+        assert_array_almost_equal(tractogram.get_affine_to_rasmm(), np.eye(4))
+        for s1, s2 in zip(tractogram.streamlines, DATA['streamlines']):
+            assert_array_almost_equal(s1, s2)
+
+        # Calling to_world twice should do nothing.
+        tractogram_world2 = transformed_tractogram.to_world()
+        assert_true(tractogram_world2 is tractogram)
+        assert_array_almost_equal(tractogram.get_affine_to_rasmm(), np.eye(4))
+        for s1, s2 in zip(tractogram.streamlines, DATA['streamlines']):
+            assert_array_almost_equal(s1, s2)
+
 
 class TestLazyTractogram(unittest.TestCase):
 
@@ -531,18 +570,47 @@ def test_lazy_tractogram_apply_affine(self):
 
         tractogram = DATA['lazy_tractogram'].copy()
 
-        tractogram.apply_affine(affine)
-        assert_true(check_iteration(tractogram))
-        assert_equal(len(tractogram), len(DATA['streamlines']))
-        for s1, s2 in zip(tractogram.streamlines, DATA['streamlines']):
-            assert_array_almost_equal(s1, s2*scaling)
+        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(),
+                           np.dot(np.eye(4), np.linalg.inv(affine)))
+        check_tractogram(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))
         assert_array_equal(transformed_tractogram.get_affine_to_rasmm(),
                            np.dot(np.eye(4), np.dot(np.linalg.inv(affine),
                                                     np.linalg.inv(affine))))
 
+    def test_tractogram_to_world(self):
+        tractogram = DATA['lazy_tractogram'].copy()
+        affine = np.random.RandomState(1234).randn(4, 4)
+        affine[-1] = [0, 0, 0, 1]  # Remove perspective projection.
+
+        # 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(),
+                           np.linalg.inv(affine))
+
+        tractogram_world = tractogram.to_world()
+        assert_true(tractogram_world is tractogram)
+        assert_array_almost_equal(tractogram.get_affine_to_rasmm(),
+                                  np.eye(4))
+        for s1, s2 in zip(tractogram.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']):
+            assert_array_almost_equal(s1, s2)
+
     def test_lazy_tractogram_copy(self):
         # Create a copy of the lazy tractogram.
         tractogram = DATA['lazy_tractogram'].copy()

nibabel/streamlines/tractogram.py

@@ -198,6 +198,12 @@ class Tractogram(object):
     Tractogram objects have three main properties: `streamlines`,
     `data_per_streamline` and `data_per_point`.
 
+    Streamlines of a tractogram can be in any coordinate system of your
+    choice as long as you provide the correct `affine_to_rasmm` matrix, at
+    construction time, that brings the streamlines back to *RAS+*, *mm* space,
+    where the coordinates (0,0,0) corresponds to the center of the voxel
+    (opposed to a corner).
+
     """
     def __init__(self, streamlines=None,
                  data_per_streamline=None,
@@ -243,9 +249,6 @@ def data_per_streamline(self):
 
     @data_per_streamline.setter
     def data_per_streamline(self, value):
-        # if is_lazy_dict(value):
-        # self._data_per_streamline = DataPerStreamlineDict(self, **value.items())
-        # else:
         self._data_per_streamline = DataPerStreamlineDict(self, value)
 
     @property
@@ -254,9 +257,6 @@ def data_per_point(self):
 
     @data_per_point.setter
     def data_per_point(self, value):
-        # if is_lazy_dict(value):
-        #     self._data_per_point = DataPerPointDict(self, **value.items())
-        # else:
         self._data_per_point = DataPerPointDict(self, value)
 
     def get_affine_to_rasmm(self):
@@ -333,6 +333,28 @@ def apply_affine(self, affine, lazy=False):
 
         return self
 
+    def to_world(self, lazy=False):
+        """ Brings the streamlines to world space (i.e. RAS+ and mm).
+
+        If `lazy` is not specified, this is performed *in-place*.
+
+        Parameters
+        ----------
+        lazy_load : {False, True}, optional
+            If True, streamlines are *not* transformed in-place and a
+            :class:`LazyTractogram` object is returned. Otherwise, streamlines
+            are modified in-place.
+
+        Returns
+        -------
+        tractogram : :class:`Tractogram` or :class:`LazyTractogram` object
+            Tractogram where the streamlines have been sent to world space.
+            If the `lazy` option is true, it returns a :class:`LazyTractogram`
+            object, otherwise it returns a reference to this
+            :class:`Tractogram` object with updated streamlines.
+        """
+        return self.apply_affine(self._affine_to_rasmm, lazy=lazy)
+
 
 class LazyTractogram(Tractogram):
     """ Class containing information about streamlines.
@@ -394,17 +416,21 @@ def from_tractogram(cls, tractogram):
         lazy_tractogram : :class:`LazyTractogram` object
             New lazy tractogram.
         """
-        data_per_streamline = {}
-        for key, value in tractogram.data_per_streamline.items():
-            data_per_streamline[key] = lambda: value
+        lazy_tractogram = cls(lambda: tractogram.streamlines.copy())
 
-        data_per_point = {}
-        for key, value in tractogram.data_per_point.items():
-                data_per_point[key] = lambda: value
+        # Set data_per_streamline using data_func
+        def _gen(key):
+            return lambda: iter(tractogram.data_per_streamline[key])
+
+        for k in tractogram.data_per_streamline:
+            lazy_tractogram._data_per_streamline[k] = _gen(k)
 
-        lazy_tractogram = cls(lambda: tractogram.streamlines.copy(),
-                              data_per_streamline,
-                              data_per_point)
+        # Set data_per_point using data_func
+        def _gen(key):
+            return lambda: iter(tractogram.data_per_point[key])
+
+        for k in tractogram.data_per_point:
+            lazy_tractogram._data_per_point[k] = _gen(k)
 
         lazy_tractogram._nb_streamlines = len(tractogram)
         lazy_tractogram._affine_to_rasmm = tractogram.get_affine_to_rasmm()
@@ -584,3 +610,16 @@ def apply_affine(self, affine):
         self._affine_to_rasmm = np.dot(self._affine_to_rasmm,
                                        np.linalg.inv(affine))
         return self
+
+    def to_world(self):
+        """ Brings the streamlines to world space (i.e. RAS+ and mm).
+
+        The transformation will be applied just before returning the
+        streamlines.
+
+        Returns
+        -------
+        lazy_tractogram : :class:`LazyTractogram` object
+            Reference to this instance of :class:`LazyTractogram`.
+        """
+        return self.apply_affine(self._affine_to_rasmm)