fix: ensure ndcoords and ndindex are NxD

Author: oesteban

nitransforms/base.py

@@ -202,30 +202,24 @@ 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"

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(
@@ -332,7 +332,7 @@ def apply(
 
         resampled = ndi.map_coordinates(
             data,
-            targets,
+            targets.T,
             order=order,
             mode=mode,
             cval=cval,

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."""