modewise_mvmd/
noise_assisted.rs1use 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#[derive(Debug, Clone, PartialEq)]
13pub struct NoiseParams {
14 pub channels: usize,
15 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#[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 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}