190 lines
7.6 KiB
Python
190 lines
7.6 KiB
Python
import math
|
|
import numpy as np
|
|
|
|
from scipy._lib._array_api import (
|
|
xp_assert_equal,
|
|
assert_array_almost_equal,
|
|
assert_almost_equal,
|
|
is_cupy,
|
|
)
|
|
|
|
import pytest
|
|
|
|
from scipy import ndimage
|
|
|
|
from scipy.conftest import array_api_compatible
|
|
skip_xp_backends = pytest.mark.skip_xp_backends
|
|
pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends"),
|
|
skip_xp_backends(cpu_only=True, exceptions=['cupy', 'jax.numpy'],)]
|
|
|
|
|
|
@skip_xp_backends('jax.numpy', reason="jax-ml/jax#23827")
|
|
class TestNdimageFourier:
|
|
|
|
@pytest.mark.parametrize('shape', [(32, 16), (31, 15), (1, 10)])
|
|
@pytest.mark.parametrize('dtype, dec', [("float32", 6), ("float64", 14)])
|
|
def test_fourier_gaussian_real01(self, shape, dtype, dec, xp):
|
|
fft = getattr(xp, 'fft')
|
|
|
|
a = np.zeros(shape, dtype=dtype)
|
|
a[0, 0] = 1.0
|
|
a = xp.asarray(a)
|
|
|
|
a = fft.rfft(a, n=shape[0], axis=0)
|
|
a = fft.fft(a, n=shape[1], axis=1)
|
|
a = ndimage.fourier_gaussian(a, [5.0, 2.5], shape[0], 0)
|
|
a = fft.ifft(a, n=shape[1], axis=1)
|
|
a = fft.irfft(a, n=shape[0], axis=0)
|
|
assert_almost_equal(ndimage.sum(a), xp.asarray(1), decimal=dec,
|
|
check_0d=False)
|
|
|
|
@pytest.mark.parametrize('shape', [(32, 16), (31, 15)])
|
|
@pytest.mark.parametrize('dtype, dec', [("complex64", 6), ("complex128", 14)])
|
|
def test_fourier_gaussian_complex01(self, shape, dtype, dec, xp):
|
|
fft = getattr(xp, 'fft')
|
|
|
|
a = np.zeros(shape, dtype=dtype)
|
|
a[0, 0] = 1.0
|
|
a = xp.asarray(a)
|
|
|
|
a = fft.fft(a, n=shape[0], axis=0)
|
|
a = fft.fft(a, n=shape[1], axis=1)
|
|
a = ndimage.fourier_gaussian(a, [5.0, 2.5], -1, 0)
|
|
a = fft.ifft(a, n=shape[1], axis=1)
|
|
a = fft.ifft(a, n=shape[0], axis=0)
|
|
assert_almost_equal(ndimage.sum(xp.real(a)), xp.asarray(1.0), decimal=dec,
|
|
check_0d=False)
|
|
|
|
@pytest.mark.parametrize('shape', [(32, 16), (31, 15), (1, 10)])
|
|
@pytest.mark.parametrize('dtype, dec', [("float32", 6), ("float64", 14)])
|
|
def test_fourier_uniform_real01(self, shape, dtype, dec, xp):
|
|
fft = getattr(xp, 'fft')
|
|
|
|
a = np.zeros(shape, dtype=dtype)
|
|
a[0, 0] = 1.0
|
|
a = xp.asarray(a)
|
|
|
|
a = fft.rfft(a, n=shape[0], axis=0)
|
|
a = fft.fft(a, n=shape[1], axis=1)
|
|
a = ndimage.fourier_uniform(a, [5.0, 2.5], shape[0], 0)
|
|
a = fft.ifft(a, n=shape[1], axis=1)
|
|
a = fft.irfft(a, n=shape[0], axis=0)
|
|
assert_almost_equal(ndimage.sum(a), xp.asarray(1.0), decimal=dec,
|
|
check_0d=False)
|
|
|
|
@pytest.mark.parametrize('shape', [(32, 16), (31, 15)])
|
|
@pytest.mark.parametrize('dtype, dec', [("complex64", 6), ("complex128", 14)])
|
|
def test_fourier_uniform_complex01(self, shape, dtype, dec, xp):
|
|
fft = getattr(xp, 'fft')
|
|
|
|
a = np.zeros(shape, dtype=dtype)
|
|
a[0, 0] = 1.0
|
|
a = xp.asarray(a)
|
|
|
|
a = fft.fft(a, n=shape[0], axis=0)
|
|
a = fft.fft(a, n=shape[1], axis=1)
|
|
a = ndimage.fourier_uniform(a, [5.0, 2.5], -1, 0)
|
|
a = fft.ifft(a, n=shape[1], axis=1)
|
|
a = fft.ifft(a, n=shape[0], axis=0)
|
|
assert_almost_equal(ndimage.sum(xp.real(a)), xp.asarray(1.0), decimal=dec,
|
|
check_0d=False)
|
|
|
|
@pytest.mark.parametrize('shape', [(32, 16), (31, 15)])
|
|
@pytest.mark.parametrize('dtype, dec', [("float32", 4), ("float64", 11)])
|
|
def test_fourier_shift_real01(self, shape, dtype, dec, xp):
|
|
fft = getattr(xp, 'fft')
|
|
|
|
expected = np.arange(shape[0] * shape[1], dtype=dtype).reshape(shape)
|
|
expected = xp.asarray(expected)
|
|
|
|
a = fft.rfft(expected, n=shape[0], axis=0)
|
|
a = fft.fft(a, n=shape[1], axis=1)
|
|
a = ndimage.fourier_shift(a, [1, 1], shape[0], 0)
|
|
a = fft.ifft(a, n=shape[1], axis=1)
|
|
a = fft.irfft(a, n=shape[0], axis=0)
|
|
assert_array_almost_equal(a[1:, 1:], expected[:-1, :-1], decimal=dec)
|
|
|
|
@pytest.mark.parametrize('shape', [(32, 16), (31, 15)])
|
|
@pytest.mark.parametrize('dtype, dec', [("complex64", 4), ("complex128", 11)])
|
|
def test_fourier_shift_complex01(self, shape, dtype, dec, xp):
|
|
fft = getattr(xp, 'fft')
|
|
|
|
expected = np.arange(shape[0] * shape[1], dtype=dtype).reshape(shape)
|
|
expected = xp.asarray(expected)
|
|
|
|
a = fft.fft(expected, n=shape[0], axis=0)
|
|
a = fft.fft(a, n=shape[1], axis=1)
|
|
a = ndimage.fourier_shift(a, [1, 1], -1, 0)
|
|
a = fft.ifft(a, n=shape[1], axis=1)
|
|
a = fft.ifft(a, n=shape[0], axis=0)
|
|
assert_array_almost_equal(xp.real(a)[1:, 1:], expected[:-1, :-1], decimal=dec)
|
|
assert_array_almost_equal(xp.imag(a), xp.zeros(shape), decimal=dec)
|
|
|
|
@pytest.mark.parametrize('shape', [(32, 16), (31, 15), (1, 10)])
|
|
@pytest.mark.parametrize('dtype, dec', [("float32", 5), ("float64", 14)])
|
|
def test_fourier_ellipsoid_real01(self, shape, dtype, dec, xp):
|
|
fft = getattr(xp, 'fft')
|
|
|
|
a = np.zeros(shape, dtype=dtype)
|
|
a[0, 0] = 1.0
|
|
a = xp.asarray(a)
|
|
|
|
a = fft.rfft(a, n=shape[0], axis=0)
|
|
a = fft.fft(a, n=shape[1], axis=1)
|
|
a = ndimage.fourier_ellipsoid(a, [5.0, 2.5], shape[0], 0)
|
|
a = fft.ifft(a, n=shape[1], axis=1)
|
|
a = fft.irfft(a, n=shape[0], axis=0)
|
|
assert_almost_equal(ndimage.sum(a), xp.asarray(1.0), decimal=dec,
|
|
check_0d=False)
|
|
|
|
@pytest.mark.parametrize('shape', [(32, 16), (31, 15)])
|
|
@pytest.mark.parametrize('dtype, dec', [("complex64", 5), ("complex128", 14)])
|
|
def test_fourier_ellipsoid_complex01(self, shape, dtype, dec, xp):
|
|
fft = getattr(xp, 'fft')
|
|
|
|
a = np.zeros(shape, dtype=dtype)
|
|
a[0, 0] = 1.0
|
|
a = xp.asarray(a)
|
|
|
|
a = fft.fft(a, n=shape[0], axis=0)
|
|
a = fft.fft(a, n=shape[1], axis=1)
|
|
a = ndimage.fourier_ellipsoid(a, [5.0, 2.5], -1, 0)
|
|
a = fft.ifft(a, n=shape[1], axis=1)
|
|
a = fft.ifft(a, n=shape[0], axis=0)
|
|
assert_almost_equal(ndimage.sum(xp.real(a)), xp.asarray(1.0), decimal=dec,
|
|
check_0d=False)
|
|
|
|
def test_fourier_ellipsoid_unimplemented_ndim(self, xp):
|
|
# arrays with ndim > 3 raise NotImplementedError
|
|
x = xp.ones((4, 6, 8, 10), dtype=xp.complex128)
|
|
with pytest.raises(NotImplementedError):
|
|
ndimage.fourier_ellipsoid(x, 3)
|
|
|
|
def test_fourier_ellipsoid_1d_complex(self, xp):
|
|
# expected result of 1d ellipsoid is the same as for fourier_uniform
|
|
for shape in [(32, ), (31, )]:
|
|
for type_, dec in zip([xp.complex64, xp.complex128], [5, 14]):
|
|
x = xp.ones(shape, dtype=type_)
|
|
a = ndimage.fourier_ellipsoid(x, 5, -1, 0)
|
|
b = ndimage.fourier_uniform(x, 5, -1, 0)
|
|
assert_array_almost_equal(a, b, decimal=dec)
|
|
|
|
@pytest.mark.parametrize('shape', [(0, ), (0, 10), (10, 0)])
|
|
@pytest.mark.parametrize('dtype', ["float32", "float64",
|
|
"complex64", "complex128"])
|
|
@pytest.mark.parametrize('test_func',
|
|
[ndimage.fourier_ellipsoid,
|
|
ndimage.fourier_gaussian,
|
|
ndimage.fourier_uniform])
|
|
def test_fourier_zero_length_dims(self, shape, dtype, test_func, xp):
|
|
if is_cupy(xp):
|
|
if (test_func.__name__ == "fourier_ellipsoid" and
|
|
math.prod(shape) == 0):
|
|
pytest.xfail(
|
|
"CuPy's fourier_ellipsoid does not accept size==0 arrays"
|
|
)
|
|
dtype = getattr(xp, dtype)
|
|
a = xp.ones(shape, dtype=dtype)
|
|
b = test_func(a, 3)
|
|
xp_assert_equal(a, b)
|