对于EEG信号SWT与Butterworth滤波器比较

import mne
import numpy as np
import scipy.io as sio
import matplotlib.pyplot as plt
from mne.decoding import CSP
from sklearn.svm import SVC
from sklearn.pipeline import Pipeline
from sklearn.model_selection import ShuffleSplit, KFold
from sklearn.model_selection import cross_val_score
import pywt
import math


load_path = "P:/Item/EEG/bci2003/dataset_BCIcomp1.mat"
load_data = sio.loadmat(load_path)
eeg_data = np.array(load_data["x_train"]).T
label = np.array(load_data["y_train"])
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}
]
#sgn函数
def sgn(num):
    if(num > 0.0):
        return 1.0
    elif(num ` 0.0):
        return 0.0
    else:
        return -1.0

def wavelet_noising(new_df):
    data = new_df
    data = data.T.tolist()  # 将np.ndarray()转为列表
    w = pywt.Wavelet('dB3')#选择dB3小波基
    maxlev = pywt.dwt_max_level(len(data), w.dec_len)
    ca3, cd3, cd2, cd1 = pywt.wavedec(data, w, level=3)  # 3层小波分解

    length1 = len(cd1)
    length0 = len(data)

    abs_cd1 = np.abs(np.array(cd1))
    median_cd1 = np.median(abs_cd1)

    sigma = (1.0 / 0.6745) * median_cd1
    lamda = sigma * math.sqrt(2.0 * math.log(float(length0 ), math.e))
    usecoeffs = []
    usecoeffs.append(ca3)
    #软阈值方法
    for k in range(length1):
        if (abs(cd1[k]) >= lamda/np.log2(2)):
            cd1[k] = sgn(cd1[k]) * (abs(cd1[k]) - lamda/np.log2(2))
        else:
            cd1[k] = 0.0

    length2 = len(cd2)
    for k in range(length2):
        if (abs(cd2[k]) >= lamda/np.log2(3)):
            cd2[k] = sgn(cd2[k]) * (abs(cd2[k]) - lamda/np.log2(3))
        else:
            cd2[k] = 0.0

    length3 = len(cd3)
    for k in range(length3):
        if (abs(cd3[k]) >= lamda/np.log2(4)):
            cd3[k] = sgn(cd3[k]) * (abs(cd3[k]) - lamda/np.log2(4))
        else:
            cd3[k] = 0.0

    usecoeffs.append(cd3)
    usecoeffs.append(cd2)
    usecoeffs.append(cd1)
    recoeffs = pywt.waverec(usecoeffs, w)#信号重构
    return recoeffs

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='db3', 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))
'''
构建一个Raw对象时,需要准备两种数据,一种是data数据,一种是Info数据,
data数据是一个二维数据,形状为(n_channels,n_times)
'''
ch_names = ['C3','Cz','C4']  # 通道名称
ch_types = ['eeg', 'eeg', 'eeg']  # 通道类型
sfreq = 128  # 采样率
info = mne.create_info(ch_names, sfreq, ch_types)  # 创建信号的信息
info.set_montage('standard_1020')

raw_0 = eeg_data[0, :, :]
for i in range(1, 140):
    raw_i = eeg_data[i, :, :]
    raw_0 = np.concatenate((raw_0, raw_i), axis=1)
raw_data = raw_0
#raw_data = wavelet_noising(raw_data)
raw_data[0] = wavelet_noising(raw_data[0])
raw_data[1] = wavelet_noising(raw_data[1])
raw_data[2] = wavelet_noising(raw_data[2])
# print(raw_data.shape)
# >>>(3, 161280)
raw = mne.io.RawArray(raw_data, info)
# print(raw)
# raw.plot(scalings={'eeg': 'auto'})
# print('数据集的形状为：',raw.get_data().shape)
# print('通道数为：',raw.info.get('nchan'))

# FIR带通滤波
#raw.filter(8., 30., fir_design='firwin', skip_by_annotation='edge')
'''
在创建Epochs对象时,必须提供一个"events"数组,
事件(event)描述的是某一种波形(症状)的起始点,其为一个三元组,形状为(n_events,3):
第一列元素以整数来描述的事件起始采样点;
第二列元素对应的是当前事件来源的刺激通道(stimulus channel)的先前值(previous value),该值大多数情况是0;
第三列元素表示的是该event的id.
'''
# 创建 events & event_id
events = np.zeros((140, 3), dtype='int')
k = sfreq * 3
for i in range(140):
    events[i, 0] = k
    k += sfreq * 9
    events[i, 2] = label[i]
# print(events)
event_id = dict(left_hand=1, right_hand=2)

# 创建epochs
tmin, tmax = -1., 4.  # 记录点的前1秒后4秒用于生成epoch数据
epochs = mne.Epochs(raw, events, event_id
                    , tmin, tmax
                    , proj=True
                    , baseline=(None, 0)
                    , preload=True)
# 想要分析的目标频带
#bandIndex = 4
firFilter = epochs.copy().filter(bandFreqs[1]['fmin'], bandFreqs[4]['fmin'])
# 计算小波包滤波后的数据并绘图
wpFilter = WP(epochs.get_data(), sfreq)
epochs.selection = wpFilter
#获取小波分解的最大层级
maxlev = pywt.swt_max_level(len(epochs.get_data()))
coeffs= pywt.swt(epochs.selection,'db3', level = maxlev) # 将信号进行小波分解
coeffs[0][0].fill(0)
coeffs[maxlev-1][1].fill(0)
datarec = pywt.iswt(coeffs,'db3') # 将信号进行小波重构




###使用滤波器
#epochs = firFilter
#epochs.selection = datarec




epochs_train = epochs.copy().crop(tmin=1., tmax=2.)  # 截取其中的1秒到2秒之间的数据，也就是提示音后
# 1秒到2秒之间的数据（这个在后面滑动窗口验证的时候有用）

labels = epochs.events[:, -1]
# print(labels)
#特征提取和分类
scores = []
epochs_data = epochs.get_data()                       #获取epochs的所有数据，主要用于后面的滑动窗口验证
#print(epochs_data.shape)
#>>>(140, 3, 641)
epochs_data_train = epochs_train.get_data()           #获取训练数据
#print(epochs_data_train.shape)
#>>>(140, 3, 129)                                     #也是多了一个
kf = KFold(n_splits=5                                 #交叉验证模型的参数
           , shuffle=True
           , random_state=42)
cv_split = kf.split(epochs_data_train)                #输出索引以将数据分为训练集和测试集
svm = SVC()                                           #支持向量机分类器
csp = CSP(n_components=2                              #2个分量的CSP
          ,reg=None
          ,log=False
          ,norm_trace=False)
clf = Pipeline([('CSP', csp), ('SVM', svm)])          #创建机器学习的Pipeline
scores = cross_val_score(clf                          #获取交叉验证模型的得分
                         , epochs_data_train
                         , labels
                         , cv=kf
                         , n_jobs=-1)
class_balance = np.mean(labels ` labels[0])          #输出结果，准确率和不同样本的占比
class_balance = max(class_balance, 1. - class_balance)
print("Classification accuracy: %f / Chance level: %f" % (np.mean(scores)
                                                          , class_balance))
