fix: final amends to ensure tests pass

Author: oesteban

nitransforms/base.py

@@ -103,7 +103,7 @@ def __init__(self, image):
         self._shape = image.shape
 
         self._ndim = getattr(image, "ndim", len(image.shape))
-        if self._ndim == 4:
+        if self._ndim >= 4:
             self._shape = image.shape[:3]
             self._ndim = 3
 

nitransforms/nonlinear.py

@@ -48,7 +48,9 @@ def __init__(self, field, reference=None):
         )
 
         try:
-            self.reference = reference if reference is not None else field
+            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"
@@ -100,6 +102,7 @@ def map(self, x, inverse=False):
         [[-6.5, -36.475167989730835, -19.5], [-1.0, -42.038356602191925, -11.25]]
 
         """
+
         if inverse is True:
             raise NotImplementedError
         ijk = self.reference.index(x)

nitransforms/tests/test_nonlinear.py

@@ -8,8 +8,9 @@
 
 import numpy as np
 import nibabel as nb
-from ..io.base import TransformFileError
-from ..nonlinear import (
+from nitransforms.base import TransformError
+from nitransforms.io.base import TransformFileError
+from nitransforms.nonlinear import (
     BSplineFieldTransform,
     DisplacementsFieldTransform,
     load as nlload,
@@ -43,7 +44,7 @@ def test_itk_disp_load(size):
 @pytest.mark.parametrize("size", [(20, 20, 20), (20, 20, 20, 2, 3)])
 def test_displacements_bad_sizes(size):
     """Checks field sizes."""
-    with pytest.raises(ValueError):
+    with pytest.raises(TransformError):
         DisplacementsFieldTransform(nb.Nifti1Image(np.zeros(size), np.eye(4), None))