enh: continue with the refactor

Author: oesteban

docs/notebooks/bold_realignment.ipynb

@@ -337,7 +337,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from nifreeze.model.base import AverageModel\n",
+    "from nifreeze.model.base import ExpectationModel\n",
     "from nifreeze.utils.iterators import random_iterator"
    ]
   },
@@ -358,7 +358,7 @@
     "        t_mask[t] = True\n",
     "\n",
     "        # Fit and predict using the model\n",
-    "        model = AverageModel()\n",
+    "        model = ExpectationModel()\n",
     "        model.fit(\n",
     "            data[..., ~t_mask],\n",
     "            stat=\"median\",\n",
@@ -376,7 +376,7 @@
     "            fixedmask_path=brainmask_path,\n",
     "            output_transform_prefix=f\"conversion-{t:02d}\",\n",
     "            num_threads=8,\n",
-    "        )\n",
+    "        ).cmdline\n",
     "\n",
     "        # Run the command\n",
     "        proc = await asyncio.create_subprocess_shell(\n",

scripts/optimize_registration.py

@@ -133,7 +133,7 @@ async def train_coro(
                 fixedmask_path=brainmask_path,
                 output_transform_prefix=f"conversion-{index:04d}",
                 **align_kwargs,
-            )
+            ).cmdline
 
             tasks.append(
                 ants(

src/nifreeze/data/dmri.py

@@ -37,6 +37,30 @@
 
 from nifreeze.data.base import BaseDataset, _cmp, _data_repr
 
+DEFAULT_CLIP_PERCENTILE = 75
+"""Upper percentile threshold for intensity clipping."""
+
+DEFAULT_MIN_S0 = 1e-5
+"""Minimum value when considering the :math:`S_{0}` DWI signal."""
+
+DEFAULT_MAX_S0 = 1.0
+"""Maximum value when considering the :math:`S_{0}` DWI signal."""
+
+DEFAULT_LOWB_THRESHOLD = 50
+"""The lower bound for the b-value so that the orientation is considered a DW volume."""
+
+DEFAULT_HIGHB_THRESHOLD = 8000
+"""A b-value cap for DWI data."""
+
+DEFAULT_NUM_BINS = 15
+"""Number of bins to classify b-values."""
+
+DEFAULT_MULTISHELL_BIN_COUNT_THR = 7
+"""Default bin count to consider a multishell scheme."""
+
+DTI_MIN_ORIENTATIONS = 6
+"""Minimum number of nonzero b-values in a DWI dataset."""
+
 
 @attr.s(slots=True)
 class DWI(BaseDataset):
@@ -221,7 +245,7 @@ def load(
     bvec_file: Path | str | None = None,
     bval_file: Path | str | None = None,
     b0_file: Path | str | None = None,
-    b0_thres: float = 50.0,
+    b0_thres: float = DEFAULT_LOWB_THRESHOLD,
 ) -> DWI:
     """
     Load DWI data and construct a DWI object.
@@ -337,3 +361,87 @@ def load(
         dwi_obj.brainmask = np.asanyarray(mask_img.dataobj, dtype=bool)
 
     return dwi_obj
+
+
+def find_shelling_scheme(
+    bvals,
+    num_bins=DEFAULT_NUM_BINS,
+    multishell_nonempty_bin_count_thr=DEFAULT_MULTISHELL_BIN_COUNT_THR,
+    bval_cap=DEFAULT_HIGHB_THRESHOLD,
+):
+    """
+    Find the shelling scheme on the given b-values.
+
+    Computes the histogram of the b-values according to ``num_bins``
+    and depending on the nonempty bin count, classify the shelling scheme
+    as single-shell if they are 2 (low-b and a shell); multi-shell if they are
+    below the ``multishell_nonempty_bin_count_thr`` value; and DSI otherwise.
+
+    Parameters
+    ----------
+    bvals : :obj:`list` or :obj:`~numpy.ndarray`
+         List or array of b-values.
+    num_bins : :obj:`int`, optional
+        Number of bins.
+    multishell_nonempty_bin_count_thr : :obj:`int`, optional
+        Bin count to consider a multi-shell scheme.
+
+    Returns
+    -------
+    scheme : :obj:`str`
+        Shelling scheme.
+    bval_groups : :obj:`list`
+        List of grouped b-values.
+    bval_estimated : :obj:`list`
+        List of 'estimated' b-values as the median value of each b-value group.
+
+    """
+
+    # Bin the b-values: use -1 as the lower bound to be able to appropriately
+    # include b0 values
+    hist, bin_edges = np.histogram(bvals, bins=num_bins, range=(-1, min(max(bvals), bval_cap)))
+
+    # Collect values in each bin
+    bval_groups = []
+    bval_estimated = []
+    for lower, upper in zip(bin_edges[:-1], bin_edges[1:], strict=False):
+        # Add only if a nonempty b-values mask
+        if (mask := (bvals > lower) & (bvals <= upper)).sum():
+            bval_groups.append(bvals[mask])
+            bval_estimated.append(np.median(bvals[mask]))
+
+    nonempty_bins = len(bval_groups)
+
+    if nonempty_bins < 2:
+        raise ValueError("DWI must have at least one high-b shell")
+
+    if nonempty_bins == 2:
+        scheme = "single-shell"
+    elif nonempty_bins < multishell_nonempty_bin_count_thr:
+        scheme = "multi-shell"
+    else:
+        scheme = "DSI"
+
+    return scheme, bval_groups, bval_estimated
+
+
+def _rasb2dipy(gradient):
+    import warnings
+
+    gradient = np.asanyarray(gradient)
+    if gradient.ndim == 1:
+        if gradient.size != 4:
+            raise ValueError("Missing gradient information.")
+        gradient = gradient[..., np.newaxis]
+
+    if gradient.shape[0] != 4:
+        gradient = gradient.T
+    elif gradient.shape == (4, 4):
+        print("Warning: make sure gradient information is not transposed!")
+
+    with warnings.catch_warnings():
+        from dipy.core.gradients import gradient_table
+
+        warnings.filterwarnings("ignore", category=UserWarning)
+        retval = gradient_table(gradient[3, :], bvecs=gradient[:3, :].T)
+    return retval

src/nifreeze/estimator.py

@@ -32,7 +32,10 @@
 
 from nifreeze.data.base import BaseDataset
 from nifreeze.model.base import BaseModel, ModelFactory
-from nifreeze.registration.ants import _prepare_registration_data, _run_registration
+from nifreeze.registration.ants import (
+    _prepare_registration_data,
+    _run_registration,
+)
 from nifreeze.utils import iterators
 
 
@@ -60,7 +63,7 @@ def run(self, dataset: BaseDataset, **kwargs):
 class Estimator:
     """Estimates rigid-body head-motion and distortions derived from eddy-currents."""
 
-    __slots__ = ("_model", "_strategy", "_dataset", "_prev", "_model_kwargs", "_align_kwargs")
+    __slots__ = ("_model", "_strategy", "_prev", "_model_kwargs", "_align_kwargs")
 
     def __init__(
         self,
@@ -111,29 +114,37 @@ def run(self, dataset: BaseDataset, **kwargs):
                 **self._model_kwargs,
             )
 
-        if self._model.is_static:
-            self._model.fit(dataset, **kwargs)
-
         kwargs["num_threads"] = kwargs.pop("omp_nthreads", None) or kwargs.pop("num_threads", None)
 
         dataset_length = len(dataset)
         with TemporaryDirectory() as tmp_dir:
             print(f"Processing in <{tmp_dir}>")
             ptmp_dir = Path(tmp_dir)
+
+            bmask_path = None
+            if dataset.brainmask is not None:
+                import nibabel as nb
+
+                bmask_path = ptmp_dir / "brainmask.nii.gz"
+                nb.Nifti1Image(
+                    dataset.brainmask.astype("uint8"), dataset.affine, None
+                ).to_filename(bmask_path)
+
             with tqdm(total=dataset_length, unit="vols.") as pbar:
                 # run a original-to-synthetic affine registration
                 for i in index_iter:
                     pbar.set_description_str(f"Fit and predict vol. <{i}>")
 
                     # fit the model
-                    reference, predicted = self._model.fit_predict(
+                    test_set = dataset[i]
+                    predicted = self._model.fit_predict(
                         i,
                         n_jobs=n_jobs,
                     )
 
                     # prepare data for running ANTs
-                    fixed, moving = _prepare_registration_data(
-                        reference,
+                    predicted_path, volume_path, init_path = _prepare_registration_data(
+                        test_set[0],
                         predicted,
                         dataset.affine,
                         i,
@@ -144,14 +155,13 @@ def run(self, dataset: BaseDataset, **kwargs):
                     pbar.set_description_str(f"Realign vol. <{i}>")
 
                     xform = _run_registration(
-                        fixed,
-                        moving,
-                        dataset.brainmask,
-                        dataset.motion_affines,
-                        dataset.affine,
-                        dataset.dataobj.shape[:3],
+                        predicted_path,
+                        volume_path,
                         i,
                         ptmp_dir,
+                        init_affine=init_path,
+                        fixedmask_path=bmask_path,
+                        output_transform_prefix=f"ants-{i:05d}",
                         **kwargs,
                     )
 

src/nifreeze/model/__init__.py

@@ -23,7 +23,7 @@
 """Data models."""
 
 from nifreeze.model.base import (
-    AverageModel,
+    ExpectationModel,
     ModelFactory,
     TrivialModel,
 )
@@ -37,7 +37,7 @@
 
 __all__ = (
     "ModelFactory",
-    "AverageModel",
+    "ExpectationModel",
     "AverageDWIModel",
     "DKIModel",
     "DTIModel",

src/nifreeze/model/_dipy.py

@@ -24,8 +24,6 @@
 
 from __future__ import annotations
 
-import warnings
-
 import numpy as np
 from dipy.core.gradients import GradientTable
 from dipy.reconst.base import ReconstModel
@@ -268,23 +266,3 @@ def predict(
 
         """
         return gp_prediction(self.model, gtab, mask=self.mask)
-
-
-def _rasb2dipy(gradient):
-    gradient = np.asanyarray(gradient)
-    if gradient.ndim == 1:
-        if gradient.size != 4:
-            raise ValueError("Missing gradient information.")
-        gradient = gradient[..., np.newaxis]
-
-    if gradient.shape[0] != 4:
-        gradient = gradient.T
-    elif gradient.shape == (4, 4):
-        print("Warning: make sure gradient information is not transposed!")
-
-    with warnings.catch_warnings():
-        from dipy.core.gradients import gradient_table
-
-        warnings.filterwarnings("ignore", category=UserWarning)
-        retval = gradient_table(gradient[3, :], bvecs=gradient[:3, :].T)
-    return retval