STFT划分数据集python代码

### 使用 Python 实现 STFT 并划分训练集与测试集 以下是基于 `numpy` 和 `scipy` 的短时傅里叶变换 (Short-Time Fourier Transform, STFT) 的实现方法，并通过 `sklearn.model_selection.train_test_split` 将数据划分为训练集和测试集。 #### 1. 短时傅里叶变换 (STFT) 的实现 可以使用列表推导式将 FFT 应用于信号的每一段，随后将其组合成二维数组[^1]。此过程类似于 MATLAB 中的 `blkproc` 功能： ```python import numpy as np from scipy.fftpack import fft def stft(signal, frame_size, hop_length): frames = [] for i in range(0, len(signal) - frame_size + 1, hop_length): frame = signal[i:i + frame_size] windowed_frame = frame * np.hamming(frame_size) # 加窗函数 frames.append(windowed_frame) spectra = np.array([fft(frame) for frame in frames]) # 对每一帧执行FFT return spectra.T # 返回频域表示 ``` 上述代码实现了窗口化操作并将每个窗口的数据传递给快速傅里叶变换 (FFT)[^2]。 #### 2. 数据预处理与分割 为了进一步分析或建模，通常需要将生成的 STFT 结果划分为训练集和测试集。这里我们借助 `sklearn` 提供的功能完成这一目标： ```python from sklearn.model_selection import train_test_split # 假设已有一个时间序列信号 data_signal frame_size = 256 # 每一帧大小 hop_length = 64 # 移动步幅长度 spectra_matrix = stft(data_signal, frame_size=frame_size, hop_length=hop_length) # 展平频率维度作为特征向量 features = np.abs(spectra_matrix).T.reshape(-1, frame_size // 2 + 1) # 创建标签（假设为随机生成） labels = np.random.randint(0, 2, size=(len(features),)) # 划分训练集和测试集 X_train, X_test, y_train, y_test = train_test_split( features, labels, test_size=0.3, random_state=42 ) ``` 在此部分中，我们将 STFT 输出的结果展平为适合机器学习模型输入的形式，并调用 `train_test_split` 方法按比例分配样本至训练集合验证集中[^3]。 #### 3. 可视化 STFT 结果 对于调试或者展示目的来说，绘制热力图是一种直观的方式查看 STFT 效果: ```python import matplotlib.pyplot as plt times = np.arange(len(X_train)) * hop_length / sampling_rate freqs = np.fft.rfftfreq(frame_size, d=1/sampling_rate)[:frame_size//2+1] plt.figure(figsize=(8, 6)) plt.pcolormesh(times[:], freqs[:], 20 * np.log10(np.abs(X_train.T)), cmap='inferno') plt.ylabel('Frequency [Hz]') plt.xlabel('Time [sec]') plt.colorbar(label="dB") plt.title("Spectrogram of Training Data") plt.show() ``` 以上绘图脚本展示了如何利用对数尺度显示幅度谱信息，并采用颜色映射增强视觉效果。 --- ###