Tách nguồn và Khử nhiễu

Tách nguồn và khử nhiễu phục hồi các tín hiệu riêng lẻ từ âm thanh hỗn hợp; bài toán cocktail party trong tính toán. File này bao quát ICA, NMF, mặt nạ thời gian-tần số, beamforming, mạng tách nguồn học sâu (Conv-TasNet, SepFormer), tăng cường giọng nói, và khử nhiễu thích ứng.

• Hãy tưởng tượng bạn đang đứng trong một bữa tiệc cocktail đông đúc. Hàng chục người đang nói chuyện cùng lúc, nhạc đang phát, ly cốc leng keng, thế mà bạn vẫn có thể tập trung vào một cuộc trò chuyện và theo dõi nó rõ ràng. Khả năng đáng kinh ngạc này, được gọi là bài toán cocktail party (Cherry, 1953), là điều mà hệ thống thính giác của con người giải quyết một cách dễ dàng nhưng máy móc lại thấy vô cùng khó khăn. File này bao quát các thuật toán cố gắng làm điều đó: tách các nguồn âm thanh hỗn hợp, khử nhiễu không mong muốn, và tăng cường giọng nói trong điều kiện bất lợi.

• Các nền tảng xử lý tín hiệu từ file 01 (STFT, spectrogram, bộ lọc) nằm dưới mọi phương pháp ở đây. Các kỹ thuật phân rã ma trận từ chương 02 (NMF, ICA, SVD) cung cấp bộ công cụ cổ điển. Các kiến trúc học sâu từ chương 06 (CNN, RNN, attention) và lý thuyết xác suất từ chương 04/05 hỗ trợ các cách tiếp cận hiện đại.

• Phát biểu bài toán: một tín hiệu hỗn hợp \(x(t)\) được quan sát tại một hoặc nhiều microphone. Hỗn hợp này là tổng (trong trường hợp đơn giản nhất) của \(C\) tín hiệu nguồn:

\[x(t) = \sum_{c=1}^{C} s_c(t) + n(t)\]

• trong đó \(s_c(t)\) là tín hiệu nguồn thứ \(c\) và \(n(t)\) là nhiễu nền. Mục tiêu là phục hồi từng \(s_c(t)\) riêng lẻ từ \(x(t)\). Trong trường hợp một microphone, bài toán này bị thiếu điều kiện nghiêm trọng: một phương trình, \(C\) ẩn số. Các giả định bổ sung (độc lập thống kê, cấu trúc phổ, prior đã học) là cần thiết để bài toán trở nên khả thi.

• Trong miền tần số (thông qua STFT từ file 01), hỗn hợp trở thành:

\[X(t, f) = \sum_{c=1}^{C} S_c(t, f) + N(t, f)\]

• Nhiều phương pháp tách nguồn làm việc trong miền thời gian-tần số bằng cách ước lượng một mặt nạ \(M_c(t, f) \in [0, 1]\) cho mỗi nguồn, sau đó phục hồi nguồn dưới dạng \(\hat{S}_c(t, f) = M_c(t, f) \cdot X(t, f)\). Mặt nạ nhị phân lý tưởng (IBM) đặt \(M_c(t, f) = 1\) nếu nguồn \(c\) chiếm ưu thế tại bin thời gian-tần số đó và 0 nếu ngược lại. Mặt nạ tỷ lệ lý tưởng (IRM) là một phiên bản mềm:

\[\text{IRM}_c(t, f) = \frac{|S_c(t, f)|^2}{\sum_{j=1}^{C} |S_j(t, f)|^2}\]

• Phân tích thành phần độc lập (ICA) là cách tiếp cận cổ điển khi số lượng microphone bằng hoặc vượt quá số lượng nguồn. ICA (chương 02) tìm một ma trận phân tách tuyến tính \(W\) sao cho \(\hat{s} = Wx\), trong đó các nguồn được phục hồi \(\hat{s}\) là độc lập thống kê tối đa. Giả định then chốt là các tín hiệu nguồn là không Gaussian và độc lập, điều này thường đúng cho giọng nói và nhạc.

• Đối với mô hình trộn tức thời nhiều microphone \(x = As\) (trong đó \(A\) là ma trận trộn), ICA phục hồi \(W \approx A^{-1}\) bằng cách tối đa hóa tính không Gaussian của đầu ra (FastICA dùng negentropy) hoặc bằng cách tối thiểu hóa thông tin tương hỗ. ICA hoạt động tốt trong các môi trường có kiểm soát nhưng thất bại khi việc trộn liên quan đến tích chập (dội âm phòng), khi số nguồn vượt quá số microphone, hoặc khi giả định độc lập bị vi phạm.

