If you have ever heard Python and Fourier nouns, chances are you’ll find this post useful: here I will explore a simple way to implement the Short-Time Fourier Transform in Python in order to run a frequency analysis for a given signal. Let’s get to work!

---

A few weeks ago I was invited to join a Kaggle’s competition with my colleagues from my research group. The idea has sounded very interesting to me since it is about seismic data, i.e. time series, which is somehow related to the signal category I was used to working with — for the curious ones, I have applied wavelets for drowsiness detection from brain signals. By the way, wavelet analysis has become famous when Stéphane Mallat applied it to seismic and geological data.

But before we start dealing with wavelets (which I do intend to get back to them in the near future), the main idea here is to get acquainted with periodic signals. I mean, a time series representing some cyclic events that repeat from time to time. And that, therefore, can be often designated as a pattern.

You may find patterns represented in the most different forms, but one, in particular, is the circle: the day and the night time, the seasons of the year, the electricity generated in hydroelectric, all of them are related to a circular source (or at least an approximated one, since the Earth is an oblate spheroid which orbits the Sun in an elliptical path).

From a mathematical perspective, any cyclic pattern can be somehow represented by the trigonometric sine and cosine functions, as demonstrated by the French mathematician and physicist Joseph Fourier, who in the early nineteenth century has established what is known today as Fourier Series and Fourier Transform.

I will assume you are already familiar with trigonometric signals and their relation with frequency. In any case, what is important to notice is that a periodic or cyclic signal will repeat itself by a given amount of time, known as period (T). The frequency (f), in turn, consists of the number of repetitions per second and is measured in hertz (Hz), in homage to the German physicist Heinrich Hertz.

In the next sections, we will explore these concepts in a very practical way (I’m also assuming you are used to Python). We will first figure out how to generate and visualize some periodic waves (or signals), moving then to a Short-Time Fourier Transform implementation. For last, we will see an example of how to apply such analysis for a non-trigonometric function but rhythmic one, as the brain signals.

1. Generating periodic signals

In data science — and here I’m considering all the disciplines related to it, such as pattern recognition, signal processing, machine learning and so on — it is always useful to have a deep understanding of how our data (a.k.a. signals) behave when inserted into an analysis or classification tool. The problem is that at first we usually know little about our data — which in the real world can be complex and noisy — and in fact, looking for analytical tools is an attempt to understand them.

In these cases, the best alternative is to explore the analysis tool with simple and well-known data. Paraphrasing the Nobel laureate Richard Feynman — who left in his blackboard the following quote, after his death in 1988: “What I can not create, I do not understand” –, we will first generate our own periodic signals.

A time series usually results from a sampling process effort. When this signal is generated, as is the case here, we must also establish the time where the signal supposedly occurs. The function below emulates the sampling process and has as input the initial and the final time samples. In other words, it returns the x-axis basis where our signal will be placed.

[code language="python"]
import numpy as np

def timeSeries(fs, ti, tf):
    '''Returns a time series based on:
        fs = the sampling rate, in hertz (e.g. 1e3 for 1kHz);
        ti = the starting time (e.g. second 0);
        tf = the final time (e.g. second 1).'''
    T = 1./fs
    timeSeries = np.arange(ti,tf,T)
    return timeSeries
[/code]

Let’s remember that a sinusoidal signal is generated by the function:
,
where f is the frequency, in hertz, of the signal y(t), and t is the time series generated by the code above. This equation is then transcripted to Python:

[code language="python"]
def genSignal(freq, amp, timeSeries):
    '''Returns the time series for a sinusoidal signal, where:
        freq = frequency in Hertz (integer);
        amp = signal amplitude (float);
        timeSeries = the linear time series (array).'''
    signal = amp*np.sin(freq*2*np.pi*timeSeries)
    return signal 
[/code]

Almost as important as generating the signal is to provide a means of visualizing it. Using the Matplotlib library, the following function can be adapted to view any time series:

[code language="python"]
import matplotlib.pyplot as plt

