Test ifft, rfft and irfft

Author: honno

array_api_tests/test_fft.py

@@ -1,15 +1,17 @@
 import math
+from typing import Optional
 
 import pytest
 from hypothesis import given
+from hypothesis import strategies as st
 
-from array_api_tests.typing import DataType
+from array_api_tests.typing import Array, DataType
 
-from . import _array_module as xp
 from . import dtype_helpers as dh
 from . import hypothesis_helpers as hh
 from . import pytest_helpers as ph
 from . import xps
+from ._array_module import mod as xp
 
 pytestmark = [
     pytest.mark.ci,
@@ -21,6 +23,22 @@
 fft_shapes_strat = hh.shapes(min_dims=1).filter(lambda s: math.prod(s) > 1)
 
 
+def n_axis_norm_kwargs(x: Array, data: st.DataObject) -> tuple:
+    size = math.prod(x.shape)
+    n = data.draw(st.none() | st.integers(size // 2, size * 2), label="n")
+    axis = data.draw(st.integers(-1, x.ndim - 1), label="axis")
+    norm = data.draw(st.sampled_from(["backward", "ortho", "forward"]), label="norm")
+    kwargs = data.draw(
+        hh.specified_kwargs(
+            ("n", n, None),
+            ("axis", axis, -1),
+            ("norm", norm, "backward"),
+        ),
+        label="kwargs",
+    )
+    return n, axis, norm, kwargs
+
+
 def assert_fft_dtype(func_name: str, *, in_dtype: DataType, out_dtype: DataType):
     if in_dtype == xp.float32:
         expected = xp.complex64
@@ -34,29 +52,80 @@ def assert_fft_dtype(func_name: str, *, in_dtype: DataType, out_dtype: DataType)
     )
 
 
-@given(x=xps.arrays(dtype=hh.all_floating_dtypes(), shape=fft_shapes_strat))
-def test_fft(x):
-    out = xp.fft.fft(x)
+def assert_n_axis_shape(
+    func_name: str, *, x: Array, n: Optional[int], axis: int, out: Array
+):
+    if n is None:
+        expected_shape = x.shape
+    else:
+        _axis = len(x.shape) - 1 if axis == -1 else axis
+        expected_shape = x.shape[:_axis] + (n,) + x.shape[_axis + 1 :]
+    ph.assert_shape(func_name, out_shape=out.shape, expected=expected_shape)
+
+
+@given(
+    x=xps.arrays(dtype=hh.all_floating_dtypes(), shape=fft_shapes_strat),
+    data=st.data(),
+)
+def test_fft(x, data):
+    n, axis, norm, kwargs = n_axis_norm_kwargs(x, data)
+
+    out = xp.fft.fft(x, **kwargs)
+
     assert_fft_dtype("fft", in_dtype=x.dtype, out_dtype=out.dtype)
-    ph.assert_shape("fft", out_shape=out.shape, expected=x.shape)
+    assert_n_axis_shape("fft", x=x, n=n, axis=axis, out=out)
+
 
+@given(
+    x=xps.arrays(dtype=hh.all_floating_dtypes(), shape=fft_shapes_strat),
+    data=st.data(),
+)
+def test_ifft(x, data):
+    n, axis, norm, kwargs = n_axis_norm_kwargs(x, data)
+
+    out = xp.fft.ifft(x, **kwargs)
 
-@given(x=xps.arrays(dtype=hh.all_floating_dtypes(), shape=fft_shapes_strat))
-def test_ifft(x):
-    out = xp.fft.ifft(x)
     assert_fft_dtype("ifft", in_dtype=x.dtype, out_dtype=out.dtype)
-    ph.assert_shape("ifft", out_shape=out.shape, expected=x.shape)
+    assert_n_axis_shape("ifft", x=x, n=n, axis=axis, out=out)
+
+
+# TODO:
+# test_fftn
+# test_ifftn
+
+
+@given(
+    x=xps.arrays(dtype=xps.floating_dtypes(), shape=fft_shapes_strat),
+    data=st.data(),
+)
+def test_rfft(x, data):
+    n, axis, norm, kwargs = n_axis_norm_kwargs(x, data)
+
+    out = xp.fft.rfft(x, **kwargs)
+
+    assert_fft_dtype("rfft", in_dtype=x.dtype, out_dtype=out.dtype)
+    assert_n_axis_shape("rfft", x=x, n=n, axis=axis, out=out)
+
+
+@given(
+    x=xps.arrays(dtype=xps.floating_dtypes(), shape=fft_shapes_strat),
+    data=st.data(),
+)
+def test_irfft(x, data):
+    n, axis, norm, kwargs = n_axis_norm_kwargs(x, data)
 
+    out = xp.fft.irfft(x, **kwargs)
 
-@given(x=xps.arrays(dtype=hh.all_floating_dtypes(), shape=fft_shapes_strat))
-def test_fftn(x):
-    out = xp.fft.fftn(x)
-    assert_fft_dtype("fftn", in_dtype=x.dtype, out_dtype=out.dtype)
-    ph.assert_shape("fftn", out_shape=out.shape, expected=x.shape)
+    assert_fft_dtype("irfft", in_dtype=x.dtype, out_dtype=out.dtype)
+    # TODO: assert shape
 
 
-@given(x=xps.arrays(dtype=hh.all_floating_dtypes(), shape=fft_shapes_strat))
-def test_ifftn(x):
-    out = xp.fft.ifftn(x)
-    assert_fft_dtype("ifftn", in_dtype=x.dtype, out_dtype=out.dtype)
-    ph.assert_shape("ifftn", out_shape=out.shape, expected=x.shape)
+# TODO:
+# test_rfftn
+# test_irfftn
+# test_hfft
+# test_ihfft
+# fftfreq
+# rfftfreq
+# fftshift
+# ifftshift