diff --git a/mkl_fft/__init__.py b/mkl_fft/__init__.py index 5c25992b..2649444a 100644 --- a/mkl_fft/__init__.py +++ b/mkl_fft/__init__.py @@ -24,11 +24,11 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -from ._pydfti import (fft, ifft, fft2, ifft2, fftn, ifftn, rfft, irfft, - rfft_numpy, irfft_numpy, rfftn_numpy, irfftn_numpy) +from ._pydfti import (fft, ifft, fft2, ifft2, fftn, ifftn, rfftpack, irfftpack, + rfft, irfft, rfft2, irfft2, rfftn, irfftn) from ._version import __version__ import mkl_fft.interfaces -__all__ = ['fft', 'ifft', 'fft2', 'ifft2', 'fftn', 'ifftn', 'rfft', 'irfft', - 'rfft_numpy', 'irfft_numpy', 'rfftn_numpy', 'irfftn_numpy', 'interfaces'] +__all__ = ['fft', 'ifft', 'fft2', 'ifft2', 'fftn', 'ifftn', 'rfftpack', 'irfftpack' + 'rfft', 'irfft', 'rfft2', 'irfft2', 'rfftn', 'irfftn','interfaces'] diff --git a/mkl_fft/_numpy_fft.py b/mkl_fft/_numpy_fft.py index 61a837c8..2c832d5e 100644 --- a/mkl_fft/_numpy_fft.py +++ b/mkl_fft/_numpy_fft.py @@ -406,7 +406,7 @@ def rfft(a, n=None, axis=-1, norm=None): fsc = ortho_sc_1d(n, x.shape[axis]) return trycall( - mkl_fft.rfft_numpy, + mkl_fft.rfft, (x,), {'n': n, 'axis': axis, 'fwd_scale': fsc}) @@ -510,7 +510,7 @@ def irfft(a, n=None, axis=-1, norm=None): fsc = ortho_sc_1d(nn, nn) return trycall( - mkl_fft.irfft_numpy, + mkl_fft.irfft, (x,), {'n': n, 'axis': axis, 'fwd_scale': fsc}) @@ -606,7 +606,7 @@ def hfft(a, n=None, axis=-1, norm=None): fsc = ortho_sc_1d(nn, nn) return trycall( - mkl_fft.irfft_numpy, + mkl_fft.irfft, (x,), {'n': n, 'axis': axis, 'fwd_scale': fsc}) @@ -682,7 +682,7 @@ def ihfft(a, n=None, axis=-1, norm=None): fsc = ortho_sc_1d(n, x.shape[axis]) output = trycall( - mkl_fft.rfft_numpy, + mkl_fft.rfft, (x,), {'n': n, 'axis': axis, 'fwd_scale': fsc}) @@ -1237,7 +1237,7 @@ def rfftn(a, s=None, axes=None, norm=None): fsc = sqrt(frwd_sc_nd(s, axes, x.shape)) return trycall( - mkl_fft.rfftn_numpy, + mkl_fft.rfftn, (x,), {'s': s, 'axes': axes, 'fwd_scale': fsc}) @@ -1394,7 +1394,7 @@ def irfftn(a, s=None, axes=None, norm=None): fsc = sqrt(frwd_sc_nd(s, axes, x.shape)) return trycall( - mkl_fft.irfftn_numpy, + mkl_fft.irfftn, (x,), {'s': s, 'axes': axes, 'fwd_scale': fsc}) diff --git a/mkl_fft/_pydfti.pyx b/mkl_fft/_pydfti.pyx index 82ed5429..47e0bfd6 100644 --- a/mkl_fft/_pydfti.pyx +++ b/mkl_fft/_pydfti.pyx @@ -408,12 +408,12 @@ def _fft1d_impl(x, n=None, axis=-1, overwrite_arg=False, direction=+1, double fs return f_arr -def rfft(x, n=None, axis=-1, overwrite_x=False, fwd_scale=1.0): +def rfftpack(x, n=None, axis=-1, overwrite_x=False, fwd_scale=1.0): """Packed real-valued harmonics of FFT of a real sequence x""" return _rr_fft1d_impl2(x, n=n, axis=axis, overwrite_arg=overwrite_x, fsc=fwd_scale) -def irfft(x, n=None, axis=-1, overwrite_x=False, fwd_scale=1.0): +def irfftpack(x, n=None, axis=-1, overwrite_x=False, fwd_scale=1.0): """Inverse FFT of a real sequence, takes packed real-valued harmonics of FFT""" return _rr_ifft1d_impl2(x, n=n, axis=axis, overwrite_arg=overwrite_x, fsc=fwd_scale) @@ -524,7 +524,7 @@ def _rr_fft1d_impl2(x, n=None, axis=-1, overwrite_arg=False, double fsc=1.0): """ Uses MKL to perform real packed 1D FFT on the input array x along the given axis. - This done by using rfft_numpy and post-processing the result. + This done by using rfft and post-processing the result. Thus overwrite_arg is effectively discarded. Functionally equivalent to scipy.fftpack.rfft @@ -580,7 +580,7 @@ def _rr_ifft1d_impl2(x, n=None, axis=-1, overwrite_arg=False, double fsc=1.0): """ Uses MKL to perform real packed 1D FFT on the input array x along the given axis. - This done by using rfft_numpy and post-processing the result. + This done by using rfft and post-processing the result. Thus overwrite_arg is effectively discarded. Functionally equivalent to scipy.fftpack.irfft @@ -793,11 +793,11 @@ def _rc_ifft1d_impl(x, n=None, axis=-1, overwrite_arg=False, double fsc=1.0): return f_arr -def rfft_numpy(x, n=None, axis=-1, fwd_scale=1.0): +def rfft(x, n=None, axis=-1, fwd_scale=1.0): return _rc_fft1d_impl(x, n=n, axis=axis, fsc=fwd_scale) -def irfft_numpy(x, n=None, axis=-1, fwd_scale=1.0): +def irfft(x, n=None, axis=-1, fwd_scale=1.0): return _rc_ifft1d_impl(x, n=n, axis=axis, fsc=fwd_scale) @@ -1121,12 +1121,12 @@ def ifftn(x, shape=None, axes=None, overwrite_x=False, fwd_scale=1.0): return _fftnd_impl(x, shape=shape, axes=axes, overwrite_x=overwrite_x, direction=-1, fsc=fwd_scale) -def rfft2_numpy(x, s=None, axes=(-2,-1), fwd_scale=1.0): - return rfftn_numpy(x, s=s, axes=axes, fsc=fwd_scale) +def rfft2(x, s=None, axes=(-2,-1), fwd_scale=1.0): + return rfftn(x, s=s, axes=axes, fsc=fwd_scale) -def irfft2_numpy(x, s=None, axes=(-2,-1), fwd_scale=1.0): - return irfftn_numpy(x, s=s, axes=axes, fsc=fwd_scale) +def irfft2(x, s=None, axes=(-2,-1), fwd_scale=1.0): + return irfftn(x, s=s, axes=axes, fsc=fwd_scale) def _remove_axis(s, axes, axis_to_remove): @@ -1181,16 +1181,15 @@ def _fix_dimensions(cnp.ndarray arr, object s, object axes): return np.pad(arr, tuple(pad_widths), 'constant') -def rfftn_numpy(x, s=None, axes=None, fwd_scale=1.0): +def rfftn(x, s=None, axes=None, fwd_scale=1.0): a = np.asarray(x) no_trim = (s is None) and (axes is None) s, axes = _cook_nd_args(a, s, axes) la = axes[-1] - # trim array, so that rfft_numpy avoids doing - # unnecessary computations + # trim array, so that rfft avoids doing unnecessary computations if not no_trim: a = _trim_array(a, s, axes) - a = rfft_numpy(a, n = s[-1], axis=la, fwd_scale=fwd_scale) + a = rfft(a, n = s[-1], axis=la, fwd_scale=fwd_scale) if len(s) > 1: if not no_trim: ss = list(s) @@ -1214,7 +1213,7 @@ def rfftn_numpy(x, s=None, axes=None, fwd_scale=1.0): return a -def irfftn_numpy(x, s=None, axes=None, fwd_scale=1.0): +def irfftn(x, s=None, axes=None, fwd_scale=1.0): a = np.asarray(x) no_trim = (s is None) and (axes is None) s, axes = _cook_nd_args(a, s, axes, invreal=True) @@ -1243,5 +1242,5 @@ def irfftn_numpy(x, s=None, axes=None, fwd_scale=1.0): for ii in range(len(axes)-1): a = ifft(a, s[ii], axes[ii], overwrite_x=ovr_x) ovr_x = True - a = irfft_numpy(a, n = s[-1], axis=la, fwd_scale=fwd_scale) + a = irfft(a, n = s[-1], axis=la, fwd_scale=fwd_scale) return a diff --git a/mkl_fft/_scipy_fft.py b/mkl_fft/_scipy_fft.py index bc613cc0..4d3c9ac9 100644 --- a/mkl_fft/_scipy_fft.py +++ b/mkl_fft/_scipy_fft.py @@ -24,47 +24,401 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -from . import _pydfti +from . import _pydfti as mkl_fft from . import _float_utils +import mkl -__all__ = ['fft', 'ifft', 'fftn', 'ifftn', 'fft2', 'ifft2', 'rfft', 'irfft'] +from numpy import (take, sqrt, prod) +import contextvars +import contextlib +import operator +import os -def fft(a, n=None, axis=-1, overwrite_x=False): - x = _float_utils.__upcast_float16_array(a) - return _pydfti.fft(x, n=n, axis=axis, overwrite_x=overwrite_x) +__doc__ = """ +This module implements interfaces mimicing `scipy.fft` module. +It also provides DftiBackend class which can be used to set mkl_fft to be used +via `scipy.fft` namespace. -def ifft(a, n=None, axis=-1, overwrite_x=False): - x = _float_utils.__upcast_float16_array(a) - return _pydfti.ifft(x, n=n, axis=axis, overwrite_x=overwrite_x) +:Example: + import scipy.fft + import mkl_fft._scipy_fft_backend as be + # Set mkl_fft to be used as backend of SciPy's FFT functions. + scipy.fft.set_global_backend(be) +""" +class _cpu_max_threads_count: + def __init__(self): + self.cpu_count = None + self.max_threads_count = None -def fftn(a, shape=None, axes=None, overwrite_x=False): - x = _float_utils.__upcast_float16_array(a) - return _pydfti.fftn(x, shape=shape, axes=axes, overwrite_x=overwrite_x) + def get_cpu_count(self): + if self.cpu_count is None: + max_threads = self.get_max_threads_count() + self.cpu_count = max_threads + return self.cpu_count + def get_max_threads_count(self): + if self.max_threads_count is None: + self.max_threads_count = mkl.get_max_threads() -def ifftn(a, shape=None, axes=None, overwrite_x=False): - x = _float_utils.__upcast_float16_array(a) - return _pydfti.ifftn(x, shape=shape, axes=axes, overwrite_x=overwrite_x) + return self.max_threads_count -def fft2(a, shape=None, axes=(-2,-1), overwrite_x=False): - x = _float_utils.__upcast_float16_array(a) - return _pydfti.fftn(x, shape=shape, axes=axes, overwrite_x=overwrite_x) +class _workers_data: + def __init__(self, workers=None): + if workers: + self.workers_ = workers + else: + self.workers_ = _cpu_max_threads_count().get_cpu_count() + self.workers_ = operator.index(self.workers_) + @property + def workers(self): + return self.workers_ -def ifft2(a, shape=None, axes=(-2,-1), overwrite_x=False): - x = _float_utils.__upcast_float16_array(a) - return _pydfti.ifftn(x, shape=shape, axes=axes, overwrite_x=overwrite_x) + @workers.setter + def workers(self, workers_val): + self.workerks_ = operator.index(workers_val) -def rfft(a, n=None, axis=-1, overwrite_x=False): - x = _float_utils.__upcast_float16_array(a) - return _pydfti.rfft(a, n=n, axis=axis, overwrite_x=overwrite_x) +_workers_global_settings = contextvars.ContextVar('scipy_backend_workers', default=_workers_data()) -def irfft(a, n=None, axis=-1, overwrite_x=False): - x = _float_utils.__upcast_float16_array(a) - return _pydfti.irfft(a, n=n, axis=axis, overwrite_x=overwrite_x) +def get_workers(): + "Gets the number of workers used by mkl_fft by default" + return _workers_global_settings.get().workers + + +@contextlib.contextmanager +def set_workers(n_workers): + "Set the value of workers used by default, returns the previous value" + nw = operator.index(n_workers) + token = None + try: + new_wd = _workers_data(nw) + token = _workers_global_settings.set(new_wd) + yield + finally: + if token: + _workers_global_settings.reset(token) + else: + raise ValueError + + +__all__ = ['fft', 'ifft', 'fft2', 'ifft2', 'fftn', 'ifftn', + 'rfft', 'irfft', 'rfft2', 'irfft2', 'rfftn', 'irfftn', + 'get_workers', 'set_workers', 'DftiBackend'] + +__ua_domain__ = "numpy.scipy.fft" + +def __ua_function__(method, args, kwargs): + """Fetch registered UA function.""" + fn = globals().get(method.__name__, None) + if fn is None: + return NotImplemented + return fn(*args, **kwargs) + + +class DftiBackend: + __ua_domain__ = "numpy.scipy.fft" + @staticmethod + def __ua_function__(method, args, kwargs): + """Fetch registered UA function.""" + fn = globals().get(method.__name__, None) + if fn is None: + return NotImplemented + return fn(*args, **kwargs) + + +def _cook_nd_args(a, s=None, axes=None, invreal=0): + if s is None: + shapeless = 1 + if axes is None: + s = list(a.shape) + else: + s = take(a.shape, axes) + else: + shapeless = 0 + s = list(s) + if axes is None: + axes = list(range(-len(s), 0)) + if len(s) != len(axes): + raise ValueError("Shape and axes have different lengths.") + if invreal and shapeless: + s[-1] = (a.shape[axes[-1]] - 1) * 2 + return s, axes + + +def _workers_to_num_threads(w): + """Handle conversion of workers to a positive number of threads in the + same way as scipy.fft.helpers._workers. + """ + if w is None: + return _workers_global_settings.get().workers + _w = operator.index(w) + if (_w == 0): + raise ValueError("Number of workers must not be zero") + if (_w < 0): + ub = os.cpu_count() + _w += ub + 1 + if _w <= 0: + raise ValueError("workers value out of range; got {}, must not be" + " less than {}".format(w, -ub)) + return _w + + +class Workers: + def __init__(self, workers): + self.workers = workers + self.n_threads = _workers_to_num_threads(workers) + + def __enter__(self): + try: + self.prev_num_threads = mkl.set_num_threads_local(self.n_threads) + except: + raise ValueError("Class argument {} result in invalid number of threads {}".format(self.workers, self.n_threads)) + + def __exit__(self, *args): + # restore old value + mkl.set_num_threads_local(self.prev_num_threads) + + +def _check_norm(norm): + if norm not in (None, "ortho", "forward", "backward"): + raise ValueError( + ("Invalid norm value {} should be None, " + "\"ortho\", \"forward\", or \"backward\".").format(norm)) + +def _check_plan(plan): + if plan is None: + return + raise NotImplementedError( + f"Passing a precomputed plan with value={plan} is currently not supported" + ) + + +def _frwd_sc_1d(n, s): + nn = n if n else s + return 1/nn if nn != 0 else 1 + + +def _frwd_sc_nd(s, axes, x_shape): + ss = s if s is not None else x_shape + if axes is not None: + nn = prod([ss[ai] for ai in axes]) + else: + nn = prod(ss) + return 1/nn if nn != 0 else 1 + + +def _ortho_sc_1d(n, s): + return sqrt(_frwd_sc_1d(n, s)) + + +def _compute_1d_fwd_scale(norm, n, s): + if norm in (None, "backward"): + fsc = 1.0 + elif norm == "forward": + fsc = _frwd_sc_1d(n, s) + elif norm == "ortho": + fsc = _ortho_sc_1d(n, s) + else: + _check_norm(norm) + return fsc + + +def _compute_nd_fwd_scale(norm, s, axes, x_shape): + if norm in (None, "backward"): + fsc = 1.0 + elif norm == "forward": + fsc = _frwd_sc_nd(s, axes, x_shape) + elif norm == "ortho": + fsc = sqrt(_frwd_sc_nd(s, axes, x_shape)) + else: + _check_norm(norm) + return fsc + + +def fft(a, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, plan=None): + try: + x = _float_utils.__supported_array_or_not_implemented(a) + except ValueError: + return NotImplemented + if x is NotImplemented: + return x + fsc = _compute_1d_fwd_scale(norm, n, x.shape[axis]) + _check_plan(plan) + with Workers(workers): + output = mkl_fft.fft(x, n=n, axis=axis, overwrite_x=overwrite_x, fwd_scale=fsc) + return output + + +def ifft(a, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, plan=None): + try: + x = _float_utils.__supported_array_or_not_implemented(a) + except ValueError: + return NotImplemented + if x is NotImplemented: + return x + fsc = _compute_1d_fwd_scale(norm, n, x.shape[axis]) + _check_plan(plan) + with Workers(workers): + output = mkl_fft.ifft(x, n=n, axis=axis, overwrite_x=overwrite_x, fwd_scale=fsc) + return output + + +def fft2(a, s=None, axes=(-2,-1), norm=None, overwrite_x=False, workers=None, plan=None): + try: + x = _float_utils.__supported_array_or_not_implemented(a) + except ValueError: + return NotImplemented + if x is NotImplemented: + return x + fsc = _compute_nd_fwd_scale(norm, s, axes, x.shape) + _check_plan(plan) + with Workers(workers): + output = mkl_fft.fftn(x, shape=s, axes=axes, overwrite_x=overwrite_x, fwd_scale=fsc) + return output + + +def ifft2(a, s=None, axes=(-2,-1), norm=None, overwrite_x=False, workers=None, plan=None): + try: + x = _float_utils.__supported_array_or_not_implemented(a) + except ValueError: + return NotImplemented + if x is NotImplemented: + return x + fsc = _compute_nd_fwd_scale(norm, s, axes, x.shape) + _check_plan(plan) + with Workers(workers): + output = mkl_fft.ifftn(x, shape=s, axes=axes, overwrite_x=overwrite_x, fwd_scale=fsc) + return output + + +def fftn(a, s=None, axes=None, norm=None, overwrite_x=False, workers=None, plan=None): + try: + x = _float_utils.__supported_array_or_not_implemented(a) + except ValueError: + return NotImplemented + if x is NotImplemented: + return x + fsc = _compute_nd_fwd_scale(norm, s, axes, x.shape) + _check_plan(plan) + with Workers(workers): + output = mkl_fft.fftn(x, shape=s, axes=axes, overwrite_x=overwrite_x, fwd_scale=fsc) + return output + + +def ifftn(a, s=None, axes=None, norm=None, overwrite_x=False, workers=None, plan=None): + try: + x = _float_utils.__supported_array_or_not_implemented(a) + except ValueError: + return NotImplemented + if x is NotImplemented: + return x + fsc = _compute_nd_fwd_scale(norm, s, axes, x.shape) + _check_plan(plan) + with Workers(workers): + output = mkl_fft.ifftn(x, shape=s, axes=axes, overwrite_x=overwrite_x, fwd_scale=fsc) + return output + + +def rfft(a, n=None, axis=-1, norm=None, workers=None, plan=None): + try: + x = _float_utils.__supported_array_or_not_implemented(a) + except ValueError: + return NotImplemented + if x is NotImplemented: + return x + fsc = _compute_1d_fwd_scale(norm, n, x.shape[axis]) + _check_plan(plan) + with Workers(workers): + output = mkl_fft.rfft(x, n=n, axis=axis, fwd_scale=fsc) + return output + + +def irfft(a, n=None, axis=-1, norm=None, workers=None, plan=None): + try: + x = _float_utils.__supported_array_or_not_implemented(a) + except ValueError: + return NotImplemented + if x is NotImplemented: + return x + nn = n if n else 2*(x.shape[axis]-1) + fsc = _compute_1d_fwd_scale(norm, nn, x.shape[axis]) + _check_plan(plan) + with Workers(workers): + output = mkl_fft.irfft(x, n=n, axis=axis, fwd_scale=fsc) + return output + + +def _compute_nd_fwd_scale_for_rfft(norm, s, axes, x, invreal=False): + if norm in (None, "backward"): + fsc = 1.0 + elif norm == "forward": + s, axes = _cook_nd_args(x, s, axes, invreal=invreal) + fsc = _frwd_sc_nd(s, axes, x.shape) + elif norm == "ortho": + s, axes = _cook_nd_args(x, s, axes, invreal=invreal) + fsc = sqrt(_frwd_sc_nd(s, axes, x.shape)) + else: + _check_norm(norm) + return s, axes, fsc + + +def rfft2(a, s=None, axes=(-2, -1), norm=None, workers=None, plan=None): + try: + x = _float_utils.__supported_array_or_not_implemented(a) + except ValueError: + return NotImplemented + if x is NotImplemented: + return x + s, axes, fsc = _compute_nd_fwd_scale_for_rfft(norm, s, axes, x) + _check_plan(plan) + with Workers(workers): + output = mkl_fft.rfftn(x, s, axes, fwd_scale=fsc) + return output + + +def irfft2(a, s=None, axes=(-2, -1), norm=None, workers=None, plan=None): + try: + x = _float_utils.__supported_array_or_not_implemented(a) + except ValueError: + return NotImplemented + if x is NotImplemented: + return x + s, axes, fsc = _compute_nd_fwd_scale_for_rfft(norm, s, axes, x, invreal=True) + _check_plan(plan) + with Workers(workers): + output = mkl_fft.irfftn(x, s, axes, fwd_scale=fsc) + return output + + +def rfftn(a, s=None, axes=None, norm=None, workers=None, plan=None): + try: + x = _float_utils.__supported_array_or_not_implemented(a) + except ValueError: + return NotImplemented + if x is NotImplemented: + return x + s, axes, fsc = _compute_nd_fwd_scale_for_rfft(norm, s, axes, x) + _check_plan(plan) + with Workers(workers): + output = mkl_fft.rfftn(x, s, axes, fwd_scale=fsc) + return output + + +def irfftn(a, s=None, axes=None, norm=None, workers=None, plan=None): + try: + x = _float_utils.__supported_array_or_not_implemented(a) + except ValueError: + return NotImplemented + if x is NotImplemented: + return x + s, axes, fsc = _compute_nd_fwd_scale_for_rfft(norm, s, axes, x, invreal=True) + _check_plan(plan) + with Workers(workers): + output = mkl_fft.irfftn(x, s, axes, fwd_scale=fsc) + return output diff --git a/mkl_fft/_scipy_fft_backend.py b/mkl_fft/_scipy_fft_backend.py deleted file mode 100644 index cea17426..00000000 --- a/mkl_fft/_scipy_fft_backend.py +++ /dev/null @@ -1,424 +0,0 @@ -#!/usr/bin/env python -# Copyright (c) 2019-2025, Intel Corporation -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# * Redistributions of source code must retain the above copyright notice, -# this list of conditions and the following disclaimer. -# * Redistributions in binary form must reproduce the above copyright -# notice, this list of conditions and the following disclaimer in the -# documentation and/or other materials provided with the distribution. -# * Neither the name of Intel Corporation nor the names of its contributors -# may be used to endorse or promote products derived from this software -# without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -from . import _pydfti -from . import _float_utils -import mkl - -from numpy import (take, sqrt, prod) -import contextvars -import contextlib -import operator -import os - - -__doc__ = """ -This module implements interfaces mimicing `scipy.fft` module. - -It also provides DftiBackend class which can be used to set mkl_fft to be used -via `scipy.fft` namespace. - -:Example: - import scipy.fft - import mkl_fft._scipy_fft_backend as be - # Set mkl_fft to be used as backend of SciPy's FFT functions. - scipy.fft.set_global_backend(be) -""" - -class _cpu_max_threads_count: - def __init__(self): - self.cpu_count = None - self.max_threads_count = None - - def get_cpu_count(self): - if self.cpu_count is None: - max_threads = self.get_max_threads_count() - self.cpu_count = max_threads - return self.cpu_count - - def get_max_threads_count(self): - if self.max_threads_count is None: - self.max_threads_count = mkl.get_max_threads() - - return self.max_threads_count - - -class _workers_data: - def __init__(self, workers=None): - if workers: - self.workers_ = workers - else: - self.workers_ = _cpu_max_threads_count().get_cpu_count() - self.workers_ = operator.index(self.workers_) - - @property - def workers(self): - return self.workers_ - - @workers.setter - def workers(self, workers_val): - self.workerks_ = operator.index(workers_val) - - -_workers_global_settings = contextvars.ContextVar('scipy_backend_workers', default=_workers_data()) - - -def get_workers(): - "Gets the number of workers used by mkl_fft by default" - return _workers_global_settings.get().workers - - -@contextlib.contextmanager -def set_workers(n_workers): - "Set the value of workers used by default, returns the previous value" - nw = operator.index(n_workers) - token = None - try: - new_wd = _workers_data(nw) - token = _workers_global_settings.set(new_wd) - yield - finally: - if token: - _workers_global_settings.reset(token) - else: - raise ValueError - - -__all__ = ['fft', 'ifft', 'fft2', 'ifft2', 'fftn', 'ifftn', - 'rfft', 'irfft', 'rfft2', 'irfft2', 'rfftn', 'irfftn', - 'get_workers', 'set_workers', 'DftiBackend'] - -__ua_domain__ = "numpy.scipy.fft" - -def __ua_function__(method, args, kwargs): - """Fetch registered UA function.""" - fn = globals().get(method.__name__, None) - if fn is None: - return NotImplemented - return fn(*args, **kwargs) - - -class DftiBackend: - __ua_domain__ = "numpy.scipy.fft" - @staticmethod - def __ua_function__(method, args, kwargs): - """Fetch registered UA function.""" - fn = globals().get(method.__name__, None) - if fn is None: - return NotImplemented - return fn(*args, **kwargs) - - -def _cook_nd_args(a, s=None, axes=None, invreal=0): - if s is None: - shapeless = 1 - if axes is None: - s = list(a.shape) - else: - s = take(a.shape, axes) - else: - shapeless = 0 - s = list(s) - if axes is None: - axes = list(range(-len(s), 0)) - if len(s) != len(axes): - raise ValueError("Shape and axes have different lengths.") - if invreal and shapeless: - s[-1] = (a.shape[axes[-1]] - 1) * 2 - return s, axes - - -def _workers_to_num_threads(w): - """Handle conversion of workers to a positive number of threads in the - same way as scipy.fft.helpers._workers. - """ - if w is None: - return _workers_global_settings.get().workers - _w = operator.index(w) - if (_w == 0): - raise ValueError("Number of workers must not be zero") - if (_w < 0): - ub = os.cpu_count() - _w += ub + 1 - if _w <= 0: - raise ValueError("workers value out of range; got {}, must not be" - " less than {}".format(w, -ub)) - return _w - - -class Workers: - def __init__(self, workers): - self.workers = workers - self.n_threads = _workers_to_num_threads(workers) - - def __enter__(self): - try: - self.prev_num_threads = mkl.set_num_threads_local(self.n_threads) - except: - raise ValueError("Class argument {} result in invalid number of threads {}".format(self.workers, self.n_threads)) - - def __exit__(self, *args): - # restore old value - mkl.set_num_threads_local(self.prev_num_threads) - - -def _check_norm(norm): - if norm not in (None, "ortho", "forward", "backward"): - raise ValueError( - ("Invalid norm value {} should be None, " - "\"ortho\", \"forward\", or \"backward\".").format(norm)) - -def _check_plan(plan): - if plan is None: - return - raise NotImplementedError( - f"Passing a precomputed plan with value={plan} is currently not supported" - ) - - -def _frwd_sc_1d(n, s): - nn = n if n else s - return 1/nn if nn != 0 else 1 - - -def _frwd_sc_nd(s, axes, x_shape): - ss = s if s is not None else x_shape - if axes is not None: - nn = prod([ss[ai] for ai in axes]) - else: - nn = prod(ss) - return 1/nn if nn != 0 else 1 - - -def _ortho_sc_1d(n, s): - return sqrt(_frwd_sc_1d(n, s)) - - -def _compute_1d_fwd_scale(norm, n, s): - if norm in (None, "backward"): - fsc = 1.0 - elif norm == "forward": - fsc = _frwd_sc_1d(n, s) - elif norm == "ortho": - fsc = _ortho_sc_1d(n, s) - else: - _check_norm(norm) - return fsc - - -def _compute_nd_fwd_scale(norm, s, axes, x_shape): - if norm in (None, "backward"): - fsc = 1.0 - elif norm == "forward": - fsc = _frwd_sc_nd(s, axes, x_shape) - elif norm == "ortho": - fsc = sqrt(_frwd_sc_nd(s, axes, x_shape)) - else: - _check_norm(norm) - return fsc - - -def fft(a, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, plan=None): - try: - x = _float_utils.__supported_array_or_not_implemented(a) - except ValueError: - return NotImplemented - if x is NotImplemented: - return x - fsc = _compute_1d_fwd_scale(norm, n, x.shape[axis]) - _check_plan(plan) - with Workers(workers): - output = _pydfti.fft(x, n=n, axis=axis, overwrite_x=overwrite_x, fwd_scale=fsc) - return output - - -def ifft(a, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, plan=None): - try: - x = _float_utils.__supported_array_or_not_implemented(a) - except ValueError: - return NotImplemented - if x is NotImplemented: - return x - fsc = _compute_1d_fwd_scale(norm, n, x.shape[axis]) - _check_plan(plan) - with Workers(workers): - output = _pydfti.ifft(x, n=n, axis=axis, overwrite_x=overwrite_x, fwd_scale=fsc) - return output - - -def fft2(a, s=None, axes=(-2,-1), norm=None, overwrite_x=False, workers=None, plan=None): - try: - x = _float_utils.__supported_array_or_not_implemented(a) - except ValueError: - return NotImplemented - if x is NotImplemented: - return x - fsc = _compute_nd_fwd_scale(norm, s, axes, x.shape) - _check_plan(plan) - with Workers(workers): - output = _pydfti.fftn(x, shape=s, axes=axes, overwrite_x=overwrite_x, fwd_scale=fsc) - return output - - -def ifft2(a, s=None, axes=(-2,-1), norm=None, overwrite_x=False, workers=None, plan=None): - try: - x = _float_utils.__supported_array_or_not_implemented(a) - except ValueError: - return NotImplemented - if x is NotImplemented: - return x - fsc = _compute_nd_fwd_scale(norm, s, axes, x.shape) - _check_plan(plan) - with Workers(workers): - output = _pydfti.ifftn(x, shape=s, axes=axes, overwrite_x=overwrite_x, fwd_scale=fsc) - return output - - -def fftn(a, s=None, axes=None, norm=None, overwrite_x=False, workers=None, plan=None): - try: - x = _float_utils.__supported_array_or_not_implemented(a) - except ValueError: - return NotImplemented - if x is NotImplemented: - return x - fsc = _compute_nd_fwd_scale(norm, s, axes, x.shape) - _check_plan(plan) - with Workers(workers): - output = _pydfti.fftn(x, shape=s, axes=axes, overwrite_x=overwrite_x, fwd_scale=fsc) - return output - - -def ifftn(a, s=None, axes=None, norm=None, overwrite_x=False, workers=None, plan=None): - try: - x = _float_utils.__supported_array_or_not_implemented(a) - except ValueError: - return NotImplemented - if x is NotImplemented: - return x - fsc = _compute_nd_fwd_scale(norm, s, axes, x.shape) - _check_plan(plan) - with Workers(workers): - output = _pydfti.ifftn(x, shape=s, axes=axes, overwrite_x=overwrite_x, fwd_scale=fsc) - return output - - -def rfft(a, n=None, axis=-1, norm=None, workers=None, plan=None): - try: - x = _float_utils.__supported_array_or_not_implemented(a) - except ValueError: - return NotImplemented - if x is NotImplemented: - return x - fsc = _compute_1d_fwd_scale(norm, n, x.shape[axis]) - _check_plan(plan) - with Workers(workers): - output = _pydfti.rfft_numpy(x, n=n, axis=axis, fwd_scale=fsc) - return output - - -def irfft(a, n=None, axis=-1, norm=None, workers=None, plan=None): - try: - x = _float_utils.__supported_array_or_not_implemented(a) - except ValueError: - return NotImplemented - if x is NotImplemented: - return x - nn = n if n else 2*(x.shape[axis]-1) - fsc = _compute_1d_fwd_scale(norm, nn, x.shape[axis]) - _check_plan(plan) - with Workers(workers): - output = _pydfti.irfft_numpy(x, n=n, axis=axis, fwd_scale=fsc) - return output - - -def _compute_nd_fwd_scale_for_rfft(norm, s, axes, x, invreal=False): - if norm in (None, "backward"): - fsc = 1.0 - elif norm == "forward": - s, axes = _cook_nd_args(x, s, axes, invreal=invreal) - fsc = _frwd_sc_nd(s, axes, x.shape) - elif norm == "ortho": - s, axes = _cook_nd_args(x, s, axes, invreal=invreal) - fsc = sqrt(_frwd_sc_nd(s, axes, x.shape)) - else: - _check_norm(norm) - return s, axes, fsc - - -def rfft2(a, s=None, axes=(-2, -1), norm=None, workers=None, plan=None): - try: - x = _float_utils.__supported_array_or_not_implemented(a) - except ValueError: - return NotImplemented - if x is NotImplemented: - return x - s, axes, fsc = _compute_nd_fwd_scale_for_rfft(norm, s, axes, x) - _check_plan(plan) - with Workers(workers): - output = _pydfti.rfftn_numpy(x, s, axes, fwd_scale=fsc) - return output - - -def irfft2(a, s=None, axes=(-2, -1), norm=None, workers=None, plan=None): - try: - x = _float_utils.__supported_array_or_not_implemented(a) - except ValueError: - return NotImplemented - if x is NotImplemented: - return x - s, axes, fsc = _compute_nd_fwd_scale_for_rfft(norm, s, axes, x, invreal=True) - _check_plan(plan) - with Workers(workers): - output = _pydfti.irfftn_numpy(x, s, axes, fwd_scale=fsc) - return output - - -def rfftn(a, s=None, axes=None, norm=None, workers=None, plan=None): - try: - x = _float_utils.__supported_array_or_not_implemented(a) - except ValueError: - return NotImplemented - if x is NotImplemented: - return x - s, axes, fsc = _compute_nd_fwd_scale_for_rfft(norm, s, axes, x) - _check_plan(plan) - with Workers(workers): - output = _pydfti.rfftn_numpy(x, s, axes, fwd_scale=fsc) - return output - - -def irfftn(a, s=None, axes=None, norm=None, workers=None, plan=None): - try: - x = _float_utils.__supported_array_or_not_implemented(a) - except ValueError: - return NotImplemented - if x is NotImplemented: - return x - s, axes, fsc = _compute_nd_fwd_scale_for_rfft(norm, s, axes, x, invreal=True) - _check_plan(plan) - with Workers(workers): - output = _pydfti.irfftn_numpy(x, s, axes, fwd_scale=fsc) - return output diff --git a/mkl_fft/_scipy_fftpack.py b/mkl_fft/_scipy_fftpack.py new file mode 100644 index 00000000..867f44c8 --- /dev/null +++ b/mkl_fft/_scipy_fftpack.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python +# Copyright (c) 2025, Intel Corporation +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of Intel Corporation nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from . import _pydfti as mkl_fft +from . import _float_utils + +__all__ = ['fft', 'ifft', 'fftn', 'ifftn', 'fft2', 'ifft2', 'rfft', 'irfft'] + + +def fft(a, n=None, axis=-1, overwrite_x=False): + x = _float_utils.__upcast_float16_array(a) + return mkl_fft.fft(x, n=n, axis=axis, overwrite_x=overwrite_x) + + +def ifft(a, n=None, axis=-1, overwrite_x=False): + x = _float_utils.__upcast_float16_array(a) + return mkl_fft.ifft(x, n=n, axis=axis, overwrite_x=overwrite_x) + + +def fftn(a, shape=None, axes=None, overwrite_x=False): + x = _float_utils.__upcast_float16_array(a) + return mkl_fft.fftn(x, shape=shape, axes=axes, overwrite_x=overwrite_x) + + +def ifftn(a, shape=None, axes=None, overwrite_x=False): + x = _float_utils.__upcast_float16_array(a) + return mkl_fft.ifftn(x, shape=shape, axes=axes, overwrite_x=overwrite_x) + + +def fft2(a, shape=None, axes=(-2,-1), overwrite_x=False): + x = _float_utils.__upcast_float16_array(a) + return mkl_fft.fftn(x, shape=shape, axes=axes, overwrite_x=overwrite_x) + + +def ifft2(a, shape=None, axes=(-2,-1), overwrite_x=False): + x = _float_utils.__upcast_float16_array(a) + return mkl_fft.ifftn(x, shape=shape, axes=axes, overwrite_x=overwrite_x) + + +def rfft(a, n=None, axis=-1, overwrite_x=False): + x = _float_utils.__upcast_float16_array(a) + return mkl_fft.rfftpack(a, n=n, axis=axis, overwrite_x=overwrite_x) + + +def irfft(a, n=None, axis=-1, overwrite_x=False): + x = _float_utils.__upcast_float16_array(a) + return mkl_fft.irfftpack(a, n=n, axis=axis, overwrite_x=overwrite_x) diff --git a/mkl_fft/interfaces/scipy_fft.py b/mkl_fft/interfaces/scipy_fft.py index 9e6d57e3..80f9368a 100644 --- a/mkl_fft/interfaces/scipy_fft.py +++ b/mkl_fft/interfaces/scipy_fft.py @@ -24,4 +24,4 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -from .._scipy_fft_backend import * +from .._scipy_fft import * diff --git a/mkl_fft/tests/test_fft1d.py b/mkl_fft/tests/test_fft1d.py index da4464e1..c226533e 100644 --- a/mkl_fft/tests/test_fft1d.py +++ b/mkl_fft/tests/test_fft1d.py @@ -341,7 +341,7 @@ def test_array6(self): assert_allclose(y1, y2, atol=2e-15) -class Test_mklfft_rfft(TestCase): +class Test_mklfft_rfftpack(TestCase): def setUp(self): rnd.seed(1234567) self.v1 = rnd.randn(16) @@ -350,14 +350,14 @@ def setUp(self): def test1(self): x = self.v1.copy() - f1 = mkl_fft.rfft(x) - f2 = mkl_fft.irfft(f1) + f1 = mkl_fft.rfftpack(x) + f2 = mkl_fft.irfftpack(f1) assert_allclose(f2,x) def test2(self): x = self.v1.copy() - f1 = mkl_fft.irfft(x) - f2 = mkl_fft.rfft(f1) + f1 = mkl_fft.irfftpack(x) + f2 = mkl_fft.rfftpack(f1) assert_allclose(f2,x) def test3(self): @@ -365,8 +365,8 @@ def test3(self): for ovwr_x in [True, False]: for dt, atol in zip([np.float32, np.float64], [2e-7, 2e-15]): x = self.m2.copy().astype(dt) - f1 = mkl_fft.rfft(x, axis=a, overwrite_x=ovwr_x) - f2 = mkl_fft.irfft(f1, axis=a, overwrite_x=ovwr_x) + f1 = mkl_fft.rfftpack(x, axis=a, overwrite_x=ovwr_x) + f2 = mkl_fft.irfftpack(f1, axis=a, overwrite_x=ovwr_x) assert_allclose(f2, self.m2.astype(dt), atol=atol, err_msg=(a, ovwr_x)) def test4(self): @@ -374,8 +374,8 @@ def test4(self): for ovwr_x in [True, False]: for dt, atol in zip([np.float32, np.float64], [2e-7, 2e-15]): x = self.m2.copy().astype(dt) - f1 = mkl_fft.irfft(x, axis=a, overwrite_x=ovwr_x) - f2 = mkl_fft.rfft(f1, axis=a, overwrite_x=ovwr_x) + f1 = mkl_fft.irfftpack(x, axis=a, overwrite_x=ovwr_x) + f2 = mkl_fft.rfftpack(f1, axis=a, overwrite_x=ovwr_x) assert_allclose(f2, self.m2.astype(dt), atol=atol) def test5(self): @@ -383,7 +383,7 @@ def test5(self): for ovwr_x in [True, False]: for dt, atol in zip([np.float32, np.float64], [4e-7, 4e-15]): x = self.t3.copy().astype(dt) - f1 = mkl_fft.irfft(x, axis=a, overwrite_x=ovwr_x) - f2 = mkl_fft.rfft(f1, axis=a, overwrite_x=ovwr_x) + f1 = mkl_fft.irfftpack(x, axis=a, overwrite_x=ovwr_x) + f2 = mkl_fft.rfftpack(f1, axis=a, overwrite_x=ovwr_x) assert_allclose(f2, self.t3.astype(dt), atol=atol) diff --git a/mkl_fft/tests/test_fftnd.py b/mkl_fft/tests/test_fftnd.py index 4c6c3d7e..20c73e4e 100644 --- a/mkl_fft/tests/test_fftnd.py +++ b/mkl_fft/tests/test_fftnd.py @@ -147,14 +147,14 @@ def test_cf_contig(self): f2 = mkl_fft.fft(d_fcont, axis=a) assert_allclose(f1, f2, rtol=r_tol, atol=a_tol) - def test_rfftn_numpy(self): - """Test that rfftn_numpy works as expected""" + def test_rfftn(self): + """Test that rfftn works as expected""" axes = [(0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0)] for x in [self.ad, self.af]: for a in axes: r_tol, a_tol = _get_rtol_atol(x) - rfft_tr = mkl_fft.rfftn_numpy(np.transpose(x, a)) - tr_rfft = np.transpose(mkl_fft.rfftn_numpy(x, axes=a), a) + rfft_tr = mkl_fft.rfftn(np.transpose(x, a)) + tr_rfft = np.transpose(mkl_fft.rfftn(x, axes=a), a) assert_allclose(rfft_tr, tr_rfft, rtol=r_tol, atol=a_tol) def test_gh64(self):