源分离与降噪

源分离与降噪

源分离与降噪从混合音频中恢复出各个独立信号——这是计算上的鸡尾酒会问题。本文件涵盖 ICA、NMF、时频掩码、波束成形、深度学习分离网络（Conv-TasNet、SepFormer）、语音增强以及自适应降噪。

• 想象你站在一个拥挤的鸡尾酒会上。几十个人同时说话，音乐在播放，酒杯在叮当作响，然而你能集中注意力于一段对话并清晰地跟上它。这种非凡的能力，即鸡尾酒会问题（Cherry，1953），是人类听觉系统毫不费力就能解决的，但机器却觉得极其困难。本文件涵盖试图解决它的算法：分离混合音频源、消除不需要的噪声、在恶劣条件下增强语音。

• 第 1 个文件中的信号处理基础（STFT、频谱图、滤波器组）支撑了这里的每一种方法。第 2 章的矩阵分解技术（NMF、ICA、SVD）提供了经典工具箱。第 6 章的深度学习架构（CNN、RNN、注意力）和第 4/5 章的概率论为现代方法提供了依据。

• 问题形式化：在一个或多个麦克风处观测到一个混合信号 x(t)。在最简单的情况下，混合是 C 个源信号之和：

• 其中 s_c(t) 是第 c 个源信号，n(t) 是背景噪声。目标是从 x(t) 中恢复出各个 s_c(t)。在单麦克风情况下，这是严重欠定的：一个方程，C 个未知数。需要额外的假设（统计独立性、频谱结构、学习到的先验）才能使问题可解。

• 在频域中（通过第 1 个文件的 STFT），混合变为：

• 许多分离方法在时频域工作，通过为每个源估计一个掩码 M_c(t, f) \in [0, 1]，然后以 \hat{S}_c(t, f) = M_c(t, f) \cdot X(t, f) 恢复源。理想二值掩码（Ideal Binary Mask，IBM）在该时频 bin 由源 c 主导时设 M_c(t, f) = 1，否则为 0。理想比值掩码（Ideal Ratio Mask，IRM）是一个软版本：

• 独立成分分析（Independent Component Analysis，ICA）是当麦克风数等于或超过源数时的经典方法。ICA（第 2 章）寻找一个线性解混矩阵 W 使得 \hat{s} = Wx，其中恢复的源 \hat{s} 最大化统计独立。关键假设是源信号是非高斯的且相互独立，这对语音和音乐通常成立。

• 对于多麦克风的瞬时混合模型 x = As（其中 A 是混合矩阵），ICA 通过最大化输出的非高斯性（FastICA 使用负熵）或最小化互信息来恢复 W \approx A^{-1}。ICA 在受控环境中效果很好，但当混合涉及卷积（房间混响）、源数超过麦克风数、或独立性假设被违反时会失效。

• 非负矩阵分解（Non-negative Matrix Factorisation，NMF）把幅度频谱图 V \in \mathbb{R}_+^{F \times T} 分解为两个非负矩阵的乘积（第 2 章）：

• 其中 W \in \mathbb{R}_+^{F \times K} 是 K 个谱基向量的字典，H \in \mathbb{R}_+^{K \times T} 包含随时间变化的激活系数。非负性约束有物理动机：幅度是非负的，声音是可加地组合的。

• 对于源分离，NMF 为每个源学习单独的字典：W_{\text{speech}} 捕捉语音的频谱模式（共振峰结构），而 W_{\text{noise}} 捕捉噪声模式。混合被分解为 V \approx W_{\text{speech}} H_{\text{speech}} + W_{\text{noise}} H_{\text{noise}}，每个源通过掩码恢复。NMF 使用乘性更新最小化，代价函数为 Frobenius 范数或 KL 散度：

\begin{aligned}
\text{Frobenius:} \quad D_F(V \| WH) &= \|V - WH\|_F^2 \\
\text{KL:} \quad D_{KL}(V \| WH) &= \sum_{f,t} \left[ V_{ft} \log \frac{V_{ft}}{(WH)_{ft}} - V_{ft} + (WH)_{ft} \right]
\end{aligned}

• 波束成形利用麦克风阵列的空间信息。当一个源信号以不同延迟到达不同麦克风时（由于空间布置），这些延迟可用于增强一个方向的信号同时抑制其他方向。

• 延时求和波束成形是最简单的方法。如果目标源相对阵列位于角度 \theta，麦克风 m 处的时间延迟为 \tau_m(\theta) = d_m \sin \theta / c，其中 d_m 是麦克风位置，c 是声速。波束成形器输出对齐并求和麦克风信号：

• 来自目标方向的信号相干地相加，而来自其他方向的信号非相干地相加，从而提供空间滤波。阵列几何决定空间分辨率：更大的阵列给出更窄的波束。

• 最小方差无失真响应（Minimum Variance Distortionless Response，MVDR）波束成形优化权重，以在无失真通过目标方向的同时最小化总输出功率：

\begin{aligned}
\min_{\mathbf{w}} \quad & \mathbf{w}^H \Phi_{nn} \mathbf{w} \\
\text{subject to} \quad & \mathbf{w}^H \mathbf{d}(\theta) = 1
\end{aligned}

• 其中 \Phi_{nn} 是噪声空间协方差矩阵，\mathbf{d}(\theta) 是方向 \theta 的导向向量。闭式解为：

• MVDR 通过使用估计的噪声协方差来适应噪声环境，提供比延时求和更好的干扰抑制。它广泛用于助听器、智能扬声器和电话会议系统。

• 用于源分离的深度学习显著提升了性能，尤其是在经典方法表现挣扎的单麦克风情况下。一般范式是：编码混合，用神经网络估计掩码或源表示，解码以恢复各个源。

