Python中FIR滤波和小波包滤波对比（MNE脑电数据处理）

小波变换有信号显微镜之称，在EEG分析中也有广泛的应用，印象中小波算法是来源于地球物理解释的。之前有介绍过小波的一些资料和实现：https://blog.csdn.net/zhoudapeng01/article/details/107025901可以参考下，这里主要分析小波和FIR滤波效果的对比。

博客对应的代码和数据

https://download.csdn.net/download/zhoudapeng01/12793085

# 短时傅里叶变换和FIR滤波效果对比

import mne
import matplotlib.pyplot as plt
from scipy import signal, fft
import numpy as np
import pywt
# 设置MNE库打印Log的级别
mne.set_log_level(False)
# 需要分析的频带及其范围
bandFreqs = [
    {'name': 'Delta', 'fmin': 1, 'fmax': 3},
    {'name': 'Theta', 'fmin': 4, 'fmax': 7},
    {'name': 'Alpha', 'fmin': 8, 'fmax': 13},
    {'name': 'Beta', 'fmin': 14, 'fmax': 31},
    {'name': 'Gamma', 'fmin': 31, 'fmax': 40}
]




def __CalcWP(data, sfreq, wavelet, maxlevel, band):
    # 如果maxlevel太小部分波段分析不到
    wp = pywt.WaveletPacket(data=data, wavelet=wavelet, mode='symmetric', maxlevel=maxlevel)
    # 频谱由低到高的对应关系，这里需要注意小波变换的频带排列默认并不是顺序排列，所以这里需要使用’freq‘排序。
    freqTree = [node.path for node in wp.get_level(maxlevel, 'freq')]
    # 计算maxlevel最小频段的带宽，采样频率的一半
    freqBand = (sfreq/2) / (2 ** maxlevel)
    bandResult = []
    #######################根据实际情况计算频谱对应关系，这里要注意系数的顺序
    for iter_freq in band:
        # 构造空的小波包
        new_wp = pywt.WaveletPacket(data=None, wavelet=wavelet, mode='symmetric', maxlevel=maxlevel)
        for i in range(len(freqTree)):
            # 第i个频段的最小频率
            bandMin = i * freqBand
            # 第i个频段的最大频率
            bandMax = bandMin + freqBand
            # 判断第i个频段是否在要分析的范围内
            if (iter_freq['fmin'] <= bandMin and iter_freq['fmax'] >= bandMax):
                # 给新构造的小波包参数赋值
                # print('freq',bandMin, bandMax,'fmin',iter_freq['fmin'],'fmax',iter_freq['fmax'])
                new_wp[freqTree[i]] = wp[freqTree[i]].data
        # 计算对应频率的数据
        bandResult.append(new_wp.reconstruct(update=True))
    return bandResult

########################################小波包变换-重构造分析不同频段的特征(注意maxlevel，如果太小可能会导致部分频段分析不到)#########################
# 定义WP函数
# epochsData:epochs的数据（mumpy格式）
# sfreq:采样频率
# wavelet:小波类型
# maxlevel:小波层数
# band:频带类型

def WP(epochsData, sfreq, wavelet='db4', maxlevel=8, band=bandFreqs):
    # 输出的维度顺序为 频率->epoch->channel->timeData
    result = []
    for epochData in epochsData:
        channel = []
        for channelData in epochData:
            # print('channel:')
            channel.append(__CalcWP(channelData, sfreq, wavelet=wavelet, maxlevel=maxlevel, band=band))
        result.append(channel)
    return np.array(result).transpose((2, 0, 1, 3))



if __name__ == '__main__':
    # 加载fif格式的数据
    epochs = mne.read_epochs(r'F:\BaiduNetdiskDownload\BCICompetition\BCICIV_2a_gdf\Train\Fif\A02T_epo.fif')
    # 绘图验证结果
    plt.figure(figsize=(15, 10))
    # 获取采样频率
    sfreq = epochs.info['sfreq']
    # 想要分析的目标频带
    bandIndex = 0
    # 想要分析的channel
    channelIndex = 0
    # 想要分析的epoch
    epochIndex = 0
    # 绘制原始数据
    plt.plot(epochs.get_data()[epochIndex][channelIndex], label='Raw')
    # 计算FIR滤波后的数据并绘图（注意这里要使用copy方法，否则会改变原始数据）
    firFilter = epochs.copy().filter(bandFreqs[bandIndex]['fmin'], bandFreqs[bandIndex]['fmax'])
    plt.plot(firFilter.get_data()[epochIndex][channelIndex], c=(1, 0, 0), label='FIR_Filter')
    # 计算小波包滤波后的数据并绘图
    wpFilter = WP(epochs.get_data(), sfreq)
    plt.plot(wpFilter[bandIndex][epochIndex][channelIndex], c=(0, 1, 0), label='WP_Filter')
    # 绘制图例和图名
    plt.legend()
    plt.title(bandFreqs[bandIndex]['name'])


    ####################################FFT对比两种方法的频谱分布
    plt.figure(figsize=(15, 10))
    # 对FIR滤波后的数据进行FFT变换
    mneFIRFreq = np.abs(fft.fft(firFilter.get_data()[epochIndex][channelIndex]))
    # 对小波包滤波后的数据进行FFT变换，需要注意小波包变换后数据的点数可能会发生变化，这里截取数据保持一致性
    pointNum = epochs.get_data()[epochIndex][channelIndex].shape[0]
    wpFreq = np.abs(fft.fft(wpFilter[bandIndex][epochIndex][channelIndex][:pointNum]))
    # 想要绘制的点数
    pointPlot = 300
    # FIR滤波后x轴对应的频率幅值范围
    FIR_X = np.linspace(0, sfreq/2, int(mneFIRFreq.shape[0]/2))
    # 小波包滤波后x轴对应的频率幅值范围
    WP_X = np.linspace(0, sfreq/2, int(wpFreq.shape[0]/2))
    # 绘制FIR滤波后的频谱分布
    plt.plot(FIR_X[:pointPlot], mneFIRFreq[:pointPlot], c=(1, 0, 0), label='FIR_Filter')
    # 绘制小波包滤波后的频谱分布
    plt.plot(WP_X[:pointPlot], wpFreq[:pointPlot], c=(0, 1, 0), label='WP_FIlter')
    # 绘制图例和图名
    plt.legend()
    plt.title(bandFreqs[bandIndex]['name'])
    plt.show()

注：小波包滤波后数据长度会发生改变，这主要和小波层数以及计算方式有关，超出原始长度的数据可以不用关心，上面的代码中在进行频谱分析前，也就是计算FFT前对数据的长度进行了处理，这样可以保证分析出频谱数据的一致性。

WP和FIR在不同频段滤波效果对比：

从滤波的结果来看，小波包的滤波效果在低频段还不错，可是到了高频段其效果一般。

https://blog.csdn.net/zhoudapeng01/article/details/108214064（FIR和STFT对比）

对比FIR、STFT以及WP这三种方法，STFT相对效果更好些。

Delta频段（1Hz-3Hz）对应的结果

Theta频段（4Hz-7Hz）对应的结果

Alpha频段（8Hz-13Hz）对应的结果

Beta频段（14Hz-31Hz）对应的结果

Gamma频段（31Hz-40Hz）对应的结果

对应的资源包下载：https://download.csdn.net/download/zhoudapeng01/12793085