• Phân rã ma trận không âm (NMF) phân rã spectrogram biên độ \(V \in \mathbb{R}_+^{F \times T}\) thành tích của hai ma trận không âm (chương 02):

\[V \approx WH\]

• trong đó \(W \in \mathbb{R}_+^{F \times K}\) là một từ điển gồm \(K\) vector cơ sở phổ và \(H \in \mathbb{R}_+^{K \times T}\) chứa các hệ số kích hoạt theo thời gian. Ràng buộc không âm có cơ sở vật lý: biên độ là không âm, và âm thanh kết hợp cộng tính.

• Để tách nguồn, NMF học các từ điển riêng cho mỗi nguồn: \(W_{\text{speech}}\) nắm bắt các mẫu phổ của giọng nói (cấu trúc formant), trong khi \(W_{\text{noise}}\) nắm bắt các mẫu nhiễu. Hỗn hợp được phân rã dưới dạng \(V \approx W_{\text{speech}} H_{\text{speech}} + W_{\text{noise}} H_{\text{noise}}\), và mỗi nguồn được phục hồi bằng cách đánh mặt nạ. NMF được tối thiểu hóa bằng cách sử dụng cập nhật nhân với chuẩn Frobenius hoặc phân kỳ KL làm hàm chi phí:

\[ \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} \]

• Beamforming khai thác thông tin không gian từ các dàn microphone. Khi một tín hiệu nguồn đến các microphone khác nhau với độ trễ khác nhau (do sự sắp xếp không gian), các độ trễ này có thể được sử dụng để tăng cường tín hiệu từ một hướng trong khi triệt tiêu các hướng khác.

• Beamforming delay-and-sum là cách tiếp cận đơn giản nhất. Nếu nguồn mong muốn ở góc \(\theta\) so với dàn, độ trễ thời gian tại microphone \(m\) là \(\tau_m(\theta) = d_m \sin \theta / c\), trong đó \(d_m\) là vị trí microphone và \(c\) là tốc độ âm thanh. Đầu ra beamformer căn chỉnh và cộng các tín hiệu microphone:

\[y(t) = \frac{1}{M} \sum_{m=1}^{M} x_m(t - \tau_m(\theta))\]

• Các tín hiệu từ hướng mục tiêu cộng hưởng một cách kết hợp, trong khi các tín hiệu từ hướng khác cộng hưởng một cách không kết hợp, tạo ra lọc không gian. Hình dạng dàn xác định độ phân giải không gian: dàn lớn hơn tạo ra chùm hẹp hơn.

• Beamforming MVDR (Minimum Variance Distortionless Response) tối ưu hóa các trọng số để tối thiểu hóa tổng công suất đầu ra trong khi vẫn giữ nguyên hướng mục tiêu không bị méo:

\[ \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} \]

• trong đó \(\Phi_{nn}\) là ma trận hiệp phương sai nhiễu không gian và \(\mathbf{d}(\theta)\) là vector steering cho hướng \(\theta\). Nghiệm dạng đóng là:

\[\mathbf{w}_{\text{MVDR}} = \frac{\Phi_{nn}^{-1} \mathbf{d}(\theta)}{\mathbf{d}(\theta)^H \Phi_{nn}^{-1} \mathbf{d}(\theta)}\]

• MVDR thích ứng với môi trường nhiễu bằng cách sử dụng hiệp phương sai nhiễu đã ước lượng, cung cấp khả năng triệt nhiễu tốt hơn delay-and-sum. Nó được sử dụng rộng rãi trong máy trợ thính, loa thông minh và hệ thống hội nghị truyền hình.

• Học sâu cho tách nguồn đã cải thiện hiệu năng một cách đáng kể, đặc biệt trong trường hợp một microphone, nơi các phương pháp cổ điển gặp khó khăn. Mô hình tổng quát là: mã hóa hỗn hợp, ước lượng mặt nạ hoặc biểu diễn nguồn bằng mạng nơ-ron, và giải mã để phục hồi các nguồn riêng lẻ.

• Deep clustering (Hershey et al., 2016) nhúng mỗi bin thời gian-tần số vào một không gian nhiều chiều nơi các bin thuộc cùng một nguồn nằm gần nhau và các bin từ các nguồn khác nhau nằm xa nhau. Một LSTM hai chiều (chương 06) ánh xạ mỗi bin T-F \((t, f)\) thành một embedding \(v_{t,f} \in \mathbb{R}^D\). Mục tiêu huấn luyện là:

\[\mathcal{L} = \|VV^T - YY^T\|_F^2\]

• trong đó \(V\) là ma trận các embedding và \(Y\) là ma trận one-hot của các gán nguồn. Tích \(VV^T\) là một ma trận affinity (mức độ giống nhau giữa các embedding của hai bin), và \(YY^T\) là affinity lý tưởng (1 nếu cùng nguồn, 0 nếu ngược lại). Khi suy luận, phân cụm K-means trên các embedding tạo ra các mặt nạ nhị phân.

• Conv-TasNet (Luo and Mesgarani, 2019) hoạt động hoàn toàn trong miền thời gian, bỏ qua STFT. Nó có ba thành phần:

• Bộ mã hóa: một tích chập 1D ánh xạ các đoạn ngắn của dạng sóng hỗn hợp thành biểu diễn tiềm ẩn. Với hỗn hợp \(x \in \mathbb{R}^T\), đầu ra bộ mã hóa là \(w = \text{ReLU}(U \ast x) \in \mathbb{R}^{N \times L}\), trong đó \(U\) là một cơ sở có thể học (tương tự như cơ sở STFT nhưng học từ dữ liệu), \(N\) là số hàm cơ sở, và \(L\) là số đoạn. Kích thước kernel và stride của bộ mã hóa (thường 2ms và 1ms) xác định độ phân giải thời gian.

• Bộ tách nguồn: một Mạng tích chập thời gian (TCN) xử lý hỗn hợp đã mã hóa và xuất ra \(C\) mặt nạ. TCN xếp chồng các tích chập tách được theo chiều sâu 1D giãn (từ chương 08 về tích chập hiệu quả) trong các khối với hệ số giãn tăng dần theo cấp số nhân \(1, 2, 4, \ldots, 2^{B-1}\), lặp lại \(R\) lần. Điều này tạo ra một trường tiếp nhận rất lớn trong khi vẫn giữ hiệu quả tính toán.

• Bộ giải mã: một tích chập 1D chuyển vị (với cơ sở học được \(V\)) chuyển đổi mỗi biểu diễn đã được đánh mặt nạ trở lại miền thời gian: \(\hat{s}_c = V^T (M_c \odot w)\).

• Conv-TasNet vượt trội đáng kể so với các phương pháp dựa trên spectrogram vì cơ sở mã hóa-giải mã học được có thể nắm bắt thông tin (đặc biệt là pha) mà STFT biên độ loại bỏ.

• DPRNN (Dual-Path RNN) (Luo et al., 2020) giải quyết vấn đề mô hình hóa chuỗi dài trong tách nguồn. Thay vì xử lý toàn bộ chuỗi đã mã hóa bằng một RNN hoặc TCN duy nhất, DPRNN chia chuỗi thành các khối chồng lấp và áp dụng RNN dọc theo hai đường: một đường nội khối (intra-chunk, mô hình hóa các mẫu cục bộ trong mỗi khối) và một đường liên khối (inter-chunk, mô hình hóa các mẫu toàn cục giữa các khối). Điều này giảm độ dài chuỗi RNN từ \(L\) xuống \(\sqrt{L}\) trong mỗi chiều:

\[ \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} \]

• trong đó \(k\) chỉ số khối và \(n\) chỉ số vị trí trong khối. LSTM nội khối xử lý theo \(n\) cho một \(k\) cố định; LSTM liên khối xử lý theo \(k\) cho một \(n\) cố định.

• SepFormer (Subakan et al., 2021) thay thế các RNN trong khung dual-path bằng các transformer (chương 07). Transformer nội khối nắm bắt các phụ thuộc cục bộ với self-attention, và transformer liên khối nắm bắt các phụ thuộc toàn cục. Khả năng mô hình hóa các phụ thuộc tầm xa của multi-head attention mà không gặp vấn đề gradient biến mất (chương 06) làm cho SepFormer đặc biệt hiệu quả cho các bản ghi dài. SepFormer đạt kết quả tiên tiến nhất trên benchmark WSJ0-2mix.

• Huấn luyện bất biến hoán vị (PIT) giải quyết một vấn đề cơ bản trong tách nguồn có giám sát: sự nhập nhằng trong gán nhãn. Nếu mạng có hai đầu ra (cho hai người nói), đầu ra nào nên tương ứng với người nói nào? Không có thứ tự tự nhiên. PIT tính toán mất mát cho mọi cách gán có thể và lấy giá trị nhỏ nhất:

\[\mathcal{L}_{\text{PIT}} = \min_{\pi \in \mathcal{P}} \sum_{c=1}^{C} \ell(\hat{s}_{\pi(c)}, s_c)\]

