From 4b67d2acd4e480294cdb4863a22a5e6b8ff12d84 Mon Sep 17 00:00:00 2001 From: ChS-YHWH Date: Tue, 25 Aug 2026 00:32:28 +0800 Subject: [PATCH] ENH: use real FFTs for real-valued CWT convolution --- pywt/_cwt.py | 16 ++++++++++++--- pywt/tests/test_cwt_wavelets.py | 35 +++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 3 deletions(-) diff --git a/pywt/_cwt.py b/pywt/_cwt.py index 0bea72c4..2cd7d577 100644 --- a/pywt/_cwt.py +++ b/pywt/_cwt.py @@ -139,6 +139,7 @@ def cwt(data, scales, wavelet, sampling_period=1., method='conv', axis=-1, if method == 'fft': size_scale0 = -1 fft_data = None + use_real_fft = data.dtype.kind != 'c' and int_psi.dtype.kind != 'c' elif method != "conv": raise ValueError("method must be 'conv' or 'fft'") @@ -179,10 +180,19 @@ def cwt(data, scales, wavelet, sampling_period=1., method='conv', axis=-1, ) if size_scale != size_scale0: # Must recompute fft_data when the padding size changes. - fft_data = np.fft.fft(data, size_scale, axis=-1) + if use_real_fft: + fft_data = np.fft.rfft(data, size_scale, axis=-1) + else: + fft_data = np.fft.fft(data, size_scale, axis=-1) size_scale0 = size_scale - fft_wav = np.fft.fft(int_psi_scale, size_scale, axis=-1) - conv = np.fft.ifft(fft_wav * fft_data, axis=-1) + if use_real_fft: + fft_wav = np.fft.rfft(int_psi_scale, size_scale, axis=-1) + conv = np.fft.irfft( + fft_wav * fft_data, n=size_scale, axis=-1 + ) + else: + fft_wav = np.fft.fft(int_psi_scale, size_scale, axis=-1) + conv = np.fft.ifft(fft_wav * fft_data, axis=-1) conv = conv[..., :data.shape[-1] + int_psi_scale.size - 1] coef = - np.sqrt(scale) * np.diff(conv, axis=-1) diff --git a/pywt/tests/test_cwt_wavelets.py b/pywt/tests/test_cwt_wavelets.py index 22113b33..ede41044 100644 --- a/pywt/tests/test_cwt_wavelets.py +++ b/pywt/tests/test_cwt_wavelets.py @@ -471,6 +471,41 @@ def test_cwt_method_fft(): assert_allclose(cfs_conv, cfs_fft, rtol=0, atol=1e-13) +@pytest.mark.parametrize('dtype, tol', [(np.float32, 1e-5), + (np.float64, 1e-13)]) +@pytest.mark.parametrize('shape, axis', [((49,), -1), + ((3, 50), 1), + ((49, 3), 0)]) +def test_cwt_method_fft_real(dtype, tol, shape, axis): + rstate = np.random.RandomState(1) + data = rstate.randn(*shape).astype(dtype) + scales = np.r_[1.0625, np.arange(1, 64)] + + cfs_conv, _ = pywt.cwt( + data, scales, 'morl', method='conv', axis=axis + ) + cfs_fft, _ = pywt.cwt( + data, scales, 'morl', method='fft', axis=axis + ) + + assert_equal(cfs_fft.dtype, dtype) + assert_allclose(cfs_conv, cfs_fft, rtol=tol, atol=tol) + + +@pytest.mark.parametrize('dtype, tol', [(np.complex64, 1e-5), + (np.complex128, 1e-13)]) +def test_cwt_method_fft_complex_data_real_wavelet(dtype, tol): + rstate = np.random.RandomState(1) + data = (rstate.randn(49) + 1j * rstate.randn(49)).astype(dtype) + scales = np.r_[1.0625, np.arange(1, 16)] + + cfs_conv, _ = pywt.cwt(data, scales, 'morl', method='conv') + cfs_fft, _ = pywt.cwt(data, scales, 'morl', method='fft') + + assert_equal(cfs_fft.dtype, dtype) + assert_allclose(cfs_conv, cfs_fft, rtol=tol, atol=tol) + + def test_continuous_wavelet_pickle(tmpdir): wavelet = pywt.ContinuousWavelet('cmor1.5-1.0') filename = os.path.join(tmpdir, 'cwav.pickle')