REF: Refactor PET model

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)"
    ]
   },
   {
@@ -429,7 +429,7 @@
    "outputs": [],
    "source": [
     "index = 2\n",
-    "predicted = model.fit_predict(pet_dataset.midframe[index])"
+    "predicted = model.fit_predict(index)"
    ]
   },
   {

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,
@@ -234,7 +234,7 @@ def __init__(self, align_kwargs: dict | None = None, strategy: str = "lofo"):
 
     def run(self, pet_dataset: PET, omp_nthreads: int | None = None) -> list:
         n_frames = len(pet_dataset)
-        frame_indices = np.arange(n_frames)
+        frame_indices = np.arange(n_frames).astype(int)
 
         if omp_nthreads:
             self.align_kwargs["num_threads"] = omp_nthreads
@@ -261,18 +261,14 @@ def run(self, pet_dataset: PET, omp_nthreads: int | None = None) -> list:
                     total_duration=pet_dataset.total_duration,
                 )
 
-                # Instantiate PETModel explicitly
-                model = PETModel(
-                    dataset=train_dataset,
-                    timepoints=train_times,
-                    xlim=pet_dataset.total_duration,
-                )
+                # Instantiate the PET model explicitly
+                model = BSplinePETModel(dataset=train_dataset)
 
                 # Fit the model once on the training dataset
                 model.fit_predict(None)
 
                 # Predict the reference volume at the test frame's timepoint
-                predicted = model.fit_predict(test_time)
+                predicted = model.fit_predict(idx)
 
                 fixed_image_path = tmp_path / f"fixed_frame_{idx:03d}.nii.gz"
                 moving_image_path = tmp_path / f"moving_frame_{idx:03d}.nii.gz"

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,6 +22,7 @@
 #
 """Models for nuclear imaging."""
 
+from abc import ABC, ABCMeta, abstractmethod
 from os import cpu_count
 from typing import Union
 
@@ -35,142 +36,183 @@
 from nifreeze.data.pet import PET
 from nifreeze.model.base import BaseModel
 
-TIMEPOINT_XLIM_DATA_MISSING_ERROR_MSG = """\
-'timepoints' and 'xlim' must be specified, found: {timepoints} and {xlim}."""
-"""PET model underspecification error."""
+PET_OBJECT_ERROR_MSG = "Dataset MUST be a PET object."
+"""PET object error message."""
 
-FIRST_TIMEPOINT_VALUE_ERROR_MSG = """\
-First frame 'timepoint' should not be zero or negative, found: {timepoints}."""
-"""PET model timepoint value error message."""
-
-LAST_TIMEPOINT_CONSISTENCY_ERROR_MSG = """\
-Last frame 'timepoints' value should not be equal or greater than 'xlim' \
-duration, found: {timepoints} and {xlim}."""
-"""PET model parameter consistency error message."""
+PET_MIDFRAME_ERROR_MSG = "Dataset MUST have a 'midframe'."
+"""PET midframe error message."""
 
 DEFAULT_TIMEPOINT_TOL = 1e-2
 """Time frame tolerance in seconds."""
 
 
