Page MenuHomec4science
Files F111601875

Utils_FFT.py
No OneTemporary
Actions

File Metadata

Created
Sat, May 3, 17:10

Utils_FFT.py
View Options

import pandas as pd
import glob
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib import cm
import ntpath
import re
from scipy import signal
import pywt
import seaborn as sns
from sklearn import preprocessing
import os
from matplotlib import mlab
import matplotlib.ticker as ticker
from matplotlib.ticker import FormatStrFormatter
#%%
def filter(signal_window, sample_rate):
"""
Apply a low-pass filter to the input signal window.
Args:
signal_window (array-like): The input signal window to be filtered.
sample_rate (float): The sample rate of the input signal.
Returns:
array-like: The filtered signal window.
"""
lowpass = 0.49 * sample_rate # Cut-off frequency of the filter
lowpass_freq = lowpass / (sample_rate / 2) # Normalize the frequency
b, a = signal.butter(5, lowpass_freq, 'low')
signal_window = signal.filtfilt(b, a, signal_window)
#plt.plot(t, lowpassfilter, label='Low-Pass')
return signal_window
def filterband(signal_window, sample_rate):
"""
Applies a low-pass filter to the given signal window.
Args:
signal_window (array-like): The input signal window to be filtered.
sample_rate (float): The sample rate of the signal.
Returns:
array-like: The filtered signal window.
"""
lowpass = 0.49 * sample_rate # Cut-off frequency of the filter
lowpass_freq = lowpass / (sample_rate / 2) # Normalize the frequency
b, a = signal.butter(5, lowpass_freq, 'low')
signal_window = signal.filtfilt(b, a, signal_window)
highpass = 0.025 * sample_rate # Normalize the frequency
highpass_freq = highpass / (sample_rate / 2)
b, a = signal.butter(5, highpass_freq, 'high')
signal_window = signal.filtfilt(b, a, signal_window)
#plt.plot(t, lowpassfilter, label='Low-Pass')
return signal_window
#%%
def STFT(data_new,time,scales,plotname,windowsize,path,lowpass):
plt.rcParams.update(plt.rcParamsDefault)
NFFT = 5000
dt = time[1] - time[0]
Fs = int(1.0 / dt)
signal_window=data_new.transpose()
fig, ax = plt.subplots(figsize=(12, 7))
# Pxx, freqs, bins, im = ax.specgram(signal_window, NFFT=NFFT, Fs=Fs, noverlap=4500)
Pxx, freqs, bins, im = ax.specgram(signal_window, Fs=Fs,cmap=cm.coolwarm)
plt.rcParams['agg.path.chunksize'] = windowsize
fig.patch.set_visible(True)
ax.axis('on')
ax.tick_params(axis='both', which='major', labelsize=20)
ax.tick_params(axis='both', which='minor', labelsize=20)
ax.ticklabel_format(axis='x', style='sci',scilimits=(0,0))
ax.ticklabel_format(axis='y', style='sci',scilimits=(0,0))
ax.yaxis.offsetText.set_fontsize(20)
ax.xaxis.offsetText.set_fontsize(20)
ax.set_ylim(0, lowpass)
cb=plt.colorbar(im)
cb.set_label(label='Intensity',fontsize=20)
cb.ax.tick_params(labelsize=20)
plottitle=('Spectrogram')
plt.suptitle(plottitle, fontsize=20)
plt.xlabel('Time(sec)',fontsize=20)
plt.ylabel('Frequency(Hz)',fontsize=20)
graphname=str(plotname)+'_Spectrogram'+'.png'
plt.savefig(os.path.join(path, graphname), bbox_inches='tight')
plt.show()
plt.clf()
#%%
def Frequencyplot(rawspace,sample_rate,plotname,windowsize,path,lowpass):
"""
Plot the frequency spectrum of a given signal.
Args:
rawspace (array-like): The input signal.
sample_rate (int): The sample rate of the signal.
plotname (str): The name of the plot.
windowsize (int): The chunk size for calculating the spectrum.
path (str): The path to save the plot.
lowpass (int): The maximum frequency to display on the plot.
Returns:
None
"""
win = 4 * sample_rate
freqs, psd = signal.welch(rawspace, sample_rate, nperseg=win)
sns.set(font_scale=1, style='white')
fig,ax = plt.subplots(figsize=(12, 7))
plt.rcParams['agg.path.chunksize'] = windowsize
ax.plot(freqs, psd, color='k', lw=0.001)
sec1, sec2 = 0, 150000
sec3, sec4 = 150000, 300000
sec5, sec6 = 300000, 450000
sec7, sec8 = 450000, 600000
sec9, sec10 = 600000, 750000
idx_delta1 = np.logical_and(freqs >= sec1, freqs <= sec2)
idx_delta2 = np.logical_and(freqs >= sec3, freqs <= sec4)
idx_delta3 = np.logical_and(freqs >= sec5, freqs <= sec6)
idx_delta4 = np.logical_and(freqs >= sec7, freqs <= sec8)
idx_delta5 = np.logical_and(freqs >= sec9, freqs <= sec10)
ylimit=psd.max() + 0.1*psd.max()
plt.ylim([0, ylimit])
ax.fill_between(freqs, psd, where=idx_delta1, color='blue')
ax.fill_between(freqs, psd, where=idx_delta2, color='red')
ax.fill_between(freqs, psd, where=idx_delta3, color='yellow')
ax.fill_between(freqs, psd, where=idx_delta4, color='green')
ax.fill_between(freqs, psd, where=idx_delta5, color='orange')
filename='Frequency Plot'
plt.suptitle(filename, fontsize=25)
# plt.title(filename)
# plt.xlim([0, freqs.max()])
plt.xlim([0, lowpass])
skyblue = mpatches.Patch(color='blue', label='0-150 kHz')
ax.legend(handles=[skyblue])
red = mpatches.Patch(color='red', label='150-300 kHz')
ax.legend(handles=[red])
yellow = mpatches.Patch(color='yellow', label='300-450 kHz')
ax.legend(handles=[yellow])
green = mpatches.Patch(color='green', label='450-600 kHz')
ax.legend(handles=[green])
cyan = mpatches.Patch(color='orange', label='600-750 kHz')
ax.legend(handles=[skyblue,red,yellow,green,cyan], prop={'size': 20})
ax.tick_params(axis='both', which='major', labelsize=25)
ax.tick_params(axis='both', which='minor', labelsize=25)
ax.ticklabel_format(axis='x', style='sci',scilimits=(0,0))
ax.ticklabel_format(axis='y', style='sci',scilimits=(0,0))
ax.yaxis.offsetText.set_fontsize(20)
ax.xaxis.offsetText.set_fontsize(20)
plt.xlabel('Frequency (Hz)',fontsize=20)
plt.ylabel('Power spectral density',fontsize=20)
graphname=str(plotname)+'_Frequency'+'.png'
plt.savefig(os.path.join(path, graphname), bbox_inches='tight',dpi=100)
plt.show()
plt.clf()
del freqs
del psd
#%%
def Specgram3d(data, sample_rate,plotname, time,path, ax=None, title=None):
"""
Generate a 3D spectrogram plot.
Args:
data (array-like): The input data for spectrogram calculation.
sample_rate (int): The sample rate of the input data.
plotname (str): The name of the plot.
time (float): The time duration of the input data.
path (str): The path to save the generated plot.
ax (Axes3D, optional): The matplotlib 3D axes object to plot on. If not provided, a new axes object will be created.
title (str, optional): The title of the plot.
Returns:
None
Raises:
None
"""
fig = plt.figure(figsize=(12, 12))
ax = plt.axes(projection='3d')
ax.view_init(azim=70, elev=40)
spec, freqs, t = mlab.specgram(data, Fs=sample_rate)
# X, Y, Z = t[None, :], freqs[:, None], 10*np.log10(spec)
X, Y, Z = t[None, :], freqs[:, None], spec
# surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm,
# linewidth=0, antialiased=False)
surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm,
linewidth=0, rstride=1, cstride=1, alpha=None, antialiased=True)
ax.set_xlabel('Time (sec)', labelpad=7)
ax.set_ylabel('Frequency (Hz)', labelpad=9)
ax.set_zlabel('Amplitude (dB)', labelpad=20)
# ax.set_zticklabels(ax.get_zticks(), fontdict={'ha':'left'})
ax.set_zticklabels(ax.get_zticks())
# ax.set_title('Surface plot')
ax.grid(False)
ax.set_facecolor('white')
num_ticks = 3
z_min, z_max = ax.get_zlim()
v = np.linspace(z_min, z_max, num_ticks)
ax.set_zticks(np.linspace(z_min, z_max, num_ticks))
print(z_min, z_max)
ax.set_zticks(np.linspace(z_min, z_max, num_ticks))
# ax.zaxis.offset_text_position= "top"
ax.set_zticklabels(ax.get_zticks(), fontdict={'ha':'right'})
ax.zaxis.set_major_formatter(FormatStrFormatter('%.2e'))
plt.gca().invert_xaxis()
ax.xaxis.pane.fill = False
ax.yaxis.pane.fill = False
ax.zaxis.pane.fill = False
fig.colorbar(surf, ax=ax,ticks=v,
shrink=0.2, aspect=4)
graphname = str(plotname)+'_3D_specgram3d.png'
plt.savefig(os.path.join(path, graphname), bbox_inches='tight', dpi=100)
plt.show()
# return X, Y, Z

Event Timeline

c4science · Help