Added imports above numpy

PyLops · Feb 12, 2023 · 5c05501 · 5c05501
1 parent 8e6cb9a
commit 5c05501
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 24 deletions.
diff --git a/pylops/signalprocessing/fft.py b/pylops/signalprocessing/fft.py
@@ -3,6 +3,9 @@
 import logging
 import warnings
 from typing import Optional, Union
+from mkl_fft import _numpy_fft as pymkl_fft
+from mkl_fft._scipy_fft_backend import fftshift as mkl_fftshift, ifftshift as mkl_ifftshift
+
 
 import numpy as np
 import numpy.typing as npt
@@ -402,10 +405,8 @@ def __init__(
 
     @reshaped
     def _matvec(self, x: NDArray) -> NDArray:
-        from mkl_fft import _numpy_fft as pymkl_fft
-        # from mkl_fft._scipy_fft_backend import fftshift as mkl_fftshift, ifftshift as mkl_iffshift
         if self.ifftshift_before:
-            x = np.fft.ifftshift(x, axes=self.axis)
+            x = mkl_ifftshift(x, axes=self.axis)
         if not self.clinear:
             x = np.real(x)
         if self.real:
@@ -419,15 +420,13 @@ def _matvec(self, x: NDArray) -> NDArray:
         if self.norm is _FFTNorms.ONE_OVER_N:
             y *= self._scale
         if self.fftshift_after:
-            y = np.fft.fftshift(y, axes=self.axis)
+            y = mkl_fftshift(y, axes=self.axis)
         return y
 
     @reshaped
     def _rmatvec(self, x: NDArray) -> NDArray:
-        from mkl_fft import _numpy_fft as pymkl_fft
-        # from mkl_fft._scipy_fft_backend import fftshift as mkl_fftshift, ifftshift as mkl_iffshift
         if self.fftshift_after:
-            x = np.fft.ifftshift(x, axes=self.axis)
+            x = mkl_ifftshift(x, axes=self.axis)
         if self.real:
             # Apply scaling to obtain a correct adjoint for this operator
             x = x.copy()
@@ -448,7 +447,7 @@ def _rmatvec(self, x: NDArray) -> NDArray:
         if not self.clinear:
             y = np.real(y)
         if self.ifftshift_before:
-            y = np.fft.fftshift(y, axes=self.axis)
+            y = mkl_fftshift(y, axes=self.axis)
         return y
 
     def __truediv__(self, y):

diff --git a/pylops/signalprocessing/fft2d.py b/pylops/signalprocessing/fft2d.py
@@ -4,9 +4,10 @@
 import warnings
 from typing import Dict, Optional, Sequence, Union
 
+from pylops.signalprocessing.fftnd import _FFTND_mklfft
+from mkl_fft._scipy_fft_backend import fftshift as mkl_fftshift, ifftshift as mkl_ifftshift
 import numpy as np
 import scipy.fft
-from pylops.signalprocessing.fftnd import _FFTND_mklfft
 
 from pylops import LinearOperator
 from pylops.signalprocessing._baseffts import _BaseFFTND, _FFTNorms
@@ -265,9 +266,8 @@ def __init__(
 
     @reshaped
     def _matvec(self, x):
-        # from mkl_fft._scipy_fft_backend import fftshift as mkl_fftshift, ifftshift as mkl_iffshift
         if self.ifftshift_before.any():
-            x = np.fft.ifftshift(x, axes=self.axes[self.ifftshift_before])
+            x = mkl_ifftshift(x, axes=self.axes[self.ifftshift_before])
         if not self.clinear:
             x = np.real(x)
         if self.real:
@@ -281,14 +281,13 @@ def _matvec(self, x):
         if self.norm is _FFTNorms.ONE_OVER_N:
             y *= self._scale
         if self.fftshift_after.any():
-            y = np.fft.fftshift(y, axes=self.axes[self.fftshift_after])
+            y = mkl_fftshift(y, axes=self.axes[self.fftshift_after])
         return y
 
     @reshaped
     def _rmatvec(self, x):
-        # from mkl_fft._scipy_fft_backend import fftshift as mkl_fftshift, ifftshift as mkl_iffshift
         if self.fftshift_after.any():
-            x = np.fft.ifftshift(x, axes=self.axes[self.fftshift_after])
+            x = mkl_ifftshift(x, axes=self.axes[self.fftshift_after])
         if self.real:
             # Apply scaling to obtain a correct adjoint for this operator
             x = x.copy()
@@ -305,7 +304,7 @@ def _rmatvec(self, x):
         if not self.clinear:
             y = np.real(y)
         if self.ifftshift_before.any():
-            y = np.fft.fftshift(y, axes=self.axes[self.ifftshift_before])
+            y = mkl_fftshift(y, axes=self.axes[self.ifftshift_before])
         return y
 
     def __truediv__(self, y):

diff --git a/pylops/signalprocessing/fftnd.py b/pylops/signalprocessing/fftnd.py
@@ -235,9 +235,10 @@ def __init__(
 
     @reshaped
     def _matvec(self, x: NDArray) -> NDArray:
-        # from mkl_fft._scipy_fft_backend import fftshift as mkl_fftshift, ifftshift as mkl_iffshift
+        from mkl_fft._scipy_fft_backend import fftshift as mkl_fftshift, ifftshift as mkl_iffshift
+
         if self.ifftshift_before.any():
-            x = np.fft.ifftshift(x, axes=self.axes[self.ifftshift_before])
+            x = mkl_fftshift(x, axes=self.axes[self.ifftshift_before])
         if not self.clinear:
             x = np.real(x)
         if self.real:
@@ -251,14 +252,15 @@ def _matvec(self, x: NDArray) -> NDArray:
         if self.norm is _FFTNorms.ONE_OVER_N:
             y *= self._scale
         if self.fftshift_after.any():
-            y = np.fft.fftshift(y, axes=self.axes[self.fftshift_after])
+            y = mkl_iffshift(y, axes=self.axes[self.fftshift_after])
         return y
 
     @reshaped
     def _rmatvec(self, x: NDArray) -> NDArray:
-        # from mkl_fft._scipy_fft_backend import fftshift as mkl_fftshift, ifftshift as mkl_iffshift
+        from mkl_fft._scipy_fft_backend import fftshift as mkl_fftshift, ifftshift as mkl_iffshift
+
         if self.fftshift_after.any():
-            x = np.fft.ifftshift(x, axes=self.axes[self.fftshift_after])
+            x = mkl_iffshift(x, axes=self.axes[self.fftshift_after])
         if self.real:
             # Apply scaling to obtain a correct adjoint for this operator
             x = x.copy()
@@ -278,12 +280,13 @@ def _rmatvec(self, x: NDArray) -> NDArray:
         if not self.clinear:
             y = np.real(y)
         if self.ifftshift_before.any():
-            y = np.fft.fftshift(y, axes=self.axes[self.ifftshift_before])
+            y = mkl_fftshift(y, axes=self.axes[self.ifftshift_before])
         return y
 
     @staticmethod
     def rfftn(a, s=None, axes=None, norm=None):
-        from mkl_fft._numpy_fft import rfft, fft, _cook_nd_args, asarray
+        from mkl_fft._numpy_fft import asarray, _cook_nd_args, fft, rfft
+
         a = asarray(a)
         s, axes = _cook_nd_args(a, s, axes)
         a = rfft(a, s[-1], axes[-1], norm)
@@ -293,7 +296,8 @@ def rfftn(a, s=None, axes=None, norm=None):
 
     @staticmethod
     def irfftn(a, s=None, axes=None, norm=None):
-        from mkl_fft._numpy_fft import irfft, ifft, asarray, _cook_nd_args
+        from mkl_fft._numpy_fft import ifft, asarray, _cook_nd_args, irfft
+
         a = asarray(a)
         s, axes = _cook_nd_args(a, s, axes, invreal=1)
         for ii in range(len(axes) - 1):
@@ -303,7 +307,8 @@ def irfftn(a, s=None, axes=None, norm=None):
 
     @staticmethod
     def fftn(a, s=None, axes=None, norm=None):
-        from mkl_fft._numpy_fft import fft, asarray, _cook_nd_args
+        from mkl_fft._numpy_fft import asarray, _cook_nd_args, fft
+
         a = asarray(a)
         s, axes = _cook_nd_args(a, s, axes)
         itl = list(range(len(axes)))
@@ -315,6 +320,7 @@ def fftn(a, s=None, axes=None, norm=None):
     @staticmethod
     def ifftn(a, s=None, axes=None, norm=None):
         from mkl_fft._numpy_fft import ifft, asarray, _cook_nd_args
+
         a = asarray(a)
         s, axes = _cook_nd_args(a, s, axes)
         itl = list(range(len(axes)))