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/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/base.py

@@ -26,8 +26,6 @@
 
 import numpy as np
 
-from nifreeze.exceptions import ModelNotFittedError
-
 
 class ModelFactory:
     """A factory for instantiating data models."""
@@ -61,7 +59,7 @@ def init(model=None, **kwargs):
             return AverageDWIModel(**kwargs)
 
         if model.lower() in ("avg", "average", "mean"):
-            return AverageModel(**kwargs)
+            return ExpectationModel(**kwargs)
 
         if model.lower() in ("dti", "dki", "pet"):
             Model = globals()[f"{model.upper()}Model"]
@@ -81,114 +79,80 @@ class BaseModel:
 
     """
 
-    __slots__ = (
-        "_model",
-        "_mask",
-        "_models",
-        "_datashape",
-        "_is_fitted",
-        "_modelargs",
-    )
+    __slots__ = {
+        "_dataset": "Reference to a :obj:`~nifreeze.data.base.BaseDataset` object.",
+    }
 
-    def __init__(self, mask=None, **kwargs):
+    def __init__(self, dataset, **kwargs):
         """Base initialization."""
 
-        # Keep model state
-        self._model = None  # "Main" model
-        self._models = None  # For parallel (chunked) execution
-
         # Setup brain mask
-        if mask is None:
+        if dataset.brainmask is None:
             warn(
                 "No mask provided; consider using a mask to avoid issues in model optimization.",
                 stacklevel=2,
             )
 
-        self._mask = mask
-
-        self._datashape = None
-        self._is_fitted = False
-
-        self._modelargs = ()
-
-    @property
-    def is_fitted(self):
-        return self._is_fitted
-
-    def fit(self, data, **kwargs):
-        """Abstract member signature of fit()."""
-        raise NotImplementedError("Cannot call fit() on a BaseModel instance.")
-
-    def predict(self, *args, **kwargs):
-        """Abstract member signature of predict()."""
-        raise NotImplementedError("Cannot call predict() on a BaseModel instance.")
+    def fit_predict(self, *_, **kwargs):
+        """Fit and predict the indicate index of the dataset (abstract signature)."""
+        raise NotImplementedError("Cannot call fit_predict() on a BaseModel instance.")
 
 
 class TrivialModel(BaseModel):
     """A trivial model that returns a given map always."""
 
-    __slots__ = ("_predicted",)
+    __slots__ = {
+        "_predicted": "A :obj:`~numpy.ndarray` with shape matching the dataset containing the map"
+        "that will always be returned as prediction (that is, a reference volume).",
+    }
 
-    def __init__(self, predicted=None, **kwargs):
+    def __init__(self, dataset, predicted=None, **kwargs):
         """Implement object initialization."""
-        if predicted is None:
-            raise TypeError("This model requires the predicted map at initialization")
 
-        super().__init__(**kwargs)
-        self._predicted = predicted
-        self._datashape = predicted.shape
+        super().__init__(dataset, **kwargs)
+        self._predicted = (
+            predicted
+            if predicted is not None
+            # Infer from dataset if not provided at initialization
+            else getattr(dataset, "reference", getattr(dataset, "bzero", None))
+        )
 
-    @property
-    def is_fitted(self):
-        return True
-
-    def fit(self, data, **kwargs):
-        """Do nothing."""
+        if self._predicted is None:
+            raise TypeError("This model requires the predicted map at initialization")
 
-    def predict(self, *_, **kwargs):
+    def fit_predict(self, *_, **kwargs):
         """Return the reference map."""
 
         # No need to check fit (if not fitted, has raised already)
         return self._predicted
 
 
-class AverageModel(BaseModel):
-    """A trivial model that returns an average map."""
+class ExpectationModel(BaseModel):
+    """A trivial model that returns an expectation map (for example, average)."""
 
-    __slots__ = ("_data",)
+    __slots__ = {"_stat": "The statistical operation to obtain the expectation map."}
 
-    def __init__(self, **kwargs):
+    def __init__(self, dataset, stat="median", **kwargs):
         """Initialize a new model."""
-        super().__init__(**kwargs)
-        self._data = None
+        super().__init__(dataset, **kwargs)
+        self._stat = stat
 
-    def fit(self, data, **kwargs):
-        """Calculate the average."""
+    def fit_predict(self, index, *_, **kwargs):
+        """
+        Return the expectation map.
 
-        # Regress out global signal differences
-        if kwargs.pop("equalize", False):
-            data = data.copy().astype("float32")
-            reshaped_data = (
-                data.reshape((-1, data.shape[-1])) if self._mask is None else data[self._mask]
-            )
-            p5 = np.percentile(reshaped_data, 5.0, axis=0)
-            p95 = np.percentile(reshaped_data, 95.0, axis=0) - p5
-            data = (data - p5) * p95.mean() / p95 + p5.mean()
+        Parameters
+        ----------
+        index : :obj:`int`
+            The volume index that is left-out in fitting, and then predicted.
 
+        """
         # Select the summary statistic
-        avg_func = getattr(np, kwargs.pop("stat", "mean"))
+        avg_func = getattr(np, kwargs.pop("stat", self._stat))
 
-        # Calculate the average
-        self._data = avg_func(data, axis=-1)
-
-    @property
-    def is_fitted(self):
-        return self._data is not None
-
-    def predict(self, *_, **kwargs):
-        """Return the average map."""
+        # Create index mask
+        mask = np.ones(len(self._dataset), dtype=bool)
+        mask[index] = False
 
-        if self._data is None:
-            raise ModelNotFittedError(f"{type(self).__name__} must be fitted before predicting")
-
-        return self._data
+        # Calculate the average
+        return avg_func(self._dataset.dataobj[mask][0], axis=-1)