make array Hermitian before calling complex-to-real FFT

CHANGELOG.md

@@ -29,6 +29,7 @@ This release achieves 100% compliance with Python Array API specification (revis
 * Removed `einsum_call` keyword from `dpnp.einsum_path` signature [#2421](https://github.com/IntelPython/dpnp/pull/2421)
 * Changed `"max dimensions"` to `None` in array API capabilities [#2432](https://github.com/IntelPython/dpnp/pull/2432)
 * Updated kernel header `i0.hpp` to expose `cyl_bessel_i0` function depending on build target [#2440](https://github.com/IntelPython/dpnp/pull/2440)
+* Updated FFT module to make input array Hermitian before calling complex-to-real FFT [#2444](https://github.com/IntelPython/dpnp/pull/2444)
 
 ### Fixed
 

dpnp/fft/dpnp_iface_fft.py

@@ -640,11 +640,11 @@ def ifft(a, n=None, axis=-1, norm=None, out=None):
     :obj:`dpnp.fft.fft`, i.e.,
 
     * ``a[0]`` should contain the zero frequency term,
-    * ``a[1:n//2]`` should contain the positive-frequency terms,
+    * ``a[1:(n+1)//2]`` should contain the positive-frequency terms,
     * ``a[n//2 + 1:]`` should contain the negative-frequency terms, in
       increasing order starting from the most negative frequency.
 
-    For an even number of input points, ``A[n//2]`` represents the sum of
+    For an even number of input points, ``a[n//2]`` represents the sum of
     the values at the positive and negative Nyquist frequencies, as the two
     are aliased together.
 
@@ -1062,7 +1062,7 @@ def irfft(a, n=None, axis=-1, norm=None, out=None):
 
     This function computes the inverse of the one-dimensional *n*-point
     discrete Fourier Transform of real input computed by :obj:`dpnp.fft.rfft`.
-    In other words, ``irfft(rfft(a), len(a)) == a`` to within numerical
+    In other words, ``irfft(rfft(a), n=len(a)) == a`` to within numerical
     accuracy. (See Notes below for why ``len(a)`` is necessary here.)
 
     The input is expected to be in the form returned by :obj:`dpnp.fft.rfft`,
@@ -1265,9 +1265,9 @@ def irfftn(a, s=None, axes=None, norm=None, out=None):
     This function computes the inverse of the *N*-dimensional discrete Fourier
     Transform for real input over any number of axes in an *M*-dimensional
     array by means of the Fast Fourier Transform (FFT). In other words,
-    ``irfftn(rfftn(a), a.shape) == a`` to within numerical accuracy. (The
-    ``a.shape`` is necessary like ``len(a)`` is for :obj:`dpnp.fft.irfft`,
-    and for the same reason.)
+    ``irfftn(rfftn(a), s=a.shape, axes=range(a.ndim)) == a`` to within
+    numerical accuracy. (The ``a.shape`` is necessary like ``len(a)`` is for
+    :obj:`dpnp.fft.irfft`, and for the same reason.)
 
     The input should be ordered in the same way as is returned by
     :obj:`dpnp.fft.rfftn`, i.e. as for :obj:`dpnp.fft.irfft` for the final

dpnp/fft/dpnp_utils_fft.py

@@ -285,25 +285,30 @@ def _copy_array(x, complex_input):
             dtype = map_dtype_to_device(dpnp.float64, x.sycl_device)
 
     if copy_flag:
-        x_copy = dpnp.empty_like(x, dtype=dtype, order="C")
-
-        exec_q = x.sycl_queue
-        _manager = dpu.SequentialOrderManager[exec_q]
-        dep_evs = _manager.submitted_events
-
-        ht_copy_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
-            src=dpnp.get_usm_ndarray(x),
-            dst=x_copy.get_array(),
-            sycl_queue=exec_q,
-            depends=dep_evs,
-        )
-        _manager.add_event_pair(ht_copy_ev, copy_ev)
-        x = x_copy
+        x = _copy_kernel(x, dtype)
 
     # if copying is done, FFT can be in-place (copy_flag = in_place flag)
     return x, copy_flag
 
 
+def _copy_kernel(x, dtype):
+    x_copy = dpnp.empty_like(x, dtype=dtype, order="C")
+
+    exec_q = x.sycl_queue
+    _manager = dpu.SequentialOrderManager[exec_q]
+    dep_evs = _manager.submitted_events
+
+    ht_copy_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
+        src=dpnp.get_usm_ndarray(x),
+        dst=x_copy.get_array(),
+        sycl_queue=exec_q,
+        depends=dep_evs,
+    )
+    _manager.add_event_pair(ht_copy_ev, copy_ev)
+
+    return x_copy
+
+
 def _extract_axes_chunk(a, s, chunk_size=3):
     """
     Classify the first input into a list of lists with each list containing
@@ -433,6 +438,41 @@ def _fft(a, norm, out, forward, in_place, c2c, axes, batch_fft=True):
     return result
 
 
+def _make_array_hermitian(a, n, copy_needed):
+    """
+    For `dpnp.fft.irfft`, the input array should be Hermitian. If it is not,
+    the behavior is undefined. This function makes necessary changes to make
+    sure the given array is Hermitian.
+
+    It is assumed that this function is called after `_cook_nd_args` and so
+    `n` is always ``None``. It is also assumed that it is called after
+    `_truncate_or_pad`, so the array has enough length.
+    """
+
+    length_is_even = n % 2 == 0
+    hermitian = dpnp.all(a[0].imag == 0)
+    assert n is not None
+    if length_is_even:
+        # Nyquist mode (n//2+1 mode) is n//2-th element
+        f_ny = n // 2
+        assert a.shape[0] > f_ny
+        hermitian = hermitian and dpnp.all(a[f_ny].imag == 0)
+    else:
+        # No Nyquist mode
+        pass
+
+    if not hermitian:
+        if copy_needed:
+            a = _copy_kernel(a, a.dtype)
+
+        a[0].imag = 0
+        if length_is_even:
+            f_ny = n // 2
+            a[f_ny].imag = 0
+
+    return a
+
+
 def _scale_result(res, a_shape, norm, forward, index):
     """Scale the result of the FFT according to `norm`."""
     if res.dtype in [dpnp.float32, dpnp.complex64]:
@@ -559,6 +599,7 @@ def dpnp_fft(a, forward, real, n=None, axis=-1, norm=None, out=None):
     """Calculates 1-D FFT of the input array along axis"""
 
     _check_norm(norm)
+    a_orig = a
     a_ndim = a.ndim
     if a_ndim == 0:
         raise ValueError("Input array must be at least 1D")
@@ -591,6 +632,14 @@ def dpnp_fft(a, forward, real, n=None, axis=-1, norm=None, out=None):
     if a.size == 0:
         return dpnp.get_result_array(a, out=out, casting="same_kind")
 
+    if c2r:
+        # input array should be Hermitian for c2r FFT
+        a = dpnp.moveaxis(a, axis, 0)
+        a = _make_array_hermitian(
+            a, a.shape[0], dpnp.are_same_logical_tensors(a, a_orig)
+        )
+        a = dpnp.moveaxis(a, 0, axis)
+
     return _fft(
         a,
         norm=norm,
@@ -607,6 +656,7 @@ def dpnp_fft(a, forward, real, n=None, axis=-1, norm=None, out=None):
 def dpnp_fftn(a, forward, real, s=None, axes=None, norm=None, out=None):
     """Calculates N-D FFT of the input array along axes"""
 
+    a_orig = a
     if isinstance(axes, Sequence) and len(axes) == 0:
         if real:
             raise IndexError("Empty axes.")
@@ -636,8 +686,14 @@ def dpnp_fftn(a, forward, real, s=None, axes=None, norm=None, out=None):
     len_axes = len(axes)
     if len_axes == 1:
         a = _truncate_or_pad(a, (s[-1],), (axes[-1],))
+        if c2r:
+            a = dpnp.moveaxis(a, axes[-1], 0)
+            a = _make_array_hermitian(
+                a, a.shape[0], dpnp.are_same_logical_tensors(a, a_orig)
+            )
+            a = dpnp.moveaxis(a, 0, axes[-1])
         return _fft(
-            a, norm, out, forward, in_place and c2c, c2c, axes[0], a.ndim != 1
+            a, norm, out, forward, in_place and c2c, c2c, axes[-1], a.ndim != 1
         )
 
     if r2c:
@@ -686,6 +742,12 @@ def dpnp_fftn(a, forward, real, s=None, axes=None, norm=None, out=None):
             batch_fft=a.ndim != len_axes - 1,
         )
         a = _truncate_or_pad(a, (s[-1],), (axes[-1],))
+        if c2r:
+            a = dpnp.moveaxis(a, axes[-1], 0)
+            a = _make_array_hermitian(
+                a, a.shape[0], dpnp.are_same_logical_tensors(a, a_orig)
+            )
+            a = dpnp.moveaxis(a, 0, axes[-1])
         return _fft(
             a, norm, out, forward, in_place and c2c, c2c, axes[-1], a.ndim != 1
         )

dpnp/tests/test_fft.py

@@ -20,32 +20,6 @@
 from .third_party.cupy import testing
 
 
-def _make_array_Hermitian(a, n):
-    """
-    For `dpnp.fft.irfft` and `dpnp.fft.hfft`, the input array should be Hermitian.
-    If it is not, the behavior is undefined. This function makes necessary changes
-    to make sure the given array is Hermitian.
-    """
-
-    a[0] = a[0].real
-    if n is None:
-        # last element is Nyquist mode
-        a[-1] = a[-1].real
-    elif n % 2 == 0:
-        # Nyquist mode (n//2+1 mode) is n//2-th element
-        f_ny = n // 2
-        if a.shape[0] > f_ny:
-            a[f_ny] = a[f_ny].real
-            a[f_ny + 1 :] = 0  # no data needed after Nyquist mode
-    else:
-        # No Nyquist mode
-        f_ny = n // 2
-        if a.shape[0] > f_ny:
-            a[f_ny + 1 :] = 0  # no data needed for the second half
-
-    return a
-
-
 class TestFft:
     @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
     @pytest.mark.parametrize("n", [None, 5, 20])
@@ -577,13 +551,14 @@ def test_basic(self, func, dtype, axes):
 
 
 class TestHfft:
-    @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
+    # TODO: include boolean dtype when mkl_fft-gh-180 is merged
+    @pytest.mark.parametrize(
+        "dtype", get_all_dtypes(no_none=True, no_bool=True)
+    )
     @pytest.mark.parametrize("n", [None, 5, 18])
     @pytest.mark.parametrize("norm", [None, "backward", "forward", "ortho"])
     def test_basic(self, dtype, n, norm):
         a = generate_random_numpy_array(11, dtype)
-        if numpy.issubdtype(dtype, numpy.complexfloating):
-            a = _make_array_Hermitian(a, n)
         ia = dpnp.array(a)
 
         result = dpnp.fft.hfft(ia, n=n, norm=norm)
@@ -626,8 +601,6 @@ class TestIrfft:
     @pytest.mark.parametrize("norm", [None, "backward", "forward", "ortho"])
     def test_basic(self, dtype, n, norm):
         a = generate_random_numpy_array(11)
-        if numpy.issubdtype(dtype, numpy.complexfloating):
-            a = _make_array_Hermitian(a, n)
         ia = dpnp.array(a)
 
         result = dpnp.fft.irfft(ia, n=n, norm=norm)
@@ -644,10 +617,6 @@ def test_basic(self, dtype, n, norm):
     def test_2d_array(self, dtype, n, axis, norm, order):
         a = generate_random_numpy_array((3, 4), dtype=dtype, order=order)
         ia = dpnp.array(a)
-        # For dpnp, each 1-D array of input should be Hermitian
-        ia = dpnp.moveaxis(ia, axis, 0)
-        ia = _make_array_Hermitian(ia, n)
-        ia = dpnp.moveaxis(ia, 0, axis)
 
         result = dpnp.fft.irfft(ia, n=n, axis=axis, norm=norm)
         expected = numpy.fft.irfft(a, n=n, axis=axis, norm=norm)
@@ -661,10 +630,6 @@ def test_2d_array(self, dtype, n, axis, norm, order):
     def test_3d_array(self, dtype, n, axis, norm, order):
         a = generate_random_numpy_array((4, 5, 6), dtype, order)
         ia = dpnp.array(a)
-        # For dpnp, each 1-D array of input should be Hermitian
-        ia = dpnp.moveaxis(ia, axis, 0)
-        ia = _make_array_Hermitian(ia, n)
-        ia = dpnp.moveaxis(ia, 0, axis)
 
         result = dpnp.fft.irfft(ia, n=n, axis=axis, norm=norm)
         expected = numpy.fft.irfft(a, n=n, axis=axis, norm=norm)
@@ -674,7 +639,6 @@ def test_3d_array(self, dtype, n, axis, norm, order):
     def test_usm_ndarray(self, n):
         a = generate_random_numpy_array(11, dtype=numpy.complex64)
         a_usm = dpt.asarray(a)
-        a_usm = _make_array_Hermitian(a_usm, n)
 
         expected = numpy.fft.irfft(a, n=n)
         out = dpt.empty(expected.shape, dtype=a_usm.real.dtype)
@@ -688,7 +652,6 @@ def test_usm_ndarray(self, n):
     def test_out(self, dtype, n, norm):
         a = generate_random_numpy_array(11, dtype=dtype)
         ia = dpnp.array(a)
-        ia = _make_array_Hermitian(ia, n)
 
         expected = numpy.fft.irfft(a, n=n, norm=norm)
         out = dpnp.empty(expected.shape, dtype=a.real.dtype)
@@ -704,9 +667,6 @@ def test_out(self, dtype, n, norm):
     def test_2d_array_out(self, dtype, n, axis, norm, order):
         a = generate_random_numpy_array((3, 4), dtype=dtype, order=order)
         ia = dpnp.array(a)
-        ia = dpnp.moveaxis(ia, axis, 0)
-        ia = _make_array_Hermitian(ia, n)
-        ia = dpnp.moveaxis(ia, 0, axis)
 
         expected = numpy.fft.irfft(a, n=n, axis=axis, norm=norm)
         out = dpnp.empty(expected.shape, dtype=expected.dtype)
@@ -994,5 +954,7 @@ def test_1d_array(self):
 
         result = dpnp.fft.irfftn(ia)
         expected = numpy.fft.irfftn(a)
-        flag = True if numpy_version() < "2.0.0" else False
+        # TODO: change to the commented line when mkl_fft-gh-180 is merged
+        flag = True
+        # flag = True if numpy_version() < "2.0.0" else False
         assert_dtype_allclose(result, expected, check_only_type_kind=flag)

dpnp/tests/third_party/cupy/fft_tests/test_fft.py

@@ -1047,10 +1047,6 @@ def test_rfft2(self, xp, dtype, order, enable_nd):
     )
     def test_irfft2(self, xp, dtype, order, enable_nd):
         # assert config.enable_nd_planning == enable_nd
-
-        if self.s is None and self.axes in [None, (-2, -1)]:
-            pytest.skip("Input is not Hermitian Symmetric")
-
         a = testing.shaped_random(self.shape, xp, dtype)
         if order == "F":
             a = xp.asfortranarray(a)
@@ -1137,10 +1133,6 @@ def test_rfftn(self, xp, dtype, order, enable_nd):
     )
     def test_irfftn(self, xp, dtype, order, enable_nd):
         # assert config.enable_nd_planning == enable_nd
-
-        if self.s is None and self.axes in [None, (-2, -1)]:
-            pytest.skip("Input is not Hermitian Symmetric")
-
         a = testing.shaped_random(self.shape, xp, dtype)
         if order == "F":
             a = xp.asfortranarray(a)