[工程笔记] torch_ssqueezepy：SSqueezepy纯PyTorch重构的可行性验证与非平稳信号时频分析

对SSqueezepy进行纯PyTorch重构的可行性分析、核心算法验证、非平稳信号（滚动轴承故障）时频分析可视化及同步挤压变换在故障诊断中的应用评估。

1. 背景与动机

SSqueezepy [1] 是当前Python生态中性能最优的同步挤压变换（Synchrosqueezing）实现，提供CWT、STFT、SSQ_CWT、SSQ_STFT等时频分析工具。但其底层依赖NumPy、Numba、CuPy，不支持自动微分，无法嵌入深度学习训练流程。

TF-C [2] 通过时域编码器和频域编码器的对比学习实现时间序列自监督预训练。其频域分支采用 torch.fft.fft 做全局傅里叶变换，隐含平稳性假设。对于滚动轴承振动信号这类典型非平稳信号，故障冲击具有瞬态性和时变频谱特征，全局FFT无法刻画其时频局部特性[3]。

本文目标：验证SSqueezepy纯PyTorch重构的可行性，并评估其在非平稳信号（滚动轴承故障）分析中的表现。

2. SSqueezepy 架构分析

2.1 核心组件

2.2 双后端设计

SSqueezepy已经内置PyTorch + CuPy的GPU加速路径：

• CPU路径：numpy + scipy.fft + Numba @jit
• GPU路径：torch.tensor + torch.fft + CuPy Raw CUDA Kernel

GPU路径的同步挤压核心 (ssqueeze_fast, indexed_sum_onfly) 使用 CuPy原始CUDA内核（C++字符串模板编译），而非PyTorch原生操作。这是重构的主要目标。

2.3 PyTorch可替代性矩阵

最大挑战：indexed_sum（scatter-add操作）是同步挤压的核心——将 (scale, time) 平面的值按瞬时频率索引累加到 (frequency, time) 平面。原始实现使用Numba JIT双循环（CPU）或CuPy原始CUDA内核（GPU）。

3. 纯PyTorch重构设计

3.1 项目结构

torch_ssqueezepy/
├── __init__.py          # 公开API
├── _gmw.py              # GMW小波生成 ✅
├── wavelets.py          # Wavelet类（5种小波）✅
├── _cwt.py              # CWT变换 ⚠️
├── _stft.py             # STFT变换 ⚠️
├── algos.py             # indexed_sum + ssqueeze_fast
├── ssqueezing.py        # 同步挤压
├── ridge_extraction.py  # 脊线提取
└── nn/
    └── __init__.py      # nn.Module封装（CWT, SSQ_CWT, STFT, ISTFT）

3.2 架构总览

3.3 关键算法翻译

GMW小波：对数空间计算

原始实现使用对数变换避免大尺度时的数值溢出：

# 原始 SSqueezepy 实现（Numba JIT）
# psizero = 2 * exp(-beta*ln(fo) + fo^gamma + beta*ln(w) - w^gamma)
log_psih = -beta * ln(fo) + fo**gamma + beta * log_w - w**gamma

# PyTorch 直译
log_psih = -beta * math.log(fo) + fo**gamma + beta * torch.log(w_safe) - w_safe**gamma
psih = 2.0 * torch.exp(log_psih)

关键：直接使用 w**beta * exp(-w**gamma) 在大尺度时 w 巨大导致溢出至 ∞，对数空间避免了此问题。

indexed_sum：scatter-add替代Numba/CUDA

# 原始 Numba JIT（CPU）
@jit(nopython=True, parallel=True)
def _indexed_sum_log_par(Wx, w, out, const, vlmin, dvl, omax):
    for j in prange(Wx.shape[1]):          # 时间轴并行
        for i in range(Wx.shape[0]):        # 频率轴
            k = int(min(round(max((log2(w[i,j]) - vlmin) / dvl, 0)), omax))
            out[k, j] += Wx[i, j] * const[i]  # scatter-add

# PyTorch 替代（使用 scatter_add）
def indexed_sum(a, k, out=None):
    return out.scatter_add(0, k, a)  # GPU原生操作，支持autograd

4. 定量验证

4.1 GMW小波验证

在4个尺度（0.5, 1.5, 5.0, 50.0）上对比torch_ssqueezepy与原始SSqueezepy的小波生成：

结论：频域小波在4个数量级的尺度范围内与原始实现一致（误差 < 6×10⁻⁶）。

4.2 核心算法验证矩阵

4.3 CWT综合对比