csp.fit_transform(epochs_data, labels)
csp.plot_patterns(epochs.info                         #绘制CSP不同分量的模式图
                  , ch_type='eeg'
                  , units='Patterns (AU)'
                  , size=1.5)
# 验证算法的性能
w_length = int(sfreq * 1.5)  # 设置滑动窗口的长度
w_step = int(sfreq * 0.1)  # 设置滑动步长
w_start = np.arange(0, epochs_data.shape[2] - w_length, w_step)
# 每次滑动窗口的起始点
scores_windows = []  # 得分列表用于保存模型得分
# 交叉验证计算模型的性能
for train_idx, test_idx in cv_split:
    y_train, y_test = labels[train_idx], labels[test_idx]  # 获取测试集和训练集数据
    X_train = csp.fit_transform(epochs_data_train[train_idx], y_train)  # 设置csp模型的参数，提取相关特征，用于后面的svm分类
    svm.fit(X_train, y_train)  # 拟合svm模型
    score_this_window = []  # 用于记录本次交叉验证的得分
    for n in w_start:
        X_test = csp.transform(epochs_data[test_idx][:, :, n:(n + w_length)])  # csp提取测试数据相关特征
        score_this_window.append(svm.score(X_test, y_test))  # 获取测试数据得分
    scores_windows.append(score_this_window)  # 添加到总得分列表
w_times = (w_start + w_length / 2.) / sfreq + epochs.tmin  # 设置绘图的时间轴，时间轴上的标志点为窗口的中间位置

# 绘制模型分类结果的性能图
plt.figure()
plt.plot(w_times, np.mean(scores_windows, 0), label='Score')
plt.axvline(0, linestyle='--', color='k', label='Onset')
plt.axhline(0.5, linestyle='-', color='k', label='Chance')
plt.xlabel('time (s)')
plt.ylabel('classification accuracy')
plt.title('Classification score over time')
plt.legend(loc='lower right')
plt.show()