# Untitled

a guest
Sep 28th, 2020
150
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. import numpy as np
2. import matplotlib.pyplot as plt
3. import pywt
4.
5. #%%###########################################################################
6. def allfft(x):
7.     coef = np.fft.fft(x)
8.     return coef, np.angle(coef), np.abs(coef)
9.
10. def wplot(x, y=None, size=0, show=False, ax_equal=False, **kw):
11.     if y is None:
12.         plt.plot(x, **kw)
13.     else:
14.         plt.plot(x, y, **kw)
15.     _scale_plot(plt.gcf(), plt.gca(), size=size, show=show, ax_equal=ax_equal)
16.
17. def wscat(x, y=None, size=0, show=False, ax_equal=False, s=18, **kw):
18.     if y is None:
19.         plt.scatter(np.arange(len(x)), x, s=s, **kw)
20.     else:
21.         plt.scatter(x, y, s=s, **kw)
22.     _scale_plot(plt.gcf(), plt.gca(), size=size, show=show, ax_equal=ax_equal)
23.
24. def imshow(data, norm=None, complex=None, show=1, **kw):
25.     kw['interpolation'] = kw.get('interpolation', 'none')
26.     if norm is None:
27.         mx = np.max(np.abs(data))
28.         vmin, vmax = -mx, mx
29.     else:
30.         vmin, vmax = norm
31.
32.     if (complex is None and np.sum(np.abs(np.imag(data))) < 1e-8) or (
33.             complex is False):
34.         plt.imshow(np.real(data), cmap='bwr', vmin=vmin, vmax=vmax, **kw)
35.     else:
36.         fig, axes = plt.subplots(1, 2)
37.         axes[0].imshow(data.real, cmap='bwr', vmin=vmin, vmax=vmax, **kw)
38.         axes[1].imshow(data.imag, cmap='bwr', vmin=vmin, vmax=vmax, **kw)
39.
40.     plt.subplots_adjust(left=0, right=1, bottom=0, top=1, wspace=0, hspace=0)
41.     if show:
42.         plt.show()
43.
44.
45. def _scale_plot(fig, ax, size=True, show=False, ax_equal=False):
46.     xmin, xmax = ax.get_xlim()
47.     rng = xmax - xmin
48.     ax.set_xlim(xmin + .018 * rng, xmax - .018 * rng)
49.     if size:
50.         fig.set_size_inches(15, 7)
51.     if ax_equal:
52.         yabsmax = max(np.abs([*ax.get_ylim()]))
53.         mx = max(yabsmax, max(np.abs([xmin, xmax])))
54.         ax.set_xlim(-mx, mx)
55.         ax.set_ylim(-mx, mx)
56.         fig.set_size_inches(8, 8)
57.     if show:
58.         plt.show()
59. #%%###########################################################################
60. def morlet(sigma=5):
61.     cs = (1 + np.exp(-sigma ** 2) - 2 * np.exp(-.75 * sigma ** 2)) ** (-.5)
62.     ks = np.exp(-.5 * sigma ** 2)
63.     return lambda t: np.exp(-.5 * t ** 2) * (np.exp(1j * sigma * t) - ks
64.                                              ) * cs * np.pi**(-.25)
65.
66. def morlet_kernel(N, minmax=(-8, 8), sigma=5, real=False):
67.     t = np.linspace(*minmax, N, False)
68.     if real:
69.         return lambda scale: np.real(morlet(sigma)(t / scale))
70.     return lambda scale: morlet(sigma)(t / scale)
71.
72. def cwt_window(N, name, real=False):
73.     if name == 'morlet':
74.         return morlet_kernel(N, real=real)
75.
76. def _transform(x, kernel, scales):
77.     coef = np.zeros(len(scales), dtype='complex128')
78.
79.     for i, scale in enumerate(scales):
80.         psi = np.conj(kernel(scale))
81.         coef[i] = np.sum(x * psi / np.sqrt(scale))
82.     return coef
83.
84. def _scales(N, nv=32):
85.     noct = np.log2(N) - 1
86.     n_scales = int(noct * nv)
87.     return np.power(2 ** (1 / nv), np.arange(1, n_scales + 1))
88. #%%###########################################################################
89. def cwt(x, win_len=None, win_inc=None, win='morlet', viz_scale=0, real=False):
90.     N = len(x)
91.     win_len = win_len or N // 8
92.     win_inc = win_inc or win_len
93.     n_wins = (N - win_len) // win_inc + 1
94.
95.     scales = _scales(N, nv=32)
96.     coef = np.zeros((n_wins, len(scales)), dtype='complex128')
97.     kernel = cwt_window(win_len, win, real=real)
98.
99.     for tau in range(n_wins):
100.         start = tau * win_inc
101.         end   = start + win_len
102.         coef[tau, :] = _transform(x[start:end], kernel, scales)
103.         if viz_scale:
104.             cwt_viz(x, kernel(viz_scale), start, end)
105.     return coef.T
106.
107.
108. def cwt2(x, win_len=None, win='morlet', real=False):
109.     N = len(x)
110.     win_len = win_len or N // 8
111.
112.     scales = _scales(N, nv=32)
113.     coef = np.zeros((len(scales), N), dtype='complex128')
114.     kernel = cwt_window(win_len, win, real=real)
115.     wl2 = win_len // 2
116.
117.     for i, scale in enumerate(scales):
118.         coef[i, :] = np.convolve(x, kernel(scale)[::-1])[wl2:-(wl2 - 1)]
119.     return coef
120.
121.
122. def cwt_viz(x, psi, start, end):
123.     wplot(x)
124.     # plt.axvline(start, color='k', linestyle='--')
125.     # plt.axvline(end,   color='k', linestyle='--')
126.     wplot(np.pad(psi.real, [start, len(x) - end]), color='purple', show=1)
127.     plt.show()
128.
129. def _t(min, max, N):
130.     return np.linspace(min, max, N, False)
131.
132. def cos_f(freqs, N=128):
133.     return np.concatenate([np.cos(2 * np.pi * f * _t(i, i + 1, N))
134.                            for i, f in enumerate(freqs)])
135. #%%###########################################################################
136. N = 128
137. x = cos_f([1, 4], N)
138. coef, theta, r = allfft(x)
139.
140. wplot(x); wscat(x, s=8, show=1)
141. wplot(coef.real); wplot(coef.imag)
142. #%%###########################################################################
143. cwt(x, win_len=96, win_inc=20, real=1, viz_scale=2)
144. coef  = cwt(x, win_len=96, win_inc=1, real=1)
145. coef2 = cwt2(x, win_len=96, win='morlet', real=1)
146. coefp = pywt.cwt(x, _scales(len(x)), 'morl', method='conv')[0]
147. #%%
148. imshow(coef)
149. imshow(coef2)
150. imshow(coefp)
151.
RAW Paste Data