fix: revising map method and tests

Author: oesteban

nitransforms/base.py

@@ -78,14 +78,15 @@ def ndcoords(self):
             )[:3, ...]
         return self._coords
 
-    def index(self, coordinates):
-        """Get the image array's indexes corresponding to coordinates."""
-        coordinates = np.array(coordinates)
-        ncoords = coordinates.shape[-1]
-        coordinates = np.vstack((coordinates, np.ones((1, ncoords))))
+    def ras(self, ijk):
+        """Get RAS+ coordinates from input indexes."""
+        ras = self._affine.dot(_as_homogeneous(ijk).T)[:3, ...]
+        return ras.T
 
-        # Back to grid coordinates
-        return np.tensordot(self._inverse, coordinates, axes=1)[:3, ...]
+    def index(self, x):
+        """Get the image array's indexes corresponding to coordinates."""
+        ijk = self._inverse.dot(_as_homogeneous(x).T)[:3, ...]
+        return ijk.T
 
     def _to_hdf5(self, group):
         group.attrs['Type'] = 'image'
@@ -114,9 +115,9 @@ def __eq__(self, other):
             return False
         return np.allclose(self.matrix, other.matrix, rtol=EQUALITY_TOL)
 
-    def __call__(self, x):
+    def __call__(self, x, inverse=False, index=0):
         """Apply y = f(x)."""
-        return self.map(x)
+        return self.map(x, inverse=inverse, index=index)
 
     @property
     def reference(self):
