Guest User

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