-class PETModel(BaseModel):
+def _exec_fit(model, data, chunk=None, **kwargs):
+    return model.fit(data, **kwargs), chunk
+
+
+def _exec_predict(model, chunk=None, **kwargs):
+    """Propagate model parameters and call predict."""
+    return np.squeeze(model.predict(**kwargs)), chunk
+
+
+class BasePETModel(BaseModel, ABC):
+    """Interface and default methods for PET models."""
+
+    __metaclass__ = ABCMeta
+
+    __slots__ = {
+        "_data_mask": "A mask for the voxels that will be fitted and predicted",
+        "_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: PET,
+        smooth_fwhm: float = 10.0,
+        thresh_pct: float = 20.0,
+        **kwargs,
+    ):
+        """Initialization.
+
+        Parameters
+        ----------
+        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(PET_OBJECT_ERROR_MSG)
+
+        if not hasattr(dataset, "midframe"):
+            raise ValueError(PET_MIDFRAME_ERROR_MSG)
+
+        self._data_mask = (
+            dataset.brainmask
+            if dataset.brainmask is not None
+            else np.ones(dataset.dataobj.shape[:3], dtype=bool)
+        )
+
+        self._smooth_fwhm = smooth_fwhm
+        self._thresh_pct = thresh_pct
+
+    def _preprocess_data(self) -> np.ndarray:
+        # 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
+            )
+            data = smoothed_img.get_fdata()
+
+        if self._thresh_pct > 0:
+            thresh_val = np.percentile(data, self._thresh_pct)
+            data[data < thresh_val] = 0
+
+        # Convert data into V (voxels) x T (timepoints)
+        return data.reshape((-1, data.shape[-1])) if brainmask is None else data[brainmask]
+
+    @property
+    def is_fitted(self) -> bool:
+        return self._locked_fit is not None
+
+    @abstractmethod
+    def fit_predict(self, index: int | None = None, **kwargs) -> Union[np.ndarray, None]:
+        """Predict the corrected volume."""
+        return None
+
+
+class BSplinePETModel(BasePETModel):
     """A PET imaging realignment model based on B-Spline approximation."""
 
     __slots__ = (
         "_t",
-        "_x",
-        "_xlim",
         "_order",
         "_n_ctrl",
-        "_datashape",
-        "_mask",
-        "_smooth_fwhm",
-        "_thresh_pct",
     )
 
     def __init__(
         self,
         dataset: PET,
-        timepoints: list | np.ndarray,
-        xlim: float,
         n_ctrl: int | None = None,
         order: int = 3,
-        smooth_fwhm: float = 10.0,
-        thresh_pct: float = 20.0,
         **kwargs,
     ):
-        """
-        Create the B-Spline interpolating matrix.
-
-        Parameters:
-        -----------
-        timepoints : :obj:`list`
-            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.]``
+        """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)
-
-        if timepoints is None or xlim is None:
-            raise ValueError(
-                TIMEPOINT_XLIM_DATA_MISSING_ERROR_MSG.format(timepoints=timepoints, xlim=xlim)
-            )
 
-        if timepoints[0] < DEFAULT_TIMEPOINT_TOL:
-            raise ValueError(FIRST_TIMEPOINT_VALUE_ERROR_MSG.format(timepoints=timepoints[0]))
-
-        if timepoints[-1] > xlim - DEFAULT_TIMEPOINT_TOL:
-            raise ValueError(
-                LAST_TIMEPOINT_CONSISTENCY_ERROR_MSG.format(timepoints=timepoints, xlim=xlim)
-            )
+        super().__init__(dataset, **kwargs)
 
         self._order = order
-        self._x = np.array(timepoints, dtype="float32")
-        self._xlim = xlim
-        self._smooth_fwhm = smooth_fwhm
-        self._thresh_pct = thresh_pct
 
         # Calculate index coordinates in the B-Spline grid
-        self._n_ctrl = n_ctrl or (len(timepoints) // 4) + 1
+        self._n_ctrl = n_ctrl or (len(self._dataset.midframe) // 4) + 1
 
         # B-Spline knots
-        self._t = np.arange(-3, float(self._n_ctrl) + 4, dtype="float32")
-
-        self._datashape = None
-        self._mask = None
-
-    @property
-    def is_fitted(self) -> bool:
-        return self._locked_fit is not None
+        self._t = np.arange(-3, self._n_ctrl + 4, dtype="float32")
 
     def _fit(self, index: int | None = None, n_jobs=None, **kwargs) -> int:
         """Fit the model."""
 
+        n_jobs = n_jobs or min(cpu_count() or 1, 8)
+
         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.asarray((np.array(timepoints, dtype="float32") / self._xlim) * self._n_ctrl)
-
-        data = self._dataset.dataobj
-        brainmask = self._dataset.brainmask
-
-        if self._smooth_fwhm > 0:
-            smoothed_img = smooth_image(
-                nb.Nifti1Image(data, self._dataset.affine), self._smooth_fwhm
-            )
-            data = smoothed_img.get_fdata()
 
-        if self._thresh_pct > 0:
-            thresh_val = np.percentile(data, self._thresh_pct)
-            data[data < thresh_val] = 0
+        data = self._preprocess_data()
 
-        # Convert data into V (voxels) x T (timepoints)
-        data = data.reshape((-1, data.shape[-1])) if brainmask is None else data[brainmask]
+        x = (
+            np.asarray(self._dataset.midframe, dtype="float32")
+            / self._dataset.total_duration
+            * self._n_ctrl
+        )
 
         # 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
         ATdotA = AT @ A
 
         # Parallelize process with joblib
-        with Parallel(n_jobs=n_jobs or min(cpu_count() or 1, 8)) as executor:
+        with Parallel(n_jobs=n_jobs) as executor:
             results = executor(delayed(cg)(ATdotA, AT @ v) for v in data)
 
-        self._locked_fit = np.array([r[0] for r in results])
+        self._locked_fit = np.asarray([r[0] for r in results])
 
         return n_jobs
 
     def fit_predict(self, index: int | None = None, **kwargs) -> Union[np.ndarray, None]:
         """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)
@@ -183,7 +225,10 @@ def fit_predict(self, index: int | None = None, **kwargs) -> Union[np.ndarray, N
             return None
 
         # Project sample timing into B-Spline coordinates
-        x = np.asarray((index / self._xlim) * self._n_ctrl)
+        # ToDo: x is not really a matrix ...
+        x = np.asarray(
+            (self._dataset.midframe[index] / self._dataset.total_duration) * self._n_ctrl
+        )
         A = BSpline.design_matrix(x, self._t, k=self._order)
 
         # A is 1 (num. timepoints) x C (num. coeff)

test/test_integration_pet.py

@@ -102,15 +102,15 @@ def test_pet_motion_estimator_run(monkeypatch, setup_random_pet_data):
     )
 
     class DummyModel:
-        def __init__(self, dataset, timepoints, xlim):
+        def __init__(self, dataset):
             self.dataset = dataset
 
         def fit_predict(self, index):
             if index is None:
                 return None
             return np.zeros(self.dataset.shape3d, dtype=np.float32)
 
-    monkeypatch.setattr("nifreeze.estimator.PETModel", DummyModel)
+    monkeypatch.setattr("nifreeze.estimator.BSplinePETModel", DummyModel)
 
     class DummyRegistration:
         def __init__(self, *args, **kwargs):