enh: allow "locking" of models with first fit

This allows having "frozen" models that are fit only once before
entering the leave-one-out loop.

These models use all the data in the fitting and then return always the
same "prediction".

This feature was lost with the refactor of the estimator. By moving it
into the models, we can use them in a more flexible way.

Author: oesteban

src/nifreeze/model/base.py

@@ -87,46 +87,59 @@ class BaseModel:
 
     """
 
-    __slots__ = ("_dataset",)
+    __slots__ = ("_dataset", "_locked_fit")
 
     def __init__(self, dataset, **kwargs):
         """Base initialization."""
 
+        self._locked_fit = None
         self._dataset = dataset
         # Warn if mask not present
         if dataset.brainmask is None:
             warn(mask_absence_warn_msg, stacklevel=2)
 
     @abstractmethod
-    def fit_predict(self, index, **kwargs) -> np.ndarray:
-        """Fit and predict the indicate index of the dataset (abstract signature)."""
+    def fit_predict(self, index: int | None = None, **kwargs) -> np.ndarray:
+        """
+        Fit and predict the indicated index of the dataset (abstract signature).
+
+        If ``index`` is ``None``, then the model is executed in *single-fit mode* meaning
+        that it will be run only once in all the data available.
+        Please note that all the predictions of this model will suffer from data leakage
+        from the original volume.
+
+        Parameters
+        ----------
+        index : :obj:`int` or ``None``
+            The index to predict.
+            If ``None``, no prediction will be executed.
+
+        """
         raise NotImplementedError("Cannot call fit_predict() on a BaseModel instance.")
 
 
 class TrivialModel(BaseModel):
     """A trivial model that returns a given map always."""
 
-    __slots__ = ("_predicted",)
-
     def __init__(self, dataset, predicted=None, **kwargs):
         """Implement object initialization."""
 
         super().__init__(dataset, **kwargs)
-        self._predicted = (
+        self._locked_fit = (
             predicted
             if predicted is not None
             # Infer from dataset if not provided at initialization
             else getattr(dataset, "reference", getattr(dataset, "bzero", None))
         )
 
-        if self._predicted is None:
+        if self._locked_fit is None:
             raise TypeError("This model requires the predicted map at initialization")
 
     def fit_predict(self, *_, **kwargs):
         """Return the reference map."""
 
         # No need to check fit (if not fitted, has raised already)
-        return self._predicted
+        return self._locked_fit
 
 
 class ExpectationModel(BaseModel):
@@ -139,7 +152,7 @@ def __init__(self, dataset, stat="median", **kwargs):
         super().__init__(dataset, **kwargs)
         self._stat = stat
 
-    def fit_predict(self, index: int, **kwargs):
+    def fit_predict(self, index: int | None = None, **kwargs):
         """
         Return the expectation map.
 
@@ -149,12 +162,20 @@ def fit_predict(self, index: int, **kwargs):
             The volume index that is left-out in fitting, and then predicted.
 
         """
+
+        if self._locked_fit is not None:
+            return self._locked_fit
+
         # Select the summary statistic
         avg_func = getattr(np, kwargs.pop("stat", self._stat))
 
         # Create index mask
         index_mask = np.ones(len(self._dataset), dtype=bool)
-        index_mask[index] = False
 
-        # Calculate the average
-        return avg_func(self._dataset[index_mask][0], axis=-1)
+        if index is not None:
+            index_mask[index] = False
+            # Calculate the average
+            return avg_func(self._dataset[index_mask][0], axis=-1)
+
+        self._locked_fit = avg_func(self._dataset[index_mask][0], axis=-1)
+        return self._locked_fit

src/nifreeze/model/dmri.py

@@ -51,6 +51,7 @@ class BaseDWIModel(BaseModel):
     __slots__ = {
         "_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):
@@ -77,13 +78,21 @@ def __init__(self, dataset: DWI, **kwargs):
 
         super().__init__(dataset, **kwargs)
 
-    def _fit(self, index, n_jobs=None, **kwargs):
+    def _fit(self, index: int | None = None, n_jobs=None, **kwargs):
         """Fit the model chunk-by-chunk asynchronously"""
+
         n_jobs = n_jobs or 1
 
+        if self._locked_fit is not None:
+            return n_jobs
+
         brainmask = self._dataset.brainmask
         idxmask = np.ones(len(self._dataset), dtype=bool)
-        idxmask[index] = False
+
+        if index is not None:
+            idxmask[index] = False
+        else:
+            self._locked_fit = True
 
         data, _, gtab = self._dataset[idxmask]
         # Select voxels within mask or just unravel 3D if no mask
@@ -96,14 +105,15 @@ def _fit(self, index, n_jobs=None, **kwargs):
 
         if model_str:
             module_name, class_name = model_str.rsplit(".", 1)
-            self._model = getattr(
+            model = getattr(
                 import_module(module_name),
                 class_name,
             )(gtab, **kwargs)
 
         # One single CPU - linear execution (full model)
         if n_jobs == 1:
-            self._model, _ = _exec_fit(self._model, data)
+            _modelfit, _ = _exec_fit(model, data)
+            self._models = [_modelfit]
             return 1
 
         # Split data into chunks of group of slices
@@ -114,15 +124,14 @@ def _fit(self, index, n_jobs=None, **kwargs):
         # Parallelize process with joblib
         with Parallel(n_jobs=n_jobs) as executor:
             results = executor(
-                delayed(_exec_fit)(self._model, dchunk, i) for i, dchunk in enumerate(data_chunks)
+                delayed(_exec_fit)(model, dchunk, i) for i, dchunk in enumerate(data_chunks)
             )
         for submodel, rindex in results:
             self._models[rindex] = submodel
 
-        self._model = None  # Preempt further actions on the model
         return n_jobs
 
-    def fit_predict(self, index: int, **kwargs):
+    def fit_predict(self, index: int | None = None, **kwargs):
         """
         Predict asynchronously chunk-by-chunk the diffusion signal.
 
