Apply suggestions from code review

Co-authored-by: Chris Markiewicz <markiewicz@stanford.edu>

Author: oesteban

nitransforms/interp/bspline.py

@@ -75,12 +75,14 @@ def grid_bspline_weights(target_nii, ctrl_nii):
     shape = target_nii.shape[:3]
     ctrl_sp = nb.affines.voxel_sizes(ctrl_nii.affine)[:3]
     ras2ijk = np.linalg.inv(ctrl_nii.affine)
+    # IJK index in the control point image of the first index in the target image
     origin = nb.affines.apply_affine(ras2ijk, [tuple(target_nii.affine[:3, 3])])[0]
 
     wd = []
     for i, (o, n, sp) in enumerate(
         zip(origin, shape, nb.affines.voxel_sizes(target_nii.affine)[:3])
     ):
+        # Locations of voxels in target image in control point image
         locations = np.arange(0, n, dtype="float16") * sp / ctrl_sp[i] + o
         knots = np.arange(0, ctrl_nii.shape[i], dtype="float16")
         distance = np.abs(locations[np.newaxis, ...] - knots[..., np.newaxis])

nitransforms/nonlinear.py

@@ -37,7 +37,7 @@ def __init__(self, field, reference=None):
         Example
         -------
         >>> DisplacementsFieldTransform(test_dir / "someones_displacement_field.nii.gz")
-        <(57, 67, 56) field of 3D displacements>
+        <DisplacementFieldTransform[3D] (57, 67, 56)>
 
         """
         super().__init__()
@@ -59,7 +59,7 @@ def __init__(self, field, reference=None):
 
     def __repr__(self):
         """Beautify the python representation."""
-        return f"<{self._field.shape[:3]} field of {self._field.shape[-1]}D displacements>"
+        return f"<DisplacementFieldTransform[{self._field.shape[-1]}D] {self._field.shape[:3]}>"
 
     def map(self, x, inverse=False):
         r"""
@@ -132,7 +132,7 @@ def __init__(self, coefficients, reference=None, order=3):
         coefficients = _ensure_image(coefficients)
 
         self._coeffs = np.asanyarray(coefficients.dataobj)
-        self._knots = ImageGrid(four_to_three(coefficients)[0])
+        self._knots = ImageGrid(coefficients)
         self._weights = None
         if reference is not None:
             self.reference = reference
@@ -158,9 +158,11 @@ def to_field(self, reference=None, dtype="float32"):
         field = np.zeros((_ref.npoints, ndim))
 
         for d in range(ndim):
+            #  1 x Nvox :                          (1 x K) @ (K x Nvox)
             field[:, d] = self._coeffs[..., d].reshape(-1) @ self._weights
 
-        return field.astype(dtype)
+        return DisplacementsFieldTransform(
+            field.astype(dtype).reshape(*_ref.shape, -1), reference=_ref)
 
     def apply(
         self,
@@ -197,10 +199,7 @@ def apply(
             )
 
         # If locations to be interpolated are on a grid, generate a displacements field
-        return DisplacementsFieldTransform(
-            self.to_field().reshape((*(_ref.shape), -1)),
-            reference=_ref,
-        ).apply(
+        return self.to_field().apply(
             spatialimage,
             reference=reference,
             order=order,
@@ -216,7 +215,7 @@ def map(self, x, inverse=False):
 
         .. math::
             \mathbf{y} = \mathbf{x} + \Psi^3(\mathbf{k}, \mathbf{x}),
-            \label{eq:2}\tag{2}
+            \label{eq:1}\tag{1}
 
         Parameters
         ----------