def plotSignal(signal, timeSeries, figsize=(6,3)):
    '''signal = numpy.array or Pandas.Series;
        timeSeries = the linear time series x-axis;
        figsize = the plot size, set as (6,3) by default.'''
    fig, axes = plt.subplots(1,1, figsize=figsize)
    plt.plot(timeSeries, signal)
    #Adjust the ylim to go 10% above and below from the signal amplitudes
    axes.set_ylim(signal.min()+signal.min()*0.1, signal.max()+signal.max()*0.1)
    axes.grid(True)
    axes.set_ylabel('Signal amplitude')
    axes.set_xlabel('Time (s)')
    return
[/code]

It also sounds practical to code a wrapper function to create and visualize a sinusoidal function getting frequency and amplitude as input parameters:

[code language="python"] def sineWave(freq, amp, timeSeries):
    '''Generates and plots a sine wave using genSignal() and plotSignal().
        freq = frequency in Hertz (int);
        amp = signal amplitude (float);
        timeSeries = a linear time series accordingly to the sample rate.'''
    signal = genSignal(freq, amp, timeSeries)
    plotSignal(signal, timeSeries)
    return signal [/code]

Finally, we will now execute these functions together in order to obtain some sinusoidal waves:

[code language='python']
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

## Sampling time for 0s to 10s, 1kHz:
time = timeSeries(1e3,0,10)

## Generating a sine wave of 1Hz and amplitude of 2:
sign1 = sineWave(1, 2, time)

## Generating a sine wave of 0.1Hz and amplitude of 1:
sign01 = sineWave(0.1, 1, time)
[/code]

Which then give us the following result:

• 

• 

It may look quite simpler, but it is important to have a signal which we know exactly how it was created — we hold its model, i.e. its formula. Observe in the first chart above that a frequency of 0.1Hz means the signal will repeat itself every 10 seconds. In the second chart, we duplicate the amplitude of the signal which now repeats its pattern every second.

Putting together the functions we have defined with the numpy.concatenate method, as shown below, we can generate more complex waveforms:

[code language='python']
# Generating two segmented time series:
time1 = timeSeries(fsLANL,0,1)
time2 = timeSeries(fsLANL,1,2)
# Concatenating them in order to plot:
time = np.concatenate((time1, time2), axis=None)

# Generating the signals for each time series:
signal1 = genSignal(20,50,time1)
# Here we add up two frequencies:
signal2 = genSignal(113,28,time2) + genSignal(70,53,time2) + genSignal(0.25,30,time2)
# Concatenating both signals
signal = np.concatenate((signal1,signal2), axis=None)

# Plotting the resulting signal:
plotSignal('[0.25, 20, 70, 113]', signal, time, figsize=(12,6)) 
[/code]

From our first charts, it becomes clear that when we are dealing with low and single frequencies we can still identify cyclic patterns only through visual inspection in the time domain signals (e.g. counting the tops during a second). However, for more complex waves (see the interval between 1s and 2s in the chart above) it becomes clear that a tool is needed to identify any cyclic pattern.

2. The Short-Time Fourier Transform

A complete understanding of time series representation in the frequency domain requires a deepening of linear algebra and calculus for expanding a signal in a series of functions.

Such an approach, though, will not be taken here: those who already understand this mathematical abstraction would find a practical code for implementing a frequency analysis in Python. In the same way, those who are facing periodic signals for the first time may start deepening their theoretical understanding from these simple examples.

---

Before we start coding, I will provide a brief definition (please don’t get scared with the symbols… If you want, you may just page down to skip this block): Fourier has proved that it is possible to expand a great number of functions into pure harmonic functions series (i.e. a sum of sines and cosines). By the way, from the Euler’s formula we know we can represent harmonic functions as a complex exponential:

.

Now consider that we can expand a given signal y(t) into a Fourier Series:

,
where are the Fourier coefficients obtained by the Fourier transform of the signal; and is the harmonic base function that represents the frequency .