• 深度聚类（Hershey 等，2016）把每个时频 bin 嵌入到高维空间，其中属于同一源的 bin 彼此靠近，来自不同源的 bin 相距很远。一个双向 LSTM（第 6 章）把每个 T-F bin (t, f) 映射到嵌入 v_{t,f} \in \mathbb{R}^D。训练目标为：

• 其中 V 是嵌入矩阵，Y 是源分配的 one-hot 矩阵。乘积 VV^T 是一个亲和矩阵（两个 bin 的嵌入有多相似），YY^T 是理想亲和（同源为 1，否则为 0）。推理时，对嵌入做 K-means 聚类产生二值掩码。

• Conv-TasNet（Luo 和 Mesgarani，2019）完全在时域操作，绕过了 STFT。它有三个组件：

• 编码器：一个一维卷积把混合波形的短段映射到潜在表示。对于混合 x \in \mathbb{R}^T，编码器输出为 w = \text{ReLU}(U \ast x) \in \mathbb{R}^{N \times L}，其中 U 是可学习的基（类似于 STFT 基但从数据中学习），N 是基函数的数量，L 是段数。编码器核大小和步长（通常 2ms 和 1ms）决定时间分辨率。

• 分离器：一个时间卷积网络（Temporal Convolutional Network，TCN）处理编码后的混合并输出 C 个掩码。TCN 以指数级增大的膨胀因子 1, 2, 4, \ldots, 2^{B-1} 堆叠膨胀的一维深度可分离卷积（来自第 8 章的高效卷积），重复 R 次。这给出了非常大的感受野，同时保持计算高效。

• 解码器：一个一维转置卷积（带学习到的基 V）把每个掩码后的表示转回时域：\hat{s}_c = V^T (M_c \odot w)。

• Conv-TasNet 显著优于基于频谱图的方法，因为学习到的编码器-解码器基能够捕捉 STFT 幅度所丢弃的信息（尤其是相位）。

• 双路径 RNN（Dual-Path RNN，DPRNN）（Luo 等，2020）解决分离中的长序列建模问题。DPRNN 不再使用单个 RNN 或 TCN 处理整个编码序列，而是把序列分成重叠的块，并沿两条路径应用 RNN：块内路径（建模每个块内的局部模式）和块间路径（建模跨块的全局模式）。这把每个维度上的 RNN 序列长度从 L 降到 \sqrt{L}：

\begin{aligned}
\text{Intra-chunk:} \quad & h_{k,n}^{\text{intra}} = \text{BiLSTM}_{\text{intra}}(z_{k,n}) \\
\text{Inter-chunk:} \quad & h_{k,n}^{\text{inter}} = \text{BiLSTM}_{\text{inter}}(h_{k,n}^{\text{intra}})
\end{aligned}

• 其中 k 索引块，n 索引块内位置。块内 LSTM 对固定 k 跨 n 处理；块间 LSTM 对固定 n 跨 k 处理。

• SepFormer（Subakan 等，2021）用 Transformer（第 7 章）取代双路径框架中的 RNN。块内 Transformer 用自注意力捕捉局部依赖，块间 Transformer 捕捉全局依赖。多头注意力在没有梯度消失问题（第 6 章）的情况下建模长距离依赖的能力，使 SepFormer 对长录音特别有效。SepFormer 在 WSJ0-2mix 基准上取得了最优结果。

• 置换不变训练（Permutation Invariant Training，PIT）解决监督源分离中的一个根本问题：标签分配歧义。如果网络有两个输出（对应两个说话人），哪个输出应对应哪个说话人？没有自然的顺序。PIT 对所有可能的分配计算损失并取最小值：

• 其中 \mathcal{P} 是 \{1, \ldots, C\} 的所有置换的集合，\ell 是每源损失（通常是尺度不变信号失真比，SI-SDR）。对于 C = 2 个源只有 2 种置换；C = 3 有 6 种。对于更大的 C，使用匈牙利算法高效计算。

• 尺度不变信号失真比（Scale-Invariant Signal-to-Distortion Ratio，SI-SDR）是源分离的标准评估指标：

\begin{aligned}
s_{\text{target}} &= \frac{\langle \hat{s}, s \rangle}{\|s\|^2} s \\
e_{\text{noise}} &= \hat{s} - s_{\text{target}} \\
\text{SI-SDR} &= 10 \log_{10} \frac{\|s_{\text{target}}\|^2}{\|e_{\text{noise}}\|^2}
\end{aligned}

• 其中 \hat{s} 是估计的源，s 是真实值。SI-SDR 对估计的整体尺度不变，这是可取的，因为绝对音量不如分离质量重要。SI-SDR（以 dB 为单位）越高越好。最优系统在 WSJ0-2mix 上达到约 20-22 dB 的 SI-SDR 改进。

• 音乐源分离把音乐录音分离成音轨：人声、鼓、贝斯和其他乐器。这使得卡拉 OK（去除人声）、混音（调整乐器电平）和转录（一次分析一种乐器）等应用成为可能。

• Open-Unmix（Stoter 等，2019）是一个参考基线，使用 3 层双向 LSTM 在幅度 STFT 域为每个源预测一个软掩码。它用专用模型独立地处理每个源。简单但有效，Open-Unmix 在 MUSDB18 上建立了可复现的基准。

• Demucs（Defossez 等，2019；更新为 Hybrid Demucs，2021）使用直接在波形上操作的 U-Net 架构（第 8 章）。编码器通过步幅卷积压缩混合，解码器通过带跳跃连接的转置卷积把它扩展回去，每个源有自己的解码器头。Hybrid Demucs 结合了时域和频域处理：编码器有并行的时域和 STFT 分支，其特征在解码器之前融合。这同时捕捉了精细的时间细节和频谱结构。

