Aluode/PerceptionLabPortable

1import numpy as np
2from scipy.signal import convolve
3
4
5def _ricker(points, a):
6    A = 2 / (np.sqrt(3 * a) * (np.pi**0.25))
7    wsq = a**2
8    vec = np.arange(0, points) - (points - 1.0) / 2
9    xsq = vec**2
10    mod = (1 - xsq / wsq)
11    gauss = np.exp(-xsq / (2 * wsq))
12    total = A * mod * gauss
13    return total
14
15
16def _cwt(data, wavelet, widths, dtype=None, **kwargs):
17    # Determine output type
18    if dtype is None:
19        if np.asarray(wavelet(1, widths[0], **kwargs)).dtype.char in 'FDG':
20            dtype = np.complex128
21        else:
22            dtype = np.float64
23
24    output = np.empty((len(widths), len(data)), dtype=dtype)
25    for ind, width in enumerate(widths):
26        N = np.min([10 * width, len(data)])
27        wavelet_data = np.conj(wavelet(N, width, **kwargs)[::-1])
28        output[ind] = convolve(data, wavelet_data, mode='same')
29    return output
30