diff --git a/src/bin/main.rs b/src/bin/main.rs index 8bbf569..1335649 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -39,6 +39,10 @@ pub struct Opts { /// Sample rate in Hertz #[arg(short, long, default_value_t = 1.0f32)] fs: f32, + + /// Exponential averaging + #[arg(short, long)] + avg: Option, } fn main() -> Result<()> { @@ -48,6 +52,7 @@ fn main() -> Result<()> { min_avg, detrend, fs, + avg, } = Opts::parse(); let logfs = fs.log10(); @@ -71,6 +76,7 @@ fn main() -> Result<()> { let mut c = PsdCascade::<{ 1 << 9 }>::default(); c.set_stage_depth(3); c.set_detrend(detrend); + c.set_avg(avg); c })); i = 0; diff --git a/src/psd.rs b/src/psd.rs index 4f324a2..386933f 100644 --- a/src/psd.rs +++ b/src/psd.rs @@ -85,6 +85,7 @@ pub struct Psd { overlap: usize, detrend: Detrend, drain: usize, + avg: Option, } impl Psd { @@ -104,12 +105,17 @@ impl Psd { overlap: 0, detrend: Detrend::None, drain: 0, + avg: None, }; s.set_overlap(N / 2); s.set_stage_depth(0); s } + pub fn set_avg(&mut self, avg: Option) { + self.avg = avg; + } + pub fn set_overlap(&mut self, o: usize) { assert_eq!(o % self.hbf.block_size().0, 0); assert!(self.hbf.block_size().1 >= o); @@ -222,7 +228,14 @@ impl PsdStage for Psd { .iter() .zip(self.spectrum[..N / 2 + 1].iter_mut()) { - *p += c.norm_sqr(); + let ci = c.norm_sqr(); + *p += if self.count == 0 { + ci + } else if let Some(avg) = self.avg { + avg * (ci - *p) + } else { + ci + }; } let start = if self.count == 0 { @@ -266,7 +279,8 @@ impl PsdStage for Psd { fn gain(&self) -> f32 { // 2 for one-sided // overlap is compensated by counting - 1.0 / ((self.count * N / 2) as f32 * self.win.nenbw * self.win.power) + let c = if self.avg.is_some() { 1 } else { self.count }; + 1.0 / ((c * N / 2) as f32 * self.win.nenbw * self.win.power) } fn buf(&self) -> &[f32] { @@ -332,6 +346,7 @@ pub struct PsdCascade { detrend: Detrend, overlap: usize, win: Arc>, + avg: Option, } impl Default for PsdCascade { @@ -350,6 +365,7 @@ impl Default for PsdCascade { detrend: Detrend::None, overlap: N / 2, win, + avg: None, } } } @@ -367,6 +383,10 @@ impl PsdCascade { } } + pub fn set_avg(&mut self, avg: Option) { + self.avg = avg; + } + pub fn set_detrend(&mut self, d: Detrend) { self.detrend = d; } @@ -377,6 +397,7 @@ impl PsdCascade { stage.set_stage_depth(self.stage_length); stage.set_detrend(self.detrend); stage.set_overlap(self.overlap); + stage.set_avg(self.avg); self.stages.push(stage); } &mut self.stages[i]