• Demucs 在 MUSDB18 上取得了最优的分离质量，人声分离尤为出色。其 U-Net 架构让人联想到第 8 章的图像分割架构，把分离问题当作一种"音频分割"来处理。

• 主动降噪（Active Noise Cancellation，ANC）通过生成一个与噪声相消干涉的反噪声信号来减少不需要的声音。想象一下降噪耳机：一个麦克风拾取环境噪声，ANC 系统生成一个反相版本，组合信号（噪声 + 反噪声）理想情况下相消为静音。

• 物理很简单：如果噪声是 n(t)，在同一空间位置生成 -n(t) 就产生静音：n(t) + (-n(t)) = 0。挑战在于反噪声必须在时间、幅度和相位上精确对齐。即使很小的误差也会产生残余噪声或伪影。

• 前馈 ANC 使用一个参考麦克风，在噪声到达听者之前拾取它。系统有时间处理噪声并生成反噪声。参考信号通过一个自适应滤波器，其输出在误差麦克风（靠近听者）处从噪声中减去。这对可预测的宽带噪声（引擎嗡嗡声、风扇噪声）效果很好。

• 反馈 ANC 只使用听者耳处的误差麦克风。系统从残余信号（听者实际听到的）估计噪声并调整反噪声。反馈 ANC 更简单（不需要参考麦克风），但带宽有限且可能变得不稳定。

• 自适应滤波是 ANC 背后的数学引擎。滤波器系数必须持续适应不断变化的噪声环境。最常见的算法是最小均方（Least Mean Squares，LMS）滤波器。

• LMS 算法：一个具有系数 \mathbf{w} = [w_0, w_1, \ldots, w_{L-1}]^T 的 FIR 滤波器处理参考信号 \mathbf{x}(n) = [x(n), x(n-1), \ldots, x(n-L+1)]^T。输出为 y(n) = \mathbf{w}^T \mathbf{x}(n)，误差为 e(n) = d(n) - y(n)（其中 d(n) 是期望/主信号），权重更新为：

• 其中 \mu 是步长（学习率）。这是对均方误差 E[e^2(n)] 的随机梯度下降步，使用瞬时梯度估计 -2 e(n) \mathbf{x}(n) 而非真实梯度（第 3 章的梯度下降和第 6 章的 SGD）。

• 步长 \mu 控制收敛速度和稳态误差之间的权衡。太大滤波器会振荡或发散；太小适应迟钝。稳定性条件为 0 < \mu < 2 / (\lambda_{\max})，其中 \lambda_{\max} 是输入自相关矩阵 R = E[\mathbf{x}\mathbf{x}^T] 的最大特征值。

• 归一化 LMS（Normalised LMS，NLMS）用输入功率归一化步长，使收敛与信号电平无关：

• 其中 \epsilon 是防止除零的小正则化常数。NLMS 比 LMS 收敛更可靠，因为有效步长会适应输入功率。

• 递归最小二乘（Recursive Least Squares，RLS）是一种收敛更快的替代方法，最小化加权最小二乘代价 \sum_{k=1}^{n} \lambda^{n-k} e^2(k)，其中 \lambda \in (0, 1] 是遗忘因子。RLS 维持对逆自相关矩阵的估计并递归更新，以每样本 O(L^2) 的计算（相比 LMS 的 O(L)）为代价实现最优收敛。

• 降噪与语音增强旨在改善嘈杂录音中的语音质量和可懂度。与分离不同源的源分离不同，语音增强专门针对语音加噪声的情况，从含噪观测中恢复干净语音。

• 谱减法是最简单的方法。在只有噪声的帧（由第 3 个文件的 VAD 检测），估计噪声谱 |\hat{N}(f)|^2。然后从每帧中减去它：

• 其中 \alpha 是过减因子（通常 1-4，激进减法去除更多噪声但引入更多伪影），\beta 是防止负值的谱底，并减少"音乐噪声"伪影（听起来像随机音符的孤立音调残余）。

• 维纳滤波提供干净语音频谱的最小均方误差估计：

• 维纳增益 G(t, f) = \text{SNR}(t, f) / (1 + \text{SNR}(t, f)) 取值范围从 0（纯噪声）到 1（纯语音），充当软掩码。挑战在于估计语音和噪声功率谱。先验信噪比 \xi(t, f) = |S(t,f)|^2 / |N(t,f)|^2 使用"决策导向"方法估计：当前帧估计和前一帧维纳滤波输出的平滑组合。

• 神经语音增强使用深度学习估计掩码（如维纳增益）或干净频谱图。架构从简单的前馈网络到 U-Net（第 8 章）、CRN（卷积循环网络）和 Transformer 不等。

• DCCRN（Deep Complex Convolutional Recurrent Network，深度复数卷积循环网络）在复数 STFT（幅度和相位都有）上操作，使用能自然处理实部和虚部的复值卷积。这避免了困扰仅幅度方法的相位估计问题。

• FullSubNet 使用双路径架构，配一个全带模型（捕捉全局频谱模式）和一个子带模型（捕捉局部谐波细节）。全带模型处理整个频谱，子带模型处理以每个频率 bin 为中心的窄频带。它们的输出组合得到最终的掩码估计。

• DNS（Deep Noise Suppression）挑战赛由微软每年为语音增强系统举办基准测试。获胜者通常使用大规模训练配以多样化噪声类型、数据增强（以各种信噪比添加噪声、混响、编解码器伪影）以及具备实时能力的架构。

