Merge pull request #120 from nipreps/fix/reenable-single-fit-models

ENH: Re-allow "locking" of models with first fit

Author: oesteban

src/nifreeze/cli/run.py

@@ -69,9 +69,11 @@ def main(argv=None) -> None:
 
     prev_model: Estimator | None = None
     for _model in args.models:
+        single_fit = _model.lower().startswith("single")
         estimator: Estimator = Estimator(
-            _model,
+            _model.lower().replace("single", ""),
             prev=prev_model,
+            single_fit=single_fit,
         )
         prev_model = estimator
 

src/nifreeze/estimator.py

@@ -70,19 +70,21 @@ def run(self, dataset: DatasetT, **kwargs) -> DatasetT:
 class Estimator:
     """Orchestrates components for a single estimation step."""
 
-    __slots__ = ("_model", "_strategy", "_prev", "_model_kwargs", "_align_kwargs")
+    __slots__ = ("_model", "_single_fit", "_strategy", "_prev", "_model_kwargs", "_align_kwargs")
 
     def __init__(
         self,
         model: BaseModel | str,
         strategy: str = "random",
         prev: Estimator | Filter | None = None,
         model_kwargs: dict | None = None,
+        single_fit: bool = False,
         **kwargs,
     ):
         self._model = model
         self._prev = prev
         self._strategy = strategy
+        self._single_fit = single_fit
         self._model_kwargs = model_kwargs or {}
         self._align_kwargs = kwargs or {}
 
@@ -115,11 +117,16 @@ def run(self, dataset: DatasetT, **kwargs) -> Self:
         # Initialize model
         if isinstance(self._model, str):
             # Factory creates the appropriate model and pipes arguments
-            self._model = ModelFactory.init(
+            model = ModelFactory.init(
                 model=self._model,
                 dataset=dataset,
                 **self._model_kwargs,
             )
+        else:
+            model = self._model
+
+        if self._single_fit:
+            model.fit_predict(None, n_jobs=n_jobs)
 
         kwargs["num_threads"] = kwargs.pop("omp_nthreads", None) or kwargs.pop("num_threads", None)
         kwargs = self._align_kwargs | kwargs
@@ -145,7 +152,7 @@ def run(self, dataset: DatasetT, **kwargs) -> Self:
 
                     # fit the model
                     test_set = dataset[i]
-                    predicted = self._model.fit_predict(  # type: ignore[union-attr]
+                    predicted = model.fit_predict(  # type: ignore[union-attr]
                         i,
                         n_jobs=n_jobs,
                     )

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]