enh: base implementation of B-Spline transforms

oesteban · oesteban · commit 4ff32acf23ea · 2022-02-23T17:13:43.000+01:00
diff --git a/nitransforms/interp/__init__.py b/nitransforms/interp/__init__.py
diff --git a/nitransforms/interp/bspline.py b/nitransforms/interp/bspline.py
@@ -0,0 +1,95 @@
+# emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*-
+# vi: set ft=python sts=4 ts=4 sw=4 et:
+### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
+#
+#   See COPYING file distributed along with the NiBabel package for the
+#   copyright and license terms.
+#
+### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
+"""Interpolate with 3D tensor-product B-Spline basis."""
+import numpy as np
+import nibabel as nb
+from scipy.sparse import csr_matrix, kron
+
+
+def _cubic_bspline(d, order=3):
+    """Evaluate the cubic bspline at distance d from the center."""
+    if order != 3:
+        raise NotImplementedError
+
+    return np.piecewise(
+        d,
+        [d < 1.0, d >= 1.0],
+        [
+            lambda d: (4.0 - 6.0 * d ** 2 + 3.0 * d ** 3) / 6.0,
+            lambda d: (2.0 - d) ** 3 / 6.0,
+        ],
+    )
+
+
+def grid_bspline_weights(target_nii, ctrl_nii):
+    r"""
+    Evaluate tensor-product B-Spline weights on a grid.
+
+    For each of the *N* input samples :math:`(s_1, s_2, s_3)` and *K* control
+    points or *knots* :math:`\mathbf{k} =(k_1, k_2, k_3)`, the tensor-product
+    cubic B-Spline kernel weights are calculated:
+
+    .. math::
+        \Psi^3(\mathbf{k}, \mathbf{s}) =
+        \beta^3(s_1 - k_1) \cdot \beta^3(s_2 - k_2) \cdot \beta^3(s_3 - k_3),
+        \label{eq:1}\tag{1}
+
+    where each :math:`\beta^3` represents the cubic B-Spline for one dimension.
+    The 1D B-Spline kernel implementation uses :obj:`numpy.piecewise`, and is based on the
+    closed-form given by Eq. (6) of [Unser1999]_.
+
+    By iterating over dimensions, the data samples that fall outside of the compact
+    support of the tensor-product kernel associated to each control point can be filtered
+    out and dismissed to lighten computation.
+
+    Finally, the resulting weights matrix :math:`\Psi^3(\mathbf{k}, \mathbf{s})`
+    can be easily identified in Eq. :math:`\eqref{eq:1}` and used as the design matrix
+    for approximation of data.
+
+    Parameters
+    ----------
+    target_nii :  :obj:`nibabel.spatialimages`
+        An spatial image object (typically, a :obj:`~nibabel.nifti1.Nifti1Image`)
+        embedding the target EPI image to be corrected.
+        Provides the location of the *N* samples (total number of voxels) in the space.
+    ctrl_nii : :obj:`nibabel.spatialimages`
+        An spatial image object (typically, a :obj:`~nibabel.nifti1.Nifti1Image`)
+        embedding the location of the control points of the B-Spline grid.
+        The data array should contain a total of :math:`K` knots (control points).
+
+    Returns
+    -------
+    weights : :obj:`numpy.ndarray` (:math:`K \times N`)
+        A sparse matrix of interpolating weights :math:`\Psi^3(\mathbf{k}, \mathbf{s})`
+        for the *N* voxels of the target EPI, for each of the total *K* knots.
+        This sparse matrix can be directly used as design matrix for the fitting
+        step of approximation/extrapolation.
+
+    """
+    shape = target_nii.shape[:3]
+    ctrl_sp = ctrl_nii.header.get_zooms()[:3]
+    ras2ijk = np.linalg.inv(ctrl_nii.affine)
+    origin = nb.affines.apply_affine(ras2ijk, [tuple(target_nii.affine[:3, 3])])[0]
+
+    wd = []
+    for i, (o, n, sp) in enumerate(
+        zip(origin, shape, target_nii.header.get_zooms()[:3])
+    ):
+        locations = np.arange(0, n, dtype="float16") * sp / ctrl_sp[i] + o
+        knots = np.arange(0, ctrl_nii.shape[i], dtype="float16")
+        distance = np.abs(locations[np.newaxis, ...] - knots[..., np.newaxis])
+
+        within_support = distance < 2.0
+        d_vals, d_idxs = np.unique(distance[within_support], return_inverse=True)
+        bs_w = _cubic_bspline(d_vals)
+        weights = np.zeros_like(distance, dtype="float32")
+        weights[within_support] = bs_w[d_idxs]
+        wd.append(csr_matrix(weights))
+
+    return kron(kron(wd[0], wd[1]), wd[2])
diff --git a/nitransforms/nonlinear.py b/nitransforms/nonlinear.py
@@ -8,12 +8,21 @@
 ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
 """Nonlinear transforms."""
 import warnings
