def stockwell_fe(x): ci = x.shape[0] cj = x.shape[-1] res = [] for i in range(ci): cur_trial = [] for j in range(cj): coef = sw.st(x[i, :, j], 8, 30) cur_trial.append(coef) res.append(np.array(cur_trial).T) return np.array(res)
def stspecgram(x, fs, lofreq=None, hifreq=None, t0=None, t1=None): """ plot out the stockwell spectrum abs(st(x)) given frequency sampling fs in Hz lofreq and hifreq are the frequency limits expressed in terms of the nyquist frequency(?) """ n = x.shape[0] if t0 == None: t0 = 0.0 if t1 == None: t1 = n / float(fs) + t0 if lofreq == None and hifreq == None: sx = stockwell.st(x) return plotspec(abs(sx), fs, t0=t0, t1=t1) lorow = stockwell.stfreq(lofreq, n, fs) hirow = stockwell.stfreq(hifreq, n, fs) sx = stockwell.st(x, lorow, hirow) return plotspec(abs(sx), fs, lofreq=lofreq, hifreq=hifreq, t0=t0, t1=t1)
def timefreqplot(x, fs, lo=0, hi=0, title=None): """ plot a signal and its spectrogram defaultis to find the entire frequency band (lo->0.0, hi-> n/2) lo (Hz) gets transformed to the sample-based lo_n for use with stockwell.st hi (Hz) """ n = len(x) if not hi: hi_n = int(n / 2) # float(fs/2) hi = fs / 2.0 else: hi_n = stockwell.stfreq(hi, n, fs) print("stfreq(hi,n,fs):", hi) print("setting hi", hi, hi_n) if not lo: lo = 0.0 lo_n = 0 else: lo_n = stockwell.stfreq(lo, n, fs) print("stfreq(lo,n,fs):", lo) print("setting low", lo, lo_n) fig, ax = plt.subplots(nrows=2, ncols=1, sharex=True) if title: ax.set_title(title) psx = abs(stockwell.st(x, lo_n, hi_n)) print("psx.shape:", psx.shape) si = 1.0 / fs ax[0].plot(np.arange(0, n) * si, x) # extents = _get_specgram_plot_extents(psx, fs=fs, lofreq=lo, hifreq=hi,t0=0, t1=n/float(fs)) extents = _get_specgram_plot_extents(psx, fs=fs, lofreq=lo, hifreq=hi, t0=0, t1=n / float(fs)) print(extents) ax[1].set_ylabel('Hz') ax[1].imshow(psx, extent=extents, aspect='auto', origin='lower') return fig, ax, psx