modewise_wavelet/
decomposer.rs

1use modewise_decomposition::{Decomposition, ModeDecomposer};
2use modewise_fourier::centroid_hz;
3use modewise_method::{MethodDescriptor, MethodError};
4use modewise_signals::MultichannelSeries;
5use ndarray::{Array3, Axis, s};
6use rayon::prelude::*;
7
8use crate::filters::FILTERS;
9use crate::modwt::{Filters, mra};
10
11/// How the series is extended past its ends.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
13pub enum Boundary {
14    /// Circular, as the MODWT is defined (MATLAB and PyWavelets default).
15    #[default]
16    Periodic,
17    /// The series followed by its reverse, then cut back to the original length (Percival and
18    /// Walden 2000; MATLAB `modwt(..., 'reflection')`). Fewer edge effects for non-periodic series.
19    Reflection,
20}
21
22/// Wavelets available, by PyWavelets name.
23pub fn wavelet_names() -> impl Iterator<Item = &'static str> {
24    FILTERS.iter().map(|(name, _)| *name)
25}
26
27/// MODWT multiresolution analysis (Percival and Walden 2000): details D_1..D_J and the smooth
28/// S_J of each channel, additive and shift-invariant. Mode 0 is S_J and the last mode D_1. The
29/// same filter bank is used for every channel, so modes are aligned across channels.
30#[derive(Debug, Clone)]
31pub struct ModwtMra {
32    wavelet: &'static str,
33    filter: &'static [f64],
34    levels: Option<usize>,
35    boundary: Boundary,
36}
37
38impl ModwtMra {
39    pub fn new(
40        wavelet: &str,
41        levels: Option<usize>,
42        boundary: Boundary,
43    ) -> Result<Self, MethodError> {
44        let Some(&(name, filter)) = FILTERS.iter().find(|(name, _)| *name == wavelet) else {
45            let known: Vec<_> = wavelet_names().collect();
46            return Err(MethodError::invalid(
47                "wavelet",
48                format!("unknown {wavelet:?}; choose one of {}", known.join(", ")),
49            ));
50        };
51        if levels == Some(0) {
52            return Err(MethodError::invalid("levels", "must be at least 1"));
53        }
54        Ok(Self {
55            wavelet: name,
56            filter,
57            levels,
58            boundary,
59        })
60    }
61
62    /// The deepest level whose equivalent filter, (2^J − 1)(L − 1) + 1 long, fits in `n`
63    /// samples (Percival and Walden 2000), and at least 1.
64    pub fn default_levels(&self, n: usize) -> usize {
65        let l = self.filter.len();
66        (1..)
67            .take_while(|&j: &usize| ((1usize << j) - 1) * (l - 1) < n)
68            .last()
69            .unwrap_or(1)
70    }
71
72    pub fn wavelet(&self) -> &str {
73        self.wavelet
74    }
75}
76
77impl ModeDecomposer for ModwtMra {
78    fn descriptor(&self) -> MethodDescriptor {
79        MethodDescriptor {
80            name: "modwt_mra",
81            version: env!("CARGO_PKG_VERSION"),
82            multivariate: false,
83        }
84    }
85
86    fn decompose(&self, signal: &MultichannelSeries) -> Result<Decomposition, MethodError> {
87        let (n_channels, n) = signal.data().dim();
88        let levels = self.levels.unwrap_or_else(|| self.default_levels(n));
89        let longest = (n.ilog2() as usize).max(1);
90        if levels > longest {
91            return Err(MethodError::invalid(
92                "levels",
93                format!("at most ⌊log₂ n⌋ = {longest} for {n} samples"),
94            ));
95        }
96        let filters = Filters::from_scaling(self.filter);
97        let rows: Vec<Vec<f64>> = signal
98            .data()
99            .rows()
100            .into_iter()
101            .map(|r| r.to_vec())
102            .collect();
103        let channels: Vec<_> = rows
104            .par_iter()
105            .map(|x| match self.boundary {
106                Boundary::Periodic => mra(x, &filters, levels),
107                Boundary::Reflection => {
108                    let mut extended = x.clone();
109                    extended.extend(x.iter().rev());
110                    let (details, smooth) = mra(&extended, &filters, levels);
111                    (
112                        details.into_iter().map(|d| d[..n].to_vec()).collect(),
113                        smooth[..n].to_vec(),
114                    )
115                }
116            })
117            .collect();
118        let mut modes = Array3::zeros((levels + 1, n_channels, n));
119        for (c, (details, smooth)) in channels.iter().enumerate() {
120            modes
121                .slice_mut(s![0, c, ..])
122                .assign(&ndarray::ArrayView1::from(smooth.as_slice()));
123            for (j, d) in details.iter().enumerate() {
124                modes
125                    .slice_mut(s![levels - j, c, ..])
126                    .assign(&ndarray::ArrayView1::from(d.as_slice()));
127            }
128        }
129        let fs = signal.sampling_rate();
130        let centres = modes
131            .axis_iter(Axis(0))
132            .map(|m| centroid_hz(m, fs.hz()))
133            .collect();
134        Decomposition::new(modes, centres, fs, signal.labels().to_vec(), true)
135    }
136}