modewise_cwt/
lib.rs

1//! Continuous wavelet transform with a complex Morlet wavelet, after Torrence & Compo (1998).
2//!
3//! Conventions (docs/time-frequency.md), matching the pycwt implementation of the same paper:
4//! - wavelet in frequency space `π^-1/4 · exp(-(sω − ω0)² / 2)`, normalised by `√(2πs / δt)`
5//! - Fourier period of scale `s`: `λ·s` with `λ = 4π / (ω0 + √(2 + ω0²))`
6//! - signal zero-padded to the next power of two before the FFT
7//! - cone of influence: e-folding time `√2·s`, i.e. a period `P` is reliable at sample `i`
8//!   when `P ≤ λ/√2 · δt · (min(i, n − 1 − i) + 0.5)`, the distance to the record boundary
9//!   half a sample beyond the first and last sample (as pycwt).
10
11use std::f64::consts::PI;
12
13use modewise_method::{MethodDescriptor, MethodError};
14use modewise_signals::MultichannelSeries;
15use modewise_spectral::{FrequencyGrid, TimeFrequencyMap, TimeFrequencyTransform};
16use ndarray::{Array2, Array3};
17use rustfft::FftPlanner;
18use rustfft::num_complex::Complex64;
19
20#[derive(Debug, Clone)]
21pub struct Cwt {
22    grid: FrequencyGrid,
23    omega0: f64,
24    rectify: bool,
25}
26
27impl Cwt {
28    /// `grid` gives the frequencies to analyse (its bin centres); `omega0` is the Morlet
29    /// wavenumber (6 in Torrence & Compo). With `rectify`, power is divided by scale
30    /// (Liu, Liang & Weisberg 2007).
31    pub fn new(grid: FrequencyGrid, omega0: f64, rectify: bool) -> Result<Self, MethodError> {
32        if !(omega0.is_finite() && omega0 >= 5.0) {
33            return Err(MethodError::invalid(
34                "omega0",
35                "must be finite and >= 5 for the Morlet to be admissible in practice",
36            ));
37        }
38        Ok(Self {
39            grid,
40            omega0,
41            rectify,
42        })
43    }
44
45    /// Fourier period of scale 1, `λ`.
46    pub fn fourier_factor(&self) -> f64 {
47        4.0 * PI / (self.omega0 + (2.0 + self.omega0 * self.omega0).sqrt())
48    }
49}
50
51impl TimeFrequencyTransform for Cwt {
52    fn descriptor(&self) -> MethodDescriptor {
53        MethodDescriptor {
54            name: "cwt",
55            version: env!("CARGO_PKG_VERSION"),
56            multivariate: false,
57        }
58    }
59
60    fn transform(&self, signal: &MultichannelSeries) -> Result<TimeFrequencyMap, MethodError> {
61        let fs = signal.sampling_rate();
62        let dt = fs.period_s();
63        let n = signal.n_samples();
64        let padded = n.next_power_of_two();
65        let lambda = self.fourier_factor();
66        let freqs = self.grid.centres();
67        let nyquist = fs.hz() / 2.0;
68        if let Some(f) = freqs.iter().find(|&&f| f <= 0.0 || f > nyquist) {
69            return Err(MethodError::invalid(
70                "frequencies",
71                format!("{f} Hz is outside (0, {nyquist}] Hz"),
72            ));
73        }
74        let scales: Vec<f64> = freqs.iter().map(|f| 1.0 / (lambda * f)).collect();
75        let omegas: Vec<f64> = (0..padded)
76            .map(|k| {
77                // The Nyquist bin counts as negative, as numpy's fftfreq (and pycwt) do.
78                let k = if k < padded / 2 {
79                    k as f64
80                } else {
81                    k as f64 - padded as f64
82                };
83                2.0 * PI * k / (padded as f64 * dt)
84            })
85            .collect();
86
87        let mut planner = FftPlanner::new();
88        let forward = planner.plan_fft_forward(padded);
89        let inverse = planner.plan_fft_inverse(padded);
90        let mut power = Array3::zeros((signal.n_channels(), scales.len(), n));
91        for (c, channel) in signal.data().rows().into_iter().enumerate() {
92            let mut spectrum: Vec<Complex64> =
93                channel.iter().map(|&v| Complex64::new(v, 0.0)).collect();
94            spectrum.resize(padded, Complex64::new(0.0, 0.0));
95            forward.process(&mut spectrum);
96            for (j, &s) in scales.iter().enumerate() {
97                let norm = (2.0 * PI * s / dt).sqrt() * PI.powf(-0.25);
98                let mut w: Vec<Complex64> = spectrum
99                    .iter()
100                    .zip(&omegas)
101                    .map(|(x, &om)| x * norm * (-0.5 * (s * om - self.omega0).powi(2)).exp())
102                    .collect();
103                inverse.process(&mut w);
104                let rectify = if self.rectify { s } else { 1.0 };
105                for t in 0..n {
106                    power[[c, j, t]] = (w[t] / padded as f64).norm_sqr() / rectify;
107                }
108            }
109        }
110
111        let coi_factor = lambda / 2f64.sqrt() * dt;
112        let reliable = Array2::from_shape_fn((freqs.len(), n), |(j, t)| {
113            1.0 / freqs[j] <= coi_factor * (t.min(n - 1 - t) as f64 + 0.5)
114        });
115        let times = (0..n).map(|t| t as f64 * dt).collect();
116        TimeFrequencyMap::new(power, freqs, times, signal.labels().to_vec())?
117            .with_sampling_rate(fs)?
118            .with_reliable(reliable)
119    }
120}
modewise_cwt/lib.rs

modewise_cwt/
lib.rs