What is important to have in mind is that each coefficient — what our computer will be calculating for us — is associated with a harmonic base function, which in turn relates this coefficient to a specific frequency. Then, if we know the value of each coefficient, we could know the intensity of each frequency component present in the signal y(t).

Another important aspect one should consider is that these base functions have infinite duration and therefore are not suitable for representing signals with discontinuities or of which the location in time or space is also desired (i.e. when we seek to know in which period of the signal a given frequency or rhythm happens the most).

The Short-Time Fourier Transform (STFT) is a way to overcome this. Its strategy consists of multiplying each basis function of the transform by a window function, w(t). This last thus being limited to a given period of time, i.e. having non-null values just for a limited interval.

Considering the function space in the Fourier transform can be defined by trigonometric functions in its complex form, for the STFT we create a new basis defined both by frequency () and position ():

,
.

The shape and length of w(t) must be chosen accordingly to the analysis interest. The shape depends on the function used to generate it and determines its capability of frequency resolution. On the other hand, the length (N) defines the window interval and therefore its temporal resolution.

Many window functions can be used and tried out (an overview of their main features and differences can be found here). For our purpose — or when in doubt about which one to pick — the Hann window can be a satisfying choice:

---

Let us now return to our Python code. Once we have understood the basic principles the STFT relies on, we can make use of the signal module from SciPy library to implement an spectrogram — which consist of plotting the squared magnitude of the STFT coefficients.

So, with the code below we will compute the STFT for our first signal (page up and seek for the sign1). Also, we must consider that the SciPy implementation will always set the frequency axis accordingly to half of the sampling frequency, the reason why I set the plt.ylim() command.

[code language='python']
## Calculating the STFT spectrogram
fs = 1e3
f, t, Zxx = signal.stft(sign1, fs)
plt.pcolormesh(t, f, np.abs(Zxx), cmap=plt.get_cmap('RdYlGn'))
plt.title('STFT Magnitude')
plt.ylim(0,160)
plt.ylabel('Frequency [Hz]')
plt.xlabel('Time [sec]')
plt.show()
[/code]

Okay, I do agree with you that this result is far short of what I expected. The good news is that (almost) everything is a matter of adjustment. To make things easier, let’s write the following function:

[code language='python']
def calcSTFT(inputSignal, samplingFreq, window='hann', nperseg=256, figsize=(9,5), cmap='magma', ylim_max=None, output=False):
    '''Calculates the STFT for a time series:
        inputSignal: numpy array for the signal (it also works for Pandas.Series);
        samplingFreq: the sampling frequency;
        window : str or tuple or array_like, optional
            Desired window to use. If `window` is a string or tuple, it is
            passed to `get_window` to generate the window values, which are
            DFT-even by default. See `get_window` for a list of windows and
            required parameters. If `window` is array_like it will be used
            directly as the window and its length must be nperseg. Defaults
            to a Hann window.
        nperseg : int, optional
            Length of each segment. Defaults to 256.
        figsize: the plot size, set as (6,3) by default;
        cmap: the color map, set as the divergence Red-Yellow-Green by default;
        ylim_max: the max frequency to be shown. By default it's the half sampling frequency;
        output: 'False', as default. If 'True', returns the STFT values.
        
        Outputs (if TRUE):
            f: the frequency values
            t: the time values
            Zxx: the STFT values"'''
    ##Calculating STFT
    f, t, Zxx = signal.stft(inputSignal, samplingFreq, window=window, nperseg=nperseg)
    ##Plotting STFT
    fig = plt.figure(figsize=figsize)
    ### Different methods can be chosen for normalization: PowerNorm; LogNorm; SymLogNorm.
    ### Reference: https://matplotlib.org/tutorials/colors/colormapnorms.html
    spec = plt.pcolormesh(t, f, np.abs(Zxx), 
                          norm=colors.PowerNorm(gamma=1./8.),
                          cmap=plt.get_cmap(cmap))
    cbar = plt.colorbar(spec)
    ##Plot adjustments
    plt.title('STFT Spectrogram')
    ax = fig.axes[0]
    ax.grid(True)
    ax.set_title('STFT Magnitude')
    if ylim_max:
        ax.set_ylim(0,ylim_max)
    ax.set_ylabel('Frequency [Hz]')
    ax.set_xlabel('Time [sec]')
    fig.show
    if output:
        return f,t,Zxx
    else:
        return
