Repository: yeatmanlab/pyAFQ (python)

Issue:
Use dipy's I/O instead of our own
So we can avoid warnings like this one: https://travis-ci.org/yeatmanlab/pyAFQ/jobs/553808933#L65

We can use functions from here: https://github.com/nipy/dipy/blob/master/dipy/io/streamline.py to replace `utils.streamlinse.read_trk` and `utils.streamlinse.save_trk`

Pull request:
Use the DIPY I/O for streamlines, instead of our own.
Resolves #127

Patch:
diff --git a/AFQ/api.py b/AFQ/api.py
index 40020b33..5bd54e44 100644
--- a/AFQ/api.py
+++ b/AFQ/api.py
@@ -14,6 +14,7 @@ from dipy.segment.mask import median_otsu
 import dipy.data as dpd
 import dipy.tracking.utils as dtu
 import dipy.tracking.streamline as dts
+from dipy.io.streamline import save_tractogram
 
 import AFQ.data as afd
 from AFQ.dti import _fit as dti_fit
@@ -372,10 +373,10 @@ def _streamlines(row, wm_labels, odf_model="DTI", directions="det",
                                 seed_mask=wm_mask,
                                 stop_mask=wm_mask)
 
-        aus.write_trk(streamlines_file,
-                      dtu.move_streamlines(streamlines,
-                                           np.linalg.inv(dwi_img.affine)),
-                      affine=dwi_img.affine)
+        save_tractogram(streamlines_file,
+                        dtu.move_streamlines(streamlines,
+                                             np.linalg.inv(dwi_img.affine)),
+                        dwi_img.affine)
 
     return streamlines_file
 