• 回声消除去除双向通信中的声学回声。当你在通话时，对端说话人的声音通过你的扬声器播放，在房间里反射，被你的麦克风拾取，产生对端说话人听到的回声。声学回声消除（Acoustic Echo Cancellation，AEC）建模从扬声器到麦克风的声学路径，并减去预测的回声。

• 声学路径被建模为一个自适应 FIR 滤波器（使用 LMS 或 NLMS），以对端信号为输入。滤波器建模房间脉冲响应，包括直射路径、早期反射和晚期混响。房间脉冲响应可能长达数百毫秒，需要具有数千个抽头的滤波器。

• 双讲检测对 AEC 至关重要：当近端和对端说话人同时说话时，自适应滤波器必须冻结（停止更新）以防止它消除近端说话人的声音。双讲检测器比较误差信号的能量与对端信号能量；误差能量的突然增加且无法由对端信号解释，表明存在近端语音。

• 对端信号 x(n) 与麦克风信号 d(n) 之间的归一化互相关提供双讲指示：

• 在单讲期间（仅对端），\xi 很高，因为 d 主要是 x 的回声。在双讲期间，\xi 下降，因为近端语音与 x 不相关。

• 现代 AEC 系统结合自适应滤波与神经网络：自适应滤波器提供初始回声估计，神经网络（类似上面的语音增强模型）清理残余回声并处理线性滤波器无法捕捉的非线性（扬声器失真）。

• 分离与增强的评估指标：

  • SI-SDR（上面定义）：源分离的标准。
  • SDR（Signal-to-Distortion Ratio，信号失真比）：来自 BSS Eval，衡量包括伪影和干扰在内的整体分离质量。
  • PESQ（Perceptual Evaluation of Speech Quality，语音质量感知评估）：ITU 标准，预测主观质量得分。范围：-0.5 到 4.5。
  • STOI（Short-Time Objective Intelligibility，短时客观可懂度）：预测语音可懂度。范围：0 到 1。
  • DNSMOS：微软的深度降噪 MOS 预测器，一个训练用于预测人类 MOS 得分的神经网络，无需干净参考音频。

编程练习（使用 CoLab 或 notebook）

• 任务 1：用于源分离的独立成分分析。 实现 FastICA 分离两个混合音频源，演示确定情况下（源数与麦克风数相等）的经典鸡尾酒会解决方案。

import jax
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt

# 生成两个源信号
sr = 8000
duration = 1.0
t = jnp.linspace(0, duration, int(sr * duration))

# 源 1：正弦波（类似一个音调）
s1 = jnp.sin(2 * jnp.pi * 440 * t) + 0.3 * jnp.sin(2 * jnp.pi * 880 * t)

# 源 2：类锯齿波（丰富的谐波）
s2 = 2 * (t * 200 % 1) - 1  # 200 Hz 锯齿波

# 归一化源
s1 = s1 / jnp.max(jnp.abs(s1))
s2 = s2 / jnp.max(jnp.abs(s2))
sources = jnp.stack([s1, s2])  # (2, T)

# 混合矩阵（对算法未知）
A = jnp.array([[0.8, 0.4],
               [0.3, 0.9]])
mixtures = A @ sources  # (2, T)

# FastICA 实现
def whiten(X):
    """中心化并白化数据。"""
    X_centered = X - jnp.mean(X, axis=1, keepdims=True)
    cov = (X_centered @ X_centered.T) / X_centered.shape[1]
    eigvals, eigvecs = jnp.linalg.eigh(cov)
    D_inv_sqrt = jnp.diag(1.0 / jnp.sqrt(eigvals + 1e-8))
    whitening = D_inv_sqrt @ eigvecs.T
    return whitening @ X_centered, whitening

def fastica(X, n_components=2, max_iter=200, tol=1e-6):
    """使用 tanh 非线性的 FastICA（负熵的近似）。"""
    X_white, whitening = whiten(X)
    n, T = X_white.shape

    key = jr.PRNGKey(42)
    W = jr.normal(key, (n_components, n))
    # 正交化 W
    U, _, Vt = jnp.linalg.svd(W, full_matrices=False)
    W = U @ Vt

    for iteration in range(max_iter):
        W_old = W.copy()

        # 对每个分量
        for i in range(n_components):
            w = W[i]
            # w^T X_white: (T,)
            wx = w @ X_white  # (T,)

            # g(u) = tanh(u), g'(u) = 1 - tanh^2(u)
            g_wx = jnp.tanh(wx)
            g_prime_wx = 1 - g_wx ** 2

            # 牛顿更新：w_new = E[X * g(w^T X)] - E[g'(w^T X)] * w
            w_new = jnp.mean(X_white * g_wx[None, :], axis=1) - \
                    jnp.mean(g_prime_wx) * w

            # 与之前的分量去相关（紧缩法）
            for j in range(i):
                w_new = w_new - jnp.dot(w_new, W[j]) * W[j]

            w_new = w_new / jnp.linalg.norm(w_new)
            W = W.at[i].set(w_new)

        # 检查收敛
        convergence = jnp.min(jnp.abs(jnp.diag(W @ W_old.T)))
        if convergence > 1 - tol:
            print(f"FastICA converged in {iteration + 1} iterations")
            break

    # 解混矩阵
    unmixing = W @ whitening
    recovered = unmixing @ X
    return recovered, unmixing

recovered, W_unmix = fastica(mixtures)

# 修复符号歧义（ICA 可能翻转符号）
for i in range(2):
    if jnp.corrcoef(recovered[i], sources[i])[0, 1] < -0.5:
        recovered = recovered.at[i].set(-recovered[i])

