Python tool for time-frequency analysis

Question

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[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:

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

Answer 1

You can find a nice tutorial for time-frequency analysis in Numerical python by Johansson, chapter 17. link to github repository.

You can also check the scipy.signal.spectrogram.

import numpy as np
from scipy import signal
from scipy.fft import fftshift
import matplotlib.pyplot as plt

# Generate a test signal, a 2 Vrms sine wave whose frequency
# is slowly modulated around 3kHz, corrupted by white noise
# of exponentially decreasing magnitude sampled at 10 kHz.

fs = 1e4
N = 1e5
amp = 2 * np.sqrt(2)
noise_power = 0.01 * fs / 2
time = np.arange(N) / float(fs)
mod = 500 * np.cos(2 * np.pi * 0.25 * time)
carrier = amp * np.sin(2 * np.pi * 3e3 * time + mod)
noise = np.random.normal(scale=np.sqrt(noise_power), size=time.shape)
noise *= np.exp(-time / 5)
x = carrier + noise

fig, ax = plt.subplots(2, figsize=(8, 7))

f, t, Sxx = signal.spectrogram(x, fs)
ax[0].pcolormesh(t, f, Sxx)
ax[0].set_ylabel('Frequency [Hz]')
ax[0].set_xlabel('Time [sec]')

# Note, if using output that is not one sided, then use the following:
f, t, Sxx = signal.spectrogram(x, fs, return_onesided=False)
ax[1].pcolormesh(t, fftshift(f), fftshift(Sxx, axes=0))
ax[1].set_ylabel('Frequency [Hz]')
ax[1].set_xlabel('Time [sec]')
plt.savefig("fig.png")
plt.show()