频谱分析—Python代码

下面是一个用python进行频谱分析的代码案例。

import matplotlib.pyplot as plt
import numpy as np


def myfft(signal, fs):
    '''
    FFT变换，用于频谱分析
    :param signal: type: ndarray, shape: (n,)
    :param fs: 采样频率， Hz
    :param one_side (default: True): 仅限制正频率
    :return:
        freq -- 频率轴
        am   -- 幅值轴
    '''

    if not isinstance(signal, np.ndarray):
        signal = np.array(signal)

    if signal.ndim != 1:
        raise ValueError('输入的信号必须是一维的！')

    # 零均值化处理，去除直流分量（也就是去除频率等于0的点，防止对结果分析造成干扰）
    signal = signal - signal.mean()

    N = len(signal)  # 信号长度
    am = np.fft.fft(signal)
    am = np.abs(am) / N
    freq = np.arange(0, N, dtype=float) * (fs/N) # 等价freq=np.fft.fftfreq(N, d=1/fs)

    am = am[: N // 2]
    freq = freq[: N // 2] # 正频率[0, fs/2]
    return freq, am

if __name__ == "__main__":
    fs = 5120 # 采样频率
    f1 = 20
    f2 = 85
    N = 4096 # 信号长度
    t = np.arange(0, N / fs, 1 / fs)
    x = np.sin(2 * np.pi * f1 * t) + np.cos(2 * np.pi * f2 * t) + 5
    x = x + np.random.normal(0, 1, size=(4096,))

    freq, am = myfft(x, fs)

    plt.figure(figsize=(8, 3))
    plt.subplot(121)
    plt.plot(t, x, )
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude')
    plt.title('Time domain')

    plt.subplot(122)
    plt.plot(freq, am, zorder=100, color='k', linewidth=1.5)
    plt.xlim([0, 500])
    plt.ylim([0, 1])
    plt.axvline(f1, color='r', linestyle='--', alpha=0.6, linewidth=1.2)
    plt.axvline(f2, color='r', linestyle='--', alpha=0.6, linewidth=1.2)
    plt.xlabel('Frequency (Hz)')
    plt.ylabel('Amplitude')
    plt.title('Frequency domain')
    plt.tight_layout()
    plt.show()

结果如下：