Updated comments and test code

This commit is contained in:
Anne de Jong 2024-07-06 19:38:22 +02:00
parent b366c47ca7
commit 0a847318f3
3 changed files with 111 additions and 74 deletions

View File

@ -1,9 +1,4 @@
//! Averaged power spectra module. Used to compute power spectra estimations on use super::timebuffer::TimeBuffer;
//! 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::*; use super::*;
use crate::config::*; use crate::config::*;
use anyhow::{bail, Result}; use anyhow::{bail, Result};

View File

@ -1,8 +1,17 @@
//! Power spectra, averaged power spectra, etc. This module contains several //!
pub mod window; //! Provides code to estimate (cross)[PowerSpectra], averaged power spectra
pub mod ps; //! [AvPowerSpectra] using
mod fft; //! Welch' method, and windows for time-windowing the data with non-rectangular
//! windows (also known as 'tapers').
//!
mod aps; mod aps;
mod fft;
mod ps;
mod timebuffer;
mod window;
use crate::config::*;
pub type CrossPowerSpecra = Array3<Cflt>;
pub use aps::{ApsMode, AvPowerSpectra, Overlap, ApsResult};
pub use ps::PowerSpectra; pub use ps::PowerSpectra;
pub use window::{Window, WindowType}; pub use window::{Window, WindowType};

View File