• trong đó \(\mathcal{P}\) là tập hợp tất cả các hoán vị của \(\{1, \ldots, C\}\) và \(\ell\) là mất mát trên mỗi nguồn (thường là tỷ lệ tín hiệu trên méo bất biến tỷ lệ, SI-SDR). Với \(C = 2\) nguồn chỉ có 2 hoán vị; với \(C = 3\) có 6. Việc này được tính hiệu quả bằng thuật toán Hungary cho \(C\) lớn hơn.

• Tỷ lệ tín hiệu trên méo bất biến tỷ lệ (SI-SDR) là metric đánh giá chuẩn cho tách nguồn:

\[ \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} \]

• trong đó \(\hat{s}\) là nguồn ước lượng và \(s\) là ground truth. SI-SDR bất biến với tỷ lệ tổng thể của ước lượng, điều này mong muốn vì âm lượng tuyệt đối ít quan trọng hơn chất lượng tách. SI-SDR càng cao (tính bằng dB) càng tốt. Các hệ thống tiên tiến nhất đạt mức cải thiện SI-SDR khoảng 20-22 dB trên WSJ0-2mix.

• Tách nguồn nhạc tách một bản ghi âm nhạc thành các phần nhạc cụ riêng: giọng hát, trống, bass, và các nhạc cụ khác. Điều này cho phép các ứng dụng như karaoke (loại bỏ giọng hát), phối lại (điều chỉnh mức nhạc cụ), và phiên âm (phân tích từng nhạc cụ một cách riêng lẻ).

• Open-Unmix (Stoter et al., 2019) là một baseline tham chiếu sử dụng LSTM hai chiều 3 lớp để dự đoán mặt nạ mềm cho mỗi nguồn trong miền STFT biên độ. Nó xử lý từng nguồn độc lập với một mô hình chuyên dụng. Đơn giản nhưng hiệu quả, Open-Unmix thiết lập các benchmark có thể tái lập trên MUSDB18.

• Demucs (Defossez et al., 2019; cập nhật thành Hybrid Demucs, 2021) sử dụng kiến trúc U-Net (chương 08) hoạt động trực tiếp trên dạng sóng. Bộ mã hóa nén hỗn hợp qua các tích chập stride, bộ giải mã mở rộng nó trở lại qua các tích chập chuyển vị với kết nối skip, và mỗi nguồn có một đầu giải mã riêng. Hybrid Demucs kết hợp xử lý miền thời gian và miền tần số: bộ mã hóa có các nhánh miền thời gian và STFT song song mà các đặc trưng của chúng được hợp nhất trước bộ giải mã. Điều này nắm bắt cả chi tiết thời gian tinh vi và cấu trúc phổ.

• Demucs đạt chất lượng tách nguồn tiên tiến nhất trên MUSDB18, đặc biệt mạnh về tách giọng hát. Kiến trúc U-Net của nó gợi nhớ đến các kiến trúc phân đoạn ảnh từ chương 08, coi bài toán tách nguồn như một dạng \"phân đoạn âm thanh\".

• Khử nhiễu chủ động (ANC) giảm âm thanh không mong muốn bằng cách tạo ra tín hiệu chống nhiễu giao thoa triệt tiêu với nhiễu. Hãy nghĩ đến tai nghe chống ồn: một microphone thu âm thanh môi trường, hệ thống ANC tạo ra một phiên bản đảo ngược, và tín hiệu kết hợp (nhiễu + chống nhiễu) lý tưởng sẽ triệt tiêu thành im lặng.

• Vật lý rất đơn giản: nếu nhiễu là \(n(t)\), tạo ra \(-n(t)\) tại cùng một điểm trong không gian tạo ra im lặng: \(n(t) + (-n(t)) = 0\). Thách thức là tín hiệu chống nhiễu phải được căn chỉnh chính xác về thời gian, biên độ và pha. Sai sót dù nhỏ cũng tạo ra nhiễu dư hoặc artefact.

• ANC feedforward sử dụng một microphone tham chiếu thu nhiễu trước khi nó đến người nghe. Hệ thống có thời gian để xử lý nhiễu và tạo ra tín hiệu chống nhiễu. Tín hiệu tham chiếu đi qua một bộ lọc thích ứng mà đầu ra của nó được trừ khỏi nhiễu tại microphone lỗi (gần người nghe). Điều này hoạt động tốt cho nhiễu băng rộng có thể dự đoán được (tiếng ồn động cơ, tiếng quạt).

