modewise_mvmd/
decomposer.rs

1use modewise_decomposition::{Decomposition, 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};
6use rustfft::num_complex::Complex64;
7use tracing::debug;
8
9use crate::params::MvmdParams;
10
11#[derive(Debug, Clone)]
12pub struct Mvmd {
13    params: MvmdParams,
14}
15
16struct Solution {
17    modes_hat: Array3<Complex64>,
18    omegas: Vec<f64>,
19}
20
21impl Mvmd {
22    pub fn new(params: MvmdParams) -> Result<Self, MethodError> {
23        params.validate()?;
24        Ok(Self { params })
25    }
26
27    pub fn params(&self) -> &MvmdParams {
28        &self.params
29    }
30
31    fn solve(&self, signal_hat: &Array2<Complex64>, n_samples: usize) -> Solution {
32        let p = &self.params;
33        let (n_channels, n_bins) = signal_hat.dim();
34        let freqs = Array1::linspace(0.0, 0.5, n_bins);
35        let mut modes_hat = Array3::<Complex64>::zeros((p.num_modes, n_channels, n_bins));
36        let mut modes_sum = Array2::<Complex64>::zeros((n_channels, n_bins));
37        let mut lambda = Array2::<Complex64>::zeros((n_channels, n_bins));
38        let mut omegas = p.init.omegas(p.num_modes);
39        // Report the frequencies the final modes were filtered with, as the reference does.
40        let mut used_omegas = omegas.clone();
41
42        let mut change = p.tolerance + f64::EPSILON;
43        let mut iteration = 0;
44        while iteration < p.max_iterations && change > p.tolerance {
45            change = 0.0;
46            let mut next_omegas = omegas.clone();
47            for (k, mut mode) in modes_hat.axis_iter_mut(Axis(0)).enumerate() {
48                let (mut power, mut weighted) = (0.0, 0.0);
49                Zip::indexed(&mut mode)
50                    .and(&mut modes_sum)
51                    .and(signal_hat)
52                    .and(&lambda)
53                    .for_each(|(_, j), m, sum, s, l| {
54                        let others = *sum - *m;
55                        let filter = 1.0 + p.alpha * (freqs[j] - omegas[k]).powi(2);
56                        let updated = (*s - others - *l * 0.5) / filter;
57                        change += (updated - *m).norm_sqr();
58                        *m = updated;
59                        *sum = others + updated;
60                        power += updated.norm_sqr();
61                        weighted += updated.norm_sqr() * freqs[j];
62                    });
63                if power > 0.0 {
64                    next_omegas[k] = weighted / power;
65                }
66            }
67            Zip::from(&mut lambda)
68                .and(&modes_sum)
69                .and(signal_hat)
70                .for_each(|l, sum, s| *l += (*sum - *s) * p.tau);
71            used_omegas = std::mem::replace(&mut omegas, next_omegas);
72            iteration += 1;
73            change /= n_samples as f64;
74        }
75        debug!(
76            iteration,
77            change,
78            converged = change <= p.tolerance,
79            "mvmd finished"
80        );
81        Solution {
82            modes_hat,
83            omegas: used_omegas,
84        }
85    }
86}
87
88impl ModeDecomposer for Mvmd {
89    fn descriptor(&self) -> MethodDescriptor {
90        MethodDescriptor {
91            name: "mvmd",
92            version: env!("CARGO_PKG_VERSION"),
93            multivariate: true,
94        }
95    }
96
97    fn decompose(&self, signal: &MultichannelSeries) -> Result<Decomposition, MethodError> {
98        let n_samples = signal.n_samples();
99        let signal_hat = to_half_spectrum(signal.data());
100        let Solution { modes_hat, omegas } = self.solve(&signal_hat, n_samples);
101
102        let mut order: Vec<usize> = (0..omegas.len()).collect();
103        order.sort_by(|&a, &b| omegas[a].total_cmp(&omegas[b]));
104
105        let mut modes = Array3::zeros((order.len(), signal.n_channels(), n_samples));
106        for (mut out, &k) in modes.axis_iter_mut(Axis(0)).zip(&order) {
107            out.assign(&from_half_spectrum(modes_hat.index_axis(Axis(0), k)));
108        }
109        let fs = signal.sampling_rate();
110        let centre_frequencies_hz = order.iter().map(|&k| omegas[k] * fs.hz()).collect();
111
112        Decomposition::new(
113            modes,
114            centre_frequencies_hz,
115            fs,
116            signal.labels().to_vec(),
117            true,
118        )
119    }
120}
modewise_mvmd/decomposer.rs

modewise_mvmd/
decomposer.rs
