diff --git a/.github/workflows/conda-package-cf.yml b/.github/workflows/conda-package-cf.yml index d5cffb8c..f95bc95d 100644 --- a/.github/workflows/conda-package-cf.yml +++ b/.github/workflows/conda-package-cf.yml @@ -117,7 +117,7 @@ jobs: - name: Install mkl_fft run: | CHANNELS="-c $GITHUB_WORKSPACE/channel ${{ env.CHANNELS }}" - conda create -n ${{ env.TEST_ENV_NAME }} python=${{ matrix.python_ver }} ${{ matrix.numpy }} $PACKAGE_NAME pytest $CHANNELS + conda create -n ${{ env.TEST_ENV_NAME }} python=${{ matrix.python_ver }} ${{ matrix.numpy }} $PACKAGE_NAME pytest scipy $CHANNELS # Test installed packages conda list -n ${{ env.TEST_ENV_NAME }} - name: Run tests @@ -268,7 +268,7 @@ jobs: SET PACKAGE_VERSION=%%F ) SET "TEST_DEPENDENCIES=pytest pytest-cov" - conda install -n ${{ env.TEST_ENV_NAME }} ${{ env.PACKAGE_NAME }}=%PACKAGE_VERSION% %TEST_DEPENDENCIES% python=${{ matrix.python }} ${{ matrix.numpy }} -c ${{ env.workdir }}/channel ${{ env.CHANNELS }} + conda install -n ${{ env.TEST_ENV_NAME }} ${{ env.PACKAGE_NAME }}=%PACKAGE_VERSION% %TEST_DEPENDENCIES% python=${{ matrix.python }} ${{ matrix.numpy }} scipy -c ${{ env.workdir }}/channel ${{ env.CHANNELS }} - name: Report content of test environment shell: cmd /C CALL {0} run: | diff --git a/.github/workflows/conda-package.yml b/.github/workflows/conda-package.yml index 3325b9f8..b2d383a1 100644 --- a/.github/workflows/conda-package.yml +++ b/.github/workflows/conda-package.yml @@ -116,7 +116,7 @@ jobs: - name: Install mkl_fft run: | CHANNELS="-c $GITHUB_WORKSPACE/channel ${{ env.CHANNELS }}" - conda create -n ${{ env.TEST_ENV_NAME }} python=${{ matrix.python }} $PACKAGE_NAME pytest $CHANNELS + conda create -n ${{ env.TEST_ENV_NAME }} python=${{ matrix.python }} $PACKAGE_NAME scipy pytest $CHANNELS # Test installed packages conda list -n ${{ env.TEST_ENV_NAME }} - name: Run tests @@ -267,7 +267,7 @@ jobs: SET PACKAGE_VERSION=%%F ) SET "TEST_DEPENDENCIES=pytest pytest-cov" - conda install -n ${{ env.TEST_ENV_NAME }} ${{ env.PACKAGE_NAME }}=%PACKAGE_VERSION% %TEST_DEPENDENCIES% python=${{ matrix.python }} -c ${{ env.workdir }}/channel ${{ env.CHANNELS }} + conda install -n ${{ env.TEST_ENV_NAME }} ${{ env.PACKAGE_NAME }}=%PACKAGE_VERSION% %TEST_DEPENDENCIES% python=${{ matrix.python }} scipy -c ${{ env.workdir }}/channel ${{ env.CHANNELS }} - name: Report content of test environment shell: cmd /C CALL {0} run: | diff --git a/.gitignore b/.gitignore index 630ac69d..921f8885 100644 --- a/.gitignore +++ b/.gitignore @@ -3,7 +3,8 @@ build/ mkl_fft.egg-info/ mkl_fft/__pycache__/ mkl_fft/_pydfti.c -mkl_fft/_pydfti.cpython-310-x86_64-linux-gnu.so +mkl_fft/_pydfti.cpython*.so mkl_fft/interfaces/__pycache__/ mkl_fft/src/mklfft.c mkl_fft/tests/__pycache__/ +mkl_fft/tests/from_scipy/__pycache__/ diff --git a/conda-recipe-cf/meta.yaml b/conda-recipe-cf/meta.yaml index fc9a9a4e..44f7de15 100644 --- a/conda-recipe-cf/meta.yaml +++ b/conda-recipe-cf/meta.yaml @@ -32,6 +32,7 @@ test: - pytest -v --pyargs mkl_fft requires: - pytest + - scipy imports: - mkl_fft - mkl_fft.interfaces diff --git a/conda-recipe/meta.yaml b/conda-recipe/meta.yaml index 98905ac5..a947b709 100644 --- a/conda-recipe/meta.yaml +++ b/conda-recipe/meta.yaml @@ -32,6 +32,7 @@ test: - pytest -v --pyargs mkl_fft requires: - pytest + - scipy imports: - mkl_fft - mkl_fft.interfaces diff --git a/mkl_fft/__init__.py b/mkl_fft/__init__.py index d2c41550..92336168 100644 --- a/mkl_fft/__init__.py +++ b/mkl_fft/__init__.py @@ -24,6 +24,7 @@ # 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. + import mkl_fft.interfaces from . import _init_helper diff --git a/mkl_fft/_fft_utils.py b/mkl_fft/_fft_utils.py new file mode 100644 index 00000000..9cfcd799 --- /dev/null +++ b/mkl_fft/_fft_utils.py @@ -0,0 +1,51 @@ +#!/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. + +import numpy as np + +__all__ = ["check_norm", "compute_fwd_scale"] + + +def check_norm(norm): + if norm not in (None, "ortho", "forward", "backward"): + raise ValueError( + f"Invalid norm value {norm} should be None, 'ortho', 'forward', " + "or 'backward'." + ) + + +def compute_fwd_scale(norm, n, shape): + check_norm(norm) + if norm in (None, "backward"): + return 1.0 + + ss = n if n is not None else shape + nn = np.prod(ss) + fsc = 1 / nn if nn != 0 else 1 + if norm == "forward": + return fsc + else: # norm == "ortho" + return np.sqrt(fsc) diff --git a/mkl_fft/_float_utils.py b/mkl_fft/_float_utils.py index 9dba61f1..864bbf4e 100644 --- a/mkl_fft/_float_utils.py +++ b/mkl_fft/_float_utils.py @@ -91,5 +91,5 @@ def __supported_array_or_not_implemented(x): if hasattr(np, "complex256"): black_list.append(np.complex256) if __x.dtype in black_list: - return NotImplemented + raise NotImplementedError return __x diff --git a/mkl_fft/_numpy_fft.py b/mkl_fft/_numpy_fft.py index 1330c60a..abec47d8 100644 --- a/mkl_fft/_numpy_fft.py +++ b/mkl_fft/_numpy_fft.py @@ -74,36 +74,59 @@ import warnings import numpy as np -from numpy import array, conjugate, prod, sqrt, take -from . import _float_utils -from . import _pydfti as mkl_fft # pylint: disable=no-name-in-module +import mkl_fft +from ._fft_utils import check_norm, compute_fwd_scale +from ._float_utils import __downcast_float128_array -def _compute_fwd_scale(norm, n, shape): - _check_norm(norm) - if norm in (None, "backward"): - return 1.0 - ss = n if n is not None else shape - nn = prod(ss) - fsc = 1 / nn if nn != 0 else 1 - if norm == "forward": - return fsc - else: # norm == "ortho" - return sqrt(fsc) - - -def _check_norm(norm): - if norm not in (None, "ortho", "forward", "backward"): - raise ValueError( - f"Invalid norm value {norm} should be None, 'ortho', 'forward', " - "or 'backward'." +# copied with modifications from: +# https://github.com/numpy/numpy/blob/main/numpy/fft/_pocketfft.py +def _cook_nd_args(a, s=None, axes=None, invreal=False): + if s is None: + shapeless = True + if axes is None: + s = list(a.shape) + else: + s = np.take(a.shape, axes) + else: + shapeless = False + s = list(s) + if axes is None: + if not shapeless and np.__version__ >= "2.0": + msg = ( + "`axes` should not be `None` if `s` is not `None` " + "(Deprecated in NumPy 2.0). In a future version of NumPy, " + "this will raise an error and `s[i]` will correspond to " + "the size along the transformed axis specified by " + "`axes[i]`. To retain current behaviour, pass a sequence " + "[0, ..., k-1] to `axes` for an array of dimension k." + ) + warnings.warn(msg, DeprecationWarning, stacklevel=3) + 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 + if None in s and np.__version__ >= "2.0": + msg = ( + "Passing an array containing `None` values to `s` is " + "deprecated in NumPy 2.0 and will raise an error in " + "a future version of NumPy. To use the default behaviour " + "of the corresponding 1-D transform, pass the value matching " + "the default for its `n` parameter. To use the default " + "behaviour for every axis, the `s` argument can be omitted." ) + warnings.warn(msg, DeprecationWarning, stacklevel=3) + # use the whole input array along axis `i` if `s[i] == -1 or None` + s = [a.shape[_a] if _s in [-1, None] else _s for _s, _a in zip(s, axes)] + + return s, axes def _swap_direction(norm): - _check_norm(norm) + check_norm(norm) _swap_direction_map = { "backward": "forward", None: "forward", @@ -218,8 +241,8 @@ def fft(a, n=None, axis=-1, norm=None): """ - x = _float_utils.__downcast_float128_array(a) - fsc = _compute_fwd_scale(norm, n, x.shape[axis]) + x = __downcast_float128_array(a) + fsc = compute_fwd_scale(norm, n, x.shape[axis]) return trycall(mkl_fft.fft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc}) @@ -311,8 +334,8 @@ def ifft(a, n=None, axis=-1, norm=None): """ - x = _float_utils.__downcast_float128_array(a) - fsc = _compute_fwd_scale(norm, n, x.shape[axis]) + x = __downcast_float128_array(a) + fsc = compute_fwd_scale(norm, n, x.shape[axis]) return trycall(mkl_fft.ifft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc}) @@ -402,8 +425,8 @@ def rfft(a, n=None, axis=-1, norm=None): """ - x = _float_utils.__downcast_float128_array(a) - fsc = _compute_fwd_scale(norm, n, x.shape[axis]) + x = __downcast_float128_array(a) + fsc = compute_fwd_scale(norm, n, x.shape[axis]) return trycall(mkl_fft.rfft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc}) @@ -495,8 +518,8 @@ def irfft(a, n=None, axis=-1, norm=None): """ - x = _float_utils.__downcast_float128_array(a) - fsc = _compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1)) + x = __downcast_float128_array(a) + fsc = compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1)) return trycall( mkl_fft.irfft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc} @@ -581,10 +604,10 @@ def hfft(a, n=None, axis=-1, norm=None): """ norm = _swap_direction(norm) - x = _float_utils.__downcast_float128_array(a) - x = array(x, copy=True, dtype=complex) - conjugate(x, out=x) - fsc = _compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1)) + x = __downcast_float128_array(a) + x = np.array(x, copy=True, dtype=complex) + np.conjugate(x, out=x) + fsc = compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1)) return trycall( mkl_fft.irfft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc} @@ -651,61 +674,18 @@ def ihfft(a, n=None, axis=-1, norm=None): # The copy may be required for multithreading. norm = _swap_direction(norm) - x = _float_utils.__downcast_float128_array(a) - x = array(x, copy=True, dtype=float) - fsc = _compute_fwd_scale(norm, n, x.shape[axis]) + x = __downcast_float128_array(a) + x = np.array(x, copy=True, dtype=float) + fsc = compute_fwd_scale(norm, n, x.shape[axis]) output = trycall( mkl_fft.rfft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc} ) - conjugate(output, out=output) + np.conjugate(output, out=output) return output -# copied from: https://github.com/numpy/numpy/blob/main/numpy/fft/_pocketfft.py -def _cook_nd_args(a, s=None, axes=None, invreal=False): - if s is None: - shapeless = True - if axes is None: - s = list(a.shape) - else: - s = take(a.shape, axes) - else: - shapeless = False - s = list(s) - if axes is None: - if not shapeless and np.__version__ >= "2.0": - msg = ( - "`axes` should not be `None` if `s` is not `None` " - "(Deprecated in NumPy 2.0). In a future version of NumPy, " - "this will raise an error and `s[i]` will correspond to " - "the size along the transformed axis specified by " - "`axes[i]`. To retain current behaviour, pass a sequence " - "[0, ..., k-1] to `axes` for an array of dimension k." - ) - warnings.warn(msg, DeprecationWarning, stacklevel=3) - 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 - if None in s and np.__version__ >= "2.0": - msg = ( - "Passing an array containing `None` values to `s` is " - "deprecated in NumPy 2.0 and will raise an error in " - "a future version of NumPy. To use the default behaviour " - "of the corresponding 1-D transform, pass the value matching " - "the default for its `n` parameter. To use the default " - "behaviour for every axis, the `s` argument can be omitted." - ) - warnings.warn(msg, DeprecationWarning, stacklevel=3) - # use the whole input array along axis `i` if `s[i] == -1 or None` - s = [a.shape[_a] if _s in [-1, None] else _s for _s, _a in zip(s, axes)] - - return s, axes - - def fftn(a, s=None, axes=None, norm=None): """ Compute the N-dimensional discrete Fourier Transform. @@ -806,9 +786,9 @@ def fftn(a, s=None, axes=None, norm=None): """ - x = _float_utils.__downcast_float128_array(a) + x = __downcast_float128_array(a) s, axes = _cook_nd_args(x, s, axes) - fsc = _compute_fwd_scale(norm, s, x.shape) + fsc = compute_fwd_scale(norm, s, x.shape) return trycall(mkl_fft.fftn, (x,), {"s": s, "axes": axes, "fwd_scale": fsc}) @@ -913,9 +893,9 @@ def ifftn(a, s=None, axes=None, norm=None): """ - x = _float_utils.__downcast_float128_array(a) + x = __downcast_float128_array(a) s, axes = _cook_nd_args(x, s, axes) - fsc = _compute_fwd_scale(norm, s, x.shape) + fsc = compute_fwd_scale(norm, s, x.shape) return trycall( mkl_fft.ifftn, (x,), {"s": s, "axes": axes, "fwd_scale": fsc} @@ -1201,9 +1181,9 @@ def rfftn(a, s=None, axes=None, norm=None): """ - x = _float_utils.__downcast_float128_array(a) + x = __downcast_float128_array(a) s, axes = _cook_nd_args(x, s, axes) - fsc = _compute_fwd_scale(norm, s, x.shape) + fsc = compute_fwd_scale(norm, s, x.shape) return trycall( mkl_fft.rfftn, (x,), {"s": s, "axes": axes, "fwd_scale": fsc} @@ -1345,9 +1325,9 @@ def irfftn(a, s=None, axes=None, norm=None): """ - x = _float_utils.__downcast_float128_array(a) + x = __downcast_float128_array(a) s, axes = _cook_nd_args(x, s, axes, invreal=True) - fsc = _compute_fwd_scale(norm, s, x.shape) + fsc = compute_fwd_scale(norm, s, x.shape) return trycall( mkl_fft.irfftn, (x,), {"s": s, "axes": axes, "fwd_scale": fsc} diff --git a/mkl_fft/_pydfti.pyx b/mkl_fft/_pydfti.pyx index 89968c5d..1f262d9f 100644 --- a/mkl_fft/_pydfti.pyx +++ b/mkl_fft/_pydfti.pyx @@ -1335,11 +1335,11 @@ def ifftn(x, s=None, axes=None, overwrite_x=False, fwd_scale=1.0): def rfft2(x, s=None, axes=(-2, -1), fwd_scale=1.0): - return rfftn(x, s=s, axes=axes, fsc=fwd_scale) + return rfftn(x, s=s, axes=axes, fwd_scale=fwd_scale) def irfft2(x, s=None, axes=(-2, -1), fwd_scale=1.0): - return irfftn(x, s=s, axes=axes, fsc=fwd_scale) + return irfftn(x, s=s, axes=axes, fwd_scale=fwd_scale) def _remove_axis(s, axes, axis_to_remove): diff --git a/mkl_fft/_scipy_fft.py b/mkl_fft/_scipy_fft.py index ec41c1b7..c194ead0 100644 --- a/mkl_fft/_scipy_fft.py +++ b/mkl_fft/_scipy_fft.py @@ -30,10 +30,12 @@ import os import mkl -from numpy import prod, sqrt, take +import numpy as np -from . import _float_utils -from . import _pydfti as mkl_fft # pylint: disable=no-name-in-module +import mkl_fft + +from ._fft_utils import compute_fwd_scale +from ._float_utils import __supported_array_or_not_implemented __doc__ = """ This module implements interfaces mimicing `scipy.fft` module. @@ -152,25 +154,6 @@ def __ua_function__(method, args, kwargs): 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. @@ -213,95 +196,65 @@ def __exit__(self, *args): 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 is not None else s - return 1 / nn if nn != 0 else 1 - - -def _frwd_sc_nd(s, x_shape): - ss = s if s is not None else x_shape - nn = prod(ss) - return 1 / nn if nn != 0 else 1 - - -def _ortho_sc_1d(n, s): - return sqrt(_frwd_sc_1d(n, s)) + if plan is not None: + raise NotImplementedError( + f"Passing a precomputed plan with value={plan} is currently not supported" + ) -def _compute_1d_fwd_scale(norm, n, s): - if norm in (None, "backward"): - return 1.0 - elif norm == "forward": - return _frwd_sc_1d(n, s) - elif norm == "ortho": - return _ortho_sc_1d(n, s) +def _cook_nd_args(a, s=None, axes=None, invreal=False): + if s is None: + shapeless = True + if axes is None: + s = list(a.shape) + else: + s = np.take(a.shape, axes) else: - _check_norm(norm) + shapeless = False + 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 _compute_nd_fwd_scale(norm, s, axes, x_shape): - if norm in (None, "backward"): - return 1.0 - elif norm == "forward": - return _frwd_sc_nd(s, x_shape) - elif norm == "ortho": - return sqrt(_frwd_sc_nd(s, x_shape)) - else: - _check_norm(norm) +def _validate_input(a): + try: + x = __supported_array_or_not_implemented(a) + except ValueError: + raise NotImplementedError + + return x 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) + x = _validate_input(a) + fsc = compute_fwd_scale(norm, n, x.shape[axis]) + with Workers(workers): - output = mkl_fft.fft( + return 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) + x = _validate_input(a) + fsc = compute_fwd_scale(norm, n, x.shape[axis]) + with Workers(workers): - output = mkl_fft.ifft( + return mkl_fft.ifft( x, n=n, axis=axis, overwrite_x=overwrite_x, fwd_scale=fsc ) - return output def fft2( @@ -313,19 +266,16 @@ def fft2( 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, s=s, axes=axes, overwrite_x=overwrite_x, fwd_scale=fsc - ) - return output + + return fftn( + a, + s=s, + axes=axes, + norm=norm, + overwrite_x=overwrite_x, + workers=workers, + plan=plan, + ) def ifft2( @@ -337,154 +287,87 @@ def ifft2( 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, s=s, axes=axes, overwrite_x=overwrite_x, fwd_scale=fsc - ) - return output + + return ifftn( + a, + s=s, + axes=axes, + norm=norm, + overwrite_x=overwrite_x, + workers=workers, + plan=plan, + ) 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) + x = _validate_input(a) + fsc = compute_fwd_scale(norm, s, x.shape) + with Workers(workers): - output = mkl_fft.fftn( + return mkl_fft.fftn( x, s=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) + x = _validate_input(a) + fsc = compute_fwd_scale(norm, s, x.shape) + with Workers(workers): - output = mkl_fft.ifftn( + return mkl_fft.ifftn( x, s=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) + x = _validate_input(a) + fsc = compute_fwd_scale(norm, n, x.shape[axis]) + with Workers(workers): - output = mkl_fft.rfft(x, n=n, axis=axis, fwd_scale=fsc) - return output + return mkl_fft.rfft(x, n=n, axis=axis, fwd_scale=fsc) 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) + x = _validate_input(a) + fsc = compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1)) + 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"): - return s, axes, 1.0 - elif norm == "forward": - s, axes = _cook_nd_args(x, s, axes, invreal=invreal) - return s, axes, _frwd_sc_nd(s, x.shape) - elif norm == "ortho": - s, axes = _cook_nd_args(x, s, axes, invreal=invreal) - return s, axes, sqrt(_frwd_sc_nd(s, x.shape)) - else: - _check_norm(norm) + return mkl_fft.irfft(x, n=n, axis=axis, fwd_scale=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 + + return rfftn(a, s=s, axes=axes, norm=norm, workers=workers, plan=plan) 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 + + return irfftn(a, s=s, axes=axes, norm=norm, workers=workers, plan=plan) 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) + x = _validate_input(a) + s, axes = _cook_nd_args(x, s, axes) + fsc = compute_fwd_scale(norm, s, x.shape) + with Workers(workers): - output = mkl_fft.rfftn(x, s, axes, fwd_scale=fsc) - return output + return mkl_fft.rfftn(x, s, axes, fwd_scale=fsc) 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) + x = _validate_input(a) + s, axes = _cook_nd_args(x, s, axes, invreal=True) + fsc = compute_fwd_scale(norm, s, x.shape) + with Workers(workers): - output = mkl_fft.irfftn(x, s, axes, fwd_scale=fsc) - return output + return mkl_fft.irfftn(x, s, axes, fwd_scale=fsc) diff --git a/mkl_fft/_scipy_fftpack.py b/mkl_fft/_scipy_fftpack.py index 330b5e51..0e29c84d 100644 --- a/mkl_fft/_scipy_fftpack.py +++ b/mkl_fft/_scipy_fftpack.py @@ -24,47 +24,48 @@ # 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 _float_utils -from . import _pydfti as mkl_fft # pylint: disable=no-name-in-module +import mkl_fft + +from ._float_utils import __upcast_float16_array __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) + x = __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) + x = __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) + x = __upcast_float16_array(a) + return mkl_fft.fftn(x, s=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) + x = __upcast_float16_array(a) + return mkl_fft.ifftn(x, s=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) + x = __upcast_float16_array(a) + return mkl_fft.fftn(x, s=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) + x = __upcast_float16_array(a) + return mkl_fft.ifftn(x, s=shape, axes=axes, overwrite_x=overwrite_x) def rfft(a, n=None, axis=-1, overwrite_x=False): - x = _float_utils.__upcast_float16_array(a) + x = __upcast_float16_array(a) return mkl_fft.rfftpack(x, 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) + x = __upcast_float16_array(a) return mkl_fft.irfftpack(x, n=n, axis=axis, overwrite_x=overwrite_x) diff --git a/mkl_fft/tests/from_scipy/test_basic.py b/mkl_fft/tests/from_scipy/test_basic.py new file mode 100644 index 00000000..9e720975 --- /dev/null +++ b/mkl_fft/tests/from_scipy/test_basic.py @@ -0,0 +1,653 @@ +# This file includes tests from scipy.fft module: +# https://github.com/scipy/scipy/blob/main/scipy/fft/tests/test_basic.py + +# TODO: remove when hfft* functions are added +# pylint: disable=no-member + +import multiprocessing +import queue +import threading + +import numpy as np +import pytest +import scipy +from numpy.random import random +from numpy.testing import assert_allclose, assert_array_almost_equal +from pytest import raises as assert_raises + +# pylint: disable=possibly-used-before-assignment +if scipy.__version__ < "1.12": + # scipy from Intel channel is 1.10 + pytest.skip( + "This test file needs scipy>=1.12", + allow_module_level=True, + ) +elif scipy.__version__ < "1.14": + # For python<=3.9, scipy<1.14 is installed + # pylint: disable=no-name-in-module + from scipy._lib._array_api import size as xp_size +else: + from scipy._lib._array_api import xp_size + +from scipy._lib._array_api import is_numpy, xp_assert_close, xp_assert_equal + +if np.__version__ < "2.0": + from numpy import transpose as permute_dims +else: + from numpy import permute_dims + +import mkl_fft.interfaces.scipy_fft as fft + +skip_xp_backends = pytest.mark.skip_xp_backends + + +@pytest.fixture(params=[np]) +def xp(request): + return request.param + + +# Expected input dtypes. Note that `scipy.fft` is more flexible for numpy, +# but for C2C transforms like `fft.fft`, the array API standard only mandates +# that complex dtypes should work, float32/float64 aren't guaranteed to. +def get_expected_input_dtype(func, xp): + if func in [ + fft.fft, + fft.fftn, + fft.fft2, + fft.ifft, + fft.ifftn, + fft.ifft2, + # TODO: fft.hfft, + # TODO: fft.hfftn, + # TODO: fft.hfft2, + fft.irfft, + fft.irfftn, + fft.irfft2, + ]: + dtype = xp.complex128 + elif func in [ + fft.rfft, + fft.rfftn, + fft.rfft2, + # TODO: fft.ihfft, + # TODO: fft.ihfftn, + # TODO: fft.ihfft2, + ]: + dtype = xp.float64 + else: + raise ValueError(f"Unknown FFT function: {func}") + + return dtype + + +def fft1(x): + L = len(x) + phase = -2j * np.pi * (np.arange(L) / float(L)) + phase = np.arange(L).reshape(-1, 1) * phase + return np.sum(x * np.exp(phase), axis=1) + + +class TestFFT: + + def test_identity(self, xp): + maxlen = 512 + x = xp.asarray(random(maxlen) + 1j * random(maxlen)) + xr = xp.asarray(random(maxlen)) + # Check some powers of 2 and some primes + for i in [1, 2, 16, 128, 512, 53, 149, 281, 397]: + xp_assert_close(fft.ifft(fft.fft(x[0:i])), x[0:i]) + xp_assert_close(fft.irfft(fft.rfft(xr[0:i]), i), xr[0:i]) + + @skip_xp_backends( + np_only=True, reason="significant overhead for some backends" + ) + def test_identity_extensive(self, xp): + maxlen = 512 + x = xp.asarray(random(maxlen) + 1j * random(maxlen)) + xr = xp.asarray(random(maxlen)) + for i in range(1, maxlen): + xp_assert_close(fft.ifft(fft.fft(x[0:i])), x[0:i]) + xp_assert_close(fft.irfft(fft.rfft(xr[0:i]), i), xr[0:i]) + + def test_fft(self, xp): + x = random(30) + 1j * random(30) + expect = xp.asarray(fft1(x)) + x = xp.asarray(x) + xp_assert_close(fft.fft(x), expect) + xp_assert_close(fft.fft(x, norm="backward"), expect) + xp_assert_close( + fft.fft(x, norm="ortho"), + expect / xp.sqrt(xp.asarray(30, dtype=xp.float64)), + ) + xp_assert_close(fft.fft(x, norm="forward"), expect / 30) + + @skip_xp_backends(np_only=True, reason="some backends allow `n=0`") + def test_fft_n(self, xp): + x = xp.asarray([1, 2, 3], dtype=xp.complex128) + assert_raises(ValueError, fft.fft, x, 0) + + def test_ifft(self, xp): + x = xp.asarray(random(30) + 1j * random(30)) + xp_assert_close(fft.ifft(fft.fft(x)), x) + for norm in ["backward", "ortho", "forward"]: + xp_assert_close(fft.ifft(fft.fft(x, norm=norm), norm=norm), x) + + def test_fft2(self, xp): + x = xp.asarray(random((30, 20)) + 1j * random((30, 20))) + expect = fft.fft(fft.fft(x, axis=1), axis=0) + xp_assert_close(fft.fft2(x), expect) + xp_assert_close(fft.fft2(x, norm="backward"), expect) + xp_assert_close( + fft.fft2(x, norm="ortho"), + expect / xp.sqrt(xp.asarray(30 * 20, dtype=xp.float64)), + ) + xp_assert_close(fft.fft2(x, norm="forward"), expect / (30 * 20)) + + def test_ifft2(self, xp): + x = xp.asarray(random((30, 20)) + 1j * random((30, 20))) + expect = fft.ifft(fft.ifft(x, axis=1), axis=0) + xp_assert_close(fft.ifft2(x), expect) + xp_assert_close(fft.ifft2(x, norm="backward"), expect) + xp_assert_close( + fft.ifft2(x, norm="ortho"), + expect * xp.sqrt(xp.asarray(30 * 20, dtype=xp.float64)), + ) + xp_assert_close(fft.ifft2(x, norm="forward"), expect * (30 * 20)) + + def test_fftn(self, xp): + x = xp.asarray(random((30, 20, 10)) + 1j * random((30, 20, 10))) + expect = fft.fft(fft.fft(fft.fft(x, axis=2), axis=1), axis=0) + xp_assert_close(fft.fftn(x), expect) + xp_assert_close(fft.fftn(x, norm="backward"), expect) + xp_assert_close( + fft.fftn(x, norm="ortho"), + expect / xp.sqrt(xp.asarray(30 * 20 * 10, dtype=xp.float64)), + ) + xp_assert_close(fft.fftn(x, norm="forward"), expect / (30 * 20 * 10)) + + def test_ifftn(self, xp): + x = xp.asarray(random((30, 20, 10)) + 1j * random((30, 20, 10))) + expect = fft.ifft(fft.ifft(fft.ifft(x, axis=2), axis=1), axis=0) + xp_assert_close(fft.ifftn(x), expect, rtol=1e-7) + xp_assert_close(fft.ifftn(x, norm="backward"), expect, rtol=1e-7) + xp_assert_close( + fft.ifftn(x, norm="ortho"), + fft.ifftn(x) * xp.sqrt(xp.asarray(30 * 20 * 10, dtype=xp.float64)), + ) + xp_assert_close( + fft.ifftn(x, norm="forward"), expect * (30 * 20 * 10), rtol=1e-7 + ) + + def test_rfft(self, xp): + x = xp.asarray(random(29), dtype=xp.float64) + for n in [xp_size(x), 2 * xp_size(x)]: + for norm in [None, "backward", "ortho", "forward"]: + xp_assert_close( + fft.rfft(x, n=n, norm=norm), + fft.fft(xp.asarray(x, dtype=xp.complex128), n=n, norm=norm)[ + : (n // 2 + 1) + ], + ) + xp_assert_close( + fft.rfft(x, n=n, norm="ortho"), + fft.rfft(x, n=n) / xp.sqrt(xp.asarray(n, dtype=xp.float64)), + ) + + def test_irfft(self, xp): + x = xp.asarray(random(30)) + xp_assert_close(fft.irfft(fft.rfft(x)), x) + for norm in ["backward", "ortho", "forward"]: + xp_assert_close(fft.irfft(fft.rfft(x, norm=norm), norm=norm), x) + + def test_rfft2(self, xp): + x = xp.asarray(random((30, 20)), dtype=xp.float64) + expect = fft.fft2(xp.asarray(x, dtype=xp.complex128))[:, :11] + xp_assert_close(fft.rfft2(x), expect) + xp_assert_close(fft.rfft2(x, norm="backward"), expect) + xp_assert_close( + fft.rfft2(x, norm="ortho"), + expect / xp.sqrt(xp.asarray(30 * 20, dtype=xp.float64)), + ) + xp_assert_close(fft.rfft2(x, norm="forward"), expect / (30 * 20)) + + def test_irfft2(self, xp): + x = xp.asarray(random((30, 20))) + xp_assert_close(fft.irfft2(fft.rfft2(x)), x) + for norm in ["backward", "ortho", "forward"]: + xp_assert_close(fft.irfft2(fft.rfft2(x, norm=norm), norm=norm), x) + + def test_rfftn(self, xp): + x = xp.asarray(random((30, 20, 10)), dtype=xp.float64) + expect = fft.fftn(xp.asarray(x, dtype=xp.complex128))[:, :, :6] + xp_assert_close(fft.rfftn(x), expect) + xp_assert_close(fft.rfftn(x, norm="backward"), expect) + xp_assert_close( + fft.rfftn(x, norm="ortho"), + expect / xp.sqrt(xp.asarray(30 * 20 * 10, dtype=xp.float64)), + ) + xp_assert_close(fft.rfftn(x, norm="forward"), expect / (30 * 20 * 10)) + + def test_irfftn(self, xp): + x = xp.asarray(random((30, 20, 10))) + xp_assert_close(fft.irfftn(fft.rfftn(x)), x) + for norm in ["backward", "ortho", "forward"]: + xp_assert_close(fft.irfftn(fft.rfftn(x, norm=norm), norm=norm), x) + + @pytest.mark.skip("hfft is not supported") + def test_hfft(self, xp): + x = random(14) + 1j * random(14) + x_herm = np.concatenate((random(1), x, random(1))) + x = np.concatenate((x_herm, x[::-1].conj())) + x = xp.asarray(x) + x_herm = xp.asarray(x_herm) + expect = xp.real(fft.fft(x)) + xp_assert_close(fft.hfft(x_herm), expect) + xp_assert_close(fft.hfft(x_herm, norm="backward"), expect) + xp_assert_close( + fft.hfft(x_herm, norm="ortho"), + expect / xp.sqrt(xp.asarray(30, dtype=xp.float64)), + ) + xp_assert_close(fft.hfft(x_herm, norm="forward"), expect / 30) + + @pytest.mark.skip("ihfft is not supported") + def test_ihfft(self, xp): + x = random(14) + 1j * random(14) + x_herm = np.concatenate((random(1), x, random(1))) + x = np.concatenate((x_herm, x[::-1].conj())) + x = xp.asarray(x) + x_herm = xp.asarray(x_herm) + xp_assert_close(fft.ihfft(fft.hfft(x_herm)), x_herm) + for norm in ["backward", "ortho", "forward"]: + xp_assert_close( + fft.ihfft(fft.hfft(x_herm, norm=norm), norm=norm), x_herm + ) + + @pytest.mark.skip("hfft2 is not supported") + def test_hfft2(self, xp): + x = xp.asarray(random((30, 20))) + xp_assert_close(fft.hfft2(fft.ihfft2(x)), x) + for norm in ["backward", "ortho", "forward"]: + xp_assert_close(fft.hfft2(fft.ihfft2(x, norm=norm), norm=norm), x) + + @pytest.mark.skip("ihfft2 is not supported") + def test_ihfft2(self, xp): + x = xp.asarray(random((30, 20)), dtype=xp.float64) + expect = fft.ifft2(xp.asarray(x, dtype=xp.complex128))[:, :11] + xp_assert_close(fft.ihfft2(x), expect) + xp_assert_close(fft.ihfft2(x, norm="backward"), expect) + xp_assert_close( + fft.ihfft2(x, norm="ortho"), + expect * xp.sqrt(xp.asarray(30 * 20, dtype=xp.float64)), + ) + xp_assert_close(fft.ihfft2(x, norm="forward"), expect * (30 * 20)) + + @pytest.mark.skip("hfftn is not supported") + def test_hfftn(self, xp): + x = xp.asarray(random((30, 20, 10))) + xp_assert_close(fft.hfftn(fft.ihfftn(x)), x) + for norm in ["backward", "ortho", "forward"]: + xp_assert_close(fft.hfftn(fft.ihfftn(x, norm=norm), norm=norm), x) + + @pytest.mark.skip("ihfftn is not supported") + def test_ihfftn(self, xp): + x = xp.asarray(random((30, 20, 10)), dtype=xp.float64) + expect = fft.ifftn(xp.asarray(x, dtype=xp.complex128))[:, :, :6] + xp_assert_close(expect, fft.ihfftn(x)) + xp_assert_close(expect, fft.ihfftn(x, norm="backward")) + xp_assert_close( + fft.ihfftn(x, norm="ortho"), + expect * xp.sqrt(xp.asarray(30 * 20 * 10, dtype=xp.float64)), + ) + xp_assert_close(fft.ihfftn(x, norm="forward"), expect * (30 * 20 * 10)) + + def _check_axes(self, op, xp): + dtype = get_expected_input_dtype(op, xp) + x = xp.asarray(random((30, 20, 10)), dtype=dtype) + axes = [ + (0, 1, 2), + (0, 2, 1), + (1, 0, 2), + (1, 2, 0), + (2, 0, 1), + (2, 1, 0), + ] + + for a in axes: + op_tr = op(permute_dims(x, axes=a)) + tr_op = permute_dims(op(x, axes=a), axes=a) + xp_assert_close(op_tr, tr_op) + + @pytest.mark.parametrize("op", [fft.fftn, fft.ifftn, fft.rfftn, fft.irfftn]) + def test_axes_standard(self, op, xp): + self._check_axes(op, xp) + + # TODO: add this test + # @pytest.mark.parametrize("op", [fft.hfftn, fft.ihfftn]) + # def test_axes_non_standard(self, op, xp): + # self._check_axes(op, xp) + + @pytest.mark.parametrize("op", [fft.fftn, fft.ifftn, fft.rfftn, fft.irfftn]) + def test_axes_subset_with_shape_standard(self, op, xp): + dtype = get_expected_input_dtype(op, xp) + x = xp.asarray(random((16, 8, 4)), dtype=dtype) + axes = [(0, 1, 2), (0, 2, 1), (1, 2, 0)] + + for a in axes: + # different shape on the first two axes + shape = tuple( + [ + 2 * x.shape[ax] if ax in a[:2] else x.shape[ax] + for ax in range(x.ndim) + ] + ) + # transform only the first two axes + op_tr = op(permute_dims(x, axes=a), s=shape[:2], axes=(0, 1)) + tr_op = permute_dims(op(x, s=shape[:2], axes=a[:2]), axes=a) + xp_assert_close(op_tr, tr_op) + + @pytest.mark.parametrize( + "op", + [ + fft.fft2, + fft.ifft2, + fft.rfft2, + fft.irfft2, + # TODO: fft.hfft2, + # TODO: fft.ihfft2, + # TODO: fft.hfftn, + # TODO: fft.ihfftn, + ], + ) + def test_axes_subset_with_shape_non_standard(self, op, xp): + dtype = get_expected_input_dtype(op, xp) + x = xp.asarray(random((16, 8, 4)), dtype=dtype) + axes = [(0, 1, 2), (0, 2, 1), (1, 2, 0)] + + for a in axes: + # different shape on the first two axes + shape = tuple( + [ + 2 * x.shape[ax] if ax in a[:2] else x.shape[ax] + for ax in range(x.ndim) + ] + ) + # transform only the first two axes + op_tr = op(permute_dims(x, axes=a), s=shape[:2], axes=(0, 1)) + tr_op = permute_dims(op(x, s=shape[:2], axes=a[:2]), axes=a) + xp_assert_close(op_tr, tr_op) + + def test_all_1d_norm_preserving(self, xp): + # verify that round-trip transforms are norm-preserving + x = xp.asarray(random(30), dtype=xp.float64) + + x_norm = np.linalg.norm(x, ord=2) + n = xp_size(x) * 2 + func_pairs = [ + (fft.rfft, fft.irfft), + # hfft: order so the first function takes x.size samples + # (necessary for comparison to x_norm above) + # TODO: (fft.ihfft, fft.hfft), + # functions that expect complex dtypes at the end + (fft.fft, fft.ifft), + ] + for forw, back in func_pairs: + if forw == fft.fft: + x = xp.asarray(x, dtype=xp.complex128) + x_norm = np.linalg.norm(x, ord=2) + for n in [xp_size(x), 2 * xp_size(x)]: + for norm in ["backward", "ortho", "forward"]: + tmp = forw(x, n=n, norm=norm) + tmp = back(tmp, n=n, norm=norm) + xp_assert_close(np.linalg.norm(tmp, ord=2), x_norm) + + @pytest.mark.skip("float16 and longdouble are not supported") + @skip_xp_backends(np_only=True) + @pytest.mark.parametrize("dtype", [np.float16, np.longdouble]) + def test_dtypes_nonstandard(self, dtype, xp): + x = random(30).astype(dtype) + out_dtypes = {np.float16: np.complex64, np.longdouble: np.clongdouble} + x_complex = x.astype(out_dtypes[dtype]) + + res_fft = fft.ifft(fft.fft(x)) + res_rfft = fft.irfft(fft.rfft(x)) + # TODO: res_hfft = fft.hfft(fft.ihfft(x), x.shape[0]) + # Check both numerical results and exact dtype matches + assert_array_almost_equal(res_fft, x_complex) + assert_array_almost_equal(res_rfft, x) + # TODO: assert_array_almost_equal(res_hfft, x) + assert res_fft.dtype == x_complex.dtype + assert res_rfft.dtype == np.result_type(np.float32, x.dtype) + # TODO: assert res_hfft.dtype == np.result_type(np.float32, x.dtype) + + @pytest.mark.parametrize("dtype", ["float32", "float64"]) + def test_dtypes_real(self, dtype, xp): + x = xp.asarray(random(30), dtype=getattr(xp, dtype)) + + res_rfft = fft.irfft(fft.rfft(x)) + # TODO: res_hfft = fft.hfft(fft.ihfft(x), x.shape[0]) + # Check both numerical results and exact dtype matches + xp_assert_close(res_rfft, x, rtol=1e-06, atol=1e-06) + # TODO: xp_assert_close(res_hfft, x) + + @pytest.mark.parametrize("dtype", ["complex64", "complex128"]) + def test_dtypes_complex(self, dtype, xp): + rng = np.random.default_rng(1234) + x = xp.asarray(rng.random(30), dtype=getattr(xp, dtype)) + + res_fft = fft.ifft(fft.fft(x)) + # Check both numerical results and exact dtype matches + xp_assert_close(res_fft, x, rtol=1e-06, atol=1e-06) + + @skip_xp_backends( + np_only=True, reason="array-likes only supported for NumPy backend" + ) + @pytest.mark.parametrize( + "op", + [ + fft.fft, + fft.ifft, + fft.fft2, + fft.ifft2, + fft.fftn, + fft.ifftn, + fft.rfft, + fft.irfft, + fft.rfft2, + fft.irfft2, + fft.rfftn, + fft.irfftn, + # TODO: fft.hfft, + # TODO: fft.ihfft, + # TODO: fft.hfft2, + # TODO: fft.ihfft2, + # TODO: fft.hfftn, + # TODO: fft.ihfftn, + ], + ) + def test_array_like(self, xp, op): + x = [ + [[1.0, 1.0], [1.0, 1.0]], + [[1.0, 1.0], [1.0, 1.0]], + [[1.0, 1.0], [1.0, 1.0]], + ] + xp_assert_close(op(x), op(xp.asarray(x))) + + +@skip_xp_backends(np_only=True) +@pytest.mark.parametrize( + "dtype", + [ + np.float32, + np.float64, + # TODO: np.longdouble, + np.complex64, + np.complex128, + # TODO: np.clongdouble, + ], +) +@pytest.mark.parametrize("order", ["F", "non-contiguous"]) +@pytest.mark.parametrize( + "fft", [fft.fft, fft.fft2, fft.fftn, fft.ifft, fft.ifft2, fft.ifftn] +) +def test_fft_with_order(dtype, order, fft, xp): + # Check that FFT/IFFT produces identical results for C, Fortran and + # non contiguous arrays + np.random.seed(42) + X = np.random.rand(8, 7, 13).astype(dtype, copy=False) + if order == "F": + Y = np.asfortranarray(X) + else: + # Make a non contiguous array + Y = X[::-1] + X = np.ascontiguousarray(X[::-1]) + + if fft.__name__.endswith("fft"): + for axis in range(3): + X_res = fft(X, axis=axis) + Y_res = fft(Y, axis=axis) + assert_array_almost_equal(X_res, Y_res) + elif fft.__name__.endswith(("fft2", "fftn")): + axes = [(0, 1), (1, 2), (0, 2)] + if fft.__name__.endswith("fftn"): + axes.extend([(0,), (1,), (2,), None]) + for ax in axes: + X_res = fft(X, axes=ax) + Y_res = fft(Y, axes=ax) + assert_array_almost_equal(X_res, Y_res, decimal=4) + else: + raise ValueError + + +@skip_xp_backends(cpu_only=True) +class TestFFTThreadSafe: + threads = 16 + input_shape = (800, 200) + + def _test_mtsame(self, func, *args, xp=None): + def worker(args, q): + q.put(func(*args)) + + q = queue.Queue() + expected = func(*args) + + # Spin off a bunch of threads to call the same function simultaneously + t = [ + threading.Thread(target=worker, args=(args, q)) + for i in range(self.threads) + ] + [x.start() for x in t] + + [x.join() for x in t] + + # Make sure all threads returned the correct value + for _ in range(self.threads): + xp_assert_equal( + q.get(timeout=5), + expected, + err_msg="Function returned wrong value in multithreaded context", + ) + + def test_fft(self, xp): + a = xp.ones(self.input_shape, dtype=xp.complex128) + self._test_mtsame(fft.fft, a, xp=xp) + + def test_ifft(self, xp): + a = xp.full(self.input_shape, 1 + 0j) + self._test_mtsame(fft.ifft, a, xp=xp) + + def test_rfft(self, xp): + a = xp.ones(self.input_shape) + self._test_mtsame(fft.rfft, a, xp=xp) + + def test_irfft(self, xp): + a = xp.full(self.input_shape, 1 + 0j) + self._test_mtsame(fft.irfft, a, xp=xp) + + @pytest.mark.skip("hfft is not supported") + def test_hfft(self, xp): + a = xp.ones(self.input_shape, dtype=xp.complex64) + self._test_mtsame(fft.hfft, a, xp=xp) + + @pytest.mark.skip("ihfft is not supported") + def test_ihfft(self, xp): + a = xp.ones(self.input_shape) + self._test_mtsame(fft.ihfft, a, xp=xp) + + +@skip_xp_backends(np_only=True) +@pytest.mark.parametrize("func", [fft.fft, fft.ifft, fft.rfft, fft.irfft]) +def test_multiprocess(func, xp): + # Test that fft still works after fork (gh-10422) + + with multiprocessing.Pool(2) as p: + res = p.map(func, [np.ones(100) for _ in range(4)]) + + expect = func(np.ones(100)) + for x in res: + assert_allclose(x, expect) + + +class TestIRFFTN: + + def test_not_last_axis_success(self, xp): + ar, ai = np.random.random((2, 16, 8, 32)) + a = ar + 1j * ai + a = xp.asarray(a) + + axes = (-2,) + + # Should not raise error + fft.irfftn(a, axes=axes) + + +@pytest.mark.parametrize( + "func", + [ + fft.fft, + fft.ifft, + fft.rfft, + fft.irfft, + fft.fftn, + fft.ifftn, + fft.rfftn, + fft.irfftn, + # TODO: fft.hfft, + # TODO: fft.ihfft, + ], +) +def test_non_standard_params(func, xp): + if func in [fft.rfft, fft.rfftn]: # TODO: fft.ihfft + dtype = xp.float64 + else: + dtype = xp.complex128 + + x = xp.asarray([1, 2, 3], dtype=dtype) + # func(x) should not raise an exception + func(x) + + if is_numpy(xp): + func(x, workers=2) + else: + assert_raises(ValueError, func, x, workers=2) + + # `plan` param is not tested since SciPy does not use it currently + # but should be tested if it comes into use + + +@pytest.mark.parametrize("dtype", ["float32", "float64"]) +@pytest.mark.parametrize( + "func", + [ + fft.fft, + fft.ifft, + fft.irfft, + fft.fftn, + fft.ifftn, + fft.irfftn, + # TODO: fft.hfft, + ], +) +def test_real_input(func, dtype, xp): + x = xp.asarray([1, 2, 3], dtype=getattr(xp, dtype)) + # func(x) should not raise an exception + func(x) diff --git a/mkl_fft/tests/from_scipy/test_multithreading.py b/mkl_fft/tests/from_scipy/test_multithreading.py new file mode 100644 index 00000000..80626dbb --- /dev/null +++ b/mkl_fft/tests/from_scipy/test_multithreading.py @@ -0,0 +1,104 @@ +# This file includes tests from scipy.fft module: +# https://github.com/scipy/scipy/blob/main/scipy/fft/tests/test_multithreading.py.py + +import multiprocessing +import os + +import numpy as np +import pytest +from numpy.testing import assert_allclose + +import mkl_fft.interfaces.scipy_fft as fft + + +@pytest.fixture(scope="module") +def x(): + return np.random.randn(512, 128) # Must be large enough to qualify for mt + + +@pytest.mark.parametrize( + "func", + [ + fft.fft, + fft.ifft, + fft.fft2, + fft.ifft2, + fft.fftn, + fft.ifftn, + fft.rfft, + fft.irfft, + fft.rfft2, + fft.irfft2, + fft.rfftn, + fft.irfftn, + # TODO: fft.hfft, fft.ihfft, fft.hfft2, fft.ihfft2, fft.hfftn, fft.ihfftn, + # TODO: fft.dct, fft.idct, fft.dctn, fft.idctn, + # TODO: fft.dst, fft.idst, fft.dstn, fft.idstn, + ], +) +@pytest.mark.parametrize("workers", [2, -1]) +def test_threaded_same(x, func, workers): + expected = func(x, workers=1) + actual = func(x, workers=workers) + assert_allclose(actual, expected) + + +def _mt_fft(x): + return fft.fft(x, workers=2) + + +@pytest.mark.slow +def test_mixed_threads_processes(x): + # Test that the fft threadpool is safe to use before & after fork + + expect = fft.fft(x, workers=2) + + with multiprocessing.Pool(2) as p: + res = p.map(_mt_fft, [x for _ in range(4)]) + + for r in res: + assert_allclose(r, expect) + + fft.fft(x, workers=2) + + +def test_invalid_workers(x): + cpus = os.cpu_count() + + fft.ifft([1], workers=-cpus) + + with pytest.raises(ValueError, match="workers must not be zero"): + fft.fft(x, workers=0) + + with pytest.raises(ValueError, match="workers value out of range"): + fft.ifft(x, workers=-cpus - 1) + + +@pytest.mark.skip() +def test_set_get_workers(): + cpus = os.cpu_count() + assert fft.get_workers() == 1 + with fft.set_workers(4): + assert fft.get_workers() == 4 + + with fft.set_workers(-1): + assert fft.get_workers() == cpus + + assert fft.get_workers() == 4 + + assert fft.get_workers() == 1 + + with fft.set_workers(-cpus): + assert fft.get_workers() == 1 + + +@pytest.mark.skip("mkl_fft does not validate workers") +def test_set_workers_invalid(): + + with pytest.raises(ValueError, match="workers must not be zero"): + with fft.set_workers(0): + pass + + with pytest.raises(ValueError, match="workers value out of range"): + with fft.set_workers(-os.cpu_count() - 1): + pass diff --git a/mkl_fft/tests/test_pocketfft.py b/mkl_fft/tests/test_from_numpy.py similarity index 100% rename from mkl_fft/tests/test_pocketfft.py rename to mkl_fft/tests/test_from_numpy.py diff --git a/mkl_fft/tests/test_interfaces.py b/mkl_fft/tests/test_interfaces.py index 6eedc120..ba45557a 100644 --- a/mkl_fft/tests/test_interfaces.py +++ b/mkl_fft/tests/test_interfaces.py @@ -25,6 +25,7 @@ import numpy as np import pytest +from numpy.testing import assert_raises import mkl_fft.interfaces as mfi @@ -143,8 +144,7 @@ def _get_blacklisted_dtypes(): @pytest.mark.parametrize("dtype", _get_blacklisted_dtypes()) def test_scipy_no_support_for(dtype): x = np.ones(16, dtype=dtype) - w = mfi.scipy_fft.fft(x) - assert w is NotImplemented + assert_raises(NotImplementedError, mfi.scipy_fft.ifft, x) def test_scipy_fft_arg_validate(): diff --git a/pyproject.toml b/pyproject.toml index 320bcbf9..eefe397e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,7 +59,7 @@ readme = {file = "README.md", content-type = "text/markdown"} requires-python = ">=3.9,<3.13" [project.optional-dependencies] -test = ["pytest"] +test = ["pytest", "scipy"] [project.urls] Download = "http://github.com/IntelPython/mkl_fft" @@ -90,4 +90,4 @@ packages = ["mkl_fft", "mkl_fft.interfaces"] version = {attr = "mkl_fft._version.__version__"} [tool.setuptools.package-data] -"mkl_fft" = ["tests/*.py"] +"mkl_fft" = ["tests/**/*.py"]