• ANC feedback chỉ sử dụng một microphone lỗi tại tai người nghe. Hệ thống ước lượng nhiễu từ tín hiệu dư (những gì người nghe thực sự nghe thấy) và điều chỉnh tín hiệu chống nhiễu. ANC feedback đơn giản hơn (không cần microphone tham chiếu) nhưng có băng thông hạn chế và có thể mất ổn định.

• Lọc thích ứng là động cơ toán học đằng sau ANC. Các hệ số bộ lọc phải liên tục thích ứng với môi trường nhiễu thay đổi. Thuật toán phổ biến nhất là bộ lọc LMS (Least Mean Squares).

• Thuật toán LMS: một bộ lọc FIR với các hệ số \(\mathbf{w} = [w_0, w_1, \ldots, w_{L-1}]^T\) xử lý tín hiệu tham chiếu \(\mathbf{x}(n) = [x(n), x(n-1), \ldots, x(n-L+1)]^T\). Đầu ra là \(y(n) = \mathbf{w}^T \mathbf{x}(n)\), lỗi là \(e(n) = d(n) - y(n)\) (trong đó \(d(n)\) là tín hiệu mong muốn/chính), và cập nhật trọng số là:

\[\mathbf{w}(n+1) = \mathbf{w}(n) + \mu \, e(n) \, \mathbf{x}(n)\]

• trong đó \(\mu\) là kích thước bước (learning rate). Đây là một bước hạ gradient ngẫu nhiên trên lỗi bình phương trung bình \(E[e^2(n)]\), sử dụng ước lượng gradient tức thời \(-2 e(n) \mathbf{x}(n)\) thay vì gradient thực (hạ gradient từ chương 03 và SGD từ chương 06).

• Kích thước bước \(\mu\) kiểm soát sự đánh đổi giữa tốc độ hội tụ và lỗi trạng thái ổn định. Quá lớn thì bộ lọc dao động hoặc phân kỳ; quá nhỏ thì sự thích ứng ì ạch. Điều kiện ổn định là \(0 < \mu < 2 / (\lambda_{\max})\), trong đó \(\lambda_{\max}\) là trị riêng lớn nhất của ma trận tự tương quan đầu vào \(R = E[\mathbf{x}\mathbf{x}^T]\).

• NLMS (Normalised LMS) chuẩn hóa kích thước bước theo công suất đầu vào, làm cho sự hội tụ độc lập với mức tín hiệu:

\[\mathbf{w}(n+1) = \mathbf{w}(n) + \frac{\mu}{\|\mathbf{x}(n)\|^2 + \epsilon} \, e(n) \, \mathbf{x}(n)\]

• trong đó \(\epsilon\) là một hằng số chính quy hóa nhỏ để ngăn chia cho 0. NLMS hội tụ đáng tin cậy hơn LMS vì kích thước bước hiệu dụng thích ứng với công suất đầu vào.

• RLS (Recursive Least Squares) là một phương pháp thay thế hội tụ nhanh hơn, tối thiểu hóa chi phí bình phương tối thiểu có trọng số \(\sum_{k=1}^{n} \lambda^{n-k} e^2(k)\), trong đó \(\lambda \in (0, 1]\) là hệ số quên (forgetting factor). RLS duy trì một ước lượng của ma trận tự tương quan nghịch đảo và cập nhật nó một cách đệ quy, đạt được sự hội tụ tối ưu với chi phí \(O(L^2)\) tính toán trên mỗi mẫu (so với \(O(L)\) cho LMS).

• Giảm nhiễu và tăng cường giọng nói nhằm cải thiện chất lượng và độ rõ của giọng nói trong các bản ghi ồn. Không giống như tách nguồn (tách các nguồn riêng biệt), tăng cường giọng nói đặc biệt nhắm vào trường hợp giọng nói cộng nhiễu, phục hồi giọng nói sạch từ quan sát ồn.

• Trừ phổ (spectral subtraction) là cách tiếp cận đơn giản nhất. Trong các frame chỉ có nhiễu (được phát hiện bởi VAD từ file 03), ước lượng phổ nhiễu \(|\hat{N}(f)|^2\). Sau đó trừ nó khỏi mỗi frame:

\[|\hat{S}(f)|^2 = \max(|X(f)|^2 - \alpha |\hat{N}(f)|^2, \beta |X(f)|^2)\]

• trong đó \(\alpha\) là hệ số trừ quá mức (thường 1-4, trừ mạnh loại bỏ nhiều nhiễu hơn nhưng tạo ra nhiều artefact hơn) và \(\beta\) là một sàn phổ ngăn giá trị âm và giảm artefact \"nhiễu nhạc\" (musical noise, các tàn dư âm sắc cô lập nghe như các nốt nhạc ngẫu nhiên).

