diff --git a/CHANGELOG.md b/CHANGELOG.md index 9253ca45..5fa176d2 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"` and `norm="ortho"` scaling in `mkl_fft.fftn`, `ifftn`, `rfftn`, `irfftn` and the `fft2`/`ifft2`/`rfft2`/`irfft2` family when only a subset of the input axes is transformed and `s` is not given. The scale factor was computed over the full array shape instead of over the transformed axes, over-normalizing the result by the product of the untransformed axis lengths (for example `fft2` on a 3-D array, or `fftn(x, axes=(0,))`). The `mkl_fft.interfaces.numpy_fft` and `mkl_fft.interfaces.scipy_fft` wrappers were unaffected, as they resolve `s` before delegating +* Fixed `norm="forward"` and `norm="ortho"` scaling in `mkl_fft.irfftn` and `irfft2`, which normalized over the input length `n` along the last transformed axis rather than the complex-to-real output length `2 * (n - 1)`. This applied even when every axis was transformed * 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..4d3e5eb1 100644 --- a/mkl_fft/_fft_utils.py +++ b/mkl_fft/_fft_utils.py @@ -78,6 +78,43 @@ def _compute_fwd_scale(norm, n, shape): return np.sqrt(fsc) +def _compute_nd_scale_shape(x, s, axes, norm=None, invreal=False): + """ + Resolve the lengths that a norm-scaled N-D transform normalizes over. + + ``_compute_fwd_scale`` falls back to the full array shape when ``s`` is + None. That over-normalizes when only a subset of axes is transformed, and + for c2r transforms the basis is the *output* length along the last + transformed axis rather than the input one. + + This mirrors what the ``numpy_fft`` and ``scipy_fft`` interfaces already + do by calling ``_cook_nd_args`` before delegating. Only the scale basis is + resolved here; ``s`` itself is left alone so that dispatch in + ``_c2c_fftnd_impl`` is unchanged. + + ``norm`` is accepted only to skip the work for the unscaled norms, whose + scale is 1.0 regardless of shape. Invalid values fall through to + ``_compute_fwd_scale``, which validates them. + """ + + 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 + 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 or empty axes; leave the scale alone and let the + # transform itself raise + 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..bf7bd49b --- /dev/null +++ b/mkl_fft/tests/test_dispatch_equivalence.py @@ -0,0 +1,230 @@ +"""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", ["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)