0
$\begingroup$

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): enter image description here

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[1] - time[0]
    [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[0], -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)

This plot from this code doesn't give me the plot I want, it looks like this: enter image description here

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.
Thank you,
Elad

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.