• Lọc Wiener cung cấp ước lượng sai số bình phương trung bình tối thiểu của phổ giọng nói sạch:

\[\hat{S}(t, f) = \frac{|S(t,f)|^2}{|S(t,f)|^2 + |N(t,f)|^2} \cdot X(t, f) = G(t, f) \cdot X(t, f)\]

• Độ lợi Wiener \(G(t, f) = \text{SNR}(t, f) / (1 + \text{SNR}(t, f))\) nằm trong khoảng từ 0 (thuần nhiễu) đến 1 (thuần giọng nói), hoạt động như một mặt nạ mềm. Thách thức là ước lượng phổ công suất giọng nói và nhiễu. SNR tiên nghiệm \(\xi(t, f) = |S(t,f)|^2 / |N(t,f)|^2\) được ước lượng bằng cách tiếp cận \"decision-directed\": một kết hợp làm mịn giữa ước lượng của frame hiện tại và đầu ra đã lọc Wiener của frame trước.

• Tăng cường giọng nói nơ-ron sử dụng học sâu để ước lượng hoặc một mặt nạ (như độ lợi Wiener) hoặc trực tiếp spectrogram sạch. Các kiến trúc bao gồm từ mạng feedforward đơn giản đến U-Net (chương 08), CRN (Convolutional Recurrent Network), và transformer.

• DCCRN (Deep Complex Convolutional Recurrent Network) hoạt động trên STFT phức (cả biên độ và pha), sử dụng các tích chập phức tự nhiên xử lý phần thực và phần ảo. Điều này tránh bài toán ước lượng pha vốn gây khó khăn cho các cách tiếp cận chỉ dùng biên độ.

• FullSubNet sử dụng kiến trúc dual-path với một mô hình băng tần đầy (full-band, nắm bắt các mẫu phổ toàn cục) và một mô hình băng con (sub-band, nắm bắt các chi tiết hài địa phương). Mô hình băng tần đầy xử lý toàn bộ phổ, trong khi mô hình băng con xử lý các dải tần hẹp xoay quanh mỗi bin tần số. Đầu ra của chúng được kết hợp để cho ước lượng mặt nạ cuối cùng.

• Thử thách DNS (Deep Noise Suppression) của Microsoft đánh giá các hệ thống tăng cường giọng nói hàng năm. Người thắng thường sử dụng huấn luyện quy mô lớn với nhiều loại nhiễu đa dạng, tăng cường dữ liệu (thêm nhiễu ở các SNR khác nhau, dội âm, artefact codec), và các kiến trúc có khả năng chạy thời gian thực.

• Khử tiếng vọng loại bỏ tiếng vọng âm thanh trong giao tiếp hai chiều. Khi bạn đang trong cuộc gọi điện thoại, giọng của người ở đầu xa phát qua loa của bạn, dội quanh phòng và được microphone của bạn thu lại, tạo ra tiếng vọng mà người ở đầu xa nghe thấy. Khử tiếng vọng âm thanh (AEC) mô hình hóa đường đi âm thanh từ loa đến microphone và trừ đi tiếng vọng dự đoán.

• Đường đi âm thanh được mô hình hóa như một bộ lọc FIR thích ứng (dùng LMS hoặc NLMS) với tín hiệu đầu xa làm đầu vào. Bộ lọc mô hình hóa đáp ứng xung phòng, bao gồm đường đi trực tiếp, phản xạ sớm và dội âm muộn. Đáp ứng xung phòng có thể dài hàng trăm mili giây, đòi hỏi các bộ lọc với hàng nghìn tap.

• Phát hiện song công (double-talk detection) rất quan trọng đối với AEC: khi cả người nói đầu gần và đầu xa nói cùng lúc, bộ lọc thích ứng phải đóng băng (ngừng cập nhật) để ngăn nó khử giọng nói của người đầu gần. Bộ phát hiện song công so sánh năng lượng của tín hiệu lỗi với năng lượng tín hiệu đầu xa; sự tăng đột ngột trong năng lượng lỗi không được giải thích bởi tín hiệu đầu xa cho thấy có giọng nói đầu gần.

• Tương quan chéo chuẩn hóa giữa tín hiệu đầu xa \(x(n)\) và tín hiệu microphone \(d(n)\) cung cấp một chỉ báo song công:

