NF: Check multi b-value

dipy/core/gradients.py

@@ -56,7 +56,6 @@ class GradientTable(object):
     using the factory function gradient_table
 
     """
-
     def __init__(self, gradients, big_delta=None, small_delta=None,
                  b0_threshold=0):
         """Constructor for GradientTable class"""
@@ -322,23 +321,66 @@ def reorient_bvecs(gtab, affines):
 def generate_bvecs(N, iters=5000):
     """Generates N bvectors.
 
-    Uses dipy.core.sphere.disperse_charges to model electrostatic repulsion on a unit sphere. 
+    Uses dipy.core.sphere.disperse_charges to model electrostatic repulsion on
+    a unit sphere.
 
     Parameters
     ----------
     N : int
-        The number of bvectors to generate. This should be equal to the number of bvals used.
+        The number of bvectors to generate. This should be equal to the number
+        of bvals used.
     iters : int
         Number of iterations to run.
 
     Returns
     -------
     bvecs : (N,3) ndarray
-        The generated directions, represented as a unit vector, of each gradient."""
-
+        The generated directions, represented as a unit vector, of each
+        gradient.
+    """
     theta = np.pi * np.random.rand(N)
     phi = 2 * np.pi * np.random.rand(N)
     hsph_initial = HemiSphere(theta=theta, phi=phi)
     hsph_updated, potential = disperse_charges(hsph_initial, iters)
     bvecs = hsph_updated.vertices
     return bvecs
+
+
+def check_multi_b(gtab, n_bvals, non_zero=True, bmag=None):
+    """
+    Check if you have enough different b-values in your gradient table
+
+    Parameters
+    ----------
+    gtab : GradientTable class instance.
+
+    n_bvals : int
+        The number of different b-values you are checking for.
+    non_zero : bool
+        Whether to check only non-zero bvalues. In this case, we will require
+        at least `n_bvals` *non-zero* b-values (where non-zero is defined
+        depending on the `gtab` object's `b0_threshold` attribute)
+    bmag : int
+        The order of magnitude of the b-values used. The function will
+        normalize the b-values relative $10^{bmag - 1}$. Default: derive this
+        value from the maximal b-value provided: $bmag=log_{10}(max(bvals))$.
+
+    Returns
+    -------
+    bool : Whether there are at least `n_bvals` different b-values in the
+    gradient table used.
+    """
+    bvals = gtab.bvals.copy()
+    if non_zero:
+        bvals = bvals[~gtab.b0s_mask]
+
+    if bmag is None:
+        bmag = int(np.log10(np.max(bvals)))
+
+    b = bvals / (10 ** (bmag - 1))  # normalize b units
+    b = b.round()
+    uniqueb = np.unique(b)
+    if uniqueb.shape[0] < n_bvals:
+        return False
+    else:
+        return True

dipy/core/tests/test_gradients.py

@@ -7,7 +7,8 @@
 from dipy.data import get_data
 from dipy.core.gradients import (gradient_table, GradientTable,
                                  gradient_table_from_bvals_bvecs,
-                                 reorient_bvecs, generate_bvecs)
+                                 reorient_bvecs, generate_bvecs,
+                                 check_multi_b)
 from dipy.io.gradients import read_bvals_bvecs
 
 
@@ -278,6 +279,25 @@ def test_generate_bvecs():
     npt.assert_almost_equal(cos_theta, 0.)
 
 
+def test_check_multi_b():
+    bvals = np.array([1000, 1000, 1000, 1000, 2000, 2000, 2000, 2000, 0])
+    bvecs = generate_bvecs(bvals.shape[-1])
+    gtab = gradient_table(bvals, bvecs)
+    npt.assert_(check_multi_b(gtab, 2, non_zero=False))
+
+    # We don't consider differences this small to be sufficient:
+    bvals = np.array([1995, 1995, 1995, 1995, 2005, 2005, 2005, 2005, 0])
+    bvecs = generate_bvecs(bvals.shape[-1])
+    gtab = gradient_table(bvals, bvecs)
+    npt.assert_(not check_multi_b(gtab, 2, non_zero=True))
+
+    # Unless you specify that you are interested in this magnitude of changes:
+    npt.assert_(check_multi_b(gtab, 2, non_zero=True, bmag=1))
+
+    # Or if you consider zero to be one of your b-values:
+    npt.assert_(check_multi_b(gtab, 2, non_zero=False))
+
+
 if __name__ == "__main__":
     from numpy.testing import run_module_suite
     run_module_suite()