diff --git a/CHANGELOG.md b/CHANGELOG.md index 9253ca45..cd65774f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed ### Fixed +* Fixed `norm="forward"`/`"ortho"` scaling in `fftn`, `ifftn`, `rfftn`, `irfftn` and the `fft2` family when only a subset of axes is transformed: the scale used the full array shape instead of the transformed axes. The `numpy_fft` and `scipy_fft` interfaces were unaffected [gh-370](https://github.com/IntelPython/mkl_fft/pull/370) +* Fixed `norm="forward"`/`"ortho"` scaling in `irfftn` and `irfft2`, which normalized over the input length `n` rather than the complex-to-real output length `2 * (n - 1)` [gh-370](https://github.com/IntelPython/mkl_fft/pull/370) * Declared `f_ndim` as a C `int` in `_allocate_result` so the buffer size is computed in C rather than through a Python object, resolving a Coverity out-of-bounds (OVERRUN) false positive [gh-364](https://github.com/IntelPython/mkl_fft/pull/364) * Silenced a Coverity `UNUSED_VALUE` finding in `__create_descriptor_1d` by marking the `DftiFreeDescriptor` status (used only by a debug-only `assert`) as intentionally unused [gh-365](https://github.com/IntelPython/mkl_fft/pull/365) diff --git a/mkl_fft/_fft_utils.py b/mkl_fft/_fft_utils.py index b3f0d0c9..24433064 100644 --- a/mkl_fft/_fft_utils.py +++ b/mkl_fft/_fft_utils.py @@ -78,6 +78,37 @@ def _compute_fwd_scale(norm, n, shape): return np.sqrt(fsc) +def _compute_nd_scale_shape(x, s, axes, norm=None, invreal=False): + """ + Lengths a norm-scaled N-D transform normalizes over. + + ``_compute_fwd_scale`` falls back to the full array shape when ``s`` is + None, which over-normalizes a subset-of-axes transform; for c2r the basis + is the output length ``2 * (n - 1)``. Mirrors what the interfaces already + do via ``_cook_nd_args``, but leaves ``s`` alone so dispatch is unchanged. + ``norm`` is only used to skip the work when the scale is 1.0 anyway. + """ + + if s is not None or norm in (None, "backward"): + return s + try: + if axes is None: + ss = list(x.shape) + last = len(ss) - 1 + elif len(axes) == 0: + # identity transform; np.prod(()) == 1 gives scale 1.0 + return () + else: + ss = [x.shape[ai] for ai in axes] + last = axes[-1] + if invreal: + ss[-1] = 2 * (x.shape[last] - 1) + except (IndexError, TypeError): + # invalid axes; let the transform raise, matching NumPy's error + return s + return tuple(ss) + + def _cook_nd_args(a, s=None, axes=None, invreal=False): if s is None: shapeless = True diff --git a/mkl_fft/_mkl_fft.py b/mkl_fft/_mkl_fft.py index 3ab60c9a..decc3f72 100644 --- a/mkl_fft/_mkl_fft.py +++ b/mkl_fft/_mkl_fft.py @@ -27,6 +27,7 @@ _c2c_fftnd_impl, _c2r_fftnd_impl, _compute_fwd_scale, + _compute_nd_scale_shape, _r2c_fftnd_impl, ) @@ -68,12 +69,14 @@ def ifft2(x, s=None, axes=(-2, -1), norm=None, out=None): def fftn(x, s=None, axes=None, norm=None, out=None): - fsc = _compute_fwd_scale(norm, s, x.shape) + ss = _compute_nd_scale_shape(x, s, axes, norm) + fsc = _compute_fwd_scale(norm, ss, x.shape) return _c2c_fftnd_impl(x, s=s, axes=axes, out=out, direction=+1, fsc=fsc) def ifftn(x, s=None, axes=None, norm=None, out=None): - fsc = _compute_fwd_scale(norm, s, x.shape) + ss = _compute_nd_scale_shape(x, s, axes, norm) + fsc = _compute_fwd_scale(norm, ss, x.shape) return _c2c_fftnd_impl(x, s=s, axes=axes, out=out, direction=-1, fsc=fsc) @@ -96,10 +99,12 @@ def irfft2(x, s=None, axes=(-2, -1), norm=None, out=None): def rfftn(x, s=None, axes=None, norm=None, out=None): - fsc = _compute_fwd_scale(norm, s, x.shape) + ss = _compute_nd_scale_shape(x, s, axes, norm) + fsc = _compute_fwd_scale(norm, ss, x.shape) return _r2c_fftnd_impl(x, s=s, axes=axes, out=out, fsc=fsc) def irfftn(x, s=None, axes=None, norm=None, out=None): - fsc = _compute_fwd_scale(norm, s, x.shape) + ss = _compute_nd_scale_shape(x, s, axes, norm, invreal=True) + fsc = _compute_fwd_scale(norm, ss, x.shape) return _c2r_fftnd_impl(x, s=s, axes=axes, out=out, fsc=fsc) diff --git a/mkl_fft/tests/test_dispatch_equivalence.py b/mkl_fft/tests/test_dispatch_equivalence.py new file mode 100644 index 00000000..12274413 --- /dev/null +++ b/mkl_fft/tests/test_dispatch_equivalence.py @@ -0,0 +1,383 @@ +# Copyright (c) 2026, 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. + +"""Cross-library equivalence checks for axis and axes dispatch. + +``third_party/scipy/test_basic.py::test_fft_with_order`` already checks that +mkl_fft agrees with *itself* across C, Fortran, and non-contiguous layouts. It +does not compare against an external reference, so a dispatch change that is +consistently wrong in every layout passes it. + +The defect recorded in ``_fft_utils._iter_complementary`` was exactly that +kind: values correct, but an element placed in the other half of the output +relative to NumPy. These tests therefore use ``numpy.fft`` as the reference. + +Two deliberate choices: + +* Every axis length differs, so an axis permutation cannot produce a + correctly shaped result and hide behind a shape assertion. +* Output dtype is asserted alongside values, so a dispatch change cannot + silently upcast. + +These cover the paths that dispatch on *which* axes are requested: full-axes +transforms reach the batched N-D descriptor, strict subsets iterate the +complementary axes, and 1-D transforms of rank > 2 arrays are batched only for +the first and last axis. +""" + +import itertools + +import numpy as np +import pytest +from numpy.testing import assert_allclose + +import mkl_fft + +_SHAPE_3D = (8, 7, 13) +_SHAPE_4D = (4, 5, 6, 7) + +_DTYPES = ["float32", "float64", "complex64", "complex128"] +_REAL_DTYPES = ["float32", "float64"] + +_ORDERS = ["C", "F", "non-contiguous"] + +# Relative tolerance by input precision. Single-precision transforms of +# random data over these lengths stay well inside 2e-5. +_TOL = { + "float32": 2e-5, + "complex64": 2e-5, + "float64": 1e-12, + "complex128": 1e-12, +} + +# every non-empty subset of the axes of a 3-D array, plus None +_AXES_3D = [ + ax for n in (1, 2, 3) for ax in itertools.combinations(range(3), n) +] + [None] + + +def _make(shape, dtype, seed=42): + rng = np.random.default_rng(seed) + dt = np.dtype(dtype) + if dt.kind == "c": + x = rng.standard_normal(shape) + 1j * rng.standard_normal(shape) + else: + x = rng.standard_normal(shape) + return x.astype(dt) + + +def _relayout(x, order): + """Return *x* laid out as requested; data content may differ by order.""" + if order == "F": + return np.asfortranarray(x) + if order == "non-contiguous": + return x[::-1] + return np.ascontiguousarray(x) + + +def _check(got, want, dtype): + assert got.dtype == want.dtype, f"dtype {got.dtype} != {want.dtype}" + assert got.shape == want.shape, f"shape {got.shape} != {want.shape}" + tol = _TOL[dtype] + assert_allclose( + got, want, rtol=tol, atol=tol * max(1.0, float(np.abs(want).max())) + ) + + +# --------------------------------------------------------------------------- +# N-D complex transforms over a subset of axes +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("func", ["fftn", "ifftn"]) +@pytest.mark.parametrize("dtype", _DTYPES) +@pytest.mark.parametrize("axes", _AXES_3D) +@pytest.mark.parametrize("order", _ORDERS) +def test_fftn_axes_subset(func, dtype, axes, order): + x = _relayout(_make(_SHAPE_3D, dtype), order) + got = getattr(mkl_fft, func)(x, axes=axes) + want = getattr(np.fft, func)(x, axes=axes) + _check(got, want, dtype) + + +@pytest.mark.parametrize("func", ["rfftn", "irfftn"]) +@pytest.mark.parametrize("dtype", _DTYPES) +@pytest.mark.parametrize("axes", _AXES_3D) +@pytest.mark.parametrize("order", _ORDERS) +def test_rfftn_axes_subset(func, dtype, axes, order): + if func == "rfftn" and dtype not in _REAL_DTYPES: + pytest.skip("rfftn takes real input") + x = _relayout(_make(_SHAPE_3D, dtype), order) + got = getattr(mkl_fft, func)(x, axes=axes) + want = getattr(np.fft, func)(x, axes=axes) + _check(got, want, dtype) + + +# --------------------------------------------------------------------------- +# 1-D transforms along each axis of a higher-rank array +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("func", ["fft", "ifft"]) +@pytest.mark.parametrize("dtype", _DTYPES) +@pytest.mark.parametrize("axis", range(len(_SHAPE_3D))) +@pytest.mark.parametrize("order", _ORDERS) +def test_fft_axis_3d(func, dtype, axis, order): + x = _relayout(_make(_SHAPE_3D, dtype), order) + got = getattr(mkl_fft, func)(x, axis=axis) + want = getattr(np.fft, func)(x, axis=axis) + _check(got, want, dtype) + + +@pytest.mark.parametrize("func", ["fft", "ifft", "rfft"]) +@pytest.mark.parametrize("dtype", ["float64", "complex128"]) +@pytest.mark.parametrize("axis", range(len(_SHAPE_4D))) +@pytest.mark.parametrize("order", _ORDERS) +def test_fft_axis_4d(func, dtype, axis, order): + """A rank-4 array has two interior axes, so the per-vector fallback in the + C backend is exercised twice within one sweep. + """ + if func == "rfft" and dtype != "float64": + pytest.skip("rfft takes real input") + x = _relayout(_make(_SHAPE_4D, dtype), order) + got = getattr(mkl_fft, func)(x, axis=axis) + want = getattr(np.fft, func)(x, axis=axis) + _check(got, want, dtype) + + +@pytest.mark.parametrize("func", ["rfft", "irfft"]) +@pytest.mark.parametrize("dtype", _DTYPES) +@pytest.mark.parametrize("axis", range(len(_SHAPE_3D))) +@pytest.mark.parametrize("order", _ORDERS) +def test_rfft_axis_3d(func, dtype, axis, order): + if func == "rfft" and dtype not in _REAL_DTYPES: + pytest.skip("rfft takes real input") + x = _relayout(_make(_SHAPE_3D, dtype), order) + got = getattr(mkl_fft, func)(x, axis=axis) + want = getattr(np.fft, func)(x, axis=axis) + _check(got, want, dtype) + + +# --------------------------------------------------------------------------- +# norm interacts with the scale factor applied at dispatch time +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("func", ["fftn", "ifftn"]) +@pytest.mark.parametrize("dtype", ["float64", "complex128"]) +@pytest.mark.parametrize("axes", [(0,), (1,), (2,), (1, 2), None]) +@pytest.mark.parametrize("norm", [None, "backward", "forward", "ortho"]) +def test_fftn_axes_subset_norm(func, dtype, axes, norm): + x = _make(_SHAPE_3D, dtype) + got = getattr(mkl_fft, func)(x, axes=axes, norm=norm) + want = getattr(np.fft, func)(x, axes=axes, norm=norm) + _check(got, want, dtype) + + +@pytest.mark.parametrize("func", ["rfftn", "irfftn"]) +@pytest.mark.parametrize("dtype", ["float64", "complex128"]) +@pytest.mark.parametrize("axes", [(0,), (1,), (2,), (1, 2), None]) +@pytest.mark.parametrize("norm", [None, "backward", "forward", "ortho"]) +def test_rfftn_axes_subset_norm(func, dtype, axes, norm): + """Includes ``axes=None``: for c2r the scale basis is the *output* length + along the last transformed axis, so a full-axes irfftn is normalized over + ``2 * (n - 1)`` rather than ``n``. + """ + if func == "rfftn" and dtype != "float64": + pytest.skip("rfftn takes real input") + x = _make(_SHAPE_3D, dtype) + got = getattr(mkl_fft, func)(x, axes=axes, norm=norm) + want = getattr(np.fft, func)(x, axes=axes, norm=norm) + _check(got, want, dtype) + + +@pytest.mark.parametrize("func", ["fftn", "ifftn", "fft2", "ifft2"]) +@pytest.mark.parametrize("dtype", ["float64", "complex128"]) +@pytest.mark.parametrize("norm", [None, "backward", "forward", "ortho"]) +def test_empty_axes_is_identity(func, dtype, norm): + """``axes=()`` transforms nothing, so the scale is 1.0 under every norm and + the input comes back untouched, as it does from NumPy. + """ + x = _make(_SHAPE_3D, dtype) + got = getattr(mkl_fft, func)(x, axes=(), norm=norm) + want = getattr(np.fft, func)(x, axes=(), norm=norm) + _check(got, want, dtype) + assert got is x, "no axes transformed, so the input should be returned" + + +@pytest.mark.parametrize("func", ["rfftn", "irfftn"]) +@pytest.mark.parametrize("norm", [None, "forward", "ortho"]) +def test_empty_axes_r2c_raises_like_numpy(func, norm): + """With no axes there is no last transformed axis to hold the half + spectrum, so both libraries raise; check the type agrees. + """ + x = _make(_SHAPE_3D, "float64") + with pytest.raises(IndexError): + getattr(np.fft, func)(x, axes=(), norm=norm) + with pytest.raises(IndexError): + getattr(mkl_fft, func)(x, axes=(), norm=norm) + + +@pytest.mark.parametrize("func", ["fft2", "ifft2", "rfft2", "irfft2"]) +@pytest.mark.parametrize("dtype", ["float64", "complex128"]) +@pytest.mark.parametrize("norm", [None, "backward", "forward", "ortho"]) +def test_fft2_on_3d_norm(func, dtype, norm): + """``fft2`` on a rank-3 array transforms 2 of 3 axes, so it is a subset + transform even though the caller passed no ``axes``. + """ + if func == "rfft2" and dtype != "float64": + pytest.skip("rfft2 takes real input") + x = _make(_SHAPE_3D, dtype) + got = getattr(mkl_fft, func)(x, norm=norm) + want = getattr(np.fft, func)(x, norm=norm) + _check(got, want, dtype) + + +@pytest.mark.parametrize("func", ["fft", "ifft"]) +@pytest.mark.parametrize("dtype", ["float64", "complex128"]) +@pytest.mark.parametrize("axis", range(len(_SHAPE_3D))) +@pytest.mark.parametrize("norm", [None, "backward", "forward", "ortho"]) +def test_fft_axis_norm(func, dtype, axis, norm): + x = _make(_SHAPE_3D, dtype) + got = getattr(mkl_fft, func)(x, axis=axis, norm=norm) + want = getattr(np.fft, func)(x, axis=axis, norm=norm) + _check(got, want, dtype) + + +# --------------------------------------------------------------------------- +# out= must not change results on any dispatch path +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", ["complex64", "complex128"]) +@pytest.mark.parametrize("axes", _AXES_3D) +def test_fftn_axes_subset_out(dtype, axes): + x = _make(_SHAPE_3D, dtype) + want = np.fft.fftn(x, axes=axes) + out = np.empty(want.shape, dtype=x.dtype) + got = mkl_fft.fftn(x, axes=axes, out=out) + assert got is out, "out= should be returned" + _check(got, want, dtype) + + +def _check_out(func, x, dtype, **kwargs): + """Run *func* with an ``out`` array shaped and typed from the reference.""" + want = getattr(np.fft, func)(x, **kwargs) + out = np.empty(want.shape, dtype=want.dtype) + got = getattr(mkl_fft, func)(x, out=out, **kwargs) + assert got is out, "out= should be returned" + _check(got, want, dtype) + + +@pytest.mark.parametrize("func", ["fftn", "ifftn"]) +@pytest.mark.parametrize("dtype", ["complex64", "complex128"]) +@pytest.mark.parametrize("axes", [(0,), (2,), (1, 2), None]) +@pytest.mark.parametrize("norm", ["forward", "ortho"]) +def test_c2c_out_with_norm(func, dtype, axes, norm): + """The scale is applied while the result is written into ``out``, which is + the path this fix changes; ``ifftn`` shares it but was never exercised. + """ + _check_out(func, _make(_SHAPE_3D, dtype), dtype, axes=axes, norm=norm) + + +@pytest.mark.parametrize("dtype", ["float32", "float64"]) +@pytest.mark.parametrize("axes", [(0,), (2,), (1, 2), None]) +@pytest.mark.parametrize("norm", [None, "forward", "ortho"]) +def test_rfftn_out_with_norm(dtype, axes, norm): + """r2c: ``out`` is complex with the last transformed axis reduced to + ``n // 2 + 1``, a different allocation from the c2c case. + """ + _check_out("rfftn", _make(_SHAPE_3D, dtype), dtype, axes=axes, norm=norm) + + +@pytest.mark.parametrize("dtype", ["complex64", "complex128"]) +@pytest.mark.parametrize("axes", [(2,), (1, 2), None]) +@pytest.mark.parametrize("norm", [None, "forward", "ortho"]) +def test_irfftn_out_with_norm(dtype, axes, norm): + """c2r: ``out`` is real with the last transformed axis expanded to + ``2 * (n - 1)`` -- the length the invreal branch of the scale helper + computes, so this ties the two together. + """ + _check_out("irfftn", _make(_SHAPE_3D, dtype), dtype, axes=axes, norm=norm) + + +# --------------------------------------------------------------------------- +# s= combined with a scaled norm +# +# The scale helper deliberately returns early when s is given, leaving +# _compute_fwd_scale to normalize over prod(s). These lock that branch down. +# axes is always passed explicitly: NumPy deprecated giving s without axes. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("func", ["fftn", "ifftn"]) +@pytest.mark.parametrize( + "axes,s", + [ + ((0,), (16,)), # pad one axis + ((0,), (4,)), # truncate one axis + ((1, 2), (10, 20)), # pad two + ((1, 2), (4, 6)), # truncate two + ((0, 1, 2), (16, 4, 20)), # pad and truncate together + ], +) +@pytest.mark.parametrize("norm", ["forward", "ortho"]) +def test_c2c_shape_arg_with_norm(func, axes, s, norm): + """The scale must come from ``prod(s)`` -- the padded or truncated length -- + not from the original axis lengths. + """ + x = _make(_SHAPE_3D, "complex128") + got = getattr(mkl_fft, func)(x, s=s, axes=axes, norm=norm) + want = getattr(np.fft, func)(x, s=s, axes=axes, norm=norm) + _check(got, want, "complex128") + + +@pytest.mark.parametrize("s", [(8, 7, 20), (8, 7, 10), (8, 7, 24)]) +@pytest.mark.parametrize("norm", ["forward", "ortho"]) +def test_irfftn_shape_arg_with_norm(s, norm): + """With ``s`` given, the invreal doubling must *not* be applied: the scale + normalizes over ``s[-1]``, not ``2 * (x.shape[-1] - 1)``. + + ``s=(8, 7, 24)`` is deliberately ``2 * (13 - 1)``, so a regression that + ignores ``s`` and falls back to the input-derived basis would still pass + that one case -- 20 and 10 are what catch it. A regression that instead + doubles ``s[-1]`` itself is caught by all three. + """ + x = _make(_SHAPE_3D, "complex128") + got = mkl_fft.irfftn(x, s=s, axes=(0, 1, 2), norm=norm) + want = np.fft.irfftn(x, s=s, axes=(0, 1, 2), norm=norm) + _check(got, want, "complex128") + + +@pytest.mark.parametrize( + "axes,s", [((1,), (10,)), ((1,), (4,)), ((1, 2), (10, 20))] +) +@pytest.mark.parametrize("norm", ["forward", "ortho"]) +def test_rfftn_shape_arg_with_norm(axes, s, norm): + """Locks the r2c ``s``-given path together with scaling.""" + x = _make(_SHAPE_3D, "float64") + got = mkl_fft.rfftn(x, s=s, axes=axes, norm=norm) + want = np.fft.rfftn(x, s=s, axes=axes, norm=norm) + _check(got, want, "float64")