fix: continue adaptations to new Estimator

Author: oesteban

src/nifreeze/estimator.py

@@ -39,11 +39,28 @@
 class Filter:
     """Alters an input data object (e.g., downsampling)."""
 
+    def run(self, dataset: BaseDataset, **kwargs):
+        """
+        Trigger execution of the designated filter.
+
+        Parameters
+        ----------
+        dataset : :obj:`~nifreeze.data.base.BaseDataset`
+            The input dataset this estimator operates on.
+
+        Returns
+        -------
+        :obj:`~nifreeze.estimator.Estimator`
+            The estimator, after fitting.
+
+        """
+        return dataset
+
 
 class Estimator:
     """Estimates rigid-body head-motion and distortions derived from eddy-currents."""
 
-    __slots__ = ("_model", "_strategy", "_prev", "_model_kwargs", "_align_kwargs")
+    __slots__ = ("_model", "_strategy", "_dataset", "_prev", "_model_kwargs", "_align_kwargs")
 
     def __init__(
         self,
@@ -56,10 +73,24 @@ def __init__(
         self._model = model
         self._prev = prev
         self._strategy = strategy
-        self._model_kwargs = model_kwargs
-        self._align_kwargs = kwargs
+        self._model_kwargs = model_kwargs or {}
+        self._align_kwargs = kwargs or {}
 
     def run(self, dataset: BaseDataset, **kwargs):
+        """
+        Trigger execution of the workflow this estimator belongs.
+
+        Parameters
+        ----------
+        dataset : :obj:`~nifreeze.data.base.BaseDataset`
+            The input dataset this estimator operates on.
+
+        Returns
+        -------
+        :obj:`~nifreeze.estimator.Estimator`
+            The estimator, after fitting.
+
+        """
         if self._prev is not None:
             result = self._prev.run(dataset, **kwargs)
             if isinstance(self._prev, Filter):
@@ -69,7 +100,7 @@ def run(self, dataset: BaseDataset, **kwargs):
 
         # Prepare iterator
         iterfunc = getattr(iterators, f"{self._strategy}_iterator")
-        index_iter = iterfunc(dataset, seed=kwargs.get("seed", None))
+        index_iter = iterfunc(len(dataset), seed=kwargs.get("seed", None))
 
         # Initialize model
         if isinstance(self._model, str):

src/nifreeze/model/base.py

@@ -33,7 +33,7 @@ class ModelFactory:
     """A factory for instantiating data models."""
 
     @staticmethod
-    def init(model="DTI", **kwargs):
+    def init(model=None, **kwargs):
         """
         Instantiate a diffusion model.
 
@@ -49,6 +49,9 @@ def init(model="DTI", **kwargs):
             A model object compliant with DIPY's interface.
 
         """
+        if model is None:
+            raise RuntimeError("No model identifier provided.")
+
         if model.lower() in ("s0", "b0"):
             return TrivialModel(predicted=kwargs.pop("S0"), gtab=kwargs.pop("gtab"))
 
@@ -143,7 +146,7 @@ def fit(self, data, **kwargs):
         """Do nothing."""
 
     def predict(self, *_, **kwargs):
-        """Return the *b=0* map."""
+        """Return the reference map."""
 
         # No need to check fit (if not fitted, has raised already)
         return self._predicted

test/test_integration.py

@@ -70,10 +70,9 @@ def test_proximity_estimator_trivial_model(datadir, tmp_path):
         brainmask=dwdata.brainmask,
     )
 
-    estimator = Estimator(dwi_motion, model="b0")
+    estimator = Estimator("b0")
     estimator.run(
-        data=dwi_motion,
-        models=("b0",),
+        dwi_motion,
         seed=None,
         align_kwargs={
             "config_file": "b0-to-b0_level0.json",