@@ -171,31 +172,47 @@ def resample(self, moving, order=3, mode='constant', cval=0.0, prefilter=True,
         if output_dtype is None:
             output_dtype = moving_data.dtype
 
-        moving_grid = ImageGrid(moving)
-
-        def _map_indexes(ijk):
-            return moving_grid.inverse.dot(self.map(self.reference.affine.dot(ijk)))
-
         moved = ndi.geometric_transform(
             moving_data,
-            mapping=_map_indexes,
+            mapping=self._map_index,
             output_shape=self.reference.shape,
             output=output_dtype,
             order=order,
             mode=mode,
             cval=cval,
             prefilter=prefilter,
-            extra_keywords={'moving': moving},
+            extra_keywords={'moving': ImageGrid(moving)},
         )
 
         moved_image = moving.__class__(moved, self.reference.affine, moving.header)
         moved_image.header.set_data_dtype(output_dtype)
         return moved_image
 
-    def map(self, x):
-        """Apply y = f(x)."""
+    def map(self, x, inverse=False, index=0):
+        r"""
+        Apply :math:`y = f(x)`.
+
+        Parameters
+        ----------
+        x : N x D numpy.ndarray
+            Input RAS+ coordinates (i.e., physical coordinates).
+        inverse : bool
+            If ``True``, apply the inverse transform :math:`x = f^{-1}(y)`.
+        index : int, optional
+            Transformation index
+
+        Returns
+        -------
+        y : N x D numpy.ndarray
+            Transformed (mapped) RAS+ coordinates (i.e., physical coordinates).
+
+        """
         raise NotImplementedError
 
+    def _map_index(self, ijk, moving):
+        x = self.reference.ras(_as_homogeneous(ijk))
+        return moving.index(self.map(x))
+
     def to_filename(self, filename, fmt='X5'):
         """Store the transform in BIDS-Transforms HDF5 file format (.x5)."""
         with h5py.File(filename, 'w') as out_file:
@@ -209,3 +226,24 @@ def to_filename(self, filename, fmt='X5'):
     def _to_hdf5(self, x5_root):
         """Serialize this object into the x5 file format."""
         raise NotImplementedError
+
+
+def _as_homogeneous(xyz, dtype='float32'):
+    """
+    Convert 2D and 3D coordinates into homogeneous coordinates.
+
+    Examples
+    --------
+    >>> _as_homogeneous((4, 5), dtype='int8').tolist()
+    [[4, 5, 1]]
+
+    >>> _as_homogeneous((4, 5, 6),dtype='int8').tolist()
+    [[4, 5, 6, 1]]
+
+    >>> _as_homogeneous([(1, 2, 3), (4, 5, 6)]).tolist()
+    [[1.0, 2.0, 3.0, 1.0], [4.0, 5.0, 6.0, 1.0]]
+
+
+    """
+    xyz = np.atleast_2d(np.array(xyz, dtype=dtype))
+    return np.hstack((xyz, np.ones((xyz.shape[0], 1), dtype=dtype)))

nitransforms/conftest.py

@@ -39,7 +39,7 @@ def data_path():
 
 
 @pytest.fixture
-def get_data():
+def get_testdata():
     """Generate data in the requested orientation."""
     global _data
 

nitransforms/linear.py

@@ -143,7 +143,7 @@ def resample(self, moving, order=3, mode='constant', cval=0.0, prefilter=True,
             nvols = moving.shape[3]
 
         movaff = moving.affine
-        movingdata = moving.get_data()
+        movingdata = np.asanyarray(moving.dataobj)
         if nvols == 1:
             movingdata = movingdata[..., np.newaxis]
 
@@ -178,39 +178,42 @@ def resample(self, moving, order=3, mode='constant', cval=0.0, prefilter=True,
 
         return moved_image
 
-    def map(self, coords, index=0, forward=True):
-        """
-        Apply y = f(x), where x is the argument `coords`.
+    def map(self, x, inverse=False, index=0):
+        r"""
+        Apply :math:`y = f(x)`.
 
         Parameters
         ----------
-        coords : array_like
-            RAS coordinates to map
+        x : N x D numpy.ndarray
+            Input RAS+ coordinates (i.e., physical coordinates).
+        inverse : bool
+            If ``True``, apply the inverse transform :math:`x = f^{-1}(y)`.
         index : int, optional
             Transformation index
-        forward: bool, optional
-            Direction of mapping. Default is set to ``True``. If ``False``,
-            the inverse transformation is applied.
 
         Returns
         -------
-        out: ndarray
-            Transformed coordinates
+        y : N x D numpy.ndarray
+            Transformed (mapped) RAS+ coordinates (i.e., physical coordinates).
 
         Examples
         --------
         >>> xfm = Affine([[1, 0, 0, 1], [0, 1, 0, 2], [0, 0, 1, 3], [0, 0, 0, 1]])
-        >>> xfm((0,0,0))
+        >>> xfm.map((0,0,0))
         array([1, 2, 3])
 
-        >>> xfm((0,0,0), forward=False)
+        >>> xfm.map((0,0,0), inverse=True)
         array([-1., -2., -3.])
 
         """
-        coords = np.array(coords)
+        coords = np.array(x)
         if coords.shape[0] == self._matrix[index].shape[0] - 1:
             coords = np.append(coords, [1])
-        affine = self._matrix[index] if forward else np.linalg.inv(self._matrix[index])
+        affine = self._matrix[index]
+
+        if inverse is True:
+            affine = np.linalg.inv(self._matrix[index])
+
         return affine.dot(coords)[:-1]
 
     def _map_voxel(self, index, nindex=0, moving=None):

nitransforms/nonlinear.py

@@ -7,6 +7,7 @@
 #
 ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
 """Nonlinear transforms."""
+import sys
 import numpy as np
 from scipy import ndimage as ndi
 # from gridbspline.maths import cubic
@@ -21,12 +22,11 @@ class DeformationFieldTransform(TransformBase):
     """Represents a dense field of displacements (one vector per voxel)."""
 
     __slots__ = ['_field', '_moving', '_moving_space']
-    __s = (slice(None), )
 
     def __init__(self, field, reference=None):
         """Create a dense deformation field transform."""
         super(DeformationFieldTransform, self).__init__()
-        self._field = field.get_data()
+        self._field = np.asanyarray(field.dataobj)
 
         ndim = self._field.ndim - 1
         if len(self._field.shape[:-1]) != ndim:
@@ -51,110 +51,46 @@ def __init__(self, field, reference=None):
 
         self.reference = reference
 
-    def _cache_moving(self, moving):
-        # Check whether input (moving) space is cached
-        moving_space = ImageGrid(moving)
-        if self._moving_space == moving_space:
-            return
+    def map(self, x, inverse=False, index=0):
+        r"""
+        Apply :math:`y = f(x)`.
 
-        # Generate grid of pixel indexes (ijk)
-        ndim = self._field.ndim - 1
-        if ndim == 2:
-            grid = np.meshgrid(
-                np.arange(self._field.shape[0]),
-                np.arange(self._field.shape[1]),
-                indexing='ij')
-        elif ndim == 3:
-            grid = np.meshgrid(
-                np.arange(self._field.shape[0]),
-                np.arange(self._field.shape[1]),
-                np.arange(self._field.shape[2]),
-                indexing='ij')
-        else:
-            raise ValueError('Wrong dimensions (%d)' % ndim)
-
-        grid = np.array(grid)
-        flatgrid = grid.reshape(ndim, -1)
-
-        # Calculate physical coords of all voxels (xyz)
-        flatxyz = np.tensordot(
-            self.reference.affine,
-            np.vstack((flatgrid, np.ones((1, flatgrid.shape[1])))),
-            axes=1
-        )
-
-        # Add field
-        newxyz = flatxyz + np.vstack((
-            np.moveaxis(self._field, -1, 0).reshape(ndim, -1),
-            np.zeros((1, flatgrid.shape[1]))))
-
-        # Back to grid coordinates
-        newijk = np.tensordot(np.linalg.inv(moving.affine),
-                              newxyz, axes=1)
-
-        # Reshape as grid
-        self._moving = np.moveaxis(
-            newijk[0:3, :].reshape((ndim, ) + self._field.shape[:-1]),
-            0, -1)
-
-        self._moving_space = moving_space
+        Parameters
+        ----------
+        x : N x D numpy.ndarray
+            Input RAS+ coordinates (i.e., physical coordinates).
+        inverse : bool
+            If ``True``, apply the inverse transform :math:`x = f^{-1}(y)`.
+        index : int, optional
+            Transformation index
 
-    def resample(self, moving, order=3, mode='constant', cval=0.0, prefilter=True,
-                 output_dtype=None):
-        """
-        Resample the ``moving`` image applying the deformation field.
+        Returns
+        -------
+        y : N x D numpy.ndarray
+            Transformed (mapped) RAS+ coordinates (i.e., physical coordinates).
 
         Examples
         --------
-        >>> ref = nb.load(testfile)
-        >>> refdata = ref.get_fdata()
-        >>> np.allclose(refdata, 0)
-        True
-
-        >>> refdata[5, 5, 5] = 1  # Set a one in the middle voxel
-        >>> moving = nb.Nifti1Image(refdata, ref.affine, ref.header)
-        >>> field = np.zeros(tuple(list(ref.shape) + [3]))
+        >>> field = np.zeros((10, 10, 10, 3))
         >>> field[..., 0] = 4.0
-        >>> fieldimg = nb.Nifti1Image(field, ref.affine, ref.header)
+        >>> fieldimg = nb.Nifti1Image(field, np.diag([2., 2., 2., 1.]))
         >>> xfm = DeformationFieldTransform(fieldimg)
-        >>> resampled = xfm.resample(moving, order=0).get_fdata()
-        >>> resampled[1, 5, 5]
-        1.0
+        >>> xfm([4.0, 4.0, 4.0]).tolist()
+        [[8.0, 4.0, 4.0]]
 
-        """
-        self._cache_moving(moving)
-        return super(DeformationFieldTransform, self).resample(
-            moving, order=order, mode=mode, cval=cval, prefilter=prefilter)
-
-    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])
+        >>> xfm([[4.0, 4.0, 4.0], [8, 2, 10]]).tolist()
+        [[8.0, 4.0, 4.0], [12.0, 2.0, 10.0]]
 
-    def map(self, x, order=3, mode='mirror', cval=0.0, prefilter=True):
-        """Apply y = f(x), where x is the argument `coords`."""
-        coordinates = np.array(x)
-        # Extract shapes and dimensions, then flatten
-        ndim = coordinates.shape[-1]
-        output_shape = coordinates.shape[:-1]
-        flatcoord = np.moveaxis(coordinates, -1, 0).reshape(ndim, -1)
-
-        # Convert coordinates to voxel indices
-        ijk = np.tensordot(
-            np.linalg.inv(self.reference.affine),
-            np.vstack((flatcoord, np.ones((1, flatcoord.shape[1])))),
-            axes=1)
-        deltas = ndi.map_coordinates(
-            self._field,
-            ijk,
-            order=order,
-            mode=mode,
-            cval=cval,
-            prefilter=prefilter)
-
-        deltas = np.moveaxis(deltas[0:3, :].reshape((ndim, ) + output_shape),
-                             0, -1)
-
-        return coordinates + deltas
+        """
+        if inverse is True:
+            raise NotImplementedError
+        ijk = self.reference.index(x)
+        indexes = np.round(ijk).astype('int')
+        if np.any(np.abs(ijk - indexes) > 0.05):
+            print('Some coordinates are off-grid of the displacements field.',
+                  file=sys.stderr)
+        indexes = tuple([tuple(i) for i in indexes.T])
+        return x + self._field[indexes]
 
 
 class BSplineFieldTransform(TransformBase):
@@ -174,7 +110,7 @@ def __init__(self, reference, coefficients, order=3):
                 'Number of components of the coefficients does '
                 'not match the number of dimensions')
 
-        self._coeffs = coefficients.get_data()
+        self._coeffs = np.asanyarray(coefficients.dataobj)
         self._knots = ImageGrid(four_to_three(coefficients)[0])
         self._cache_moving()
 

nitransforms/tests/test_base.py

@@ -6,12 +6,21 @@
 
 
 @pytest.mark.parametrize('image_orientation', ['RAS', 'LAS', 'LPS', 'oblique'])
-def test_ImageGrid(get_data, image_orientation):
+def test_ImageGrid(get_testdata, image_orientation):
     """Check the grid object."""
-    im = get_data[image_orientation]
+    im = get_testdata[image_orientation]
 
     img = ImageGrid(im)
-    assert np.all(img.affine == np.linalg.inv(img.inverse))
+    assert np.allclose(img.affine, np.linalg.inv(img.inverse))
+
+    # Test ras2vox and vox2ras conversions
+    ijk = [[10, 10, 10], [40, 4, 20], [0, 0, 0], [s - 1 for s in im.shape[:3]]]
+    xyz = [img._affine.dot(idx + [1])[:-1] for idx in ijk]
+
+    assert np.allclose(img.ras(ijk[0]), xyz[0])
+    assert np.allclose(np.round(img.index(xyz[0])), ijk[0])
+    assert np.allclose(img.ras(ijk), xyz)
+    assert np.allclose(np.round(img.index(xyz)), ijk)
 
     # nd index / coords
     idxs = img.ndindex