Aluode/PerceptionLabPortable

1from math import ceil, floor
2
3from ._extensions._pywt import (
4    ContinuousWavelet,
5    DiscreteContinuousWavelet,
6    Wavelet,
7    _check_dtype,
8)
9from ._functions import integrate_wavelet, scale2frequency
10from ._utils import AxisError
11
12__all__ = ["cwt"]
13
14
15import numpy as np
16
17try:
18    import scipy
19    fftmodule = scipy.fft
20    next_fast_len = fftmodule.next_fast_len
21except ImportError:
22    fftmodule = np.fft
23
24    # provide a fallback so scipy is an optional requirement
25    # note: numpy.fft in numpy 2.0 is as fast as scipy.fft, so could be used
26    # unconditionally once the minimum supported numpy version is >=2.0
27    def next_fast_len(n):
28        """Round up size to the nearest power of two.
29
30        Given a number of samples `n`, returns the next power of two
31        following this number to take advantage of FFT speedup.
32        This fallback is less efficient than `scipy.fftpack.next_fast_len`
33        """
34        return 2**ceil(np.log2(n))
35
36
37def cwt(data, scales, wavelet, sampling_period=1., method='conv', axis=-1):
38    """
39    cwt(data, scales, wavelet)
40
41    One dimensional Continuous Wavelet Transform.
42
43    Parameters
44    ----------
45    data : array_like
46        Input signal
47    scales : array_like
48        The wavelet scales to use. One can use
49        ``f = scale2frequency(wavelet, scale)/sampling_period`` to determine
50        what physical frequency, ``f``. Here, ``f`` is in hertz when the
51        ``sampling_period`` is given in seconds.
52    wavelet : Wavelet object or name
53        Wavelet to use
54    sampling_period : float
55        Sampling period for the frequencies output (optional).
56        The values computed for ``coefs`` are independent of the choice of
57        ``sampling_period`` (i.e. ``scales`` is not scaled by the sampling
58        period).
59    method : {'conv', 'fft'}, optional
60        The method used to compute the CWT. Can be any of:
61            - ``conv`` uses ``numpy.convolve``.
62            - ``fft`` uses frequency domain convolution.
63            - ``auto`` uses automatic selection based on an estimate of the
64              computational complexity at each scale.
65
66        The ``conv`` method complexity is ``O(len(scale) * len(data))``.
67        The ``fft`` method is ``O(N * log2(N))`` with
68        ``N = len(scale) + len(data) - 1``. It is well suited for large size
69        signals but slightly slower than ``conv`` on small ones.
70    axis: int, optional
71        Axis over which to compute the CWT. If not given, the last axis is
72        used.
73
74    Returns
75    -------
76    coefs : array_like
77        Continuous wavelet transform of the input signal for the given scales
78        and wavelet. The first axis of ``coefs`` corresponds to the scales.
79        The remaining axes match the shape of ``data``.
80    frequencies : array_like
81        If the unit of sampling period are seconds and given, then frequencies
82        are in hertz. Otherwise, a sampling period of 1 is assumed.
83
84    Notes
85    -----
86    Size of coefficients arrays depends on the length of the input array and
87    the length of given scales.
88
89    Examples
90    --------
91    >>> import pywt
92    >>> import numpy as np
93    >>> import matplotlib.pyplot as plt
94    >>> x = np.arange(512)
95    >>> y = np.sin(2*np.pi*x/32)
96    >>> coef, freqs=pywt.cwt(y,np.arange(1,129),'gaus1')
97    >>> plt.matshow(coef)
98    >>> plt.show()
99
100    >>> import pywt
101    >>> import numpy as np
102    >>> import matplotlib.pyplot as plt
103    >>> t = np.linspace(-1, 1, 200, endpoint=False)
104    >>> sig  = np.cos(2 * np.pi * 7 * t) + np.real(np.exp(-7*(t-0.4)**2)*np.exp(1j*2*np.pi*2*(t-0.4)))
105    >>> widths = np.arange(1, 31)
106    >>> cwtmatr, freqs = pywt.cwt(sig, widths, 'mexh')
107    >>> plt.imshow(cwtmatr, extent=[-1, 1, 1, 31], cmap='PRGn', aspect='auto',
108    ...            vmax=abs(cwtmatr).max(), vmin=-abs(cwtmatr).max())
109    >>> plt.show()
110    """
111
112    # accept array_like input; make a copy to ensure a contiguous array
113    dt = _check_dtype(data)
114    data = np.asarray(data, dtype=dt)
115    dt_cplx = np.result_type(dt, np.complex64)
116    if not isinstance(wavelet, (ContinuousWavelet, Wavelet)):
117        wavelet = DiscreteContinuousWavelet(wavelet)
118
119    scales = np.atleast_1d(scales)
120    if np.any(scales <= 0):
121        raise ValueError("`scales` must only include positive values")
122
123    if not np.isscalar(axis):
124        raise AxisError("axis must be a scalar.")
125
126    dt_out = dt_cplx if wavelet.complex_cwt else dt
127    out = np.empty((np.size(scales),) + data.shape, dtype=dt_out)
128    precision = 10
129    int_psi, x = integrate_wavelet(wavelet, precision=precision)
130    int_psi = np.conj(int_psi) if wavelet.complex_cwt else int_psi
131
132    # convert int_psi, x to the same precision as the data
133    dt_psi = dt_cplx if int_psi.dtype.kind == 'c' else dt
134    int_psi = np.asarray(int_psi, dtype=dt_psi)
135    x = np.asarray(x, dtype=data.real.dtype)
136
137    if method == 'fft':
138        size_scale0 = -1
139        fft_data = None
140    elif method != "conv":
141        raise ValueError("method must be 'conv' or 'fft'")
142
143    if data.ndim > 1:
144        # move axis to be transformed last (so it is contiguous)
145        data = data.swapaxes(-1, axis)
146
147        # reshape to (n_batch, data.shape[-1])
148        data_shape_pre = data.shape
149        data = data.reshape((-1, data.shape[-1]))
150
151    for i, scale in enumerate(scales):
152        step = x[1] - x[0]
153        j = np.arange(scale * (x[-1] - x[0]) + 1) / (scale * step)
154        j = j.astype(int)  # floor
155        if j[-1] >= int_psi.size:
156            j = np.extract(j < int_psi.size, j)
157        int_psi_scale = int_psi[j][::-1]
158
159        if method == 'conv':
160            if data.ndim == 1:
161                conv = np.convolve(data, int_psi_scale)
162            else:
163                # batch convolution via loop
164                conv_shape = list(data.shape)
165                conv_shape[-1] += int_psi_scale.size - 1
166                conv_shape = tuple(conv_shape)
167                conv = np.empty(conv_shape, dtype=dt_out)
168                for n in range(data.shape[0]):
169                    conv[n, :] = np.convolve(data[n], int_psi_scale)
170        else:
171            # The padding is selected for:
172            # - optimal FFT complexity
173            # - to be larger than the two signals length to avoid circular
174            #   convolution
175            size_scale = next_fast_len(
176                data.shape[-1] + int_psi_scale.size - 1
177            )
178            if size_scale != size_scale0:
179                # Must recompute fft_data when the padding size changes.
180                fft_data = fftmodule.fft(data, size_scale, axis=-1)
181            size_scale0 = size_scale
182            fft_wav = fftmodule.fft(int_psi_scale, size_scale, axis=-1)
183            conv = fftmodule.ifft(fft_wav * fft_data, axis=-1)
184            conv = conv[..., :data.shape[-1] + int_psi_scale.size - 1]
185
186        coef = - np.sqrt(scale) * np.diff(conv, axis=-1)
187        if out.dtype.kind != 'c':
188            coef = coef.real
189        # transform axis is always -1 due to the data reshape above
190        d = (coef.shape[-1] - data.shape[-1]) / 2.
191        if d > 0:
192            coef = coef[..., floor(d):-ceil(d)]
193        elif d < 0:
194            raise ValueError(
195                f"Selected scale of {scale} too small.")
196        if data.ndim > 1:
197            # restore original data shape and axis position
198            coef = coef.reshape(data_shape_pre)
199            coef = coef.swapaxes(axis, -1)
200        out[i, ...] = coef
201
202    frequencies = scale2frequency(wavelet, scales, precision)
203    if np.isscalar(frequencies):
204        frequencies = np.array([frequencies])
205    frequencies /= sampling_period
206    return out, frequencies
207