modewise_decomposition/
lib.rs

1use modewise_method::{MethodDescriptor, MethodError};
2use modewise_signals::{MultichannelSeries, SamplingRate};
3use ndarray::{Array2, Array3, ArrayView2, ArrayView3, Axis};
4
5mod loadings;
6
7pub use loadings::Loadings;
8
9/// Instantaneous amplitude and frequency of every mode, laid out as mode × channel × sample,
10/// with the centre frequencies of the modes they came from.
11#[derive(Debug, Clone, PartialEq)]
12pub struct InstantaneousModes {
13    amplitude: Array3<f64>,
14    frequency_hz: Array3<f64>,
15    centre_frequencies_hz: Vec<f64>,
16    sampling_rate: SamplingRate,
17    labels: Vec<String>,
18}
19
20impl InstantaneousModes {
21    pub fn new(
22        amplitude: Array3<f64>,
23        frequency_hz: Array3<f64>,
24        centre_frequencies_hz: Vec<f64>,
25        sampling_rate: SamplingRate,
26        labels: Vec<String>,
27    ) -> Result<Self, MethodError> {
28        let (n_modes, n_channels, _) = amplitude.dim();
29        if frequency_hz.dim() != amplitude.dim() {
30            return Err(MethodError::ShapeMismatch(format!(
31                "amplitude {:?} and frequency {:?} differ",
32                amplitude.dim(),
33                frequency_hz.dim()
34            )));
35        }
36        if centre_frequencies_hz.len() != n_modes || labels.len() != n_channels {
37            return Err(MethodError::ShapeMismatch(format!(
38                "{n_modes} modes and {n_channels} channels, but {} centre frequencies and {} labels",
39                centre_frequencies_hz.len(),
40                labels.len()
41            )));
42        }
43        Ok(Self {
44            amplitude,
45            frequency_hz,
46            centre_frequencies_hz,
47            sampling_rate,
48            labels,
49        })
50    }
51
52    pub fn amplitude(&self) -> ArrayView3<'_, f64> {
53        self.amplitude.view()
54    }
55
56    pub fn frequency_hz(&self) -> ArrayView3<'_, f64> {
57        self.frequency_hz.view()
58    }
59
60    pub fn centre_frequencies_hz(&self) -> &[f64] {
61        &self.centre_frequencies_hz
62    }
63
64    pub fn sampling_rate(&self) -> SamplingRate {
65        self.sampling_rate
66    }
67
68    pub fn labels(&self) -> &[String] {
69        &self.labels
70    }
71
72    pub fn n_samples(&self) -> usize {
73        self.amplitude.dim().2
74    }
75}
76
77fn median(mut values: Vec<f64>) -> f64 {
78    if values.is_empty() {
79        return f64::NAN;
80    }
81    values.sort_by(f64::total_cmp);
82    let mid = values.len() / 2;
83    if values.len() % 2 == 1 {
84        values[mid]
85    } else {
86        (values[mid - 1] + values[mid]) / 2.0
87    }
88}
89
90/// Splits a signal into time-domain modes.
91pub trait ModeDecomposer {
92    fn descriptor(&self) -> MethodDescriptor;
93    fn decompose(&self, signal: &MultichannelSeries) -> Result<Decomposition, MethodError>;
94}
95
96/// Modes laid out as mode × channel × sample, ordered by ascending centre frequency.
97#[derive(Debug, Clone, PartialEq)]
98pub struct Decomposition {
99    modes: Array3<f64>,
100    centre_frequencies_hz: Vec<f64>,
101    sampling_rate: SamplingRate,
102    labels: Vec<String>,
103    aligned_across_channels: bool,
104    /// Mode × channel, for methods that decompose each channel on its own.
105    channel_centre_frequencies_hz: Option<Array2<f64>>,
106}
107
108impl Decomposition {
109    pub fn new(
110        modes: Array3<f64>,
111        centre_frequencies_hz: Vec<f64>,
112        sampling_rate: SamplingRate,
113        labels: Vec<String>,
114        aligned_across_channels: bool,
115    ) -> Result<Self, MethodError> {
116        let (n_modes, n_channels, _) = modes.dim();
117        if centre_frequencies_hz.len() != n_modes {
118            return Err(MethodError::ShapeMismatch(format!(
119                "{n_modes} modes but {} centre frequencies",
120                centre_frequencies_hz.len()
121            )));
122        }
123        if labels.len() != n_channels {
124            return Err(MethodError::ShapeMismatch(format!(
125                "{n_channels} channels but {} labels",
126                labels.len()
127            )));
128        }
129        Ok(Self {
130            modes,
131            centre_frequencies_hz,
132            sampling_rate,
133            labels,
134            aligned_across_channels,
135            channel_centre_frequencies_hz: None,
136        })
137    }
138
139    /// Modes from a method that decomposes each channel on its own (VMD, EMD). Mode k need not
140    /// mean the same oscillation in every channel, so the decomposition is not aligned. Each
141    /// mode's centre frequency is the median of its channels' finite values; NaN marks a channel
142    /// without that mode.
143    pub fn per_channel(
144        modes: Array3<f64>,
145        channel_centre_frequencies_hz: Array2<f64>,
146        sampling_rate: SamplingRate,
147        labels: Vec<String>,
148    ) -> Result<Self, MethodError> {
149        let centres = channel_centre_frequencies_hz
150            .rows()
151            .into_iter()
152            .map(|row| median(row.iter().copied().filter(|v| v.is_finite()).collect()))
153            .collect();
154        Self::new(modes, centres, sampling_rate, labels, false)?
155            .with_channel_centre_frequencies(channel_centre_frequencies_hz)
156    }
157
158    /// Attaches per-channel centre frequencies (mode × channel).
159    pub fn with_channel_centre_frequencies(
160        mut self,
161        frequencies_hz: Array2<f64>,
162    ) -> Result<Self, MethodError> {
163        let (n_modes, n_channels, _) = self.modes.dim();
164        if frequencies_hz.dim() != (n_modes, n_channels) {
165            return Err(MethodError::ShapeMismatch(format!(
166                "{n_modes} modes and {n_channels} channels, but channel centre frequencies are {:?}",
167                frequencies_hz.dim()
168            )));
169        }
170        self.channel_centre_frequencies_hz = Some(frequencies_hz);
171        Ok(self)
172    }
173
174    /// Mode × channel centre frequencies, when the method gives them per channel.
175    pub fn channel_centre_frequencies_hz(&self) -> Option<ArrayView2<'_, f64>> {
176        self.channel_centre_frequencies_hz
177            .as_ref()
178            .map(|f| f.view())
179    }
180
181    pub fn modes(&self) -> ArrayView3<'_, f64> {
182        self.modes.view()
183    }
184
185    pub fn mode(&self, index: usize) -> ArrayView2<'_, f64> {
186        self.modes.index_axis(Axis(0), index)
187    }
188
189    pub fn centre_frequencies_hz(&self) -> &[f64] {
190        &self.centre_frequencies_hz
191    }
192
193    pub fn sampling_rate(&self) -> SamplingRate {
194        self.sampling_rate
195    }
196
197    pub fn labels(&self) -> &[String] {
198        &self.labels
199    }
200
201    /// True when mode `k` means the same oscillation in every channel.
202    pub fn aligned_across_channels(&self) -> bool {
203        self.aligned_across_channels
204    }
205
206    pub fn n_modes(&self) -> usize {
207        self.modes.len_of(Axis(0))
208    }
209
210    /// Sums the chosen modes back into a signal, for example to keep one frequency band.
211    pub fn reconstruct(&self, modes: &[usize]) -> Result<MultichannelSeries, MethodError> {
212        if modes.is_empty() {
213            return Err(MethodError::invalid("modes", "choose at least one mode"));
214        }
215        let (_, n_channels, n_samples) = self.modes.dim();
216        let mut sum = Array2::zeros((n_channels, n_samples));
217        for (i, &k) in modes.iter().enumerate() {
218            if k >= self.n_modes() {
219                return Err(MethodError::invalid(
220                    "modes",
221                    format!(
222                        "mode {k} does not exist; there are {} modes",
223                        self.n_modes()
224                    ),
225                ));
226            }
227            if modes[..i].contains(&k) {
228                return Err(MethodError::invalid(
229                    "modes",
230                    format!("mode {k} is listed twice"),
231                ));
232            }
233            sum += &self.mode(k);
234        }
235        MultichannelSeries::new(sum, self.sampling_rate, self.labels.clone())
236            .map_err(|e| MethodError::ShapeMismatch(e.to_string()))
237    }
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243
244    fn fs() -> SamplingRate {
245        SamplingRate::new(1.0).unwrap()
246    }
247
248    #[test]
249    fn rejects_frequency_count_mismatch() {
250        let modes = Array3::zeros((2, 1, 4));
251        let err = Decomposition::new(modes, vec![0.1], fs(), vec!["a".into()], true);
252        assert!(matches!(err, Err(MethodError::ShapeMismatch(_))));
253    }
254
255    #[test]
256    fn rejects_label_count_mismatch() {
257        let modes = Array3::zeros((1, 2, 4));
258        let err = Decomposition::new(modes, vec![0.1], fs(), vec!["a".into()], true);
259        assert!(matches!(err, Err(MethodError::ShapeMismatch(_))));
260    }
261
262    #[test]
263    fn reconstruct_sums_chosen_modes() {
264        let mut modes = Array3::zeros((3, 1, 2));
265        modes[[0, 0, 0]] = 1.0;
266        modes[[1, 0, 0]] = 2.0;
267        modes[[2, 0, 1]] = 4.0;
268        let d =
269            Decomposition::new(modes, vec![0.1, 0.2, 0.3], fs(), vec!["a".into()], true).unwrap();
270        let series = d.reconstruct(&[0, 2]).unwrap();
271        assert_eq!(series.data().row(0).to_vec(), [1.0, 4.0]);
272        assert_eq!(series.labels(), ["a"]);
273    }
274
275    #[test]
276    fn reconstruct_rejects_bad_selection() {
277        let d = Decomposition::new(
278            Array3::zeros((2, 1, 2)),
279            vec![0.1, 0.2],
280            fs(),
281            vec!["a".into()],
282            true,
283        )
284        .unwrap();
285        for selection in [&[][..], &[2], &[1, 1]] {
286            assert!(d.reconstruct(selection).is_err(), "accepted {selection:?}");
287        }
288    }
289
290    #[test]
291    fn getters_return_what_was_stored() {
292        let d = Decomposition::new(
293            Array3::zeros((2, 1, 3)),
294            vec![0.1, 0.2],
295            fs(),
296            vec!["amy".into()],
297            false,
298        )
299        .unwrap();
300        assert_eq!(d.centre_frequencies_hz(), [0.1, 0.2]);
301        assert_eq!(d.labels(), ["amy"]);
302        assert!(!d.aligned_across_channels());
303        assert_eq!(d.sampling_rate(), fs());
304        let aligned = Decomposition::new(
305            Array3::zeros((1, 1, 3)),
306            vec![0.1],
307            fs(),
308            vec!["a".into()],
309            true,
310        )
311        .unwrap();
312        assert!(aligned.aligned_across_channels());
313    }
314
315    #[test]
316    fn instantaneous_modes_check_shapes() {
317        let ok = InstantaneousModes::new(
318            Array3::zeros((2, 1, 4)),
319            Array3::zeros((2, 1, 4)),
320            vec![0.1, 0.2],
321            fs(),
322            vec!["a".into()],
323        )
324        .unwrap();
325        assert_eq!(
326            (ok.n_samples(), ok.centre_frequencies_hz(), ok.labels()),
327            (4, [0.1, 0.2].as_slice(), ["a".to_string()].as_slice())
328        );
329        assert_eq!(ok.sampling_rate(), fs());
330        assert_eq!(ok.frequency_hz().dim(), ok.amplitude().dim());
331        assert!(
332            InstantaneousModes::new(
333                Array3::zeros((2, 1, 4)),
334                Array3::zeros((2, 1, 5)),
335                vec![0.1, 0.2],
336                fs(),
337                vec!["a".into()]
338            )
339            .is_err()
340        );
341        assert!(
342            InstantaneousModes::new(
343                Array3::zeros((2, 1, 4)),
344                Array3::zeros((2, 1, 4)),
345                vec![0.1],
346                fs(),
347                vec!["a".into()]
348            )
349            .is_err()
350        );
351        assert!(
352            InstantaneousModes::new(
353                Array3::zeros((2, 1, 4)),
354                Array3::zeros((2, 1, 4)),
355                vec![0.1, 0.2],
356                fs(),
357                vec![]
358            )
359            .is_err()
360        );
361    }
362
363    #[test]
364    fn exposes_single_mode() {
365        let mut modes = Array3::zeros((2, 1, 3));
366        modes[[1, 0, 2]] = 5.0;
367        let d = Decomposition::new(modes, vec![0.1, 0.2], fs(), vec!["a".into()], true).unwrap();
368        assert_eq!(d.n_modes(), 2);
369        assert_eq!(d.mode(1)[[0, 2]], 5.0);
370    }
371}
372
373#[cfg(test)]
374mod per_channel {
375    use super::*;
376
377    #[test]
378    fn mode_frequencies_are_channel_medians_ignoring_missing_modes() {
379        let fs = SamplingRate::new(1.0).unwrap();
380        let channel = ndarray::array![
381            [0.1, 0.3, 0.2],
382            [0.05, f64::NAN, 0.25],
383            [f64::NAN, f64::NAN, f64::NAN]
384        ];
385        let labels = vec!["a".into(), "b".into(), "c".into()];
386        let d = Decomposition::per_channel(
387            Array3::zeros((3, 3, 4)),
388            channel.clone(),
389            fs,
390            labels.clone(),
391        )
392        .unwrap();
393        assert_eq!(d.centre_frequencies_hz()[..2], [0.2, 0.15]);
394        assert!(d.centre_frequencies_hz()[2].is_nan());
395        assert!(!d.aligned_across_channels());
396        let stored = d.channel_centre_frequencies_hz().unwrap();
397        assert!(
398            stored
399                .iter()
400                .zip(&channel)
401                .all(|(a, b)| a.to_bits() == b.to_bits())
402        );
403        assert!(
404            Decomposition::per_channel(Array3::zeros((3, 2, 4)), channel, fs, labels[..2].to_vec())
405                .is_err()
406        );
407        let plain = Decomposition::new(
408            Array3::zeros((1, 1, 4)),
409            vec![0.1],
410            fs,
411            vec!["a".into()],
412            true,
413        )
414        .unwrap();
415        assert!(plain.channel_centre_frequencies_hz().is_none());
416    }
417}