# -*- coding: utf-8 -*-
import matplotlib.pyplot as plt
import numpy as np
import scipy.signal
from ..signal.signal_detrend import signal_detrend
[docs]
def signal_timefrequency(
signal,
sampling_rate=1000,
min_frequency=0.04,
max_frequency=None,
method="stft",
window=None,
window_type="hann",
mode="psd",
nfreqbin=None,
overlap=None,
analytical_signal=True,
show=True,
):
"""**Quantify changes of a nonstationary signal’s frequency over time**
The objective of time-frequency analysis is to offer a more informative description of the
signal which reveals the temporal variation of its frequency contents.
There are many different Time-Frequency Representations (TFRs) available:
* Linear TFRs: efficient but create tradeoff between time and frequency resolution
* Short Time Fourier Transform (STFT): the time-domain signal is windowed into short
segments and FT is applied to each segment, mapping the signal into the TF plane. This
method assumes that the signal is quasi-stationary (stationary over the duration of the
window). The width of the window is the trade-off between good time (requires short
duration window) versus good frequency resolution (requires long duration windows)
* Wavelet Transform (WT): similar to STFT but instead of a fixed duration window function,
a varying window length by scaling the axis of the window is used. At low frequency, WT
proves high spectral resolution but poor temporal resolution. On the other hand, for high
frequencies, the WT provides high temporal resolution but poor spectral resolution.
* Quadratic TFRs: better resolution but computationally expensive and suffers from having
cross terms between multiple signal components
* Wigner Ville Distribution (WVD): while providing very good resolution in time and
frequency of the underlying signal structure, because of its bilinear nature, existence
of negative values, the WVD has misleading TF results in the case of multi-component
signals such as EEG due to the presence of cross terms and inference terms. Cross WVD
terms can be reduced by using smoothing kernel functions as well as analyzing the
analytic signal (instead of the original signal)
* Smoothed Pseudo Wigner Ville Distribution (SPWVD): to address the problem of cross-terms
suppression, SPWVD allows two independent analysis windows, one in time and the other in
frequency domains.
Parameters
----------
signal : Union[list, np.array, pd.Series]
The signal (i.e., a time series) in the form of a vector of values.
sampling_rate : int
The sampling frequency of the signal (in Hz, i.e., samples/second).
method : str
Time-Frequency decomposition method.
min_frequency : float
The minimum frequency.
max_frequency : float
The maximum frequency.
window : int
Length of each segment in seconds. If ``None`` (default), window will be automatically
calculated. For ``"STFT" method``.
window_type : str
Type of window to create, defaults to ``"hann"``. See :func:`.scipy.signal.get_window` to
see full options of windows. For ``"STFT" method``.
mode : str
Type of return values for ``"STFT" method``. Can be ``"psd"``, ``"complex"`` (default,
equivalent to output of ``"STFT"`` with no padding or boundary extension), ``"magnitude"``,
``"angle"``, ``"phase"``. Defaults to ``"psd"``.
nfreqbin : int, float
Number of frequency bins. If ``None`` (default), nfreqbin will be set to
``0.5*sampling_rate``.
overlap : int
Number of points to overlap between segments. If ``None``, ``noverlap = nperseg // 8``.
Defaults to ``None``.
analytical_signal : bool
If ``True``, analytical signal instead of actual signal is used in `Wigner Ville
Distribution` methods.
show : bool
If ``True``, will return two PSD plots.
Returns
-------
frequency : np.array
Frequency.
time : np.array
Time array.
stft : np.array
Short Term Fourier Transform. Time increases across its columns and frequency increases
down the rows.
Examples
-------
.. ipython:: python
import neurokit2 as nk
sampling_rate = 100
signal = nk.signal_simulate(100, sampling_rate, frequency=[3, 10])
# STFT Method
@savefig p_signal_timefrequency1.png scale=100%
f, t, stft = nk.signal_timefrequency(signal,
sampling_rate,
max_frequency=20,
method="stft",
show=True)
@suppress
plt.close()
.. ipython:: python
# CWTM Method
@savefig p_signal_timefrequency2.png scale=100%
f, t, cwtm = nk.signal_timefrequency(signal,
sampling_rate,
max_frequency=20,
method="cwt",
show=True)
@suppress
plt.close()
.. ipython:: python
# WVD Method
@savefig p_signal_timefrequency3.png scale=100%
f, t, wvd = nk.signal_timefrequency(signal,
sampling_rate,
max_frequency=20,
method="wvd",
show=True)
@suppress
plt.close()
.. ipython:: python
# PWVD Method
@savefig p_signal_timefrequency4.png scale=100%
f, t, pwvd = nk.signal_timefrequency(signal,
sampling_rate,
max_frequency=20,
method="pwvd",
show=True)
@suppress
plt.close()
"""
# Initialize empty container for results
# Define window length
if min_frequency == 0:
min_frequency = 0.04 # sanitize lowest frequency to lf
if max_frequency is None:
max_frequency = sampling_rate // 2 # nyquist
# STFT
if method.lower() in ["stft"]:
frequency, time, tfr = short_term_ft(
signal,
sampling_rate=sampling_rate,
overlap=overlap,
window=window,
mode=mode,
min_frequency=min_frequency,
window_type=window_type,
)
# CWT
elif method.lower() in ["cwt", "wavelet"]:
frequency, time, tfr = continuous_wt(
signal,
sampling_rate=sampling_rate,
min_frequency=min_frequency,
max_frequency=max_frequency,
)
# WVD
elif method in ["WignerVille", "wvd"]:
frequency, time, tfr = wvd(
signal,
sampling_rate=sampling_rate,
n_freqbins=nfreqbin,
analytical_signal=analytical_signal,
method="WignerVille",
)
# pseudoWVD
elif method in ["pseudoWignerVille", "pwvd"]:
frequency, time, tfr = wvd(
signal,
sampling_rate=sampling_rate,
n_freqbins=nfreqbin,
analytical_signal=analytical_signal,
method="pseudoWignerVille",
)
# Sanitize output
lower_bound = len(frequency) - len(frequency[frequency >= min_frequency])
f = frequency[(frequency >= min_frequency) & (frequency <= max_frequency)]
z = tfr[lower_bound : lower_bound + len(f)]
if show is True:
plot_timefrequency(
z,
time,
f,
signal=signal,
method=method,
)
return f, time, z
# =============================================================================
# Short-Time Fourier Transform (STFT)
# =============================================================================
def short_term_ft(
signal,
sampling_rate=1000,
min_frequency=0.04,
overlap=None,
window=None,
window_type="hann",
mode="psd",
):
"""Short-term Fourier Transform."""
if window is not None:
nperseg = int(window * sampling_rate)
else:
# to capture at least 5 times slowest wave-length
nperseg = int((2 / min_frequency) * sampling_rate)
frequency, time, tfr = scipy.signal.spectrogram(
signal,
fs=sampling_rate,
window=window_type,
scaling="density",
nperseg=nperseg,
nfft=None,
detrend=False,
noverlap=overlap,
mode=mode,
)
return frequency, time, np.abs(tfr)
# =============================================================================
# Continuous Wavelet Transform (CWT) - Morlet
# =============================================================================
def continuous_wt(
signal, sampling_rate=1000, min_frequency=0.04, max_frequency=None, nfreqbin=None
):
"""**Continuous Wavelet Transform**
References
----------
* Neto, O. P., Pinheiro, A. O., Pereira Jr, V. L., Pereira, R., Baltatu, O. C., & Campos, L.
A. (2016). Morlet wavelet transforms of heart rate variability for autonomic nervous system
activity. Applied and Computational Harmonic Analysis, 40(1), 200-206.
* Wachowiak, M. P., Wachowiak-Smolíková, R., Johnson, M. J., Hay, D. C., Power, K. E.,
& Williams-Bell, F. M. (2018). Quantitative feature analysis of continuous analytic wavelet
transforms of electrocardiography and electromyography. Philosophical Transactions of the
Royal Society A: Mathematical, Physical and Engineering Sciences, 376(2126), 20170250.
"""
# central frequency
w = 6.0 # recommended
if nfreqbin is None:
nfreqbin = sampling_rate // 2
# frequency
frequency = np.linspace(min_frequency, max_frequency, nfreqbin)
# time
time = np.arange(len(signal)) / sampling_rate
widths = w * sampling_rate / (2 * frequency * np.pi)
# Mother wavelet = Morlet
tfr = scipy.signal.cwt(signal, scipy.signal.morlet2, widths, w=w)
return frequency, time, np.abs(tfr)
# =============================================================================
# Wigner-Ville Distribution
# =============================================================================
def wvd(signal, sampling_rate=1000, n_freqbins=None, analytical_signal=True, method="WignerVille"):
"""Wigner Ville Distribution and Pseudo-Wigner Ville Distribution."""
# Compute the analytical signal
if analytical_signal:
signal = scipy.signal.hilbert(signal_detrend(signal))
# Pre-processing
if n_freqbins is None:
n_freqbins = 256
if method in ["pseudoWignerVille", "pwvd"]:
fwindows = np.zeros(n_freqbins + 1)
fwindows_mpts = len(fwindows) // 2
windows_length = n_freqbins // 4
windows_length = windows_length - windows_length % 2 + 1
windows = np.hamming(windows_length)
fwindows[fwindows_mpts + np.arange(-windows_length // 2, windows_length // 2)] = windows
else:
fwindows = np.ones(n_freqbins + 1)
fwindows_mpts = len(fwindows) // 2
time = np.arange(len(signal)) * 1.0 / sampling_rate
# This is discrete frequency (should we return?)
if n_freqbins % 2 == 0:
frequency = np.hstack((np.arange(n_freqbins / 2), np.arange(-n_freqbins / 2, 0)))
else:
frequency = np.hstack(
(np.arange((n_freqbins - 1) / 2), np.arange(-(n_freqbins - 1) / 2, 0))
)
tfr = np.zeros((n_freqbins, time.shape[0]), dtype=complex) # the time-frequency matrix
tausec = round(n_freqbins / 2.0)
winlength = tausec - 1
# taulens: len of tau for each step
taulens = np.min(
np.c_[
np.arange(signal.shape[0]),
signal.shape[0] - np.arange(signal.shape[0]) - 1,
winlength * np.ones(time.shape),
],
axis=1,
)
conj_signal = np.conj(signal)
# iterate and compute the wv for each indices
for idx in range(time.shape[0]):
tau = np.arange(-taulens[idx], taulens[idx] + 1).astype(int)
# this step is required to use the efficient DFT
indices = np.remainder(n_freqbins + tau, n_freqbins).astype(int)
tfr[indices, idx] = (
fwindows[fwindows_mpts + tau] * signal[idx + tau] * conj_signal[idx - tau]
)
if (idx < signal.shape[0] - tausec) and (idx >= tausec + 1):
tfr[tausec, idx] = (
fwindows[fwindows_mpts + tausec]
* signal[idx + tausec]
* np.conj(signal[idx - tausec])
+ fwindows[fwindows_mpts - tausec]
* signal[idx - tausec]
* conj_signal[idx + tausec]
)
tfr[tausec, idx] *= 0.5
# Now tfr contains the product of the signal segments and its conjugate.
# To find wd we need to apply fft one more time.
tfr = np.fft.fft(tfr, axis=0)
tfr = np.real(tfr)
# continuous time frequency
frequency = 0.5 * np.arange(n_freqbins, dtype=float) / n_freqbins * sampling_rate
return frequency, time, tfr
# =============================================================================
# Smooth Pseudo-Wigner-Ville Distribution
# =============================================================================
def smooth_pseudo_wvd(
signal,
sampling_rate=1000,
freq_length=None,
time_length=None,
segment_step=1,
nfreqbin=None,
window_method="hamming",
):
"""**Smoothed Pseudo Wigner Ville Distribution**
Parameters
----------
signal : Union[list, np.array, pd.Series]
The signal (i.e., a time series) in the form of a vector of values.
sampling_rate : int
The sampling frequency of the signal (in Hz, i.e., samples/second).
freq_length : np.ndarray
Lenght of frequency smoothing window.
time_length: np.array
Lenght of time smoothing window
segment_step : int
The step between samples in ``time_array``. Default to 1.
nfreqbin : int
Number of Frequency bins.
window_method : str
Method used to create smoothing windows. Can be "hanning"/ "hamming" or "gaussian".
Returns
-------
frequency_array : np.ndarray
Frequency array.
time_array : np.ndarray
Time array.
pwvd : np.ndarray
SPWVD. Time increases across its columns and frequency increases
down the rows.
References
----------
* J. M. O' Toole, M. Mesbah, and B. Boashash, (2008), "A New Discrete Analytic Signal for
Reducing Aliasing in the Discrete Wigner-Ville Distribution", IEEE Trans.
"""
# Define parameters
N = len(signal)
# sample_spacing = 1 / sampling_rate
if nfreqbin is None:
nfreqbin = 300
# Zero-padded signal to length 2N
signal_padded = np.append(signal, np.zeros_like(signal))
# DFT
signal_fft = np.fft.fft(signal_padded)
signal_fft[1 : N - 1] = signal_fft[1 : N - 1] * 2
signal_fft[N:] = 0
# Inverse FFT
signal_ifft = np.fft.ifft(signal_fft)
signal_ifft[N:] = 0
# Make analytic signal
signal = scipy.signal.hilbert(signal_detrend(signal_ifft))
# Create smoothing windows in time and frequency
if freq_length is None:
freq_length = np.floor(N / 4.0)
# Plus one if window length is not odd
if freq_length % 2 == 0:
freq_length += 1
elif len(freq_length) % 2 == 0:
raise ValueError("The length of frequency smoothing window must be odd.")
if time_length is None:
time_length = np.floor(N / 10.0)
# Plus one if window length is not odd
if time_length % 2 == 0:
time_length += 1
elif len(time_length) % 2 == 0:
raise ValueError("The length of time smoothing window must be odd.")
if window_method == "hamming":
freq_window = scipy.signal.hamming(int(freq_length)) # normalize by max
time_window = scipy.signal.hamming(int(time_length)) # normalize by max
elif window_method == "gaussian":
std_freq = freq_length / (6 * np.sqrt(2 * np.log(2)))
freq_window = scipy.signal.gaussian(freq_length, std_freq)
freq_window /= max(freq_window)
std_time = time_length / (6 * np.sqrt(2 * np.log(2)))
time_window = scipy.signal.gaussian(time_length, std_time)
time_window /= max(time_window)
# to add warning if method is not one of the supported methods
# Mid-point index of windows
midpt_freq = (len(freq_window) - 1) // 2
midpt_time = (len(time_window) - 1) // 2
# Create arrays
time_array = np.arange(start=0, stop=N, step=segment_step, dtype=int) / sampling_rate
# frequency_array = np.fft.fftfreq(nfreqbin, sample_spacing)[0:nfreqbin / 2]
frequency_array = 0.5 * np.arange(nfreqbin, dtype=float) / N
pwvd = np.zeros((nfreqbin, len(time_array)), dtype=complex)
# Calculate pwvd
for i, t in enumerate(time_array):
# time shift
tau_max = np.min(
[t + midpt_time - 1, N - t + midpt_time, np.round(N / 2.0) - 1, midpt_freq]
)
# time-lag list
tau = np.arange(
start=-np.min([midpt_time, N - t]), stop=np.min([midpt_time, t - 1]) + 1, dtype="int"
)
time_pts = (midpt_time + tau).astype(int)
g2 = time_window[time_pts]
g2 = g2 / np.sum(g2)
signal_pts = (t - tau - 1).astype(int)
# zero frequency
pwvd[0, i] = np.sum(g2 * signal[signal_pts] * np.conjugate(signal[signal_pts]))
# other frequencies
for m in range(int(tau_max)):
tau = np.arange(
start=-np.min([midpt_time, N - t - m]),
stop=np.min([midpt_time, t - m - 1]) + 1,
dtype="int",
)
time_pts = (midpt_time + tau).astype(int)
g2 = time_window[time_pts]
g2 = g2 / np.sum(g2)
signal_pt1 = (t + m - tau - 1).astype(int)
signal_pt2 = (t - m - tau - 1).astype(int)
# compute positive half
rmm = np.sum(g2 * signal[signal_pt1] * np.conjugate(signal[signal_pt2]))
pwvd[m + 1, i] = freq_window[midpt_freq + m + 1] * rmm
# compute negative half
rmm = np.sum(g2 * signal[signal_pt2] * np.conjugate(signal[signal_pt1]))
pwvd[nfreqbin - m - 1, i] = freq_window[midpt_freq - m + 1] * rmm
m = np.round(N / 2.0)
if t <= N - m and t >= m + 1 and m <= midpt_freq:
tau = np.arange(
start=-np.min([midpt_time, N - t - m]),
stop=np.min([midpt_time, t - 1 - m]) + 1,
dtype="int",
)
time_pts = (midpt_time + tau + 1).astype(int)
g2 = time_window[time_pts]
g2 = g2 / np.sum(g2)
signal_pt1 = (t + m - tau).astype(int)
signal_pt2 = (t - m - tau).astype(int)
x = np.sum(g2 * signal[signal_pt1] * np.conjugate(signal[signal_pt2]))
x *= freq_window[midpt_freq + m + 1]
y = np.sum(g2 * signal[signal_pt2] * np.conjugate(signal[signal_pt1]))
y *= freq_window[midpt_freq - m + 1]
pwvd[m, i] = 0.5 * (x + y)
pwvd = np.real(np.fft.fft(pwvd, axis=0))
# Visualization
return frequency_array, time_array, pwvd
# =============================================================================
# Plot function
# =============================================================================
def plot_timefrequency(z, time, f, signal=None, method="stft"):
"""Visualize a time-frequency matrix."""
if method == "stft":
figure_title = "Short-time Fourier Transform Magnitude"
fig, ax = plt.subplots()
for i in range(len(time)):
ax.plot(f, z[:, i], label="Segment" + str(np.arange(len(time))[i] + 1))
ax.legend()
ax.set_title("Signal Spectrogram")
ax.set_ylabel("STFT Magnitude")
ax.set_xlabel("Frequency (Hz)")
elif method == "cwt":
figure_title = "Continuous Wavelet Transform Magnitude"
elif method == "wvd":
figure_title = "Wigner Ville Distrubution Spectrogram"
fig = plt.figure()
plt.plot(time, signal)
plt.xlabel("Time (sec)")
plt.ylabel("Signal")
elif method == "pwvd":
figure_title = "Pseudo Wigner Ville Distribution Spectrogram"
fig, ax = plt.subplots()
spec = ax.pcolormesh(time, f, z, cmap=plt.get_cmap("magma"), shading="auto")
plt.colorbar(spec)
ax.set_title(figure_title)
ax.set_ylabel("Frequency (Hz)")
ax.set_xlabel("Time (sec)")
return fig