# 如果源被交换，修复置换
corr_00 = jnp.abs(jnp.corrcoef(recovered[0], sources[0])[0, 1])
corr_01 = jnp.abs(jnp.corrcoef(recovered[0], sources[1])[0, 1])
if corr_01 > corr_00:
    recovered = recovered[::-1]

# 归一化以便显示
recovered = recovered / jnp.max(jnp.abs(recovered), axis=1, keepdims=True)

fig, axes = plt.subplots(3, 2, figsize=(14, 9))

axes[0, 0].plot(t[:1000], s1[:1000], color='#3498db', linewidth=0.8)
axes[0, 0].set_title('Source 1 (Original)')
axes[0, 0].set_ylabel('Amplitude')

axes[0, 1].plot(t[:1000], s2[:1000], color='#e74c3c', linewidth=0.8)
axes[0, 1].set_title('Source 2 (Original)')

axes[1, 0].plot(t[:1000], mixtures[0, :1000], color='#9b59b6', linewidth=0.8)
axes[1, 0].set_title('Mixture 1 (Microphone 1)')
axes[1, 0].set_ylabel('Amplitude')

axes[1, 1].plot(t[:1000], mixtures[1, :1000], color='#9b59b6', linewidth=0.8)
axes[1, 1].set_title('Mixture 2 (Microphone 2)')

axes[2, 0].plot(t[:1000], recovered[0, :1000], color='#27ae60', linewidth=0.8)
axes[2, 0].set_title('Recovered Source 1 (FastICA)')
axes[2, 0].set_ylabel('Amplitude')
axes[2, 0].set_xlabel('Time (s)')

axes[2, 1].plot(t[:1000], recovered[1, :1000], color='#f39c12', linewidth=0.8)
axes[2, 1].set_title('Recovered Source 2 (FastICA)')
axes[2, 1].set_xlabel('Time (s)')

plt.tight_layout()
plt.show()

# 报告与原始信号的相关性
for i in range(2):
    corr = jnp.corrcoef(recovered[i], sources[i])[0, 1]
    print(f"Source {i+1} recovery correlation: {corr:.4f}")
​

• 任务 2：基于 NMF 的频谱图源分离。 使用非负矩阵分解（第 2 章）把频谱图分离为两个分量，演示 NMF 如何为每个源学习频谱字典。

import jax
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt

# 生成两个具有不同频谱特征的信号
sr = 8000
duration = 1.0
t = jnp.linspace(0, duration, int(sr * duration))

# 源 1：低频谐波（模拟贝斯）
src1 = (jnp.sin(2 * jnp.pi * 100 * t) +
        0.5 * jnp.sin(2 * jnp.pi * 200 * t) +
        0.3 * jnp.sin(2 * jnp.pi * 300 * t))

# 源 2：高频谐波（模拟长笛）
src2 = (jnp.sin(2 * jnp.pi * 800 * t) +
        0.4 * jnp.sin(2 * jnp.pi * 1600 * t))

# 时变振幅（源在不同时间活跃）
env1 = jnp.where(t < 0.5, 1.0, 0.3)
env2 = jnp.where(t > 0.3, 1.0, 0.2)
src1 = src1 * env1
src2 = src2 * env2

mixture = src1 + src2

# 计算幅度频谱图（STFT）
n_fft = 512
hop = 128
window = jnp.hanning(n_fft)

def compute_stft(signal, n_fft, hop, window):
    n_frames = 1 + (len(signal) - n_fft) // hop
    frames = jnp.stack([
        signal[i * hop : i * hop + n_fft] * window
        for i in range(n_frames)
    ])
    return jnp.fft.rfft(frames, n=n_fft)

S_mix = compute_stft(mixture, n_fft, hop, window)
V = jnp.abs(S_mix).T  # (F, T) - 频率 x 时间
phase = jnp.angle(S_mix).T

F, T = V.shape
print(f"Spectrogram shape: {F} freq bins x {T} time frames")

# NMF：V ≈ WH，使用乘性更新规则
def nmf(V, K, n_iter=200, key=jr.PRNGKey(0)):
    """Frobenius 范数下的非负矩阵分解。"""
    k1, k2 = jr.split(key)
    W = jnp.abs(jr.normal(k1, (F, K))) * 0.1 + 0.01  # (F, K)
    H = jnp.abs(jr.normal(k2, (K, T))) * 0.1 + 0.01  # (K, T)

    costs = []
    for i in range(n_iter):
        # H 的乘性更新
        WtV = W.T @ V
        WtWH = W.T @ W @ H + 1e-8
        H = H * (WtV / WtWH)

        # W 的乘性更新
        VHt = V @ H.T
        WHHt = W @ H @ H.T + 1e-8
        W = W * (VHt / WHHt)

        cost = jnp.sum((V - W @ H) ** 2)
        costs.append(float(cost))

    return W, H, costs

# 用 K=2 个分量运行 NMF
K = 2
W, H, costs = nmf(V, K, n_iter=300)

# 用软掩码重建每个源
V_hat = W @ H
mask1 = (W[:, 0:1] @ H[0:1, :]) / (V_hat + 1e-8)
mask2 = (W[:, 1:2] @ H[1:2, :]) / (V_hat + 1e-8)

V_src1 = mask1 * V
V_src2 = mask2 * V

# 可视化
fig, axes = plt.subplots(3, 2, figsize=(14, 10))

# 混合频谱图
axes[0, 0].imshow(jnp.log1p(V), aspect='auto', origin='lower', cmap='magma')
axes[0, 0].set_title('Mixture Spectrogram |X|')
axes[0, 0].set_ylabel('Frequency bin')

# NMF 收敛
axes[0, 1].plot(costs, color='#3498db', linewidth=1.5)
axes[0, 1].set_title('NMF Convergence')
axes[0, 1].set_xlabel('Iteration')
axes[0, 1].set_ylabel('Frobenius cost')
axes[0, 1].set_yscale('log')

