diff --git a/src/ps/aps.rs b/src/ps/aps.rs index 3de982c..72dd5e2 100644 --- a/src/ps/aps.rs +++ b/src/ps/aps.rs @@ -1,9 +1,4 @@ -//! Averaged power spectra module. Used to compute power spectra estimations on -//! long datasets, where nfft << length of data. This way, the variance of a -//! single periodogram is suppressed with increasing number of averages. -//! -//! For more information, see the book on numerical recipes. -//! +use super::timebuffer::TimeBuffer; use super::*; use crate::config::*; use anyhow::{bail, Result}; diff --git a/src/ps/mod.rs b/src/ps/mod.rs index e186ff2..21d66d2 100644 --- a/src/ps/mod.rs +++ b/src/ps/mod.rs @@ -1,8 +1,17 @@ -//! Power spectra, averaged power spectra, etc. This module contains several -pub mod window; -pub mod ps; -mod fft; +//! +//! Provides code to estimate (cross)[PowerSpectra], averaged power spectra +//! [AvPowerSpectra] using +//! Welch' method, and windows for time-windowing the data with non-rectangular +//! windows (also known as 'tapers'). +//! mod aps; +mod fft; +mod ps; +mod timebuffer; +mod window; +use crate::config::*; +pub type CrossPowerSpecra = Array3; +pub use aps::{ApsMode, AvPowerSpectra, Overlap, ApsResult}; pub use ps::PowerSpectra; -pub use window::{Window, WindowType}; \ No newline at end of file +pub use window::{Window, WindowType}; diff --git a/src/ps/ps.rs b/src/ps/ps.rs index 45fa83c..7366963 100644 --- a/src/ps/ps.rs +++ b/src/ps/ps.rs @@ -1,5 +1,3 @@ -//! Power spectra estimator, that uses a Windowed FFT to estimate cross-power -//! spectra. Window functions are documented in the `window` module. use crate::config::*; use ndarray::parallel::prelude::*; use num::pow::Pow; @@ -9,20 +7,25 @@ use std::usize; use crate::Dcol; -use super::fft::FFT; +use super::{fft::FFT, CrossPowerSpecra}; use super::window::*; use std::mem::MaybeUninit; use realfft::{RealFftPlanner, RealToComplex}; -/// Singlesided cross-Power spectra computation engine. -/// -/// Computes the signal(s) auto power and cross-power spectrum in each frequency -/// bin. +/// Single-sided (cross)power spectra estimator, that uses a Windowed FFT to estimate cross-power +/// spectra. Window functions are documented in the `window` module. Note that +/// directly using this power spectra estimator is generally not useful as it is +/// basically the periodogram estimator, with its high variance. +/// +/// This power spectrum estimator is instead used as a building block for for +/// example the computations of spectrograms, or Welch' method of spectral +/// estimation. +/// pub struct PowerSpectra { - // Window used in estimator + /// Window used in estimator pub window: Window, - // The window power, is corrected for in power spectra estimants + /// The window power, is corrected for in power spectra estimants pub sqrt_win_pwr: Flt, ffts: Vec, @@ -34,7 +37,7 @@ pub struct PowerSpectra { } impl PowerSpectra { - /// Return the FFT length used in power spectra computations + /// Returns the FFT length used in power spectra computations pub fn nfft(&self) -> usize { self.window.win.len() } @@ -44,9 +47,14 @@ impl PowerSpectra { /// /// - If win.len() != nfft /// - if nfft == 0 + /// + /// # Args + /// + /// - `window` - A `Window` struct, from which NFFT is also used. + /// pub fn newFromWindow(window: Window) -> PowerSpectra { let nfft = window.win.len(); - let win_pwr = window.win.mapv(|w| w.powi(2)).sum()/(nfft as Flt); + let win_pwr = window.win.mapv(|w| w.powi(2)).sum() / (nfft as Flt); assert!(nfft > 0); assert!(nfft % 2 == 0); @@ -64,7 +72,8 @@ impl PowerSpectra { } } - // Compute FFTs of input channel data. + /// Compute FFTs of input channel data. Stores the scaled FFT data in + /// self.freqdata. fn compute_ffts(&mut self, timedata: ArrayView2) -> &Array2 { let (n, nch) = timedata.dim(); let nfft = self.nfft(); @@ -76,9 +85,7 @@ impl PowerSpectra { self.freqdata .push_column(Ccol::from_vec(vec![Cflt::new(0., 0.); nfft / 2 + 1]).view()) .unwrap(); - self.timedata - .push_column(Dcol::zeros(nfft).view()) - .unwrap(); + self.timedata.push_column(Dcol::zeros(nfft).view()).unwrap(); } assert!(n == self.nfft()); @@ -105,25 +112,39 @@ impl PowerSpectra { /// Compute cross power spectra from input time data. First axis is /// frequency, second axis is channel i, third axis is channel j. - /// - /// # Argument - /// + /// + /// # Panics + /// + /// - When `timedata.nrows() != self.nfft()` + /// + /// # Args + /// /// * `tdata` - Input time data. This is a 2D array, where the first axis is /// time and the second axis is the channel number. - pub fn compute<'a, T>(&mut self, tdata: T) -> Array3 + /// + /// # Returns + /// + /// - 3D complex array of signal cross-powers with the following shape + /// (nfft/2+1,timedata.ncols(), timedata.ncols()). Its content is: + /// [freq_index, chi, chj] = crosspower: chi*conj(chj) + /// + pub fn compute<'a, T>(&mut self, tdata: T) -> CrossPowerSpecra where - T: Into>, + T: AsArray<'a, Flt, Ix2>, { let tdata = tdata.into(); - let clen = self.nfft() / 2 + 1; - let nchannel = tdata.ncols(); - let win_pwr = self.sqrt_win_pwr; + let nfft = self.nfft(); + let clen = nfft / 2 + 1; + if tdata.nrows() != nfft { + panic!("Invalid timedata length! Should be equal to nfft={nfft}"); + } + let nchannels = tdata.ncols(); // Compute fft of input data, and store in self.freqdata let fd = self.compute_ffts(tdata); let fdconj = fd.mapv(|c| c.conj()); - let result = Array3::uninit((clen, nchannel, nchannel)); + let result = Array3::uninit((clen, nchannels, nchannels)); let mut result: Array3 = unsafe { result.assume_init() }; // Loop over result axis one and channel i IN PARALLEL @@ -131,7 +152,7 @@ impl PowerSpectra { .and(fd.axis_iter(Axis(1))) .par_for_each(|mut out, chi| { // out: channel i of output 3D array, channel j all - // chi: channel i + // chi: channel i Zip::from(out.axis_iter_mut(Axis(1))) .and(fdconj.axis_iter(Axis(1))) .for_each(|mut out, chj| { @@ -140,18 +161,16 @@ impl PowerSpectra { Zip::from(&mut out) .and(chi) .and(chj) - .for_each(|out, chi, chjc|{ + .for_each(|out, chi, chjc| { // Loop over frequency components *out = 0.5 * chi * chjc; - } - ); - + }); + // The DC component has no 0.5 correction, as it only // occurs ones in a (double-sided) power spectrum. So // here we undo the 0.5 of 4 lines above here. out[0] *= 2.; - out[clen-1] *= 2.; - + out[clen - 1] *= 2.; }); }); result @@ -166,16 +185,16 @@ mod test { use rand_distr::StandardNormal; /// Generate a sine wave at the order i - fn generate_sinewave(nfft: usize,order: usize) -> Dcol { - Dcol::from_iter((0..nfft).map(|i| { - Flt::sin(i as Flt/(nfft) as Flt * order as Flt * 2.*pi) - })) + fn generate_sinewave(nfft: usize, order: usize) -> Dcol { + Dcol::from_iter( + (0..nfft).map(|i| Flt::sin(i as Flt / (nfft) as Flt * order as Flt * 2. * pi)), + ) } /// Generate a sine wave at the order i - fn generate_cosinewave(nfft: usize,order: usize) -> Dcol { - Dcol::from_iter((0..nfft).map(|i| { - Flt::cos(i as Flt/(nfft) as Flt * order as Flt * 2.*pi) - })) + fn generate_cosinewave(nfft: usize, order: usize) -> Dcol { + Dcol::from_iter( + (0..nfft).map(|i| Flt::cos(i as Flt / (nfft) as Flt * order as Flt * 2. * pi)), + ) } use super::*; @@ -206,8 +225,7 @@ mod test { // Start with a time signal let mut t: Dmat = Dmat::default((nfft, 0)); - t.push_column(generate_sinewave(nfft,1).view()) - .unwrap(); + t.push_column(generate_sinewave(nfft, 1).view()).unwrap(); // println!("{:?}", t); let fd = ps.compute_ffts(t.view()); @@ -227,37 +245,44 @@ mod test { ); } - /// Thest whether power spectra scale properly. Signals with amplitude of 1 /// should come back with a power of 0.5. DC offsets should come in as /// value^2 at frequency index 0. #[test] fn test_ps_scale() { - const nfft: usize = 124; let rect = Window::new(WindowType::Rect, nfft); let mut ps = PowerSpectra::newFromWindow(rect); // Start with a time signal let mut t: Dmat = Dmat::default((nfft, 0)); - t.push_column(generate_cosinewave(nfft,1).view()) - .unwrap(); + t.push_column(generate_cosinewave(nfft, 1).view()).unwrap(); let dc_component = 0.25; let dc_power = dc_component.pow(2); t.mapv_inplace(|t| t + dc_component); let power = ps.compute(t.view()); - assert_relative_eq!(power[(0, 0,0)].re, dc_power, epsilon = Flt::EPSILON * nfft as Flt); - assert_relative_eq!(power[(1, 0,0)].re, 0.5, epsilon = Flt::EPSILON * nfft as Flt); - assert_relative_eq!(power[(1, 0,0)].im, 0.0, epsilon = Flt::EPSILON * nfft as Flt); - + assert_relative_eq!( + power[(0, 0, 0)].re, + dc_power, + epsilon = Flt::EPSILON * nfft as Flt + ); + assert_relative_eq!( + power[(1, 0, 0)].re, + 0.5, + epsilon = Flt::EPSILON * nfft as Flt + ); + assert_relative_eq!( + power[(1, 0, 0)].im, + 0.0, + epsilon = Flt::EPSILON * nfft as Flt + ); } use ndarray_rand::RandomExt; // Test parseval's theorem for some random data #[test] fn test_parseval() { - const nfft: usize = 512; let rect = Window::new(WindowType::Rect, nfft); let mut ps = PowerSpectra::newFromWindow(rect); @@ -265,11 +290,11 @@ mod test { // Start with a time signal let t: Dmat = Dmat::random((nfft, 1), StandardNormal); - let tavg = t.sum()/(nfft as Flt); + let tavg = t.sum() / (nfft as Flt); let t_dc_power = tavg.powi(2); // println!("dc power in time domain: {:?}", t_dc_power); - let signal_pwr = t.mapv(|t| t.powi(2)).sum()/(nfft as Flt); + let signal_pwr = t.mapv(|t| t.powi(2)).sum() / (nfft as Flt); // println!("Total signal power in time domain: {:?} ", signal_pwr); let power = ps.compute(t.view()); @@ -277,15 +302,21 @@ mod test { let fpower = power.sum().abs(); - assert_ulps_eq!(t_dc_power, power[(0,0,0)].abs(), epsilon = Flt::EPSILON * (nfft as Flt).powi(2)); - assert_ulps_eq!(signal_pwr, fpower, epsilon = Flt::EPSILON * (nfft as Flt).powi(2)); - + assert_ulps_eq!( + t_dc_power, + power[(0, 0, 0)].abs(), + epsilon = Flt::EPSILON * (nfft as Flt).powi(2) + ); + assert_ulps_eq!( + signal_pwr, + fpower, + epsilon = Flt::EPSILON * (nfft as Flt).powi(2) + ); } // Test parseval's theorem for some random data #[test] fn test_parseval_with_window() { - // A sufficiently high value is required here, to show that it works. const nfft: usize = 2usize.pow(20); let window = Window::new(WindowType::Hann, nfft); @@ -293,13 +324,13 @@ mod test { let mut ps = PowerSpectra::newFromWindow(window); // Start with a time signal - let t: Dmat = 2.*Dmat::random((nfft, 1), StandardNormal); + let t: Dmat = 2. * Dmat::random((nfft, 1), StandardNormal); - let tavg = t.sum()/(nfft as Flt); + let tavg = t.sum() / (nfft as Flt); let t_dc_power = tavg.powi(2); // println!("dc power in time domain: {:?}", t_dc_power); - let signal_pwr = t.mapv(|t| t.powi(2)).sum()/(nfft as Flt); + let signal_pwr = t.mapv(|t| t.powi(2)).sum() / (nfft as Flt); // println!("Total signal power in time domain: {:?} ", signal_pwr); let power = ps.compute(t.view()); @@ -307,11 +338,13 @@ mod test { let fpower = power.sum().abs(); - assert_ulps_eq!(t_dc_power, power[(0,0,0)].abs(), epsilon = Flt::EPSILON * (nfft as Flt).powi(2)); + assert_ulps_eq!( + t_dc_power, + power[(0, 0, 0)].abs(), + epsilon = Flt::EPSILON * (nfft as Flt).powi(2) + ); // This one fails when nfft is too short. - assert_ulps_eq!(signal_pwr, fpower, epsilon = 1e-2); - + assert_ulps_eq!(signal_pwr, fpower, epsilon = 2e-2); } - }