Repository: yeatmanlab/pyAFQ
Pull request #335 (2020-08-07 17:15:13)

Adds ParticleFilteringTractography.
Addresses #334. For now, I have done the naive thing of assuming that these objects are duck-typing. From my local testing, it looks like it's going to be a bit more complicated than that. We might need to do some more work in the `_track` function to be able to use both objects.

pep8speaks: Hello @arokem! Thanks for opening this PR. We checked the lines you've touched for [PEP 8](https://www.python.org/dev/peps/pep-0008) issues, and found:

* In the file [`AFQ/tests/test_tractography.py`](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py):

> [Line 37:25](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L37): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 38:25](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L38): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 39:25](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L39): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 40:25](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L40): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 41:25](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L41): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 42:25](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L42): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 43:25](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L43): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 44:25](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L44): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 45:25](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L45): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 55:21](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L55): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 56:21](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L56): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 57:21](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L57): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 58:21](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L58): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 59:21](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L59): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 60:21](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L60): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent
> [Line 61:21](https://github.com/yeatmanlab/pyAFQ/blob/18e964a8deb70707612ea96e7c595c0dbe90c886/AFQ/tests/test_tractography.py#L61): [E128](https://duckduckgo.com/?q=pep8%20E128) continuation line under-indented for visual indent




arokem: This does not currently work, because PFT requires using an `AnatomicalStoppingCriterion` class as its stopping criterion. Because this interacts with work happening in #330, I will wait to merge that one and then revisit. For now, I am marking this one for 0.5, so we don't have to worry about it for the 0.4 milestone.  
arokem: I have removed the WIP label. I think that this is ready for review. 

Related to #370 : for PFT, this now requires a string as the `stop_mask`: either "ACT" or "CMC" for the strategy to use in determining where to stop and a tuple as the `stop_threshold` that is to be interpreted as three volumes: `(pve_wm, pve_gm, pve_csf)`. 

And a `tracker` input is added, which defaults to `"local"`. When `tracker="pft"`, `stop_mask` needs to be a string and `stop_threshold` needs to be a tuple. 
arokem: In describing this, I can already see that better error-handling would be good.
36000: does this mean we should get rid of the wm_mask parameter altogether?
36000: Where does the user get pve_wm, pve_gm, pve_csf data from? It seems to me like there should be a special PFTmask where you give it the bids suffix for each of these files, and get_mask returns these in its tuple instead of what it currently returns. 
And then maybe stop_threshold should be "ACT" or "CMC" instead.
36000: So I've been thinking about this, and this is what I think the ui should look like (for example):
seed_mask = (afm.MaskFile(f"WM_probseg"), afm.MaskFile(f"GM_probseg"), afm.MaskFile(f"CSF_probseg"))
stop_threshold = "ACT"

Then the api finds the relevant mask for each subject.
arokem: Would love to have #380 merged before moving on here, so I am switching to that one for now. 
36000: I have been thinking about this, and I think all of my suggestions should be implemented and handled in the api. the track function should not use mask objects, only the API. This can be another PR.

So I think the only outstanding issue on this PR is how to handle average_voxel_size
arokem: Agreed! I'll try to get that done today. 

Patch:
diff --git a/AFQ/_fixes.py b/AFQ/_fixes.py
index 98457d2a..1003dfe9 100644
--- a/AFQ/_fixes.py
+++ b/AFQ/_fixes.py
@@ -4,7 +4,8 @@ from scipy.special import lpmv, gammaln
 
 from tqdm import tqdm
 from dipy.align import Bunch
-from dipy.tracking.local_tracking import LocalTracking
+from dipy.tracking.local_tracking import (LocalTracking,
+                                          ParticleFilteringTracking)
 from dipy.align.imaffine import AffineMap
 import random
 
@@ -42,59 +43,66 @@ def spherical_harmonics(m, n, theta, phi):
 TissueTypes = Bunch(OUTSIDEIMAGE=-1, INVALIDPOINT=0, TRACKPOINT=1, ENDPOINT=2)
 
 
+def _verbose_generate_streamlines(self):
+    """A streamline generator"""
+    # Get inverse transform (lin/offset) for seeds
+    inv_A = np.linalg.inv(self.affine)
+    lin = inv_A[:3, :3]
+    offset = inv_A[:3, 3]
+
+    F = np.empty((self.max_length + 1, 3), dtype=float)
+    B = F.copy()
+    for s in tqdm(self.seeds):
+        s = np.dot(lin, s) + offset
+        # Set the random seed in numpy and random
+        if self.random_seed is not None:
+            s_random_seed = hash(np.abs((np.sum(s)) + self.random_seed)) \
+                % (2**32 - 1)
+            random.seed(s_random_seed)
+            np.random.seed(s_random_seed)
+        directions = self.direction_getter.initial_direction(s)
+        if directions.size == 0 and self.return_all:
+            continue
+        directions = directions[:self.max_cross]
+        for first_step in directions:
+            stepsF, tissue_class = self._tracker(s, first_step, F)
+            if not (self.return_all
+                    or tissue_class == TissueTypes.ENDPOINT
+                    or tissue_class == TissueTypes.OUTSIDEIMAGE):
+                continue
+            first_step = -first_step
+            stepsB, tissue_class = self._tracker(s, first_step, B)
+            if not (self.return_all
+                    or tissue_class == TissueTypes.ENDPOINT
+                    or tissue_class == TissueTypes.OUTSIDEIMAGE):
+                continue
+            if stepsB == 1:
+                streamline = F[:stepsF].copy()
+            else:
+                parts = (B[stepsB - 1:0:-1], F[:stepsF])
+                streamline = np.concatenate(parts, axis=0)
+
+            len_sl = len(streamline)
+            if len_sl < self.min_length * self.step_size \
+                    or len_sl > self.max_length * self.step_size:
+                continue
+            else:
+                yield streamline
+
+
 class VerboseLocalTracking(LocalTracking):
     def __init__(self, *args, min_length=10, max_length=1000, **kwargs):
         super().__init__(*args, **kwargs)
         self.min_length = min_length
         self.max_length = max_length
+    _generate_streamlines = _verbose_generate_streamlines
 
-    def _generate_streamlines(self):
-        """A streamline generator"""
-
-        # Get inverse transform (lin/offset) for seeds
-        inv_A = np.linalg.inv(self.affine)
-        lin = inv_A[:3, :3]
-        offset = inv_A[:3, 3]
-
-        F = np.empty((self.max_length + 1, 3), dtype=float)
-        B = F.copy()
-        for s in tqdm(self.seeds):
-            s = np.dot(lin, s) + offset
-            # Set the random seed in numpy and random
-            if self.random_seed is not None:
-                s_random_seed = hash(np.abs((np.sum(s)) + self.random_seed)) \
-                    % (2**32 - 1)
-                random.seed(s_random_seed)
-                np.random.seed(s_random_seed)
-            directions = self.direction_getter.initial_direction(s)
-            if directions.size == 0 and self.return_all:
-                # only the seed position
-                yield [s]
-            directions = directions[:self.max_cross]
-            for first_step in directions:
-                stepsF, tissue_class = self._tracker(s, first_step, F)
-                if not (self.return_all
-                        or tissue_class == TissueTypes.ENDPOINT
-                        or tissue_class == TissueTypes.OUTSIDEIMAGE):
-                    continue
-                first_step = -first_step
-                stepsB, tissue_class = self._tracker(s, first_step, B)
-                if not (self.return_all
-                        or tissue_class == TissueTypes.ENDPOINT
-                        or tissue_class == TissueTypes.OUTSIDEIMAGE):
-                    continue
-                if stepsB == 1:
-                    streamline = F[:stepsF].copy()
-                else:
-                    parts = (B[stepsB - 1:0:-1], F[:stepsF])
-                    streamline = np.concatenate(parts, axis=0)
-
-                len_sl = len(streamline)
-                if len_sl < self.min_length * self.step_size \
-                        or len_sl > self.max_length * self.step_size:
-                    continue
-                else:
-                    yield streamline
+class VerboseParticleFilteringTracking(ParticleFilteringTracking):
+    def __init__(self, *args, min_length=10, max_length=1000, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.min_length = min_length
+        self.max_length = max_length
+    _generate_streamlines = _verbose_generate_streamlines
 
 
 def in_place_norm(vec, axis=-1, keepdims=False, delvec=True):
diff --git a/AFQ/tractography.py b/AFQ/tractography.py
index 249d9f98..79ea94d2 100644
--- a/AFQ/tractography.py
+++ b/AFQ/tractography.py
@@ -8,15 +8,20 @@ from dipy.direction import (DeterministicMaximumDirectionGetter,
                             ProbabilisticDirectionGetter)
 import dipy.tracking.utils as dtu
 from dipy.io.stateful_tractogram import StatefulTractogram, Space
-from dipy.tracking.stopping_criterion import ThresholdStoppingCriterion
+from dipy.tracking.stopping_criterion import (ThresholdStoppingCriterion,
+                                              CmcStoppingCriterion,
+                                              ActStoppingCriterion)
 
-from AFQ._fixes import VerboseLocalTracking, tensor_odf
+
+from AFQ._fixes import (VerboseLocalTracking, VerboseParticleFilteringTracking,
+                        tensor_odf)
 
 
 def track(params_file, directions="det", max_angle=30., sphere=None,
           seed_mask=None, seed_threshold=0, n_seeds=1, random_seeds=False,
           rng_seed=None, stop_mask=None, stop_threshold=0, step_size=0.5,
-          min_length=10, max_length=1000, odf_model="DTI"):
+          min_length=10, max_length=1000, odf_model="DTI",
+          tracker="local"):
     """
     Tractography
 
@@ -48,17 +53,26 @@ def track(params_file, directions="det", max_angle=30., sphere=None,
         to generate within the mask.
     random_seeds : bool
         Whether to generate a total of n_seeds random seeds in the mask.
-        Default: XXX.
+        Default: False.
     rng_seed : int
         random seed used to generate random seeds if random_seeds is
         set to True. Default: None
-    stop_mask : array, optional.
-        A float or binary mask that determines a stopping criterion (e.g. FA).
-        Default to no stopping (all ones).
-    stop_threshold : float, optional.
-        A value of the stop_mask below which tracking is terminated. Default to
+    stop_mask : array or str, optional.
+        If array: A float or binary mask that determines a stopping criterion
+        (e.g. FA).
+        If str: "CMC" for Continuous Map Criterion [Girard2014]_.
+                "ACT" for Anatomically-constrained tractography [Smith2012]_.
+        Defaults to no stopping (all ones). A string is required if
+        the tracker is set to "pft".
+    stop_threshold : float or tuple, optional.
+        If float, this a value of the stop_mask below which tracking is
+        terminated (and stop_mask has to be an array). Defaults to
         0 (this means that if no stop_mask is passed, we will stop only at
-        the edge of the image)
+        the edge of the image). If this is a tuple, it contains a sequence
+        that is interpreted as: (pve_wm, pve_gm, pve_csf), each item of
+        which is either a string (full path) or a nibabel img to be used
+        in particle filtering tractography. A tuple is required if tracker
+        is set to "pft".
     step_size : float, optional.
         The size (in mm) of a step of tractography. Default: 1.0
     min_length: int, optional
@@ -67,9 +81,20 @@ def track(params_file, directions="det", max_angle=30., sphere=None,
         The miminal length (mm) in a streamline. Default: 250
     odf_model : str, optional
         One of {"DTI", "CSD", "DKI"}. Defaults to use "DTI"
+    tracker : str, optional
+        Which strategy to use in tracking. This can be the standard local
+        tracking ("local") or Particle Filtering Tracking ([Girard2014]_).
+        One of {"local", "pft"}. Default: "local"
+
     Returns
     -------
     list of streamlines ()
+
+    References
+    ----------
+    .. [Girard2014] Girard, G., Whittingstall, K., Deriche, R., &
+        Descoteaux, M. Towards quantitative connectivity analysis: reducing
+        tractography biases. NeuroImage, 98, 266-278, 2014.
     """
     logger = logging.getLogger('AFQ.tractography')
 
@@ -117,38 +142,74 @@ def track(params_file, directions="det", max_angle=30., sphere=None,
     elif odf_model == "CSD":
         dg = dg.from_shcoeff(model_params, max_angle=max_angle, sphere=sphere)
 
-    if stop_mask is None:
-        stop_mask = np.ones(params_img.shape[:3])
+    if tracker == "local":
+        if stop_mask is None:
+            stop_mask = np.ones(params_img.shape[:3])
 
-    if stop_mask.dtype == 'bool':
-        stopping_criterion = ThresholdStoppingCriterion(stop_mask,
-                                                        0.5)
-    else:
-        stopping_criterion = ThresholdStoppingCriterion(stop_mask,
-                                                        stop_threshold)
+        if stop_mask.dtype == 'bool':
+            stopping_criterion = ThresholdStoppingCriterion(stop_mask,
+                                                            0.5)
+        else:
+            stopping_criterion = ThresholdStoppingCriterion(stop_mask,
+                                                            stop_threshold)
+
+        my_tracker = VerboseLocalTracking
+
+    elif tracker == "pft":
+        if not isinstance(stop_mask, str):
+            raise RuntimeError(
+            "You are using PFT tracking, but did not provide a string 'stop_mask' input. ",
+            "Possible inputs are: 'CMC' or 'ACT'")
+        if not isinstance(stop_threshold, tuple):
+            raise RuntimeError(
+            "You are using PFT tracking, but did not provide a tuple for `stop_threshold`",
+            "input. Expected a (pve_wm, pve_gm, pve_csf) tuple.")
+        pves = []
+        vox_sizes = []
+        for ii, pve in enumerate(stop_threshold):
+            if isinstance(pve, str):
+                pve = nib.load(pve)
+            pves.append(pve.get_fdata())
+            vox_sizes.append(np.mean(pve.header.get_zooms()[:3]))
+        average_voxel_size = np.mean(vox_sizes)
+        pve_wm_data, pve_gm_data, pve_csf_data = pves
+        my_tracker = VerboseParticleFilteringTracking
+        if stop_mask == "CMC":
+            stopping_criterion = CmcStoppingCriterion.from_pve(
+                pve_wm_data,
+                pve_gm_data,
+                pve_csf_data,
+                step_size=step_size,
+                average_voxel_size=average_voxel_size)
+        elif stop_mask == "ACT":
+            stopping_criterion = ActStoppingCriterion.from_pve(
+                pve_wm_data,
+                pve_gm_data,
+                pve_csf_data)
 
     logger.info("Tracking...")
 
-    return _local_tracking(seeds, dg, stopping_criterion, params_img,
-                           step_size=step_size, min_length=min_length,
-                           max_length=max_length, random_seed=rng_seed)
+    return _tracking(my_tracker, seeds, dg, stopping_criterion, params_img,
+                     step_size=step_size, min_length=min_length,
+                     max_length=max_length, random_seed=rng_seed)
 
-
-def _local_tracking(seeds, dg, stopping_criterion, params_img,
-                    step_size=0.5, min_length=10, max_length=1000,
-                    random_seed=None):
+def _tracking(tracker, seeds, dg, stopping_criterion, params_img,
+              step_size=0.5, min_length=10, max_length=1000,
+              random_seed=None):
     """
     Helper function
     """
     if len(seeds.shape) == 1:
         seeds = seeds[None, ...]
-    tracker = VerboseLocalTracking(dg,
-                                   stopping_criterion,
-                                   seeds,
-                                   params_img.affine,
-                                   step_size=step_size,
-                                   min_length=min_length,
-                                   max_length=max_length,
-                                   random_seed=random_seed)
+
+    tracker = tracker(
+        dg,
+        stopping_criterion,
+        seeds,
+        params_img.affine,
+        step_size=step_size,
+        min_length=min_length,
+        max_length=max_length,
+        random_seed=random_seed)
 
     return StatefulTractogram(tracker, params_img, Space.RASMM)


Test patch:
diff --git a/AFQ/tests/test_tractography.py b/AFQ/tests/test_tractography.py
index 1d473780..fbbbae3a 100644
--- a/AFQ/tests/test_tractography.py
+++ b/AFQ/tests/test_tractography.py
@@ -1,7 +1,9 @@
 import os.path as op
 import numpy as np
 import numpy.testing as npt
+import pytest
 
+import nibabel as nib
 import nibabel.tmpdirs as nbtmp
 
 from AFQ.models.csd import fit_csd
@@ -25,35 +27,100 @@ min_length = 20
 step_size = 0.5
 
 
-def test_csd_tracking():
+def test_csd_local_tracking():
     for sh_order in [4, 8, 10]:
         fname = fit_csd(fdata, fbval, fbvec,
                         response=((0.0015, 0.0003, 0.0003), 100),
-                        sh_order=8, lambda_=1, tau=0.1, mask=None,
+                        sh_order=sh_order, lambda_=1, tau=0.1, mask=None,
                         out_dir=tmpdir.name)
         for directions in ["det", "prob"]:
-            sl = track(fname, directions,
-                       odf_model="CSD",
-                       max_angle=30.,
-                       sphere=None,
-                       seed_mask=None,
-                       n_seeds=seeds,
-                       stop_mask=None,
-                       step_size=step_size,
-                       min_length=min_length).streamlines
+            sl = track(
+                fname,
+                directions,
+                odf_model="CSD",
+                max_angle=30.,
+                sphere=None,
+                seed_mask=None,
+                n_seeds=seeds,
+                stop_mask=None,
+                step_size=step_size,
+                min_length=min_length,
+                tracker="local").streamlines
 
             npt.assert_(len(sl[0]) >= step_size * min_length)
 
 
-def test_dti_tracking():
+def test_dti_local_tracking():
     fdict = fit_dti(fdata, fbval, fbvec)
     for directions in ["det", "prob"]:
-        sl = track(fdict['params'],
-                   directions,
-                   max_angle=30.,
-                   sphere=None,
-                   seed_mask=None,
-                   n_seeds=1,
-                   step_size=step_size,
-                   min_length=min_length).streamlines
+        sl = track(
+            fdict['params'],
+            directions,
+            max_angle=30.,
+            sphere=None,
+            seed_mask=None,
+            n_seeds=1,
+            step_size=step_size,
+            min_length=min_length,
+            tracker="local").streamlines
         npt.assert_(len(sl[0]) >= min_length * step_size)
+
+
+def test_pft_tracking():
+    for fname in [fit_dti(fdata, fbval, fbvec)['params'],
+                  fit_csd(fdata, fbval, fbvec,
+                        response=((0.0015, 0.0003, 0.0003), 100),
+                        sh_order=8, lambda_=1, tau=0.1, mask=None,
+                        out_dir=tmpdir.name)]:
+        img = nib.load(fdata)
+        data_shape  = img.shape
+        data_affine = img.affine
+        pve_wm_data = nib.Nifti1Image(np.ones(data_shape[:3]), img.affine)
+        pve_gm_data = nib.Nifti1Image(np.zeros(data_shape[:3]), img.affine)
+        pve_csf_data = nib.Nifti1Image(np.zeros(data_shape[:3]), img.affine)
+        stop_threshold = (pve_wm_data, pve_gm_data, pve_csf_data)
+
+        for directions in ["det", "prob"]:
+            for stop_mask in ["ACT", "CMC"]:
+                sl = track(
+                    fname,
+                    directions,
+                    max_angle=30.,
+                    sphere=None,
+                    seed_mask=None,
+                    stop_mask=stop_mask,
+                    stop_threshold=stop_threshold,
+                    n_seeds=1,
+                    step_size=step_size,
+                    min_length=min_length,
+                    tracker="pft").streamlines
+                npt.assert_(len(sl[0]) >= min_length * step_size)
+
+    # Test error handling:
+    with pytest.raises(RuntimeError):
+        track(
+            fname,
+            directions,
+            max_angle=30.,
+            sphere=None,
+            seed_mask=None,
+            stop_mask=None,  # Stop mask needs to be a string!
+            stop_threshold=stop_threshold,
+            n_seeds=1,
+            step_size=step_size,
+            min_length=min_length,
+            tracker="pft")
+
+    with pytest.raises(RuntimeError):
+        track(
+            fname,
+            directions,
+            max_angle=30.,
+            sphere=None,
+            seed_mask=None,
+            stop_mask=stop_mask,
+            stop_threshold=0,  # Stop threshold needs to be a tuple!
+            n_seeds=1,
+            step_size=step_size,
+            min_length=min_length,
+            tracker="pft")
\ No newline at end of file