# 频谱基向量 W
freq_hz = jnp.arange(F) * sr / n_fft
axes[1, 0].plot(freq_hz, W[:, 0], color='#27ae60', linewidth=1.5,
                label='Basis 1 (low freq)')
axes[1, 0].plot(freq_hz, W[:, 1], color='#e74c3c', linewidth=1.5,
                label='Basis 2 (high freq)')
axes[1, 0].set_title('Learned Spectral Bases W')
axes[1, 0].set_xlabel('Frequency (Hz)')
axes[1, 0].set_ylabel('Magnitude')
axes[1, 0].legend()

# 时间激活 H
time_s = jnp.arange(T) * hop / sr
axes[1, 1].plot(time_s, H[0], color='#27ae60', linewidth=1.5,
                label='Activation 1')
axes[1, 1].plot(time_s, H[1], color='#e74c3c', linewidth=1.5,
                label='Activation 2')
axes[1, 1].set_title('Temporal Activations H')
axes[1, 1].set_xlabel('Time (s)')
axes[1, 1].set_ylabel('Activation')
axes[1, 1].legend()

# 分离后的频谱图
axes[2, 0].imshow(jnp.log1p(V_src1), aspect='auto', origin='lower', cmap='magma')
axes[2, 0].set_title('Separated Source 1 (low-frequency)')
axes[2, 0].set_ylabel('Frequency bin')
axes[2, 0].set_xlabel('Time frame')

axes[2, 1].imshow(jnp.log1p(V_src2), aspect='auto', origin='lower', cmap='magma')
axes[2, 1].set_title('Separated Source 2 (high-frequency)')
axes[2, 1].set_xlabel('Time frame')

plt.tight_layout()
plt.show()

print(f"Reconstruction error: {jnp.sum((V - W @ H)**2):.2f}")
print(f"NMF learns spectral bases that capture each source's frequency profile.")
​

• 任务 3：用于降噪的 LMS 自适应滤波器。 实现 LMS 和 NLMS 算法用于回声/噪声消除，展示收敛行为和步长的影响。

import jax
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt

# 模拟回声消除场景
# 对端信号 -> 房间脉冲响应 -> 麦克风处的回声
# 近端语音是我们想要保留的期望信号

sr = 8000
duration = 2.0
n_samples = int(sr * duration)
key = jr.PRNGKey(42)
keys = jr.split(key, 5)

# 对端信号（参考）：随机的类语音信号
far_end = jr.normal(keys[0], (n_samples,)) * 0.5

# 房间脉冲响应（对算法未知）
rir_length = 64
rir = jnp.zeros(rir_length)
rir = rir.at[0].set(0.8)   # 直射路径
rir = rir.at[5].set(0.3)   # 早期反射
rir = rir.at[12].set(-0.2) # 反射
rir = rir.at[25].set(0.1)  # 晚期反射
rir = rir.at[40].set(-0.05)

# 回声：对端信号与 RIR 的卷积
echo = jnp.convolve(far_end, rir)[:n_samples]

# 近端语音（在信号的一部分中活跃）
near_end = jnp.zeros(n_samples)
start, end = n_samples // 3, 2 * n_samples // 3
near_speech = 0.3 * jnp.sin(
    2 * jnp.pi * 300 * jnp.linspace(0, (end - start) / sr, end - start)
)
near_end = near_end.at[start:end].set(near_speech)

# 麦克风信号：回声 + 近端 + 噪声
noise = jr.normal(keys[1], (n_samples,)) * 0.01
mic_signal = echo + near_end + noise

# LMS 自适应滤波器
def lms_filter(reference, desired, filter_length, mu):
    """标准 LMS 自适应滤波器。"""
    n = len(reference)
    w = jnp.zeros(filter_length)
    output = jnp.zeros(n)
    error = jnp.zeros(n)
    w_history = []

    for i in range(filter_length, n):
        x = reference[i:i-filter_length:-1]  # 反转段
        if len(x) < filter_length:
            x = jnp.pad(x, (0, filter_length - len(x)))
        x = reference[max(0, i-filter_length+1):i+1][::-1]

        y = jnp.dot(w, x)
        e = desired[i] - y
        w = w + mu * e * x

        output = output.at[i].set(y)
        error = error.at[i].set(e)

        if i % 500 == 0:
            w_history.append(w.copy())

    return output, error, w_history

# NLMS 自适应滤波器
def nlms_filter(reference, desired, filter_length, mu, eps=1e-6):
    """归一化 LMS 自适应滤波器。"""
    n = len(reference)
    w = jnp.zeros(filter_length)
    output = jnp.zeros(n)
    error = jnp.zeros(n)

    for i in range(filter_length, n):
        x = reference[max(0, i-filter_length+1):i+1][::-1]

        y = jnp.dot(w, x)
        e = desired[i] - y
        norm_factor = jnp.dot(x, x) + eps
        w = w + (mu / norm_factor) * e * x

        output = output.at[i].set(y)
        error = error.at[i].set(e)

    return output, error

# 用不同步长运行 LMS
filter_len = 64
mu_values = [0.001, 0.01, 0.05]
colors_mu = ['#3498db', '#e74c3c', '#27ae60']

fig, axes = plt.subplots(2, 2, figsize=(14, 10))

# 原始信号
t = jnp.arange(n_samples) / sr
axes[0, 0].plot(t, mic_signal, color='#9b59b6', linewidth=0.5, alpha=0.7,
                label='Mic (echo + near-end)')
axes[0, 0].plot(t, echo, color='#e74c3c', linewidth=0.5, alpha=0.7,
                label='Echo (to cancel)')
