cute_dsp/
windows.rs

1//! Window functions for spectral analysis
2//!
3//! This module provides window functions commonly used in spectral analysis,
4//! such as the Kaiser window and Approximate Confined Gaussian window.
5
6#![allow(unused_imports)]
7
8#[cfg(feature = "std")]
9use std::f32::consts::PI;
10
11#[cfg(not(feature = "std"))]
12use core::f32::consts::PI;
13
14use num_traits::{Float, FromPrimitive, NumCast};
15
16/// The Kaiser window (almost) maximizes the energy in the main-lobe compared to the side-lobes.
17pub struct Kaiser<T: Float> {
18    beta: T,
19    inv_b0: T,
20}
21
22impl<T: Float + FromPrimitive + NumCast> Kaiser<T> {
23    /// Create a Kaiser window with a given shape parameter.
24    pub fn new(beta: T) -> Self {
25        Self {
26            beta,
27            inv_b0: T::from_f32(1.0).unwrap() / Self::bessel0(beta),
28        }
29    }
30
31    /// Create a Kaiser window with a specified bandwidth
32    pub fn with_bandwidth(bandwidth: T, heuristic_optimal: bool) -> Self {
33        Self::new(Self::bandwidth_to_beta(bandwidth, heuristic_optimal))
34    }
35
36    /// Convert bandwidth to beta parameter
37    pub fn bandwidth_to_beta(bandwidth: T, heuristic_optimal: bool) -> T {
38        let bandwidth = if heuristic_optimal {
39            Self::heuristic_bandwidth(bandwidth)
40        } else {
41            bandwidth
42        };
43        let bandwidth = bandwidth.max(T::from_f32(2.0).unwrap());
44        let alpha = (bandwidth * bandwidth * T::from_f32(0.25).unwrap() - T::from_f32(1.0).unwrap()).sqrt();
45        alpha * T::from_f32(PI).unwrap()
46    }
47
48    /// Fills a slice with a Kaiser window
49    pub fn fill(&self, data: &mut [T]) {
50        let size = data.len();
51        let inv_size = T::from_f32(1.0).unwrap() / T::from_usize(size).unwrap();
52        for i in 0..size {
53            let r = T::from_usize(2 * i + 1).unwrap() * inv_size - T::from_f32(1.0).unwrap();
54            let arg = (T::from_f32(1.0).unwrap() - r * r).sqrt();
55            data[i] = Self::bessel0(self.beta * arg) * self.inv_b0;
56        }
57    }
58
59    // Modified Bessel function of the first kind, order 0
60    fn bessel0(x: T) -> T {
61        const SIGNIFICANCE_LIMIT: f32 = 1e-4;
62        let mut result = T::from_f32(0.0).unwrap();
63        let mut term = T::from_f32(1.0).unwrap();
64        let mut m = T::from_f32(0.0).unwrap();
65        while term > T::from_f32(SIGNIFICANCE_LIMIT).unwrap() {
66            result = result + term;
67            m = m + T::from_f32(1.0).unwrap();
68            term = term * (x * x) / (T::from_f32(4.0).unwrap() * m * m);
69        }
70        result
71    }
72
73    // Heuristic for optimal bandwidth
74    fn heuristic_bandwidth(bandwidth: T) -> T {
75        bandwidth + T::from_f32(8.0).unwrap() / ((bandwidth + T::from_f32(3.0).unwrap()) * (bandwidth + T::from_f32(3.0).unwrap()))
76            + T::from_f32(0.25).unwrap() * (T::from_f32(3.0).unwrap() - bandwidth).max(T::from_f32(0.0).unwrap())
77    }
78}
79
80/// The Approximate Confined Gaussian window
81pub struct ApproximateConfinedGaussian<T: Float> {
82    gaussian_factor: T,
83}
84
85impl <T: Float> ApproximateConfinedGaussian<T> {
86    /// Create an ACG window with a given shape parameter
87    pub fn new(sigma: T) -> Self {
88        Self {
89            gaussian_factor: T::from(0.0625).unwrap() / (sigma * sigma),
90        }
91    }
92
93    /// Create an ACG window with a specified bandwidth
94    pub fn with_bandwidth(bandwidth: T) -> Self {
95        Self::new(Self::bandwidth_to_sigma(bandwidth))
96    }
97
98    /// Heuristic map from bandwidth to sigma
99    pub fn bandwidth_to_sigma(bandwidth: T) -> T {
100        T::from(0.3).unwrap() / bandwidth.sqrt()
101    }
102
103    /// Fills a slice with an ACG window
104    pub fn fill(&self, data: &mut [T]) {
105        let size = data.len();
106        let inv_size = T::from(1.0).unwrap() / T::from(size as f32).unwrap();
107        let offset_scale = self.gaussian(T::from(1.0).unwrap()) / (self.gaussian(T::from(3.0).unwrap()) + self.gaussian(T::from(-1.0).unwrap()));
108        let norm = T::from(1.0).unwrap() / (self.gaussian(T::from(0.0).unwrap()) - T::from(2.0).unwrap() * offset_scale * self.gaussian(T::from(2.0).unwrap()));
109        
110        for i in 0..size {
111            let r = (T::from(2.0).unwrap() * T::from(i as f32).unwrap() + T::from(1.0).unwrap()) * inv_size - T::from(1.0).unwrap();
112            let value = norm * (self.gaussian(r) - offset_scale * (self.gaussian(r - T::from(2.0).unwrap()) + self.gaussian(r + T::from(2.0).unwrap())));
113            data[i] = value;
114        }
115    }
116
117    // Gaussian function
118    fn gaussian(&self, x: T) -> T {
119        (-T::from(x * x).unwrap() * self.gaussian_factor).exp()
120    }
121}
122
123/// Forces STFT perfect-reconstruction on an existing window
124pub fn force_perfect_reconstruction<T: Float>(
125    data: &mut [T],
126    window_length: usize,
127    interval: usize,
128) {
129    for i in 0..interval {
130        let mut sum2 = 0.0;
131        let mut index = i;
132        while index < window_length {
133            let val = data[index].to_f32().unwrap();
134            sum2 += val * val;
135            index += interval;
136        }
137        let factor = 1.0 / sum2.sqrt();
138        
139        index = i;
140        while index < window_length {
141            data[index] = data[index] * T::from(factor).unwrap();
142            index += interval;
143        }
144    }
145}
146
147//test
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn test_kaiser_window() {
154        let mut window = vec![0.0f32; 8];
155        let kaiser = Kaiser::new(3.0);
156        kaiser.fill(&mut window);
157        assert_eq!(window.len(), 8);
158        assert!(window.iter().all(|&x| x >= 0.0 && x <= 1.0));
159    }
160
161    #[test]
162    fn test_acg_window() {
163        let mut window = vec![0.0f32; 8];
164        let acg = ApproximateConfinedGaussian::new(1.0);
165        acg.fill(&mut window);
166        assert_eq!(window.len(), 8);
167        assert!(window.iter().all(|&x| x >= 0.0 && x <= 1.0));
168    }
169    
170    #[test]
171    fn test_force_perfect_reconstruction() {
172        let mut window = vec![0.1f32; 8];
173        force_perfect_reconstruction(&mut window, 8, 2);
174        assert_eq!(window.len(), 8);
175        assert!(window.iter().all(|&x| x >= 0.0 && x <= 1.0));
176    }
177    
178    #[test]
179    fn test_bandwidth_to_beta() {
180        let beta = Kaiser::<f32>::bandwidth_to_beta(4.0, true);
181        assert!(beta > 0.0);
182        let kaiser = Kaiser::new(beta);
183        let mut window = vec![0.0f32; 8];
184        kaiser.fill(&mut window);
185        assert_eq!(window.len(), 8);
186        assert!(window.iter().all(|&x| x >= 0.0 && x <= 1.0));
187    }
188    
189    #[test]
190    fn test_bandwidth_to_sigma() {
191        let sigma = ApproximateConfinedGaussian::<f32>::bandwidth_to_sigma(4.0);
192        assert!(sigma > 0.0);
193        let acg = ApproximateConfinedGaussian::new(sigma);
194        let mut window = vec![0.0f32; 8];
195        acg.fill(&mut window);
196        assert_eq!(window.len(), 8);
197        assert!(window.iter().all(|&x| x >= 0.0 && x <= 1.0));
198    }
199    
200    #[test]
201    fn test_heuristic_bandwidth() {
202        let bandwidth = Kaiser::<f32>::heuristic_bandwidth(4.0);
203        assert!(bandwidth > 0.0);
204        let beta = Kaiser::<f32>::bandwidth_to_beta(bandwidth, true);
205        assert!(beta > 0.0);
206        let kaiser = Kaiser::new(beta);
207        let mut window = vec![0.0f32; 8];
208        kaiser.fill(&mut window);
209        assert_eq!(window.len(), 8);
210        assert!(window.iter().all(|&x| x >= 0.0 && x <= 1.0));
211    }
212    
213    
214}
cute_dsp/windows.rs

cute_dsp/
windows.rs
