用Python和Matlab玩转东南大学齿轮箱数据集：从数据读取到故障分类实战

用Python和Matlab玩转东南大学齿轮箱数据集从数据读取到故障分类实战当你第一次打开东南大学齿轮箱数据集时面对几十个CSV文件和8通道的振动信号数据可能会感到无从下手。这份数据集包含了五种齿轮故障类型和两种工况是研究机械故障诊断的绝佳素材。本文将带你用Python和Matlab两种工具从数据清洗到特征提取最终构建一个能够自动识别齿轮故障类型的机器学习模型。1. 数据集理解与预处理东南大学齿轮箱数据集包含了五种齿轮故障状态齿缺损、断齿、根部裂纹、齿面磨损以及正常运行状态。每种故障类型又分为两种工况转速20Hz-负载0V和转速30Hz-负载2V这使得数据集具有很好的多样性。1.1 数据结构解析每个CSV文件包含8列数据分别对应电机振动信号行星齿轮x方向振动行星齿轮y方向振动行星齿轮z方向振动电机扭矩减速器x方向振动减速器y方向振动减速器z方向振动采样频率为5120Hz这意味着每秒钟记录了5120个数据点。理解这些通道的物理意义对后续特征提取至关重要。1.2 Python数据读取方法使用Python的pandas库可以高效读取这些CSV文件import pandas as pd import os def load_gear_data(folder_path): data_dict {} for file in os.listdir(folder_path): if file.endswith(.csv): # 从文件名解析故障类型和工况 fault_type file.split(_)[0] condition file.split(_)[1] # 读取CSV跳过可能的文件头 df pd.read_csv(os.path.join(folder_path, file), headerNone) # 添加列名 df.columns [motor_vib, planet_x, planet_y, planet_z, motor_torque, reducer_x, reducer_y, reducer_z] data_dict[f{fault_type}_{condition}] df return data_dict1.3 Matlab数据读取优化虽然原始资料提供了Matlab读取代码但我们可以做一些改进function [data_cell, labels] loadGearData(folder_path) files dir(fullfile(folder_path, *.csv)); data_cell cell(length(files), 1); labels cell(length(files), 1); for i 1:length(files) filename files(i).name; % 解析故障类型和工况 parts strsplit(filename, _); fault_type parts{1}; condition parts{2}(1:end-4); % 去掉.csv % 读取数据 data readmatrix(fullfile(folder_path, filename)); % 存储数据和标签 data_cell{i} data; labels{i} [fault_type _ condition]; end end2. 数据可视化与探索性分析2.1 时域信号可视化在Python中我们可以使用matplotlib快速绘制各通道信号import matplotlib.pyplot as plt def plot_time_domain(data_dict, sample_key, start0, end1000): data data_dict[sample_key].iloc[start:end] plt.figure(figsize(15, 10)) for i, col in enumerate(data.columns, 1): plt.subplot(4, 2, i) plt.plot(data[col]) plt.title(col) plt.xlabel(Sample) plt.ylabel(Amplitude) plt.tight_layout() plt.show()2.2 频域分析技巧振动信号的频域特征往往比时域特征更具区分度。使用快速傅里叶变换(FFT)进行分析import numpy as np from scipy.fft import fft def compute_fft(signal, fs5120): n len(signal) yf fft(signal) xf np.linspace(0, fs/2, n//2) return xf, 2/n * np.abs(yf[0:n//2]) def plot_frequency_domain(data_dict, sample_key, channelplanet_x): signal data_dict[sample_key][channel].values xf, yf compute_fft(signal) plt.figure(figsize(10, 4)) plt.plot(xf, yf) plt.title(fFrequency Domain - {channel}) plt.xlabel(Frequency (Hz)) plt.ylabel(Magnitude) plt.grid() plt.show()2.3 多通道信号相关性分析8个通道信号之间存在一定的相关性计算相关系数矩阵可以帮助我们理解它们之间的关系def plot_correlation_matrix(data_dict, sample_key): data data_dict[sample_key] corr data.corr() plt.figure(figsize(10, 8)) sns.heatmap(corr, annotTrue, cmapcoolwarm, center0) plt.title(Channel Correlation Matrix) plt.show()3. 特征工程从原始信号到特征向量3.1 时域特征提取时域特征计算简单且具有明确的物理意义。以下是一些常用的时域特征def extract_time_features(signal): features { mean: np.mean(signal), std: np.std(signal), rms: np.sqrt(np.mean(signal**2)), peak: np.max(np.abs(signal)), kurtosis: scipy.stats.kurtosis(signal), skewness: scipy.stats.skew(signal), crest_factor: np.max(np.abs(signal)) / np.sqrt(np.mean(signal**2)), impulse_factor: np.max(np.abs(signal)) / np.mean(np.abs(signal)) } return features3.2 频域特征提取频域特征通常能更好地反映机械故障特征def extract_freq_features(signal, fs5120): xf, yf compute_fft(signal, fs) features { freq_mean: np.mean(yf), freq_std: np.std(yf), freq_peak: np.max(yf), peak_freq: xf[np.argmax(yf)], band_energy_0_500: np.sum(yf[(xf 0) (xf 500)]), band_energy_500_1000: np.sum(yf[(xf 500) (xf 1000)]), band_energy_1000_2000: np.sum(yf[(xf 1000) (xf 2000)]) } return features3.3 多通道特征融合策略如何有效利用8个通道的信息是一个关键问题。以下是几种融合策略通道选择根据相关性分析选择信息量最大的几个通道特征级融合计算所有通道特征的均值或最大值决策级融合为每个通道训练单独的分类器然后投票决定最终结果def extract_all_features(data_dict): all_features [] labels [] for key, df in data_dict.items(): # 为每个样本提取特征 sample_features {} # 为每个通道提取时域和频域特征 for channel in df.columns: signal df[channel].values time_feat extract_time_features(signal) freq_feat extract_freq_features(signal) # 添加通道前缀 for k, v in time_feat.items(): sample_features[f{channel}_{k}] v for k, v in freq_feat.items(): sample_features[f{channel}_{k}] v all_features.append(sample_features) labels.append(key.split(_)[0]) # 只取故障类型 # 转换为DataFrame feature_df pd.DataFrame(all_features) feature_df[label] labels return feature_df4. 故障分类模型构建4.1 数据准备与划分在构建模型前我们需要准备好特征矩阵和标签from sklearn.model_selection import train_test_split from sklearn.preprocessing import StandardScaler, LabelEncoder # 加载特征 feature_df extract_all_features(data_dict) # 编码标签 le LabelEncoder() y le.fit_transform(feature_df[label]) X feature_df.drop(label, axis1) # 标准化特征 scaler StandardScaler() X_scaled scaler.fit_transform(X) # 划分训练测试集 X_train, X_test, y_train, y_test train_test_split( X_scaled, y, test_size0.3, random_state42, stratifyy)4.2 随机森林模型随机森林适合处理高维特征且对参数不敏感from sklearn.ensemble import RandomForestClassifier from sklearn.metrics import classification_report rf RandomForestClassifier(n_estimators200, max_depth10, random_state42) rf.fit(X_train, y_train) y_pred rf.predict(X_test) print(classification_report(y_test, y_pred, target_namesle.classes_))4.3 SVM模型SVM在小样本情况下表现优异但对特征缩放敏感from sklearn.svm import SVC svm SVC(C1.0, kernelrbf, gammascale, random_state42) svm.fit(X_train, y_train) y_pred_svm svm.predict(X_test) print(classification_report(y_test, y_pred_svm, target_namesle.classes_))4.4 模型评估与特征重要性分析哪些特征对分类最有帮助# 获取特征重要性 importances rf.feature_importances_ indices np.argsort(importances)[::-1] # 打印最重要的10个特征 print(Feature ranking:) for f in range(10): print(f{f1}. {X.columns[indices[f]]}: {importances[indices[f]]})5. 高级技巧与实战建议5.1 处理类别不平衡齿轮故障数据通常存在类别不平衡问题可以采用以下方法from imblearn.over_sampling import SMOTE smote SMOTE(random_state42) X_resampled, y_resampled smote.fit_resample(X_train, y_train)5.2 时频分析结合结合时域和频域分析可以提取更丰富的特征from scipy.signal import spectrogram def compute_spectrogram(signal, fs5120): f, t, Sxx spectrogram(signal, fsfs, nperseg256) return f, t, Sxx def plot_spectrogram(signal, fs5120): f, t, Sxx compute_spectrogram(signal, fs) plt.pcolormesh(t, f, 10*np.log10(Sxx), shadinggouraud) plt.ylabel(Frequency [Hz]) plt.xlabel(Time [sec]) plt.colorbar(labelPower/Frequency [dB/Hz]) plt.show()5.3 模型集成与优化结合多个模型可以提升分类性能from sklearn.ensemble import VotingClassifier # 定义多个分类器 estimators [ (rf, RandomForestClassifier(n_estimators200, random_state42)), (svm, SVC(probabilityTrue, random_state42)), (xgb, XGBClassifier(random_state42)) ] # 创建投票分类器 ensemble VotingClassifier(estimatorsestimators, votingsoft) ensemble.fit(X_train, y_train) # 评估 y_pred_ens ensemble.predict(X_test) print(classification_report(y_test, y_pred_ens, target_namesle.classes_))5.4 实际应用中的注意事项采样一致性确保所有数据使用相同的采样频率处理工况影响不同转速和负载下的信号特征可能有显著差异特征选择不是所有提取的特征都有用需要进行筛选实时性考虑在实际应用中需要考虑计算效率