modewise_mvmd/
noise_assisted.rs

1use modewise_decomposition::{Decomposition, ModeDecomposer};
2use modewise_method::{MethodDescriptor, MethodError};
3use modewise_signals::MultichannelSeries;
4use ndarray::{Array2, Axis, s};
5use rand::SeedableRng;
6use rand_chacha::ChaCha8Rng;
7use rand_distr::{Distribution, StandardNormal};
8
9use crate::{Mvmd, MvmdParams};
10
11/// White-noise channels added before MVMD. Defaults follow the reference `na_mvmd.m`.
12#[derive(Debug, Clone, PartialEq)]
13pub struct NoiseParams {
14    pub channels: usize,
15    /// Noise SD, relative to the unit-SD signal channels.
16    pub sd: f64,
17    pub seed: u64,
18}
19
20impl Default for NoiseParams {
21    fn default() -> Self {
22        Self {
23            channels: 1,
24            sd: 0.8,
25            seed: 0,
26        }
27    }
28}
29
30/// Noise-assisted MVMD (Zisou, Apostolidis and Hadjileontiadis, ICASSP 2021). Each channel is
31/// scaled to unit SD (ddof = 1), white-noise channels are appended, MVMD runs on the whole set,
32/// the noise channels are dropped and the modes scaled back. The paper's generalised
33/// cross-spectrum centroid for one snapshot is Σ_c Σ_f f|û_c|² / Σ_c Σ_f |û_c|², the update
34/// MVMD already uses.
35#[derive(Debug, Clone)]
36pub struct NaMvmd {
37    inner: Mvmd,
38    noise: NoiseParams,
39}
40
41impl NaMvmd {
42    pub fn new(params: MvmdParams, noise: NoiseParams) -> Result<Self, MethodError> {
43        if noise.channels == 0 {
44            return Err(MethodError::invalid("noise_channels", "must be at least 1"));
45        }
46        if !(noise.sd.is_finite() && noise.sd > 0.0) {
47            return Err(MethodError::invalid(
48                "noise_sd",
49                "must be finite and positive",
50            ));
51        }
52        Ok(Self {
53            inner: Mvmd::new(params)?,
54            noise,
55        })
56    }
57
58    /// The unit-SD signal with the noise channels appended, and each channel's original SD.
59    pub fn augment(
60        &self,
61        signal: &MultichannelSeries,
62    ) -> Result<(MultichannelSeries, Vec<f64>), MethodError> {
63        let (n_channels, n_samples) = signal.data().dim();
64        let sds: Vec<f64> = signal
65            .data()
66            .axis_iter(Axis(0))
67            .map(|row| row.std(1.0))
68            .collect();
69        if let Some(c) = sds.iter().position(|&sd| !(sd > 0.0 && sd.is_finite())) {
70            return Err(MethodError::invalid(
71                "signal",
72                format!("channel {:?} is constant", signal.labels()[c]),
73            ));
74        }
75        let mut data = Array2::zeros((n_channels + self.noise.channels, n_samples));
76        for (c, sd) in sds.iter().enumerate() {
77            data.row_mut(c)
78                .assign(&signal.data().row(c).mapv(|v| v / sd));
79        }
80        let mut rng = ChaCha8Rng::seed_from_u64(self.noise.seed);
81        for v in data.slice_mut(s![n_channels.., ..]).iter_mut() {
82            let z: f64 = StandardNormal.sample(&mut rng);
83            *v = self.noise.sd * z;
84        }
85        let mut labels = signal.labels().to_vec();
86        labels.extend((0..self.noise.channels).map(|i| format!("noise{i}")));
87        Ok((
88            MultichannelSeries::new(data, signal.sampling_rate(), labels)
89                .map_err(|e| MethodError::ShapeMismatch(e.to_string()))?,
90            sds,
91        ))
92    }
93}
94
95impl ModeDecomposer for NaMvmd {
96    fn descriptor(&self) -> MethodDescriptor {
97        MethodDescriptor {
98            name: "na_mvmd",
99            version: env!("CARGO_PKG_VERSION"),
100            multivariate: true,
101        }
102    }
103
104    fn decompose(&self, signal: &MultichannelSeries) -> Result<Decomposition, MethodError> {
105        let (augmented, sds) = self.augment(signal)?;
106        let full = self.inner.decompose(&augmented)?;
107        let mut modes = full
108            .modes()
109            .slice(s![.., ..signal.n_channels(), ..])
110            .to_owned();
111        for (c, sd) in sds.iter().enumerate() {
112            modes.index_axis_mut(Axis(1), c).mapv_inplace(|v| v * sd);
113        }
114        Decomposition::new(
115            modes,
116            full.centre_frequencies_hz().to_vec(),
117            signal.sampling_rate(),
118            signal.labels().to_vec(),
119            true,
120        )
121    }
122}