\[\xi(n) = \frac{|\sum_{k=0}^{L-1} x(n-k) d(n-k)|}{\sqrt{\sum_{k} x^2(n-k)} \sqrt{\sum_{k} d^2(n-k)}}\]

• Trong suốt đơn công (chỉ đầu xa), \(\xi\) cao vì \(d\) hầu hết là tiếng vọng của \(x\). Trong suốt song công, \(\xi\) giảm vì giọng nói đầu gần không tương quan với \(x\).

• Các hệ thống AEC hiện đại kết hợp lọc thích ứng với mạng nơ-ron: bộ lọc thích ứng cung cấp ước lượng tiếng vọng ban đầu, và một mạng nơ-ron (tương tự các mô hình tăng cường giọng nói ở trên) làm sạch tiếng vọng dư và xử lý các phi tuyến (méo loa) mà bộ lọc tuyến tính không thể nắm bắt.

• Các metric đánh giá cho tách nguồn và tăng cường:

  • SI-SDR (đã định nghĩa ở trên): chuẩn cho tách nguồn.
  • SDR (Signal-to-Distortion Ratio): từ BSS Eval, đo chất lượng tách tổng thể bao gồm artefact và nhiễu.
  • PESQ (Perceptual Evaluation of Speech Quality): tiêu chuẩn ITU dự đoán điểm chất lượng chủ quan. Khoảng: -0.5 đến 4.5.
  • STOI (Short-Time Objective Intelligibility): dự đoán độ rõ giọng nói. Khoảng: 0 đến 1.
  • DNSMOS: bộ dự đoán MOS khử nhiễu sâu của Microsoft, một mạng nơ-ron được huấn luyện để dự đoán điểm MOS của con người mà không cần âm thanh tham chiếu sạch.

Coding Tasks (dùng CoLab hoặc notebook)

• Task 1: Independent Component Analysis để tách nguồn. Cài đặt FastICA để tách hai nguồn âm thanh hỗn hợp, minh họa giải pháp cocktail party cổ điển cho trường hợp xác định (số nguồn bằng số microphone).

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

# Generate two source signals
sr = 8000
duration = 1.0
t = jnp.linspace(0, duration, int(sr * duration))

# Source 1: sinusoidal (like a tone)
s1 = jnp.sin(2 * jnp.pi * 440 * t) + 0.3 * jnp.sin(2 * jnp.pi * 880 * t)

# Source 2: sawtooth-like (rich harmonics)
s2 = 2 * (t * 200 % 1) - 1  # sawtooth at 200 Hz

# Normalise sources
s1 = s1 / jnp.max(jnp.abs(s1))
s2 = s2 / jnp.max(jnp.abs(s2))
sources = jnp.stack([s1, s2])  # (2, T)

# Mixing matrix (unknown to the algorithm)
A = jnp.array([[0.8, 0.4],
               [0.3, 0.9]])
mixtures = A @ sources  # (2, T)

# FastICA implementation
def whiten(X):
    """Centre and whiten the data."""
    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):
    """FastICA using tanh non-linearity (approximation to negentropy)."""
    X_white, whitening = whiten(X)
    n, T = X_white.shape

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

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

        # For each component
        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

            # Newton update: 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

            # Decorrelate from previous components (deflation)
            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)

        # Check convergence
        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 matrix
    unmixing = W @ whitening
    recovered = unmixing @ X
    return recovered, unmixing

recovered, W_unmix = fastica(mixtures)

# Fix sign ambiguity (ICA can flip signs)
for i in range(2):
    if jnp.corrcoef(recovered[i], sources[i])[0, 1] < -0.5:
        recovered = recovered.at[i].set(-recovered[i])

# If sources are swapped, fix permutation
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]

# Normalise for display
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()

# Report correlation with originals
for i in range(2):
    corr = jnp.corrcoef(recovered[i], sources[i])[0, 1]
    print(f"Source {i+1} recovery correlation: {corr:.4f}")

• Task 2: Tách nguồn dựa trên NMF trên spectrogram. Sử dụng phân rã ma trận không âm (NMF, chương 02) để tách một spectrogram thành hai thành phần, minh họa cách NMF học từ điển phổ cho mỗi nguồn.

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

# Generate two signals with distinct spectral characteristics
sr = 8000
duration = 1.0
t = jnp.linspace(0, duration, int(sr * duration))

# Source 1: low-frequency harmonic (simulating bass)
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))

# Source 2: high-frequency harmonic (simulating a flute)
src2 = (jnp.sin(2 * jnp.pi * 800 * t) +
        0.4 * jnp.sin(2 * jnp.pi * 1600 * t))

# Time-varying amplitudes (sources active at different times)
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

