refactor: use the base operator safe_squeeze.

Author: paquiteau

src/mrinufft/operators/base.py

@@ -192,6 +192,19 @@ def check_shape(self, *, image=None, ksp=None):
         if image is None and ksp is None:
             raise ValueError("Nothing to check, provides image or ksp arguments")
 
+    def _safe_squeeze(self, arr):
+        """Squeeze the first two dimensions of shape of the operator."""
+        if self.squeeze_dims:
+            try:
+                arr = arr.squeeze(axis=1)
+            except ValueError:
+                pass
+            try:
+                arr = arr.squeeze(axis=0)
+            except ValueError:
+                pass
+        return arr
+
     @abstractmethod
     def op(self, data):
         """Compute operator transform.
@@ -795,19 +808,6 @@ def _grad_calibless(self, image_data, obs_data):
         grad /= self.norm_factor
         return grad.reshape(B, C, *XYZ)
 
-    def _safe_squeeze(self, arr):
-        """Squeeze the first two dimensions of shape of the operator."""
-        if self.squeeze_dims:
-            try:
-                arr = arr.squeeze(axis=1)
-            except ValueError:
-                pass
-            try:
-                arr = arr.squeeze(axis=0)
-            except ValueError:
-                pass
-        return arr
-
 
 def power_method(
     max_iter: int,

src/mrinufft/operators/interfaces/cufinufft.py

@@ -749,19 +749,6 @@ def _dc_calibless_device(self, image_data, obs_data):
         grad /= self.norm_factor
         return grad
 
-    def _safe_squeeze(self, arr):
-        """Squeeze the first two dimensions of shape of the operator."""
-        if self.squeeze_dims:
-            try:
-                arr = arr.squeeze(axis=1)
-            except ValueError:
-                pass
-            try:
-                arr = arr.squeeze(axis=0)
-            except ValueError:
-                pass
-        return arr
-
     @property
     def eps(self):
         """Return the underlying precision parameter."""

src/mrinufft/operators/interfaces/gpunufft.py

@@ -656,19 +656,6 @@ def get_lipschitz_cst(self, max_iter=10, tolerance=1e-5, **kwargs):
             max_iter=max_iter, tolerance=tolerance
         )
 
-    def _safe_squeeze(self, arr):
-        """Squeeze the first two dimensions of shape of the operator."""
-        if self.squeeze_dims:
-            try:
-                arr = arr.squeeze(axis=1)
-            except ValueError:
-                pass
-            try:
-                arr = arr.squeeze(axis=0)
-            except ValueError:
-                pass
-        return arr
-
     @with_numpy_cupy
     def data_consistency(self, image_data, obs_data):
         """Compute the data consistency estimation directly on gpu.

src/mrinufft/operators/interfaces/torchkbnufft.py

@@ -166,19 +166,6 @@ def adj_op(self, coeffs, out=None):
         img /= self.norm_factor
         return self._safe_squeeze(img)
 
-    def _safe_squeeze(self, arr):
-        """Squeeze the first two dimensions of shape of the operator."""
-        if self.squeeze_dims:
-            try:
-                arr = arr.squeeze(axis=1)
-            except ValueError:
-                pass
-            try:
-                arr = arr.squeeze(axis=0)
-            except ValueError:
-                pass
-        return arr
-
     @with_torch
     def data_consistency(self, data, obs_data):
         """Compute the data consistency.

src/mrinufft/operators/stacked.py

@@ -283,19 +283,6 @@ def _adj_op_calibless(self, coeffs, img):
         img = self._ifftz(imgz)
         return img
 
-    def _safe_squeeze(self, arr):
-        """Squeeze the first two dimensions of shape of the operator."""
-        if self.squeeze_dims:
-            try:
-                arr = arr.squeeze(axis=1)
-            except ValueError:
-                pass
-            try:
-                arr = arr.squeeze(axis=0)
-            except ValueError:
-                pass
-        return arr
-
     def get_lipschitz_cst(self, max_iter=10):
         """Return the Lipschitz constant of the operator.