REF: Refactor PET model

Refactor PET model: use a base class that contains the essential
properties for a PET model and create a derived class that implements
the B-Spline correction.

- Remove the unnecessary explicit `float` casting around the number of
  control points: the `dtype="float32"` specifier creates a float array.
  Fixes:
  ```
  Expected type 'str | Buffer | SupportsFloat | SupportsIndex', got 'Literal[0] | None | {__eq__} | int' instead
  ```
  raised by the IDE.

Author: jhlegarreta

docs/notebooks/pet_motion_estimation.ipynb

@@ -406,9 +406,9 @@
     }
    ],
    "source": [
-    "from nifreeze.model import PETModel\n",
+    "from nifreeze.model import BSplinePETModel\n",
     "\n",
-    "model = PETModel(dataset=pet_dataset, timepoints=pet_dataset.midframe, xlim=7000)"
+    "model = BSplinePETModel(dataset=pet_dataset, timepoints=pet_dataset.midframe, xlim=7000)"
    ]
   },
   {
@@ -2258,7 +2258,7 @@
     "from nifreeze.estimator import PETMotionEstimator\n",
     "\n",
     "# Instantiate with a PETModel or appropriate model instance\n",
-    "model = PETModel(\n",
+    "model = BSplinePETModel(\n",
     "    dataset=pet_dataset, timepoints=pet_dataset.midframe, xlim=pet_dataset.total_duration\n",
     ")\n",
     "estimator = PETMotionEstimator(model=model)\n",

src/nifreeze/estimator.py

@@ -40,7 +40,7 @@
 from nifreeze.data.base import BaseDataset
 from nifreeze.data.pet import PET
 from nifreeze.model.base import BaseModel, ModelFactory
-from nifreeze.model.pet import PETModel
+from nifreeze.model.pet import BSplinePETModel
 from nifreeze.registration.ants import (
     Registration,
     _prepare_registration_data,
@@ -255,8 +255,8 @@ def run(self, pet_dataset, omp_nthreads=None):
                     total_duration=pet_dataset.total_duration,
                 )
 
-                # Instantiate PETModel explicitly
-                model = PETModel(
+                # Instantiate the PET model explicitly
+                model = BSplinePETModel(
                     dataset=train_dataset,
                     timepoints=train_times,
                     xlim=pet_dataset.total_duration,

src/nifreeze/model/__init__.py

@@ -33,7 +33,7 @@
     DTIModel,
     GPModel,
 )
-from nifreeze.model.pet import PETModel
+from nifreeze.model.pet import BSplinePETModel
 
 __all__ = (
     "ModelFactory",
@@ -43,5 +43,5 @@
     "DTIModel",
     "GPModel",
     "TrivialModel",
-    "PETModel",
+    "BSplinePETModel",
 )

src/nifreeze/model/pet.py

@@ -22,7 +22,9 @@
 #
 """Models for nuclear imaging."""
 
+from importlib import import_module
 from os import cpu_count
+from typing import Any
 
 import nibabel as nb
 import numpy as np
@@ -31,60 +33,86 @@
 from scipy.interpolate import BSpline
 from scipy.sparse.linalg import cg
 
+from nifreeze.data.pet import PET
 from nifreeze.model.base import BaseModel
 
 DEFAULT_TIMEFRAME_MIDPOINT_TOL = 1e-2
 """Time frame tolerance in seconds."""
 
 
-class PETModel(BaseModel):
-    """A PET imaging realignment model based on B-Spline approximation."""
+def _exec_fit(model, data, chunk=None, **kwargs):
+    return model.fit(data, **kwargs), chunk
 
-    __slots__ = (
-        "_t",
-        "_x",
-        "_xlim",
-        "_order",
-        "_n_ctrl",
-        "_datashape",
-        "_mask",
-        "_smooth_fwhm",
-        "_thresh_pct",
-    )
+
+def _exec_predict(model, chunk=None, **kwargs):
+    """Propagate model parameters and call predict."""
+    return np.squeeze(model.predict(**kwargs)), chunk
+
+
+class BasePETModel(BaseModel):
+    """Interface and default methods for PET models."""
+
+    __slots__ = {
+        "_data_mask": "A mask for the voxels that will be fitted and predicted",
+        "_x" : "",
+        "_xlim" : "",
+        "_smooth_fwhm" : "FWHM in mm over which to smooth",
+        "_thresh_pct" : "Thresholding percentile for the signal",
+        "_model_class": "Defining a model class",
+        "_modelargs": "Arguments acceptable by the underlying model",
+        "_models": "List with one or more (if parallel execution) model instances",
+    }
 
     def __init__(
         self,
-        dataset,
-        timepoints=None,
-        xlim=None,
-        n_ctrl=None,
-        order=3,
-        smooth_fwhm=10,
-        thresh_pct=20,
+        dataset: PET,
+        timepoints: list | np.ndarray = None,  ## Is there a way to use array-like
+        xlim: list | np.ndarray = None,
+        smooth_fwhm: float = 10.0,
+        thresh_pct : float = 20.0,
         **kwargs,
     ):
-        """
-        Create the B-Spline interpolating matrix.
+        """Initialization.
 
-        Parameters:
-        -----------
-        timepoints : :obj:`list`
+        Parameters
+        ----------
+        timepoints : :obj:`list` or :obj:`~np.ndarray`
             The timing (in sec) of each PET volume.
             E.g., ``[15.,   45.,   75.,  105.,  135.,  165.,  210.,  270.,  330.,
             420.,  540.,  750., 1050., 1350., 1650., 1950., 2250., 2550.]``
-
-        n_ctrl : :obj:`int`
-            Number of B-Spline control points. If `None`, then one control point every
-            six timepoints will be used. The less control points, the smoother is the
-            model.
-
+        xlim : .
+            .
+        smooth_fwhm : obj:`float`
+            FWHM in mm over which to smooth the signal.
+        thresh_pct : obj:`float`
+            Thresholding percentile for the signal.
         """
+
         super().__init__(dataset, **kwargs)
 
+        # Duck typing, instead of explicitly testing for PET type
+        if not hasattr(dataset, "total_duration"):
+            raise TypeError("Dataset MUST be a PET object.")
+
+        if not hasattr(dataset, "midframe"):
+            raise ValueError("Dataset MUST have a midframe.")
+
+        # ToDO
+        # Are the timepoints your "gradients" ??? If so, can they be computed
+        # from frame_time or frame_duration
+        # Or else frame_time and frame_duration ????
+
+        self._data_mask = (
+            dataset.brainmask
+            if dataset.brainmask is not None
+            else np.ones(dataset.dataobj.shape[:3], dtype=bool)
+        )
+
+        # ToDo
+        # Are timepoints and xlim features that all PET models require ??
         if timepoints is None or xlim is None:
-            raise TypeError("timepoints must be provided in initialization")
+            raise ValueError("`timepoints` and `xlim` must be specified and have a nonzero value.")
 
-        self._order = order
         self._x = np.array(timepoints, dtype="float32")
         self._xlim = xlim
         self._smooth_fwhm = smooth_fwhm
@@ -95,14 +123,7 @@ def __init__(
         if self._x[-1] > (self._xlim - DEFAULT_TIMEFRAME_MIDPOINT_TOL):
             raise ValueError("Last frame midpoint should not be equal or greater than duration")
 
-        # Calculate index coordinates in the B-Spline grid
-        self._n_ctrl = n_ctrl or (len(timepoints) // 4) + 1
-
-        # B-Spline knots
-        self._t = np.arange(-3, float(self._n_ctrl) + 4, dtype="float32")
-
-        self._datashape = None
-        self._mask = None
+        super().__init__(dataset, **kwargs)
 
     @property
     def is_fitted(self):
@@ -111,17 +132,24 @@ def is_fitted(self):
     def _fit(self, index: int | None = None, n_jobs=None, **kwargs):
         """Fit the model."""
 
+        n_jobs = n_jobs or 1
+
         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")
+
+        # ToDo
+        # Does not make sense to make timepoints be a kwarg if it is provided as a named parameter to __init__
         timepoints = kwargs.get("timepoints", None) or self._x
-        x = (np.array(timepoints, dtype="float32") / self._xlim) * self._n_ctrl
 
+        # ToDo
+        # data, _, gtab = self._dataset[idxmask]  ### This needs the PET data model to be changed
         data = self._dataset.dataobj
         brainmask = self._dataset.brainmask
 
+        # Preprocess the data
         if self._smooth_fwhm > 0:
             smoothed_img = smooth_image(
                 nb.Nifti1Image(data, self._dataset.affine), self._smooth_fwhm
@@ -135,6 +163,135 @@ def _fit(self, index: int | None = None, n_jobs=None, **kwargs):
         # Convert data into V (voxels) x T (timepoints)
         data = data.reshape((-1, data.shape[-1])) if brainmask is None else data[brainmask]
 
+        # ToDo
+        # What is the gtab equivalent of PET ?
+        model_str = getattr(self, "_model_class", "")
+        module_name, class_name = model_str.rsplit(".", 1)
+        model = getattr(
+                import_module(module_name),
+                class_name,
+            )(gtab, **kwargs)
+
+        fit_kwargs: dict[str, Any] = {}  # Add here keyword arguments
+
+        # Split data into chunks of group of slices
+        data_chunks = np.array_split(data, n_jobs)
+
+        self._models = [None] * n_jobs
+
+        # Parallelize process with joblib
+        with Parallel(n_jobs=n_jobs) as executor:
+            results = executor(
+                delayed(_exec_fit)(model, dchunk, i, **fit_kwargs)
+                for i, dchunk in enumerate(data_chunks)
+            )
+        for submodel, rindex in results:
+            self._models[rindex] = submodel
+
+        return n_jobs
+
+
+    def fit_predict(self, index: int | None = None, **kwargs):
+        """Return the corrected volume using B-spline interpolation."""
+
+        n_models = self._fit(
+            index,
+            n_jobs=kwargs.pop("n_jobs"),
+            **kwargs,
+        )
+
+        if index is None:  # If no index, just fit the data.
+            return None
+
+        # ToDo
+        # What are the gtab (and S0 if any) equivalent of PET ?
+        if n_models == 1:
+            predicted, _ = _exec_predict(
+                self._models[0], **(kwargs | {"gtab": gradient, "S0": self._S0})
+            )
+        else:
+            predicted = [None] * n_models
+            S0 = np.array_split(self._S0, n_models)
+
+            # Parallelize process with joblib
+            with Parallel(n_jobs=n_models) as executor:
+                results = executor(
+                    delayed(_exec_predict)(
+                        model,
+                        chunk=i,
+                        **(kwargs | {"gtab": gradient, "S0": S0[i]}),
+                    )
+                    for i, model in enumerate(self._models)
+                )
+            for subprediction, index in results:
+                predicted[index] = subprediction
+
+            predicted = np.hstack(predicted)
+
+        retval = np.zeros_like(self._data_mask, dtype=self._dataset.dataobj.dtype)
+        retval[self._data_mask, ...] = predicted
+        return retval
+
+
+class BSplinePETModel(BasePETModel):
+    """A PET imaging realignment model based on B-Spline approximation."""
+
+    __slots__ = (
+        "_t",
+        "_order",
+        "_n_ctrl",
+    )
+
+    def __init__(
+        self,
+        dataset: PET,
+        n_ctrl: int = None,
+        order: int = 3,
+        **kwargs,
+    ):
+        """Create the B-Spline interpolating matrix.
+
+        Parameters
+        ----------
+        n_ctrl : :obj:`int`
+            Number of B-Spline control points. If `None`, then one control point every
+            six timepoints will be used. The less control points, the smoother is the
+            model.
+        order : :obj:`int`
+            Order of the B-Spline approximation.
+        """
+
+        super().__init__(dataset, **kwargs)
+
+        self._order = order
+
+        # Calculate index coordinates in the B-Spline grid
+        self._n_ctrl = n_ctrl or (len(self._x) // 4) + 1
+
+        # B-Spline knots
+        self._t = np.arange(-3, self._n_ctrl + 4, dtype="float32")
+
+
+    @property
+    def is_fitted(self):
+        return self._locked_fit is not None
+
+    def _fit(self, index: int | None = None, n_jobs=None, **kwargs):
+        """Fit the model."""
+
+        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")
+
+        # ToDo
+        # Does not make sense to make timepoints be a kwarg if it is provided as a named parameter to __init__
+        timepoints = kwargs.get("timepoints", None) or self._x
+        x = (np.array(timepoints, dtype="float32") / self._xlim) * self._n_ctrl
+
+        data = self._dataset.dataobj
+
         # A.shape = (T, K - 4); T= n. timepoints, K= n. knots (with padding)
         A = BSpline.design_matrix(x, self._t, k=self._order)
         AT = A.T
@@ -149,6 +306,12 @@ def _fit(self, index: int | None = None, n_jobs=None, **kwargs):
     def fit_predict(self, index: int | None = None, **kwargs):
         """Return the corrected volume using B-spline interpolation."""
 
+        # ToDo
+        # Does the below apply to PET ? Martin has the return None statement
+        # if index is None:
+        #    raise RuntimeError(
+        #        f"Model {self.__class__.__name__} does not allow locking.")
+
         # Fit the BSpline basis on all data
         if self._locked_fit is None:
             self._fit(index, n_jobs=kwargs.pop("n_jobs", None), **kwargs)