88 lines
3.8 KiB
Python
88 lines
3.8 KiB
Python
|
import numpy as np
|
||
|
|
import scipy.fft
|
||
|
|
from scipy.fft import set_backend
|
||
|
|
from scipy.fft import _pocketfft
|
||
|
|
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')
|
||
|
|
|
||
|
|
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)
|
||
|
|
|
||
|
|
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)
|
||
|
|
|
||
|
|
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)
|
||
|
|
|
||
|
|
|
||
|
|
@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)
|
||
|
|
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')
|