MCC5-THU齿轮箱多模式故障数据集(含变转速、变负载、轴承齿轮复合故障)

## 导入包from matplotlib import pyplot as pltfrom matplotlib import rcParamsimport numpy as npimport pandas as pdimport os
config = {    "font.family": 'serif', # 衬线字体    "font.size": 14, # 相当于小四大小    "font.serif": ['SimSun'], # 宋体    "mathtext.fontset": 'stix', # matplotlib渲染数学字体时使用的字体，和Times New Roman差别不大    'axes.unicode_minus': False # 处理负号，即-号}rcParams.update(config)
def data_read(file_path):    """    :fun: 读取csv数据    :param file_path: 文件路径    :return accl_data_df: Dataframe，八列分别为'speed', 'torque', 'motor_vibration_x', 'motor_vibration_y', 'motor_vibration_z', 'gearbox_vibration_x', 'gearbox_vibration_y', 'gearbox_vibration_z'    """    file_path = os.path.join(base_dir, file_name)    accl_data_df = pd.read_csv(file_path)    return accl_data_df    file_path = r'E:\03-公开数据集\MCC5-THU gearbox fault diagnosis datasets with variable working conditions\原始数据/gear_pitting_M_torque_circulation_3000rpm_20Nm.csv'df = data_read(file_path)acc_z_arr= df['gearbox_vibration_z']  # 选择变速箱z轴数据df

##========绘制时域信号图========##def plt_time_domain(arr, fs=1600, ylabel='Amp($m/s^2$)', title='原始数据时域图', img_save_path=None, x_vline=None, y_hline=None):    """    :fun: 绘制时域图模板    :param arr: 输入一维数组数据    :param fs: 采样频率    :param ylabel: y轴标签    :param title: 图标题    :return: None    """    import matplotlib.pyplot as plt    plt.rcParams['font.sans-serif'] = ['SimHei']  # 显示中文    plt.rcParams['axes.unicode_minus'] = False  # 显示负号    font = {'family': 'Times New Roman', 'size': '20', 'color': '0.5', 'weight': 'bold'}
    plt.figure(figsize=(12,4))    length = len(arr)    t = np.linspace(0, length/fs, length)    plt.plot(t, arr, c='g')    plt.xlabel('t(s)')    plt.ylabel(ylabel)    plt.title(title)    if x_vline:        plt.vlines(x=x_vline, ymin=np.min(arr), ymax=np.max(arr), linestyle='--', colors='r')    if y_hline:        plt.hlines(y=0.2, xmin=np.min(t), xmax=np.max(t), linestyle=':', colors='y')    #===保存图片====#    if img_save_path:        plt.savefig(img_save_path, dpi=500, bbox_inches = 'tight')    plt.show()

fs = 12800file_name = r'miss_teeth_torque_circulation_3000rpm_20Nm.csv'##=====绘制时域数据====##time_img_save_path = file_path.replace('.csv', '时域_.png')plt_time_domain(acc_z_arr, fs=fs, title=file_name, img_save_path=time_img_save_path, ylabel = 'Amp(g)')

##========绘制频域信号图========##def plt_fft_img(arr, fs, ylabel='Amp(mg)', title='频域图', img_save_path=None, vline=None, hline=None, xlim=None):    """    :fun: 绘制频域图模板    :param arr: 输入一维时域数组数据    :param fs: 采样频率    :param ylabel: y轴标签    :param title: 图标题    :return: None    """    # 计算频域幅值    length = len(arr)    t = np.linspace(0, length/fs, length)    fft_result = np.fft.fft(arr)    fft_freq= np.fft.fftfreq(len(arr), d=t[1]-t[0])  # FFT频率    fft_amp= 2*np.abs(fft_result)/len(t)                     # FFT幅值
    # 绘制频域图    plt.figure(figsize=(12,4))    plt.title(title)    plt.plot(fft_freq[0: int(len(t)/2)], fft_amp[0: int(len(t)/2)], label='Frequency Spectrum', color='b')    plt.xlabel('频率 (Hz)')    plt.ylabel('幅值')    plt.legend()    if vline:        plt.vlines(x=vline, ymin=np.min(fft_amp), ymax=np.max(fft_amp), linestyle='--', colors='r')    if hline:        plt.hlines(y=hline, xmin=np.min(fft_freq), xmax=np.max(fft_freq), linestyle=':', colors='y')    #===保存图片====#    if img_save_path:        plt.savefig(img_save_path, dpi=500, bbox_inches = 'tight')    if xlim: # 图片横坐标是否设置xlim        plt.xlim(0, xlim)      plt.tight_layout()    plt.show()

fs = 12800file_name = r'miss_teeth_torque_circulation_3000rpm_20Nm.csv'##=====绘制频域数据====##fft_img_save_path = file_path.replace('.csv', '频域_.png')plt_fft_img(acc_z_arr, fs=fs, title=file_name, img_save_path=fft_img_save_path)

def plt_stft_img(arr, fs, ylabel='Amp(mg)', title='频域图', img_save_path=None, vline=None, hline=None, xlim=None):    """    :fun: 绘制stft时频域图模板    :param arr: 输入一维时域数组数据    :param fs: 采样频率    :param ylabel: y轴标签    :param title: 图标题    :return: None    """    import scipy.signal as signal    import numpy as np    import matplotlib.pyplot as plt
    f, t, nd = signal.stft(arr, fs=fs, window='hann', nperseg=128, noverlap=64,nfft=None,                           detrend=False, return_onesided=True, boundary='odd', padded=False, axis=-1)    #  fs:时间序列的采样频率,  nperseg:每个段的长度，默认为256(2^n)   noverlap:段之间重叠的点数。如果没有则noverlap=nperseg/2
    #window ：字符串或元组或数组，可选需要使用的窗。    # #如果window是一个字符串或元组，则传递给它window是数组类型，直接以其为窗，其长度必须是nperseg。    # 常用的窗函数有boxcar，triang，hamming， hann等，默认为Hann窗。
    #nfft ：int，可选。如果需要零填充FFT，则为使用FFT的长度。如果为 None，则FFT长度为nperseg。默认为无
    # detrend ：str或function或False，可选    # 指定如何去除每个段的趋势。如果类型参数传递给False，则不进行去除趋势。默认为False。
    # return_onesided ：bool，可选    # 如果为True，则返回实际数据的单侧频谱。如果 False返回双侧频谱。默认为 True。请注意，对于复杂数据，始终返回双侧频谱。
    # boundary ：str或None，可选    # 指定输入信号是否在两端扩展，以及如何生成新值，以使第一个窗口段在第一个输入点上居中。    # 这具有当所采用的窗函数从零开始时能够重建第一输入点的益处。    # 有效选项是['even', 'odd', 'constant', 'zeros', None].    # 默认为‘zeros’,对于补零操作[1, 2, 3, 4]变成[0, 1, 2, 3, 4, 0] 当nperseg=3.
    # padded：bool，可选    # 指定输入信号在末尾是否填充零以使信号精确地拟合为整数个窗口段，以便所有信号都包含在输出中。默认为True。    # 填充发生在边界扩展之后，如果边界不是None，则填充为True，默认情况下也是如此。    # axis ：int，可选
    # 绘制STFT时频域图    plt.figure(figsize=(12,4))    plt.pcolormesh(t, f, np.abs(nd), vmin = np.min(np.abs(nd)), vmax = np.max(np.abs(nd)))    plt.title(title)    plt.xlabel('时间（t）')    plt.ylabel('频率 (Hz)')    if vline:        plt.vlines(x=vline, ymin=np.min(fft_amp), ymax=np.max(fft_amp), linestyle='--', colors='r')    if hline:        plt.hlines(y=hline, xmin=np.min(fft_freq), xmax=np.max(fft_freq), linestyle=':', colors='y')    #===保存图片====#    if img_save_path:        plt.savefig(img_save_path, dpi=500, bbox_inches = 'tight')    if xlim: # 图片横坐标是否设置xlim        plt.xlim(0, xlim)      plt.tight_layout()    plt.show()

fs = 12800file_name = r'miss_teeth_torque_circulation_3000rpm_20Nm.csv'##=====绘制STFT时频域数据====##stft_img_save_path = file_path.replace('.csv', '时频域_.png')    plt_stft_img(acc_z_arr, fs=fs, img_save_path=stft_img_save_path)