From 2c4340a9de8fb77f0b46e3820e7ce2a732086337 Mon Sep 17 00:00:00 2001 From: Vahid Tavanashad Date: Wed, 2 Apr 2025 10:33:38 -0700 Subject: [PATCH 01/11] update _scipy_fft.py --- .github/workflows/build_pip.yaml | 1 + .github/workflows/conda-package-cf.yml | 4 +- .github/workflows/conda-package.yml | 4 +- conda-recipe-cf/meta.yaml | 1 + conda-recipe/meta.yaml | 1 + mkl_fft/_fft_utils.py | 51 ++ mkl_fft/_float_utils.py | 2 +- mkl_fft/_numpy_fft.py | 135 ++-- mkl_fft/_scipy_fft.py | 294 +++----- .../{test_pocketfft.py => test_from_numpy.py} | 0 mkl_fft/tests/test_from_scipy.py | 653 ++++++++++++++++++ mkl_fft/tests/test_interfaces.py | 4 +- pyproject.toml | 2 +- 13 files changed, 860 insertions(+), 292 deletions(-) create mode 100644 mkl_fft/_fft_utils.py rename mkl_fft/tests/{test_pocketfft.py => test_from_numpy.py} (100%) create mode 100644 mkl_fft/tests/test_from_scipy.py diff --git a/.github/workflows/build_pip.yaml b/.github/workflows/build_pip.yaml index d9c11e93..05e41488 100644 --- a/.github/workflows/build_pip.yaml +++ b/.github/workflows/build_pip.yaml @@ -54,6 +54,7 @@ jobs: run: | pip install --no-cache-dir cython pytest hypothesis pip install --no-cache-dir numpy ${{ matrix.use_pre }} + pip install --no-cache-dir scipy echo "CONDA_PREFFIX is '${CONDA_PREFIX}'" export MKLROOT=${CONDA_PREFIX} pip install -e . --no-build-isolation --verbose --no-deps diff --git a/.github/workflows/conda-package-cf.yml b/.github/workflows/conda-package-cf.yml index e87c62f8..c387fdbe 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 @@ -269,7 +269,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 d01f2dee..1caa9fee 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/conda-recipe-cf/meta.yaml b/conda-recipe-cf/meta.yaml index ebedddbb..7745b86c 100644 --- a/conda-recipe-cf/meta.yaml +++ b/conda-recipe-cf/meta.yaml @@ -33,6 +33,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 f5b042d0..60aefaf7 100644 --- a/conda-recipe/meta.yaml +++ b/conda-recipe/meta.yaml @@ -33,6 +33,7 @@ test: - pytest -v --pyargs mkl_fft requires: - pytest + - scipy imports: - mkl_fft - mkl_fft.interfaces diff --git a/mkl_fft/_fft_utils.py b/mkl_fft/_fft_utils.py new file mode 100644 index 00000000..58c152bf --- /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..f0e56bc9 100644 --- a/mkl_fft/_numpy_fft.py +++ b/mkl_fft/_numpy_fft.py @@ -74,32 +74,54 @@ 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 +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): @@ -218,7 +240,7 @@ def fft(a, n=None, axis=-1, norm=None): """ - x = _float_utils.__downcast_float128_array(a) + 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,7 +333,7 @@ def ifft(a, n=None, axis=-1, norm=None): """ - x = _float_utils.__downcast_float128_array(a) + 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,7 +424,7 @@ def rfft(a, n=None, axis=-1, norm=None): """ - x = _float_utils.__downcast_float128_array(a) + 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,7 +517,7 @@ def irfft(a, n=None, axis=-1, norm=None): """ - x = _float_utils.__downcast_float128_array(a) + x = __downcast_float128_array(a) fsc = _compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1)) return trycall( @@ -581,9 +603,9 @@ 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) + 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( @@ -651,61 +673,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) + 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,7 +785,7 @@ 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) @@ -913,7 +892,7 @@ 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) @@ -1201,7 +1180,7 @@ 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) @@ -1345,7 +1324,7 @@ 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) diff --git a/mkl_fft/_scipy_fft.py b/mkl_fft/_scipy_fft.py index ec41c1b7..3eae75c1 100644 --- a/mkl_fft/_scipy_fft.py +++ b/mkl_fft/_scipy_fft.py @@ -30,10 +30,11 @@ 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 +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 +153,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 +195,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 +265,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 +286,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/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_from_scipy.py b/mkl_fft/tests/test_from_scipy.py new file mode 100644 index 00000000..74e969d2 --- /dev/null +++ b/mkl_fft/tests/test_from_scipy.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) + # 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/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 d6106ec5..087ddd04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,7 +50,7 @@ classifiers = [ "Operating System :: POSIX", "Operating System :: Unix" ] -dependencies = ["numpy >=1.26.4", "mkl", "mkl-service"] +dependencies = ["numpy>=1.26.4", "mkl", "mkl-service"] description = "MKL-based FFT transforms for NumPy arrays" dynamic = ["version"] keywords = ["DFTI", "FFT", "Fourier", "MKL"] From 4edc66c1d8ef4acd4d5a288133a370e33b18caf2 Mon Sep 17 00:00:00 2001 From: Vahid Tavanashad Date: Thu, 3 Apr 2025 21:35:19 -0700 Subject: [PATCH 02/11] update CHaNGES.rst and README.md --- CHANGES.rst | 17 +++++++++++++++++ README.md | 12 ++++++------ 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/CHANGES.rst b/CHANGES.rst index e2f379bb..5e422b4b 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -2,6 +2,23 @@ mkl_fft changelog ================= +1.3.13 +====== + +migrate from `setup.py` to `pyproject.toml` + +includes support in virtual environment out of the box + +the original `mkl_fft.rfft` and `mkl_fft.irfft` are renamed to `mkl_fft.rfftpack` and `mkl_fft.irfftpack`, since they +replicate the behavior from the deprectaed `scipy.fftapck` module. + +`mkl_fft.rfft_numpy`, `mkl_fft.irfft_numpy`, `mkl_fft.rfft2_numpy`, `mkl_fft.irfft2_numpy`, `mkl_fft.rfftn_numpy` +and `mkl_fft.irfftn_numpy` are renamed to `mkl_fft.rfft`, `mkl_fft.irfft`, `mkl_fft.rfft2`, `mkl_fft.irfft2`, `mkl_fft.rfftn` +and `mkl_fft.irfftn`, respectively. (consistent with `numpy.fft` and `scipy.fft` modules) + +file `_scipy_fft_backend.py` is renamed to `_scipy_fft.py` since it replicates `scipy.fft` module (similar to file +`_numpy_fft.py` which replicates `numpy.fft` module) + 1.3.11 ====== diff --git a/README.md b/README.md index 859f2487..82d696e2 100644 --- a/README.md +++ b/README.md @@ -47,9 +47,9 @@ More details can be found in SciPy 2017 conference proceedings: --- -It implements the following functions: +`mkl_fft` implements the following functions: -### Complex transforms, similar to those in `scipy.fftpack`: +### Complex transforms, similar to those in `scipy.fft`: `fft(x, n=None, axis=-1, overwrite_x=False)` @@ -65,13 +65,13 @@ It implements the following functions: ### Real transforms -`rfft(x, n=None, axis=-1, overwrite_x=False)` - real 1D Fourier transform, like `scipy.fftpack.rfft` +`rfftpack(x, n=None, axis=-1, overwrite_x=False)` - real 1D Fourier transform, like `scipy.fftpack.rfft` -`rfft_numpy(x, n=None, axis=-1)` - real 1D Fourier transform, like `numpy.fft.rfft` +`rfft(x, n=None, axis=-1)` - real 1D Fourier transform, like `numpy.fft.rfft` -`rfft2_numpy(x, s=None, axes=(-2,-1))` - real 2D Fourier transform, like `numpy.fft.rfft2` +`rfft2(x, s=None, axes=(-2,-1))` - real 2D Fourier transform, like `numpy.fft.rfft2` -`rfftn_numpy(x, s=None, axes=None)` - real ND Fourier transform, like `numpy.fft.rfftn` +`rfftn(x, s=None, axes=None)` - real ND Fourier transform, like `numpy.fft.rfftn` ... and similar `irfft*` functions. From 1a9e5d091d0599e51248a267c6dd41829f567e1d Mon Sep 17 00:00:00 2001 From: Vahid Tavanashad Date: Thu, 3 Apr 2025 21:37:43 -0700 Subject: [PATCH 03/11] fix typo --- CHANGES.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGES.rst b/CHANGES.rst index 5e422b4b..eb8ab213 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -10,7 +10,7 @@ migrate from `setup.py` to `pyproject.toml` includes support in virtual environment out of the box the original `mkl_fft.rfft` and `mkl_fft.irfft` are renamed to `mkl_fft.rfftpack` and `mkl_fft.irfftpack`, since they -replicate the behavior from the deprectaed `scipy.fftapck` module. +replicate the behavior from the deprectaed `scipy.fftpack` module. `mkl_fft.rfft_numpy`, `mkl_fft.irfft_numpy`, `mkl_fft.rfft2_numpy`, `mkl_fft.irfft2_numpy`, `mkl_fft.rfftn_numpy` and `mkl_fft.irfftn_numpy` are renamed to `mkl_fft.rfft`, `mkl_fft.irfft`, `mkl_fft.rfft2`, `mkl_fft.irfft2`, `mkl_fft.rfftn` From 560afd3b97233f922fd78c0ecfa135bea0457613 Mon Sep 17 00:00:00 2001 From: Vahid Tavanashad Date: Fri, 4 Apr 2025 05:51:30 -0700 Subject: [PATCH 04/11] update formatting CHANGE.rst --- CHANGES.rst | 52 ++++++++++++++++++++++++++++------------------------ 1 file changed, 28 insertions(+), 24 deletions(-) diff --git a/CHANGES.rst b/CHANGES.rst index eb8ab213..103a6191 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -5,19 +5,20 @@ mkl_fft changelog 1.3.13 ====== -migrate from `setup.py` to `pyproject.toml` +migrate from :code:`setup.py` to :code:`pyproject.toml` includes support in virtual environment out of the box -the original `mkl_fft.rfft` and `mkl_fft.irfft` are renamed to `mkl_fft.rfftpack` and `mkl_fft.irfftpack`, since they -replicate the behavior from the deprectaed `scipy.fftpack` module. +the original :code:`mkl_fft.rfft` and :code:`mkl_fft.irfft` are renamed to :code:`mkl_fft.rfftpack` and :code:`mkl_fft.irfftpack`, +since they replicate the behavior from the deprectaed :code:`scipy.fftpack` module. -`mkl_fft.rfft_numpy`, `mkl_fft.irfft_numpy`, `mkl_fft.rfft2_numpy`, `mkl_fft.irfft2_numpy`, `mkl_fft.rfftn_numpy` -and `mkl_fft.irfftn_numpy` are renamed to `mkl_fft.rfft`, `mkl_fft.irfft`, `mkl_fft.rfft2`, `mkl_fft.irfft2`, `mkl_fft.rfftn` -and `mkl_fft.irfftn`, respectively. (consistent with `numpy.fft` and `scipy.fft` modules) +:code:`mkl_fft.rfft_numpy`, :code:`mkl_fft.irfft_numpy`, :code:`mkl_fft.rfft2_numpy`, :code:`mkl_fft.irfft2_numpy`, +:code:`mkl_fft.rfftn_numpy`, and :code:`mkl_fft.irfftn_numpy` are renamed to :code:`mkl_fft.rfft`, :code:`mkl_fft.irfft`, +:code:`mkl_fft.rfft2`, :code:`mkl_fft.irfft2`, `mkl_fft.rfftn`, and :code:`mkl_fft.irfftn`, respectively. +(consistent with :code:`numpy.fft` and :code:`scipy.fft` modules) -file `_scipy_fft_backend.py` is renamed to `_scipy_fft.py` since it replicates `scipy.fft` module (similar to file -`_numpy_fft.py` which replicates `numpy.fft` module) +file :code:`_scipy_fft_backend.py` is renamed to :code:`_scipy_fft.py` since it replicates :code:`scipy.fft` module +(similar to file :code:`_numpy_fft.py` which replicates :code:`numpy.fft` module) 1.3.11 ====== @@ -39,7 +40,7 @@ Updated code and build system to support NumPy 2.0 1.3.8 ===== -Added vendored `conv_template.py` from NumPy's distutils submodule to enable building of `mkl_fft` with +Added vendored :code:`conv_template.py` from NumPy's distutils submodule to enable building of :code:`mkl_fft` with NumPy >=1.25 and Python 3.12 1.3.7 @@ -54,7 +55,7 @@ Transitioned to Cython 3.0. ===== Updated numpy interface to support new in NumPy 1.20 supported values of norm keyword, such as "forward" and "backward". -To enable this, `mkl_fft` functions now support `forward_scale` parameter that defaults to 1. +To enable this, :code:`mkl_fft` functions now support `forward_scale` parameter that defaults to 1. Fixed issue #48. @@ -66,16 +67,17 @@ Includes bug fix #54 1.2.0 ===== -Due to removal of deprecated real-to-real FFT with `DFTI_CONJUGATE_EVEN_STORAGE=DFTI_COMPLEX_REAL` and `DFTI_PACKED_FORMAT=DFTI_PACK` -from Intel(R) Math Kernel Library, reimplemented `mkl_fft.rfft` and `mkl_fft.irfft` to use real-to-complex functionality with subsequent -copying to rearange the transform as expected of `mkl_fft.rfft`, with the associated performance penalty. The use of the real-to-complex +Due to removal of deprecated real-to-real FFT with :code:`DFTI_CONJUGATE_EVEN_STORAGE=DFTI_COMPLEX_REAL` and +:code:`DFTI_PACKED_FORMAT=DFTI_PACK` from Intel(R) Math Kernel Library, reimplemented :code:`mkl_fft.rfft` and +:code:`mkl_fft.irfft` to use real-to-complex functionality with subsequent copying to rearange the transform as expected +of :code:`mkl_fft.rfft`, with the associated performance penalty. The use of the real-to-complex transform improves multi-core utilization which may offset the performance loss incurred due to copying. 1.1.0 ===== -Added `scipy.fft` backend, see #42. Fixed #46. +Added :code:`scipy.fft` backend, see #42. Fixed #46. ``` Python 3.7.5 (default, Nov 23 2019, 04:02:01) @@ -102,16 +104,16 @@ Out[4]: True 1.0.15 ====== -Changed tests to not compare against numpy fft, as this broke due to renaming of `np.fft.pocketfft` to -`np.fft._pocketfft`. Instead compare against naive realization of 1D FFT as a sum. +Changed tests to not compare against numpy fft, as this broke due to renaming of :code:`np.fft.pocketfft` to +:code:`np.fft._pocketfft`. Instead compare against naive realization of 1D FFT as a sum. -Setup script is now aware of `MKLROOT` environment variable. If unset, NumPy's mkl_info will be queried. +Setup script is now aware of :code:`MKLROOT` environment variable. If unset, NumPy's mkl_info will be queried. 1.0.14 ====== -Fixed unreferenced bug in `irfftn_numpy`, and adjusted NumPy interfaces to change to pocketfft in NumPy 1.17 +Fixed unreferenced bug in :code:`irfftn_numpy`, and adjusted NumPy interfaces to change to pocketfft in NumPy 1.17 1.0.13 @@ -124,8 +126,8 @@ Issue #39 fixed (memory leak with complex FFT on real arrays) ====== Issue #37 fixed. -Inhibited vectorization of short loops computing pointer to memory referenced by a multi-iterator by Intel (R) C Compiler, improving -performance of ND `fft` and `ifft` on real input arrays. +Inhibited vectorization of short loops computing pointer to memory referenced by a multi-iterator by Intel (R) C Compiler, +improving performance of ND :code:`fft` and :code:`ifft` on real input arrays. 1.0.11 @@ -155,7 +157,7 @@ Fixed issues #21, and addressed NumPy 1.15 deprecation warnings from using lists ===== Fixed issues #7, #17, #18. -Consolidated version specification into a single file `mkl_fft/_version.py`. +Consolidated version specification into a single file :code:`mkl_fft/_version.py`. 1.0.4 ===== @@ -169,13 +171,15 @@ This is a bug fix release. It fixes issues #9, and #13. -As part of fixing issue #13, out-of-place 1D FFT calls such as `fft`, `ifft`, `rfft_numpy` and `irfftn_numpy` will allocate Fortran layout array for the output is the input is a Fotran array. +As part of fixing issue #13, out-of-place 1D FFT calls such as :code:`fft`, :code:`ifft`, :code:`rfft_numpy` +and :code:`irfftn_numpy` will allocate Fortran layout array for the output is the input is a Fotran array. 1.0.2 ===== -Minor update of `mkl_fft`, reflecting renaming of `numpy.core.multiarray_tests` module to `numpy.core._multiarray_tests` as well as fixing #4. +Minor update of :code:`mkl_fft`, reflecting renaming of :code:`numpy.core.multiarray_tests` module to +:code:`numpy.core._multiarray_tests` as well as fixing #4. 1.0.1 @@ -186,4 +190,4 @@ Bug fix release. 1.0.0 ===== -Initial release of `mkl_fft`. +Initial release of :code:`mkl_fft`. From 79575bb9563f45599cfd92d73319eabe942330dd Mon Sep 17 00:00:00 2001 From: Vahid Tavanashad Date: Fri, 4 Apr 2025 07:52:00 -0700 Subject: [PATCH 05/11] rename functions --- mkl_fft/_fft_utils.py | 8 ++++---- mkl_fft/_numpy_fft.py | 24 ++++++++++++------------ mkl_fft/_pydfti.pyx | 4 ++-- mkl_fft/_scipy_fft.py | 18 +++++++++--------- 4 files changed, 27 insertions(+), 27 deletions(-) diff --git a/mkl_fft/_fft_utils.py b/mkl_fft/_fft_utils.py index 58c152bf..9cfcd799 100644 --- a/mkl_fft/_fft_utils.py +++ b/mkl_fft/_fft_utils.py @@ -26,10 +26,10 @@ import numpy as np -__all__ = ["_check_norm", "_compute_fwd_scale"] +__all__ = ["check_norm", "compute_fwd_scale"] -def _check_norm(norm): +def check_norm(norm): if norm not in (None, "ortho", "forward", "backward"): raise ValueError( f"Invalid norm value {norm} should be None, 'ortho', 'forward', " @@ -37,8 +37,8 @@ def _check_norm(norm): ) -def _compute_fwd_scale(norm, n, shape): - _check_norm(norm) +def compute_fwd_scale(norm, n, shape): + check_norm(norm) if norm in (None, "backward"): return 1.0 diff --git a/mkl_fft/_numpy_fft.py b/mkl_fft/_numpy_fft.py index f0e56bc9..64fe832a 100644 --- a/mkl_fft/_numpy_fft.py +++ b/mkl_fft/_numpy_fft.py @@ -76,7 +76,7 @@ import numpy as np from . import _pydfti as mkl_fft # pylint: disable=no-name-in-module -from ._fft_utils import _check_norm, _compute_fwd_scale +from ._fft_utils import check_norm, compute_fwd_scale from ._float_utils import __downcast_float128_array @@ -125,7 +125,7 @@ def _cook_nd_args(a, s=None, axes=None, invreal=False): def _swap_direction(norm): - _check_norm(norm) + check_norm(norm) _swap_direction_map = { "backward": "forward", None: "forward", @@ -241,7 +241,7 @@ def fft(a, n=None, axis=-1, norm=None): """ x = __downcast_float128_array(a) - fsc = _compute_fwd_scale(norm, n, x.shape[axis]) + fsc = compute_fwd_scale(norm, n, x.shape[axis]) return trycall(mkl_fft.fft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc}) @@ -334,7 +334,7 @@ def ifft(a, n=None, axis=-1, norm=None): """ x = __downcast_float128_array(a) - fsc = _compute_fwd_scale(norm, n, x.shape[axis]) + fsc = compute_fwd_scale(norm, n, x.shape[axis]) return trycall(mkl_fft.ifft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc}) @@ -425,7 +425,7 @@ def rfft(a, n=None, axis=-1, norm=None): """ x = __downcast_float128_array(a) - fsc = _compute_fwd_scale(norm, n, x.shape[axis]) + fsc = compute_fwd_scale(norm, n, x.shape[axis]) return trycall(mkl_fft.rfft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc}) @@ -518,7 +518,7 @@ def irfft(a, n=None, axis=-1, norm=None): """ x = __downcast_float128_array(a) - fsc = _compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1)) + fsc = compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1)) return trycall( mkl_fft.irfft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc} @@ -606,7 +606,7 @@ def hfft(a, n=None, axis=-1, norm=None): 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)) + fsc = compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1)) return trycall( mkl_fft.irfft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc} @@ -675,7 +675,7 @@ def ihfft(a, n=None, axis=-1, norm=None): norm = _swap_direction(norm) x = __downcast_float128_array(a) x = np.array(x, copy=True, dtype=float) - fsc = _compute_fwd_scale(norm, n, x.shape[axis]) + fsc = compute_fwd_scale(norm, n, x.shape[axis]) output = trycall( mkl_fft.rfft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc} @@ -787,7 +787,7 @@ def fftn(a, s=None, axes=None, norm=None): 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}) @@ -894,7 +894,7 @@ def ifftn(a, s=None, axes=None, norm=None): 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} @@ -1182,7 +1182,7 @@ def rfftn(a, s=None, axes=None, norm=None): 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} @@ -1326,7 +1326,7 @@ def irfftn(a, s=None, axes=None, norm=None): 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 3eae75c1..f2140f52 100644 --- a/mkl_fft/_scipy_fft.py +++ b/mkl_fft/_scipy_fft.py @@ -33,7 +33,7 @@ import numpy as np from . import _pydfti as mkl_fft # pylint: disable=no-name-in-module -from ._fft_utils import _compute_fwd_scale +from ._fft_utils import compute_fwd_scale from ._float_utils import __supported_array_or_not_implemented __doc__ = """ @@ -235,7 +235,7 @@ def fft( ): _check_plan(plan) x = _validate_input(a) - fsc = _compute_fwd_scale(norm, n, x.shape[axis]) + fsc = compute_fwd_scale(norm, n, x.shape[axis]) with Workers(workers): return mkl_fft.fft( @@ -248,7 +248,7 @@ def ifft( ): _check_plan(plan) x = _validate_input(a) - fsc = _compute_fwd_scale(norm, n, x.shape[axis]) + fsc = compute_fwd_scale(norm, n, x.shape[axis]) with Workers(workers): return mkl_fft.ifft( @@ -303,7 +303,7 @@ def fftn( ): _check_plan(plan) x = _validate_input(a) - fsc = _compute_fwd_scale(norm, s, x.shape) + fsc = compute_fwd_scale(norm, s, x.shape) with Workers(workers): return mkl_fft.fftn( @@ -316,7 +316,7 @@ def ifftn( ): _check_plan(plan) x = _validate_input(a) - fsc = _compute_fwd_scale(norm, s, x.shape) + fsc = compute_fwd_scale(norm, s, x.shape) with Workers(workers): return mkl_fft.ifftn( @@ -327,7 +327,7 @@ def ifftn( def rfft(a, n=None, axis=-1, norm=None, workers=None, plan=None): _check_plan(plan) x = _validate_input(a) - fsc = _compute_fwd_scale(norm, n, x.shape[axis]) + fsc = compute_fwd_scale(norm, n, x.shape[axis]) with Workers(workers): return mkl_fft.rfft(x, n=n, axis=axis, fwd_scale=fsc) @@ -336,7 +336,7 @@ def rfft(a, n=None, axis=-1, norm=None, workers=None, plan=None): def irfft(a, n=None, axis=-1, norm=None, workers=None, plan=None): _check_plan(plan) x = _validate_input(a) - fsc = _compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1)) + fsc = compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1)) with Workers(workers): return mkl_fft.irfft(x, n=n, axis=axis, fwd_scale=fsc) @@ -356,7 +356,7 @@ def rfftn(a, s=None, axes=None, norm=None, workers=None, plan=None): _check_plan(plan) x = _validate_input(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) with Workers(workers): return mkl_fft.rfftn(x, s, axes, fwd_scale=fsc) @@ -366,7 +366,7 @@ def irfftn(a, s=None, axes=None, norm=None, workers=None, plan=None): _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) + fsc = compute_fwd_scale(norm, s, x.shape) with Workers(workers): return mkl_fft.irfftn(x, s, axes, fwd_scale=fsc) From 8b63b0e39625c6dfd61ccfd7d0dff4644cf1c6ef Mon Sep 17 00:00:00 2001 From: Vahid Tavanashad Date: Fri, 4 Apr 2025 11:33:54 -0700 Subject: [PATCH 06/11] add multithreading tests --- .gitignore | 1 + CHANGES.rst | 41 +++---- .../test_basic.py} | 4 +- .../tests/from_scipy/test_multithreading.py | 104 ++++++++++++++++++ 4 files changed, 128 insertions(+), 22 deletions(-) rename mkl_fft/tests/{test_from_scipy.py => from_scipy/test_basic.py} (99%) create mode 100644 mkl_fft/tests/from_scipy/test_multithreading.py diff --git a/.gitignore b/.gitignore index 630ac69d..8638c23c 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,4 @@ mkl_fft/_pydfti.cpython-310-x86_64-linux-gnu.so mkl_fft/interfaces/__pycache__/ mkl_fft/src/mklfft.c mkl_fft/tests/__pycache__/ +mkl_fft/tests/from_scipy/__pycache__/ diff --git a/CHANGES.rst b/CHANGES.rst index 103a6191..48a77a58 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -79,26 +79,27 @@ transform improves multi-core utilization which may offset the performance loss Added :code:`scipy.fft` backend, see #42. Fixed #46. -``` -Python 3.7.5 (default, Nov 23 2019, 04:02:01) -Type 'copyright', 'credits' or 'license' for more information -IPython 7.11.1 -- An enhanced Interactive Python. Type '?' for help. - -In [1]: import numpy as np, mkl_fft, mkl_fft._scipy_fft_backend as mkl_be, scipy, scipy.fft, mkl - -In [2]: mkl.verbose(1) -Out[2]: True - -In [3]: x = np.random.randn(8*7).reshape((7, 8)) -...: with scipy.fft.set_backend(mkl_be, only=True): -...: ff = scipy.fft.fft2(x, workers=4) -...: ff2 = scipy.fft.fft2(x) -MKL_VERBOSE Intel(R) MKL 2020.0 Product build 20191102 for Intel(R) 64 architecture Intel(R) Advanced Vector Extensions 2 (Intel(R) AVX2) enabled processors, Lnx 2.40GHz intel_thread -MKL_VERBOSE FFT(drfo7:8:8x8:1:1,bScale:0.0178571,tLim:1,desc:0x5629ad31b800) 24.85ms CNR:OFF Dyn:1 FastMM:1 TID:0 NThr:16,FFT:4 - -In [4]: np.allclose(ff, ff2) -Out[4]: True -``` + +.. code-block:: python + + Python 3.7.5 (default, Nov 23 2019, 04:02:01) + Type 'copyright', 'credits' or 'license' for more information + IPython 7.11.1 -- An enhanced Interactive Python. Type '?' for help. + + In [1]: import numpy as np, mkl_fft, mkl_fft._scipy_fft as mkl_be, scipy, scipy.fft, mkl + + In [2]: mkl.verbose(1) + Out[2]: True + + In [3]: x = np.random.randn(8*7).reshape((7, 8)) + ...: with scipy.fft.set_backend(mkl_be, only=True): + ...: ff = scipy.fft.fft2(x, workers=4) + ...: ff2 = scipy.fft.fft2(x) + MKL_VERBOSE Intel(R) MKL 2020.0 Product build 20191102 for Intel(R) 64 architecture Intel(R) Advanced Vector Extensions 2 (Intel(R) AVX2) enabled processors, Lnx 2.40GHz intel_thread + MKL_VERBOSE FFT(drfo7:8:8x8:1:1,bScale:0.0178571,tLim:1,desc:0x5629ad31b800) 24.85ms CNR:OFF Dyn:1 FastMM:1 TID:0 NThr:16,FFT:4 + + In [4]: np.allclose(ff, ff2) + Out[4]: True 1.0.15 diff --git a/mkl_fft/tests/test_from_scipy.py b/mkl_fft/tests/from_scipy/test_basic.py similarity index 99% rename from mkl_fft/tests/test_from_scipy.py rename to mkl_fft/tests/from_scipy/test_basic.py index 74e969d2..9e720975 100644 --- a/mkl_fft/tests/test_from_scipy.py +++ b/mkl_fft/tests/from_scipy/test_basic.py @@ -1,7 +1,7 @@ # 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 +# TODO: remove when hfft* functions are added # pylint: disable=no-member import multiprocessing @@ -426,7 +426,7 @@ def test_dtypes_real(self, dtype, xp): 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) + 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"]) 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 From 67e068fe8b092a0e137ac76ed46cf9f65080f45e Mon Sep 17 00:00:00 2001 From: Vahid Tavanashad Date: Sat, 5 Apr 2025 18:03:36 -0700 Subject: [PATCH 07/11] update scipy.fftpack file --- mkl_fft/_numpy_fft.py | 3 ++- mkl_fft/_scipy_fft.py | 3 ++- mkl_fft/_scipy_fftpack.py | 29 +++++++++++++++-------------- 3 files changed, 19 insertions(+), 16 deletions(-) diff --git a/mkl_fft/_numpy_fft.py b/mkl_fft/_numpy_fft.py index 64fe832a..abec47d8 100644 --- a/mkl_fft/_numpy_fft.py +++ b/mkl_fft/_numpy_fft.py @@ -75,7 +75,8 @@ import numpy as np -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 diff --git a/mkl_fft/_scipy_fft.py b/mkl_fft/_scipy_fft.py index f2140f52..c194ead0 100644 --- a/mkl_fft/_scipy_fft.py +++ b/mkl_fft/_scipy_fft.py @@ -32,7 +32,8 @@ import mkl import numpy as np -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 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) From c5a636c18706f314c07c8b5e47c788e6b7d04c5d Mon Sep 17 00:00:00 2001 From: Vahid Tavanashad Date: Mon, 7 Apr 2025 12:20:29 -0700 Subject: [PATCH 08/11] address comments --- .github/workflows/build_pip.yaml | 2 +- pyproject.toml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/build_pip.yaml b/.github/workflows/build_pip.yaml index 05e41488..acc3a95c 100644 --- a/.github/workflows/build_pip.yaml +++ b/.github/workflows/build_pip.yaml @@ -54,8 +54,8 @@ jobs: run: | pip install --no-cache-dir cython pytest hypothesis pip install --no-cache-dir numpy ${{ matrix.use_pre }} - pip install --no-cache-dir scipy echo "CONDA_PREFFIX is '${CONDA_PREFIX}'" export MKLROOT=${CONDA_PREFIX} pip install -e . --no-build-isolation --verbose --no-deps + pip install --no-cache-dir scipy python -m pytest -v mkl_fft/tests diff --git a/pyproject.toml b/pyproject.toml index 087ddd04..339b2fd6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,7 +60,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" @@ -91,4 +91,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"] From 3fc8fecee75a21b616c5c55ba768ca0e60d4e80c Mon Sep 17 00:00:00 2001 From: Vahid Tavanashad Date: Tue, 8 Apr 2025 12:00:47 -0700 Subject: [PATCH 09/11] add mkl-devel to required packages and use optional depndencies in github work flows --- .github/workflows/build_pip.yaml | 5 ++--- CHANGES.rst | 6 ++++-- pyproject.toml | 2 +- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/.github/workflows/build_pip.yaml b/.github/workflows/build_pip.yaml index acc3a95c..594a5ec1 100644 --- a/.github/workflows/build_pip.yaml +++ b/.github/workflows/build_pip.yaml @@ -52,10 +52,9 @@ jobs: - name: Build conda package run: | - pip install --no-cache-dir cython pytest hypothesis + pip install --no-cache-dir cython pip install --no-cache-dir numpy ${{ matrix.use_pre }} echo "CONDA_PREFFIX is '${CONDA_PREFIX}'" export MKLROOT=${CONDA_PREFIX} - pip install -e . --no-build-isolation --verbose --no-deps - pip install --no-cache-dir scipy + pip install -e .[test] --no-build-isolation --verbose python -m pytest -v mkl_fft/tests diff --git a/CHANGES.rst b/CHANGES.rst index 48a77a58..edd440c3 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -2,8 +2,10 @@ mkl_fft changelog ================= -1.3.13 -====== +1.3.13 (03/25/2025) +=================== + +Supported python versions are 3.9, 3.10, 3.11, 3.12 migrate from :code:`setup.py` to :code:`pyproject.toml` diff --git a/pyproject.toml b/pyproject.toml index 339b2fd6..ee20e854 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ [build-system] build-backend = "setuptools.build_meta" -requires = ["setuptools>=64", "Cython", "numpy"] +requires = ["setuptools>=64", "Cython", "numpy", "mkl-devel"] [project] authors = [ From 4d329f2e3bd9c1413e6ebf0bb65a7207e60681dd Mon Sep 17 00:00:00 2001 From: Vahid Tavanashad Date: Wed, 9 Apr 2025 17:27:24 -0700 Subject: [PATCH 10/11] skip a test --- mkl_fft/tests/from_scipy/test_basic.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mkl_fft/tests/from_scipy/test_basic.py b/mkl_fft/tests/from_scipy/test_basic.py index 9e720975..182bb6d4 100644 --- a/mkl_fft/tests/from_scipy/test_basic.py +++ b/mkl_fft/tests/from_scipy/test_basic.py @@ -574,6 +574,7 @@ def test_ihfft(self, xp): self._test_mtsame(fft.ihfft, a, xp=xp) +@pytest.mark.skip("skip due to gh-issue-151") @skip_xp_backends(np_only=True) @pytest.mark.parametrize("func", [fft.fft, fft.ifft, fft.rfft, fft.irfft]) def test_multiprocess(func, xp): From 3f52adfcde2cb8af7f373b49468fcce1656cdbf4 Mon Sep 17 00:00:00 2001 From: Vahid Tavanashad Date: Thu, 10 Apr 2025 11:35:30 -0700 Subject: [PATCH 11/11] include a test --- .gitignore | 2 +- mkl_fft/__init__.py | 1 + mkl_fft/tests/from_scipy/test_basic.py | 1 - 3 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 8638c23c..921f8885 100644 --- a/.gitignore +++ b/.gitignore @@ -3,7 +3,7 @@ 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__/ 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/tests/from_scipy/test_basic.py b/mkl_fft/tests/from_scipy/test_basic.py index 182bb6d4..9e720975 100644 --- a/mkl_fft/tests/from_scipy/test_basic.py +++ b/mkl_fft/tests/from_scipy/test_basic.py @@ -574,7 +574,6 @@ def test_ihfft(self, xp): self._test_mtsame(fft.ihfft, a, xp=xp) -@pytest.mark.skip("skip due to gh-issue-151") @skip_xp_backends(np_only=True) @pytest.mark.parametrize("func", [fft.fft, fft.ifft, fft.rfft, fft.irfft]) def test_multiprocess(func, xp):