From 5c05501d44487505f487bb3dbf69ea77944a0b76 Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Sun, 12 Feb 2023 21:56:22 +0530 Subject: [PATCH] Added imports above numpy --- pylops/signalprocessing/fft.py | 15 +++++++-------- pylops/signalprocessing/fft2d.py | 13 ++++++------- pylops/signalprocessing/fftnd.py | 24 +++++++++++++++--------- 3 files changed, 28 insertions(+), 24 deletions(-) diff --git a/pylops/signalprocessing/fft.py b/pylops/signalprocessing/fft.py index 006631ca..7a3051b3 100644 --- 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 index cf150994..e1957b25 100644 --- 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 index 9ce04af7..ae00b6e1 100644 --- 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)))