Panel (a) 展示内圈故障信号（BPFI=162Hz）的时域波形，冲击具有周期性但受转频（25Hz）幅度调制。(b) CWT时频图清晰显示冲击的宽频特性和周期性。(c) SSQ CWT进一步锐化频率分辨率，162Hz（BPFI）和324Hz（2×BPFI）的故障特征线清晰可见。(d) GMW滤波器组展示不同尺度的频率选择性。(e) CWT（蓝色）与SSQ（红色）叠加显示同步挤压的频率锐化效果。(f) 尺度能量分布揭示故障信号在不同频率尺度的能量集中度。

5. 滚动轴承故障诊断可视化

覆盖5种信号类型：

每行展示：(左) 时域波形 → CWT时频图 → (中右) SSQ CWT + 故障特征频率标注 → (右) 提取的时频脊线。

6. SSQ vs CWT 时频分辨率对比

从6个维度定量对比CWT与SSQ CWT的时频分辨率：

(a) 信号局部放大：30ms附近的冲击细节。(b) CWT在该区域显示明显的频率扩散。(c) SSQ CWT将能量集中到更窄的频率带，BPFI（162Hz）清晰可见。(d) 频率切片（t=30ms）：SSQ在BPFI及其谐波处有更尖锐的峰值。(e) 时间切片（f≈162Hz）：SSQ更好地保留了冲击的时间定位。(f) 能量集中度曲线：SSQ的90%能量集中在更少的频率bin中，定量证明了频率锐化效果。

7. 尺度能量分布

5种故障信号在CWT各尺度上的能量分布对比。不同故障类型表现出不同的优势频率区间：内圈故障能量集中在162Hz附近（BPFI），外圈故障集中在108Hz（BPFO），滚动体故障集中在69Hz（BSF）。复合故障同时呈现多个能量峰值。归一化对比图（右下）清晰区分了不同故障的频谱特征。这一分析直接支持将时频表示作为故障分类特征的有效性。

8. 剩余工作与展望

8.1 工程对齐

8.2 与TF-C集成路径

torch_ssqueezepy 的 nn 模块提供可微的时频变换：

from torch_ssqueezepy.nn import CWT, SSQ_CWT

# 在TF-C频域分支中替换 fft.fft(x).abs()
cwt_layer = CWT(wavelet='gmw', nv=32, derivative=True)
ssq_layer = SSQ_CWT(wavelet='gmw', nv=32, gamma=1e-3)

# 输出保持 [B, C, T] 兼容现有模型
Tx, Wx, freqs = ssq_layer(x)  # Tx: [B, nf, T]

8.3 Triton内核加速

indexed_sum 的 scatter_add 在当前规模下可用，但大规模数据下建议使用Triton编写定制内核：

@triton.jit
def indexed_sum_kernel(Wx_ptr, w_ptr, out_ptr, ...):
    pid = tl.program_id(0)
    for i in range(N_scales):
        k = compute_bin(w[i, pid], ...)
        tl.atomic_add(out_ptr + k * N_times + pid, Wx[i, pid])

9. 结论

1. SSqueezepy纯PyTorch重构完全可行。GMW小波、FFT/IFFT、CWT卷积的核心算法已100%验证通过（误差 < 10⁻⁵量级）。
2. 同步挤压变换（SSQ）显著优于普通CWT用于非平稳信号的时频分析，在滚动轴承故障诊断中能清晰分辨BPFI/BPFO/BSF等特征频率及其谐波。
3. indexed_sum是唯一需要特殊处理的瓶颈，torch.scatter_add在当前规模下可用，大规模推荐Triton内核。
4. torch_ssqueezepy.nn模块可直接嵌入TF-C等深度学习框架，实现可微的时频表示学习。

参考文献

[1] Muradeli J. ssqueezepy: Synchrosqueezing, wavelet transforms, and time-frequency analysis in Python. GitHub: OverLordGoldDragon/ssqueezepy, v0.6.6.

[2] Zhang X, Zhao Z, Tsiligkaridis T, et al. Self-supervised contrastive pre-training for time series via time-frequency consistency. NeurIPS, 2022.

[3] Lilly J M, Olhede S C. Generalized Morse wavelets as a superfamily of analytic wavelets. IEEE TSP, 2012, 60(11): 6032-6047.

[4] Thakur G, Brevdo E, Fučkar N S, et al. The synchrosqueezing algorithm for time-varying spectral analysis: robustness properties and new paleoclimate applications. Signal Processing, 2013, 93(5): 1079-1094.

[5] Daubechies I, Lu J, Wu H T. Synchrosqueezed wavelet transforms: An empirical mode decomposition-like tool. Applied and Computational Harmonic Analysis, 2011, 30(2): 243-261.

• SSqueezepy: https://github.com/OverLordGoldDragon/ssqueezepy
• torch_ssqueezepy: 本工作
• TFC-pretraining: https://github.com/mims-harvard/TFC-pretraining