modewise_vlmd/
decomposer.rs

1use modewise_decomposition::{Decomposition, Loadings, ModeDecomposer};
2use modewise_fourier::{from_half_spectrum, to_half_spectrum};
3use modewise_method::{MethodDescriptor, MethodError};
4use modewise_signals::MultichannelSeries;
5use ndarray::{Array1, Array2, Array3, Axis, Zip, s};
6use rustfft::num_complex::Complex64;
7use tracing::{debug, warn};
8
9use crate::lasso::lasso;
10use crate::params::VlmdParams;
11
12/// Below this total spectral energy a mode is treated as empty, as in the reference.
13const EMPTY_MODE: f64 = 1e-10;
14
15#[derive(Debug, Clone)]
16pub struct Vlmd {
17    params: VlmdParams,
18}
19
20/// Everything VLMD estimates. `modes` are the channel modes `U⁽ᵏ⁾ = Θ⁽ᵏ⁾A` (Eq. 7);
21/// they sum to the latent fit `ZA`, not exactly to the signal.
22#[derive(Debug, Clone, PartialEq)]
23pub struct VlmdOutput {
24    pub modes: Decomposition,
25    pub latent_modes: Decomposition,
26    pub loadings: Loadings,
27    pub iterations: usize,
28}
29
30struct State {
31    loadings: Array2<f64>,
32    latents: Array2<f64>,
33    latents_hat: Array2<Complex64>,
34    modes_hat: Array3<Complex64>,
35    modes_sum: Array2<Complex64>,
36    dual: Array2<Complex64>,
37}
38
39impl Vlmd {
40    pub fn new(params: VlmdParams) -> Result<Self, MethodError> {
41        params.validate()?;
42        Ok(Self { params })
43    }
44
45    pub fn params(&self) -> &VlmdParams {
46        &self.params
47    }
48
49    /// Channel modes, latent modes and loadings, with modes ordered by centre frequency.
50    pub fn decompose_latent(&self, signal: &MultichannelSeries) -> Result<VlmdOutput, MethodError> {
51        let p = &self.params;
52        let n_channels = signal.n_channels();
53        let n_latents = p.latents_for(n_channels)?;
54        let x = signal.data();
55        let x_hat = to_half_spectrum(x);
56        let n_bins = x_hat.ncols();
57        let freqs = Array1::linspace(0.0, 0.5, n_bins);
58        let penalty = p.reg_lambda * 1e-3;
59        let fit = 2.0 / p.reg_rho;
60        let narrowness = 4.0 * p.alpha / p.reg_rho;
61
62        let latents = x.slice(s![..n_latents, ..]).to_owned();
63        let mut state = State {
64            loadings: Array2::eye(n_latents.max(n_channels))
65                .slice_move(s![..n_latents, ..n_channels]),
66            latents_hat: to_half_spectrum(latents.view()),
67            latents,
68            modes_hat: Array3::zeros((p.num_modes, n_latents, n_bins)),
69            modes_sum: Array2::zeros((n_latents, n_bins)),
70            dual: Array2::zeros((n_latents, n_bins)),
71        };
72        let mut omegas = vec![0.0; p.num_modes];
73        let mut used = omegas.clone();
74        let mut change = p.tolerance + f64::EPSILON;
75        let mut iterations = 0;
76        while iterations < p.max_iterations && change > p.tolerance {
77            lasso(state.latents.view(), x, penalty, &mut state.loadings);
78            bound_loadings(&mut state.loadings);
79            update_latents(&mut state, &x_hat, fit);
80            let next = update_modes(&mut state, &freqs, &omegas, narrowness);
81            Zip::from(&mut state.dual)
82                .and(&state.latents_hat)
83                .and(&state.modes_sum)
84                .for_each(|g, z, m| *g += (z - m) * p.tau);
85            state.latents = from_half_spectrum(state.latents_hat.view());
86            change = omegas
87                .iter()
88                .zip(&next)
89                .map(|(old, new)| ((new - old) / (old + 1e-10)).powi(2))
90                .sum();
91            used = std::mem::replace(&mut omegas, next);
92            iterations += 1;
93        }
94        debug!(
95            iterations,
96            change,
97            converged = change <= p.tolerance,
98            "vlmd finished"
99        );
100        self.output(signal, state, &used, iterations)
101    }
102
103    fn output(
104        &self,
105        signal: &MultichannelSeries,
106        state: State,
107        omegas: &[f64],
108        iterations: usize,
109    ) -> Result<VlmdOutput, MethodError> {
110        let fs = signal.sampling_rate();
111        let (n_latents, n_channels) = state.loadings.dim();
112        let n_samples = signal.n_samples();
113        let mut order: Vec<usize> = (0..omegas.len()).collect();
114        order.sort_by(|&a, &b| omegas[a].total_cmp(&omegas[b]));
115
116        let mut latent_modes = Array3::zeros((order.len(), n_latents, n_samples));
117        let mut modes = Array3::zeros((order.len(), n_channels, n_samples));
118        for (i, &k) in order.iter().enumerate() {
119            let theta = from_half_spectrum(state.modes_hat.index_axis(Axis(0), k));
120            modes
121                .index_axis_mut(Axis(0), i)
122                .assign(&state.loadings.t().dot(&theta));
123            latent_modes.index_axis_mut(Axis(0), i).assign(&theta);
124        }
125        let centres: Vec<f64> = order.iter().map(|&k| omegas[k] * fs.hz()).collect();
126        let latent_labels: Vec<String> = (0..n_latents).map(|l| format!("latent_{l}")).collect();
127        Ok(VlmdOutput {
128            modes: Decomposition::new(modes, centres.clone(), fs, signal.labels().to_vec(), true)?,
129            latent_modes: Decomposition::new(
130                latent_modes,
131                centres,
132                fs,
133                latent_labels.clone(),
134                true,
135            )?,
136            loadings: Loadings::new(state.loadings, latent_labels, signal.labels().to_vec())?,
137            iterations,
138        })
139    }
140}
141
142/// Keeps |a| ≤ 1 by rescaling each latent's row, which the latent signal can absorb.
143fn bound_loadings(loadings: &mut Array2<f64>) {
144    for mut row in loadings.rows_mut() {
145        let largest = row.iter().fold(0.0f64, |m, v| m.max(v.abs()));
146        if largest > 1.0 {
147            row /= largest;
148        }
149    }
150}
151
152/// Eq. (16), one latent at a time against the residual left by the others.
153fn update_latents(state: &mut State, x_hat: &Array2<Complex64>, fit: f64) {
154    let mut residual = x_hat
155        - &state
156            .loadings
157            .t()
158            .mapv(|v| Complex64::new(v, 0.0))
159            .dot(&state.latents_hat);
160    for l in 0..state.loadings.nrows() {
161        let a = state.loadings.row(l);
162        let mut z = state.latents_hat.row_mut(l);
163        for (c, &acl) in a.iter().enumerate() {
164            if acl != 0.0 {
165                Zip::from(residual.row_mut(c))
166                    .and(&z)
167                    .for_each(|r, &zl| *r += zl * acl);
168            }
169        }
170        let mut projection = Array1::<Complex64>::zeros(z.len());
171        for (c, &acl) in a.iter().enumerate() {
172            if acl != 0.0 {
173                projection.scaled_add(Complex64::new(acl, 0.0), &residual.row(c));
174            }
175        }
176        let denominator = 1.0 + fit * a.dot(&a);
177        Zip::from(&mut z)
178            .and(&projection)
179            .and(state.modes_sum.row(l))
180            .and(state.dual.row(l))
181            .for_each(|z, &proj, &modes, &dual| *z = (proj * fit + modes - dual) / denominator);
182        for (c, &acl) in a.iter().enumerate() {
183            if acl != 0.0 {
184                Zip::from(residual.row_mut(c))
185                    .and(&z)
186                    .for_each(|r, &zl| *r -= zl * acl);
187            }
188        }
189    }
190}
191
192/// Eqs. (18) and (21): each mode against the others, then its centre frequency.
193/// Returns the next frequencies; the filters use the current ones.
194fn update_modes(
195    state: &mut State,
196    freqs: &Array1<f64>,
197    omegas: &[f64],
198    narrowness: f64,
199) -> Vec<f64> {
200    let mut next = omegas.to_vec();
201    for (k, mut mode) in state.modes_hat.axis_iter_mut(Axis(0)).enumerate() {
202        let (mut power, mut weighted) = (0.0, 0.0);
203        Zip::indexed(&mut mode)
204            .and(&mut state.modes_sum)
205            .and(&state.latents_hat)
206            .and(&state.dual)
207            .for_each(|(_, j), m, sum, &z, &g| {
208                let others = *sum - *m;
209                let updated =
210                    (z - others + g) / (1.0 + narrowness * (freqs[j] - omegas[k]).powi(2));
211                *m = updated;
212                *sum = others + updated;
213                power += updated.norm_sqr();
214                weighted += updated.norm_sqr() * freqs[j];
215            });
216        next[k] = if power < EMPTY_MODE {
217            warn!(
218                mode = k,
219                "vlmd mode has no energy; its frequency is set to 0"
220            );
221            0.0
222        } else {
223            weighted / power
224        };
225    }
226    next
227}
228
229impl ModeDecomposer for Vlmd {
230    fn descriptor(&self) -> MethodDescriptor {
231        MethodDescriptor {
232            name: "vlmd",
233            version: env!("CARGO_PKG_VERSION"),
234            multivariate: true,
235        }
236    }
237
238    fn decompose(&self, signal: &MultichannelSeries) -> Result<Decomposition, MethodError> {
239        Ok(self.decompose_latent(signal)?.modes)
240    }
241}
242
243#[cfg(test)]
244mod tests {
245    use super::*;
246
247    #[test]
248    fn an_empty_mode_gets_frequency_zero_not_nan() {
249        let mut state = State {
250            loadings: Array2::eye(1),
251            latents: Array2::zeros((1, 4)),
252            latents_hat: Array2::zeros((1, 5)),
253            modes_hat: Array3::zeros((2, 1, 5)),
254            modes_sum: Array2::zeros((1, 5)),
255            dual: Array2::zeros((1, 5)),
256        };
257        let freqs = Array1::linspace(0.0, 0.5, 5);
258        assert_eq!(
259            update_modes(&mut state, &freqs, &[0.1, 0.3], 10.0),
260            [0.0, 0.0]
261        );
262        state.latents_hat[[0, 2]] = Complex64::new(1e-6, 0.0);
263        assert_eq!(
264            update_modes(&mut state, &freqs, &[0.1, 0.3], 10.0),
265            [0.0, 0.0],
266            "energy below 1e-10"
267        );
268    }
269
270    #[test]
271    fn rows_above_one_are_rescaled_and_others_left_alone() {
272        let mut a = ndarray::array![[2.0, -4.0], [0.5, -1.0]];
273        bound_loadings(&mut a);
274        assert_eq!(a, ndarray::array![[0.5, -1.0], [0.5, -1.0]]);
275    }
276}
modewise_vlmd/decomposer.rs

modewise_vlmd/
decomposer.rs
