triforce/
lib.rs

1// SPDX-License-Identifier: GPL-2.0-or-later
2/*
3 * An attempt at an MVDR beamformer for equilateral triangle mic arrays
4 * as found on newer Apple Silicon Macs.
5 *
6 * Currently mono, but could probably be extended to be stereo somewhat
7 * easily. I think.
8 *
9 * Copyright (C) 2024 James Calligeros <jcalligeros99@gmail.com>
10 */
11
12use std::f32::consts::PI;
13
14use lv2::prelude::*;
15
16use nalgebra::{linalg::SVD, Matrix3, Vector3};
17
18use rustfft::{num_complex::Complex, num_traits::Zero, FftPlanner};
19
20use std::sync::Mutex;
21
22const C: f32 = 343.00; /* m*s^-1 */
23
24/// The distance of a given element in the array from the zeroth
25/// element
26#[derive(Copy, Clone)]
27struct ElemDistance {
28    x: f32,
29    y: f32,
30}
31
32/// Perform a Hilbert transform on a slice of f32s to give us the analytic signal of
33/// our input sample buffer. This is necessary to extract phase information from the
34/// signal, and to make matrix operations a bit easier.
35fn analytic_signal(planner: &mut FftPlanner<f32>, signal: &[f32], len: usize) -> Vec<Complex<f32>> {
36    // Convert each real sample into a complex sample
37    let mut complex_signal: Vec<Complex<f32>> =
38        signal.iter().map(|&x| Complex::new(x, 0.0)).collect();
39
40    // Set up the fft and inverse fft
41    let fft = planner.plan_fft_forward(len);
42    let ifft = planner.plan_fft_inverse(len);
43
44    // Mutate the output buffer into the forward FFT
45    fft.process(&mut complex_signal);
46
47    // Perform the Hilbert transform on the FFT. To do this, we multiply every
48    // positive sample under the Nyquist limit by 2+0j, and destroy every sample
49    // above it.
50    for i in 0..len {
51        if i > 0 && i < len / 2 {
52            complex_signal[i] *= Complex::new(2.0, 0.0);
53        } else if i >= len / 2 {
54            complex_signal[i] = Complex::zero();
55        }
56    }
57
58    // Turn the original complex buffer into the inverse FFT and then normalise
59    ifft.process(&mut complex_signal);
60    complex_signal.iter_mut().for_each(|x| *x /= len as f32);
61
62    complex_signal
63}
64
65/// The steering vector is a representation of the phase delays at each microphone.
66/// It is calculated by taking the dot product of the array geometry matrix and the
67/// unit vector of the direction of arrival.
68fn steering_vec(theta: f32, phi: f32, f: f32, elems: [ElemDistance; 3]) -> Vector3<Complex<f32>> {
69    // Mic positions are relative to Left/Top to preserve x/y axis semantics
70    let mic_positions: Vec<Vector3<f32>> =
71        elems.iter().map(|e| Vector3::new(e.x, e.y, 0f32)).collect();
72
73    // Calculate angular repetency (2pi/lambda)
74    let repetency = (2f32 * PI) / (f / C);
75
76    // Compute the unit vector of the DOA
77    let u_dir = Vector3::new(phi.sin() * theta.cos(), phi.sin() * theta.sin(), phi.cos());
78
79    // Calculate the steering vector by taking the array geometry, speed of sound,
80    // and DOA unit vector
81    let mut steering_vector = Vector3::from_element(Complex::new(0f32, 0f32));
82
83    for (i, mic_pos) in mic_positions.iter().enumerate() {
84        let delay = mic_pos.dot(&u_dir) / C;
85        let phase = -repetency * delay;
86        steering_vector[i] = Complex::new(phase.cos(), phase.sin());
87    }
88
89    steering_vector
90}
91
92/// There's nothing special about this, it's just a covariance matrix. It is always
93/// square.
94fn covariance(signals: &Vec<Vec<Complex<f32>>>, n_samples: usize) -> Matrix3<Complex<f32>> {
95    let mut covar = Matrix3::zeros();
96
97    for t in 0..n_samples {
98        let discrete: Vector3<Complex<f32>> = Vector3::from_iterator(signals.iter().map(|s| s[t]));
99        covar += &discrete * discrete.adjoint();
100    }
101
102    // Our samples are shit, so we can't get a very nice covariance matrix.
103    // Regularise the shit covariance matrix by introducing a constant value
104    // across the identity
105    let reg = Matrix3::identity().map(|x: f32| Complex::new(x * 1e-4f32, 0f32));
106
107    covar /= Complex::new(n_samples as f32, 0f32);
108    covar + reg
109}
110
111/// To calculate the weighting vector for the beamformer, we need to invert
112/// the covariance matrix, multiply it by the steering vector, then divide that
113/// by itself multiplied by the Hermitian transpose of the steering vector
114/// w = (cov^-1 * sv) / (sv.adjoint() * (cov^-1 * sv)). Note that the denominator
115/// is the same as the conjugate-linear dot product of the steering vector and
116/// the numerator.
117fn mvdr_weights(cov: &Matrix3<Complex<f32>>, sv: &Vector3<Complex<f32>>) -> Vector3<Complex<f32>> {
118    // Since we have a numerically unstable covariance matrix, we can't take the
119    // true inverse of it. Let's instead decompose it and take the pseudoinverse.
120    let svd = SVD::new(cov.to_owned(), true, true);
121    let r_inv = svd.pseudo_inverse(1e-4f32).unwrap();
122
123    let num = r_inv * sv;
124    let den = sv.dotc(&num); // Conjugate-linear dot product
125
126    num / den
127}
128
129/*
130 * Input and output ports used by the plugin
131 *
132 *
133 * Ports:
134 *      in_1: channel 1 input (left/top)
135 *      in_2: channel 2 input (right/bottom)
136 *      in_3: channel 3 input (vertex)
137 *      out: output
138 *      h_angle: horizontal steering angle in degrees (relative to input 1)
139 *      v_angle: vertical steering angle in degrees
140 *      opt_freq: frequency to optimise for
141 *      t_win: covariance matrix time window
142 */
143#[derive(PortCollection)]
144pub struct Ports {
145    in_1: InputPort<Audio>,
146    in_2: InputPort<Audio>,
147    in_3: InputPort<Audio>,
148    out: OutputPort<Audio>,
149    h_angle: InputPort<Control>,
150    v_angle: InputPort<Control>,
151    opt_freq: InputPort<Control>,
152    t_win: InputPort<Control>,
153    mic2_x: InputPort<Control>,
154    mic2_y: InputPort<Control>,
155    mic3_x: InputPort<Control>,
156    mic3_y: InputPort<Control>,
157}
158
159/*
160 * Plugin state
161 */
162#[uri("https://chadmed.au/triforce")]
163pub struct Triforce {
164    hangle_curr: f32,
165    vangle_curr: f32,
166    freq_curr: f32,
167    sample_rate: f32,
168    samples_since_last_update: usize,
169    covar_window: Vec<Vec<Complex<f32>>>,
170    steering_vector: Vector3<Complex<f32>>,
171    covar: Matrix3<Complex<f32>>,
172    array_geom: [ElemDistance; 3],
173    fft_planner: Mutex<FftPlanner<f32>>,
174    weights: Vector3<Complex<f32>>,
175}
176
177trait Beamformer: Plugin {
178    fn update_params(&mut self, ports: &mut Ports);
179}
180
181impl Triforce {
182    pub fn with_sample_rate(sample_rate: f32) -> Self {
183        Self {
184            hangle_curr: 0f32,
185            vangle_curr: 0f32,
186            freq_curr: 1000f32,
187            samples_since_last_update: usize::max_value(),
188            sample_rate,
189            covar_window: vec![
190                vec![Complex::new(0f32, 0f32); 256],
191                vec![Complex::new(0f32, 0f32); 256],
192                vec![Complex::new(0f32, 0f32); 256],
193            ],
194            array_geom: [ElemDistance { x: 0f32, y: 0f32 }; 3],
195            steering_vector: steering_vec(
196                90f32.to_radians(),
197                45f32.to_radians(),
198                1000f32,
199                [ElemDistance { x: 0f32, y: 0f32 }; 3],
200            ),
201            covar: Matrix3::zeros(),
202            fft_planner: Mutex::new(FftPlanner::new()),
203            weights: Vector3::zeros(),
204        }
205    }
206
207    pub fn process_slice(
208        &mut self,
209        mic1: &[f32],
210        mic2: &[f32],
211        mic3: &[f32],
212        output: &mut [f32],
213        t_win: f32,
214        buf_len: usize
215    ) {
216        // Steering vector is relative to Left/Top mic
217        let inputs = {
218            let mut planner = self.fft_planner.lock().unwrap();
219            vec![
220                analytic_signal(&mut planner, mic1, buf_len),
221                analytic_signal(&mut planner, mic2, buf_len),
222                analytic_signal(&mut planner, mic3, buf_len),
223            ]
224        };
225
226        // Update the covariance matrix. We use an overlapping window to smooth over
227        // the transitions.
228        if self.samples_since_last_update as f32 >= (t_win / 1000f32) * self.sample_rate {
229            self.samples_since_last_update = 0;
230            // We want a 1/3 overlap
231            let i = buf_len / 3;
232            self.covar_window[0].extend_from_slice(&inputs[0][0..i]);
233            self.covar_window[1].extend_from_slice(&inputs[1][0..i]);
234            self.covar_window[2].extend_from_slice(&inputs[2][0..i]);
235            self.covar = covariance(&self.covar_window, self.covar_window[0].len());
236            self.covar_window[0] = inputs[0][i + 1..buf_len].to_vec();
237            self.covar_window[1] = inputs[1][i + 1..buf_len].to_vec();
238            self.covar_window[2] = inputs[2][i + 1..buf_len].to_vec();
239            self.weights = mvdr_weights(&self.covar, &self.steering_vector);
240        } else {
241            self.samples_since_last_update += buf_len;
242        }
243
244        for t in 0..buf_len {
245            let discrete: Vector3<Complex<f32>> =
246                Vector3::from_iterator(inputs.iter().map(|s| s[t]));
247
248            let out =
249                // Conjugate-linear dot product
250                self.weights.dotc(&discrete)
251                // // Now we need to revert the Hilbert transform and output the signal
252                .re;
253
254            // Do all of our NFP and clamping here
255            output[t] = if out.is_finite() && !out.is_nan() {
256                out.clamp(-10f32, 10f32)
257            } else {
258                0f32
259            };
260        }
261    }
262}
263
264impl Plugin for Triforce {
265    type Ports = Ports;
266
267    type InitFeatures = ();
268    type AudioFeatures = ();
269
270    fn new(info: &PluginInfo, _features: &mut ()) -> Option<Self> {
271        Some(Self::with_sample_rate(info.sample_rate() as f32))
272    }
273
274    fn run(&mut self, ports: &mut Ports, _features: &mut (), samples: u32) {
275        Beamformer::update_params(self, ports);
276        self.process_slice(
277            &ports.in_1,
278            &ports.in_2,
279            &ports.in_3,
280            &mut ports.out,
281            *ports.t_win,
282            samples as usize
283        );
284    }
285}
286
287impl Beamformer for Triforce {
288    fn update_params(&mut self, ports: &mut Ports) {
289        if self.hangle_curr != *ports.h_angle
290            || self.freq_curr != *ports.opt_freq
291            || self.vangle_curr != *ports.v_angle
292            || self.array_geom[1].x != *ports.mic2_x
293            || self.array_geom[1].y != *ports.mic2_y
294            || self.array_geom[2].x != *ports.mic3_x
295            || self.array_geom[2].y != *ports.mic3_y
296        {
297            self.hangle_curr = *ports.h_angle;
298            self.vangle_curr = *ports.v_angle;
299            self.freq_curr = *ports.opt_freq;
300            self.array_geom = [
301                ElemDistance { x: 0f32, y: 0f32 },
302                ElemDistance {
303                    x: *ports.mic2_x,
304                    y: *ports.mic2_y,
305                },
306                ElemDistance {
307                    x: *ports.mic3_x,
308                    y: *ports.mic3_y,
309                },
310            ];
311            self.steering_vector = steering_vec(
312                self.hangle_curr.to_radians(),
313                self.vangle_curr.to_radians(),
314                self.freq_curr,
315                self.array_geom,
316            );
317
318            // The steering vector has changed
319            self.weights = mvdr_weights(&self.covar, &self.steering_vector);
320        }
321    }
322}
323
324lv2_descriptors!(Triforce);
triforce/lib.rs

triforce/
lib.rs
