asd
This commit is contained in:
@ -0,0 +1,92 @@
|
||||
import numpy as np
|
||||
import scipy.fft
|
||||
|
||||
class _MockFunction:
|
||||
def __init__(self, return_value = None):
|
||||
self.number_calls = 0
|
||||
self.return_value = return_value
|
||||
self.last_args = ([], {})
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
self.number_calls += 1
|
||||
self.last_args = (args, kwargs)
|
||||
return self.return_value
|
||||
|
||||
|
||||
fft = _MockFunction(np.random.random(10))
|
||||
fft2 = _MockFunction(np.random.random(10))
|
||||
fftn = _MockFunction(np.random.random(10))
|
||||
|
||||
ifft = _MockFunction(np.random.random(10))
|
||||
ifft2 = _MockFunction(np.random.random(10))
|
||||
ifftn = _MockFunction(np.random.random(10))
|
||||
|
||||
rfft = _MockFunction(np.random.random(10))
|
||||
rfft2 = _MockFunction(np.random.random(10))
|
||||
rfftn = _MockFunction(np.random.random(10))
|
||||
|
||||
irfft = _MockFunction(np.random.random(10))
|
||||
irfft2 = _MockFunction(np.random.random(10))
|
||||
irfftn = _MockFunction(np.random.random(10))
|
||||
|
||||
hfft = _MockFunction(np.random.random(10))
|
||||
hfft2 = _MockFunction(np.random.random(10))
|
||||
hfftn = _MockFunction(np.random.random(10))
|
||||
|
||||
ihfft = _MockFunction(np.random.random(10))
|
||||
ihfft2 = _MockFunction(np.random.random(10))
|
||||
ihfftn = _MockFunction(np.random.random(10))
|
||||
|
||||
dct = _MockFunction(np.random.random(10))
|
||||
idct = _MockFunction(np.random.random(10))
|
||||
dctn = _MockFunction(np.random.random(10))
|
||||
idctn = _MockFunction(np.random.random(10))
|
||||
|
||||
dst = _MockFunction(np.random.random(10))
|
||||
idst = _MockFunction(np.random.random(10))
|
||||
dstn = _MockFunction(np.random.random(10))
|
||||
idstn = _MockFunction(np.random.random(10))
|
||||
|
||||
fht = _MockFunction(np.random.random(10))
|
||||
ifht = _MockFunction(np.random.random(10))
|
||||
|
||||
|
||||
__ua_domain__ = "numpy.scipy.fft"
|
||||
|
||||
|
||||
_implements = {
|
||||
scipy.fft.fft: fft,
|
||||
scipy.fft.fft2: fft2,
|
||||
scipy.fft.fftn: fftn,
|
||||
scipy.fft.ifft: ifft,
|
||||
scipy.fft.ifft2: ifft2,
|
||||
scipy.fft.ifftn: ifftn,
|
||||
scipy.fft.rfft: rfft,
|
||||
scipy.fft.rfft2: rfft2,
|
||||
scipy.fft.rfftn: rfftn,
|
||||
scipy.fft.irfft: irfft,
|
||||
scipy.fft.irfft2: irfft2,
|
||||
scipy.fft.irfftn: irfftn,
|
||||
scipy.fft.hfft: hfft,
|
||||
scipy.fft.hfft2: hfft2,
|
||||
scipy.fft.hfftn: hfftn,
|
||||
scipy.fft.ihfft: ihfft,
|
||||
scipy.fft.ihfft2: ihfft2,
|
||||
scipy.fft.ihfftn: ihfftn,
|
||||
scipy.fft.dct: dct,
|
||||
scipy.fft.idct: idct,
|
||||
scipy.fft.dctn: dctn,
|
||||
scipy.fft.idctn: idctn,
|
||||
scipy.fft.dst: dst,
|
||||
scipy.fft.idst: idst,
|
||||
scipy.fft.dstn: dstn,
|
||||
scipy.fft.idstn: idstn,
|
||||
scipy.fft.fht: fht,
|
||||
scipy.fft.ifht: ifht
|
||||
}
|
||||
|
||||
|
||||
def __ua_function__(method, args, kwargs):
|
||||
fn = _implements.get(method)
|
||||
return (fn(*args, **kwargs) if fn is not None
|
||||
else NotImplemented)
|
||||
@ -0,0 +1,98 @@
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
import scipy.fft
|
||||
from scipy.fft import _fftlog, _pocketfft, set_backend
|
||||
from scipy.fft.tests import mock_backend
|
||||
|
||||
from numpy.testing import assert_allclose, assert_equal
|
||||
import pytest
|
||||
|
||||
fnames = ('fft', 'fft2', 'fftn',
|
||||
'ifft', 'ifft2', 'ifftn',
|
||||
'rfft', 'rfft2', 'rfftn',
|
||||
'irfft', 'irfft2', 'irfftn',
|
||||
'dct', 'idct', 'dctn', 'idctn',
|
||||
'dst', 'idst', 'dstn', 'idstn',
|
||||
'fht', 'ifht')
|
||||
|
||||
np_funcs = (np.fft.fft, np.fft.fft2, np.fft.fftn,
|
||||
np.fft.ifft, np.fft.ifft2, np.fft.ifftn,
|
||||
np.fft.rfft, np.fft.rfft2, np.fft.rfftn,
|
||||
np.fft.irfft, np.fft.irfft2, np.fft.irfftn,
|
||||
np.fft.hfft, _pocketfft.hfft2, _pocketfft.hfftn, # np has no hfftn
|
||||
np.fft.ihfft, _pocketfft.ihfft2, _pocketfft.ihfftn,
|
||||
_pocketfft.dct, _pocketfft.idct, _pocketfft.dctn, _pocketfft.idctn,
|
||||
_pocketfft.dst, _pocketfft.idst, _pocketfft.dstn, _pocketfft.idstn,
|
||||
# must provide required kwargs for fht, ifht
|
||||
partial(_fftlog.fht, dln=2, mu=0.5),
|
||||
partial(_fftlog.ifht, dln=2, mu=0.5))
|
||||
|
||||
funcs = (scipy.fft.fft, scipy.fft.fft2, scipy.fft.fftn,
|
||||
scipy.fft.ifft, scipy.fft.ifft2, scipy.fft.ifftn,
|
||||
scipy.fft.rfft, scipy.fft.rfft2, scipy.fft.rfftn,
|
||||
scipy.fft.irfft, scipy.fft.irfft2, scipy.fft.irfftn,
|
||||
scipy.fft.hfft, scipy.fft.hfft2, scipy.fft.hfftn,
|
||||
scipy.fft.ihfft, scipy.fft.ihfft2, scipy.fft.ihfftn,
|
||||
scipy.fft.dct, scipy.fft.idct, scipy.fft.dctn, scipy.fft.idctn,
|
||||
scipy.fft.dst, scipy.fft.idst, scipy.fft.dstn, scipy.fft.idstn,
|
||||
# must provide required kwargs for fht, ifht
|
||||
partial(scipy.fft.fht, dln=2, mu=0.5),
|
||||
partial(scipy.fft.ifht, dln=2, mu=0.5))
|
||||
|
||||
mocks = (mock_backend.fft, mock_backend.fft2, mock_backend.fftn,
|
||||
mock_backend.ifft, mock_backend.ifft2, mock_backend.ifftn,
|
||||
mock_backend.rfft, mock_backend.rfft2, mock_backend.rfftn,
|
||||
mock_backend.irfft, mock_backend.irfft2, mock_backend.irfftn,
|
||||
mock_backend.hfft, mock_backend.hfft2, mock_backend.hfftn,
|
||||
mock_backend.ihfft, mock_backend.ihfft2, mock_backend.ihfftn,
|
||||
mock_backend.dct, mock_backend.idct,
|
||||
mock_backend.dctn, mock_backend.idctn,
|
||||
mock_backend.dst, mock_backend.idst,
|
||||
mock_backend.dstn, mock_backend.idstn,
|
||||
mock_backend.fht, mock_backend.ifht)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("func, np_func, mock", zip(funcs, np_funcs, mocks))
|
||||
def test_backend_call(func, np_func, mock):
|
||||
x = np.arange(20).reshape((10,2))
|
||||
answer = np_func(x.astype(np.float64))
|
||||
assert_allclose(func(x), answer, atol=1e-10)
|
||||
|
||||
with set_backend(mock_backend, only=True):
|
||||
mock.number_calls = 0
|
||||
y = func(x)
|
||||
assert_equal(y, mock.return_value)
|
||||
assert_equal(mock.number_calls, 1)
|
||||
|
||||
assert_allclose(func(x), answer, atol=1e-10)
|
||||
|
||||
|
||||
plan_funcs = (scipy.fft.fft, scipy.fft.fft2, scipy.fft.fftn,
|
||||
scipy.fft.ifft, scipy.fft.ifft2, scipy.fft.ifftn,
|
||||
scipy.fft.rfft, scipy.fft.rfft2, scipy.fft.rfftn,
|
||||
scipy.fft.irfft, scipy.fft.irfft2, scipy.fft.irfftn,
|
||||
scipy.fft.hfft, scipy.fft.hfft2, scipy.fft.hfftn,
|
||||
scipy.fft.ihfft, scipy.fft.ihfft2, scipy.fft.ihfftn)
|
||||
|
||||
plan_mocks = (mock_backend.fft, mock_backend.fft2, mock_backend.fftn,
|
||||
mock_backend.ifft, mock_backend.ifft2, mock_backend.ifftn,
|
||||
mock_backend.rfft, mock_backend.rfft2, mock_backend.rfftn,
|
||||
mock_backend.irfft, mock_backend.irfft2, mock_backend.irfftn,
|
||||
mock_backend.hfft, mock_backend.hfft2, mock_backend.hfftn,
|
||||
mock_backend.ihfft, mock_backend.ihfft2, mock_backend.ihfftn)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("func, mock", zip(plan_funcs, plan_mocks))
|
||||
def test_backend_plan(func, mock):
|
||||
x = np.arange(20).reshape((10, 2))
|
||||
|
||||
with pytest.raises(NotImplementedError, match='precomputed plan'):
|
||||
func(x, plan='foo')
|
||||
|
||||
with set_backend(mock_backend, only=True):
|
||||
mock.number_calls = 0
|
||||
y = func(x, plan='foo')
|
||||
assert_equal(y, mock.return_value)
|
||||
assert_equal(mock.number_calls, 1)
|
||||
assert_equal(mock.last_args[1]['plan'], 'foo')
|
||||
491
venv/lib/python3.12/site-packages/scipy/fft/tests/test_basic.py
Normal file
491
venv/lib/python3.12/site-packages/scipy/fft/tests/test_basic.py
Normal file
@ -0,0 +1,491 @@
|
||||
import queue
|
||||
import threading
|
||||
import multiprocessing
|
||||
import numpy as np
|
||||
import pytest
|
||||
from numpy.random import random
|
||||
from numpy.testing import assert_array_almost_equal, assert_allclose
|
||||
from pytest import raises as assert_raises
|
||||
import scipy.fft as fft
|
||||
from scipy.conftest import array_api_compatible
|
||||
from scipy._lib._array_api import (
|
||||
array_namespace, size, xp_assert_close, xp_assert_equal
|
||||
)
|
||||
|
||||
pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends")]
|
||||
skip_xp_backends = pytest.mark.skip_xp_backends
|
||||
|
||||
|
||||
# Expected input dtypes. Note that `scipy.fft` is more flexible for numpy,
|
||||
# but for C2C transforms like `fft.fft`, the array API standard only mandates
|
||||
# that complex dtypes should work, float32/float64 aren't guaranteed to.
|
||||
def get_expected_input_dtype(func, xp):
|
||||
if func in [fft.fft, fft.fftn, fft.fft2,
|
||||
fft.ifft, fft.ifftn, fft.ifft2,
|
||||
fft.hfft, fft.hfftn, fft.hfft2,
|
||||
fft.irfft, fft.irfftn, fft.irfft2]:
|
||||
dtype = xp.complex128
|
||||
elif func in [fft.rfft, fft.rfftn, fft.rfft2,
|
||||
fft.ihfft, fft.ihfftn, fft.ihfft2]:
|
||||
dtype = xp.float64
|
||||
else:
|
||||
raise ValueError(f'Unknown FFT function: {func}')
|
||||
|
||||
return dtype
|
||||
|
||||
|
||||
def fft1(x):
|
||||
L = len(x)
|
||||
phase = -2j*np.pi*(np.arange(L)/float(L))
|
||||
phase = np.arange(L).reshape(-1, 1) * phase
|
||||
return np.sum(x*np.exp(phase), axis=1)
|
||||
|
||||
class TestFFT:
|
||||
|
||||
def test_identity(self, xp):
|
||||
maxlen = 512
|
||||
x = xp.asarray(random(maxlen) + 1j*random(maxlen))
|
||||
xr = xp.asarray(random(maxlen))
|
||||
# Check some powers of 2 and some primes
|
||||
for i in [1, 2, 16, 128, 512, 53, 149, 281, 397]:
|
||||
xp_assert_close(fft.ifft(fft.fft(x[0:i])), x[0:i])
|
||||
xp_assert_close(fft.irfft(fft.rfft(xr[0:i]), i), xr[0:i])
|
||||
|
||||
@skip_xp_backends(np_only=True, reasons=['significant overhead for some backends'])
|
||||
def test_identity_extensive(self, xp):
|
||||
maxlen = 512
|
||||
x = xp.asarray(random(maxlen) + 1j*random(maxlen))
|
||||
xr = xp.asarray(random(maxlen))
|
||||
for i in range(1, maxlen):
|
||||
xp_assert_close(fft.ifft(fft.fft(x[0:i])), x[0:i])
|
||||
xp_assert_close(fft.irfft(fft.rfft(xr[0:i]), i), xr[0:i])
|
||||
|
||||
def test_fft(self, xp):
|
||||
x = random(30) + 1j*random(30)
|
||||
expect = xp.asarray(fft1(x))
|
||||
x = xp.asarray(x)
|
||||
xp_assert_close(fft.fft(x), expect)
|
||||
xp_assert_close(fft.fft(x, norm="backward"), expect)
|
||||
xp_assert_close(fft.fft(x, norm="ortho"),
|
||||
expect / xp.sqrt(xp.asarray(30, dtype=xp.float64)),)
|
||||
xp_assert_close(fft.fft(x, norm="forward"), expect / 30)
|
||||
|
||||
@skip_xp_backends(np_only=True, reasons=['some backends allow `n=0`'])
|
||||
def test_fft_n(self, xp):
|
||||
x = xp.asarray([1, 2, 3], dtype=xp.complex128)
|
||||
assert_raises(ValueError, fft.fft, x, 0)
|
||||
|
||||
def test_ifft(self, xp):
|
||||
x = xp.asarray(random(30) + 1j*random(30))
|
||||
xp_assert_close(fft.ifft(fft.fft(x)), x)
|
||||
for norm in ["backward", "ortho", "forward"]:
|
||||
xp_assert_close(fft.ifft(fft.fft(x, norm=norm), norm=norm), x)
|
||||
|
||||
def test_fft2(self, xp):
|
||||
x = xp.asarray(random((30, 20)) + 1j*random((30, 20)))
|
||||
expect = fft.fft(fft.fft(x, axis=1), axis=0)
|
||||
xp_assert_close(fft.fft2(x), expect)
|
||||
xp_assert_close(fft.fft2(x, norm="backward"), expect)
|
||||
xp_assert_close(fft.fft2(x, norm="ortho"),
|
||||
expect / xp.sqrt(xp.asarray(30 * 20, dtype=xp.float64)))
|
||||
xp_assert_close(fft.fft2(x, norm="forward"), expect / (30 * 20))
|
||||
|
||||
def test_ifft2(self, xp):
|
||||
x = xp.asarray(random((30, 20)) + 1j*random((30, 20)))
|
||||
expect = fft.ifft(fft.ifft(x, axis=1), axis=0)
|
||||
xp_assert_close(fft.ifft2(x), expect)
|
||||
xp_assert_close(fft.ifft2(x, norm="backward"), expect)
|
||||
xp_assert_close(fft.ifft2(x, norm="ortho"),
|
||||
expect * xp.sqrt(xp.asarray(30 * 20, dtype=xp.float64)))
|
||||
xp_assert_close(fft.ifft2(x, norm="forward"), expect * (30 * 20))
|
||||
|
||||
def test_fftn(self, xp):
|
||||
x = xp.asarray(random((30, 20, 10)) + 1j*random((30, 20, 10)))
|
||||
expect = fft.fft(fft.fft(fft.fft(x, axis=2), axis=1), axis=0)
|
||||
xp_assert_close(fft.fftn(x), expect)
|
||||
xp_assert_close(fft.fftn(x, norm="backward"), expect)
|
||||
xp_assert_close(fft.fftn(x, norm="ortho"),
|
||||
expect / xp.sqrt(xp.asarray(30 * 20 * 10, dtype=xp.float64)))
|
||||
xp_assert_close(fft.fftn(x, norm="forward"), expect / (30 * 20 * 10))
|
||||
|
||||
def test_ifftn(self, xp):
|
||||
x = xp.asarray(random((30, 20, 10)) + 1j*random((30, 20, 10)))
|
||||
expect = fft.ifft(fft.ifft(fft.ifft(x, axis=2), axis=1), axis=0)
|
||||
xp_assert_close(fft.ifftn(x), expect, rtol=1e-7)
|
||||
xp_assert_close(fft.ifftn(x, norm="backward"), expect, rtol=1e-7)
|
||||
xp_assert_close(
|
||||
fft.ifftn(x, norm="ortho"),
|
||||
fft.ifftn(x) * xp.sqrt(xp.asarray(30 * 20 * 10, dtype=xp.float64))
|
||||
)
|
||||
xp_assert_close(fft.ifftn(x, norm="forward"),
|
||||
expect * (30 * 20 * 10),
|
||||
rtol=1e-7)
|
||||
|
||||
def test_rfft(self, xp):
|
||||
x = xp.asarray(random(29), dtype=xp.float64)
|
||||
for n in [size(x), 2*size(x)]:
|
||||
for norm in [None, "backward", "ortho", "forward"]:
|
||||
xp_assert_close(fft.rfft(x, n=n, norm=norm),
|
||||
fft.fft(xp.asarray(x, dtype=xp.complex128),
|
||||
n=n, norm=norm)[:(n//2 + 1)])
|
||||
xp_assert_close(
|
||||
fft.rfft(x, n=n, norm="ortho"),
|
||||
fft.rfft(x, n=n) / xp.sqrt(xp.asarray(n, dtype=xp.float64))
|
||||
)
|
||||
|
||||
def test_irfft(self, xp):
|
||||
x = xp.asarray(random(30))
|
||||
xp_assert_close(fft.irfft(fft.rfft(x)), x)
|
||||
for norm in ["backward", "ortho", "forward"]:
|
||||
xp_assert_close(fft.irfft(fft.rfft(x, norm=norm), norm=norm), x)
|
||||
|
||||
def test_rfft2(self, xp):
|
||||
x = xp.asarray(random((30, 20)), dtype=xp.float64)
|
||||
expect = fft.fft2(xp.asarray(x, dtype=xp.complex128))[:, :11]
|
||||
xp_assert_close(fft.rfft2(x), expect)
|
||||
xp_assert_close(fft.rfft2(x, norm="backward"), expect)
|
||||
xp_assert_close(fft.rfft2(x, norm="ortho"),
|
||||
expect / xp.sqrt(xp.asarray(30 * 20, dtype=xp.float64)))
|
||||
xp_assert_close(fft.rfft2(x, norm="forward"), expect / (30 * 20))
|
||||
|
||||
def test_irfft2(self, xp):
|
||||
x = xp.asarray(random((30, 20)))
|
||||
xp_assert_close(fft.irfft2(fft.rfft2(x)), x)
|
||||
for norm in ["backward", "ortho", "forward"]:
|
||||
xp_assert_close(fft.irfft2(fft.rfft2(x, norm=norm), norm=norm), x)
|
||||
|
||||
def test_rfftn(self, xp):
|
||||
x = xp.asarray(random((30, 20, 10)), dtype=xp.float64)
|
||||
expect = fft.fftn(xp.asarray(x, dtype=xp.complex128))[:, :, :6]
|
||||
xp_assert_close(fft.rfftn(x), expect)
|
||||
xp_assert_close(fft.rfftn(x, norm="backward"), expect)
|
||||
xp_assert_close(fft.rfftn(x, norm="ortho"),
|
||||
expect / xp.sqrt(xp.asarray(30 * 20 * 10, dtype=xp.float64)))
|
||||
xp_assert_close(fft.rfftn(x, norm="forward"), expect / (30 * 20 * 10))
|
||||
|
||||
def test_irfftn(self, xp):
|
||||
x = xp.asarray(random((30, 20, 10)))
|
||||
xp_assert_close(fft.irfftn(fft.rfftn(x)), x)
|
||||
for norm in ["backward", "ortho", "forward"]:
|
||||
xp_assert_close(fft.irfftn(fft.rfftn(x, norm=norm), norm=norm), x)
|
||||
|
||||
def test_hfft(self, xp):
|
||||
x = random(14) + 1j*random(14)
|
||||
x_herm = np.concatenate((random(1), x, random(1)))
|
||||
x = np.concatenate((x_herm, x[::-1].conj()))
|
||||
x = xp.asarray(x)
|
||||
x_herm = xp.asarray(x_herm)
|
||||
expect = xp.real(fft.fft(x))
|
||||
xp_assert_close(fft.hfft(x_herm), expect)
|
||||
xp_assert_close(fft.hfft(x_herm, norm="backward"), expect)
|
||||
xp_assert_close(fft.hfft(x_herm, norm="ortho"),
|
||||
expect / xp.sqrt(xp.asarray(30, dtype=xp.float64)))
|
||||
xp_assert_close(fft.hfft(x_herm, norm="forward"), expect / 30)
|
||||
|
||||
def test_ihfft(self, xp):
|
||||
x = random(14) + 1j*random(14)
|
||||
x_herm = np.concatenate((random(1), x, random(1)))
|
||||
x = np.concatenate((x_herm, x[::-1].conj()))
|
||||
x = xp.asarray(x)
|
||||
x_herm = xp.asarray(x_herm)
|
||||
xp_assert_close(fft.ihfft(fft.hfft(x_herm)), x_herm)
|
||||
for norm in ["backward", "ortho", "forward"]:
|
||||
xp_assert_close(fft.ihfft(fft.hfft(x_herm, norm=norm), norm=norm), x_herm)
|
||||
|
||||
def test_hfft2(self, xp):
|
||||
x = xp.asarray(random((30, 20)))
|
||||
xp_assert_close(fft.hfft2(fft.ihfft2(x)), x)
|
||||
for norm in ["backward", "ortho", "forward"]:
|
||||
xp_assert_close(fft.hfft2(fft.ihfft2(x, norm=norm), norm=norm), x)
|
||||
|
||||
def test_ihfft2(self, xp):
|
||||
x = xp.asarray(random((30, 20)), dtype=xp.float64)
|
||||
expect = fft.ifft2(xp.asarray(x, dtype=xp.complex128))[:, :11]
|
||||
xp_assert_close(fft.ihfft2(x), expect)
|
||||
xp_assert_close(fft.ihfft2(x, norm="backward"), expect)
|
||||
xp_assert_close(
|
||||
fft.ihfft2(x, norm="ortho"),
|
||||
expect * xp.sqrt(xp.asarray(30 * 20, dtype=xp.float64))
|
||||
)
|
||||
xp_assert_close(fft.ihfft2(x, norm="forward"), expect * (30 * 20))
|
||||
|
||||
def test_hfftn(self, xp):
|
||||
x = xp.asarray(random((30, 20, 10)))
|
||||
xp_assert_close(fft.hfftn(fft.ihfftn(x)), x)
|
||||
for norm in ["backward", "ortho", "forward"]:
|
||||
xp_assert_close(fft.hfftn(fft.ihfftn(x, norm=norm), norm=norm), x)
|
||||
|
||||
def test_ihfftn(self, xp):
|
||||
x = xp.asarray(random((30, 20, 10)), dtype=xp.float64)
|
||||
expect = fft.ifftn(xp.asarray(x, dtype=xp.complex128))[:, :, :6]
|
||||
xp_assert_close(expect, fft.ihfftn(x))
|
||||
xp_assert_close(expect, fft.ihfftn(x, norm="backward"))
|
||||
xp_assert_close(
|
||||
fft.ihfftn(x, norm="ortho"),
|
||||
expect * xp.sqrt(xp.asarray(30 * 20 * 10, dtype=xp.float64))
|
||||
)
|
||||
xp_assert_close(fft.ihfftn(x, norm="forward"), expect * (30 * 20 * 10))
|
||||
|
||||
def _check_axes(self, op, xp):
|
||||
dtype = get_expected_input_dtype(op, xp)
|
||||
x = xp.asarray(random((30, 20, 10)), dtype=dtype)
|
||||
axes = [(0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0)]
|
||||
xp_test = array_namespace(x)
|
||||
for a in axes:
|
||||
op_tr = op(xp_test.permute_dims(x, axes=a))
|
||||
tr_op = xp_test.permute_dims(op(x, axes=a), axes=a)
|
||||
xp_assert_close(op_tr, tr_op)
|
||||
|
||||
@pytest.mark.parametrize("op", [fft.fftn, fft.ifftn, fft.rfftn, fft.irfftn])
|
||||
def test_axes_standard(self, op, xp):
|
||||
self._check_axes(op, xp)
|
||||
|
||||
@pytest.mark.parametrize("op", [fft.hfftn, fft.ihfftn])
|
||||
def test_axes_non_standard(self, op, xp):
|
||||
self._check_axes(op, xp)
|
||||
|
||||
@pytest.mark.parametrize("op", [fft.fftn, fft.ifftn,
|
||||
fft.rfftn, fft.irfftn])
|
||||
def test_axes_subset_with_shape_standard(self, op, xp):
|
||||
dtype = get_expected_input_dtype(op, xp)
|
||||
x = xp.asarray(random((16, 8, 4)), dtype=dtype)
|
||||
axes = [(0, 1, 2), (0, 2, 1), (1, 2, 0)]
|
||||
xp_test = array_namespace(x)
|
||||
for a in axes:
|
||||
# different shape on the first two axes
|
||||
shape = tuple([2*x.shape[ax] if ax in a[:2] else x.shape[ax]
|
||||
for ax in range(x.ndim)])
|
||||
# transform only the first two axes
|
||||
op_tr = op(xp_test.permute_dims(x, axes=a),
|
||||
s=shape[:2], axes=(0, 1))
|
||||
tr_op = xp_test.permute_dims(op(x, s=shape[:2], axes=a[:2]),
|
||||
axes=a)
|
||||
xp_assert_close(op_tr, tr_op)
|
||||
|
||||
@pytest.mark.parametrize("op", [fft.fft2, fft.ifft2,
|
||||
fft.rfft2, fft.irfft2,
|
||||
fft.hfft2, fft.ihfft2,
|
||||
fft.hfftn, fft.ihfftn])
|
||||
def test_axes_subset_with_shape_non_standard(self, op, xp):
|
||||
dtype = get_expected_input_dtype(op, xp)
|
||||
x = xp.asarray(random((16, 8, 4)), dtype=dtype)
|
||||
axes = [(0, 1, 2), (0, 2, 1), (1, 2, 0)]
|
||||
xp_test = array_namespace(x)
|
||||
for a in axes:
|
||||
# different shape on the first two axes
|
||||
shape = tuple([2*x.shape[ax] if ax in a[:2] else x.shape[ax]
|
||||
for ax in range(x.ndim)])
|
||||
# transform only the first two axes
|
||||
op_tr = op(xp_test.permute_dims(x, axes=a), s=shape[:2], axes=(0, 1))
|
||||
tr_op = xp_test.permute_dims(op(x, s=shape[:2], axes=a[:2]), axes=a)
|
||||
xp_assert_close(op_tr, tr_op)
|
||||
|
||||
def test_all_1d_norm_preserving(self, xp):
|
||||
# verify that round-trip transforms are norm-preserving
|
||||
x = xp.asarray(random(30), dtype=xp.float64)
|
||||
xp_test = array_namespace(x)
|
||||
x_norm = xp_test.linalg.vector_norm(x)
|
||||
n = size(x) * 2
|
||||
func_pairs = [(fft.rfft, fft.irfft),
|
||||
# hfft: order so the first function takes x.size samples
|
||||
# (necessary for comparison to x_norm above)
|
||||
(fft.ihfft, fft.hfft),
|
||||
# functions that expect complex dtypes at the end
|
||||
(fft.fft, fft.ifft),
|
||||
]
|
||||
for forw, back in func_pairs:
|
||||
if forw == fft.fft:
|
||||
x = xp.asarray(x, dtype=xp.complex128)
|
||||
x_norm = xp_test.linalg.vector_norm(x)
|
||||
for n in [size(x), 2*size(x)]:
|
||||
for norm in ['backward', 'ortho', 'forward']:
|
||||
tmp = forw(x, n=n, norm=norm)
|
||||
tmp = back(tmp, n=n, norm=norm)
|
||||
xp_assert_close(xp_test.linalg.vector_norm(tmp), x_norm)
|
||||
|
||||
@skip_xp_backends(np_only=True)
|
||||
@pytest.mark.parametrize("dtype", [np.float16, np.longdouble])
|
||||
def test_dtypes_nonstandard(self, dtype):
|
||||
x = random(30).astype(dtype)
|
||||
out_dtypes = {np.float16: np.complex64, np.longdouble: np.clongdouble}
|
||||
x_complex = x.astype(out_dtypes[dtype])
|
||||
|
||||
res_fft = fft.ifft(fft.fft(x))
|
||||
res_rfft = fft.irfft(fft.rfft(x))
|
||||
res_hfft = fft.hfft(fft.ihfft(x), x.shape[0])
|
||||
# Check both numerical results and exact dtype matches
|
||||
assert_array_almost_equal(res_fft, x_complex)
|
||||
assert_array_almost_equal(res_rfft, x)
|
||||
assert_array_almost_equal(res_hfft, x)
|
||||
assert res_fft.dtype == x_complex.dtype
|
||||
assert res_rfft.dtype == np.result_type(np.float32, x.dtype)
|
||||
assert res_hfft.dtype == np.result_type(np.float32, x.dtype)
|
||||
|
||||
@pytest.mark.parametrize("dtype", ["float32", "float64"])
|
||||
def test_dtypes_real(self, dtype, xp):
|
||||
x = xp.asarray(random(30), dtype=getattr(xp, dtype))
|
||||
|
||||
res_rfft = fft.irfft(fft.rfft(x))
|
||||
res_hfft = fft.hfft(fft.ihfft(x), x.shape[0])
|
||||
# Check both numerical results and exact dtype matches
|
||||
xp_assert_close(res_rfft, x)
|
||||
xp_assert_close(res_hfft, x)
|
||||
|
||||
@pytest.mark.parametrize("dtype", ["complex64", "complex128"])
|
||||
def test_dtypes_complex(self, dtype, xp):
|
||||
x = xp.asarray(random(30), dtype=getattr(xp, dtype))
|
||||
|
||||
res_fft = fft.ifft(fft.fft(x))
|
||||
# Check both numerical results and exact dtype matches
|
||||
xp_assert_close(res_fft, x)
|
||||
|
||||
@skip_xp_backends(np_only=True,
|
||||
reasons=['array-likes only supported for NumPy backend'])
|
||||
@pytest.mark.parametrize("op", [fft.fft, fft.ifft,
|
||||
fft.fft2, fft.ifft2,
|
||||
fft.fftn, fft.ifftn,
|
||||
fft.rfft, fft.irfft,
|
||||
fft.rfft2, fft.irfft2,
|
||||
fft.rfftn, fft.irfftn,
|
||||
fft.hfft, fft.ihfft,
|
||||
fft.hfft2, fft.ihfft2,
|
||||
fft.hfftn, fft.ihfftn,])
|
||||
def test_array_like(self, xp, op):
|
||||
x = [[[1.0, 1.0], [1.0, 1.0]],
|
||||
[[1.0, 1.0], [1.0, 1.0]],
|
||||
[[1.0, 1.0], [1.0, 1.0]]]
|
||||
xp_assert_close(op(x), op(xp.asarray(x)))
|
||||
|
||||
|
||||
@skip_xp_backends(np_only=True)
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[np.float32, np.float64, np.longdouble,
|
||||
np.complex64, np.complex128, np.clongdouble])
|
||||
@pytest.mark.parametrize("order", ["F", 'non-contiguous'])
|
||||
@pytest.mark.parametrize(
|
||||
"fft",
|
||||
[fft.fft, fft.fft2, fft.fftn,
|
||||
fft.ifft, fft.ifft2, fft.ifftn])
|
||||
def test_fft_with_order(dtype, order, fft):
|
||||
# Check that FFT/IFFT produces identical results for C, Fortran and
|
||||
# non contiguous arrays
|
||||
rng = np.random.RandomState(42)
|
||||
X = rng.rand(8, 7, 13).astype(dtype, copy=False)
|
||||
if order == 'F':
|
||||
Y = np.asfortranarray(X)
|
||||
else:
|
||||
# Make a non contiguous array
|
||||
Y = X[::-1]
|
||||
X = np.ascontiguousarray(X[::-1])
|
||||
|
||||
if fft.__name__.endswith('fft'):
|
||||
for axis in range(3):
|
||||
X_res = fft(X, axis=axis)
|
||||
Y_res = fft(Y, axis=axis)
|
||||
assert_array_almost_equal(X_res, Y_res)
|
||||
elif fft.__name__.endswith(('fft2', 'fftn')):
|
||||
axes = [(0, 1), (1, 2), (0, 2)]
|
||||
if fft.__name__.endswith('fftn'):
|
||||
axes.extend([(0,), (1,), (2,), None])
|
||||
for ax in axes:
|
||||
X_res = fft(X, axes=ax)
|
||||
Y_res = fft(Y, axes=ax)
|
||||
assert_array_almost_equal(X_res, Y_res)
|
||||
else:
|
||||
raise ValueError
|
||||
|
||||
|
||||
@skip_xp_backends(cpu_only=True)
|
||||
class TestFFTThreadSafe:
|
||||
threads = 16
|
||||
input_shape = (800, 200)
|
||||
|
||||
def _test_mtsame(self, func, *args, xp=None):
|
||||
def worker(args, q):
|
||||
q.put(func(*args))
|
||||
|
||||
q = queue.Queue()
|
||||
expected = func(*args)
|
||||
|
||||
# Spin off a bunch of threads to call the same function simultaneously
|
||||
t = [threading.Thread(target=worker, args=(args, q))
|
||||
for i in range(self.threads)]
|
||||
[x.start() for x in t]
|
||||
|
||||
[x.join() for x in t]
|
||||
|
||||
# Make sure all threads returned the correct value
|
||||
for i in range(self.threads):
|
||||
xp_assert_equal(
|
||||
q.get(timeout=5), expected,
|
||||
err_msg='Function returned wrong value in multithreaded context'
|
||||
)
|
||||
|
||||
def test_fft(self, xp):
|
||||
a = xp.ones(self.input_shape, dtype=xp.complex128)
|
||||
self._test_mtsame(fft.fft, a, xp=xp)
|
||||
|
||||
def test_ifft(self, xp):
|
||||
a = xp.full(self.input_shape, 1+0j)
|
||||
self._test_mtsame(fft.ifft, a, xp=xp)
|
||||
|
||||
def test_rfft(self, xp):
|
||||
a = xp.ones(self.input_shape)
|
||||
self._test_mtsame(fft.rfft, a, xp=xp)
|
||||
|
||||
def test_irfft(self, xp):
|
||||
a = xp.full(self.input_shape, 1+0j)
|
||||
self._test_mtsame(fft.irfft, a, xp=xp)
|
||||
|
||||
def test_hfft(self, xp):
|
||||
a = xp.ones(self.input_shape, dtype=xp.complex64)
|
||||
self._test_mtsame(fft.hfft, a, xp=xp)
|
||||
|
||||
def test_ihfft(self, xp):
|
||||
a = xp.ones(self.input_shape)
|
||||
self._test_mtsame(fft.ihfft, a, xp=xp)
|
||||
|
||||
|
||||
@skip_xp_backends(np_only=True)
|
||||
@pytest.mark.parametrize("func", [fft.fft, fft.ifft, fft.rfft, fft.irfft])
|
||||
def test_multiprocess(func):
|
||||
# Test that fft still works after fork (gh-10422)
|
||||
|
||||
with multiprocessing.Pool(2) as p:
|
||||
res = p.map(func, [np.ones(100) for _ in range(4)])
|
||||
|
||||
expect = func(np.ones(100))
|
||||
for x in res:
|
||||
assert_allclose(x, expect)
|
||||
|
||||
|
||||
class TestIRFFTN:
|
||||
|
||||
def test_not_last_axis_success(self, xp):
|
||||
ar, ai = np.random.random((2, 16, 8, 32))
|
||||
a = ar + 1j*ai
|
||||
a = xp.asarray(a)
|
||||
|
||||
axes = (-2,)
|
||||
|
||||
# Should not raise error
|
||||
fft.irfftn(a, axes=axes)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("func", [fft.fft, fft.ifft, fft.rfft, fft.irfft,
|
||||
fft.fftn, fft.ifftn,
|
||||
fft.rfftn, fft.irfftn, fft.hfft, fft.ihfft])
|
||||
def test_non_standard_params(func, xp):
|
||||
if func in [fft.rfft, fft.rfftn, fft.ihfft]:
|
||||
dtype = xp.float64
|
||||
else:
|
||||
dtype = xp.complex128
|
||||
|
||||
if xp.__name__ != 'numpy':
|
||||
x = xp.asarray([1, 2, 3], dtype=dtype)
|
||||
# func(x) should not raise an exception
|
||||
func(x)
|
||||
assert_raises(ValueError, func, x, workers=2)
|
||||
# `plan` param is not tested since SciPy does not use it currently
|
||||
# but should be tested if it comes into use
|
||||
179
venv/lib/python3.12/site-packages/scipy/fft/tests/test_fftlog.py
Normal file
179
venv/lib/python3.12/site-packages/scipy/fft/tests/test_fftlog.py
Normal file
@ -0,0 +1,179 @@
|
||||
import warnings
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from scipy.fft._fftlog import fht, ifht, fhtoffset
|
||||
from scipy.special import poch
|
||||
|
||||
from scipy.conftest import array_api_compatible
|
||||
from scipy._lib._array_api import xp_assert_close
|
||||
|
||||
pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends"),]
|
||||
skip_xp_backends = pytest.mark.skip_xp_backends
|
||||
|
||||
|
||||
def test_fht_agrees_with_fftlog(xp):
|
||||
# check that fht numerically agrees with the output from Fortran FFTLog,
|
||||
# the results were generated with the provided `fftlogtest` program,
|
||||
# after fixing how the k array is generated (divide range by n-1, not n)
|
||||
|
||||
# test function, analytical Hankel transform is of the same form
|
||||
def f(r, mu):
|
||||
return r**(mu+1)*np.exp(-r**2/2)
|
||||
|
||||
r = np.logspace(-4, 4, 16)
|
||||
|
||||
dln = np.log(r[1]/r[0])
|
||||
mu = 0.3
|
||||
offset = 0.0
|
||||
bias = 0.0
|
||||
|
||||
a = xp.asarray(f(r, mu))
|
||||
|
||||
# test 1: compute as given
|
||||
ours = fht(a, dln, mu, offset=offset, bias=bias)
|
||||
theirs = [-0.1159922613593045E-02, +0.1625822618458832E-02,
|
||||
-0.1949518286432330E-02, +0.3789220182554077E-02,
|
||||
+0.5093959119952945E-03, +0.2785387803618774E-01,
|
||||
+0.9944952700848897E-01, +0.4599202164586588E+00,
|
||||
+0.3157462160881342E+00, -0.8201236844404755E-03,
|
||||
-0.7834031308271878E-03, +0.3931444945110708E-03,
|
||||
-0.2697710625194777E-03, +0.3568398050238820E-03,
|
||||
-0.5554454827797206E-03, +0.8286331026468585E-03]
|
||||
theirs = xp.asarray(theirs, dtype=xp.float64)
|
||||
xp_assert_close(ours, theirs)
|
||||
|
||||
# test 2: change to optimal offset
|
||||
offset = fhtoffset(dln, mu, bias=bias)
|
||||
ours = fht(a, dln, mu, offset=offset, bias=bias)
|
||||
theirs = [+0.4353768523152057E-04, -0.9197045663594285E-05,
|
||||
+0.3150140927838524E-03, +0.9149121960963704E-03,
|
||||
+0.5808089753959363E-02, +0.2548065256377240E-01,
|
||||
+0.1339477692089897E+00, +0.4821530509479356E+00,
|
||||
+0.2659899781579785E+00, -0.1116475278448113E-01,
|
||||
+0.1791441617592385E-02, -0.4181810476548056E-03,
|
||||
+0.1314963536765343E-03, -0.5422057743066297E-04,
|
||||
+0.3208681804170443E-04, -0.2696849476008234E-04]
|
||||
theirs = xp.asarray(theirs, dtype=xp.float64)
|
||||
xp_assert_close(ours, theirs)
|
||||
|
||||
# test 3: positive bias
|
||||
bias = 0.8
|
||||
offset = fhtoffset(dln, mu, bias=bias)
|
||||
ours = fht(a, dln, mu, offset=offset, bias=bias)
|
||||
theirs = [-7.3436673558316850E+00, +0.1710271207817100E+00,
|
||||
+0.1065374386206564E+00, -0.5121739602708132E-01,
|
||||
+0.2636649319269470E-01, +0.1697209218849693E-01,
|
||||
+0.1250215614723183E+00, +0.4739583261486729E+00,
|
||||
+0.2841149874912028E+00, -0.8312764741645729E-02,
|
||||
+0.1024233505508988E-02, -0.1644902767389120E-03,
|
||||
+0.3305775476926270E-04, -0.7786993194882709E-05,
|
||||
+0.1962258449520547E-05, -0.8977895734909250E-06]
|
||||
theirs = xp.asarray(theirs, dtype=xp.float64)
|
||||
xp_assert_close(ours, theirs)
|
||||
|
||||
# test 4: negative bias
|
||||
bias = -0.8
|
||||
offset = fhtoffset(dln, mu, bias=bias)
|
||||
ours = fht(a, dln, mu, offset=offset, bias=bias)
|
||||
theirs = [+0.8985777068568745E-05, +0.4074898209936099E-04,
|
||||
+0.2123969254700955E-03, +0.1009558244834628E-02,
|
||||
+0.5131386375222176E-02, +0.2461678673516286E-01,
|
||||
+0.1235812845384476E+00, +0.4719570096404403E+00,
|
||||
+0.2893487490631317E+00, -0.1686570611318716E-01,
|
||||
+0.2231398155172505E-01, -0.1480742256379873E-01,
|
||||
+0.1692387813500801E+00, +0.3097490354365797E+00,
|
||||
+2.7593607182401860E+00, 10.5251075070045800E+00]
|
||||
theirs = xp.asarray(theirs, dtype=xp.float64)
|
||||
xp_assert_close(ours, theirs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('optimal', [True, False])
|
||||
@pytest.mark.parametrize('offset', [0.0, 1.0, -1.0])
|
||||
@pytest.mark.parametrize('bias', [0, 0.1, -0.1])
|
||||
@pytest.mark.parametrize('n', [64, 63])
|
||||
def test_fht_identity(n, bias, offset, optimal, xp):
|
||||
rng = np.random.RandomState(3491349965)
|
||||
|
||||
a = xp.asarray(rng.standard_normal(n))
|
||||
dln = rng.uniform(-1, 1)
|
||||
mu = rng.uniform(-2, 2)
|
||||
|
||||
if optimal:
|
||||
offset = fhtoffset(dln, mu, initial=offset, bias=bias)
|
||||
|
||||
A = fht(a, dln, mu, offset=offset, bias=bias)
|
||||
a_ = ifht(A, dln, mu, offset=offset, bias=bias)
|
||||
|
||||
xp_assert_close(a_, a, rtol=1.5e-7)
|
||||
|
||||
|
||||
def test_fht_special_cases(xp):
|
||||
rng = np.random.RandomState(3491349965)
|
||||
|
||||
a = xp.asarray(rng.standard_normal(64))
|
||||
dln = rng.uniform(-1, 1)
|
||||
|
||||
# let x = (mu+1+q)/2, y = (mu+1-q)/2, M = {0, -1, -2, ...}
|
||||
|
||||
# case 1: x in M, y in M => well-defined transform
|
||||
mu, bias = -4.0, 1.0
|
||||
with warnings.catch_warnings(record=True) as record:
|
||||
fht(a, dln, mu, bias=bias)
|
||||
assert not record, 'fht warned about a well-defined transform'
|
||||
|
||||
# case 2: x not in M, y in M => well-defined transform
|
||||
mu, bias = -2.5, 0.5
|
||||
with warnings.catch_warnings(record=True) as record:
|
||||
fht(a, dln, mu, bias=bias)
|
||||
assert not record, 'fht warned about a well-defined transform'
|
||||
|
||||
# case 3: x in M, y not in M => singular transform
|
||||
mu, bias = -3.5, 0.5
|
||||
with pytest.warns(Warning) as record:
|
||||
fht(a, dln, mu, bias=bias)
|
||||
assert record, 'fht did not warn about a singular transform'
|
||||
|
||||
# case 4: x not in M, y in M => singular inverse transform
|
||||
mu, bias = -2.5, 0.5
|
||||
with pytest.warns(Warning) as record:
|
||||
ifht(a, dln, mu, bias=bias)
|
||||
assert record, 'ifht did not warn about a singular transform'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('n', [64, 63])
|
||||
def test_fht_exact(n, xp):
|
||||
rng = np.random.RandomState(3491349965)
|
||||
|
||||
# for a(r) a power law r^\gamma, the fast Hankel transform produces the
|
||||
# exact continuous Hankel transform if biased with q = \gamma
|
||||
|
||||
mu = rng.uniform(0, 3)
|
||||
|
||||
# convergence of HT: -1-mu < gamma < 1/2
|
||||
gamma = rng.uniform(-1-mu, 1/2)
|
||||
|
||||
r = np.logspace(-2, 2, n)
|
||||
a = xp.asarray(r**gamma)
|
||||
|
||||
dln = np.log(r[1]/r[0])
|
||||
|
||||
offset = fhtoffset(dln, mu, initial=0.0, bias=gamma)
|
||||
|
||||
A = fht(a, dln, mu, offset=offset, bias=gamma)
|
||||
|
||||
k = np.exp(offset)/r[::-1]
|
||||
|
||||
# analytical result
|
||||
At = xp.asarray((2/k)**gamma * poch((mu+1-gamma)/2, gamma))
|
||||
|
||||
xp_assert_close(A, At)
|
||||
|
||||
@skip_xp_backends(np_only=True,
|
||||
reasons=['array-likes only supported for NumPy backend'])
|
||||
@pytest.mark.parametrize("op", [fht, ifht])
|
||||
def test_array_like(xp, op):
|
||||
x = [[[1.0, 1.0], [1.0, 1.0]],
|
||||
[[1.0, 1.0], [1.0, 1.0]],
|
||||
[[1.0, 1.0], [1.0, 1.0]]]
|
||||
xp_assert_close(op(x, 1.0, 2.0), op(xp.asarray(x), 1.0, 2.0))
|
||||
570
venv/lib/python3.12/site-packages/scipy/fft/tests/test_helper.py
Normal file
570
venv/lib/python3.12/site-packages/scipy/fft/tests/test_helper.py
Normal file
@ -0,0 +1,570 @@
|
||||
"""Includes test functions for fftpack.helper module
|
||||
|
||||
Copied from fftpack.helper by Pearu Peterson, October 2005
|
||||
Modified for Array API, 2023
|
||||
|
||||
"""
|
||||
from scipy.fft._helper import next_fast_len, prev_fast_len, _init_nd_shape_and_axes
|
||||
from numpy.testing import assert_equal
|
||||
from pytest import raises as assert_raises
|
||||
import pytest
|
||||
import numpy as np
|
||||
import sys
|
||||
from scipy.conftest import array_api_compatible
|
||||
from scipy._lib._array_api import (
|
||||
xp_assert_close, get_xp_devices, device, array_namespace
|
||||
)
|
||||
from scipy import fft
|
||||
|
||||
pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends")]
|
||||
skip_xp_backends = pytest.mark.skip_xp_backends
|
||||
|
||||
_5_smooth_numbers = [
|
||||
2, 3, 4, 5, 6, 8, 9, 10,
|
||||
2 * 3 * 5,
|
||||
2**3 * 3**5,
|
||||
2**3 * 3**3 * 5**2,
|
||||
]
|
||||
|
||||
def test_next_fast_len():
|
||||
for n in _5_smooth_numbers:
|
||||
assert_equal(next_fast_len(n), n)
|
||||
|
||||
|
||||
def _assert_n_smooth(x, n):
|
||||
x_orig = x
|
||||
if n < 2:
|
||||
assert False
|
||||
|
||||
while True:
|
||||
q, r = divmod(x, 2)
|
||||
if r != 0:
|
||||
break
|
||||
x = q
|
||||
|
||||
for d in range(3, n+1, 2):
|
||||
while True:
|
||||
q, r = divmod(x, d)
|
||||
if r != 0:
|
||||
break
|
||||
x = q
|
||||
|
||||
assert x == 1, \
|
||||
f'x={x_orig} is not {n}-smooth, remainder={x}'
|
||||
|
||||
|
||||
@skip_xp_backends(np_only=True)
|
||||
class TestNextFastLen:
|
||||
|
||||
def test_next_fast_len(self):
|
||||
np.random.seed(1234)
|
||||
|
||||
def nums():
|
||||
yield from range(1, 1000)
|
||||
yield 2**5 * 3**5 * 4**5 + 1
|
||||
|
||||
for n in nums():
|
||||
m = next_fast_len(n)
|
||||
_assert_n_smooth(m, 11)
|
||||
assert m == next_fast_len(n, False)
|
||||
|
||||
m = next_fast_len(n, True)
|
||||
_assert_n_smooth(m, 5)
|
||||
|
||||
def test_np_integers(self):
|
||||
ITYPES = [np.int16, np.int32, np.int64, np.uint16, np.uint32, np.uint64]
|
||||
for ityp in ITYPES:
|
||||
x = ityp(12345)
|
||||
testN = next_fast_len(x)
|
||||
assert_equal(testN, next_fast_len(int(x)))
|
||||
|
||||
def testnext_fast_len_small(self):
|
||||
hams = {
|
||||
1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 8, 8: 8, 14: 15, 15: 15,
|
||||
16: 16, 17: 18, 1021: 1024, 1536: 1536, 51200000: 51200000
|
||||
}
|
||||
for x, y in hams.items():
|
||||
assert_equal(next_fast_len(x, True), y)
|
||||
|
||||
@pytest.mark.xfail(sys.maxsize < 2**32,
|
||||
reason="Hamming Numbers too large for 32-bit",
|
||||
raises=ValueError, strict=True)
|
||||
def testnext_fast_len_big(self):
|
||||
hams = {
|
||||
510183360: 510183360, 510183360 + 1: 512000000,
|
||||
511000000: 512000000,
|
||||
854296875: 854296875, 854296875 + 1: 859963392,
|
||||
196608000000: 196608000000, 196608000000 + 1: 196830000000,
|
||||
8789062500000: 8789062500000, 8789062500000 + 1: 8796093022208,
|
||||
206391214080000: 206391214080000,
|
||||
206391214080000 + 1: 206624260800000,
|
||||
470184984576000: 470184984576000,
|
||||
470184984576000 + 1: 470715894135000,
|
||||
7222041363087360: 7222041363087360,
|
||||
7222041363087360 + 1: 7230196133913600,
|
||||
# power of 5 5**23
|
||||
11920928955078125: 11920928955078125,
|
||||
11920928955078125 - 1: 11920928955078125,
|
||||
# power of 3 3**34
|
||||
16677181699666569: 16677181699666569,
|
||||
16677181699666569 - 1: 16677181699666569,
|
||||
# power of 2 2**54
|
||||
18014398509481984: 18014398509481984,
|
||||
18014398509481984 - 1: 18014398509481984,
|
||||
# above this, int(ceil(n)) == int(ceil(n+1))
|
||||
19200000000000000: 19200000000000000,
|
||||
19200000000000000 + 1: 19221679687500000,
|
||||
288230376151711744: 288230376151711744,
|
||||
288230376151711744 + 1: 288325195312500000,
|
||||
288325195312500000 - 1: 288325195312500000,
|
||||
288325195312500000: 288325195312500000,
|
||||
288325195312500000 + 1: 288555831593533440,
|
||||
}
|
||||
for x, y in hams.items():
|
||||
assert_equal(next_fast_len(x, True), y)
|
||||
|
||||
def test_keyword_args(self):
|
||||
assert next_fast_len(11, real=True) == 12
|
||||
assert next_fast_len(target=7, real=False) == 7
|
||||
|
||||
@skip_xp_backends(np_only=True)
|
||||
class TestPrevFastLen:
|
||||
|
||||
def test_prev_fast_len(self):
|
||||
np.random.seed(1234)
|
||||
|
||||
def nums():
|
||||
yield from range(1, 1000)
|
||||
yield 2**5 * 3**5 * 4**5 + 1
|
||||
|
||||
for n in nums():
|
||||
m = prev_fast_len(n)
|
||||
_assert_n_smooth(m, 11)
|
||||
assert m == prev_fast_len(n, False)
|
||||
|
||||
m = prev_fast_len(n, True)
|
||||
_assert_n_smooth(m, 5)
|
||||
|
||||
def test_np_integers(self):
|
||||
ITYPES = [np.int16, np.int32, np.int64, np.uint16, np.uint32,
|
||||
np.uint64]
|
||||
for ityp in ITYPES:
|
||||
x = ityp(12345)
|
||||
testN = prev_fast_len(x)
|
||||
assert_equal(testN, prev_fast_len(int(x)))
|
||||
|
||||
testN = prev_fast_len(x, real=True)
|
||||
assert_equal(testN, prev_fast_len(int(x), real=True))
|
||||
|
||||
def testprev_fast_len_small(self):
|
||||
hams = {
|
||||
1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 6, 8: 8, 14: 12, 15: 15,
|
||||
16: 16, 17: 16, 1021: 1000, 1536: 1536, 51200000: 51200000
|
||||
}
|
||||
for x, y in hams.items():
|
||||
assert_equal(prev_fast_len(x, True), y)
|
||||
|
||||
hams = {
|
||||
1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8, 9: 9, 10: 10,
|
||||
11: 11, 12: 12, 13: 12, 14: 14, 15: 15, 16: 16, 17: 16, 18: 18,
|
||||
19: 18, 20: 20, 21: 21, 22: 22, 120: 120, 121: 121, 122: 121,
|
||||
1021: 1008, 1536: 1536, 51200000: 51200000
|
||||
}
|
||||
for x, y in hams.items():
|
||||
assert_equal(prev_fast_len(x, False), y)
|
||||
|
||||
@pytest.mark.xfail(sys.maxsize < 2**32,
|
||||
reason="Hamming Numbers too large for 32-bit",
|
||||
raises=ValueError, strict=True)
|
||||
def testprev_fast_len_big(self):
|
||||
hams = {
|
||||
# 2**6 * 3**13 * 5**1
|
||||
510183360: 510183360,
|
||||
510183360 + 1: 510183360,
|
||||
510183360 - 1: 509607936, # 2**21 * 3**5
|
||||
# 2**6 * 5**6 * 7**1 * 73**1
|
||||
511000000: 510183360,
|
||||
511000000 + 1: 510183360,
|
||||
511000000 - 1: 510183360, # 2**6 * 3**13 * 5**1
|
||||
# 3**7 * 5**8
|
||||
854296875: 854296875,
|
||||
854296875 + 1: 854296875,
|
||||
854296875 - 1: 850305600, # 2**6 * 3**12 * 5**2
|
||||
# 2**22 * 3**1 * 5**6
|
||||
196608000000: 196608000000,
|
||||
196608000000 + 1: 196608000000,
|
||||
196608000000 - 1: 195910410240, # 2**13 * 3**14 * 5**1
|
||||
# 2**5 * 3**2 * 5**15
|
||||
8789062500000: 8789062500000,
|
||||
8789062500000 + 1: 8789062500000,
|
||||
8789062500000 - 1: 8748000000000, # 2**11 * 3**7 * 5**9
|
||||
# 2**24 * 3**9 * 5**4
|
||||
206391214080000: 206391214080000,
|
||||
206391214080000 + 1: 206391214080000,
|
||||
206391214080000 - 1: 206158430208000, # 2**39 * 3**1 * 5**3
|
||||
# 2**18 * 3**15 * 5**3
|
||||
470184984576000: 470184984576000,
|
||||
470184984576000 + 1: 470184984576000,
|
||||
470184984576000 - 1: 469654673817600, # 2**33 * 3**7 **5**2
|
||||
# 2**25 * 3**16 * 5**1
|
||||
7222041363087360: 7222041363087360,
|
||||
7222041363087360 + 1: 7222041363087360,
|
||||
7222041363087360 - 1: 7213895789838336, # 2**40 * 3**8
|
||||
# power of 5 5**23
|
||||
11920928955078125: 11920928955078125,
|
||||
11920928955078125 + 1: 11920928955078125,
|
||||
11920928955078125 - 1: 11901557422080000, # 2**14 * 3**19 * 5**4
|
||||
# power of 3 3**34
|
||||
16677181699666569: 16677181699666569,
|
||||
16677181699666569 + 1: 16677181699666569,
|
||||
16677181699666569 - 1: 16607531250000000, # 2**7 * 3**12 * 5**12
|
||||
# power of 2 2**54
|
||||
18014398509481984: 18014398509481984,
|
||||
18014398509481984 + 1: 18014398509481984,
|
||||
18014398509481984 - 1: 18000000000000000, # 2**16 * 3**2 * 5**15
|
||||
# 2**20 * 3**1 * 5**14
|
||||
19200000000000000: 19200000000000000,
|
||||
19200000000000000 + 1: 19200000000000000,
|
||||
19200000000000000 - 1: 19131876000000000, # 2**11 * 3**14 * 5**9
|
||||
# 2**58
|
||||
288230376151711744: 288230376151711744,
|
||||
288230376151711744 + 1: 288230376151711744,
|
||||
288230376151711744 - 1: 288000000000000000, # 2**20 * 3**2 * 5**15
|
||||
# 2**5 * 3**10 * 5**16
|
||||
288325195312500000: 288325195312500000,
|
||||
288325195312500000 + 1: 288325195312500000,
|
||||
288325195312500000 - 1: 288230376151711744, # 2**58
|
||||
}
|
||||
for x, y in hams.items():
|
||||
assert_equal(prev_fast_len(x, True), y)
|
||||
|
||||
def test_keyword_args(self):
|
||||
assert prev_fast_len(11, real=True) == 10
|
||||
assert prev_fast_len(target=7, real=False) == 7
|
||||
|
||||
|
||||
@skip_xp_backends(cpu_only=True)
|
||||
class Test_init_nd_shape_and_axes:
|
||||
|
||||
def test_py_0d_defaults(self, xp):
|
||||
x = xp.asarray(4)
|
||||
shape = None
|
||||
axes = None
|
||||
|
||||
shape_expected = ()
|
||||
axes_expected = []
|
||||
|
||||
shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes)
|
||||
|
||||
assert shape_res == shape_expected
|
||||
assert axes_res == axes_expected
|
||||
|
||||
def test_xp_0d_defaults(self, xp):
|
||||
x = xp.asarray(7.)
|
||||
shape = None
|
||||
axes = None
|
||||
|
||||
shape_expected = ()
|
||||
axes_expected = []
|
||||
|
||||
shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes)
|
||||
|
||||
assert shape_res == shape_expected
|
||||
assert axes_res == axes_expected
|
||||
|
||||
def test_py_1d_defaults(self, xp):
|
||||
x = xp.asarray([1, 2, 3])
|
||||
shape = None
|
||||
axes = None
|
||||
|
||||
shape_expected = (3,)
|
||||
axes_expected = [0]
|
||||
|
||||
shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes)
|
||||
|
||||
assert shape_res == shape_expected
|
||||
assert axes_res == axes_expected
|
||||
|
||||
def test_xp_1d_defaults(self, xp):
|
||||
x = xp.arange(0, 1, .1)
|
||||
shape = None
|
||||
axes = None
|
||||
|
||||
shape_expected = (10,)
|
||||
axes_expected = [0]
|
||||
|
||||
shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes)
|
||||
|
||||
assert shape_res == shape_expected
|
||||
assert axes_res == axes_expected
|
||||
|
||||
def test_py_2d_defaults(self, xp):
|
||||
x = xp.asarray([[1, 2, 3, 4],
|
||||
[5, 6, 7, 8]])
|
||||
shape = None
|
||||
axes = None
|
||||
|
||||
shape_expected = (2, 4)
|
||||
axes_expected = [0, 1]
|
||||
|
||||
shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes)
|
||||
|
||||
assert shape_res == shape_expected
|
||||
assert axes_res == axes_expected
|
||||
|
||||
def test_xp_2d_defaults(self, xp):
|
||||
x = xp.arange(0, 1, .1)
|
||||
x = xp.reshape(x, (5, 2))
|
||||
shape = None
|
||||
axes = None
|
||||
|
||||
shape_expected = (5, 2)
|
||||
axes_expected = [0, 1]
|
||||
|
||||
shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes)
|
||||
|
||||
assert shape_res == shape_expected
|
||||
assert axes_res == axes_expected
|
||||
|
||||
def test_xp_5d_defaults(self, xp):
|
||||
x = xp.zeros([6, 2, 5, 3, 4])
|
||||
shape = None
|
||||
axes = None
|
||||
|
||||
shape_expected = (6, 2, 5, 3, 4)
|
||||
axes_expected = [0, 1, 2, 3, 4]
|
||||
|
||||
shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes)
|
||||
|
||||
assert shape_res == shape_expected
|
||||
assert axes_res == axes_expected
|
||||
|
||||
def test_xp_5d_set_shape(self, xp):
|
||||
x = xp.zeros([6, 2, 5, 3, 4])
|
||||
shape = [10, -1, -1, 1, 4]
|
||||
axes = None
|
||||
|
||||
shape_expected = (10, 2, 5, 1, 4)
|
||||
axes_expected = [0, 1, 2, 3, 4]
|
||||
|
||||
shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes)
|
||||
|
||||
assert shape_res == shape_expected
|
||||
assert axes_res == axes_expected
|
||||
|
||||
def test_xp_5d_set_axes(self, xp):
|
||||
x = xp.zeros([6, 2, 5, 3, 4])
|
||||
shape = None
|
||||
axes = [4, 1, 2]
|
||||
|
||||
shape_expected = (4, 2, 5)
|
||||
axes_expected = [4, 1, 2]
|
||||
|
||||
shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes)
|
||||
|
||||
assert shape_res == shape_expected
|
||||
assert axes_res == axes_expected
|
||||
|
||||
def test_xp_5d_set_shape_axes(self, xp):
|
||||
x = xp.zeros([6, 2, 5, 3, 4])
|
||||
shape = [10, -1, 2]
|
||||
axes = [1, 0, 3]
|
||||
|
||||
shape_expected = (10, 6, 2)
|
||||
axes_expected = [1, 0, 3]
|
||||
|
||||
shape_res, axes_res = _init_nd_shape_and_axes(x, shape, axes)
|
||||
|
||||
assert shape_res == shape_expected
|
||||
assert axes_res == axes_expected
|
||||
|
||||
def test_shape_axes_subset(self, xp):
|
||||
x = xp.zeros((2, 3, 4, 5))
|
||||
shape, axes = _init_nd_shape_and_axes(x, shape=(5, 5, 5), axes=None)
|
||||
|
||||
assert shape == (5, 5, 5)
|
||||
assert axes == [1, 2, 3]
|
||||
|
||||
def test_errors(self, xp):
|
||||
x = xp.zeros(1)
|
||||
with assert_raises(ValueError, match="axes must be a scalar or "
|
||||
"iterable of integers"):
|
||||
_init_nd_shape_and_axes(x, shape=None, axes=[[1, 2], [3, 4]])
|
||||
|
||||
with assert_raises(ValueError, match="axes must be a scalar or "
|
||||
"iterable of integers"):
|
||||
_init_nd_shape_and_axes(x, shape=None, axes=[1., 2., 3., 4.])
|
||||
|
||||
with assert_raises(ValueError,
|
||||
match="axes exceeds dimensionality of input"):
|
||||
_init_nd_shape_and_axes(x, shape=None, axes=[1])
|
||||
|
||||
with assert_raises(ValueError,
|
||||
match="axes exceeds dimensionality of input"):
|
||||
_init_nd_shape_and_axes(x, shape=None, axes=[-2])
|
||||
|
||||
with assert_raises(ValueError,
|
||||
match="all axes must be unique"):
|
||||
_init_nd_shape_and_axes(x, shape=None, axes=[0, 0])
|
||||
|
||||
with assert_raises(ValueError, match="shape must be a scalar or "
|
||||
"iterable of integers"):
|
||||
_init_nd_shape_and_axes(x, shape=[[1, 2], [3, 4]], axes=None)
|
||||
|
||||
with assert_raises(ValueError, match="shape must be a scalar or "
|
||||
"iterable of integers"):
|
||||
_init_nd_shape_and_axes(x, shape=[1., 2., 3., 4.], axes=None)
|
||||
|
||||
with assert_raises(ValueError,
|
||||
match="when given, axes and shape arguments"
|
||||
" have to be of the same length"):
|
||||
_init_nd_shape_and_axes(xp.zeros([1, 1, 1, 1]),
|
||||
shape=[1, 2, 3], axes=[1])
|
||||
|
||||
with assert_raises(ValueError,
|
||||
match="invalid number of data points"
|
||||
r" \(\[0\]\) specified"):
|
||||
_init_nd_shape_and_axes(x, shape=[0], axes=None)
|
||||
|
||||
with assert_raises(ValueError,
|
||||
match="invalid number of data points"
|
||||
r" \(\[-2\]\) specified"):
|
||||
_init_nd_shape_and_axes(x, shape=-2, axes=None)
|
||||
|
||||
|
||||
class TestFFTShift:
|
||||
|
||||
def test_definition(self, xp):
|
||||
x = xp.asarray([0., 1, 2, 3, 4, -4, -3, -2, -1])
|
||||
y = xp.asarray([-4., -3, -2, -1, 0, 1, 2, 3, 4])
|
||||
xp_assert_close(fft.fftshift(x), y)
|
||||
xp_assert_close(fft.ifftshift(y), x)
|
||||
x = xp.asarray([0., 1, 2, 3, 4, -5, -4, -3, -2, -1])
|
||||
y = xp.asarray([-5., -4, -3, -2, -1, 0, 1, 2, 3, 4])
|
||||
xp_assert_close(fft.fftshift(x), y)
|
||||
xp_assert_close(fft.ifftshift(y), x)
|
||||
|
||||
def test_inverse(self, xp):
|
||||
for n in [1, 4, 9, 100, 211]:
|
||||
x = xp.asarray(np.random.random((n,)))
|
||||
xp_assert_close(fft.ifftshift(fft.fftshift(x)), x)
|
||||
|
||||
def test_axes_keyword(self, xp):
|
||||
freqs = xp.asarray([[0., 1, 2], [3, 4, -4], [-3, -2, -1]])
|
||||
shifted = xp.asarray([[-1., -3, -2], [2, 0, 1], [-4, 3, 4]])
|
||||
xp_assert_close(fft.fftshift(freqs, axes=(0, 1)), shifted)
|
||||
xp_assert_close(fft.fftshift(freqs, axes=0), fft.fftshift(freqs, axes=(0,)))
|
||||
xp_assert_close(fft.ifftshift(shifted, axes=(0, 1)), freqs)
|
||||
xp_assert_close(fft.ifftshift(shifted, axes=0),
|
||||
fft.ifftshift(shifted, axes=(0,)))
|
||||
xp_assert_close(fft.fftshift(freqs), shifted)
|
||||
xp_assert_close(fft.ifftshift(shifted), freqs)
|
||||
|
||||
def test_uneven_dims(self, xp):
|
||||
""" Test 2D input, which has uneven dimension sizes """
|
||||
freqs = xp.asarray([
|
||||
[0, 1],
|
||||
[2, 3],
|
||||
[4, 5]
|
||||
], dtype=xp.float64)
|
||||
|
||||
# shift in dimension 0
|
||||
shift_dim0 = xp.asarray([
|
||||
[4, 5],
|
||||
[0, 1],
|
||||
[2, 3]
|
||||
], dtype=xp.float64)
|
||||
xp_assert_close(fft.fftshift(freqs, axes=0), shift_dim0)
|
||||
xp_assert_close(fft.ifftshift(shift_dim0, axes=0), freqs)
|
||||
xp_assert_close(fft.fftshift(freqs, axes=(0,)), shift_dim0)
|
||||
xp_assert_close(fft.ifftshift(shift_dim0, axes=[0]), freqs)
|
||||
|
||||
# shift in dimension 1
|
||||
shift_dim1 = xp.asarray([
|
||||
[1, 0],
|
||||
[3, 2],
|
||||
[5, 4]
|
||||
], dtype=xp.float64)
|
||||
xp_assert_close(fft.fftshift(freqs, axes=1), shift_dim1)
|
||||
xp_assert_close(fft.ifftshift(shift_dim1, axes=1), freqs)
|
||||
|
||||
# shift in both dimensions
|
||||
shift_dim_both = xp.asarray([
|
||||
[5, 4],
|
||||
[1, 0],
|
||||
[3, 2]
|
||||
], dtype=xp.float64)
|
||||
xp_assert_close(fft.fftshift(freqs, axes=(0, 1)), shift_dim_both)
|
||||
xp_assert_close(fft.ifftshift(shift_dim_both, axes=(0, 1)), freqs)
|
||||
xp_assert_close(fft.fftshift(freqs, axes=[0, 1]), shift_dim_both)
|
||||
xp_assert_close(fft.ifftshift(shift_dim_both, axes=[0, 1]), freqs)
|
||||
|
||||
# axes=None (default) shift in all dimensions
|
||||
xp_assert_close(fft.fftshift(freqs, axes=None), shift_dim_both)
|
||||
xp_assert_close(fft.ifftshift(shift_dim_both, axes=None), freqs)
|
||||
xp_assert_close(fft.fftshift(freqs), shift_dim_both)
|
||||
xp_assert_close(fft.ifftshift(shift_dim_both), freqs)
|
||||
|
||||
|
||||
@skip_xp_backends("cupy", "jax.numpy",
|
||||
reasons=["CuPy has not implemented the `device` param",
|
||||
"JAX has not implemented the `device` param"])
|
||||
class TestFFTFreq:
|
||||
|
||||
def test_definition(self, xp):
|
||||
x = xp.asarray([0, 1, 2, 3, 4, -4, -3, -2, -1], dtype=xp.float64)
|
||||
x2 = xp.asarray([0, 1, 2, 3, 4, -5, -4, -3, -2, -1], dtype=xp.float64)
|
||||
|
||||
# default dtype varies across backends
|
||||
|
||||
y = 9 * fft.fftfreq(9, xp=xp)
|
||||
xp_assert_close(y, x, check_dtype=False, check_namespace=True)
|
||||
|
||||
y = 9 * xp.pi * fft.fftfreq(9, xp.pi, xp=xp)
|
||||
xp_assert_close(y, x, check_dtype=False)
|
||||
|
||||
y = 10 * fft.fftfreq(10, xp=xp)
|
||||
xp_assert_close(y, x2, check_dtype=False)
|
||||
|
||||
y = 10 * xp.pi * fft.fftfreq(10, xp.pi, xp=xp)
|
||||
xp_assert_close(y, x2, check_dtype=False)
|
||||
|
||||
def test_device(self, xp):
|
||||
xp_test = array_namespace(xp.empty(0))
|
||||
devices = get_xp_devices(xp)
|
||||
for d in devices:
|
||||
y = fft.fftfreq(9, xp=xp, device=d)
|
||||
x = xp_test.empty(0, device=d)
|
||||
assert device(y) == device(x)
|
||||
|
||||
|
||||
@skip_xp_backends("cupy", "jax.numpy",
|
||||
reasons=["CuPy has not implemented the `device` param",
|
||||
"JAX has not implemented the `device` param"])
|
||||
class TestRFFTFreq:
|
||||
|
||||
def test_definition(self, xp):
|
||||
x = xp.asarray([0, 1, 2, 3, 4], dtype=xp.float64)
|
||||
x2 = xp.asarray([0, 1, 2, 3, 4, 5], dtype=xp.float64)
|
||||
|
||||
# default dtype varies across backends
|
||||
|
||||
y = 9 * fft.rfftfreq(9, xp=xp)
|
||||
xp_assert_close(y, x, check_dtype=False, check_namespace=True)
|
||||
|
||||
y = 9 * xp.pi * fft.rfftfreq(9, xp.pi, xp=xp)
|
||||
xp_assert_close(y, x, check_dtype=False)
|
||||
|
||||
y = 10 * fft.rfftfreq(10, xp=xp)
|
||||
xp_assert_close(y, x2, check_dtype=False)
|
||||
|
||||
y = 10 * xp.pi * fft.rfftfreq(10, xp.pi, xp=xp)
|
||||
xp_assert_close(y, x2, check_dtype=False)
|
||||
|
||||
def test_device(self, xp):
|
||||
xp_test = array_namespace(xp.empty(0))
|
||||
devices = get_xp_devices(xp)
|
||||
for d in devices:
|
||||
y = fft.rfftfreq(9, xp=xp, device=d)
|
||||
x = xp_test.empty(0, device=d)
|
||||
assert device(y) == device(x)
|
||||
@ -0,0 +1,84 @@
|
||||
from scipy import fft
|
||||
import numpy as np
|
||||
import pytest
|
||||
from numpy.testing import assert_allclose
|
||||
import multiprocessing
|
||||
import os
|
||||
|
||||
|
||||
@pytest.fixture(scope='module')
|
||||
def x():
|
||||
return np.random.randn(512, 128) # Must be large enough to qualify for mt
|
||||
|
||||
|
||||
@pytest.mark.parametrize("func", [
|
||||
fft.fft, fft.ifft, fft.fft2, fft.ifft2, fft.fftn, fft.ifftn,
|
||||
fft.rfft, fft.irfft, fft.rfft2, fft.irfft2, fft.rfftn, fft.irfftn,
|
||||
fft.hfft, fft.ihfft, fft.hfft2, fft.ihfft2, fft.hfftn, fft.ihfftn,
|
||||
fft.dct, fft.idct, fft.dctn, fft.idctn,
|
||||
fft.dst, fft.idst, fft.dstn, fft.idstn,
|
||||
])
|
||||
@pytest.mark.parametrize("workers", [2, -1])
|
||||
def test_threaded_same(x, func, workers):
|
||||
expected = func(x, workers=1)
|
||||
actual = func(x, workers=workers)
|
||||
assert_allclose(actual, expected)
|
||||
|
||||
|
||||
def _mt_fft(x):
|
||||
return fft.fft(x, workers=2)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_mixed_threads_processes(x):
|
||||
# Test that the fft threadpool is safe to use before & after fork
|
||||
|
||||
expect = fft.fft(x, workers=2)
|
||||
|
||||
with multiprocessing.Pool(2) as p:
|
||||
res = p.map(_mt_fft, [x for _ in range(4)])
|
||||
|
||||
for r in res:
|
||||
assert_allclose(r, expect)
|
||||
|
||||
fft.fft(x, workers=2)
|
||||
|
||||
|
||||
def test_invalid_workers(x):
|
||||
cpus = os.cpu_count()
|
||||
|
||||
fft.ifft([1], workers=-cpus)
|
||||
|
||||
with pytest.raises(ValueError, match='workers must not be zero'):
|
||||
fft.fft(x, workers=0)
|
||||
|
||||
with pytest.raises(ValueError, match='workers value out of range'):
|
||||
fft.ifft(x, workers=-cpus-1)
|
||||
|
||||
|
||||
def test_set_get_workers():
|
||||
cpus = os.cpu_count()
|
||||
assert fft.get_workers() == 1
|
||||
with fft.set_workers(4):
|
||||
assert fft.get_workers() == 4
|
||||
|
||||
with fft.set_workers(-1):
|
||||
assert fft.get_workers() == cpus
|
||||
|
||||
assert fft.get_workers() == 4
|
||||
|
||||
assert fft.get_workers() == 1
|
||||
|
||||
with fft.set_workers(-cpus):
|
||||
assert fft.get_workers() == 1
|
||||
|
||||
|
||||
def test_set_workers_invalid():
|
||||
|
||||
with pytest.raises(ValueError, match='workers must not be zero'):
|
||||
with fft.set_workers(0):
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError, match='workers value out of range'):
|
||||
with fft.set_workers(-os.cpu_count()-1):
|
||||
pass
|
||||
@ -0,0 +1,249 @@
|
||||
import numpy as np
|
||||
from numpy.testing import assert_allclose, assert_array_equal
|
||||
import pytest
|
||||
import math
|
||||
|
||||
from scipy.fft import dct, idct, dctn, idctn, dst, idst, dstn, idstn
|
||||
import scipy.fft as fft
|
||||
from scipy import fftpack
|
||||
from scipy.conftest import array_api_compatible
|
||||
from scipy._lib._array_api import copy, xp_assert_close
|
||||
|
||||
pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends")]
|
||||
skip_xp_backends = pytest.mark.skip_xp_backends
|
||||
|
||||
SQRT_2 = math.sqrt(2)
|
||||
|
||||
# scipy.fft wraps the fftpack versions but with normalized inverse transforms.
|
||||
# So, the forward transforms and definitions are already thoroughly tested in
|
||||
# fftpack/test_real_transforms.py
|
||||
|
||||
|
||||
@skip_xp_backends(cpu_only=True)
|
||||
@pytest.mark.parametrize("forward, backward", [(dct, idct), (dst, idst)])
|
||||
@pytest.mark.parametrize("type", [1, 2, 3, 4])
|
||||
@pytest.mark.parametrize("n", [2, 3, 4, 5, 10, 16])
|
||||
@pytest.mark.parametrize("axis", [0, 1])
|
||||
@pytest.mark.parametrize("norm", [None, 'backward', 'ortho', 'forward'])
|
||||
@pytest.mark.parametrize("orthogonalize", [False, True])
|
||||
def test_identity_1d(forward, backward, type, n, axis, norm, orthogonalize, xp):
|
||||
# Test the identity f^-1(f(x)) == x
|
||||
x = xp.asarray(np.random.rand(n, n))
|
||||
|
||||
y = forward(x, type, axis=axis, norm=norm, orthogonalize=orthogonalize)
|
||||
z = backward(y, type, axis=axis, norm=norm, orthogonalize=orthogonalize)
|
||||
xp_assert_close(z, x)
|
||||
|
||||
pad = [(0, 0)] * 2
|
||||
pad[axis] = (0, 4)
|
||||
|
||||
y2 = xp.asarray(np.pad(np.asarray(y), pad, mode='edge'))
|
||||
z2 = backward(y2, type, n, axis, norm, orthogonalize=orthogonalize)
|
||||
xp_assert_close(z2, x)
|
||||
|
||||
|
||||
@skip_xp_backends(np_only=True,
|
||||
reasons=['`overwrite_x` only supported for NumPy backend.'])
|
||||
@pytest.mark.parametrize("forward, backward", [(dct, idct), (dst, idst)])
|
||||
@pytest.mark.parametrize("type", [1, 2, 3, 4])
|
||||
@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64,
|
||||
np.complex64, np.complex128])
|
||||
@pytest.mark.parametrize("axis", [0, 1])
|
||||
@pytest.mark.parametrize("norm", [None, 'backward', 'ortho', 'forward'])
|
||||
@pytest.mark.parametrize("overwrite_x", [True, False])
|
||||
def test_identity_1d_overwrite(forward, backward, type, dtype, axis, norm,
|
||||
overwrite_x):
|
||||
# Test the identity f^-1(f(x)) == x
|
||||
x = np.random.rand(7, 8).astype(dtype)
|
||||
x_orig = x.copy()
|
||||
|
||||
y = forward(x, type, axis=axis, norm=norm, overwrite_x=overwrite_x)
|
||||
y_orig = y.copy()
|
||||
z = backward(y, type, axis=axis, norm=norm, overwrite_x=overwrite_x)
|
||||
if not overwrite_x:
|
||||
assert_allclose(z, x, rtol=1e-6, atol=1e-6)
|
||||
assert_array_equal(x, x_orig)
|
||||
assert_array_equal(y, y_orig)
|
||||
else:
|
||||
assert_allclose(z, x_orig, rtol=1e-6, atol=1e-6)
|
||||
|
||||
|
||||
@skip_xp_backends(cpu_only=True)
|
||||
@pytest.mark.parametrize("forward, backward", [(dctn, idctn), (dstn, idstn)])
|
||||
@pytest.mark.parametrize("type", [1, 2, 3, 4])
|
||||
@pytest.mark.parametrize("shape, axes",
|
||||
[
|
||||
((4, 4), 0),
|
||||
((4, 4), 1),
|
||||
((4, 4), None),
|
||||
((4, 4), (0, 1)),
|
||||
((10, 12), None),
|
||||
((10, 12), (0, 1)),
|
||||
((4, 5, 6), None),
|
||||
((4, 5, 6), 1),
|
||||
((4, 5, 6), (0, 2)),
|
||||
])
|
||||
@pytest.mark.parametrize("norm", [None, 'backward', 'ortho', 'forward'])
|
||||
@pytest.mark.parametrize("orthogonalize", [False, True])
|
||||
def test_identity_nd(forward, backward, type, shape, axes, norm,
|
||||
orthogonalize, xp):
|
||||
# Test the identity f^-1(f(x)) == x
|
||||
|
||||
x = xp.asarray(np.random.random(shape))
|
||||
|
||||
if axes is not None:
|
||||
shape = np.take(shape, axes)
|
||||
|
||||
y = forward(x, type, axes=axes, norm=norm, orthogonalize=orthogonalize)
|
||||
z = backward(y, type, axes=axes, norm=norm, orthogonalize=orthogonalize)
|
||||
xp_assert_close(z, x)
|
||||
|
||||
if axes is None:
|
||||
pad = [(0, 4)] * x.ndim
|
||||
elif isinstance(axes, int):
|
||||
pad = [(0, 0)] * x.ndim
|
||||
pad[axes] = (0, 4)
|
||||
else:
|
||||
pad = [(0, 0)] * x.ndim
|
||||
|
||||
for a in axes:
|
||||
pad[a] = (0, 4)
|
||||
|
||||
# TODO write an array-agnostic pad
|
||||
y2 = xp.asarray(np.pad(np.asarray(y), pad, mode='edge'))
|
||||
z2 = backward(y2, type, shape, axes, norm, orthogonalize=orthogonalize)
|
||||
xp_assert_close(z2, x)
|
||||
|
||||
|
||||
@skip_xp_backends(np_only=True,
|
||||
reasons=['`overwrite_x` only supported for NumPy backend.'])
|
||||
@pytest.mark.parametrize("forward, backward", [(dctn, idctn), (dstn, idstn)])
|
||||
@pytest.mark.parametrize("type", [1, 2, 3, 4])
|
||||
@pytest.mark.parametrize("shape, axes",
|
||||
[
|
||||
((4, 5), 0),
|
||||
((4, 5), 1),
|
||||
((4, 5), None),
|
||||
])
|
||||
@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64,
|
||||
np.complex64, np.complex128])
|
||||
@pytest.mark.parametrize("norm", [None, 'backward', 'ortho', 'forward'])
|
||||
@pytest.mark.parametrize("overwrite_x", [False, True])
|
||||
def test_identity_nd_overwrite(forward, backward, type, shape, axes, dtype,
|
||||
norm, overwrite_x):
|
||||
# Test the identity f^-1(f(x)) == x
|
||||
|
||||
x = np.random.random(shape).astype(dtype)
|
||||
x_orig = x.copy()
|
||||
|
||||
if axes is not None:
|
||||
shape = np.take(shape, axes)
|
||||
|
||||
y = forward(x, type, axes=axes, norm=norm)
|
||||
y_orig = y.copy()
|
||||
z = backward(y, type, axes=axes, norm=norm)
|
||||
if overwrite_x:
|
||||
assert_allclose(z, x_orig, rtol=1e-6, atol=1e-6)
|
||||
else:
|
||||
assert_allclose(z, x, rtol=1e-6, atol=1e-6)
|
||||
assert_array_equal(x, x_orig)
|
||||
assert_array_equal(y, y_orig)
|
||||
|
||||
|
||||
@skip_xp_backends(cpu_only=True)
|
||||
@pytest.mark.parametrize("func", ['dct', 'dst', 'dctn', 'dstn'])
|
||||
@pytest.mark.parametrize("type", [1, 2, 3, 4])
|
||||
@pytest.mark.parametrize("norm", [None, 'backward', 'ortho', 'forward'])
|
||||
def test_fftpack_equivalience(func, type, norm, xp):
|
||||
x = np.random.rand(8, 16)
|
||||
fftpack_res = xp.asarray(getattr(fftpack, func)(x, type, norm=norm))
|
||||
x = xp.asarray(x)
|
||||
fft_res = getattr(fft, func)(x, type, norm=norm)
|
||||
|
||||
xp_assert_close(fft_res, fftpack_res)
|
||||
|
||||
|
||||
@skip_xp_backends(cpu_only=True)
|
||||
@pytest.mark.parametrize("func", [dct, dst, dctn, dstn])
|
||||
@pytest.mark.parametrize("type", [1, 2, 3, 4])
|
||||
def test_orthogonalize_default(func, type, xp):
|
||||
# Test orthogonalize is the default when norm="ortho", but not otherwise
|
||||
x = xp.asarray(np.random.rand(100))
|
||||
|
||||
for norm, ortho in [
|
||||
("forward", False),
|
||||
("backward", False),
|
||||
("ortho", True),
|
||||
]:
|
||||
a = func(x, type=type, norm=norm, orthogonalize=ortho)
|
||||
b = func(x, type=type, norm=norm)
|
||||
xp_assert_close(a, b)
|
||||
|
||||
|
||||
@skip_xp_backends(cpu_only=True)
|
||||
@pytest.mark.parametrize("norm", ["backward", "ortho", "forward"])
|
||||
@pytest.mark.parametrize("func, type", [
|
||||
(dct, 4), (dst, 1), (dst, 4)])
|
||||
def test_orthogonalize_noop(func, type, norm, xp):
|
||||
# Transforms where orthogonalize is a no-op
|
||||
x = xp.asarray(np.random.rand(100))
|
||||
y1 = func(x, type=type, norm=norm, orthogonalize=True)
|
||||
y2 = func(x, type=type, norm=norm, orthogonalize=False)
|
||||
xp_assert_close(y1, y2)
|
||||
|
||||
|
||||
@skip_xp_backends('jax.numpy',
|
||||
reasons=['jax arrays do not support item assignment'],
|
||||
cpu_only=True)
|
||||
@pytest.mark.parametrize("norm", ["backward", "ortho", "forward"])
|
||||
def test_orthogonalize_dct1(norm, xp):
|
||||
x = xp.asarray(np.random.rand(100))
|
||||
|
||||
x2 = copy(x, xp=xp)
|
||||
x2[0] *= SQRT_2
|
||||
x2[-1] *= SQRT_2
|
||||
|
||||
y1 = dct(x, type=1, norm=norm, orthogonalize=True)
|
||||
y2 = dct(x2, type=1, norm=norm, orthogonalize=False)
|
||||
|
||||
y2[0] /= SQRT_2
|
||||
y2[-1] /= SQRT_2
|
||||
xp_assert_close(y1, y2)
|
||||
|
||||
|
||||
@skip_xp_backends('jax.numpy',
|
||||
reasons=['jax arrays do not support item assignment'],
|
||||
cpu_only=True)
|
||||
@pytest.mark.parametrize("norm", ["backward", "ortho", "forward"])
|
||||
@pytest.mark.parametrize("func", [dct, dst])
|
||||
def test_orthogonalize_dcst2(func, norm, xp):
|
||||
x = xp.asarray(np.random.rand(100))
|
||||
y1 = func(x, type=2, norm=norm, orthogonalize=True)
|
||||
y2 = func(x, type=2, norm=norm, orthogonalize=False)
|
||||
|
||||
y2[0 if func == dct else -1] /= SQRT_2
|
||||
xp_assert_close(y1, y2)
|
||||
|
||||
|
||||
@skip_xp_backends('jax.numpy',
|
||||
reasons=['jax arrays do not support item assignment'],
|
||||
cpu_only=True)
|
||||
@pytest.mark.parametrize("norm", ["backward", "ortho", "forward"])
|
||||
@pytest.mark.parametrize("func", [dct, dst])
|
||||
def test_orthogonalize_dcst3(func, norm, xp):
|
||||
x = xp.asarray(np.random.rand(100))
|
||||
x2 = copy(x, xp=xp)
|
||||
x2[0 if func == dct else -1] *= SQRT_2
|
||||
|
||||
y1 = func(x, type=3, norm=norm, orthogonalize=True)
|
||||
y2 = func(x2, type=3, norm=norm, orthogonalize=False)
|
||||
xp_assert_close(y1, y2)
|
||||
|
||||
@skip_xp_backends(np_only=True,
|
||||
reasons=['array-likes only supported for NumPy backend'])
|
||||
@pytest.mark.parametrize("func", [dct, idct, dctn, idctn, dst, idst, dstn, idstn])
|
||||
def test_array_like(xp, func):
|
||||
x = [[[1.0, 1.0], [1.0, 1.0]],
|
||||
[[1.0, 1.0], [1.0, 1.0]],
|
||||
[[1.0, 1.0], [1.0, 1.0]]]
|
||||
xp_assert_close(func(x), func(xp.asarray(x)))
|
||||
Reference in New Issue
Block a user