Pytorch 实现 PSD 功率谱计算：periodogram 和 welch 方法

最新推荐文章于 2025-09-17 15:20:26 发布

原创

最新推荐文章于 2025-09-17 15:20:26 发布 · 6.7k 阅读

44 ·

CC 4.0 BY-SA版权

文章标签：

#Pytorch #DSP #功率谱 #FFT #PSD

本文介绍了如何使用PyTorch实现功率谱估计中的Periodogram和Welch方法，代码参考了scipy的API设计。Periodogram函数用于计算单边功率谱，Welch方法则通过窗口函数进行分段处理以减小频谱泄漏。提供了详细的函数定义和测试案例，适用于信号处理和数据分析场景。

前言

由于跑代码需要，本人用 Pytorch 简单实现了 PSD 估计中的最常用的 periodogram 和 welch 函数。这里把我的实现分享给大家。

在实现方案和参数的设计上，本人参考了 scipy 的 API 设计。

`periodogram` 的实现

函数体如下。

def _periodogram(X: torch.Tensor, fs, detrend, scaling):
    if X.dim() > 2:
        X = torch.squeeze(X)
    elif X.dim() == 1:
        X = X.unsqueeze(0)

    if detrend:
        X -= X.mean(-1, keepdim=True)

    N = X.size(-1)
    assert N % 2 == 0

    df = fs / N
    dt = df
    f = torch.arange(0, N / 2 + 1) * df  # [0:df:f/2]

    dual_side = fft.fft(X)  # 双边谱
    half_idx = int(N / 2 + 1)
    single_side = dual_side[:, 0:half_idx]
    win = torch.abs(single_side)

    ps = win ** 2
    if scaling == 'density':  # 计算功率谱密度
        scale = N * fs
    elif scaling == 's