expose bmag parameter and check number of b0s consistent after rounding

Author: josephmje

dmriprep/interfaces/vectors.py

@@ -20,6 +20,7 @@ class _CheckGradientTableInputSpec(BaseInterfaceInputSpec):
     in_rasb = File(exists=True, xor=["in_bval", "in_bvec"])
     b0_threshold = traits.Float(B0_THRESHOLD, usedefault=True)
     bvec_norm_epsilon = traits.Float(BVEC_NORM_EPSILON, usedefault=True)
+    b_mag = traits.Int(None, usedefault=True)
     b_scale = traits.Bool(True, usedefault=True)
 
 
@@ -84,6 +85,7 @@ def _run_interface(self, runtime):
             bvecs=_undefined(self.inputs, "in_bvec"),
             bvals=_undefined(self.inputs, "in_bval"),
             rasb_file=rasb_file,
+            b_mag=self.inputs.b_mag,
             b_scale=self.inputs.b_scale,
             bvec_norm_epsilon=self.inputs.bvec_norm_epsilon,
             b0_threshold=self.inputs.b0_threshold,

dmriprep/utils/vectors.py

@@ -17,6 +17,7 @@ class DiffusionGradientTable:
     __slots__ = [
         "_affine",
         "_b0_thres",
+        "_b_mag",
         "_b_scale",
         "_bvals",
         "_bvec_norm_epsilon",
@@ -30,6 +31,7 @@ class DiffusionGradientTable:
     def __init__(
         self,
         b0_threshold=B0_THRESHOLD,
+        b_mag=None,
         b_scale=True,
         bvals=None,
         bvec_norm_epsilon=BVEC_NORM_EPSILON,
@@ -46,12 +48,14 @@ def __init__(
         ----------
         b0_threshold : :obj:`float`
             The upper threshold to consider a low-b shell as :math:`b=0`.
+        b_mag : :obj:`int`
+            The order of magnitude to round the b-values.
         b_scale : :obj:`bool`
             Automatically scale the *b*-values with the norm of the corresponding
             *b*-vectors before the latter are normalized.
         bvals : str or os.pathlike or numpy.ndarray
             File path of the b-values.
-        b_vec_norm_epsilon : :obj:`float`
+        bvec_norm_epsilon : :obj:`float`
             The minimum difference in the norm of two *b*-vectors to consider them different.
         bvecs : str or os.pathlike or numpy.ndarray
             File path of the b-vectors.
@@ -93,6 +97,7 @@ def __init__(
         """
         self._affine = None
         self._b0_thres = b0_threshold
+        self._b_mag = b_mag
         self._b_scale = b_scale
         self._bvals = None
         self._bvec_norm_epsilon = bvec_norm_epsilon
@@ -195,6 +200,7 @@ def normalize(self):
             self.bvals,
             b0_threshold=self._b0_thres,
             bvec_norm_epsilon=self._bvec_norm_epsilon,
+            b_mag=self._b_mag,
             b_scale=self._b_scale,
             raise_error=self._raise_inconsistent,
         )
@@ -290,6 +296,7 @@ def normalize_gradients(
     bvals,
     b0_threshold=B0_THRESHOLD,
     bvec_norm_epsilon=BVEC_NORM_EPSILON,
+    b_mag=None,
     b_scale=True,
     raise_error=False,
 ):
@@ -361,17 +368,25 @@ def normalize_gradients(
             raise ValueError(msg)
         config.loggers.cli.warning(msg)
 
-    # Rescale b-vals if requested
+    # Rescale bvals if requested
     if b_scale:
         bvals[~b0s] *= np.linalg.norm(bvecs[~b0s], axis=1) ** 2
 
     # Ensure b0s have (0, 0, 0) vectors
     bvecs[b0s, :3] = np.zeros(3)
 
     # Round bvals
-    bvals = round_bvals(bvals)
+    bvals = round_bvals(bvals, bmag=b_mag)
 
-    # Rescale b-vecs, skipping b0's, on the appropriate axis to unit-norm length.
+    # Ensure rounding bvals doesn't change the number of b0s
+    rounded_b0s = bvals == 0
+    if not np.all(b0s == rounded_b0s):
+        msg = f"Inconsistent b0s before ({b0s.sum()}) and after rounding ({rounded_b0s.sum()})."
+        if raise_error:
+            raise ValueError(msg)
+        config.loggers.cli.warning(msg)
+
+    # Rescale bvecs, skipping b0's, on the appropriate axis to unit-norm length.
     bvecs[~b0s] /= np.linalg.norm(bvecs[~b0s], axis=1)[..., np.newaxis]
     return bvecs, bvals.astype("uint16")