Merge branch 'fix/267' into fix/itk-displacements-field-round-trip

oesteban · oesteban · commit 8e83ccd2aec6 · 2025-07-25T20:04:55.000+02:00
diff --git a/nitransforms/base.py b/nitransforms/base.py
@@ -202,30 +202,26 @@ def inverse(self):
     def ndindex(self):
         """List the indexes corresponding to the space grid."""
         if self._ndindex is None:
-            indexes = tuple([np.arange(s) for s in self._shape])
-            self._ndindex = np.array(np.meshgrid(*indexes, indexing="ij")).reshape(
-                self._ndim, self._npoints
-            )
+            indexes = np.mgrid[
+                0:self._shape[0], 0:self._shape[1], 0:self._shape[2]
+            ]
+            self._ndindex = indexes.reshape((indexes.shape[0], -1)).T
         return self._ndindex
 
     @property
     def ndcoords(self):
         """List the physical coordinates of this gridded space samples."""
         if self._coords is None:
-            self._coords = np.tensordot(
-                self._affine,
-                np.vstack((self.ndindex, np.ones((1, self._npoints)))),
-                axes=1,
-            )[:3, ...]
+            self._coords = self.ras(self.ndindex)
         return self._coords
 
     def ras(self, ijk):
         """Get RAS+ coordinates from input indexes."""
-        return _apply_affine(ijk, self._affine, self._ndim)
+        return _apply_affine(ijk, self._affine, self._ndim).T
 
     def index(self, x):
         """Get the image array's indexes corresponding to coordinates."""
-        return _apply_affine(x, self._inverse, self._ndim)
+        return _apply_affine(x, self._inverse, self._ndim).T
 
     def _to_hdf5(self, group):
         group.attrs["Type"] = "image"
diff --git a/nitransforms/nonlinear.py b/nitransforms/nonlinear.py
@@ -65,50 +65,47 @@ def __init__(self, field=None, is_deltas=True, reference=None):
         <DenseFieldTransform[3D] (57, 67, 56)>
 
         """
+
         if field is None and reference is None:
-            raise TransformError("DenseFieldTransforms require a spatial reference")
+            raise TransformError("cannot initialize field")
 
         super().__init__()
 
-        self._is_deltas = is_deltas
+        if field is not None:
+            field = _ensure_image(field)
+            # Extract data if nibabel object otherwise assume numpy array
+            _data = np.squeeze(
+                np.asanyarray(field.dataobj)
+                if hasattr(field, "dataobj")
+                else field.copy()
+            )
 
         try:
             self.reference = ImageGrid(reference if reference is not None else field)
         except AttributeError:
             raise TransformError(
-                "Field must be a spatial image if reference is not provided"
+                "field must be a spatial image if reference is not provided"
                 if reference is None
-                else "Reference is not a spatial image"
+                else "reference is not a spatial image"
             )
 
         fieldshape = (*self.reference.shape, self.reference.ndim)
-        if field is not None:
-            field = _ensure_image(field)
-            self._field = np.squeeze(
-                np.asanyarray(field.dataobj) if hasattr(field, "dataobj") else field
-            )
-            if fieldshape != self._field.shape:
-                raise TransformError(
-                    f"Shape of the field ({'x'.join(str(i) for i in self._field.shape)}) "
-                    f"doesn't match that of the reference({'x'.join(str(i) for i in fieldshape)})"
-                )
-        else:
-            self._field = np.zeros(fieldshape, dtype="float32")
-            self._is_deltas = True
-
-        if self._field.shape[-1] != self.ndim:
+        if field is None:
+            _data = np.zeros(fieldshape)
+        elif fieldshape != _data.shape:
             raise TransformError(
-                "The number of components of the field (%d) does not match "
-                "the number of dimensions (%d)" % (self._field.shape[-1], self.ndim)
+                f"Shape of the field ({'x'.join(str(i) for i in _data.shape)}) "
+                f"doesn't match that of the reference({'x'.join(str(i) for i in fieldshape)})"
             )
 
+        self._is_deltas = is_deltas
+        self._field = self.reference.ndcoords.reshape(fieldshape)
+
         if self.is_deltas:
-            self._deltas = (
-                self._field.copy()
-            )  # IMPORTANT: you don't want to update deltas
-            # Convert from displacements (deltas) to deformations fields
-            # (just add its origin to each delta vector)
-            self._field += self.reference.ndcoords.T.reshape(fieldshape)
+            self._deltas = _data.copy()
+            self._field += self._deltas
+        else:
+            self._field = _data.copy()
 
     def __repr__(self):
         """Beautify the python representation."""
@@ -153,7 +150,7 @@ def map(self, x, inverse=False):
         ...     test_dir / "someones_displacement_field.nii.gz",
         ...     is_deltas=False,
         ... )
-        >>> xfm.map([-6.5, -36., -19.5]).tolist()
+        >>> 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()
@@ -170,8 +167,8 @@ def map(self, x, inverse=False):
         ...     test_dir / "someones_displacement_field.nii.gz",
         ...     is_deltas=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]]
+        >>> xfm.map([[-6.5, -36., -19.5], [-1., -41.5, -11.25]]).tolist()  # doctest: +ELLIPSIS
+        [[-6.5, -36.475..., -19.5], [-1.0, -42.038..., -11.25]]
 
         >>> np.array_str(
         ...     xfm.map([[-6.7, -36.3, -19.2], [-1., -41.5, -11.25]]),
@@ -185,19 +182,19 @@ def map(self, x, inverse=False):
         if inverse is True:
             raise NotImplementedError
 
-        ijk = self.reference.index(x)
+        ijk = self.reference.index(np.array(x, dtype="float32"))
         indexes = np.round(ijk).astype("int")
+        ongrid = np.where(np.linalg.norm(ijk - indexes, axis=1) < 1e-3)[0]
 
-        import pdb; pdb.set_trace()
-        if np.all(np.abs(ijk - indexes) < 1e-3):
-            indexes = tuple(tuple(i) for i in indexes)
-            return self._field[indexes]
+        if ongrid.size == np.shape(x)[0]:
+            # return self._field[*indexes.T, :]  # From Python 3.11
+            return self._field[tuple(indexes.T) + (np.s_[:],)]
 
-        new_map = np.vstack(
+        mapped_coords = np.vstack(
             tuple(
                 map_coordinates(
                     self._field[..., i],
-                    ijk,
+                    ijk.T,
                     order=3,
                     mode="constant",
                     cval=np.nan,
@@ -208,8 +205,8 @@ def map(self, x, inverse=False):
         ).T
 
         # Set NaN values back to the original coordinates value = no displacement
-        new_map[np.isnan(new_map)] = np.array(x)[np.isnan(new_map)]
-        return new_map
+        mapped_coords[np.isnan(mapped_coords)] = np.array(x)[np.isnan(mapped_coords)]
+        return mapped_coords
 
     def __matmul__(self, b):
         """
