Deform streamlines

dipy/tracking/streamline.py

@@ -231,6 +231,51 @@ def center_streamlines(streamlines):
     return [s - center for s in streamlines], center
 
 
+def deform_streamlines(streamlines,
+                       deform_field,
+                       stream_to_current_grid,
+                       current_grid_to_world,
+                       stream_to_ref_grid,
+                       ref_grid_to_world):
+    """ Apply deformation field to streamlines
+
+    Parameters
+    ----------
+    streamlines : list
+        List of 2D ndarrays of shape[-1]==3
+    deform_field : 4D numpy array
+        x,y,z displacements stored in volume, shape[-1]==3
+    stream_to_current_grid : array, (4, 4)
+        transform matrix voxmm space to original grid space
+    current_grid_to_world : array (4, 4)
+        transform matrix original grid space to world coordinates
+    stream_to_ref_grid : array (4, 4)
+        transform matrix voxmm space to new grid space
+    ref_grid_to_world : array(4, 4)
+        transform matrix new grid space to world coordinates
+
+    Returns
+    -------
+    new_streamlines : list
+        List of the transformed 2D ndarrays of shape[-1]==3
+    """
+
+    if deform_field.shape[-1] != 3:
+        raise ValueError("Last dimension of deform_field needs shape==3")
+
+    stream_in_curr_grid = transform_streamlines(streamlines,
+                                                stream_to_current_grid)
+    displacements = values_from_volume(deform_field, stream_in_curr_grid)
+    stream_in_world = transform_streamlines(stream_in_curr_grid,
+                                            current_grid_to_world)
+    new_streams_in_world = list(np.add(displacements, stream_in_world))
+    new_streams_grid = transform_streamlines(new_streams_in_world,
+                                             np.linalg.inv(ref_grid_to_world))
+    new_streamlines = transform_streamlines(new_streams_grid,
+                                            np.linalg.inv(stream_to_ref_grid))
+    return new_streamlines
+
+
 def transform_streamlines(streamlines, mat):
     """ Apply affine transformation to streamlines
 

dipy/tracking/tests/test_streamline.py

@@ -10,7 +10,7 @@
 
 from nose.tools import assert_true, assert_equal, assert_almost_equal
 from numpy.testing import (assert_array_equal, assert_array_almost_equal,
-                           assert_raises, run_module_suite)
+                           assert_raises, run_module_suite, assert_allclose)
 
 from dipy.tracking.streamline import Streamlines
 import dipy.tracking.utils as ut
@@ -24,7 +24,8 @@
                                       compress_streamlines,
                                       select_by_rois,
                                       orient_by_rois,
-                                      values_from_volume)
+                                      values_from_volume,
+                                      deform_streamlines)
 
 
 streamline = np.array([[82.20181274,  91.36505890,  43.15737152],
@@ -516,6 +517,56 @@ def test_unlist_relist_streamlines():
         assert_array_equal(streamlines[i], streamlines2[i])
 
 
+def test_deform_streamlines():
+    # streamlines needs to be a list of streamlines
+    A = np.array([[1, 2, 3], [1, 2, 3.]])
+
+    # Create Random deformation field
+    deformation_field = np.random.randn(200, 200, 200, 3)
+    # shape = deformation_field.shape
+    # for i in range(shape[0]):
+    #     for j in range(shape[1]):
+    #         for k in range(shape[2]):
+    #             deformation_field[i][j][k] = np.random.randn(3)
+    # Specify stream2grid and grid2world
+    stream2grid = np.array([[np.random.randn(1)[0], 0, 0, 0],
+                            [0, np.random.randn(1)[0], 0, 0],
+                            [0, 0, np.random.randn(1)[0], 0],
+                            [0, 0, 0, 1]])
+    grid2world = np.array([[np.random.randn(1)[0], 0, 0, 0],
+                           [0, np.random.randn(1)[0], 0, 0],
+                           [0, 0, np.random.randn(1)[0], 0],
+                           [0, 0, 0, 1]])
+    stream2world = np.dot(stream2grid, grid2world)
+
+    # Deform streamlines (let two grid spaces be the same for simplicity)
+    new_streamlines = deform_streamlines(streamlines,
+                                         deformation_field,
+                                         stream2grid,
+                                         grid2world,
+                                         stream2grid,
+                                         grid2world)
+
+    # Interpolate displacements onto original streamlines
+    streamlines_in_grid = transform_streamlines(streamlines, stream2grid)
+    displacements = values_from_volume(deformation_field, streamlines_in_grid)
+
+    # Put new_streamlines into world space
+    new_streamlines_world = transform_streamlines(new_streamlines,
+                                                  stream2world)
+
+    # Subtract displacements from new_streamlines in world space
+    orig_streamlines_world = list(np.subtract(new_streamlines_world,
+                                              displacements))
+
+    # Put orig_streamlines_world into voxmm
+    orig_streamlines = transform_streamlines(orig_streamlines_world,
+                                             np.linalg.inv(stream2world))
+    # All close because of floating pt inprecision
+    for o, s in zip(orig_streamlines, streamlines):
+        assert_allclose(s, o, rtol=1e-10, atol=0)
+
+
 def test_center_and_transform():
     A = np.array([[1, 2, 3], [1, 2, 3.]])
     streamlines = [A for i in range(10)]