Updated comments and test code

src/ps/aps.rs

@@ -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 crate::config::*;
 use anyhow::{bail, Result};

src/ps/mod.rs

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

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.
+/// 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 {
-    // 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
+            self.timedata.push_column(Dcol::zeros(nfft).view()).unwrap();
-                .push_column(Dcol::zeros(nfft).view())
-                .unwrap();
         }
 
         assert!(n == self.nfft());
@@ -106,24 +113,38 @@ 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 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
         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
@@ -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 {
+    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
-    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::*;
@@ -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())
+        t.push_column(generate_sinewave(nfft, 1).view()).unwrap();
-            .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())
+        t.push_column(generate_cosinewave(nfft, 1).view()).unwrap();
-            .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!(
-        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;
     // 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!(
-        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]
     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);
-
     }
-
 }