diff --git a/nitransforms/resampling.py b/nitransforms/resampling.py
@@ -253,7 +253,7 @@ def apply(
     serialize_4d = n_resamplings >= serialize_nvols
 
     targets = None
-    ref_ndcoords = _ref.ndcoords.T
+    ref_ndcoords = _ref.ndcoords
     if hasattr(transform, "to_field") and callable(transform.to_field):
         targets = ImageGrid(spatialimage).index(
             _as_homogeneous(
@@ -271,11 +271,8 @@ def apply(
             else targets
         )
 
-    if targets.ndim == 3:
-        targets = np.rollaxis(targets, targets.ndim - 1, 0)
-    else:
-        assert targets.ndim == 2
-        targets = targets[np.newaxis, ...]
+    if targets.ndim == 2:
+        targets = targets.T[np.newaxis, ...]
 
     if serialize_4d:
         data = (
@@ -290,6 +287,9 @@ def apply(
             (len(ref_ndcoords), n_resamplings), dtype=input_dtype, order="F"
         )
 
+        if targets.ndim == 3:
+            targets = np.rollaxis(targets, targets.ndim - 1, 1)
+
         resampled = asyncio.run(
             _apply_serial(
                 data,
@@ -311,6 +311,9 @@ def apply(
     else:
         data = np.asanyarray(spatialimage.dataobj, dtype=input_dtype)
 
+        if targets.ndim == 3:
+            targets = np.rollaxis(targets, targets.ndim - 1, 0)
+
         if data_nvols == 1 and xfm_nvols == 1:
             targets = np.squeeze(targets)
             assert targets.ndim == 2
@@ -320,15 +323,19 @@ def apply(
 
         if xfm_nvols > 1:
             assert targets.ndim == 3
-            n_time, n_dim, n_vox = targets.shape
+
+            # Targets must have shape (n_dim x n_time x n_vox)
+            n_dim, n_time, n_vox = targets.shape
             # Reshape to (3, n_time x n_vox)
-            ijk_targets = np.rollaxis(targets, 0, 2).reshape((n_dim, -1))
+            ijk_targets = targets.reshape((n_dim, -1))
             time_row = np.repeat(np.arange(n_time), n_vox)[None, :]
 
             # Now targets is (4, n_vox x n_time), with indexes (t, i, j, k)
             # t is the slowest-changing axis, so we put it first
             targets = np.vstack((time_row, ijk_targets))
             data = np.rollaxis(data, data.ndim - 1, 0)
+        else:
+            targets = targets.T
 
         resampled = ndi.map_coordinates(
             data,
diff --git a/nitransforms/tests/test_base.py b/nitransforms/tests/test_base.py
@@ -55,20 +55,24 @@ def test_ImageGrid(get_testdata, image_orientation):
 
     assert np.allclose(np.squeeze(img.ras(ijk[0])), xyz[0])
     assert np.allclose(np.round(img.index(xyz[0])), ijk[0])
-    assert np.allclose(img.ras(ijk).T, xyz)
-    assert np.allclose(np.round(img.index(xyz)).T, ijk)
+    assert np.allclose(img.ras(ijk), xyz)
+    assert np.allclose(np.round(img.index(xyz)), ijk)
 
     # nd index / coords
     idxs = img.ndindex
     coords = img.ndcoords
     assert len(idxs.shape) == len(coords.shape) == 2
-    assert idxs.shape[0] == coords.shape[0] == img.ndim == 3
-    assert idxs.shape[1] == coords.shape[1] == img.npoints == np.prod(im.shape)
+    assert idxs.shape[1] == coords.shape[1] == img.ndim == 3
+    assert idxs.shape[0] == coords.shape[0] == img.npoints == np.prod(im.shape)
 
     img2 = ImageGrid(img)
     assert img2 == img
     assert (img2 != img) is False
 
+    # Test indexing round trip
+    np.testing.assert_allclose(img.ndcoords, img.ras(img.ndindex))
+    np.testing.assert_allclose(img.ndindex, np.round(img.index(img.ndcoords)))
+
 
 def test_ImageGrid_utils(tmpdir, testdata_path, get_testdata):
     """Check that images can be objects or paths and equality."""
diff --git a/nitransforms/tests/test_io.py b/nitransforms/tests/test_io.py
@@ -17,13 +17,13 @@
 from nibabel.affines import from_matvec
 from scipy.io import loadmat
 from nitransforms.linear import Affine
-from nitransforms import nonlinear as nitnl
+from nitransforms.nonlinear import DenseFieldTransform, BSplineFieldTransform
 from nitransforms.io import (
     afni,
     fsl,
     lta as fs,
     itk,
-    x5
+    x5,
 )
 from nitransforms.io.lta import (
     VolumeGeometry as VG,
@@ -773,7 +773,7 @@ def test_densefield_x5_roundtrip(tmp_path, is_deltas):
     ref = nb.Nifti1Image(np.zeros((2, 2, 2), dtype="uint8"), np.eye(4))
     disp = nb.Nifti1Image(np.random.rand(2, 2, 2, 3).astype("float32"), np.eye(4))
 
-    xfm = nitnl.DenseFieldTransform(disp, is_deltas=is_deltas, reference=ref)
+    xfm = DenseFieldTransform(disp, is_deltas=is_deltas, reference=ref)
 
     node = xfm.to_x5(metadata={"GeneratedBy": "pytest"})
     assert node.type == "nonlinear"
@@ -785,7 +785,7 @@ def test_densefield_x5_roundtrip(tmp_path, is_deltas):
     fname = tmp_path / "test.x5"
     x5.to_filename(fname, [node])
 
-    xfm2 = nitnl.DenseFieldTransform.from_filename(fname, fmt="X5")
+    xfm2 = DenseFieldTransform.from_filename(fname, fmt="X5")
 
     assert xfm2.reference.shape == ref.shape
     assert np.allclose(xfm2.reference.affine, ref.affine)
@@ -797,7 +797,7 @@ def test_bspline_to_x5(tmp_path):
     coeff = nb.Nifti1Image(np.zeros((2, 2, 2, 3), dtype="float32"), np.eye(4))
     ref = nb.Nifti1Image(np.zeros((2, 2, 2), dtype="uint8"), np.eye(4))
 
-    xfm = nitnl.BSplineFieldTransform(coeff, reference=ref)
+    xfm = BSplineFieldTransform(coeff, reference=ref)
     node = xfm.to_x5(metadata={"tool": "pytest"})
     assert node.type == "nonlinear"
     assert node.subtype == "bspline"
@@ -807,7 +807,7 @@ def test_bspline_to_x5(tmp_path):
     fname = tmp_path / "bspline.x5"
     x5.to_filename(fname, [node])
 
-    xfm2 = nitnl.BSplineFieldTransform.from_filename(fname, fmt="X5")
+    xfm2 = BSplineFieldTransform.from_filename(fname, fmt="X5")
     assert np.allclose(xfm._coeffs, xfm2._coeffs)
     assert xfm2.reference.shape == ref.shape
     assert np.allclose(xfm2.reference.affine, ref.affine)
diff --git a/nitransforms/tests/test_nonlinear.py b/nitransforms/tests/test_nonlinear.py
diff --git a/nitransforms/tests/test_resampling.py b/nitransforms/tests/test_resampling.py
