ASCEE/lasprs
5
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Updated comments and test code

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
src/ps/aps.rs
View File

@ -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};

19
src/ps/mod.rs
View File

@ -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<Cflt>;
pub use aps::{ApsMode, AvPowerSpectra, Overlap, ApsResult};
pub use ps::PowerSpectra;
pub use window::{Window, WindowType};
pub use window::{Window, WindowType};

159
src/ps/ps.rs
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 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<FFT>,
@ -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<Flt>) -> &Array2<Cflt> {
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<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
T: Into<ArrayView<'a, Flt, Ix2>>,
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<Cflt> = 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);
}
}