[ENH] Update DiffusionGradientTable interface to support vector reorientation

Author: dPys

dmriprep/interfaces/vectors.py

@@ -97,3 +97,71 @@ def _undefined(objekt, name, default=None):
     if not isdefined(value):
         return default
     return value
+
+
+class _ReorientVectorsInputSpec(BaseInterfaceInputSpec):
+    dwi_file = File(exists=True)
+    rasb_file = File(exists=True)
+    affines = traits.List()
+    b0_threshold = traits.Float(B0_THRESHOLD, usedefault=True)
+
+
+class _ReorientVectorsOutputSpec(TraitedSpec):
+    out_rasb = File(exists=True)
+
+
+class ReorientVectors(SimpleInterface):
+    """
+    Reorient Vectors
+
+    Example
+    -------
+    >>> os.chdir(tmpdir)
+    >>> oldrasb = str(data_dir / 'dwi.tsv')
+    >>> oldrasb_mat = np.loadtxt(str(data_dir / 'dwi.tsv'), skiprows=1)
+    >>> # The simple case: all affines are identity
+    >>> affine_list = np.zeros((len(oldrasb_mat[:, 3][oldrasb_mat[:, 3] != 0]), 4, 4))
+    >>> for i in range(4):
+    >>>     affine_list[:, i, i] = 1
+    >>>     reor_vecs = ReorientVectors()
+    >>> reor_vecs = ReorientVectors()
+    >>> reor_vecs.inputs.affines = affine_list
+    >>> reor_vecs.inputs.in_rasb = oldrasb
+    >>> res = reor_vecs.run()
+    >>> out_rasb = res.outputs.out_rasb
+    >>> out_rasb_mat = np.loadtxt(out_rasb, skiprows=1)
+    >>> assert oldrasb_mat == out_rasb_mat
+    True
+    """
+
+    input_spec = _ReorientVectorsInputSpec
+    output_spec = _ReorientVectorsOutputSpec
+
+    def _run_interface(self, runtime):
+        from nipype.utils.filemanip import fname_presuffix
+
+        table = DiffusionGradientTable(
+            dwi_file=self.inputs.dwi_file,
+            rasb_file=self.inputs.rasb_file,
+            transforms=self.inputs.affines,
+        )
+        table.generate_vecval()
+        reor_table = table.reorient_rasb()
+
+        cwd = Path(runtime.cwd).absolute()
+        reor_rasb_file = fname_presuffix(
+            self.inputs.rasb_file,
+            use_ext=False,
+            suffix="_reoriented.tsv",
+            newpath=str(cwd),
+        )
+        np.savetxt(
+            str(reor_rasb_file),
+            reor_table,
+            delimiter="\t",
+            header="\t".join("RASB"),
+            fmt=["%.8f"] * 3 + ["%g"],
+        )
+
+        self._results["out_rasb"] = reor_rasb_file
+        return runtime

dmriprep/utils/vectors.py

@@ -12,11 +12,29 @@
 class DiffusionGradientTable:
     """Data structure for DWI gradients."""
 
-    __slots__ = ['_affine', '_gradients', '_b_scale', '_bvecs', '_bvals', '_normalized',
-                 '_b0_thres', '_bvec_norm_epsilon']
-
-    def __init__(self, dwi_file=None, bvecs=None, bvals=None, rasb_file=None,
-                 b_scale=True, b0_threshold=B0_THRESHOLD, bvec_norm_epsilon=BVEC_NORM_EPSILON):
+    __slots__ = [
+        "_affine",
+        "_gradients",
+        "_b_scale",
+        "_bvecs",
+        "_bvals",
+        "_normalized",
+        "_transforms",
+        "_b0_thres",
+        "_bvec_norm_epsilon",
+    ]
+
+    def __init__(
+        self,
+        dwi_file=None,
+        bvecs=None,
+        bvals=None,
+        rasb_file=None,
+        b_scale=True,
+        transforms=None,
+        b0_threshold=B0_THRESHOLD,
+        bvec_norm_epsilon=BVEC_NORM_EPSILON,
+    ):
         """
         Create a new table of diffusion gradients.
 
@@ -34,8 +52,8 @@ def __init__(self, dwi_file=None, bvecs=None, bvals=None, rasb_file=None,
             then bvecs and bvals will be dismissed.
         b_scale : bool
             Whether b-values should be normalized.
-
         """
