756 lines
22 KiB
Rust
756 lines
22 KiB
Rust
//! Time-Frequency Representations (TFR)
|
||
//!
|
||
//! This module provides time-frequency analysis methods for MEG/EEG data:
|
||
//!
|
||
//! - **Morlet Wavelets**: Complex Morlet wavelet transform
|
||
//! - **STFT**: Short-Time Fourier Transform with configurable windows
|
||
//! - **PSD**: Power Spectral Density (Welch, Multitaper)
|
||
//!
|
||
//! # Example
|
||
//!
|
||
//! ```ignore
|
||
//! use rtx_neuro_signal::tfr::{tfr_morlet, TfrOutput};
|
||
//!
|
||
//! let data = vec![0.0; 1000]; // [n_channels x n_samples]
|
||
//! let sfreq = 256.0;
|
||
//! let freqs = vec![8.0, 10.0, 12.0, 14.0]; // Alpha band
|
||
//!
|
||
//! let tfr = tfr_morlet(&[data], sfreq, &freqs, 7.0, TfrOutput::Power)?;
|
||
//! ```
|
||
|
||
use num_complex::Complex64;
|
||
use rayon::prelude::*;
|
||
use rustfft::{Fft, FftPlanner};
|
||
use std::f64::consts::PI;
|
||
use std::sync::Arc;
|
||
|
||
use crate::{SignalError, SignalResult};
|
||
|
||
/// Output type for time-frequency representations
|
||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||
pub enum TfrOutput {
|
||
/// Power (magnitude squared)
|
||
Power,
|
||
/// Phase angle (radians)
|
||
Phase,
|
||
/// Complex values
|
||
Complex,
|
||
/// Inter-trial coherence (requires multiple trials)
|
||
Itc,
|
||
}
|
||
|
||
/// Cycle specification for wavelets
|
||
#[derive(Debug, Clone)]
|
||
pub enum CycleSpec {
|
||
/// Fixed number of cycles for all frequencies
|
||
Fixed(f64),
|
||
/// Frequency-adaptive cycles (min at low freq, max at high freq)
|
||
Adaptive {
|
||
/// Minimum cycles at lowest frequency
|
||
min: f64,
|
||
/// Maximum cycles at highest frequency
|
||
max: f64,
|
||
},
|
||
}
|
||
|
||
impl CycleSpec {
|
||
/// Get number of cycles for a given frequency
|
||
pub fn get_cycles(&self, freq: f64, freq_min: f64, freq_max: f64) -> f64 {
|
||
match self {
|
||
CycleSpec::Fixed(n) => *n,
|
||
CycleSpec::Adaptive { min, max } => {
|
||
let t = (freq - freq_min) / (freq_max - freq_min + 1e-10);
|
||
min + t * (max - min)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Result of time-frequency analysis
|
||
#[derive(Debug, Clone)]
|
||
pub struct TfrResult {
|
||
/// TFR data: [n_channels x n_freqs x n_times] for single trial
|
||
/// or [n_epochs x n_channels x n_freqs x n_times] for epochs
|
||
pub data: Vec<Vec<Vec<f64>>>,
|
||
/// Frequencies analyzed
|
||
pub freqs: Vec<f64>,
|
||
/// Time points (relative to epoch start)
|
||
pub times: Vec<f64>,
|
||
/// Sampling frequency
|
||
pub sfreq: f64,
|
||
/// Output type
|
||
pub output: TfrOutput,
|
||
}
|
||
|
||
/// Generate a complex Morlet wavelet
|
||
///
|
||
/// The Morlet wavelet is defined as:
|
||
/// w(t) = A * exp(-t²/(2σ²)) * exp(2πif₀t)
|
||
///
|
||
/// where σ = n_cycles / (2πf₀)
|
||
fn morlet_wavelet(freq: f64, sfreq: f64, n_cycles: f64) -> Vec<Complex64> {
|
||
// Gaussian standard deviation in samples
|
||
let sigma_t = n_cycles / (2.0 * PI * freq);
|
||
let sigma_samples = sigma_t * sfreq;
|
||
|
||
// Wavelet length: 5 sigma on each side
|
||
let half_len = (5.0 * sigma_samples).ceil() as usize;
|
||
let len = 2 * half_len + 1;
|
||
|
||
let mut wavelet = Vec::with_capacity(len);
|
||
let norm = 1.0 / (sigma_t * (2.0 * PI).sqrt()).sqrt();
|
||
|
||
for i in 0..len {
|
||
let t = (i as f64 - half_len as f64) / sfreq;
|
||
let gaussian = (-t * t / (2.0 * sigma_t * sigma_t)).exp();
|
||
let oscillation = Complex64::new(0.0, 2.0 * PI * freq * t).exp();
|
||
wavelet.push(norm * gaussian * oscillation);
|
||
}
|
||
|
||
wavelet
|
||
}
|
||
|
||
/// Compute Morlet wavelet time-frequency representation
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `data` - Input data [n_channels][n_samples]
|
||
/// * `sfreq` - Sampling frequency in Hz
|
||
/// * `freqs` - Frequencies to analyze
|
||
/// * `n_cycles` - Number of wavelet cycles (typically 7)
|
||
/// * `output` - Output type (Power, Phase, Complex)
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// TfrResult with data [n_channels x n_freqs x n_times]
|
||
pub fn tfr_morlet(
|
||
data: &[Vec<f64>],
|
||
sfreq: f64,
|
||
freqs: &[f64],
|
||
n_cycles: f64,
|
||
output: TfrOutput,
|
||
) -> SignalResult<TfrResult> {
|
||
tfr_morlet_adaptive(data, sfreq, freqs, CycleSpec::Fixed(n_cycles), output, 1)
|
||
}
|
||
|
||
/// Compute Morlet wavelet TFR with adaptive cycles and decimation
|
||
pub fn tfr_morlet_adaptive(
|
||
data: &[Vec<f64>],
|
||
sfreq: f64,
|
||
freqs: &[f64],
|
||
cycles: CycleSpec,
|
||
output: TfrOutput,
|
||
decim: usize,
|
||
) -> SignalResult<TfrResult> {
|
||
if data.is_empty() || data[0].is_empty() {
|
||
return Err(SignalError::InvalidLength("Empty data".to_string()));
|
||
}
|
||
if freqs.is_empty() {
|
||
return Err(SignalError::InvalidParameters(
|
||
"No frequencies specified".to_string(),
|
||
));
|
||
}
|
||
|
||
let n_samples = data[0].len();
|
||
let n_freqs = freqs.len();
|
||
let freq_min = freqs.iter().copied().fold(f64::INFINITY, f64::min);
|
||
let freq_max = freqs.iter().copied().fold(f64::NEG_INFINITY, f64::max);
|
||
|
||
// Decimated output length
|
||
let n_times = n_samples.div_ceil(decim);
|
||
|
||
// Generate wavelets for each frequency
|
||
let wavelets: Vec<Vec<Complex64>> = freqs
|
||
.iter()
|
||
.map(|&f| {
|
||
let nc = cycles.get_cycles(f, freq_min, freq_max);
|
||
morlet_wavelet(f, sfreq, nc)
|
||
})
|
||
.collect();
|
||
|
||
// Process each channel in parallel
|
||
let result: Vec<Vec<Vec<f64>>> = data
|
||
.par_iter()
|
||
.map(|channel| {
|
||
let mut freq_results = Vec::with_capacity(n_freqs);
|
||
|
||
for wavelet in &wavelets {
|
||
// Convolve channel with wavelet
|
||
let convolved = convolve_complex(channel, wavelet);
|
||
|
||
// Decimate and compute output
|
||
let mut times_result = Vec::with_capacity(n_times);
|
||
for t in (0..n_samples).step_by(decim) {
|
||
let c = convolved[t + wavelet.len() / 2];
|
||
let val = match output {
|
||
TfrOutput::Power => c.norm_sqr(),
|
||
TfrOutput::Phase => c.arg(),
|
||
TfrOutput::Complex => c.re, // Store real part; use separate fn for complex
|
||
TfrOutput::Itc => c.norm(), // For single trial, just amplitude
|
||
};
|
||
times_result.push(val);
|
||
}
|
||
freq_results.push(times_result);
|
||
}
|
||
|
||
freq_results
|
||
})
|
||
.collect();
|
||
|
||
// Generate time vector
|
||
let times: Vec<f64> = (0..n_times).map(|t| (t * decim) as f64 / sfreq).collect();
|
||
|
||
Ok(TfrResult {
|
||
data: result,
|
||
freqs: freqs.to_vec(),
|
||
times,
|
||
sfreq,
|
||
output,
|
||
})
|
||
}
|
||
|
||
/// Convolve signal with complex wavelet
|
||
fn convolve_complex(signal: &[f64], wavelet: &[Complex64]) -> Vec<Complex64> {
|
||
let n = signal.len();
|
||
let w_len = wavelet.len();
|
||
let out_len = n + w_len - 1;
|
||
|
||
// Use FFT for efficient convolution
|
||
let fft_len = out_len.next_power_of_two();
|
||
|
||
let mut planner = FftPlanner::new();
|
||
let fft = planner.plan_fft_forward(fft_len);
|
||
let ifft = planner.plan_fft_inverse(fft_len);
|
||
|
||
// Zero-pad signal
|
||
let mut signal_fft: Vec<Complex64> = signal
|
||
.iter()
|
||
.map(|&x| Complex64::new(x, 0.0))
|
||
.chain(std::iter::repeat(Complex64::new(0.0, 0.0)))
|
||
.take(fft_len)
|
||
.collect();
|
||
|
||
// Zero-pad wavelet
|
||
let mut wavelet_fft: Vec<Complex64> = wavelet
|
||
.iter()
|
||
.copied()
|
||
.chain(std::iter::repeat(Complex64::new(0.0, 0.0)))
|
||
.take(fft_len)
|
||
.collect();
|
||
|
||
// Forward FFT
|
||
fft.process(&mut signal_fft);
|
||
fft.process(&mut wavelet_fft);
|
||
|
||
// Multiply in frequency domain
|
||
for i in 0..fft_len {
|
||
signal_fft[i] *= wavelet_fft[i];
|
||
}
|
||
|
||
// Inverse FFT
|
||
ifft.process(&mut signal_fft);
|
||
|
||
// Normalize and take valid portion
|
||
let scale = 1.0 / fft_len as f64;
|
||
signal_fft
|
||
.into_iter()
|
||
.take(out_len)
|
||
.map(|c| c * scale)
|
||
.collect()
|
||
}
|
||
|
||
/// Window function types for STFT
|
||
#[derive(Debug, Clone, Copy)]
|
||
pub enum Window {
|
||
/// Rectangular (no windowing)
|
||
Rectangular,
|
||
/// Hann window
|
||
Hann,
|
||
/// Hamming window
|
||
Hamming,
|
||
/// Blackman window
|
||
Blackman,
|
||
}
|
||
|
||
impl Window {
|
||
/// Generate window function of given length
|
||
pub fn generate(&self, len: usize) -> Vec<f64> {
|
||
match self {
|
||
Window::Rectangular => vec![1.0; len],
|
||
Window::Hann => (0..len)
|
||
.map(|i| 0.5 * (1.0 - (2.0 * PI * i as f64 / (len - 1) as f64).cos()))
|
||
.collect(),
|
||
Window::Hamming => (0..len)
|
||
.map(|i| 0.54 - 0.46 * (2.0 * PI * i as f64 / (len - 1) as f64).cos())
|
||
.collect(),
|
||
Window::Blackman => (0..len)
|
||
.map(|i| {
|
||
let a0 = 0.42;
|
||
let a1 = 0.5;
|
||
let a2 = 0.08;
|
||
let t = 2.0 * PI * i as f64 / (len - 1) as f64;
|
||
a0 - a1 * t.cos() + a2 * (2.0 * t).cos()
|
||
})
|
||
.collect(),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Result of STFT computation
|
||
#[derive(Debug, Clone)]
|
||
pub struct StftResult {
|
||
/// Complex STFT data: [n_channels x n_freqs x n_frames]
|
||
pub data: Vec<Vec<Vec<Complex64>>>,
|
||
/// Frequency bins
|
||
pub freqs: Vec<f64>,
|
||
/// Time points (frame centers)
|
||
pub times: Vec<f64>,
|
||
/// Sampling frequency
|
||
pub sfreq: f64,
|
||
}
|
||
|
||
impl StftResult {
|
||
/// Get power spectrogram (magnitude squared)
|
||
pub fn power(&self) -> Vec<Vec<Vec<f64>>> {
|
||
self.data
|
||
.iter()
|
||
.map(|ch| {
|
||
ch.iter()
|
||
.map(|f| f.iter().map(num_complex::Complex::norm_sqr).collect())
|
||
.collect()
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// Get phase spectrogram
|
||
pub fn phase(&self) -> Vec<Vec<Vec<f64>>> {
|
||
self.data
|
||
.iter()
|
||
.map(|ch| {
|
||
ch.iter()
|
||
.map(|f| f.iter().map(|c| c.arg()).collect())
|
||
.collect()
|
||
})
|
||
.collect()
|
||
}
|
||
}
|
||
|
||
/// Compute Short-Time Fourier Transform
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `data` - Input data [n_channels][n_samples]
|
||
/// * `sfreq` - Sampling frequency
|
||
/// * `n_fft` - FFT size
|
||
/// * `hop_length` - Hop length between frames (default: n_fft / 4)
|
||
/// * `window` - Window function
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// StftResult with complex STFT data
|
||
pub fn stft(
|
||
data: &[Vec<f64>],
|
||
sfreq: f64,
|
||
n_fft: usize,
|
||
hop_length: Option<usize>,
|
||
window: Window,
|
||
) -> SignalResult<StftResult> {
|
||
if data.is_empty() || data[0].is_empty() {
|
||
return Err(SignalError::InvalidLength("Empty data".to_string()));
|
||
}
|
||
if n_fft == 0 || !n_fft.is_power_of_two() {
|
||
return Err(SignalError::InvalidParameters(
|
||
"n_fft must be a power of 2".to_string(),
|
||
));
|
||
}
|
||
|
||
let n_samples = data[0].len();
|
||
let hop = hop_length.unwrap_or(n_fft / 4);
|
||
let win = window.generate(n_fft);
|
||
|
||
// Number of frames
|
||
let n_frames = (n_samples.saturating_sub(n_fft)) / hop + 1;
|
||
let n_freqs = n_fft / 2 + 1; // One-sided spectrum
|
||
|
||
let mut planner = FftPlanner::new();
|
||
let fft: Arc<dyn Fft<f64>> = planner.plan_fft_forward(n_fft);
|
||
|
||
// Process each channel in parallel
|
||
let result: Vec<Vec<Vec<Complex64>>> = data
|
||
.par_iter()
|
||
.map(|channel| {
|
||
let mut freq_data = vec![vec![Complex64::new(0.0, 0.0); n_frames]; n_freqs];
|
||
|
||
for frame_idx in 0..n_frames {
|
||
let start = frame_idx * hop;
|
||
|
||
// Apply window
|
||
let mut frame: Vec<Complex64> = (0..n_fft)
|
||
.map(|i| {
|
||
let sample = if start + i < channel.len() {
|
||
channel[start + i]
|
||
} else {
|
||
0.0
|
||
};
|
||
Complex64::new(sample * win[i], 0.0)
|
||
})
|
||
.collect();
|
||
|
||
// FFT
|
||
fft.process(&mut frame);
|
||
|
||
// Store one-sided spectrum
|
||
for (f_idx, freq_bin) in frame.iter().take(n_freqs).enumerate() {
|
||
freq_data[f_idx][frame_idx] = *freq_bin;
|
||
}
|
||
}
|
||
|
||
freq_data
|
||
})
|
||
.collect();
|
||
|
||
// Generate frequency and time vectors
|
||
let freqs: Vec<f64> = (0..n_freqs)
|
||
.map(|i| i as f64 * sfreq / n_fft as f64)
|
||
.collect();
|
||
let times: Vec<f64> = (0..n_frames)
|
||
.map(|i| (i * hop + n_fft / 2) as f64 / sfreq)
|
||
.collect();
|
||
|
||
Ok(StftResult {
|
||
data: result,
|
||
freqs,
|
||
times,
|
||
sfreq,
|
||
})
|
||
}
|
||
|
||
/// Result of PSD computation
|
||
#[derive(Debug, Clone)]
|
||
pub struct PsdResult {
|
||
/// Power spectral density [n_channels x n_freqs]
|
||
pub psd: Vec<Vec<f64>>,
|
||
/// Frequency bins
|
||
pub freqs: Vec<f64>,
|
||
}
|
||
|
||
/// Compute Power Spectral Density using Welch's method
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `data` - Input data [n_channels][n_samples]
|
||
/// * `sfreq` - Sampling frequency
|
||
/// * `n_fft` - FFT size (default: 256)
|
||
/// * `n_overlap` - Overlap between segments (default: n_fft / 2)
|
||
/// * `window` - Window function
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// PsdResult with PSD values
|
||
pub fn psd_welch(
|
||
data: &[Vec<f64>],
|
||
sfreq: f64,
|
||
n_fft: Option<usize>,
|
||
n_overlap: Option<usize>,
|
||
window: Window,
|
||
) -> SignalResult<PsdResult> {
|
||
let n_fft = n_fft.unwrap_or(256);
|
||
let n_overlap = n_overlap.unwrap_or(n_fft / 2);
|
||
let hop = n_fft - n_overlap;
|
||
|
||
// Compute STFT
|
||
let stft_result = stft(data, sfreq, n_fft, Some(hop), window)?;
|
||
|
||
// Average power across time
|
||
let psd: Vec<Vec<f64>> = stft_result
|
||
.data
|
||
.iter()
|
||
.map(|ch| {
|
||
ch.iter()
|
||
.map(|freq_frames| {
|
||
let n_frames = freq_frames.len();
|
||
if n_frames == 0 {
|
||
return 0.0;
|
||
}
|
||
let sum: f64 = freq_frames.iter().map(num_complex::Complex::norm_sqr).sum();
|
||
sum / n_frames as f64
|
||
})
|
||
.collect()
|
||
})
|
||
.collect();
|
||
|
||
Ok(PsdResult {
|
||
psd,
|
||
freqs: stft_result.freqs,
|
||
})
|
||
}
|
||
|
||
/// Generate DPSS (Discrete Prolate Spheroidal Sequences) tapers
|
||
///
|
||
/// Simplified implementation using Slepian sequences approximation
|
||
fn dpss_tapers(n: usize, _bandwidth: f64, n_tapers: usize) -> Vec<Vec<f64>> {
|
||
// Simplified: use sine tapers as approximation
|
||
// Full DPSS would require eigenvalue computation
|
||
let mut tapers = Vec::with_capacity(n_tapers);
|
||
|
||
for k in 0..n_tapers {
|
||
let taper: Vec<f64> = (0..n)
|
||
.map(|i| {
|
||
let t = (i as f64 + 0.5) / n as f64;
|
||
((k + 1) as f64 * PI * t).sin() * (2.0 / n as f64).sqrt()
|
||
})
|
||
.collect();
|
||
tapers.push(taper);
|
||
}
|
||
|
||
tapers
|
||
}
|
||
|
||
/// Compute Power Spectral Density using multitaper method
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `data` - Input data [n_channels][n_samples]
|
||
/// * `sfreq` - Sampling frequency
|
||
/// * `bandwidth` - Frequency bandwidth (Hz)
|
||
/// * `fmin` - Minimum frequency (optional)
|
||
/// * `fmax` - Maximum frequency (optional)
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// PsdResult with PSD values
|
||
pub fn psd_multitaper(
|
||
data: &[Vec<f64>],
|
||
sfreq: f64,
|
||
bandwidth: f64,
|
||
fmin: Option<f64>,
|
||
fmax: Option<f64>,
|
||
) -> SignalResult<PsdResult> {
|
||
if data.is_empty() || data[0].is_empty() {
|
||
return Err(SignalError::InvalidLength("Empty data".to_string()));
|
||
}
|
||
|
||
let n_samples = data[0].len();
|
||
let n_fft = n_samples.next_power_of_two();
|
||
let n_freqs = n_fft / 2 + 1;
|
||
|
||
// Calculate number of tapers
|
||
let nw = bandwidth * n_samples as f64 / sfreq;
|
||
let n_tapers = ((2.0 * nw - 1.0).floor() as usize).max(1);
|
||
|
||
// Generate tapers
|
||
let tapers = dpss_tapers(n_samples, bandwidth, n_tapers);
|
||
|
||
let mut planner = FftPlanner::new();
|
||
let fft: Arc<dyn Fft<f64>> = planner.plan_fft_forward(n_fft);
|
||
|
||
// Process each channel
|
||
let psd: Vec<Vec<f64>> = data
|
||
.par_iter()
|
||
.map(|channel| {
|
||
let mut spectrum = vec![0.0; n_freqs];
|
||
|
||
// Average across tapers
|
||
for taper in &tapers {
|
||
// Apply taper
|
||
let mut tapered: Vec<Complex64> = channel
|
||
.iter()
|
||
.zip(taper.iter())
|
||
.map(|(&x, &t)| Complex64::new(x * t, 0.0))
|
||
.chain(std::iter::repeat(Complex64::new(0.0, 0.0)))
|
||
.take(n_fft)
|
||
.collect();
|
||
|
||
// FFT
|
||
fft.process(&mut tapered);
|
||
|
||
// Accumulate power
|
||
for (i, c) in tapered.iter().take(n_freqs).enumerate() {
|
||
spectrum[i] += c.norm_sqr();
|
||
}
|
||
}
|
||
|
||
// Average and normalize
|
||
for val in &mut spectrum {
|
||
*val /= n_tapers as f64;
|
||
*val /= sfreq; // Normalize to PSD units
|
||
}
|
||
|
||
spectrum
|
||
})
|
||
.collect();
|
||
|
||
// Generate frequencies
|
||
let freqs: Vec<f64> = (0..n_freqs)
|
||
.map(|i| i as f64 * sfreq / n_fft as f64)
|
||
.collect();
|
||
|
||
// Apply frequency limits
|
||
let fmin = fmin.unwrap_or(0.0);
|
||
let fmax = fmax.unwrap_or(sfreq / 2.0);
|
||
|
||
let (psd_filtered, freqs_filtered): (Vec<Vec<f64>>, Vec<f64>) = {
|
||
let mask: Vec<bool> = freqs.iter().map(|&f| f >= fmin && f <= fmax).collect();
|
||
|
||
let freqs_f: Vec<f64> = freqs
|
||
.iter()
|
||
.zip(&mask)
|
||
.filter_map(|(&f, &m)| if m { Some(f) } else { None })
|
||
.collect();
|
||
|
||
let psd_f: Vec<Vec<f64>> = psd
|
||
.iter()
|
||
.map(|ch| {
|
||
ch.iter()
|
||
.zip(&mask)
|
||
.filter_map(|(&p, &m)| if m { Some(p) } else { None })
|
||
.collect()
|
||
})
|
||
.collect();
|
||
|
||
(psd_f, freqs_f)
|
||
};
|
||
|
||
Ok(PsdResult {
|
||
psd: psd_filtered,
|
||
freqs: freqs_filtered,
|
||
})
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use approx::assert_relative_eq;
|
||
|
||
fn generate_sine(freq: f64, sfreq: f64, duration: f64) -> Vec<f64> {
|
||
let n_samples = (duration * sfreq) as usize;
|
||
(0..n_samples)
|
||
.map(|i| (2.0 * PI * freq * i as f64 / sfreq).sin())
|
||
.collect()
|
||
}
|
||
|
||
#[test]
|
||
fn test_morlet_wavelet() {
|
||
let wavelet = morlet_wavelet(10.0, 256.0, 7.0);
|
||
assert!(!wavelet.is_empty());
|
||
// Wavelet should be roughly symmetric
|
||
let mid = wavelet.len() / 2;
|
||
assert_relative_eq!(wavelet[mid].norm(), wavelet[mid].norm(), epsilon = 1e-10);
|
||
}
|
||
|
||
#[test]
|
||
fn test_tfr_morlet_sine() {
|
||
let sfreq = 256.0;
|
||
let signal = generate_sine(10.0, sfreq, 2.0);
|
||
let data = vec![signal];
|
||
let freqs = vec![5.0, 10.0, 15.0, 20.0];
|
||
|
||
let result = tfr_morlet(&data, sfreq, &freqs, 7.0, TfrOutput::Power).unwrap();
|
||
|
||
assert_eq!(result.data.len(), 1); // 1 channel
|
||
assert_eq!(result.data[0].len(), 4); // 4 frequencies
|
||
assert_eq!(result.freqs.len(), 4);
|
||
|
||
// 10 Hz should have highest power
|
||
let powers: Vec<f64> = result.data[0]
|
||
.iter()
|
||
.map(|f| f.iter().sum::<f64>() / f.len() as f64)
|
||
.collect();
|
||
|
||
let max_idx = powers
|
||
.iter()
|
||
.enumerate()
|
||
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
|
||
.map(|(i, _)| i)
|
||
.unwrap();
|
||
|
||
assert_eq!(max_idx, 1); // Index 1 = 10 Hz
|
||
}
|
||
|
||
#[test]
|
||
fn test_stft() {
|
||
let sfreq = 256.0;
|
||
let signal = generate_sine(20.0, sfreq, 1.0);
|
||
let data = vec![signal];
|
||
|
||
let result = stft(&data, sfreq, 64, None, Window::Hann).unwrap();
|
||
|
||
assert_eq!(result.data.len(), 1);
|
||
assert!(!result.freqs.is_empty());
|
||
assert!(!result.times.is_empty());
|
||
|
||
// Check power has peak at 20 Hz
|
||
let power = result.power();
|
||
let mean_power: Vec<f64> = power[0]
|
||
.iter()
|
||
.map(|f| f.iter().sum::<f64>() / f.len() as f64)
|
||
.collect();
|
||
|
||
let max_idx = mean_power
|
||
.iter()
|
||
.enumerate()
|
||
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
|
||
.map(|(i, _)| i)
|
||
.unwrap();
|
||
|
||
let peak_freq = result.freqs[max_idx];
|
||
assert!(peak_freq >= 18.0 && peak_freq <= 22.0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_psd_welch() {
|
||
let sfreq = 256.0;
|
||
let signal = generate_sine(30.0, sfreq, 2.0);
|
||
let data = vec![signal];
|
||
|
||
let result = psd_welch(&data, sfreq, Some(128), None, Window::Hann).unwrap();
|
||
|
||
assert_eq!(result.psd.len(), 1);
|
||
assert!(!result.freqs.is_empty());
|
||
|
||
// Find peak frequency
|
||
let max_idx = result.psd[0]
|
||
.iter()
|
||
.enumerate()
|
||
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
|
||
.map(|(i, _)| i)
|
||
.unwrap();
|
||
|
||
let peak_freq = result.freqs[max_idx];
|
||
assert!(peak_freq >= 28.0 && peak_freq <= 32.0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_psd_multitaper() {
|
||
let sfreq = 256.0;
|
||
let signal = generate_sine(15.0, sfreq, 2.0);
|
||
let data = vec![signal];
|
||
|
||
let result = psd_multitaper(&data, sfreq, 4.0, Some(1.0), Some(50.0)).unwrap();
|
||
|
||
assert_eq!(result.psd.len(), 1);
|
||
assert!(!result.freqs.is_empty());
|
||
|
||
// All frequencies should be in range
|
||
assert!(result.freqs.iter().all(|&f| f >= 1.0 && f <= 50.0));
|
||
}
|
||
|
||
#[test]
|
||
fn test_window_functions() {
|
||
let len = 64;
|
||
|
||
let hann = Window::Hann.generate(len);
|
||
assert_eq!(hann.len(), len);
|
||
assert_relative_eq!(hann[0], 0.0, epsilon = 1e-10);
|
||
// Hann window maximum is near 1.0 but not exactly due to discrete sampling
|
||
assert!(hann[len / 2] > 0.99);
|
||
|
||
let hamming = Window::Hamming.generate(len);
|
||
assert_eq!(hamming.len(), len);
|
||
assert!(hamming[0] > 0.0); // Hamming doesn't go to 0
|
||
|
||
let blackman = Window::Blackman.generate(len);
|
||
assert_eq!(blackman.len(), len);
|
||
}
|
||
}
|