# Compute magnitude spectrogram (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) - frequency x time
phase = jnp.angle(S_mix).T

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

# NMF: V ≈ WH using multiplicative update rules
def nmf(V, K, n_iter=200, key=jr.PRNGKey(0)):
    """Non-negative Matrix Factorisation with Frobenius norm."""
    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):
        # Multiplicative update for H
        WtV = W.T @ V
        WtWH = W.T @ W @ H + 1e-8
        H = H * (WtV / WtWH)

        # Multiplicative update for 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

# Run NMF with K=2 components
K = 2
W, H, costs = nmf(V, K, n_iter=300)

# Reconstruct each source using soft masks
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

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

# Mixture spectrogram
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 convergence
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')

# Spectral basis vectors 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()

# Temporal activations 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()

# Separated spectrograms
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.")

• Task 3: Bộ lọc thích ứng LMS cho khử nhiễu. Cài đặt các thuật toán LMS và NLMS để khử tiếng vọng/nhiễu, thể hiện hành vi hội tụ và ảnh hưởng của kích thước bước.

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

# Simulate an echo cancellation scenario
# Far-end signal -> room impulse response -> echo at microphone
# Near-end speech is the desired signal we want to preserve

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

# Far-end signal (reference): random speech-like signal
far_end = jr.normal(keys[0], (n_samples,)) * 0.5

# Room impulse response (unknown to the algorithm)
rir_length = 64
rir = jnp.zeros(rir_length)
rir = rir.at[0].set(0.8)   # direct path
rir = rir.at[5].set(0.3)   # early reflection
rir = rir.at[12].set(-0.2) # reflection
rir = rir.at[25].set(0.1)  # late reflection
rir = rir.at[40].set(-0.05)

# Echo: convolution of far-end with RIR
echo = jnp.convolve(far_end, rir)[:n_samples]

# Near-end speech (active in a portion of the signal)
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)

# Microphone signal: echo + near-end + noise
noise = jr.normal(keys[1], (n_samples,)) * 0.01
mic_signal = echo + near_end + noise

# LMS adaptive filter
def lms_filter(reference, desired, filter_length, mu):
    """Standard LMS adaptive filter."""
    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]  # reversed segment
        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 adaptive filter
def nlms_filter(reference, desired, filter_length, mu, eps=1e-6):
    """Normalised LMS adaptive filter."""
    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

# Run LMS with different step sizes
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))

# Original signals
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 convergence for different step sizes
for mu, color in zip(mu_values, colors_mu):
    _, err, _ = lms_filter(far_end, mic_signal, filter_len, mu)
    # Smoothed squared error
    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()

# Best LMS result
_, 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 result
_, 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()

# Measure echo reduction
echo_power = jnp.mean(echo ** 2)
lms_residual = jnp.mean(err_lms[n_samples//2:] ** 2)  # after convergence
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)")

• Task 4: Đánh mặt nạ thời gian-tần số để tăng cường giọng nói. Cài đặt một cách tiếp cận đánh mặt nạ phổ đơn giản (ideal ratio mask) và so sánh với trừ phổ, trực quan hóa chất lượng tách trên một tín hiệu giọng nói nhiễu tổng hợp.

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

# Create synthetic "speech" and "noise" signals
sr = 8000
duration = 2.0
t = jnp.linspace(0, duration, int(sr * duration))

# Speech: harmonic series with time-varying amplitude (simulating speech)
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 modulation
    speech = speech + (0.5 / (f0/150)) * amp_env * jnp.sin(2 * jnp.pi * f0 * t)
speech = speech / jnp.max(jnp.abs(speech))

# Noise: band-limited noise
key = jr.PRNGKey(42)
noise_raw = jr.normal(key, t.shape) * 0.4

# Mix at a given SNR
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)

# Method 1: Ideal Ratio Mask (oracle - upper bound)
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))

# Method 2: Spectral subtraction
# Estimate noise from first 0.2s (assumed silence)
noise_frames = int(0.2 * sr / hop)
noise_est = jnp.mean(mag_mix[:noise_frames] ** 2, axis=0, keepdims=True)
alpha = 2.0  # over-subtraction factor
beta = 0.02  # spectral floor
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))

# Method 3: Wiener filter
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))

# Compute SI-SDR for each method
def si_sdr(estimate, reference):
    """Scale-invariant signal-to-distortion ratio."""
    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)

# Visualisation
fig, axes = plt.subplots(3, 2, figsize=(14, 12))

# Spectrograms
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)')

# Masks
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')

# Enhanced waveforms comparison
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 bar chart
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)")