I am trying to perform time-frequency analyses using the PyWavelets (pywt) toolkit for python. My ultimate goal is to perform time-frequency analyses for EEG signals but I am starting with something simpler.
For a sanity test, I am creating a simple signal of length 2 seconds, with sample rate 250Hz, containing 2 sine waves - one of 3Hz and one of 10Hz. I would like to create a time-frequency plot that has two horizontal lines - one for the 3Hz and one for the 10Hz, which looks like this (only for illustration purposes):
For this purpose, I tried using code from the following tutorial : http://ataspinar.com/2018/12/21/a-guide-for-using-the-wavelet-transform-in-machine-learning/, specifically in section 3.1 of the tutorial.
This is a minimal example based on the code from the tutorial:
from UliEngineering.SignalProcessing.Simulation import sine_wave import pywt import numpy as np import matplotlib.pyplot as plt def plot_wavelet(time, signal, scales, waveletname='cmor', cmap=plt.cm.seismic, title='Wavelet Transform (Power Spectrum) of signal', ylabel='Period (seconds)', xlabel='Time'): dt = time - time [coefficients, frequencies] = pywt.cwt(signal, scales, waveletname, dt) power = (abs(coefficients)) ** 2 period = 1. / frequencies levels = [0.0625, 0.125, 0.25, 0.5, 1, 2, 4, 8] contourlevels = np.log2(levels) fig, ax = plt.subplots(figsize=(15, 10)) im = ax.contourf(time, np.log2(period), np.log2(power), contourlevels, extend='both', cmap=cmap) ax.set_title(title, fontsize=20) ax.set_ylabel(ylabel, fontsize=18) ax.set_xlabel(xlabel, fontsize=18) yticks = 2 ** np.arange(np.ceil(np.log2(period.min())), np.ceil(np.log2(period.max()))) ax.set_yticks(np.log2(yticks)) ax.set_yticklabels(yticks) ax.invert_yaxis() ylim = ax.get_ylim() ax.set_ylim(ylim, -1) cbar_ax = fig.add_axes([0.95, 0.5, 0.03, 0.25]) fig.colorbar(im, cax=cbar_ax, orientation="vertical") plt.show() def generate_sine_wave(length, samplerate, frequencies): wave = np.zeros(int(length * samplerate)) for frequency in frequencies: wave += sine_wave(frequency=frequency, samplerate=samplerate, length=length) return wave signal = generate_sine_wave(2, 250, [3, 10]) N = len(signal) t0 = 0 dt = 1/250 time = np.arange(0, N) * dt +t0 scales = np.arange(1, 256) plot_wavelet(time, signal, scales)
I tried many modifications for this code but none gave me the result I want. And there are a couple of things I don't understand about the code:
- What is the purpose of the "period" variable in the "plot_wavelet" function and how do I make the y-axis show frequencies instead?
- What is the purpose of the "scales" variable?
- How do I define a frequency range that I want the result to include?
- How do I use linear scaling for the frequencies instead of log scale?
If anyone can give some pointers regarding this I will be very happy. Been spending some time trying to plot normal time-frequency plots but still haven't been able to find a python tool that performs this simple plot which makes sense to me.