Merge pull request #153 from nipreps/enh/dmri-models-bzero

ENH: Improved handling of the *b=0* and a max *b* in dMRI datasets

Author: oesteban

src/nifreeze/cli/parser.py

@@ -143,6 +143,12 @@ def build_parser() -> ArgumentParser:
         help="A NIfTI file containing the b-zero reference",
     )
 
+    g_dmri.add_argument(
+        "--ignore-b0",
+        action="store_true",
+        help="Ignore the low-b reference and use the robust signal maximum",
+    )
+
     g_pet = parser.add_argument_group("Options for PET inputs")
     g_pet.add_argument(
         "--timing-file",

src/nifreeze/cli/run.py

@@ -67,13 +67,19 @@ def main(argv=None) -> None:
         **extra_kwargs,
     )
 
+    model_kwargs = {}
+
+    if args.ignore_b0:
+        model_kwargs["ignore_bzero"] = True
+
     prev_model: Estimator | None = None
     for _model in args.models:
         single_fit = _model.lower().startswith("single")
         estimator: Estimator = Estimator(
             _model.lower().replace("single", ""),
             prev=prev_model,
             single_fit=single_fit,
+            model_kwargs=model_kwargs,
         )
         prev_model = estimator
 

src/nifreeze/model/dmri.py

@@ -34,6 +34,9 @@
 )
 from nifreeze.model.base import BaseModel, ExpectationModel
 
+S0_EPSILON = 1e-6
+B_MIN = 50
+
 
 def _exec_fit(model, data, chunk=None):
     retval = model.fit(data)
@@ -49,12 +52,15 @@ class BaseDWIModel(BaseModel):
     """Interface and default methods for DWI models."""
 
     __slots__ = {
+        "_max_b": "The maximum b-value supported by the model",
+        "_data_mask": "A mask for the voxels that will be fitted and predicted",
+        "_S0": "The S0 (b=0 reference signal) that will be fed into DIPY models",
         "_model_class": "Defining a model class, DIPY models are instantiated automagically",
         "_modelargs": "Arguments acceptable by the underlying DIPY-like model.",
         "_models": "List with one or more (if parallel execution) model instances",
     }
 
-    def __init__(self, dataset: DWI, **kwargs):
+    def __init__(self, dataset: DWI, max_b: float | int | None = None, **kwargs):
         r"""Initialization.
 
         Parameters
@@ -76,6 +82,26 @@ def __init__(self, dataset: DWI, **kwargs):
                 f"DWI dataset is too small ({dataset.gradients.shape[0]} directions)."
             )
 
+        if max_b is not None and max_b > B_MIN:
+            self._max_b = max_b
+
+        self._data_mask = (
+            dataset.brainmask
+            if dataset.brainmask is not None
+            else np.ones(dataset.dataobj.shape[:3], dtype=bool)
+        )
+
+        # By default, set S0 to the 98% percentile of the DWI data within mask
+        self._S0 = np.full(
+            self._data_mask.sum(),
+            np.round(np.percentile(dataset.dataobj[self._data_mask, ...], 98)),
+        )
+
+        # If b=0 is present and not to be ignored, update brain mask and set
+        if not kwargs.pop("ignore_bzero", False) and dataset.bzero is not None:
+            self._data_mask[dataset.bzero < S0_EPSILON] = False
+            self._S0 = dataset.bzero[self._data_mask]
+
         super().__init__(dataset, **kwargs)
 
     def _fit(self, index: int | None = None, n_jobs=None, **kwargs):
@@ -151,26 +177,20 @@ def fit_predict(self, index: int | None = None, **kwargs):
         if index is None:
             return None
 
-        brainmask = self._dataset.brainmask
         gradient = self._dataset.gradients[:, index]
 
         if "dipy" in getattr(self, "_model_class", ""):
             gradient = gradient_table_from_bvals_bvecs(
                 gradient[np.newaxis, -1], gradient[np.newaxis, :-1]
             )
 
-        S0 = self._dataset.bzero
-        if S0 is not None:
-            S0 = S0[brainmask, ...] if brainmask is not None else S0.reshape(-1)
-
         if n_models == 1:
             predicted, _ = _exec_predict(
-                self._models[0], **(kwargs | {"gtab": gradient, "S0": S0})
+                self._models[0], **(kwargs | {"gtab": gradient, "S0": self._S0})
             )
         else:
-            S0 = np.array_split(S0, n_models) if S0 is not None else np.full(n_models, None)
-
             predicted = [None] * n_models
+            S0 = np.array_split(self._S0, n_models)
 
             # Parallelize process with joblib
             with Parallel(n_jobs=n_models) as executor:
@@ -187,12 +207,8 @@ def fit_predict(self, index: int | None = None, **kwargs):
 
             predicted = np.hstack(predicted)
 
-        if brainmask is not None:
-            retval = np.zeros_like(brainmask, dtype="float32")
-            retval[brainmask, ...] = predicted
-        else:
-            retval = predicted.reshape(self._dataset.dataobj.shape[:-1])
-
+        retval = np.zeros_like(self._data_mask, dtype=self._dataset.dataobj.dtype)
+        retval[self._data_mask, ...] = predicted
         return retval