rlx_aec/
fdaf.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Frequency-domain adaptive filter (FDAF) with overlap-save NLMS.
17
18use crate::dtd::DoubleTalkDetector;
19use crate::residual::ResidualWeights;
20use anyhow::{Result, ensure};
21use rlx_fft::reference::{fft_real_batch, ifft_complex_batch};
22
23#[derive(Debug, Clone)]
24pub struct FdafConfig {
25    pub n_fft: usize,
26    pub frame_samples: usize,
27    pub step_size: f32,
28    pub adapt: bool,
29    pub use_residual: bool,
30}
31
32impl Default for FdafConfig {
33    fn default() -> Self {
34        Self {
35            n_fft: 1024,
36            frame_samples: 160,
37            step_size: 0.05,
38            adapt: true,
39            use_residual: true,
40        }
41    }
42}
43
44pub struct FdafNlms {
45    cfg: FdafConfig,
46    w_re: Vec<f32>,
47    w_im: Vec<f32>,
48    psd: Vec<f32>,
49    far_block: Vec<f32>,
50    dtd: DoubleTalkDetector,
51    residual: Option<ResidualWeights>,
52    far_power_smooth: f32,
53    psd_smooth: f32,
54}
55
56impl FdafNlms {
57    pub fn new(cfg: FdafConfig, residual: Option<ResidualWeights>) -> Result<Self> {
58        let n = cfg.n_fft;
59        let hop = cfg.frame_samples;
60        ensure!(hop > 0 && hop < n, "frame_samples must be in (0, n_fft)");
61        let use_residual = cfg.use_residual;
62        Ok(Self {
63            cfg,
64            w_re: vec![0.0; n],
65            w_im: vec![0.0; n],
66            psd: vec![1.0; n],
67            far_block: vec![0.0; n],
68            dtd: DoubleTalkDetector::default(),
69            residual: residual.filter(|r| use_residual && r.n_fft == n),
70            far_power_smooth: 0.0,
71            psd_smooth: 0.98,
72        })
73    }
74
75    pub fn reset(&mut self) {
76        self.w_re.fill(0.0);
77        self.w_im.fill(0.0);
78        self.psd.fill(1.0);
79        self.far_block.fill(0.0);
80        self.far_power_smooth = 0.0;
81    }
82
83    pub fn config(&self) -> &FdafConfig {
84        &self.cfg
85    }
86
87    fn fft_block(&self, time: &[f32]) -> Result<Vec<f32>> {
88        ensure!(time.len() == self.cfg.n_fft);
89        fft_real_batch(time, 1, self.cfg.n_fft)
90    }
91
92    fn ifft_block(&self, spec: &[f32]) -> Result<Vec<f32>> {
93        ifft_complex_batch(spec, 1, self.cfg.n_fft)
94    }
95
96    /// Process one hop of mic and aligned far-end samples.
97    pub fn process_frame(&mut self, mic: &[f32], far: &[f32], out: &mut [f32]) -> Result<()> {
98        let hop = self.cfg.frame_samples;
99        let n = self.cfg.n_fft;
100        ensure!(mic.len() >= hop && far.len() >= hop && out.len() >= hop);
101
102        self.far_block.copy_within(hop..n, 0);
103        for i in 0..hop {
104            self.far_block[n - hop + i] = far[i];
105        }
106
107        let xf = self.fft_block(&self.far_block)?;
108
109        for k in 0..n {
110            let xr = xf[k * 2];
111            let xi = xf[k * 2 + 1];
112            let power = xr * xr + xi * xi;
113            self.psd[k] = self.psd_smooth * self.psd[k] + (1.0 - self.psd_smooth) * power;
114        }
115
116        let mut y_spec = vec![0.0f32; n * 2];
117        for k in 0..n {
118            let xr = xf[k * 2];
119            let xi = xf[k * 2 + 1];
120            let wr = self.w_re[k];
121            let wi = self.w_im[k];
122            y_spec[k * 2] = wr * xr - wi * xi;
123            y_spec[k * 2 + 1] = wr * xi + wi * xr;
124        }
125
126        let y_time = self.ifft_block(&y_spec)?;
127        let scale = 1.0 / n as f32;
128
129        let mut echo_est = vec![0.0f32; hop];
130        for i in 0..hop {
131            echo_est[i] = y_time[(n - hop + i) * 2] * scale;
132        }
133
134        let mut error = vec![0.0f32; hop];
135        for i in 0..hop {
136            error[i] = mic[i] - echo_est[i];
137        }
138
139        let mic_power = DoubleTalkDetector::frame_power(&mic[..hop]);
140        let echo_power = DoubleTalkDetector::frame_power(&echo_est);
141        self.far_power_smooth =
142            0.9 * self.far_power_smooth + 0.1 * (self.psd.iter().sum::<f32>() / n as f32);
143        let pause = self
144            .dtd
145            .pause_adaptation(mic_power, self.far_power_smooth, echo_power);
146
147        let mut e_time_buf = vec![0.0f32; n];
148        for i in 0..hop {
149            e_time_buf[n - hop + i] = error[i];
150        }
151        let mut ef = self.fft_block(&e_time_buf)?;
152
153        if let Some(res) = &self.residual {
154            res.apply_spectrum(&mut ef);
155            let e_time = self.ifft_block(&ef)?;
156            for i in 0..hop {
157                error[i] = e_time[(n - hop + i) * 2] * scale;
158            }
159        }
160
161        out[..hop].copy_from_slice(&error[..hop]);
162
163        if self.cfg.adapt && !pause {
164            let mu = self.cfg.step_size;
165            for k in 0..n {
166                let xr = xf[k * 2];
167                let xi = xf[k * 2 + 1];
168                let er = ef[k * 2];
169                let ei = ef[k * 2 + 1];
170                let gr = (xr * er + xi * ei) / (self.psd[k] + 1e-10);
171                let gi = (xr * ei - xi * er) / (self.psd[k] + 1e-10);
172                self.w_re[k] += mu * gr;
173                self.w_im[k] += mu * gi;
174            }
175        }
176
177        Ok(())
178    }
179
180    pub fn process_buffer(&mut self, mic: &[f32], far: &[f32], out: &mut [f32]) -> Result<()> {
181        ensure!(mic.len() == far.len() && mic.len() == out.len());
182        let hop = self.cfg.frame_samples;
183        let mut pos = 0;
184        while pos < mic.len() {
185            let end = (pos + hop).min(mic.len());
186            let chunk = end - pos;
187            let mut mp = vec![0.0f32; hop];
188            let mut fp = vec![0.0f32; hop];
189            let mut op = vec![0.0f32; hop];
190            mp[..chunk].copy_from_slice(&mic[pos..end]);
191            fp[..chunk].copy_from_slice(&far[pos..end]);
192            self.process_frame(&mp, &fp, &mut op)?;
193            out[pos..end].copy_from_slice(&op[..chunk]);
194            pos += chunk;
195        }
196        Ok(())
197    }
198}
rlx_aec/fdaf.rs

rlx_aec/
fdaf.rs
