如何创建X轴为Frequency、Y轴为Magnitude的2D热力图以可视化多组FFT输出？

需求背景 #

原代码仅能绘制少量FFT的幅值-频率散点图，当需要处理1000组及以上FFT输出时，散点图会因过度重叠失去可读性。以下提供两种基于numpy和matplotlib的可视化方案，分别用于展示频谱分布密度和批量频谱轮廓。

方案1：2D直方图热力图（展示分布密度） #

通过numpy.histogram2d统计所有FFT结果中（频率，幅值）组合的出现频次，用颜色深浅表示密度，适合分析大量信号的频谱分布规律。

完整代码 #

import numpy as np
import matplotlib.pyplot as plt

# 基础参数配置
fs = 1000
t = np.arange(0, 0.1, 1/fs)
N = len(t)
f_bin = fs / N
# 仅保留正频率部分（FFT结果对称）
f = np.arange(0, fs, f_bin)[:N//2]

# 生成1000组随机频率的正弦波信号（模拟真实场景的多组数据）
num_groups = 1000
random_freqs = np.random.uniform(50, 450, num_groups)

# 批量计算FFT幅值
all_magnitudes = []
for freq in random_freqs:
    signal = np.sin(2 * np.pi * freq * t)
    fft_result = np.fft.fft(signal)
    mag = np.abs(fft_result)[:N//2]
    all_magnitudes.append(mag)

# 将所有数据展平为一维数组，用于二维直方图统计
f_flat = np.tile(f, num_groups)
mag_flat = np.array(all_magnitudes).flatten()

# 计算二维直方图分箱
bins_x = 50  # 频率轴分箱数
bins_y = 50  # 幅值轴分箱数
hist, x_edges, y_edges = np.histogram2d(f_flat, mag_flat, bins=(bins_x, bins_y))

# 绘制热力图
fig, ax = plt.subplots(figsize=(10, 6))
# 转置直方图数据以匹配imshow的轴顺序
im = ax.pcolormesh(x_edges, y_edges, hist.T, cmap='viridis')

ax.set_xlabel('Frequency (Hz)')
ax.set_ylabel('Magnitude')
ax.set_title('2D Histogram of 1000 FFT Results')
plt.colorbar(im, label='Occurrence Count')
plt.grid(False)
plt.show()

说明 #

• 颜色越深表示对应（频率，幅值）组合出现的次数越多，能快速定位高频出现的频谱特征
• 可根据数据调整bins_x和bins_y的数值，平衡分辨率和计算效率

方案2：批量频谱轮廓热力图（展示每组细节） #

将每组FFT的频谱按行排列，用颜色表示幅值，Y轴为组索引，适合观察批量信号的频谱变化趋势。

完整代码 #

import numpy as np
import matplotlib.pyplot as plt

# 基础参数配置
fs = 1000
t = np.arange(0, 0.1, 1/fs)
N = len(t)
f = np.arange(0, fs, fs/N)[:N//2]

# 生成1000组带噪声的信号（模拟真实场景数据）
num_groups = 1000
random_freqs = np.random.uniform(50, 450, num_groups)
magnitude_matrix = []

for freq in random_freqs:
    # 添加高斯噪声模拟真实信号
    signal = np.sin(2 * np.pi * freq * t) + 0.1 * np.random.randn(len(t))
    fft_mag = np.abs(np.fft.fft(signal))[:N//2]
    magnitude_matrix.append(fft_mag)

magnitude_matrix = np.array(magnitude_matrix)

# 绘制热力图
fig, ax = plt.subplots(figsize=(10, 6))
# 设置extent参数匹配频率轴范围，aspect='auto'自适应比例
im = ax.imshow(magnitude_matrix, aspect='auto', extent=[f[0], f[-1], num_groups, 0], cmap='viridis')

ax.set_xlabel('Frequency (Hz)')
ax.set_ylabel('FFT Group Index')
ax.set_title('1000 FFT Spectra (Magnitude vs Frequency)')
plt.colorbar(im, label='Magnitude')
plt.show()

说明 #

• 若幅值范围差异过大，可对幅值取对数（np.log10(magnitude_matrix)），提升低幅值区域的细节可见度
• Y轴的组索引可替换为实际业务中的分组标识（如时间、实验编号等）

内容的提问来源于stack exchange，提问作者ge45mue