use std::f32::consts::PI;
pub type Float = super::Float;
pub type Complex = super::Complex;
pub trait Wavelet {
fn generate_mother(&mut self, size: usize) -> Vec<Float>;
fn generate(&self, size: usize, scale: Float) -> Vec<Complex>;
fn bandwidth(&self) -> Float;
fn four_wavelen(&self) -> Float;
fn imag_frequency(&self) -> bool;
fn is_double_sided(&self) -> bool;
fn mother(&self) -> &[Float];
fn get_support(&self, scale: Float) -> isize {
(self.bandwidth() * scale * 3.0) as isize
}
}
pub struct MorletWavelet {
four_wavelen: Float,
imag_frequency: bool,
double_sided: bool,
mother: Vec<Float>,
fb: Float,
ifb: Float,
fb2: Float,
}
impl MorletWavelet {
pub fn new(bandwidth: Float) -> Self {
Self {
four_wavelen: 0.9876,
fb: bandwidth,
fb2: 2.0 * bandwidth * bandwidth,
ifb: 1.0 / bandwidth,
imag_frequency: false,
double_sided: false,
mother: vec![],
}
}
}
impl Wavelet for MorletWavelet {
fn generate_mother(&mut self, size: usize) -> Vec<Float> {
let mut mother = Vec::with_capacity(size);
let torad = (2.0*PI)/size as Float;
let norm = (2.0*PI).sqrt()*super::IPI4;
for i in 0..size {
let mut tmp = 2.0 * (i as Float * torad) * self.fb - 2.0 * PI * self.fb;
tmp = -(tmp * tmp) / 2.0;
mother.push(norm * tmp.exp());
}
self.mother = mother.clone();
mother
}
fn generate(&self, size: usize, scale: Float) -> Vec<Complex> {
let width = self.get_support(scale);
let norm = size as Float * self.ifb * super::IPI4;
let mut output: Vec<Complex> = Vec::with_capacity((width*2+1) as usize);
for i in 0..width*2+1 {
let tmp1 = (i - width) as Float / scale;
let tmp2 = (-(tmp1*tmp1) / self.fb2).exp();
let real = norm * tmp2 * (tmp1*2.0*PI).cos() / scale;
let imag = norm * tmp2 * (tmp1*2.0*PI).sin() / scale;
output.push(Complex::new(real, imag));
}
output
}
fn bandwidth(&self) -> Float {
self.fb
}
fn four_wavelen(&self) -> Float {
self.four_wavelen
}
fn imag_frequency(&self) -> bool {
self.imag_frequency
}
fn is_double_sided(&self) -> bool {
self.double_sided
}
fn mother(&self) -> &[Float] {
self.mother.as_slice()
}
}
#[cfg(test)]
mod tests {
use super::*;
const EPS: Float = 1e-5;
#[test]
fn test_morlet_wavelet_new() {
let morlet = MorletWavelet::new(1.0);
assert_eq!(morlet.fb, 1.0);
assert_eq!(morlet.fb2, 2.0);
assert_eq!(morlet.ifb, 1.0);
assert_eq!(morlet.four_wavelen, 0.9876);
assert!(!morlet.imag_frequency);
assert!(!morlet.double_sided);
}
#[test]
fn test_morlet_wavelet_generate_mother() {
let mut morlet = MorletWavelet::new(1.0);
morlet.mother = morlet.generate_mother(4);
assert_eq!(morlet.mother.len(), 4);
}
#[test]
fn test_morlet_wavelet_generate_time() {
let morlet = MorletWavelet::new(2.0);
let size = 25;
let scale = 2.0;
let result = morlet.generate(size, scale);
assert_eq!(result.len(), (morlet.get_support(scale) as Float * 2.0 + 1.0) as usize);
assert!(result[0].re - 0.05215157 < EPS);
assert!(result[0].im - -4.9752100e-09 < EPS);
assert!(result[1].re - -0.10700808 < EPS);
assert!(result[1].im - 2.8149339e-07 < EPS);
for (i, complex) in result.iter().enumerate() {
let tmp1 = (i as isize - morlet.get_support(scale)) as f32 / scale;
let tmp2 = (-(tmp1 * tmp1) / morlet.fb2).exp();
let expected_real = size as Float * morlet.ifb * crate::IPI4 * tmp2 * (tmp1 * 2.0 * PI).cos() / scale;
let expected_imag = size as Float * morlet.ifb * crate::IPI4 * tmp2 * (tmp1 * 2.0 * PI).sin() / scale;
assert!((complex.re - expected_real).abs() < EPS);
assert!((complex.im - expected_imag).abs() < EPS);
}
}
#[test]
fn test_morlet_wavelet_bandwidth() {
let morlet = MorletWavelet::new(2.5);
assert_eq!(morlet.bandwidth(), 2.5);
}
#[test]
fn test_morlet_wavelet_get_support() {
let morlet = MorletWavelet::new(1.0);
assert_eq!(morlet.get_support(1.0), 3);
}
}