+from pathlib import Path
 import numpy as np
-from .base import TransformBase
-from . import io
+from scipy.sparse import sparse_vstack
+from scipy import ndimage as ndi
+from nibabel.funcs import four_to_three
+from nibabel.loadsave import load as _nbload
 
-# from .base import ImageGrid
-# from nibabel.funcs import four_to_three
+from . import io
+from .interp.bspline import grid_bspline_weights
+from .base import (
+    TransformBase,
+    ImageGrid,
+    SpatialReference,
+    _as_homogeneous,
+)
 
 
 class DisplacementsFieldTransform(TransformBase):
@@ -90,70 +99,118 @@ def from_filename(cls, filename, fmt="X5"):
 
 load = DisplacementsFieldTransform.from_filename
 
-# class BSplineFieldTransform(TransformBase):
-#     """Represent a nonlinear transform parameterized by BSpline basis."""
-
-#     __slots__ = ['_coeffs', '_knots', '_refknots', '_order', '_moving']
-#     __s = (slice(None), )
-
-#     def __init__(self, reference, coefficients, order=3):
-#         """Create a smooth deformation field using B-Spline basis."""
-#         super(BSplineFieldTransform, self).__init__()
-#         self._order = order
-#         self.reference = reference
-
-#         if coefficients.shape[-1] != self.ndim:
-#             raise ValueError(
-#                 'Number of components of the coefficients does '
-#                 'not match the number of dimensions')
-
-#         self._coeffs = np.asanyarray(coefficients.dataobj)
-#         self._knots = ImageGrid(four_to_three(coefficients)[0])
-#         self._cache_moving()
-
-#     def _cache_moving(self):
-#         self._moving = np.zeros((self.reference.shape) + (3, ),
-#                                 dtype='float32')
-#         ijk = np.moveaxis(self.reference.ndindex, 0, -1).reshape(-1, self.ndim)
-#         xyz = np.moveaxis(self.reference.ndcoords, 0, -1).reshape(-1, self.ndim)
-#         print(np.shape(xyz))
-
-#         for i in range(np.shape(xyz)[0]):
-#             print(i, xyz[i, :])
-#             self._moving[tuple(ijk[i]) + self.__s] = self._interp_transform(xyz[i, :])
-
-#     def _interp_transform(self, coords):
-#         # Calculate position in the grid of control points
-#         knots_ijk = self._knots.inverse.dot(np.hstack((coords, 1)))[:3]
-#         neighbors = []
-#         offset = 0.0 if self._order & 1 else 0.5
-#         # Calculate neighbors along each dimension
-#         for dim in range(self.ndim):
-#             first = int(np.floor(knots_ijk[dim] + offset) - self._order // 2)
-#             neighbors.append(list(range(first, first + self._order + 1)))
-
-#         # Get indexes of the neighborings clique
-#         ndindex = np.moveaxis(
-#             np.array(np.meshgrid(*neighbors, indexing='ij')), 0, -1).reshape(
-#             -1, self.ndim)
-
-#         # Calculate the tensor B-spline weights of each neighbor
-#         # weights = np.prod(vbspl(ndindex - knots_ijk), axis=-1)
-#         ndindex = [tuple(v) for v in ndindex]
-
-#         # Retrieve coefficients and deal with boundary conditions
-#         zero = np.zeros(self.ndim)
-#         shape = np.array(self._knots.shape)
-#         coeffs = []
-#         for ijk in ndindex:
-#             offbounds = (zero > ijk) | (shape <= ijk)
-#             coeffs.append(
-#                 self._coeffs[ijk] if not np.any(offbounds)
-#                 else [0.0] * self.ndim)
-
-#         # coords[:3] += weights.dot(np.array(coeffs, dtype=float))
-#         return self.reference.inverse.dot(np.hstack((coords, 1)))[:3]
-
-#     def _map_voxel(self, index, moving=None):
-#         """Apply ijk' = f_ijk((i, j, k)), equivalent to the above with indexes."""
-#         return tuple(self._moving[index + self.__s])
+
+class BSplineFieldTransform(TransformBase):
+    """Represent a nonlinear transform parameterized by BSpline basis."""
+
+    __slots__ = ['_coeffs', '_knots', '_weights', '_order', '_moving']
+    __s = (slice(None), )
+
+    def __init__(self, reference, coefficients, order=3):
+        """Create a smooth deformation field using B-Spline basis."""
+        super(BSplineFieldTransform, self).__init__()
+        self._order = order
+        self.reference = reference
+
+        if coefficients.shape[-1] != self.ndim:
+            raise ValueError(
+                'Number of components of the coefficients does '
+                'not match the number of dimensions')
+
+        self._coeffs = np.asanyarray(coefficients.dataobj)
+        self._knots = ImageGrid(four_to_three(coefficients)[0])
+        self._weights = None
+
+    def apply(
+        self,
+        spatialimage,
+        reference=None,
+        order=3,
+        mode="constant",
+        cval=0.0,
+        prefilter=True,
+        output_dtype=None,
+    ):
+        """Apply a B-Spline transform on input data."""
+
+        if reference is not None and isinstance(reference, (str, Path)):
+            reference = _nbload(str(reference))
+
+        _ref = (
+            self.reference if reference is None else SpatialReference.factory(reference)
+        )
+
+        if isinstance(spatialimage, (str, Path)):
+            spatialimage = _nbload(str(spatialimage))
+
+        if not isinstance(_ref, ImageGrid):
+            return super().apply(
+                spatialimage,
+                reference=reference,
+                order=order,
+                mode=mode,
+                cval=cval,
+                prefilter=prefilter,
+                output_dtype=output_dtype,
+            )
+
+        # If locations to be interpolated are on a grid, use faster tensor-bspline calculation
+        if self._weights is None:
+            self._weights = grid_bspline_weights(_ref, self._knots)
+
+        ycoords = _ref.ndcoords.T + (
+            np.squeeze(np.hstack(self._coeffs).T) @ sparse_vstack(self._weights)
+        )
+
+        data = np.squeeze(np.asanyarray(spatialimage.dataobj))
+        output_dtype = output_dtype or data.dtype
+        targets = ImageGrid(spatialimage).index(  # data should be an image
+            _as_homogeneous(np.vstack(ycoords), dim=_ref.ndim)
+        )
+
+        if data.ndim == 4:
+            if len(self) != data.shape[-1]:
+                raise ValueError(
+                    "Attempting to apply %d transforms on a file with "
+                    "%d timepoints" % (len(self), data.shape[-1])
+                )
+            targets = targets.reshape((len(self), -1, targets.shape[-1]))
+            resampled = np.stack(
+                [
+                    ndi.map_coordinates(
+                        data[..., t],
+                        targets[t, ..., : _ref.ndim].T,
+                        output=output_dtype,
+                        order=order,
+                        mode=mode,
+                        cval=cval,
+                        prefilter=prefilter,
+                    )
+                    for t in range(data.shape[-1])
+                ],
+                axis=0,
+            )
+        elif data.ndim in (2, 3):
+            resampled = ndi.map_coordinates(
+                data,
+                targets[..., : _ref.ndim].T,
+                output=output_dtype,
+                order=order,
+                mode=mode,
+                cval=cval,
+                prefilter=prefilter,
+            )
+
+        newdata = resampled.reshape((len(self), *_ref.shape))
+        moved = spatialimage.__class__(
+            np.moveaxis(newdata, 0, -1), _ref.affine, spatialimage.header
+        )
+        moved.header.set_data_dtype(output_dtype)
+        return moved
+
+    def map(self, x, inverse=False):
+        raise NotImplementedError
+
+    def _map_voxel(self, index, moving=None):
+        """Apply ijk' = f_ijk((i, j, k)), equivalent to the above with indexes."""
+        raise NotImplementedError
