# Python tool for time-frequency analysis

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)


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,