@ -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 crate::config::*;
use ndarray::parallel::prelude::*; use ndarray::parallel::prelude::*;
use num::pow::Pow; use num::pow::Pow;
@ -9,20 +7,25 @@ use std::usize;
use crate::Dcol; use crate::Dcol;
use super::fft::FFT; use super::{fft::FFT, CrossPowerSpecra};
use super::window::*; use super::window::*;
use std::mem::MaybeUninit; use std::mem::MaybeUninit;
use realfft::{RealFftPlanner, RealToComplex}; use realfft::{RealFftPlanner, RealToComplex};
/// Singlesided cross-Power spectra computation engine. /// 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.
/// ///
/// Computes the signal(s) auto power and cross-power spectrum in each frequency
/// bin.
pub struct PowerSpectra { pub struct PowerSpectra {
// Window used in estimator /// Window used in estimator
pub window: Window, 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, pub sqrt_win_pwr: Flt,
ffts: Vec<FFT>, ffts: Vec<FFT>,
@ -34,7 +37,7 @@ pub struct PowerSpectra {
} }
impl 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 { pub fn nfft(&self) -> usize {
self.window.win.len() self.window.win.len()
} }
@ -44,9 +47,14 @@ impl PowerSpectra {
/// ///
/// - If win.len() != nfft /// - If win.len() != nfft
/// - if nfft == 0 /// - if nfft == 0
///
/// # Args
///
/// - `window` - A `Window` struct, from which NFFT is also used.
///
pub fn newFromWindow(window: Window) -> PowerSpectra { pub fn newFromWindow(window: Window) -> PowerSpectra {
let nfft = window.win.len(); 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 > 0);
assert!(nfft % 2 == 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<Flt>) -> &Array2<Cflt> { fn compute_ffts(&mut self, timedata: ArrayView2<Flt>) -> &Array2<Cflt> {
let (n, nch) = timedata.dim(); let (n, nch) = timedata.dim();
let nfft = self.nfft(); let nfft = self.nfft();
@ -76,9 +85,7 @@ impl PowerSpectra {
self.freqdata self.freqdata
.push_column(Ccol::from_vec(vec![Cflt::new(0., 0.); nfft / 2 + 1]).view()) .push_column(Ccol::from_vec(vec![Cflt::new(0., 0.); nfft / 2 + 1]).view())
.unwrap(); .unwrap();
self.timedata self.timedata.push_column(Dcol::zeros(nfft).view()).unwrap();
.push_column(Dcol::zeros(nfft).view())
.unwrap();
} }
assert!(n == self.nfft()); assert!(n == self.nfft());
@ -106,24 +113,38 @@ impl PowerSpectra {
/// Compute cross power spectra from input time data. First axis is /// Compute cross power spectra from input time data. First axis is
/// frequency, second axis is channel i, third axis is channel j. /// 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 /// * `tdata` - Input time data. This is a 2D array, where the first axis is
/// time and the second axis is the channel number. /// time and the second axis is the channel number.
pub fn compute<'a, T>(&mut self, tdata: T) -> Array3<Cflt> ///
/// # 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 where
T: Into<ArrayView<'a, Flt, Ix2>>, T: AsArray<'a, Flt, Ix2>,
{ {
let tdata = tdata.into(); let tdata = tdata.into();
let clen = self.nfft() / 2 + 1; let nfft = self.nfft();
let nchannel = tdata.ncols(); let clen = nfft / 2 + 1;
let win_pwr = self.sqrt_win_pwr; 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 // Compute fft of input data, and store in self.freqdata
let fd = self.compute_ffts(tdata); let fd = self.compute_ffts(tdata);
let fdconj = fd.mapv(|c| c.conj()); 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<Cflt> = unsafe { result.assume_init() }; let mut result: Array3<Cflt> = unsafe { result.assume_init() };
// Loop over result axis one and channel i IN PARALLEL // Loop over result axis one and channel i IN PARALLEL
@ -140,18 +161,16 @@ impl PowerSpectra {
Zip::from(&mut out) Zip::from(&mut out)
.and(chi) .and(chi)
.and(chj) .and(chj)
.for_each(|out, chi, chjc|{ .for_each(|out, chi, chjc| {
// Loop over frequency components // Loop over frequency components
*out = 0.5 * chi * chjc; *out = 0.5 * chi * chjc;
} });
);
// The DC component has no 0.5 correction, as it only // The DC component has no 0.5 correction, as it only
// occurs ones in a (double-sided) power spectrum. So // occurs ones in a (double-sided) power spectrum. So
// here we undo the 0.5 of 4 lines above here. // here we undo the 0.5 of 4 lines above here.
out[0] *= 2.; out[0] *= 2.;
out[clen-1] *= 2.; out[clen - 1] *= 2.;
}); });
}); });
result result
@ -166,16 +185,16 @@ mod test {
use rand_distr::StandardNormal; use rand_distr::StandardNormal;
/// Generate a sine wave at the order i /// Generate a sine wave at the order i
fn generate_sinewave(nfft: usize,order: usize) -> Dcol { fn generate_sinewave(nfft: usize, order: usize) -> Dcol {
Dcol::from_iter((0..nfft).map(|i| { Dcol::from_iter(
Flt::sin(i as Flt/(nfft) as Flt * order as Flt * 2.*pi) (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 /// Generate a sine wave at the order i
fn generate_cosinewave(nfft: usize,order: usize) -> Dcol { fn generate_cosinewave(nfft: usize, order: usize) -> Dcol {
Dcol::from_iter((0..nfft).map(|i| { Dcol::from_iter(
Flt::cos(i as Flt/(nfft) as Flt * order as Flt * 2.*pi) (0..nfft).map(|i| Flt::cos(i as Flt / (nfft) as Flt * order as Flt * 2. * pi)),
})) )
} }
use super::*; use super::*;
@ -206,8 +225,7 @@ mod test {
// Start with a time signal // Start with a time signal
let mut t: Dmat = Dmat::default((nfft, 0)); let mut t: Dmat = Dmat::default((nfft, 0));
t.push_column(generate_sinewave(nfft,1).view()) t.push_column(generate_sinewave(nfft, 1).view()).unwrap();
.unwrap();
// println!("{:?}", t); // println!("{:?}", t);
let fd = ps.compute_ffts(t.view()); let fd = ps.compute_ffts(t.view());
@ -227,37 +245,44 @@ mod test {
); );
} }
/// Thest whether power spectra scale properly. Signals with amplitude of 1 /// 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 /// should come back with a power of 0.5. DC offsets should come in as
/// value^2 at frequency index 0. /// value^2 at frequency index 0.
#[test] #[test]
fn test_ps_scale() { fn test_ps_scale() {
const nfft: usize = 124; const nfft: usize = 124;
let rect = Window::new(WindowType::Rect, nfft); let rect = Window::new(WindowType::Rect, nfft);
let mut ps = PowerSpectra::newFromWindow(rect); let mut ps = PowerSpectra::newFromWindow(rect);
// Start with a time signal // Start with a time signal
let mut t: Dmat = Dmat::default((nfft, 0)); let mut t: Dmat = Dmat::default((nfft, 0));
t.push_column(generate_cosinewave(nfft,1).view()) t.push_column(generate_cosinewave(nfft, 1).view()).unwrap();
.unwrap();
let dc_component = 0.25; let dc_component = 0.25;
let dc_power = dc_component.pow(2); let dc_power = dc_component.pow(2);
t.mapv_inplace(|t| t + dc_component); t.mapv_inplace(|t| t + dc_component);
let power = ps.compute(t.view()); 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!(
assert_relative_eq!(power[(1, 0,0)].re, 0.5, epsilon = Flt::EPSILON * nfft as Flt); power[(0, 0, 0)].re,
assert_relative_eq!(power[(1, 0,0)].im, 0.0, epsilon = Flt::EPSILON * nfft as Flt); 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; use ndarray_rand::RandomExt;
// Test parseval's theorem for some random data // Test parseval's theorem for some random data
#[test] #[test]
fn test_parseval() { fn test_parseval() {
const nfft: usize = 512; const nfft: usize = 512;
let rect = Window::new(WindowType::Rect, nfft); let rect = Window::new(WindowType::Rect, nfft);
let mut ps = PowerSpectra::newFromWindow(rect); let mut ps = PowerSpectra::newFromWindow(rect);
@ -265,11 +290,11 @@ mod test {
// Start with a time signal // Start with a time signal
let t: Dmat = Dmat::random((nfft, 1), StandardNormal); 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); let t_dc_power = tavg.powi(2);
// println!("dc power in time domain: {:?}", t_dc_power); // 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); // println!("Total signal power in time domain: {:?} ", signal_pwr);
let power = ps.compute(t.view()); let power = ps.compute(t.view());
@ -277,15 +302,21 @@ mod test {
let fpower = power.sum().abs(); 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!(
assert_ulps_eq!(signal_pwr, fpower, epsilon = Flt::EPSILON * (nfft as Flt).powi(2)); 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 parseval's theorem for some random data
#[test] #[test]
fn test_parseval_with_window() { fn test_parseval_with_window() {
// A sufficiently high value is required here, to show that it works. // A sufficiently high value is required here, to show that it works.
const nfft: usize = 2usize.pow(20); const nfft: usize = 2usize.pow(20);
let window = Window::new(WindowType::Hann, nfft); let window = Window::new(WindowType::Hann, nfft);
@ -293,13 +324,13 @@ mod test {
let mut ps = PowerSpectra::newFromWindow(window); let mut ps = PowerSpectra::newFromWindow(window);
// Start with a time signal // 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); let t_dc_power = tavg.powi(2);
// println!("dc power in time domain: {:?}", t_dc_power); // 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); // println!("Total signal power in time domain: {:?} ", signal_pwr);
let power = ps.compute(t.view()); let power = ps.compute(t.view());
@ -307,11 +338,13 @@ mod test {
let fpower = power.sum().abs(); 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. // 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);
} }
} }