@@ -607,12 +608,12 @@ def _export_bundles(row, wm_labels, bundle_dict, reg_template,
             idx = np.where(tg.data_per_streamline['bundle'] == uid)[0]
             this_sl = (streamlines[idx])
             fname = op.join(bundles_dir, '%s.trk' % bundle)
-            aus.write_trk(
+            save_tractogram(
                 fname,
                 dtu.move_streamlines(
                     this_sl,
                     np.linalg.inv(row['dwi_affine'])),
-                affine=row['dwi_affine'])
+                row['dwi_affine'])
 
 
 def _get_affine(fname):
diff --git a/AFQ/registration.py b/AFQ/registration.py
index c35cbc3d..51a90dce 100644
--- a/AFQ/registration.py
+++ b/AFQ/registration.py
@@ -23,6 +23,7 @@ import dipy.data as dpd
 from dipy.align.streamlinear import StreamlineLinearRegistration
 from dipy.tracking.streamline import set_number_of_points
 from dipy.tracking.utils import move_streamlines
+from dipy.io.streamline import load_tractogram
 
 import AFQ.utils.models as mut
 import AFQ.utils.streamlines as sut
@@ -45,6 +46,7 @@ def syn_registration(moving, static,
                      dim=3,
                      level_iters=[10, 10, 5],
                      sigma_diff=2.0,
+                     radius=4,
                      prealign=None):
     """Register a source image (moving) to a target image (static).
 
@@ -67,6 +69,8 @@ def syn_registration(moving, static,
         the number of iterations at each level of the Gaussian Pyramid (the
         length of the list defines the number of pyramid levels to be
         used).
+    sigma_diff, radius : float
+        Parameters for initialization of the metric.
 
     Returns
     -------
@@ -79,7 +83,8 @@ def syn_registration(moving, static,
         The vector field describing the backward warping from the target to the
         source.
     """
-    use_metric = syn_metric_dict[metric](dim, sigma_diff=sigma_diff)
+    use_metric = syn_metric_dict[metric](dim, sigma_diff=sigma_diff,
+                                         radius=radius)
 
     sdr = SymmetricDiffeomorphicRegistration(use_metric, level_iters,
                                              step_length=step_length)
@@ -423,9 +428,9 @@ def streamline_registration(moving, static, n_points=100,
     """
     # Load the streamlines, if you were given a file-name
     if isinstance(moving, str):
-        moving = sut.read_trk(moving)
+        moving = load_tractogram(moving)[0]
     if isinstance(static, str):
-        static = sut.read_trk(static)
+        static = load_tractogram(static)[0]
 
     srr = StreamlineLinearRegistration()
     srm = srr.optimize(static=set_number_of_points(static, n_points),
diff --git a/AFQ/utils/streamlines.py b/AFQ/utils/streamlines.py
index 4e55faf7..d359e9d8 100644
--- a/AFQ/utils/streamlines.py
+++ b/AFQ/utils/streamlines.py
@@ -2,32 +2,6 @@ import numpy as np
 import nibabel as nib
 from nibabel import trackvis
 import dipy.tracking.utils as dtu
-import dipy.tracking.streamline as dts
-
-
-def read_trk(fname):
-    """
-    Read from a .trk file, return streamlines and header
-
-    Parameters
-    ----------
-    fname : str
-        Full path to a trk file containing
-
-    Returns
-    -------
-    list : list of streamlines (3D coordinates)
-
-    Notes
-    -----
-    We assume that all streamlines are provided with the "rasmm" points_space.
-    That is, they have been transformed to the space reported by the affine
-    associated with the image from whence it came, and saved with this affine
-    (e.g., using `write_trk`).
-
-    """
-    streams, hdr = trackvis.read(fname, points_space="rasmm")
-    return [s[0] for s in streams]
 
 
 def write_trk(fname, streamlines, affine=None, shape=None):
diff --git a/examples/plot_tract_profile.py b/examples/plot_tract_profile.py
index 6faa49b2..21548274 100644
--- a/examples/plot_tract_profile.py
+++ b/examples/plot_tract_profile.py
@@ -15,6 +15,7 @@ import dipy.data as dpd
 from dipy.data import fetcher
 import dipy.tracking.utils as dtu
 import dipy.tracking.streamline as dts
+from dipy.io.streamline import save_tractogram
 
 import AFQ.utils.streamlines as aus
 import AFQ.data as afd
@@ -44,7 +45,7 @@ else:
 print("Tracking...")
 if not op.exists('dti_streamlines.trk'):
     streamlines = list(aft.track(dti_params['params']))
-    aus.write_trk('./dti_streamlines.trk', streamlines, affine=img.affine)
+    save_tractogram('./dti_streamlines.trk', streamlines, img.affine)
 else:
     tg = nib.streamlines.load('./dti_streamlines.trk').tractogram
     streamlines = tg.apply_affine(np.linalg.inv(img.affine)).streamlines


Test patch:
diff --git a/AFQ/tests/test_api.py b/AFQ/tests/test_api.py
index 27f5109a..e28d929a 100644
--- a/AFQ/tests/test_api.py
+++ b/AFQ/tests/test_api.py
@@ -16,6 +16,7 @@ import dipy.tracking.utils as dtu
 import dipy.tracking.streamline as dts
 import dipy.data as dpd
 from dipy.data import fetcher
+from dipy.io.streamline import save_tractogram
 
 from AFQ import api
 import AFQ.data as afd
@@ -114,7 +115,7 @@ def test_AFQ_data_planes():
 
     sl_file = op.join(myafq.data_frame.results_dir[0],
                       'sub-01_sess-01_dwiDTI_det_streamlines.trk')
-    aus.write_trk(sl_file, streamlines, affine=myafq.dwi_affine[0])
+    save_tractogram(sl_file, streamlines, myafq.dwi_affine[0])
 
     mapping_file = op.join(myafq.data_frame.results_dir[0],
                                 'sub-01_sess-01_dwi_mapping.nii.gz')
diff --git a/AFQ/tests/test_registration.py b/AFQ/tests/test_registration.py
index 11ab3a9b..9a5c6640 100644
--- a/AFQ/tests/test_registration.py
+++ b/AFQ/tests/test_registration.py
@@ -15,8 +15,8 @@ from AFQ.registration import (syn_registration, register_series, register_dwi,
                               read_mapping, syn_register_dwi, DiffeomorphicMap)
 
 from dipy.tracking.utils import move_streamlines
+from dipy.io.streamline import save_tractogram
 
-from AFQ.utils.streamlines import write_trk
 
 MNI_T2 = dpd.read_mni_template()
 hardi_img, gtab = dpd.read_stanford_hardi()
@@ -45,8 +45,9 @@ def test_syn_registration():
                                                   step_length=0.1,
                                                   metric='CC',
                                                   dim=3,
-                                                  level_iters=[10, 10, 5],
+                                                  level_iters=[5, 5, 5],
                                                   sigma_diff=2.0,
+                                                  radius=1,
                                                   prealign=None)
 
         npt.assert_equal(warped_moving.shape, subset_t2.shape)
@@ -67,7 +68,9 @@ def test_syn_registration():
 
 
 def test_syn_register_dwi():
-    warped_b0, mapping = syn_register_dwi(subset_dwi_data, gtab, template=subset_t2_img)
+    warped_b0, mapping = syn_register_dwi(subset_dwi_data, gtab,
+                                          template=subset_t2_img,
+                                          radius=1)
     npt.assert_equal(isinstance(mapping, DiffeomorphicMap), True)
     npt.assert_equal(warped_b0.shape, subset_t2_img.shape)
 
@@ -127,15 +130,15 @@ def test_streamline_registration():
             fname2 = op.join(tmpdir, 'sl2.trk')
             if use_aff is not None:
                 # Move the streamlines to this other space, and report it:
-                write_trk(fname1,
-                          move_streamlines(sl1, np.linalg.inv(use_aff)),
-                          use_aff)
-                write_trk(fname2,
-                          move_streamlines(sl2, np.linalg.inv(use_aff)),
-                          use_aff)
+                save_tractogram(fname1,
+                                move_streamlines(sl1, np.linalg.inv(use_aff)),
+                                use_aff)
+                save_tractogram(fname2,
+                                move_streamlines(sl2, np.linalg.inv(use_aff)),
+                                use_aff)
             else:
-                write_trk(fname1, sl1)
-                write_trk(fname2, sl2)
+                save_tractogram(fname1, sl1, np.eye(4))
+                save_tractogram(fname2, sl2, np.eye(4))
 
             aligned, matrix = streamline_registration(fname2, fname1)
             npt.assert_almost_equal(aligned[0], sl1[0], decimal=5)
diff --git a/AFQ/utils/tests/test_streamlines.py b/AFQ/utils/tests/test_streamlines.py
index ce637cad..59013ab8 100644
--- a/AFQ/utils/tests/test_streamlines.py
+++ b/AFQ/utils/tests/test_streamlines.py
@@ -8,30 +8,6 @@ import dipy.tracking.utils as dtu
 import dipy.tracking.streamline as dts
 
 
-def test_read_write_trk():
-    sl = [np.array([[0, 0, 0], [0, 0, 0.5], [0, 0, 1], [0, 0, 1.5]]),
-          np.array([[0, 0, 0], [0, 0.5, 0.5], [0, 1, 1]])]
-
-    with nbtmp.InTemporaryDirectory() as tmpdir:
-        fname = op.join(tmpdir, 'sl.trk')
-        aus.write_trk(fname, sl)
-        new_sl = aus.read_trk(fname)
-        npt.assert_equal(list(new_sl), sl)
-
-        # What happens if this set of streamlines has some funky affine
-        # associated with it?
-        aff = np.eye(4) * np.random.rand()
-        aff[:3, 3] = np.array([1, 2, 3])
-        aff[3, 3] = 1
-        # We move the streamlines, and report the inverse of the affine:
-        aus.write_trk(fname, dtu.move_streamlines(sl, aff),
-                      affine=np.linalg.inv(aff))
-        # When we read this, we get back what we put in:
-        new_sl = aus.read_trk(fname)
-        # Compare each streamline:
-        for new, old in zip(new_sl, sl):
-            npt.assert_almost_equal(new, old, decimal=4)
-
 def test_bundles_to_tgram():
     bundles = {'b1': dts.Streamlines([np.array([[0, 0, 0],
                                                 [0, 0, 0.5],