audio_lab/dsp/
convolve.rs

1use ndarray::{Array2, s};
2use ndrustfft::{R2cFftHandler, ndfft_r2c, ndifft_r2c};
3use num::complex::Complex64;
4use thiserror::Error;
5
6use crate::signal::{SignalError, TimeSignal};
7
8#[derive(Debug, Error)]
9pub enum ConvolveError {
10    #[error("sample rates do not match: {left} Hz vs {right} Hz")]
11    SampleRateMismatch { left: f64, right: f64 },
12    #[error("channel counts do not match: {left} vs {right}")]
13    ChannelMismatch { left: usize, right: usize },
14    #[error("signals must not be empty")]
15    EmptySignal,
16    #[error(transparent)]
17    Signal(#[from] SignalError),
18}
19
20#[derive(Clone, Copy, Debug, PartialEq, Eq)]
21pub enum ConvolveMode {
22    Full,
23    Cut,
24    Cyclic,
25}
26
27pub fn convolve(
28    signal1: &TimeSignal,
29    signal2: &TimeSignal,
30    mode: ConvolveMode,
31) -> Result<TimeSignal, ConvolveError> {
32    if signal1.sample_rate() != signal2.sample_rate() {
33        return Err(ConvolveError::SampleRateMismatch {
34            left: signal1.sample_rate(),
35            right: signal2.sample_rate(),
36        });
37    }
38
39    if signal1.num_channels() != signal2.num_channels() {
40        return Err(ConvolveError::ChannelMismatch {
41            left: signal1.num_channels(),
42            right: signal2.num_channels(),
43        });
44    }
45
46    let n1 = signal1.num_time_steps();
47    let n2 = signal2.num_time_steps();
48    if n1 == 0 || n2 == 0 {
49        return Err(ConvolveError::EmptySignal);
50    }
51
52    let conv_len = n1 + n2 - 1;
53    // Power-of-two FFT lengths are much faster than arbitrary (possibly prime)
54    // lengths. The extra zero padding does not change the linear convolution.
55    let fft_len = conv_len.next_power_of_two();
56    let num_freq_bins = fft_len / 2 + 1;
57    let num_channels = signal1.num_channels();
58    let fft_handler = R2cFftHandler::<f64>::new(fft_len);
59
60    let mut buffer = Array2::<f64>::zeros((num_channels, fft_len));
61    let mut spectrum1 = Array2::<Complex64>::zeros((num_channels, num_freq_bins));
62    let mut spectrum2 = Array2::<Complex64>::zeros((num_channels, num_freq_bins));
63
64    buffer.slice_mut(s![.., ..n1]).assign(&signal1.time_data());
65    ndfft_r2c(&buffer.view(), &mut spectrum1.view_mut(), &fft_handler, 1);
66
67    buffer.slice_mut(s![.., ..n1]).fill(0.0);
68    buffer.slice_mut(s![.., ..n2]).assign(&signal2.time_data());
69    ndfft_r2c(&buffer.view(), &mut spectrum2.view_mut(), &fft_handler, 1);
70
71    spectrum1 *= &spectrum2;
72    ndifft_r2c(&spectrum1.view(), &mut buffer.view_mut(), &fft_handler, 1);
73
74    let out_len = match mode {
75        ConvolveMode::Full => conv_len,
76        ConvolveMode::Cut | ConvolveMode::Cyclic => n1.max(n2),
77    };
78
79    if mode == ConvolveMode::Cyclic {
80        // Wrap the tail (out_len..conv_len) back onto the start.
81        let wrap_len = conv_len - out_len;
82        for mut channel in buffer.outer_iter_mut() {
83            for i in 0..wrap_len {
84                channel[i] += channel[out_len + i];
85            }
86        }
87    }
88
89    Ok(TimeSignal::new(
90        buffer.slice(s![.., ..out_len]).to_owned(),
91        signal1.sample_rate(),
92    )?)
93}
94
95#[cfg(test)]
96mod tests {
97    use crate::signal::{ImpulseConfig, generate_impulse};
98    use approx::assert_abs_diff_eq;
99    use ndarray::{Array1, arr1, arr2};
100
101    use super::*;
102
103    #[test]
104    fn convolve_full_matches_known_result() {
105        let s1 = TimeSignal::new(arr2(&[[1.0, 2.0, 3.0]]), 1.0).unwrap();
106        let s2 = TimeSignal::new(arr2(&[[1.0, 1.0]]), 1.0).unwrap();
107
108        let result = convolve(&s1, &s2, ConvolveMode::Full).unwrap();
109
110        assert_eq!(result.num_time_steps(), 4);
111        assert_abs_diff_eq!(
112            result.channel(0),
113            arr1(&[1.0, 3.0, 5.0, 3.0]),
114            epsilon = 1e-10
115        );
116    }
117
118    #[test]
119    fn convolve_cut_truncates_to_longer_input() {
120        let s1 = TimeSignal::new(arr2(&[[1.0, 2.0, 3.0]]), 1.0).unwrap();
121        let s2 = TimeSignal::new(arr2(&[[1.0, 1.0]]), 1.0).unwrap();
122
123        let result = convolve(&s1, &s2, ConvolveMode::Cut).unwrap();
124
125        assert_eq!(result.num_time_steps(), 3);
126        assert_abs_diff_eq!(result.channel(0), arr1(&[1.0, 3.0, 5.0]), epsilon = 1e-10);
127    }
128
129    #[test]
130    fn convolve_cyclic_wraps_tail_onto_start() {
131        let s1 = TimeSignal::new(arr2(&[[1.0, 2.0, 3.0]]), 1.0).unwrap();
132        let s2 = TimeSignal::new(arr2(&[[1.0, 1.0]]), 1.0).unwrap();
133
134        let result = convolve(&s1, &s2, ConvolveMode::Cyclic).unwrap();
135
136        assert_eq!(result.num_time_steps(), 3);
137        assert_abs_diff_eq!(result.channel(0), arr1(&[4.0, 3.0, 5.0]), epsilon = 1e-10);
138    }
139
140    #[test]
141    fn convolve_matches_direct_convolution_for_non_power_of_two_lengths() {
142        let n1 = 37;
143        let n2 = 13;
144        let x: Vec<f64> = (0..n1).map(|i| ((i * 7 % 11) as f64) - 5.0).collect();
145        let h: Vec<f64> = (0..n2).map(|i| ((i * 3 % 5) as f64) * 0.5 - 1.0).collect();
146
147        let mut full = vec![0.0; n1 + n2 - 1];
148        for (i, &xi) in x.iter().enumerate() {
149            for (j, &hj) in h.iter().enumerate() {
150                full[i + j] += xi * hj;
151            }
152        }
153        let mut cyclic = full[..n1].to_vec();
154        for (i, &tail) in full[n1..].iter().enumerate() {
155            cyclic[i] += tail;
156        }
157
158        let s1 = TimeSignal::new(Array2::from_shape_vec((1, n1), x).unwrap(), 1.0).unwrap();
159        let s2 = TimeSignal::new(Array2::from_shape_vec((1, n2), h).unwrap(), 1.0).unwrap();
160
161        let result = convolve(&s1, &s2, ConvolveMode::Full).unwrap();
162        assert_abs_diff_eq!(
163            result.channel(0),
164            Array1::from(full.clone()),
165            epsilon = 1e-9
166        );
167
168        let result = convolve(&s1, &s2, ConvolveMode::Cut).unwrap();
169        assert_abs_diff_eq!(
170            result.channel(0),
171            Array1::from(full[..n1].to_vec()),
172            epsilon = 1e-9
173        );
174
175        let result = convolve(&s2, &s1, ConvolveMode::Cyclic).unwrap();
176        assert_abs_diff_eq!(result.channel(0), Array1::from(cyclic), epsilon = 1e-9);
177    }
178
179    #[test]
180    fn convolve_rejects_sample_rate_mismatch() {
181        let s1 = TimeSignal::new(arr2(&[[1.0, 2.0]]), 44_100.0).unwrap();
182        let s2 = TimeSignal::new(arr2(&[[1.0, 1.0]]), 48_000.0).unwrap();
183
184        let result = convolve(&s1, &s2, ConvolveMode::Full);
185
186        assert!(matches!(
187            result,
188            Err(ConvolveError::SampleRateMismatch {
189                left: 44_100.0,
190                right: 48_000.0
191            })
192        ));
193    }
194
195    #[test]
196    fn convolve_rejects_channel_mismatch() {
197        let mono = TimeSignal::new(arr2(&[[1.0, 2.0]]), 1.0).unwrap();
198        let stereo = TimeSignal::new(arr2(&[[1.0, 2.0], [3.0, 4.0]]), 1.0).unwrap();
199
200        let result = convolve(&mono, &stereo, ConvolveMode::Full);
201
202        assert!(matches!(
203            result,
204            Err(ConvolveError::ChannelMismatch { left: 1, right: 2 })
205        ));
206    }
207
208    #[test]
209    fn convolve_applies_same_channel_processing_for_multichannel_input() {
210        let s1 = TimeSignal::new(arr2(&[[1.0, 2.0], [3.0, 4.0]]), 1.0).unwrap();
211        let s2 = TimeSignal::new(arr2(&[[1.0, 0.0], [0.0, 1.0]]), 1.0).unwrap();
212
213        let result = convolve(&s1, &s2, ConvolveMode::Full).unwrap();
214
215        assert_eq!(result.num_channels(), 2);
216        assert_eq!(result.num_time_steps(), 3);
217        assert_abs_diff_eq!(result.channel(0), arr1(&[1.0, 2.0, 0.0]), epsilon = 1e-10);
218        assert_abs_diff_eq!(result.channel(1), arr1(&[0.0, 3.0, 4.0]), epsilon = 1e-10);
219    }
220
221    #[test]
222    fn convolve_rejects_empty_input() {
223        let empty = TimeSignal::zeros(1, 0, 1.0).unwrap();
224        let signal = TimeSignal::new(arr2(&[[1.0, 2.0]]), 1.0).unwrap();
225
226        let result = convolve(&empty, &signal, ConvolveMode::Full);
227
228        assert!(matches!(result, Err(ConvolveError::EmptySignal)));
229    }
230
231    #[test]
232    fn convolving_with_generated_delta_impulse_is_identity() {
233        let signal = TimeSignal::new(arr2(&[[0.25, -0.5, 1.0, 0.75, -0.125]]), 48_000.0).unwrap();
234        let impulse = generate_impulse(
235            1,
236            ImpulseConfig {
237                sample_rate: 48_000.0,
238                ..Default::default()
239            },
240        )
241        .unwrap();
242
243        let result = convolve(&signal, &impulse, ConvolveMode::Full).unwrap();
244
245        assert_eq!(result.num_time_steps(), signal.num_time_steps());
246        assert_abs_diff_eq!(result.channel(0), signal.channel(0), epsilon = 1e-10);
247    }
248}
audio_lab/dsp/convolve.rs

audio_lab/dsp/
convolve.rs