+        self._transforms = transforms
         self._b_scale = b_scale
         self._b0_thres = b0_threshold
         self._bvec_norm_epsilon = bvec_norm_epsilon
@@ -87,7 +105,7 @@ def affine(self, value):
             dwi_file = nb.load(str(value))
             self._affine = dwi_file.affine.copy()
             return
-        if hasattr(value, 'affine'):
+        if hasattr(value, "affine"):
             self._affine = value.affine
         self._affine = np.array(value)
 
@@ -102,20 +120,20 @@ def bvecs(self, value):
         if isinstance(value, (str, Path)):
             value = np.loadtxt(str(value)).T
         else:
-            value = np.array(value, dtype='float32')
+            value = np.array(value, dtype="float32")
 
         # Correct any b0's in bvecs misstated as 10's.
         value[np.any(abs(value) >= 10, axis=1)] = np.zeros(3)
         if self.bvals is not None and value.shape[0] != self.bvals.shape[0]:
-            raise ValueError('The number of b-vectors and b-values do not match')
+            raise ValueError("The number of b-vectors and b-values do not match")
         self._bvecs = value
 
     @bvals.setter
     def bvals(self, value):
         if isinstance(value, (str, Path)):
             value = np.loadtxt(str(value)).flatten()
         if self.bvecs is not None and value.shape[0] != self.bvecs.shape[0]:
-            raise ValueError('The number of b-vectors and b-values do not match')
+            raise ValueError("The number of b-vectors and b-values do not match")
         self._bvals = np.array(value)
 
     @property
@@ -129,10 +147,12 @@ def normalize(self):
             return
 
         self._bvecs, self._bvals = normalize_gradients(
-            self.bvecs, self.bvals,
+            self.bvecs,
+            self.bvals,
             b0_threshold=self._b0_thres,
             bvec_norm_epsilon=self._bvec_norm_epsilon,
-            b_scale=self._b_scale)
+            b_scale=self._b_scale,
+        )
         self._normalized = True
 
     def generate_rasb(self):
@@ -142,14 +162,50 @@ def generate_rasb(self):
             _ras = bvecs2ras(self.affine, self.bvecs)
             self.gradients = np.hstack((_ras, self.bvals[..., np.newaxis]))
 
+    def reorient_rasb(self):
+        """Reorient the vectors based o a list of affine transforms."""
+        from dipy.core.gradients import gradient_table_from_bvals_bvecs, reorient_bvecs
+
+        affines = self._transforms.copy()
+        bvals = self._bvals
+        bvecs = self._bvecs
+
+        # Verify that number of non-B0 volumes corresponds to the number of affines.
+        # If not, raise an error.
+        if len(self._bvals[self._bvals >= self._b0_thres]) != len(affines):
+            b0_indices = np.where(self._bvals <= self._b0_thres)[0].tolist()
+            if len(self._bvals[self._bvals >= self._b0_thres]) < len(affines):
+                for i in sorted(b0_indices, reverse=True):
+                    del affines[i]
+            if len(self._bvals[self._bvals >= self._b0_thres]) > len(affines):
+                ras_b_mat = self._gradients.copy()
+                ras_b_mat = np.delete(ras_b_mat, tuple(b0_indices), axis=0)
+                bvals = ras_b_mat[:, 3]
+                bvecs = ras_b_mat[:, 0:3]
+            if len(self._bvals[self._bvals > self._b0_thres]) != len(affines):
+                raise ValueError(
+                    "Affine transformations do not correspond to gradients"
+                )
+
+        # Build gradient table object
+        gt = gradient_table_from_bvals_bvecs(bvals, bvecs, b0_threshold=self._b0_thres)
+
+        # Reorient table
+        new_gt = reorient_bvecs(gt, [np.load(aff) for aff in affines])
+
+        return np.hstack((new_gt.bvecs, new_gt.bvals[..., np.newaxis]))
+
     def generate_vecval(self):
         """Compose a bvec/bval pair in image coordinates."""
         if self.bvecs is None or self.bvals is None:
             if self.affine is None:
                 raise TypeError(
                     "Cannot generate b-vectors & b-values in image coordinates. "
-                    "Please set the corresponding DWI image's affine matrix.")
-            self._bvecs = bvecs2ras(np.linalg.inv(self.affine), self.gradients[..., :-1])
+                    "Please set the corresponding DWI image's affine matrix."
+                )
+            self._bvecs = bvecs2ras(
+                np.linalg.inv(self.affine), self.gradients[..., :-1]
+            )
             self._bvals = self.gradients[..., -1].flatten()
 
     @property
