enh: propose a new API for dense fields

Author: oesteban

nitransforms/__init__.py

@@ -18,7 +18,7 @@
 """
 from . import linear, manip, nonlinear
 from .linear import Affine, LinearTransformsMapping
-from .nonlinear import DisplacementsFieldTransform
+from .nonlinear import DenseFieldTransform
 from .manip import TransformChain
 
 try:
@@ -42,7 +42,7 @@
     "nonlinear",
     "Affine",
     "LinearTransformsMapping",
-    "DisplacementsFieldTransform",
+    "DenseFieldTransform",
     "TransformChain",
     "__copyright__",
     "__packagename__",

nitransforms/manip.py

@@ -15,7 +15,7 @@
     TransformError,
 )
 from .linear import Affine
-from .nonlinear import DisplacementsFieldTransform
+from .nonlinear import DenseFieldTransform
 
 
 class TransformChain(TransformBase):
@@ -197,7 +197,7 @@ def from_filename(cls, filename, fmt="X5", reference=None, moving=None):
                 if isinstance(xfmobj, itk.ITKLinearTransform):
                     retval.insert(0, Affine(xfmobj.to_ras(), reference=reference))
                 else:
-                    retval.insert(0, DisplacementsFieldTransform(xfmobj))
+                    retval.insert(0, DenseFieldTransform(xfmobj))
 
             return TransformChain(retval)
 

nitransforms/nonlinear.py

@@ -20,31 +20,43 @@
     ImageGrid,
     SpatialReference,
     _as_homogeneous,
-    EQUALITY_TOL,
 )
 
 
-class DeformationFieldTransform(TransformBase):
-    """Represents a dense field of deformed locations (corresponding to each voxel)."""
+class DenseFieldTransform(TransformBase):
+    """Represents dense field (voxel-wise) transforms."""
 
-    __slots__ = ["_field"]
+    __slots__ = ("_field", "_displacements")
 
-    def __init__(self, field, reference=None):
+    def __init__(self, field=None, displacements=True, reference=None):
         """
-        Create a dense deformation field transform.
+        Create a dense field transform.
+
+        Converting to a field of deformations is straightforward by just adding the corresponding
+        displacement to the :math:`(x, y, z)` coordinates of each voxel.
+        Numerically, deformation fields are less susceptible to rounding errors
+        than displacements fields.
+        SPM generally prefers deformations for that reason.
 
         Example
         -------
-        >>> DeformationFieldTransform(test_dir / "someones_displacement_field.nii.gz")
-        <DeformationFieldTransform[3D] (57, 67, 56)>
+        >>> DenseFieldTransform(test_dir / "someones_displacement_field.nii.gz")
+        <DenseFieldTransform[3D] (57, 67, 56)>
 
         """
+        if field is None and reference is None:
+            raise TransformError("DenseFieldTransforms require a spatial reference")
+
         super().__init__()
 
-        field = _ensure_image(field)
-        self._field = np.squeeze(
-            np.asanyarray(field.dataobj) if hasattr(field, "dataobj") else field
-        )
+        if field is not None:
+            field = _ensure_image(field)
+            self._field = np.squeeze(
+                np.asanyarray(field.dataobj) if hasattr(field, "dataobj") else field
+            )
+        else:
+            self._field = np.zeros((*reference.shape, reference.ndim), dtype="float32")
+            displacements = True
 
         try:
             self.reference = ImageGrid(
@@ -64,6 +76,12 @@ def __init__(self, field, reference=None):
                 "the number of dimensions (%d)" % (self._field.shape[-1], ndim)
             )
 
+        if displacements:
+            self._displacements = self._field
+            # Convert from displacements to deformations fields
+            # (just add the origin to the displacements vector)
+            self._field += self.reference.ndcoords.T.reshape(self._field.shape)
+
     def __repr__(self):
         """Beautify the python representation."""
         return f"<{self.__class__.__name__}[{self._field.shape[-1]}D] {self._field.shape[:3]}>"
@@ -93,13 +111,23 @@ def map(self, x, inverse=False):
 
         Examples
         --------
-        >>> xfm = DeformationFieldTransform(test_dir / "someones_displacement_field.nii.gz")
+        >>> xfm = DenseFieldTransform(
+        ...     test_dir / "someones_displacement_field.nii.gz",
+        ...     displacements=False,
+        ... )
         >>> xfm.map([-6.5, -36., -19.5]).tolist()
         [[0.0, -0.47516798973083496, 0.0]]
 
         >>> xfm.map([[-6.5, -36., -19.5], [-1., -41.5, -11.25]]).tolist()
         [[0.0, -0.47516798973083496, 0.0], [0.0, -0.538356602191925, 0.0]]
 
+        >>> xfm = DenseFieldTransform(
+        ...     test_dir / "someones_displacement_field.nii.gz",
+        ...     displacements=True,
+        ... )
+        >>> xfm.map([[-6.5, -36., -19.5], [-1., -41.5, -11.25]]).tolist()
+        [[-6.5, -36.47516632080078, -19.5], [-1.0, -42.03835678100586, -11.25]]
+
         """
 
         if inverse is True:
@@ -117,25 +145,33 @@ def __matmul__(self, b):
 
         Examples
         --------