@@ -133,8 +142,14 @@ def fit_predict(self, index: int, **kwargs):
 
         """
 
-        n_models = self._fit(index, **kwargs)
-        kwargs.pop("n_jobs")
+        n_models = self._fit(
+            index,
+            n_jobs=kwargs.pop("n_jobs"),
+            **kwargs,
+        )
+
+        if index is None:
+            return None
 
         brainmask = self._dataset.brainmask
         gradient = self._dataset.gradients[:, index]
@@ -149,9 +164,10 @@ def fit_predict(self, index: int, **kwargs):
             S0 = S0[brainmask, ...] if brainmask is not None else S0.reshape(-1)
 
         if n_models == 1:
-            predicted, _ = _exec_predict(self._model, **(kwargs | {"gtab": gradient, "S0": S0}))
+            predicted, _ = _exec_predict(
+                self._models[0], **(kwargs | {"gtab": gradient, "S0": S0})
+            )
         else:
-            print(n_models, S0)
             S0 = np.array_split(S0, n_models) if S0 is not None else np.full(n_models, None)
 
             predicted = [None] * n_models
@@ -221,9 +237,12 @@ def __init__(
         self._th_high = th_high
         self._detrend = detrend
 
-    def fit_predict(self, index, *_, **kwargs):
+    def fit_predict(self, index: int | None = None, *_, **kwargs):
         """Return the average map."""
 
+        if index is None:
+            raise RuntimeError(f"Model {self.__class__.__name__} does not allow locking.")
+
         bvalues = self._dataset.gradients[:, -1]
         bcenter = bvalues[index]
 

src/nifreeze/model/pet.py

@@ -36,7 +36,7 @@
 class PETModel(BaseModel):
     """A PET imaging realignment model based on B-Spline approximation."""
 
-    __slots__ = ("_t", "_x", "_xlim", "_order", "_coeff", "_n_ctrl")
+    __slots__ = ("_t", "_x", "_xlim", "_order", "_n_ctrl")
 
     def __init__(self, timepoints=None, xlim=None, n_ctrl=None, order=3, **kwargs):
         """
@@ -76,13 +76,17 @@ def __init__(self, timepoints=None, xlim=None, n_ctrl=None, order=3, **kwargs):
         # B-Spline knots
         self._t = np.arange(-3, float(self._n_ctrl) + 4, dtype="float32")
 
-        self._coeff = None
-
-    def _fit(self, n_jobs=None, **kwargs):
+    def _fit(self, index: int | None = None, n_jobs=None, **kwargs):
         """Fit the model."""
         from scipy.interpolate import BSpline
         from scipy.sparse.linalg import cg
 
+        if self._locked_fit is not None:
+            return n_jobs
+
+        if index is not None:
+            raise NotImplementedError("Fitting with held-out data is not supported")
+
         timepoints = kwargs.get("timepoints", None) or self._x
         x = (np.array(timepoints, dtype="float32") / self._xlim) * self._n_ctrl
 
@@ -101,15 +105,15 @@ def _fit(self, n_jobs=None, **kwargs):
         with Parallel(n_jobs=n_jobs or min(cpu_count() or 1, 8)) as executor:
             results = executor(delayed(cg)(ATdotA, AT @ v) for v in data)
 
-        self._coeff = np.array([r[0] for r in results])
+        self._locked_fit = np.array([r[0] for r in results])
 
     def fit_predict(self, index: int | None = None, **kwargs):
         """Return the corrected volume using B-spline interpolation."""
         from scipy.interpolate import BSpline
 
         # Fit the BSpline basis on all data
-        if self._coeff is None:
-            self._fit(n_jobs=kwargs.pop("n_jobs", None))
+        if self._locked_fit is None:
+            self._fit(index, n_jobs=kwargs.pop("n_jobs", None), **kwargs)
 
         if index is None:  # If no index, just fit the data.
             return None
@@ -120,7 +124,7 @@ def fit_predict(self, index: int | None = None, **kwargs):
 
         # A is 1 (num. timepoints) x C (num. coeff)
         # self._coeff is V (num. voxels) x K - 4
-        predicted = np.squeeze(A @ self._coeff.T)
+        predicted = np.squeeze(A @ self._locked_fit.T)
 
         brainmask = self._dataset.brainmask
         datashape = self._dataset.dataobj.shape[:3]