[/code]

To test it, we will create a new signal whose values are easy to understand and manipulate, a 100Hz sine wave with amplitude 1000. Since we will keep it lasting 10s, with the same sampling rate of 1kHz, we can use the array time defined before.

[code language='python']
fs = 1e3  #The sampling frequency
test = genSignal(100,1000,time)  #Generating the signal

calcSTFT(test, fs)  #Calling our function in its simplest form!
[/code]

The purpose to start with a simple signal is that it allows us to check whether there is a gross error in our code or not. Fortunately, it’s not the case! The plot frequency range goes up to 500Hz, half the sampling frequency, as expected. Also, we can precisely see the 100Hz component, which values are normalized (you will find the details on the normalization strategy here).

Going further, we can specify other parameters to our STFT. To create the charts below I set different values for the nperseg parameter, which correspond to the window function size. For the left image, nperseg = 64, while for the right one nperseg = 2048. Note how a larger window results in a much more accurate frequency resolution. And vice versa.

---

There is no specific rule to determine the number of samples per segment. However, it’s useful to consider 1/4 of the sampling rate. E.g. for a 1kHz sampling rate, the nperseg parameter is set by default to 256. This tradeoff can become clearer if we analyze our complex signal, which we generated with components from 0.25 to 123Hz:

[code language='python']
calcSTFT(signal, fs, nperseg=2**19, ylim_max=300)
[/code]

In the plot above all the known frequency components of the signal become evident, although there is lower time resolution. For some applications, 250ms are insignificant. For others, it can be limiting. I must anticipate to you that there is no analytical tool that may offer accuracy both in time and frequency, as demonstrated by the Heisenberg’s principle — but this will be subject for a next post.

3. Brain signal analysis

To exemplify how the STFT analysis can be useful for “real life” situations, we will now apply the functions we built on a brain signal segment obtained from electroencephalography (EEG). If you want to try it on your own, the data is available here.

---

Since the signal was previously processed on Matlab, we will first have to convert the data structure to a Python dictionary using the loadmat method. Knowing that the signal was labeled as sD_W, we can move on to loading it:

[code language='python']
import scipy.io
mat = scipy.io.loadmat('brainSignal.mat')

brain_raw = mat['sD_W']
## Checking the shape of the variable:
brain_raw.shape

## Adequating the signal to an np.array
brain = brain_raw.T[0]

## Plotting the brain signal:
plt.plot(brain)

## Calculating its STFT
calcSTFT(brain, 200, nperseg=2**8, figsize=(9,4), ylim_max=100, cmap='inferno')
[/code]

Unlike when dealing with sinusoidal signals, now we can’t deduce the rhythms (i.e. frequency components) present in the EEG signal only by looking to its time domain plot. We could even say that something is happening after 20s, but we could not give any detail.

On the other hand, the spectrogram above shows us that after the event around 20s, a well established low-frequency rhythm appears from 20s to 36s and then again from 44s to 56s. As with any data science activity, the interpretation is always context-dependent. In this case, this EEG signal is from a subject submitted to polysomnography. The identified frequency range corresponds to the alpha rhythm, indicating that the patient is probably awake and relaxed, but with closed eyes. For last, the dark region around 50Hz corresponds to a pass-band filter applied during the signal acquisition in order to filter the 50Hz interference from AC power supply.

4. Conclusions

The goal of this post was not only to show how to implement STFT in Python, but also to bring a brief introduction to the theory behind this powerful analytical tool — which supports the more intricate ideas of mathematics.

Also, by being able to understand the main parameters that should be observed in such analysis, and drawing on the example with brain signals, I hope the reader may tailor the code and apply such a tool to any cyclic signal he/she may find: from biomedical or seismic signals going even to stock market prices or the movement of stars.