-        >>> xfm = DeformationFieldTransform(test_dir / "someones_displacement_field.nii.gz")
-        >>> xfm2 = xfm @ TransformBase()
-        >>> xfm == xfm2
+        >>> deff = DenseFieldTransform(
+        ...     test_dir / "someones_displacement_field.nii.gz",
+        ...     displacements=False,
+        ... )
+        >>> deff2 = deff @ TransformBase()
+        >>> deff == deff2
+        True
+
+        >>> disp = DenseFieldTransform(test_dir / "someones_displacement_field.nii.gz")
+        >>> disp2 = disp @ TransformBase()
+        >>> disp == disp2
         True
 
         """
         retval = b.map(
             self._field.reshape((-1, self._field.shape[-1]))
         ).reshape(self._field.shape)
-        return DeformationFieldTransform(retval, reference=self.reference)
+        return DenseFieldTransform(retval, displacements=False, reference=self.reference)
 
     def __eq__(self, other):
         """
         Overload equals operator.
 
         Examples
         --------
-        >>> xfm1 = DeformationFieldTransform(test_dir / "someones_displacement_field.nii.gz")
-        >>> xfm2 = DeformationFieldTransform(test_dir / "someones_displacement_field.nii.gz")
+        >>> xfm1 = DenseFieldTransform(test_dir / "someones_displacement_field.nii.gz")
+        >>> xfm2 = DenseFieldTransform(test_dir / "someones_displacement_field.nii.gz")
         >>> xfm1 == xfm2
         True
 
@@ -145,41 +181,6 @@ def __eq__(self, other):
             warnings.warn("Fields are equal, but references do not match.")
         return _eq
 
-
-class DisplacementsFieldTransform(DeformationFieldTransform):
-    """
-    Represents a dense field of displacements (one vector per voxel).
-
-    Converting to a field of deformations is straightforward by just adding the corresponding
-    displacement to the :math:`(x, y, z)` coordinates of each voxel.
-    Numerically, deformation fields are less susceptible to rounding errors
-    than displacements fields.
-    SPM generally prefers deformations for that reason.
-
-    """
-
-    __slots__ = ["_displacements"]
-
-    def __init__(self, field, reference=None):
-        """
-        Create a transform supported by a field of voxel-wise displacements.
-
-        Example
-        -------
-        >>> xfm = DisplacementsFieldTransform(test_dir / "someones_displacement_field.nii.gz")
-        >>> xfm
-        <DisplacementsFieldTransform[3D] (57, 67, 56)>
-
-        >>> xfm.map([[-6.5, -36., -19.5], [-1., -41.5, -11.25]]).tolist()
-        [[-6.5, -36.47516632080078, -19.5], [-1.0, -42.03835678100586, -11.25]]
-
-        """
-        super().__init__(field, reference=reference)
-        self._displacements = self._field
-        # Convert from displacements to deformations fields
-        # (just add the origin to the displacements vector)
-        self._field += self.reference.ndcoords.T.reshape(self._field.shape)
-
     @classmethod
     def from_filename(cls, filename, fmt="X5"):
         _factory = {
@@ -193,7 +194,7 @@ def from_filename(cls, filename, fmt="X5"):
         return cls(_factory[fmt].from_filename(filename))
 
 
-load = DisplacementsFieldTransform.from_filename
+load = DenseFieldTransform.from_filename
 
 
 class BSplineFieldTransform(TransformBase):
@@ -239,8 +240,9 @@ def to_field(self, reference=None, dtype="float32"):
             #  1 x Nvox :                          (1 x K) @ (K x Nvox)
             field[:, d] = self._coeffs[..., d].reshape(-1) @ self._weights
 
-        return DisplacementsFieldTransform(
-            field.astype(dtype).reshape(*_ref.shape, -1), reference=_ref)
+        return DenseFieldTransform(
+            field.astype(dtype).reshape(*_ref.shape, -1), reference=_ref
+        )
 
     def apply(
         self,

nitransforms/tests/test_nonlinear.py

@@ -12,7 +12,7 @@
 from nitransforms.io.base import TransformFileError
 from nitransforms.nonlinear import (
     BSplineFieldTransform,
-    DisplacementsFieldTransform,
+    DenseFieldTransform,
     load as nlload,
 )
 from ..io.itk import ITKDisplacementsField
@@ -45,7 +45,7 @@ def test_itk_disp_load(size):
 def test_displacements_bad_sizes(size):
     """Checks field sizes."""
     with pytest.raises(TransformError):
-        DisplacementsFieldTransform(nb.Nifti1Image(np.zeros(size), np.eye(4), None))
+        DenseFieldTransform(nb.Nifti1Image(np.zeros(size), np.eye(4), None))
 
 
 def test_itk_disp_load_intent():
@@ -59,15 +59,15 @@ def test_itk_disp_load_intent():
 
 
 def test_displacements_init():
-    DisplacementsFieldTransform(
+    DenseFieldTransform(
         np.zeros((10, 10, 10, 3)),
         reference=nb.Nifti1Image(np.zeros((10, 10, 10, 3)), np.eye(4), None),
     )
 
     with pytest.raises(TransformError):
-        DisplacementsFieldTransform(np.zeros((10, 10, 10, 3)))
+        DenseFieldTransform(np.zeros((10, 10, 10, 3)))
     with pytest.raises(TransformError):
-        DisplacementsFieldTransform(
+        DenseFieldTransform(
             np.zeros((10, 10, 10, 3)),
             reference=np.zeros((10, 10, 10, 3)),
         )
@@ -237,7 +237,7 @@ def test_bspline(tmp_path, testdata_path):
     bs_name = testdata_path / "someones_bspline_coefficients.nii.gz"
 
     bsplxfm = BSplineFieldTransform(bs_name, reference=img_name)
-    dispxfm = DisplacementsFieldTransform(disp_name)
+    dispxfm = DenseFieldTransform(disp_name)
 
     out_disp = dispxfm.apply(img_name)
     out_bspl = bsplxfm.apply(img_name)