@@ -161,22 +217,28 @@ def pole(self):
 
         """
         self.generate_rasb()
-        return calculate_pole(self.gradients[..., :-1], bvec_norm_epsilon=self._bvec_norm_epsilon)
+        return calculate_pole(
+            self.gradients[..., :-1], bvec_norm_epsilon=self._bvec_norm_epsilon
+        )
 
-    def to_filename(self, filename, filetype='rasb'):
+    def to_filename(self, filename, filetype="rasb"):
         """Write files (RASB, bvecs/bvals) to a given path."""
-        if filetype.lower() == 'rasb':
+        if filetype.lower() == "rasb":
             self.generate_rasb()
-            np.savetxt(str(filename), self.gradients,
-                       delimiter='\t', header='\t'.join('RASB'),
-                       fmt=['%.8f'] * 3 + ['%g'])
-        elif filetype.lower() == 'fsl':
+            np.savetxt(
+                str(filename),
+                self.gradients,
+                delimiter="\t",
+                header="\t".join("RASB"),
+                fmt=["%.8f"] * 3 + ["%g"],
+            )
+        elif filetype.lower() == "fsl":
             self.generate_vecval()
-            np.savetxt('%s.bvec' % filename, self.bvecs.T, fmt='%.6f')
-            np.savetxt('%s.bval' % filename, self.bvals, fmt='%.6f')
+            np.savetxt("%s.bvec" % filename, self.bvecs.T, fmt="%.6f")
+            np.savetxt("%s.bval" % filename, self.bvals, fmt="%.6f")
         else:
             raise ValueError('Unknown filetype "%s"' % filetype)
-
+        
 
 def normalize_gradients(bvecs, bvals, b0_threshold=B0_THRESHOLD,
                         bvec_norm_epsilon=BVEC_NORM_EPSILON, b_scale=True):
@@ -375,32 +437,3 @@ def bvecs2ras(affine, bvecs, norm=True, bvec_norm_epsilon=0.2):
         rotated_bvecs[~b0s] /= norms_bvecs[~b0s, np.newaxis]
         rotated_bvecs[b0s] = np.zeros(3)
     return rotated_bvecs
-
-
-def reorient_bvecs_from_ras_b(ras_b, affines):
-    """
-    Reorient the vectors from a rasb .tsv file.
-    When correcting for motion, rotation of the diffusion-weighted volumes
-    might cause systematic bias in rotationally invariant measures, such as FA
-    and MD, and also cause characteristic biases in tractography, unless the
-    gradient directions are appropriately reoriented to compensate for this
-    effect [Leemans2009]_.
-
-    Parameters
-    ----------
-    rasb_file : str or os.pathlike
-        File path to a RAS-B gradient table. If rasb_file is provided,
-        then bvecs and bvals will be dismissed.
-
-    affines : list or ndarray of shape (n, 4, 4) or (n, 3, 3)
-        Each entry in this list or array contain either an affine
-        transformation (4,4) or a rotation matrix (3, 3).
-        In both cases, the transformations encode the rotation that was applied
-        to the image corresponding to one of the non-zero gradient directions.
-    """
-    from dipy.core.gradients import gradient_table_from_bvals_bvecs, reorient_bvecs
-
-    ras_b_mat = np.genfromtxt(ras_b, delimiter='\t')
-    gt = gradient_table_from_bvals_bvecs(ras_b_mat[:,3], ras_b_mat[:,0:3], b0_threshold=50)
-
-    return reorient_bvecs(gt, affines)