correct calling to invert_vector_field_2d

dipy/align/imwarp.py

@@ -1320,40 +1320,72 @@ def __invert_models(self, current_disp_world2grid, current_disp_spacing):
         """
 
         # Invert the forward model's forward field
-        self.static_to_ref.backward = np.array(
-            self.invert_vector_field(self.static_to_ref.forward,
-                                     current_disp_world2grid,
-                                     current_disp_spacing,
-                                     self.inv_iter, self.inv_tol,
-                                     self.static_to_ref.backward,
-                                     self.num_threads))
-
-        # Invert the backward model's forward field
-        self.moving_to_ref.backward = np.array(
-            self.invert_vector_field(self.moving_to_ref.forward,
-                                     current_disp_world2grid,
-                                     current_disp_spacing,
-                                     self.inv_iter, self.inv_tol,
-                                     self.moving_to_ref.backward,
-                                     self.num_threads))
-
-        # Invert the forward model's backward field
-        self.static_to_ref.forward = np.array(
-            self.invert_vector_field(self.static_to_ref.backward,
-                                     current_disp_world2grid,
-                                     current_disp_spacing,
-                                     self.inv_iter, self.inv_tol,
-                                     self.static_to_ref.forward,
-                                     self.num_threads))
-
-        # Invert the backward model's backward field
-        self.moving_to_ref.forward = np.array(
-            self.invert_vector_field(self.moving_to_ref.backward,
-                                     current_disp_world2grid,
-                                     current_disp_spacing,
-                                     self.inv_iter, self.inv_tol,
-                                     self.moving_to_ref.forward,
-                                     self.num_threads))
+        if self.dim == 2:
+            self.static_to_ref.backward = np.array(
+                self.invert_vector_field(self.static_to_ref.forward,
+                                         current_disp_world2grid,
+                                         current_disp_spacing,
+                                         self.inv_iter, self.inv_tol,
+                                         self.static_to_ref.backward))
+
+            # Invert the backward model's forward field
+            self.moving_to_ref.backward = np.array(
+                self.invert_vector_field(self.moving_to_ref.forward,
+                                         current_disp_world2grid,
+                                         current_disp_spacing,
+                                         self.inv_iter, self.inv_tol,
+                                         self.moving_to_ref.backward))
+
+            # Invert the forward model's backward field
+            self.static_to_ref.forward = np.array(
+                self.invert_vector_field(self.static_to_ref.backward,
+                                         current_disp_world2grid,
+                                         current_disp_spacing,
+                                         self.inv_iter, self.inv_tol,
+                                         self.static_to_ref.forward))
+
+            # Invert the backward model's backward field
+            self.moving_to_ref.forward = np.array(
+                self.invert_vector_field(self.moving_to_ref.backward,
+                                         current_disp_world2grid,
+                                         current_disp_spacing,
+                                         self.inv_iter, self.inv_tol,
+                                         self.moving_to_ref.forward))
+        else:
+            self.static_to_ref.backward = np.array(
+                self.invert_vector_field(self.static_to_ref.forward,
+                                         current_disp_world2grid,
+                                         current_disp_spacing,
+                                         self.inv_iter, self.inv_tol,
+                                         self.static_to_ref.backward,
+                                         self.num_threads))
+
+            # Invert the backward model's forward field
+            self.moving_to_ref.backward = np.array(
+                self.invert_vector_field(self.moving_to_ref.forward,
+                                         current_disp_world2grid,
+                                         current_disp_spacing,
+                                         self.inv_iter, self.inv_tol,
+                                         self.moving_to_ref.backward,
+                                         self.num_threads))
+
+            # Invert the forward model's backward field
+            self.static_to_ref.forward = np.array(
+                self.invert_vector_field(self.static_to_ref.backward,
+                                         current_disp_world2grid,
+                                         current_disp_spacing,
+                                         self.inv_iter, self.inv_tol,
+                                         self.static_to_ref.forward,
+                                         self.num_threads))
+
+            # Invert the backward model's backward field
+            self.moving_to_ref.forward = np.array(
+                self.invert_vector_field(self.moving_to_ref.backward,
+                                         current_disp_world2grid,
+                                         current_disp_spacing,
+                                         self.inv_iter, self.inv_tol,
+                                         self.moving_to_ref.forward,
+                                         self.num_threads))
 
     def _approximate_derivative_direct(self, x, y):
         r"""Derivative of the degree-2 polynomial fit of the given x, y pairs