axes[0, 0].plot(t, near_end, color='#27ae60', linewidth=0.8,
                label='Near-end speech (to preserve)')
axes[0, 0].set_title('Signal Components')
axes[0, 0].set_xlabel('Time (s)')
axes[0, 0].set_ylabel('Amplitude')
axes[0, 0].legend(fontsize=8)

# 不同步长的 LMS 收敛
for mu, color in zip(mu_values, colors_mu):
    _, err, _ = lms_filter(far_end, mic_signal, filter_len, mu)
    # 平方的误差平滑
    sq_err = err ** 2
    window_size = 200
    smoothed = jnp.convolve(sq_err, jnp.ones(window_size)/window_size,
                             mode='valid')
    axes[0, 1].plot(smoothed, color=color, linewidth=1.2,
                    label=f'mu={mu}')

axes[0, 1].set_title('LMS Convergence (smoothed MSE)')
axes[0, 1].set_xlabel('Sample')
axes[0, 1].set_ylabel('Squared Error')
axes[0, 1].set_yscale('log')
axes[0, 1].legend()

# 最佳 LMS 结果
_, err_lms, w_hist = lms_filter(far_end, mic_signal, filter_len, 0.01)
axes[1, 0].plot(t, mic_signal, color='#9b59b6', linewidth=0.5, alpha=0.4,
                label='Before cancellation')
axes[1, 0].plot(t, err_lms, color='#3498db', linewidth=0.5, alpha=0.8,
                label='After LMS cancellation')
axes[1, 0].plot(t, near_end, color='#27ae60', linewidth=0.8, alpha=0.5,
                label='True near-end')
axes[1, 0].set_title('LMS Echo Cancellation Result (mu=0.01)')
axes[1, 0].set_xlabel('Time (s)')
axes[1, 0].set_ylabel('Amplitude')
axes[1, 0].legend(fontsize=8)

# NLMS 结果
_, err_nlms = nlms_filter(far_end, mic_signal, filter_len, 0.5)
axes[1, 1].plot(t, mic_signal, color='#9b59b6', linewidth=0.5, alpha=0.4,
                label='Before cancellation')
axes[1, 1].plot(t, err_nlms, color='#f39c12', linewidth=0.5, alpha=0.8,
                label='After NLMS cancellation')
axes[1, 1].plot(t, near_end, color='#27ae60', linewidth=0.8, alpha=0.5,
                label='True near-end')
axes[1, 1].set_title('NLMS Echo Cancellation Result (mu=0.5)')
axes[1, 1].set_xlabel('Time (s)')
axes[1, 1].set_ylabel('Amplitude')
axes[1, 1].legend(fontsize=8)

plt.tight_layout()
plt.show()

