|
|
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): |
|
|
|
|
|
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): |
|
|
|
|
|
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) |
|
|
|