# 衡量回声衰减
echo_power = jnp.mean(echo ** 2)
lms_residual = jnp.mean(err_lms[n_samples//2:] ** 2)  # 收敛后
nlms_residual = jnp.mean(err_nlms[n_samples//2:] ** 2)
print(f"Echo power: {10*jnp.log10(echo_power):.1f} dB")
print(f"LMS residual: {10*jnp.log10(lms_residual):.1f} dB "
      f"(ERLE: {10*jnp.log10(echo_power/lms_residual):.1f} dB)")
print(f"NLMS residual: {10*jnp.log10(nlms_residual):.1f} dB "
      f"(ERLE: {10*jnp.log10(echo_power/nlms_residual):.1f} dB)")
​

• 任务 4：用于语音增强的时频掩码。 实现一种简单的频谱掩码方法（理想比值掩码）并与谱减法对比，可视化在合成含噪语音信号上的分离质量。

import jax
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt

# 创建合成的"语音"和"噪声"信号
sr = 8000
duration = 2.0
t = jnp.linspace(0, duration, int(sr * duration))

# 语音：带时变振幅的谐波序列（模拟语音）
speech = jnp.zeros_like(t)
for f0 in [150, 300, 450, 600, 900]:
    amp_env = 0.5 + 0.5 * jnp.sin(2 * jnp.pi * 2.0 * t)  # 2 Hz 调制
    speech = speech + (0.5 / (f0/150)) * amp_env * jnp.sin(2 * jnp.pi * f0 * t)
speech = speech / jnp.max(jnp.abs(speech))

# 噪声：带限噪声
key = jr.PRNGKey(42)
noise_raw = jr.normal(key, t.shape) * 0.4

# 以给定信噪比混合
snr_db = 5.0
speech_power = jnp.mean(speech ** 2)
noise_power = jnp.mean(noise_raw ** 2)
noise_scale = jnp.sqrt(speech_power / (noise_power * 10 ** (snr_db / 10)))
noise = noise_raw * noise_scale
mixture = speech + noise

# STFT
n_fft = 512
hop = 128
window = jnp.hanning(n_fft)

def stft(signal, n_fft, hop, window):
    n_frames = 1 + (len(signal) - n_fft) // hop
    frames = jnp.stack([
        signal[i * hop : i * hop + n_fft] * window
        for i in range(n_frames)
    ])
    return jnp.fft.rfft(frames, n=n_fft)

def istft(S, hop, window, length):
    n_fft = (S.shape[1] - 1) * 2
    n_frames = S.shape[0]
    frames = jnp.fft.irfft(S, n=n_fft) * window[None, :]
    output = jnp.zeros(length)
    window_sum = jnp.zeros(length)
    for i in range(n_frames):
        start = i * hop
        end = start + n_fft
        if end <= length:
            output = output.at[start:end].add(frames[i])
            window_sum = window_sum.at[start:end].add(window ** 2)
    window_sum = jnp.maximum(window_sum, 1e-8)
    return output / window_sum

S_speech = stft(speech, n_fft, hop, window)
S_noise = stft(noise, n_fft, hop, window)
S_mix = stft(mixture, n_fft, hop, window)

mag_speech = jnp.abs(S_speech)
mag_noise = jnp.abs(S_noise)
mag_mix = jnp.abs(S_mix)
phase_mix = jnp.angle(S_mix)

# 方法 1：理想比值掩码（oracle - 上界）
irm = mag_speech ** 2 / (mag_speech ** 2 + mag_noise ** 2 + 1e-8)
S_irm = (irm * mag_mix) * jnp.exp(1j * phase_mix)
enhanced_irm = istft(S_irm, hop, window, len(mixture))

# 方法 2：谱减法
# 从前 0.2 秒估计噪声（假设为静音）
noise_frames = int(0.2 * sr / hop)
noise_est = jnp.mean(mag_mix[:noise_frames] ** 2, axis=0, keepdims=True)
alpha = 2.0  # 过减因子
beta = 0.02  # 谱底
mag_sub = jnp.maximum(mag_mix ** 2 - alpha * noise_est, beta * mag_mix ** 2)
mag_sub = jnp.sqrt(mag_sub)
S_sub = mag_sub * jnp.exp(1j * phase_mix)
enhanced_sub = istft(S_sub, hop, window, len(mixture))

# 方法 3：维纳滤波
snr_est = mag_mix ** 2 / (noise_est + 1e-8)
wiener_gain = snr_est / (1 + snr_est)
S_wiener = (wiener_gain * mag_mix) * jnp.exp(1j * phase_mix)
enhanced_wiener = istft(S_wiener, hop, window, len(mixture))

# 计算每种方法的 SI-SDR
def si_sdr(estimate, reference):
    """尺度不变信号失真比。"""
    ref = reference[:len(estimate)]
    est = estimate[:len(reference)]
    s_target = (jnp.dot(est, ref) / (jnp.dot(ref, ref) + 1e-8)) * ref
    e_noise = est - s_target
    return 10 * jnp.log10(jnp.dot(s_target, s_target) /
                           (jnp.dot(e_noise, e_noise) + 1e-8))

si_sdr_mix = si_sdr(mixture, speech)
si_sdr_irm_val = si_sdr(enhanced_irm, speech)
si_sdr_sub_val = si_sdr(enhanced_sub, speech)
si_sdr_wiener_val = si_sdr(enhanced_wiener, speech)

# 可视化
fig, axes = plt.subplots(3, 2, figsize=(14, 12))

# 频谱图
axes[0, 0].imshow(jnp.log1p(mag_speech.T), aspect='auto', origin='lower',
                   cmap='magma')
axes[0, 0].set_title('Clean Speech Spectrogram')
axes[0, 0].set_ylabel('Frequency bin')

axes[0, 1].imshow(jnp.log1p(mag_mix.T), aspect='auto', origin='lower',
                   cmap='magma')
axes[0, 1].set_title(f'Noisy Mixture ({snr_db:.0f} dB SNR)')

# 掩码
axes[1, 0].imshow(irm.T, aspect='auto', origin='lower', cmap='RdYlGn')
axes[1, 0].set_title('Ideal Ratio Mask (Oracle)')
axes[1, 0].set_ylabel('Frequency bin')

axes[1, 1].imshow(wiener_gain.T, aspect='auto', origin='lower', cmap='RdYlGn',
                   vmin=0, vmax=1)
axes[1, 1].set_title('Estimated Wiener Gain')

# 增强后波形对比
n_show = 3000
axes[2, 0].plot(t[:n_show], speech[:n_show], color='#27ae60', linewidth=0.8,
                alpha=0.5, label='Clean')
axes[2, 0].plot(t[:n_show], mixture[:n_show], color='#e74c3c', linewidth=0.5,
                alpha=0.4, label='Noisy')
axes[2, 0].plot(t[:n_show], enhanced_irm[:n_show], color='#3498db',
                linewidth=0.8, label='IRM enhanced')
axes[2, 0].set_title('Waveform Comparison (IRM)')
axes[2, 0].set_xlabel('Time (s)')
axes[2, 0].set_ylabel('Amplitude')
axes[2, 0].legend(fontsize=8)

# SI-SDR 柱状图
methods = ['Mixture', 'Spectral\nSubtraction', 'Wiener\nFilter', 'Ideal Ratio\nMask']
sdr_values = [float(si_sdr_mix), float(si_sdr_sub_val),
              float(si_sdr_wiener_val), float(si_sdr_irm_val)]
bar_colors = ['#e74c3c', '#f39c12', '#9b59b6', '#27ae60']
bars = axes[2, 1].bar(methods, sdr_values, color=bar_colors, alpha=0.8)
axes[2, 1].set_ylabel('SI-SDR (dB)')
axes[2, 1].set_title('Enhancement Quality Comparison')
for bar, val in zip(bars, sdr_values):
    axes[2, 1].text(bar.get_x() + bar.get_width()/2., bar.get_height() + 0.3,
                    f'{val:.1f}', ha='center', fontsize=10)
axes[2, 1].axhline(0, color='gray', linestyle='--', linewidth=0.8)

plt.tight_layout()
plt.show()

print(f"SI-SDR (noisy mixture):        {si_sdr_mix:.2f} dB")
print(f"SI-SDR (spectral subtraction): {si_sdr_sub_val:.2f} dB")
print(f"SI-SDR (Wiener filter):        {si_sdr_wiener_val:.2f} dB")
print(f"SI-SDR (ideal ratio mask):     {si_sdr_irm_val:.2f} dB (oracle upper bound)")
​