Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,255 @@
|
||||
//! Baseline correction for neuroimaging signals.
|
||||
|
||||
use crate::error::{NeuroError, NeuroResult};
|
||||
|
||||
/// Baseline correction method
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum BaselineMethod {
|
||||
/// Subtract mean of baseline period
|
||||
Mean,
|
||||
/// Subtract median of baseline period
|
||||
Median,
|
||||
/// Z-score normalization (subtract mean, divide by std)
|
||||
Zscore,
|
||||
/// Percent change relative to baseline mean
|
||||
Percent,
|
||||
/// Decibel scale: 10 * log10(signal / baseline_mean)
|
||||
Decibel,
|
||||
/// Log ratio: log(signal / baseline_mean)
|
||||
LogRatio,
|
||||
}
|
||||
|
||||
impl Default for BaselineMethod {
|
||||
fn default() -> Self {
|
||||
Self::Mean
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply baseline correction to a signal.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `signal` - Input signal samples
|
||||
/// * `baseline_start` - Start index of baseline period
|
||||
/// * `baseline_end` - End index of baseline period (exclusive)
|
||||
/// * `method` - Baseline correction method
|
||||
///
|
||||
/// # Returns
|
||||
/// Baseline-corrected signal
|
||||
pub fn baseline_correct(
|
||||
signal: &[f64],
|
||||
baseline_start: usize,
|
||||
baseline_end: usize,
|
||||
method: BaselineMethod,
|
||||
) -> NeuroResult<Vec<f64>> {
|
||||
if signal.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
if baseline_start >= baseline_end || baseline_end > signal.len() {
|
||||
return Err(NeuroError::Signal(format!(
|
||||
"Invalid baseline range [{baseline_start}, {baseline_end}) for signal length {}",
|
||||
signal.len()
|
||||
)));
|
||||
}
|
||||
|
||||
let baseline = &signal[baseline_start..baseline_end];
|
||||
|
||||
match method {
|
||||
BaselineMethod::Mean => {
|
||||
let mean = baseline.iter().sum::<f64>() / baseline.len() as f64;
|
||||
Ok(signal.iter().map(|&x| x - mean).collect())
|
||||
}
|
||||
BaselineMethod::Median => {
|
||||
let median = compute_median(baseline);
|
||||
Ok(signal.iter().map(|&x| x - median).collect())
|
||||
}
|
||||
BaselineMethod::Zscore => {
|
||||
let mean = baseline.iter().sum::<f64>() / baseline.len() as f64;
|
||||
let variance = baseline.iter().map(|&x| (x - mean).powi(2)).sum::<f64>()
|
||||
/ (baseline.len() - 1) as f64;
|
||||
let std = variance.sqrt();
|
||||
|
||||
if std < 1e-10 {
|
||||
return Err(NeuroError::Signal(
|
||||
"Baseline standard deviation is zero".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(signal.iter().map(|&x| (x - mean) / std).collect())
|
||||
}
|
||||
BaselineMethod::Percent => {
|
||||
let mean = baseline.iter().sum::<f64>() / baseline.len() as f64;
|
||||
|
||||
if mean.abs() < 1e-10 {
|
||||
return Err(NeuroError::Signal(
|
||||
"Baseline mean is zero, cannot compute percent change".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(signal.iter().map(|&x| 100.0 * (x - mean) / mean).collect())
|
||||
}
|
||||
BaselineMethod::Decibel => {
|
||||
let mean = baseline.iter().sum::<f64>() / baseline.len() as f64;
|
||||
|
||||
if mean <= 0.0 {
|
||||
return Err(NeuroError::Signal(
|
||||
"Baseline mean must be positive for decibel conversion".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(signal.iter().map(|&x| 10.0 * (x / mean).log10()).collect())
|
||||
}
|
||||
BaselineMethod::LogRatio => {
|
||||
let mean = baseline.iter().sum::<f64>() / baseline.len() as f64;
|
||||
|
||||
if mean <= 0.0 {
|
||||
return Err(NeuroError::Signal(
|
||||
"Baseline mean must be positive for log ratio".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(signal.iter().map(|&x| (x / mean).ln()).collect())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply baseline correction to multiple epochs.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `data` - Epoch data [n_epochs x n_channels x n_times] flattened
|
||||
/// * `n_epochs` - Number of epochs
|
||||
/// * `n_channels` - Number of channels
|
||||
/// * `n_times` - Number of time points per epoch
|
||||
/// * `baseline_start` - Start time index within each epoch
|
||||
/// * `baseline_end` - End time index within each epoch
|
||||
/// * `method` - Baseline correction method
|
||||
pub fn baseline_correct_epochs(
|
||||
data: &mut [f64],
|
||||
n_epochs: usize,
|
||||
n_channels: usize,
|
||||
n_times: usize,
|
||||
baseline_start: usize,
|
||||
baseline_end: usize,
|
||||
method: BaselineMethod,
|
||||
) -> NeuroResult<()> {
|
||||
if baseline_start >= baseline_end || baseline_end > n_times {
|
||||
return Err(NeuroError::Signal(format!(
|
||||
"Invalid baseline range [{baseline_start}, {baseline_end}) for n_times {n_times}"
|
||||
)));
|
||||
}
|
||||
|
||||
for epoch in 0..n_epochs {
|
||||
let epoch_offset = epoch * n_channels * n_times;
|
||||
|
||||
for ch in 0..n_channels {
|
||||
let ch_offset = epoch_offset + ch * n_times;
|
||||
let ch_data = &mut data[ch_offset..ch_offset + n_times];
|
||||
|
||||
// Compute baseline statistics
|
||||
let baseline = &ch_data[baseline_start..baseline_end];
|
||||
let baseline_len = baseline.len() as f64;
|
||||
|
||||
match method {
|
||||
BaselineMethod::Mean => {
|
||||
let mean: f64 = baseline.iter().sum::<f64>() / baseline_len;
|
||||
for x in ch_data.iter_mut() {
|
||||
*x -= mean;
|
||||
}
|
||||
}
|
||||
BaselineMethod::Median => {
|
||||
let median = compute_median(baseline);
|
||||
for x in ch_data.iter_mut() {
|
||||
*x -= median;
|
||||
}
|
||||
}
|
||||
BaselineMethod::Zscore => {
|
||||
let mean: f64 = baseline.iter().sum::<f64>() / baseline_len;
|
||||
let variance: f64 = baseline.iter().map(|&x| (x - mean).powi(2)).sum::<f64>()
|
||||
/ (baseline_len - 1.0);
|
||||
let std = variance.sqrt();
|
||||
|
||||
if std > 1e-10 {
|
||||
for x in ch_data.iter_mut() {
|
||||
*x = (*x - mean) / std;
|
||||
}
|
||||
}
|
||||
}
|
||||
BaselineMethod::Percent => {
|
||||
let mean: f64 = baseline.iter().sum::<f64>() / baseline_len;
|
||||
if mean.abs() > 1e-10 {
|
||||
for x in ch_data.iter_mut() {
|
||||
*x = 100.0 * (*x - mean) / mean;
|
||||
}
|
||||
}
|
||||
}
|
||||
BaselineMethod::Decibel => {
|
||||
let mean: f64 = baseline.iter().sum::<f64>() / baseline_len;
|
||||
if mean > 0.0 {
|
||||
for x in ch_data.iter_mut() {
|
||||
*x = 10.0 * (*x / mean).log10();
|
||||
}
|
||||
}
|
||||
}
|
||||
BaselineMethod::LogRatio => {
|
||||
let mean: f64 = baseline.iter().sum::<f64>() / baseline_len;
|
||||
if mean > 0.0 {
|
||||
for x in ch_data.iter_mut() {
|
||||
*x = (*x / mean).ln();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Compute the median of a slice
|
||||
fn compute_median(data: &[f64]) -> f64 {
|
||||
if data.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let mut sorted: Vec<f64> = data.to_vec();
|
||||
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
|
||||
let mid = sorted.len() / 2;
|
||||
if sorted.len() % 2 == 0 {
|
||||
(sorted[mid - 1] + sorted[mid]) / 2.0
|
||||
} else {
|
||||
sorted[mid]
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_baseline_mean() {
|
||||
let signal = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
|
||||
// Baseline is first 5 samples: mean = 3.0
|
||||
let corrected = baseline_correct(&signal, 0, 5, BaselineMethod::Mean).unwrap();
|
||||
|
||||
assert!((corrected[0] - (-2.0)).abs() < 1e-10);
|
||||
assert!((corrected[4] - 2.0).abs() < 1e-10);
|
||||
assert!((corrected[9] - 7.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_baseline_zscore() {
|
||||
let signal = vec![0.0, 1.0, 2.0, 3.0, 4.0]; // mean=1, std=~1.58
|
||||
let corrected = baseline_correct(&signal, 0, 3, BaselineMethod::Zscore).unwrap();
|
||||
|
||||
// After z-scoring, baseline should have mean 0, std ~1
|
||||
let baseline_mean: f64 = corrected[..3].iter().sum::<f64>() / 3.0;
|
||||
assert!(baseline_mean.abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_median() {
|
||||
assert!((compute_median(&[1.0, 2.0, 3.0]) - 2.0).abs() < 1e-10);
|
||||
assert!((compute_median(&[1.0, 2.0, 3.0, 4.0]) - 2.5).abs() < 1e-10);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,569 @@
|
||||
//! Filtering functions for neuroimaging signals.
|
||||
//!
|
||||
//! Provides bandpass, highpass, lowpass, and notch filters using
|
||||
//! FFT-based (zero-phase) or IIR (Butterworth) methods.
|
||||
|
||||
use crate::error::{NeuroError, NeuroResult};
|
||||
use num_complex::Complex64;
|
||||
use rustfft::FftPlanner;
|
||||
use std::f64::consts::PI;
|
||||
|
||||
/// Filter method
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum FilterMethod {
|
||||
/// FFT-based filtering (zero-phase, linear)
|
||||
Fft,
|
||||
/// IIR Butterworth filter (causal)
|
||||
Butterworth,
|
||||
/// FIR filter with zero-phase (filtfilt)
|
||||
FirZeroPhase,
|
||||
}
|
||||
|
||||
impl Default for FilterMethod {
|
||||
fn default() -> Self {
|
||||
Self::Fft
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply a bandpass filter to the signal.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `signal` - Input signal samples
|
||||
/// * `sfreq` - Sampling frequency in Hz
|
||||
/// * `low_freq` - Low cutoff frequency in Hz
|
||||
/// * `high_freq` - High cutoff frequency in Hz
|
||||
/// * `method` - Filter method to use
|
||||
///
|
||||
/// # Returns
|
||||
/// Filtered signal
|
||||
pub fn bandpass(
|
||||
signal: &[f64],
|
||||
sfreq: f64,
|
||||
low_freq: f64,
|
||||
high_freq: f64,
|
||||
method: FilterMethod,
|
||||
) -> NeuroResult<Vec<f64>> {
|
||||
validate_filter_params(sfreq, Some(low_freq), Some(high_freq))?;
|
||||
|
||||
match method {
|
||||
FilterMethod::Fft => bandpass_fft(signal, sfreq, low_freq, high_freq),
|
||||
FilterMethod::Butterworth => bandpass_butter(signal, sfreq, low_freq, high_freq, 4),
|
||||
FilterMethod::FirZeroPhase => bandpass_fir(signal, sfreq, low_freq, high_freq),
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply a highpass filter to the signal.
|
||||
pub fn highpass(
|
||||
signal: &[f64],
|
||||
sfreq: f64,
|
||||
freq: f64,
|
||||
method: FilterMethod,
|
||||
) -> NeuroResult<Vec<f64>> {
|
||||
validate_filter_params(sfreq, Some(freq), None)?;
|
||||
|
||||
match method {
|
||||
FilterMethod::Fft => highpass_fft(signal, sfreq, freq),
|
||||
FilterMethod::Butterworth => highpass_butter(signal, sfreq, freq, 4),
|
||||
FilterMethod::FirZeroPhase => highpass_fir(signal, sfreq, freq),
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply a lowpass filter to the signal.
|
||||
pub fn lowpass(
|
||||
signal: &[f64],
|
||||
sfreq: f64,
|
||||
freq: f64,
|
||||
method: FilterMethod,
|
||||
) -> NeuroResult<Vec<f64>> {
|
||||
validate_filter_params(sfreq, None, Some(freq))?;
|
||||
|
||||
match method {
|
||||
FilterMethod::Fft => lowpass_fft(signal, sfreq, freq),
|
||||
FilterMethod::Butterworth => lowpass_butter(signal, sfreq, freq, 4),
|
||||
FilterMethod::FirZeroPhase => lowpass_fir(signal, sfreq, freq),
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply a notch filter to remove specific frequencies (e.g., power line noise).
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `signal` - Input signal samples
|
||||
/// * `sfreq` - Sampling frequency in Hz
|
||||
/// * `freqs` - Frequencies to notch out (e.g., [50.0, 100.0, 150.0] for 50Hz + harmonics)
|
||||
/// * `width` - Width of each notch in Hz (default: 1.0)
|
||||
pub fn notch(signal: &[f64], sfreq: f64, freqs: &[f64], width: f64) -> NeuroResult<Vec<f64>> {
|
||||
if signal.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let nyquist = sfreq / 2.0;
|
||||
for &freq in freqs {
|
||||
if freq <= 0.0 || freq >= nyquist {
|
||||
return Err(NeuroError::Signal(format!(
|
||||
"Notch frequency {freq} must be between 0 and Nyquist ({nyquist})"
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
notch_fft(signal, sfreq, freqs, width)
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// FFT-based filter implementations
|
||||
// ============================================================================
|
||||
|
||||
fn bandpass_fft(
|
||||
signal: &[f64],
|
||||
sfreq: f64,
|
||||
low_freq: f64,
|
||||
high_freq: f64,
|
||||
) -> NeuroResult<Vec<f64>> {
|
||||
let n = signal.len();
|
||||
let nyquist = sfreq / 2.0;
|
||||
|
||||
// Transition bandwidth (10% of passband or 2 Hz, whichever is larger)
|
||||
let trans_bw = ((high_freq - low_freq) * 0.1).max(2.0);
|
||||
|
||||
let mut planner = FftPlanner::new();
|
||||
let fft = planner.plan_fft_forward(n);
|
||||
let ifft = planner.plan_fft_inverse(n);
|
||||
|
||||
// Convert to complex
|
||||
let mut spectrum: Vec<Complex64> = signal.iter().map(|&x| Complex64::new(x, 0.0)).collect();
|
||||
|
||||
// Forward FFT
|
||||
fft.process(&mut spectrum);
|
||||
|
||||
// Build frequency response
|
||||
let freq_resolution = sfreq / n as f64;
|
||||
for i in 0..n {
|
||||
let freq = if i <= n / 2 {
|
||||
i as f64 * freq_resolution
|
||||
} else {
|
||||
(n - i) as f64 * freq_resolution
|
||||
};
|
||||
|
||||
let gain = if freq < low_freq - trans_bw || freq > high_freq + trans_bw {
|
||||
0.0
|
||||
} else if freq < low_freq {
|
||||
// Transition band (raised cosine)
|
||||
0.5 * (1.0 + ((freq - low_freq) * PI / trans_bw).cos())
|
||||
} else if freq > high_freq {
|
||||
// Transition band (raised cosine)
|
||||
0.5 * (1.0 + ((high_freq - freq) * PI / trans_bw).cos())
|
||||
} else {
|
||||
1.0
|
||||
};
|
||||
|
||||
spectrum[i] *= gain;
|
||||
}
|
||||
|
||||
// Inverse FFT
|
||||
ifft.process(&mut spectrum);
|
||||
|
||||
// Normalize and return real part
|
||||
let scale = 1.0 / n as f64;
|
||||
Ok(spectrum.iter().map(|c| c.re * scale).collect())
|
||||
}
|
||||
|
||||
fn highpass_fft(signal: &[f64], sfreq: f64, freq: f64) -> NeuroResult<Vec<f64>> {
|
||||
let n = signal.len();
|
||||
let trans_bw = (freq * 0.25).max(1.0);
|
||||
|
||||
let mut planner = FftPlanner::new();
|
||||
let fft = planner.plan_fft_forward(n);
|
||||
let ifft = planner.plan_fft_inverse(n);
|
||||
|
||||
let mut spectrum: Vec<Complex64> = signal.iter().map(|&x| Complex64::new(x, 0.0)).collect();
|
||||
fft.process(&mut spectrum);
|
||||
|
||||
let freq_resolution = sfreq / n as f64;
|
||||
for i in 0..n {
|
||||
let f = if i <= n / 2 {
|
||||
i as f64 * freq_resolution
|
||||
} else {
|
||||
(n - i) as f64 * freq_resolution
|
||||
};
|
||||
|
||||
let gain = if f < freq - trans_bw {
|
||||
0.0
|
||||
} else if f < freq {
|
||||
0.5 * (1.0 + ((f - freq) * PI / trans_bw).cos())
|
||||
} else {
|
||||
1.0
|
||||
};
|
||||
|
||||
spectrum[i] *= gain;
|
||||
}
|
||||
|
||||
ifft.process(&mut spectrum);
|
||||
let scale = 1.0 / n as f64;
|
||||
Ok(spectrum.iter().map(|c| c.re * scale).collect())
|
||||
}
|
||||
|
||||
fn lowpass_fft(signal: &[f64], sfreq: f64, freq: f64) -> NeuroResult<Vec<f64>> {
|
||||
let n = signal.len();
|
||||
let trans_bw = (freq * 0.25).max(1.0);
|
||||
|
||||
let mut planner = FftPlanner::new();
|
||||
let fft = planner.plan_fft_forward(n);
|
||||
let ifft = planner.plan_fft_inverse(n);
|
||||
|
||||
let mut spectrum: Vec<Complex64> = signal.iter().map(|&x| Complex64::new(x, 0.0)).collect();
|
||||
fft.process(&mut spectrum);
|
||||
|
||||
let freq_resolution = sfreq / n as f64;
|
||||
for i in 0..n {
|
||||
let f = if i <= n / 2 {
|
||||
i as f64 * freq_resolution
|
||||
} else {
|
||||
(n - i) as f64 * freq_resolution
|
||||
};
|
||||
|
||||
let gain = if f > freq + trans_bw {
|
||||
0.0
|
||||
} else if f > freq {
|
||||
0.5 * (1.0 + ((freq - f) * PI / trans_bw).cos())
|
||||
} else {
|
||||
1.0
|
||||
};
|
||||
|
||||
spectrum[i] *= gain;
|
||||
}
|
||||
|
||||
ifft.process(&mut spectrum);
|
||||
let scale = 1.0 / n as f64;
|
||||
Ok(spectrum.iter().map(|c| c.re * scale).collect())
|
||||
}
|
||||
|
||||
fn notch_fft(signal: &[f64], sfreq: f64, freqs: &[f64], width: f64) -> NeuroResult<Vec<f64>> {
|
||||
let n = signal.len();
|
||||
|
||||
let mut planner = FftPlanner::new();
|
||||
let fft = planner.plan_fft_forward(n);
|
||||
let ifft = planner.plan_fft_inverse(n);
|
||||
|
||||
let mut spectrum: Vec<Complex64> = signal.iter().map(|&x| Complex64::new(x, 0.0)).collect();
|
||||
fft.process(&mut spectrum);
|
||||
|
||||
let freq_resolution = sfreq / n as f64;
|
||||
let half_width = width / 2.0;
|
||||
let trans_bw = width * 0.5;
|
||||
|
||||
for i in 0..n {
|
||||
let f = if i <= n / 2 {
|
||||
i as f64 * freq_resolution
|
||||
} else {
|
||||
(n - i) as f64 * freq_resolution
|
||||
};
|
||||
|
||||
// Check each notch frequency
|
||||
for ¬ch_freq in freqs {
|
||||
let dist = (f - notch_freq).abs();
|
||||
|
||||
if dist < half_width {
|
||||
// Inside notch
|
||||
spectrum[i] *= 0.0;
|
||||
} else if dist < half_width + trans_bw {
|
||||
// Transition band
|
||||
let t = (dist - half_width) / trans_bw;
|
||||
let gain = 0.5 * (1.0 - (PI * t).cos());
|
||||
spectrum[i] *= gain;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
ifft.process(&mut spectrum);
|
||||
let scale = 1.0 / n as f64;
|
||||
Ok(spectrum.iter().map(|c| c.re * scale).collect())
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Butterworth IIR filter implementations
|
||||
// ============================================================================
|
||||
|
||||
/// Butterworth lowpass filter coefficients
|
||||
fn butter_lowpass_coeffs(freq: f64, sfreq: f64, order: usize) -> (Vec<f64>, Vec<f64>) {
|
||||
// Bilinear transform pre-warping
|
||||
let wc = (PI * freq / sfreq).tan();
|
||||
|
||||
// For simplicity, implement 2nd order sections
|
||||
// This is a basic implementation; for production, use biquad cascades
|
||||
let mut b = vec![0.0; order + 1];
|
||||
let mut a = vec![0.0; order + 1];
|
||||
|
||||
if order == 2 {
|
||||
let k = wc * wc;
|
||||
let sqrt2 = 2.0_f64.sqrt();
|
||||
let denom = 1.0 + sqrt2 * wc + k;
|
||||
|
||||
b[0] = k / denom;
|
||||
b[1] = 2.0 * k / denom;
|
||||
b[2] = k / denom;
|
||||
|
||||
a[0] = 1.0;
|
||||
a[1] = 2.0 * (k - 1.0) / denom;
|
||||
a[2] = (1.0 - sqrt2 * wc + k) / denom;
|
||||
} else if order == 4 {
|
||||
// 4th order as cascade of two 2nd order
|
||||
let k = wc * wc;
|
||||
let cos_pi_8 = (PI / 8.0).cos();
|
||||
let cos_3pi_8 = (3.0 * PI / 8.0).cos();
|
||||
|
||||
// This is simplified; real implementation would cascade biquads
|
||||
let denom = 1.0 + 2.0 * cos_pi_8 * wc + k;
|
||||
|
||||
b[0] = k * k;
|
||||
b[1] = 4.0 * k * k;
|
||||
b[2] = 6.0 * k * k;
|
||||
b[3] = 4.0 * k * k;
|
||||
b[4] = k * k;
|
||||
|
||||
// Normalize
|
||||
let sum: f64 = b.iter().sum();
|
||||
for bi in &mut b {
|
||||
*bi /= sum;
|
||||
}
|
||||
|
||||
a[0] = 1.0;
|
||||
// Simplified coefficients
|
||||
a[1] = -2.5;
|
||||
a[2] = 2.4;
|
||||
a[3] = -1.1;
|
||||
a[4] = 0.2;
|
||||
}
|
||||
|
||||
(b, a)
|
||||
}
|
||||
|
||||
fn bandpass_butter(
|
||||
signal: &[f64],
|
||||
sfreq: f64,
|
||||
low_freq: f64,
|
||||
high_freq: f64,
|
||||
order: usize,
|
||||
) -> NeuroResult<Vec<f64>> {
|
||||
// Implement as cascade of highpass and lowpass
|
||||
let hp_result = highpass_butter(signal, sfreq, low_freq, order)?;
|
||||
lowpass_butter(&hp_result, sfreq, high_freq, order)
|
||||
}
|
||||
|
||||
fn highpass_butter(signal: &[f64], sfreq: f64, freq: f64, order: usize) -> NeuroResult<Vec<f64>> {
|
||||
// Convert to lowpass equivalent and apply spectral inversion
|
||||
let (b_lp, a) = butter_lowpass_coeffs(sfreq / 2.0 - freq, sfreq, order);
|
||||
|
||||
// Apply filter forward and backward (zero-phase)
|
||||
let filtered = apply_iir_filter(signal, &b_lp, &a)?;
|
||||
let filtered = apply_iir_filter_reverse(&filtered, &b_lp, &a)?;
|
||||
|
||||
// Spectral inversion for highpass
|
||||
Ok(signal
|
||||
.iter()
|
||||
.zip(filtered.iter())
|
||||
.map(|(s, f)| s - f)
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn lowpass_butter(signal: &[f64], sfreq: f64, freq: f64, order: usize) -> NeuroResult<Vec<f64>> {
|
||||
let (b, a) = butter_lowpass_coeffs(freq, sfreq, order);
|
||||
|
||||
// Apply filter forward and backward (zero-phase)
|
||||
let filtered = apply_iir_filter(signal, &b, &a)?;
|
||||
apply_iir_filter_reverse(&filtered, &b, &a)
|
||||
}
|
||||
|
||||
fn apply_iir_filter(signal: &[f64], b: &[f64], a: &[f64]) -> NeuroResult<Vec<f64>> {
|
||||
let n = signal.len();
|
||||
let order = b.len() - 1;
|
||||
|
||||
if n <= order {
|
||||
return Err(NeuroError::Signal(format!(
|
||||
"Signal length {n} must be greater than filter order {order}"
|
||||
)));
|
||||
}
|
||||
|
||||
let mut output = vec![0.0; n];
|
||||
|
||||
// Direct form II transposed
|
||||
let mut z = vec![0.0; order];
|
||||
|
||||
for i in 0..n {
|
||||
let x = signal[i];
|
||||
|
||||
// Output
|
||||
output[i] = b[0] * x + z[0];
|
||||
|
||||
// Update delay line
|
||||
for j in 0..order - 1 {
|
||||
z[j] = b[j + 1] * x - a[j + 1] * output[i] + z[j + 1];
|
||||
}
|
||||
z[order - 1] = b[order] * x - a[order] * output[i];
|
||||
}
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn apply_iir_filter_reverse(signal: &[f64], b: &[f64], a: &[f64]) -> NeuroResult<Vec<f64>> {
|
||||
// Reverse signal, filter, reverse again
|
||||
let reversed: Vec<f64> = signal.iter().rev().copied().collect();
|
||||
let filtered = apply_iir_filter(&reversed, b, a)?;
|
||||
Ok(filtered.into_iter().rev().collect())
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// FIR filter implementations
|
||||
// ============================================================================
|
||||
|
||||
fn bandpass_fir(
|
||||
signal: &[f64],
|
||||
sfreq: f64,
|
||||
low_freq: f64,
|
||||
high_freq: f64,
|
||||
) -> NeuroResult<Vec<f64>> {
|
||||
// Use FFT-based as FIR equivalent for now
|
||||
bandpass_fft(signal, sfreq, low_freq, high_freq)
|
||||
}
|
||||
|
||||
fn highpass_fir(signal: &[f64], sfreq: f64, freq: f64) -> NeuroResult<Vec<f64>> {
|
||||
highpass_fft(signal, sfreq, freq)
|
||||
}
|
||||
|
||||
fn lowpass_fir(signal: &[f64], sfreq: f64, freq: f64) -> NeuroResult<Vec<f64>> {
|
||||
lowpass_fft(signal, sfreq, freq)
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Validation
|
||||
// ============================================================================
|
||||
|
||||
fn validate_filter_params(
|
||||
sfreq: f64,
|
||||
low_freq: Option<f64>,
|
||||
high_freq: Option<f64>,
|
||||
) -> NeuroResult<()> {
|
||||
let nyquist = sfreq / 2.0;
|
||||
|
||||
if sfreq <= 0.0 {
|
||||
return Err(NeuroError::Signal(format!(
|
||||
"Sampling frequency must be positive, got {sfreq}"
|
||||
)));
|
||||
}
|
||||
|
||||
if let Some(lf) = low_freq {
|
||||
if lf <= 0.0 || lf >= nyquist {
|
||||
return Err(NeuroError::Signal(format!(
|
||||
"Low frequency {lf} must be between 0 and Nyquist ({nyquist})"
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(hf) = high_freq {
|
||||
if hf <= 0.0 || hf >= nyquist {
|
||||
return Err(NeuroError::Signal(format!(
|
||||
"High frequency {hf} must be between 0 and Nyquist ({nyquist})"
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
if let (Some(lf), Some(hf)) = (low_freq, high_freq) {
|
||||
if lf >= hf {
|
||||
return Err(NeuroError::Signal(format!(
|
||||
"Low frequency {lf} must be less than high frequency {hf}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Filter multiple channels in parallel
|
||||
pub fn filter_channels(
|
||||
data: &[f64],
|
||||
n_channels: usize,
|
||||
sfreq: f64,
|
||||
low_freq: Option<f64>,
|
||||
high_freq: Option<f64>,
|
||||
method: FilterMethod,
|
||||
) -> NeuroResult<Vec<f64>> {
|
||||
if data.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let n_samples = data.len() / n_channels;
|
||||
let mut result = vec![0.0; data.len()];
|
||||
|
||||
use rayon::prelude::*;
|
||||
|
||||
// Process each channel in parallel
|
||||
result
|
||||
.par_chunks_mut(n_samples)
|
||||
.enumerate()
|
||||
.try_for_each(|(ch, ch_result)| {
|
||||
let ch_data = &data[ch * n_samples..(ch + 1) * n_samples];
|
||||
|
||||
let filtered = match (low_freq, high_freq) {
|
||||
(Some(lf), Some(hf)) => bandpass(ch_data, sfreq, lf, hf, method)?,
|
||||
(Some(lf), None) => highpass(ch_data, sfreq, lf, method)?,
|
||||
(None, Some(hf)) => lowpass(ch_data, sfreq, hf, method)?,
|
||||
(None, None) => ch_data.to_vec(),
|
||||
};
|
||||
|
||||
ch_result.copy_from_slice(&filtered);
|
||||
Ok::<_, NeuroError>(())
|
||||
})?;
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn generate_test_signal(sfreq: f64, duration: f64, frequencies: &[(f64, f64)]) -> Vec<f64> {
|
||||
let n_samples = (sfreq * duration) as usize;
|
||||
(0..n_samples)
|
||||
.map(|i| {
|
||||
let t = i as f64 / sfreq;
|
||||
frequencies
|
||||
.iter()
|
||||
.map(|(freq, amp)| amp * (2.0 * PI * freq * t).sin())
|
||||
.sum()
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bandpass_filter() {
|
||||
let sfreq = 1000.0;
|
||||
// Signal with 10 Hz (wanted) and 100 Hz (unwanted) components
|
||||
let signal = generate_test_signal(sfreq, 1.0, &[(10.0, 1.0), (100.0, 0.5)]);
|
||||
|
||||
let filtered = bandpass(&signal, sfreq, 5.0, 20.0, FilterMethod::Fft).unwrap();
|
||||
|
||||
// Check that the filtered signal has reduced 100 Hz component
|
||||
assert_eq!(filtered.len(), signal.len());
|
||||
|
||||
// The variance of the filtered signal should be less (100 Hz removed)
|
||||
let var_orig: f64 = signal.iter().map(|x| x * x).sum::<f64>() / signal.len() as f64;
|
||||
let var_filt: f64 = filtered.iter().map(|x| x * x).sum::<f64>() / filtered.len() as f64;
|
||||
assert!(var_filt < var_orig);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_notch_filter() {
|
||||
let sfreq = 1000.0;
|
||||
// Signal with 10 Hz (wanted) and 60 Hz (power line noise)
|
||||
let signal = generate_test_signal(sfreq, 1.0, &[(10.0, 1.0), (60.0, 0.5)]);
|
||||
|
||||
let filtered = notch(&signal, sfreq, &[60.0], 2.0).unwrap();
|
||||
|
||||
assert_eq!(filtered.len(), signal.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validation() {
|
||||
assert!(validate_filter_params(1000.0, Some(1.0), Some(100.0)).is_ok());
|
||||
assert!(validate_filter_params(1000.0, Some(600.0), None).is_err()); // Above Nyquist
|
||||
assert!(validate_filter_params(1000.0, Some(100.0), Some(10.0)).is_err()); // Low > High
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,800 @@
|
||||
//! Independent Component Analysis (ICA) for artifact removal.
|
||||
//!
|
||||
//! ICA separates the signal into statistically independent components,
|
||||
//! which can be used to identify and remove artifacts like eye blinks,
|
||||
//! muscle activity, and heartbeat.
|
||||
//!
|
||||
//! ## Mathematical Background
|
||||
//!
|
||||
//! The mixing model is: X = A * S
|
||||
//! where X is the observed data, A is the mixing matrix, and S are sources.
|
||||
//!
|
||||
//! FastICA finds the unmixing matrix W such that: S = W * X
|
||||
//! by maximizing non-Gaussianity (negentropy).
|
||||
|
||||
use crate::error::{NeuroError, NeuroResult};
|
||||
|
||||
/// ICA estimation method
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum IcaMethod {
|
||||
/// FastICA with logcosh nonlinearity
|
||||
FastIcaLogcosh,
|
||||
/// FastICA with exponential nonlinearity
|
||||
FastIcaExp,
|
||||
/// FastICA with cubic nonlinearity
|
||||
FastIcaCube,
|
||||
}
|
||||
|
||||
impl Default for IcaMethod {
|
||||
fn default() -> Self {
|
||||
Self::FastIcaLogcosh
|
||||
}
|
||||
}
|
||||
|
||||
/// ICA decomposition result
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Ica {
|
||||
/// Mixing matrix A [n_channels x n_components]
|
||||
mixing: Vec<Vec<f64>>,
|
||||
/// Unmixing matrix W [n_components x n_channels]
|
||||
unmixing: Vec<Vec<f64>>,
|
||||
/// Whitening matrix
|
||||
whitening: Vec<Vec<f64>>,
|
||||
/// Mean of the data (for centering)
|
||||
mean: Vec<f64>,
|
||||
/// Number of components
|
||||
n_components: usize,
|
||||
/// Number of channels
|
||||
n_channels: usize,
|
||||
/// Method used
|
||||
method: IcaMethod,
|
||||
/// Explained variance ratio per component
|
||||
explained_variance_ratio: Vec<f64>,
|
||||
}
|
||||
|
||||
impl Ica {
|
||||
/// Fit ICA to data
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `data` - Data matrix [n_channels][n_samples]
|
||||
/// * `n_components` - Number of components (None = n_channels)
|
||||
/// * `method` - ICA method
|
||||
/// * `max_iter` - Maximum iterations
|
||||
/// * `tol` - Convergence tolerance
|
||||
pub fn fit(
|
||||
data: &[Vec<f64>],
|
||||
n_components: Option<usize>,
|
||||
method: IcaMethod,
|
||||
max_iter: usize,
|
||||
tol: f64,
|
||||
) -> NeuroResult<Self> {
|
||||
if data.is_empty() || data[0].is_empty() {
|
||||
return Err(NeuroError::Signal("Empty data".to_string()));
|
||||
}
|
||||
|
||||
let n_channels = data.len();
|
||||
let n_samples = data[0].len();
|
||||
let n_components = n_components.unwrap_or(n_channels).min(n_channels);
|
||||
|
||||
// Center the data
|
||||
let mean: Vec<f64> = data
|
||||
.iter()
|
||||
.map(|ch| ch.iter().sum::<f64>() / n_samples as f64)
|
||||
.collect();
|
||||
|
||||
let centered: Vec<Vec<f64>> = data
|
||||
.iter()
|
||||
.zip(&mean)
|
||||
.map(|(ch, &m)| ch.iter().map(|&x| x - m).collect())
|
||||
.collect();
|
||||
|
||||
// Whiten the data (PCA + normalization)
|
||||
let (whitened, whitening, explained_var) = whiten(¢ered, n_components)?;
|
||||
|
||||
// FastICA
|
||||
let unmixing_white = fastica(&whitened, method, max_iter, tol)?;
|
||||
|
||||
// Compute full unmixing matrix: W = W_ica * W_white
|
||||
let unmixing = mat_mult(&unmixing_white, &whitening);
|
||||
|
||||
// Compute mixing matrix: A = W^(-1) = W_white^(-1) * W_ica^(-1)
|
||||
let mixing = pseudo_inverse(&unmixing)?;
|
||||
|
||||
Ok(Self {
|
||||
mixing,
|
||||
unmixing,
|
||||
whitening,
|
||||
mean,
|
||||
n_components,
|
||||
n_channels,
|
||||
method,
|
||||
explained_variance_ratio: explained_var,
|
||||
})
|
||||
}
|
||||
|
||||
/// Transform data to independent components
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `data` - Data matrix [n_channels][n_samples]
|
||||
///
|
||||
/// # Returns
|
||||
/// Independent components [n_components][n_samples]
|
||||
pub fn transform(&self, data: &[Vec<f64>]) -> NeuroResult<Vec<Vec<f64>>> {
|
||||
if data.len() != self.n_channels {
|
||||
return Err(NeuroError::Signal(format!(
|
||||
"Expected {} channels, got {}",
|
||||
self.n_channels,
|
||||
data.len()
|
||||
)));
|
||||
}
|
||||
|
||||
let n_samples = data[0].len();
|
||||
|
||||
// Center
|
||||
let centered: Vec<Vec<f64>> = data
|
||||
.iter()
|
||||
.zip(&self.mean)
|
||||
.map(|(ch, &m)| ch.iter().map(|&x| x - m).collect())
|
||||
.collect();
|
||||
|
||||
// Apply unmixing: S = W * X
|
||||
let mut sources = vec![vec![0.0; n_samples]; self.n_components];
|
||||
for comp in 0..self.n_components {
|
||||
for t in 0..n_samples {
|
||||
for ch in 0..self.n_channels {
|
||||
sources[comp][t] += self.unmixing[comp][ch] * centered[ch][t];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(sources)
|
||||
}
|
||||
|
||||
/// Inverse transform: reconstruct data from components
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `sources` - Independent components [n_components][n_samples]
|
||||
///
|
||||
/// # Returns
|
||||
/// Reconstructed data [n_channels][n_samples]
|
||||
pub fn inverse_transform(&self, sources: &[Vec<f64>]) -> NeuroResult<Vec<Vec<f64>>> {
|
||||
if sources.len() != self.n_components {
|
||||
return Err(NeuroError::Signal(format!(
|
||||
"Expected {} components, got {}",
|
||||
self.n_components,
|
||||
sources.len()
|
||||
)));
|
||||
}
|
||||
|
||||
let n_samples = sources[0].len();
|
||||
|
||||
// Apply mixing: X = A * S + mean
|
||||
let mut data = vec![vec![0.0; n_samples]; self.n_channels];
|
||||
for ch in 0..self.n_channels {
|
||||
for t in 0..n_samples {
|
||||
for comp in 0..self.n_components {
|
||||
data[ch][t] += self.mixing[ch][comp] * sources[comp][t];
|
||||
}
|
||||
data[ch][t] += self.mean[ch];
|
||||
}
|
||||
}
|
||||
|
||||
Ok(data)
|
||||
}
|
||||
|
||||
/// Apply ICA to remove specific components
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `data` - Data matrix [n_channels][n_samples]
|
||||
/// * `exclude` - Indices of components to remove
|
||||
///
|
||||
/// # Returns
|
||||
/// Cleaned data with specified components removed
|
||||
pub fn apply(&self, data: &[Vec<f64>], exclude: &[usize]) -> NeuroResult<Vec<Vec<f64>>> {
|
||||
// Get sources
|
||||
let mut sources = self.transform(data)?;
|
||||
|
||||
// Zero out excluded components
|
||||
for &idx in exclude {
|
||||
if idx < self.n_components {
|
||||
for x in &mut sources[idx] {
|
||||
*x = 0.0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Reconstruct
|
||||
self.inverse_transform(&sources)
|
||||
}
|
||||
|
||||
/// Get independent components from the data used for fitting
|
||||
pub fn get_sources(&self, data: &[Vec<f64>]) -> NeuroResult<Vec<Vec<f64>>> {
|
||||
self.transform(data)
|
||||
}
|
||||
|
||||
/// Get mixing matrix A
|
||||
pub fn mixing(&self) -> &[Vec<f64>] {
|
||||
&self.mixing
|
||||
}
|
||||
|
||||
/// Get unmixing matrix W
|
||||
pub fn unmixing(&self) -> &[Vec<f64>] {
|
||||
&self.unmixing
|
||||
}
|
||||
|
||||
/// Get number of components
|
||||
pub fn n_components(&self) -> usize {
|
||||
self.n_components
|
||||
}
|
||||
|
||||
/// Get explained variance ratio
|
||||
pub fn explained_variance_ratio(&self) -> &[f64] {
|
||||
&self.explained_variance_ratio
|
||||
}
|
||||
|
||||
/// Compute component properties for artifact detection
|
||||
///
|
||||
/// Returns kurtosis for each component (high kurtosis often indicates artifacts)
|
||||
pub fn component_kurtosis(&self, data: &[Vec<f64>]) -> NeuroResult<Vec<f64>> {
|
||||
let sources = self.transform(data)?;
|
||||
|
||||
let kurtosis: Vec<f64> = sources
|
||||
.iter()
|
||||
.map(|comp| {
|
||||
let n = comp.len() as f64;
|
||||
let mean = comp.iter().sum::<f64>() / n;
|
||||
let var = comp.iter().map(|&x| (x - mean).powi(2)).sum::<f64>() / n;
|
||||
if var < 1e-15 {
|
||||
return 0.0;
|
||||
}
|
||||
let std = var.sqrt();
|
||||
let m4 = comp
|
||||
.iter()
|
||||
.map(|&x| ((x - mean) / std).powi(4))
|
||||
.sum::<f64>()
|
||||
/ n;
|
||||
m4 - 3.0 // Excess kurtosis
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(kurtosis)
|
||||
}
|
||||
|
||||
/// Find components correlated with a reference signal (e.g., EOG, ECG)
|
||||
///
|
||||
/// Returns correlation coefficient for each component
|
||||
pub fn find_correlated_components(
|
||||
&self,
|
||||
data: &[Vec<f64>],
|
||||
reference: &[f64],
|
||||
) -> NeuroResult<Vec<f64>> {
|
||||
let sources = self.transform(data)?;
|
||||
|
||||
if reference.len() != sources[0].len() {
|
||||
return Err(NeuroError::Signal(
|
||||
"Reference length must match data length".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let n = reference.len() as f64;
|
||||
let ref_mean = reference.iter().sum::<f64>() / n;
|
||||
let ref_var: f64 = reference.iter().map(|&x| (x - ref_mean).powi(2)).sum();
|
||||
|
||||
let correlations: Vec<f64> = sources
|
||||
.iter()
|
||||
.map(|comp| {
|
||||
let comp_mean = comp.iter().sum::<f64>() / n;
|
||||
let comp_var: f64 = comp.iter().map(|&x| (x - comp_mean).powi(2)).sum();
|
||||
|
||||
if comp_var < 1e-15 || ref_var < 1e-15 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let cov: f64 = comp
|
||||
.iter()
|
||||
.zip(reference)
|
||||
.map(|(&c, &r)| (c - comp_mean) * (r - ref_mean))
|
||||
.sum();
|
||||
|
||||
cov / (comp_var * ref_var).sqrt()
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(correlations)
|
||||
}
|
||||
}
|
||||
|
||||
/// Whiten data using PCA
|
||||
fn whiten(
|
||||
data: &[Vec<f64>],
|
||||
n_components: usize,
|
||||
) -> NeuroResult<(Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<f64>)> {
|
||||
let n_channels = data.len();
|
||||
let n_samples = data[0].len();
|
||||
|
||||
// Compute covariance matrix
|
||||
let mut cov = vec![vec![0.0; n_channels]; n_channels];
|
||||
for i in 0..n_channels {
|
||||
for j in i..n_channels {
|
||||
let c: f64 =
|
||||
(0..n_samples).map(|t| data[i][t] * data[j][t]).sum::<f64>() / n_samples as f64;
|
||||
cov[i][j] = c;
|
||||
cov[j][i] = c;
|
||||
}
|
||||
}
|
||||
|
||||
// Eigendecomposition
|
||||
let (eigenvectors, eigenvalues) = symmetric_eigen(&cov)?;
|
||||
|
||||
// Compute explained variance
|
||||
let total_var: f64 = eigenvalues.iter().sum();
|
||||
let explained_var: Vec<f64> = eigenvalues
|
||||
.iter()
|
||||
.take(n_components)
|
||||
.map(|&v| v / total_var)
|
||||
.collect();
|
||||
|
||||
// Build whitening matrix: W = D^(-1/2) * V^T
|
||||
let mut whitening = vec![vec![0.0; n_channels]; n_components];
|
||||
for i in 0..n_components {
|
||||
let scale = if eigenvalues[i] > 1e-10 {
|
||||
1.0 / eigenvalues[i].sqrt()
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
for j in 0..n_channels {
|
||||
whitening[i][j] = scale * eigenvectors[i][j];
|
||||
}
|
||||
}
|
||||
|
||||
// Apply whitening
|
||||
let mut whitened = vec![vec![0.0; n_samples]; n_components];
|
||||
for comp in 0..n_components {
|
||||
for t in 0..n_samples {
|
||||
for ch in 0..n_channels {
|
||||
whitened[comp][t] += whitening[comp][ch] * data[ch][t];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok((whitened, whitening, explained_var))
|
||||
}
|
||||
|
||||
/// FastICA algorithm
|
||||
fn fastica(
|
||||
whitened: &[Vec<f64>],
|
||||
method: IcaMethod,
|
||||
max_iter: usize,
|
||||
tol: f64,
|
||||
) -> NeuroResult<Vec<Vec<f64>>> {
|
||||
let n_components = whitened.len();
|
||||
let n_samples = whitened[0].len();
|
||||
|
||||
// Initialize unmixing matrix with random-ish values
|
||||
let mut w = vec![vec![0.0; n_components]; n_components];
|
||||
for i in 0..n_components {
|
||||
for j in 0..n_components {
|
||||
w[i][j] = ((i * 7 + j * 13 + 1) as f64).sin();
|
||||
}
|
||||
}
|
||||
|
||||
// Orthogonalize
|
||||
orthogonalize(&mut w);
|
||||
|
||||
// FastICA iteration
|
||||
for _ in 0..max_iter {
|
||||
let w_old = w.clone();
|
||||
|
||||
for p in 0..n_components {
|
||||
// Compute w^T * x for all samples
|
||||
let wx: Vec<f64> = (0..n_samples)
|
||||
.map(|t| (0..n_components).map(|i| w[p][i] * whitened[i][t]).sum())
|
||||
.collect();
|
||||
|
||||
// Compute g(w^T * x) and g'(w^T * x)
|
||||
let (g, g_prime) = match method {
|
||||
IcaMethod::FastIcaLogcosh => {
|
||||
let g: Vec<f64> = wx.iter().map(|&x| x.tanh()).collect();
|
||||
let g_p: Vec<f64> = wx.iter().map(|&x| 1.0 - x.tanh().powi(2)).collect();
|
||||
(g, g_p)
|
||||
}
|
||||
IcaMethod::FastIcaExp => {
|
||||
let g: Vec<f64> = wx.iter().map(|&x| x * (-x * x / 2.0).exp()).collect();
|
||||
let g_p: Vec<f64> = wx
|
||||
.iter()
|
||||
.map(|&x| (1.0 - x * x) * (-x * x / 2.0).exp())
|
||||
.collect();
|
||||
(g, g_p)
|
||||
}
|
||||
IcaMethod::FastIcaCube => {
|
||||
let g: Vec<f64> = wx.iter().map(|&x| x * x * x).collect();
|
||||
let g_p: Vec<f64> = wx.iter().map(|&x| 3.0 * x * x).collect();
|
||||
(g, g_p)
|
||||
}
|
||||
};
|
||||
|
||||
// Update rule: w_new = E{x * g(w^T * x)} - E{g'(w^T * x)} * w
|
||||
let e_g_prime: f64 = g_prime.iter().sum::<f64>() / n_samples as f64;
|
||||
|
||||
for i in 0..n_components {
|
||||
let e_xg: f64 =
|
||||
(0..n_samples).map(|t| whitened[i][t] * g[t]).sum::<f64>() / n_samples as f64;
|
||||
|
||||
w[p][i] = e_xg - e_g_prime * w[p][i];
|
||||
}
|
||||
}
|
||||
|
||||
// Orthogonalize
|
||||
orthogonalize(&mut w);
|
||||
|
||||
// Check convergence
|
||||
let mut max_diff = 0.0f64;
|
||||
for p in 0..n_components {
|
||||
let dot: f64 = (0..n_components)
|
||||
.map(|i| w[p][i] * w_old[p][i])
|
||||
.sum::<f64>();
|
||||
let diff = 1.0 - dot.abs();
|
||||
max_diff = max_diff.max(diff);
|
||||
}
|
||||
|
||||
if max_diff < tol {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(w)
|
||||
}
|
||||
|
||||
/// Orthogonalize matrix using symmetric decorrelation
|
||||
fn orthogonalize(w: &mut [Vec<f64>]) {
|
||||
let n = w.len();
|
||||
|
||||
// W = W * (W^T * W)^(-1/2)
|
||||
// First compute W^T * W
|
||||
let mut wtw = vec![vec![0.0; n]; n];
|
||||
for i in 0..n {
|
||||
for j in 0..n {
|
||||
wtw[i][j] = (0..n).map(|k| w[i][k] * w[j][k]).sum();
|
||||
}
|
||||
}
|
||||
|
||||
// Eigendecomposition
|
||||
if let Ok((eigvecs, eigvals)) = symmetric_eigen(&wtw) {
|
||||
// Compute (W^T * W)^(-1/2) = V * D^(-1/2) * V^T
|
||||
let mut inv_sqrt = vec![vec![0.0; n]; n];
|
||||
for i in 0..n {
|
||||
for j in 0..n {
|
||||
for k in 0..n {
|
||||
let scale = if eigvals[k] > 1e-10 {
|
||||
1.0 / eigvals[k].sqrt()
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
inv_sqrt[i][j] += eigvecs[k][i] * scale * eigvecs[k][j];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// W = W * inv_sqrt
|
||||
let w_old = w.to_vec();
|
||||
for i in 0..n {
|
||||
for j in 0..n {
|
||||
w[i][j] = (0..n).map(|k| w_old[i][k] * inv_sqrt[k][j]).sum();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Symmetric eigendecomposition using power iteration
|
||||
fn symmetric_eigen(matrix: &[Vec<f64>]) -> NeuroResult<(Vec<Vec<f64>>, Vec<f64>)> {
|
||||
let n = matrix.len();
|
||||
if n == 0 {
|
||||
return Ok((vec![], vec![]));
|
||||
}
|
||||
|
||||
let mut eigenvectors = Vec::with_capacity(n);
|
||||
let mut eigenvalues = Vec::with_capacity(n);
|
||||
let mut work = matrix.to_vec();
|
||||
|
||||
for _ in 0..n {
|
||||
// Power iteration
|
||||
let mut v: Vec<f64> = (0..n).map(|i| ((i + 1) as f64).sin()).collect();
|
||||
let mut norm: f64 = v.iter().map(|&x| x * x).sum::<f64>().sqrt();
|
||||
for x in &mut v {
|
||||
*x /= norm;
|
||||
}
|
||||
|
||||
let mut lambda = 0.0;
|
||||
for _ in 0..100 {
|
||||
// w = A * v
|
||||
let w: Vec<f64> = (0..n)
|
||||
.map(|i| (0..n).map(|j| work[i][j] * v[j]).sum())
|
||||
.collect();
|
||||
|
||||
lambda = (0..n).map(|i| v[i] * w[i]).sum();
|
||||
norm = w.iter().map(|&x| x * x).sum::<f64>().sqrt();
|
||||
|
||||
if norm < 1e-15 {
|
||||
break;
|
||||
}
|
||||
|
||||
let diff: f64 = (0..n).map(|i| (w[i] / norm - v[i]).abs()).sum();
|
||||
|
||||
v = w.iter().map(|&x| x / norm).collect();
|
||||
|
||||
if diff < 1e-10 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if lambda.abs() < 1e-15 {
|
||||
break;
|
||||
}
|
||||
|
||||
eigenvalues.push(lambda);
|
||||
eigenvectors.push(v.clone());
|
||||
|
||||
// Deflate
|
||||
for i in 0..n {
|
||||
for j in 0..n {
|
||||
work[i][j] -= lambda * v[i] * v[j];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok((eigenvectors, eigenvalues))
|
||||
}
|
||||
|
||||
/// Matrix multiplication
|
||||
fn mat_mult(a: &[Vec<f64>], b: &[Vec<f64>]) -> Vec<Vec<f64>> {
|
||||
let m = a.len();
|
||||
let n = b[0].len();
|
||||
let k = b.len();
|
||||
|
||||
let mut result = vec![vec![0.0; n]; m];
|
||||
for i in 0..m {
|
||||
for j in 0..n {
|
||||
for l in 0..k {
|
||||
result[i][j] += a[i][l] * b[l][j];
|
||||
}
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// Pseudo-inverse using direct computation for square-ish matrices
|
||||
fn pseudo_inverse(matrix: &[Vec<f64>]) -> NeuroResult<Vec<Vec<f64>>> {
|
||||
let m = matrix.len();
|
||||
let n = matrix[0].len();
|
||||
|
||||
// For the mixing matrix case (m x n where m >= n), compute (A^T A)^-1 A^T
|
||||
// Compute A^T * A
|
||||
let mut ata = vec![vec![0.0; n]; n];
|
||||
for i in 0..n {
|
||||
for j in 0..n {
|
||||
ata[i][j] = (0..m).map(|k| matrix[k][i] * matrix[k][j]).sum();
|
||||
}
|
||||
}
|
||||
|
||||
// Regularize for numerical stability
|
||||
let trace: f64 = (0..n).map(|i| ata[i][i]).sum();
|
||||
let reg = 1e-6 * trace / n as f64;
|
||||
for i in 0..n {
|
||||
ata[i][i] += reg;
|
||||
}
|
||||
|
||||
// Invert A^T * A using Cholesky-like method (symmetric positive definite)
|
||||
let ata_inv = invert_symmetric(&ata)?;
|
||||
|
||||
// Compute (A^T A)^-1 * A^T
|
||||
let mut pinv = vec![vec![0.0; m]; n];
|
||||
for i in 0..n {
|
||||
for j in 0..m {
|
||||
for k in 0..n {
|
||||
pinv[i][j] += ata_inv[i][k] * matrix[j][k];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(pinv)
|
||||
}
|
||||
|
||||
/// Invert a symmetric positive definite matrix
|
||||
fn invert_symmetric(matrix: &[Vec<f64>]) -> NeuroResult<Vec<Vec<f64>>> {
|
||||
let n = matrix.len();
|
||||
|
||||
// LDL decomposition
|
||||
let mut l = vec![vec![0.0; n]; n];
|
||||
let mut d = vec![0.0; n];
|
||||
|
||||
for i in 0..n {
|
||||
// Compute D[i]
|
||||
let mut sum = matrix[i][i];
|
||||
for k in 0..i {
|
||||
sum -= l[i][k] * l[i][k] * d[k];
|
||||
}
|
||||
d[i] = sum;
|
||||
|
||||
if d[i].abs() < 1e-15 {
|
||||
d[i] = 1e-10; // Regularize
|
||||
}
|
||||
|
||||
l[i][i] = 1.0;
|
||||
|
||||
// Compute L[j][i] for j > i
|
||||
for j in (i + 1)..n {
|
||||
let mut sum = matrix[j][i];
|
||||
for k in 0..i {
|
||||
sum -= l[j][k] * l[i][k] * d[k];
|
||||
}
|
||||
l[j][i] = sum / d[i];
|
||||
}
|
||||
}
|
||||
|
||||
// Invert L
|
||||
let mut l_inv = vec![vec![0.0; n]; n];
|
||||
for i in 0..n {
|
||||
l_inv[i][i] = 1.0;
|
||||
for j in (i + 1)..n {
|
||||
let mut sum = 0.0;
|
||||
for k in i..j {
|
||||
sum -= l[j][k] * l_inv[k][i];
|
||||
}
|
||||
l_inv[j][i] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
// Compute A^-1 = L^-T * D^-1 * L^-1
|
||||
let mut result = vec![vec![0.0; n]; n];
|
||||
for i in 0..n {
|
||||
for j in 0..n {
|
||||
for k in 0..n {
|
||||
result[i][j] += l_inv[k][i] * (1.0 / d[k]) * l_inv[k][j];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::f64::consts::PI;
|
||||
|
||||
fn create_test_data() -> Vec<Vec<f64>> {
|
||||
let n = 1000;
|
||||
let sfreq = 100.0;
|
||||
|
||||
// Two independent source signals
|
||||
let s1: Vec<f64> = (0..n)
|
||||
.map(|i| (2.0 * PI * 3.0 * i as f64 / sfreq).sin())
|
||||
.collect();
|
||||
|
||||
let s2: Vec<f64> = (0..n)
|
||||
.map(|i| {
|
||||
let t = i as f64 / sfreq;
|
||||
if (t * 2.0) as usize % 2 == 0 {
|
||||
1.0
|
||||
} else {
|
||||
-1.0
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Mixing matrix
|
||||
let a = vec![vec![0.8, 0.6], vec![0.4, 0.9], vec![0.7, 0.3]];
|
||||
|
||||
// Mixed signals: X = A * S
|
||||
let mut data = vec![vec![0.0; n]; 3];
|
||||
for ch in 0..3 {
|
||||
for t in 0..n {
|
||||
data[ch][t] = a[ch][0] * s1[t] + a[ch][1] * s2[t];
|
||||
}
|
||||
}
|
||||
|
||||
data
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ica_fit() {
|
||||
let data = create_test_data();
|
||||
|
||||
let ica = Ica::fit(&data, Some(2), IcaMethod::FastIcaLogcosh, 200, 1e-4);
|
||||
assert!(ica.is_ok());
|
||||
|
||||
let ica = ica.unwrap();
|
||||
assert_eq!(ica.n_components(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ica_transform() {
|
||||
let data = create_test_data();
|
||||
let ica = Ica::fit(&data, Some(2), IcaMethod::FastIcaLogcosh, 200, 1e-4).unwrap();
|
||||
|
||||
let sources = ica.transform(&data).unwrap();
|
||||
assert_eq!(sources.len(), 2);
|
||||
assert_eq!(sources[0].len(), 1000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ica_inverse() {
|
||||
let data = create_test_data();
|
||||
let ica = Ica::fit(&data, Some(2), IcaMethod::FastIcaLogcosh, 200, 1e-4).unwrap();
|
||||
|
||||
let sources = ica.transform(&data).unwrap();
|
||||
let reconstructed = ica.inverse_transform(&sources).unwrap();
|
||||
|
||||
// Check basic structure
|
||||
assert_eq!(reconstructed.len(), 3);
|
||||
assert_eq!(reconstructed[0].len(), 1000);
|
||||
|
||||
// Verify that ICA decomposition and reconstruction work
|
||||
// (exact reconstruction is limited by numerical precision and component number)
|
||||
// The key test is that the inverse_transform doesn't crash and produces valid output
|
||||
for ch in 0..3 {
|
||||
// Check no NaN or Inf
|
||||
assert!(
|
||||
reconstructed[ch].iter().all(|x| x.is_finite()),
|
||||
"Reconstruction contains non-finite values in channel {}",
|
||||
ch
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ica_apply() {
|
||||
let data = create_test_data();
|
||||
let ica = Ica::fit(&data, Some(2), IcaMethod::FastIcaLogcosh, 200, 1e-4).unwrap();
|
||||
|
||||
// Remove first component
|
||||
let cleaned = ica.apply(&data, &[0]).unwrap();
|
||||
assert_eq!(cleaned.len(), 3);
|
||||
|
||||
// Data should be different from original
|
||||
let diff: f64 = data[0]
|
||||
.iter()
|
||||
.zip(&cleaned[0])
|
||||
.map(|(&a, &b)| (a - b).abs())
|
||||
.sum();
|
||||
assert!(diff > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_component_kurtosis() {
|
||||
let data = create_test_data();
|
||||
let ica = Ica::fit(&data, Some(2), IcaMethod::FastIcaLogcosh, 200, 1e-4).unwrap();
|
||||
|
||||
let kurtosis = ica.component_kurtosis(&data).unwrap();
|
||||
assert_eq!(kurtosis.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_find_correlated() {
|
||||
let data = create_test_data();
|
||||
let ica = Ica::fit(&data, Some(2), IcaMethod::FastIcaLogcosh, 200, 1e-4).unwrap();
|
||||
|
||||
// Use first channel as reference
|
||||
let reference = data[0].clone();
|
||||
let correlations = ica.find_correlated_components(&data, &reference).unwrap();
|
||||
|
||||
assert_eq!(correlations.len(), 2);
|
||||
// All correlations should be bounded
|
||||
assert!(correlations.iter().all(|&c| c.abs() <= 1.0 + 1e-6));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ica_methods() {
|
||||
let data = create_test_data();
|
||||
|
||||
// Test all methods
|
||||
for method in [
|
||||
IcaMethod::FastIcaLogcosh,
|
||||
IcaMethod::FastIcaExp,
|
||||
IcaMethod::FastIcaCube,
|
||||
] {
|
||||
let ica = Ica::fit(&data, Some(2), method, 200, 1e-4);
|
||||
assert!(ica.is_ok(), "ICA {:?} failed", method);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
//! # Signal Processing Module
|
||||
//!
|
||||
//! Signal processing for MEG/EEG neuroimaging data.
|
||||
//!
|
||||
//! ## Features
|
||||
//!
|
||||
//! - **Filtering**: Bandpass, highpass, lowpass, notch filters
|
||||
//! - **Artifact Removal**: SSP, ICA
|
||||
//! - **Resampling**: Up/downsampling with anti-aliasing
|
||||
//! - **Baseline Correction**: Mean, median, zscore normalization
|
||||
//! - **Time-Frequency**: Morlet wavelets, STFT, PSD
|
||||
|
||||
#![warn(missing_docs)]
|
||||
|
||||
pub mod baseline;
|
||||
pub mod filter;
|
||||
pub mod ica;
|
||||
pub mod resample;
|
||||
pub mod ssp;
|
||||
pub mod tfr;
|
||||
|
||||
// Re-export main types
|
||||
pub use baseline::{BaselineMethod, baseline_correct};
|
||||
pub use filter::{FilterMethod, bandpass, highpass, lowpass, notch};
|
||||
pub use ica::{Ica, IcaMethod};
|
||||
pub use ssp::SspProjector;
|
||||
pub use tfr::{
|
||||
CycleSpec, PsdResult, StftResult, TfrOutput, TfrResult, Window, psd_multitaper, psd_welch,
|
||||
stft, tfr_morlet, tfr_morlet_adaptive,
|
||||
};
|
||||
@@ -0,0 +1,220 @@
|
||||
//! Resampling functions for neuroimaging signals.
|
||||
|
||||
use super::filter::{FilterMethod, lowpass};
|
||||
use crate::error::{NeuroError, NeuroResult};
|
||||
use std::f64::consts::PI;
|
||||
|
||||
/// Resample a signal to a new sampling frequency.
|
||||
///
|
||||
/// Applies anti-aliasing lowpass filter when downsampling.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `signal` - Input signal samples
|
||||
/// * `sfreq_in` - Original sampling frequency in Hz
|
||||
/// * `sfreq_out` - Target sampling frequency in Hz
|
||||
///
|
||||
/// # Returns
|
||||
/// Resampled signal
|
||||
pub fn resample(signal: &[f64], sfreq_in: f64, sfreq_out: f64) -> NeuroResult<Vec<f64>> {
|
||||
if signal.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
if sfreq_in <= 0.0 || sfreq_out <= 0.0 {
|
||||
return Err(NeuroError::Signal(
|
||||
"Sampling frequencies must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Calculate the resampling ratio
|
||||
let ratio = sfreq_out / sfreq_in;
|
||||
|
||||
if (ratio - 1.0).abs() < 1e-10 {
|
||||
return Ok(signal.to_vec());
|
||||
}
|
||||
|
||||
// If downsampling, apply anti-aliasing filter first
|
||||
let filtered = if ratio < 1.0 {
|
||||
let cutoff = sfreq_out / 2.0 * 0.9; // 90% of new Nyquist
|
||||
lowpass(signal, sfreq_in, cutoff, FilterMethod::Fft)?
|
||||
} else {
|
||||
signal.to_vec()
|
||||
};
|
||||
|
||||
// Calculate new length
|
||||
let n_out = (filtered.len() as f64 * ratio).round() as usize;
|
||||
|
||||
// Resample using linear interpolation (simple but effective)
|
||||
let mut output = Vec::with_capacity(n_out);
|
||||
|
||||
for i in 0..n_out {
|
||||
let src_idx = i as f64 / ratio;
|
||||
let idx_low = src_idx.floor() as usize;
|
||||
let idx_high = (idx_low + 1).min(filtered.len() - 1);
|
||||
let frac = src_idx - idx_low as f64;
|
||||
|
||||
let value = filtered[idx_low] * (1.0 - frac) + filtered[idx_high] * frac;
|
||||
output.push(value);
|
||||
}
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
/// Resample using sinc interpolation (higher quality).
|
||||
///
|
||||
/// Uses a windowed sinc kernel for interpolation, providing
|
||||
/// better frequency response than linear interpolation.
|
||||
pub fn resample_sinc(
|
||||
signal: &[f64],
|
||||
sfreq_in: f64,
|
||||
sfreq_out: f64,
|
||||
num_taps: usize,
|
||||
) -> NeuroResult<Vec<f64>> {
|
||||
if signal.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
if sfreq_in <= 0.0 || sfreq_out <= 0.0 {
|
||||
return Err(NeuroError::Signal(
|
||||
"Sampling frequencies must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let ratio = sfreq_out / sfreq_in;
|
||||
|
||||
if (ratio - 1.0).abs() < 1e-10 {
|
||||
return Ok(signal.to_vec());
|
||||
}
|
||||
|
||||
// Anti-aliasing filter for downsampling
|
||||
let filtered = if ratio < 1.0 {
|
||||
let cutoff = sfreq_out / 2.0 * 0.9;
|
||||
lowpass(signal, sfreq_in, cutoff, FilterMethod::Fft)?
|
||||
} else {
|
||||
signal.to_vec()
|
||||
};
|
||||
|
||||
let n_in = filtered.len();
|
||||
let n_out = (n_in as f64 * ratio).round() as usize;
|
||||
let half_taps = num_taps / 2;
|
||||
|
||||
let mut output = Vec::with_capacity(n_out);
|
||||
|
||||
for i in 0..n_out {
|
||||
let src_idx = i as f64 / ratio;
|
||||
let src_int = src_idx.floor() as isize;
|
||||
|
||||
let mut sum = 0.0;
|
||||
let mut weight_sum = 0.0;
|
||||
|
||||
for j in -(half_taps as isize)..=(half_taps as isize) {
|
||||
let idx = src_int + j;
|
||||
if idx < 0 || idx >= n_in as isize {
|
||||
continue;
|
||||
}
|
||||
|
||||
let t = src_idx - idx as f64;
|
||||
let weight = sinc(t) * blackman_window(t, half_taps as f64);
|
||||
|
||||
sum += filtered[idx as usize] * weight;
|
||||
weight_sum += weight;
|
||||
}
|
||||
|
||||
let value = if weight_sum.abs() > 1e-10 {
|
||||
sum / weight_sum
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
output.push(value);
|
||||
}
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
/// Normalized sinc function: sinc(x) = sin(pi*x) / (pi*x)
|
||||
fn sinc(x: f64) -> f64 {
|
||||
if x.abs() < 1e-10 {
|
||||
1.0
|
||||
} else {
|
||||
let px = PI * x;
|
||||
px.sin() / px
|
||||
}
|
||||
}
|
||||
|
||||
/// Blackman window function
|
||||
fn blackman_window(x: f64, half_width: f64) -> f64 {
|
||||
if x.abs() > half_width {
|
||||
0.0
|
||||
} else {
|
||||
let t = x / half_width; // Normalize to [-1, 1]
|
||||
let theta = PI * (t + 1.0); // Map to [0, 2*pi]
|
||||
0.42 - 0.5 * theta.cos() + 0.08 * (2.0 * theta).cos()
|
||||
}
|
||||
}
|
||||
|
||||
/// Decimate a signal by an integer factor.
|
||||
///
|
||||
/// More efficient than resample for integer factors.
|
||||
pub fn decimate(signal: &[f64], sfreq: f64, factor: usize) -> NeuroResult<Vec<f64>> {
|
||||
if factor == 0 {
|
||||
return Err(NeuroError::Signal(
|
||||
"Decimation factor must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
if factor == 1 {
|
||||
return Ok(signal.to_vec());
|
||||
}
|
||||
|
||||
// Apply anti-aliasing filter
|
||||
let cutoff = sfreq / factor as f64 / 2.0 * 0.9;
|
||||
let filtered = lowpass(signal, sfreq, cutoff, FilterMethod::Fft)?;
|
||||
|
||||
// Decimate
|
||||
Ok(filtered.iter().step_by(factor).copied().collect())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::f64::consts::PI;
|
||||
|
||||
fn generate_sine(sfreq: f64, duration: f64, freq: f64) -> Vec<f64> {
|
||||
let n = (sfreq * duration) as usize;
|
||||
(0..n)
|
||||
.map(|i| (2.0 * PI * freq * i as f64 / sfreq).sin())
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resample_downsample() {
|
||||
let signal = generate_sine(1000.0, 1.0, 10.0);
|
||||
let resampled = resample(&signal, 1000.0, 500.0).unwrap();
|
||||
|
||||
assert_eq!(resampled.len(), 500);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resample_upsample() {
|
||||
let signal = generate_sine(500.0, 1.0, 10.0);
|
||||
let resampled = resample(&signal, 500.0, 1000.0).unwrap();
|
||||
|
||||
assert_eq!(resampled.len(), 1000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decimate() {
|
||||
let signal: Vec<f64> = (0..100).map(|x| x as f64).collect();
|
||||
let decimated = decimate(&signal, 1000.0, 2).unwrap();
|
||||
|
||||
assert_eq!(decimated.len(), 50);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sinc() {
|
||||
assert!((sinc(0.0) - 1.0).abs() < 1e-10);
|
||||
assert!(sinc(1.0).abs() < 1e-10);
|
||||
assert!(sinc(2.0).abs() < 1e-10);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,690 @@
|
||||
//! Signal Space Projection (SSP) for artifact removal.
|
||||
//!
|
||||
//! SSP is a spatial filtering technique that removes artifacts by projecting
|
||||
//! out directions in sensor space that contain artifact activity.
|
||||
//!
|
||||
//! ## Mathematical Background
|
||||
//!
|
||||
//! Given artifact data matrix A [n_channels x n_samples], we compute the
|
||||
//! covariance C = A * A^T and find its principal components via SVD.
|
||||
//!
|
||||
//! The projection operator is: P = I - Σ(u_i * u_i^T)
|
||||
//! where u_i are the artifact subspace vectors to remove.
|
||||
//!
|
||||
//! Clean data is obtained by: M_clean = P * M
|
||||
|
||||
use crate::error::{NeuroError, NeuroResult};
|
||||
|
||||
/// SSP projector for artifact removal
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SspProjector {
|
||||
/// Projection vectors [n_projectors x n_channels]
|
||||
/// Each row is a unit vector defining an artifact direction
|
||||
vectors: Vec<Vec<f64>>,
|
||||
/// Whether each projector is active
|
||||
active: Vec<bool>,
|
||||
/// Explained variance by each projector
|
||||
explained_var: Vec<f64>,
|
||||
/// Names/descriptions of projectors
|
||||
names: Vec<String>,
|
||||
/// Number of channels
|
||||
n_channels: usize,
|
||||
}
|
||||
|
||||
impl SspProjector {
|
||||
/// Create an SSP projector from artifact epochs
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `epochs` - Artifact epochs [n_epochs][n_channels][n_samples]
|
||||
/// * `n_components` - Number of projectors to compute
|
||||
/// * `name` - Name for this projector set (e.g., "EOG", "ECG")
|
||||
///
|
||||
/// # Returns
|
||||
/// SSP projector with the requested number of components
|
||||
pub fn from_epochs(
|
||||
epochs: &[Vec<Vec<f64>>],
|
||||
n_components: usize,
|
||||
name: &str,
|
||||
) -> NeuroResult<Self> {
|
||||
if epochs.is_empty() {
|
||||
return Err(NeuroError::Signal("No epochs provided".to_string()));
|
||||
}
|
||||
|
||||
let n_channels = epochs[0].len();
|
||||
if n_channels == 0 {
|
||||
return Err(NeuroError::Signal("Empty channels".to_string()));
|
||||
}
|
||||
|
||||
// Concatenate all epochs into a data matrix [n_channels x total_samples]
|
||||
let total_samples: usize = epochs.iter().map(|e| e[0].len()).sum();
|
||||
let mut data = vec![vec![0.0; total_samples]; n_channels];
|
||||
|
||||
let mut offset = 0;
|
||||
for epoch in epochs {
|
||||
let n_times = epoch[0].len();
|
||||
for ch in 0..n_channels {
|
||||
for t in 0..n_times {
|
||||
data[ch][offset + t] = epoch[ch][t];
|
||||
}
|
||||
}
|
||||
offset += n_times;
|
||||
}
|
||||
|
||||
// Compute covariance matrix: C = (1/n) * X * X^T
|
||||
let mut cov = vec![vec![0.0; n_channels]; n_channels];
|
||||
for i in 0..n_channels {
|
||||
for j in i..n_channels {
|
||||
let mut sum = 0.0;
|
||||
for t in 0..total_samples {
|
||||
sum += data[i][t] * data[j][t];
|
||||
}
|
||||
cov[i][j] = sum / total_samples as f64;
|
||||
cov[j][i] = cov[i][j];
|
||||
}
|
||||
}
|
||||
|
||||
// Compute SVD of covariance matrix
|
||||
let (eigenvectors, eigenvalues) = symmetric_eigen(&cov)?;
|
||||
|
||||
// Total variance
|
||||
let total_var: f64 = eigenvalues.iter().sum();
|
||||
|
||||
// Extract top n_components
|
||||
let n_proj = n_components.min(n_channels);
|
||||
let mut vectors = Vec::with_capacity(n_proj);
|
||||
let mut explained_var = Vec::with_capacity(n_proj);
|
||||
let mut names = Vec::with_capacity(n_proj);
|
||||
let mut active = Vec::with_capacity(n_proj);
|
||||
|
||||
for i in 0..n_proj {
|
||||
vectors.push(eigenvectors[i].clone());
|
||||
explained_var.push(eigenvalues[i] / total_var);
|
||||
names.push(format!("{}-{}", name, i + 1));
|
||||
active.push(true);
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
vectors,
|
||||
active,
|
||||
explained_var,
|
||||
names,
|
||||
n_channels,
|
||||
})
|
||||
}
|
||||
|
||||
/// Create SSP projector from EOG data
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `data` - Continuous data [n_channels][n_samples]
|
||||
/// * `eog_channels` - Indices of EOG channels
|
||||
/// * `n_components` - Number of projectors (typically 1-2)
|
||||
pub fn from_eog(
|
||||
data: &[Vec<f64>],
|
||||
eog_channels: &[usize],
|
||||
n_components: usize,
|
||||
) -> NeuroResult<Self> {
|
||||
if data.is_empty() {
|
||||
return Err(NeuroError::Signal("Empty data".to_string()));
|
||||
}
|
||||
|
||||
let n_channels = data.len();
|
||||
let n_samples = data[0].len();
|
||||
|
||||
// Create artifact epochs from EOG threshold crossings
|
||||
let mut artifact_epochs = Vec::new();
|
||||
|
||||
for &eog_ch in eog_channels {
|
||||
if eog_ch >= n_channels {
|
||||
return Err(NeuroError::Signal(format!(
|
||||
"EOG channel {} out of range",
|
||||
eog_ch
|
||||
)));
|
||||
}
|
||||
|
||||
// Find threshold (use robust estimate: 5 * median absolute deviation)
|
||||
let eog_data = &data[eog_ch];
|
||||
let threshold = compute_robust_threshold(eog_data);
|
||||
|
||||
// Find peaks above threshold
|
||||
let peaks = find_peaks(eog_data, threshold);
|
||||
|
||||
// Extract epochs around each peak (±100 samples or available)
|
||||
let half_window = 100.min(n_samples / 2);
|
||||
for peak in peaks {
|
||||
let start = peak.saturating_sub(half_window);
|
||||
let end = (peak + half_window).min(n_samples);
|
||||
|
||||
let epoch: Vec<Vec<f64>> = data.iter().map(|ch| ch[start..end].to_vec()).collect();
|
||||
artifact_epochs.push(epoch);
|
||||
}
|
||||
}
|
||||
|
||||
if artifact_epochs.is_empty() {
|
||||
// Create a projector from the raw EOG channel correlation
|
||||
return Self::from_eog_correlation(data, eog_channels, n_components);
|
||||
}
|
||||
|
||||
Self::from_epochs(&artifact_epochs, n_components, "EOG")
|
||||
}
|
||||
|
||||
/// Create SSP from EOG channel correlation (fallback method)
|
||||
fn from_eog_correlation(
|
||||
data: &[Vec<f64>],
|
||||
eog_channels: &[usize],
|
||||
n_components: usize,
|
||||
) -> NeuroResult<Self> {
|
||||
let n_channels = data.len();
|
||||
let n_samples = data[0].len();
|
||||
|
||||
// Compute correlation of each channel with EOG channels
|
||||
let mut vectors = Vec::new();
|
||||
let mut explained_var = Vec::new();
|
||||
|
||||
for (comp_idx, &eog_ch) in eog_channels.iter().take(n_components).enumerate() {
|
||||
let eog_data = &data[eog_ch];
|
||||
let eog_var: f64 = eog_data.iter().map(|&x| x * x).sum::<f64>() / n_samples as f64;
|
||||
|
||||
if eog_var < 1e-15 {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Compute correlation vector
|
||||
let mut corr_vec = vec![0.0; n_channels];
|
||||
for ch in 0..n_channels {
|
||||
let cov: f64 = (0..n_samples)
|
||||
.map(|t| data[ch][t] * eog_data[t])
|
||||
.sum::<f64>()
|
||||
/ n_samples as f64;
|
||||
corr_vec[ch] = cov / eog_var.sqrt();
|
||||
}
|
||||
|
||||
// Normalize
|
||||
let norm: f64 = corr_vec.iter().map(|&x| x * x).sum::<f64>().sqrt();
|
||||
if norm > 1e-10 {
|
||||
for v in &mut corr_vec {
|
||||
*v /= norm;
|
||||
}
|
||||
vectors.push(corr_vec);
|
||||
explained_var.push(1.0 / (comp_idx + 1) as f64);
|
||||
}
|
||||
}
|
||||
|
||||
let n_proj = vectors.len();
|
||||
Ok(Self {
|
||||
vectors,
|
||||
active: vec![true; n_proj],
|
||||
explained_var,
|
||||
names: (0..n_proj).map(|i| format!("EOG-{}", i + 1)).collect(),
|
||||
n_channels,
|
||||
})
|
||||
}
|
||||
|
||||
/// Create SSP projector from ECG data
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `data` - Continuous data [n_channels][n_samples]
|
||||
/// * `ecg_channel` - Index of ECG channel
|
||||
/// * `n_components` - Number of projectors (typically 1-3)
|
||||
pub fn from_ecg(
|
||||
data: &[Vec<f64>],
|
||||
ecg_channel: usize,
|
||||
n_components: usize,
|
||||
) -> NeuroResult<Self> {
|
||||
if data.is_empty() {
|
||||
return Err(NeuroError::Signal("Empty data".to_string()));
|
||||
}
|
||||
|
||||
let n_channels = data.len();
|
||||
let n_samples = data[0].len();
|
||||
|
||||
if ecg_channel >= n_channels {
|
||||
return Err(NeuroError::Signal(format!(
|
||||
"ECG channel {} out of range",
|
||||
ecg_channel
|
||||
)));
|
||||
}
|
||||
|
||||
let ecg_data = &data[ecg_channel];
|
||||
|
||||
// Find R-peaks using simple threshold detection
|
||||
let threshold = compute_robust_threshold(ecg_data);
|
||||
let r_peaks = find_peaks(ecg_data, threshold);
|
||||
|
||||
if r_peaks.len() < 3 {
|
||||
return Err(NeuroError::Signal(
|
||||
"Not enough ECG peaks detected".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Extract epochs around each R-peak (±50 samples)
|
||||
let half_window = 50.min(n_samples / 2);
|
||||
let mut artifact_epochs = Vec::new();
|
||||
|
||||
for peak in r_peaks {
|
||||
let start = peak.saturating_sub(half_window);
|
||||
let end = (peak + half_window).min(n_samples);
|
||||
|
||||
let epoch: Vec<Vec<f64>> = data.iter().map(|ch| ch[start..end].to_vec()).collect();
|
||||
artifact_epochs.push(epoch);
|
||||
}
|
||||
|
||||
Self::from_epochs(&artifact_epochs, n_components, "ECG")
|
||||
}
|
||||
|
||||
/// Apply the SSP projector to data
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `data` - Data to clean [n_channels][n_samples]
|
||||
///
|
||||
/// # Returns
|
||||
/// Cleaned data with artifacts removed
|
||||
pub fn apply(&self, data: &[Vec<f64>]) -> NeuroResult<Vec<Vec<f64>>> {
|
||||
if data.len() != self.n_channels {
|
||||
return Err(NeuroError::Signal(format!(
|
||||
"Expected {} channels, got {}",
|
||||
self.n_channels,
|
||||
data.len()
|
||||
)));
|
||||
}
|
||||
|
||||
let n_samples = data[0].len();
|
||||
|
||||
// Compute projection matrix: P = I - Σ(u_i * u_i^T)
|
||||
// For efficiency, we apply it directly: y = x - Σ((u_i^T * x) * u_i)
|
||||
let mut result = data.to_vec();
|
||||
|
||||
for (i, vec) in self.vectors.iter().enumerate() {
|
||||
if !self.active[i] {
|
||||
continue;
|
||||
}
|
||||
|
||||
for t in 0..n_samples {
|
||||
// Compute dot product: u^T * x
|
||||
let dot: f64 = (0..self.n_channels).map(|ch| vec[ch] * data[ch][t]).sum();
|
||||
|
||||
// Subtract: x = x - (u^T * x) * u
|
||||
for ch in 0..self.n_channels {
|
||||
result[ch][t] -= dot * vec[ch];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Apply SSP to epoched data
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `epochs` - Epoched data [n_epochs][n_channels][n_samples]
|
||||
///
|
||||
/// # Returns
|
||||
/// Cleaned epochs
|
||||
pub fn apply_epochs(&self, epochs: &[Vec<Vec<f64>>]) -> NeuroResult<Vec<Vec<Vec<f64>>>> {
|
||||
epochs.iter().map(|epoch| self.apply(epoch)).collect()
|
||||
}
|
||||
|
||||
/// Activate or deactivate a projector
|
||||
pub fn set_active(&mut self, index: usize, active: bool) -> NeuroResult<()> {
|
||||
if index >= self.vectors.len() {
|
||||
return Err(NeuroError::Signal(format!(
|
||||
"Projector index {} out of range",
|
||||
index
|
||||
)));
|
||||
}
|
||||
self.active[index] = active;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get number of projectors
|
||||
pub fn n_projectors(&self) -> usize {
|
||||
self.vectors.len()
|
||||
}
|
||||
|
||||
/// Get number of active projectors
|
||||
pub fn n_active(&self) -> usize {
|
||||
self.active.iter().filter(|&&a| a).count()
|
||||
}
|
||||
|
||||
/// Get projector vectors
|
||||
pub fn vectors(&self) -> &[Vec<f64>] {
|
||||
&self.vectors
|
||||
}
|
||||
|
||||
/// Get explained variance ratios
|
||||
pub fn explained_var(&self) -> &[f64] {
|
||||
&self.explained_var
|
||||
}
|
||||
|
||||
/// Get projector names
|
||||
pub fn names(&self) -> &[String] {
|
||||
&self.names
|
||||
}
|
||||
|
||||
/// Check if projector is active
|
||||
pub fn is_active(&self, index: usize) -> bool {
|
||||
self.active.get(index).copied().unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Compute the projection matrix P = I - Σ(u_i * u_i^T)
|
||||
pub fn compute_projection_matrix(&self) -> Vec<Vec<f64>> {
|
||||
let n = self.n_channels;
|
||||
let mut proj = vec![vec![0.0; n]; n];
|
||||
|
||||
// Start with identity
|
||||
for i in 0..n {
|
||||
proj[i][i] = 1.0;
|
||||
}
|
||||
|
||||
// Subtract outer products of active projectors
|
||||
for (i, vec) in self.vectors.iter().enumerate() {
|
||||
if !self.active[i] {
|
||||
continue;
|
||||
}
|
||||
for row in 0..n {
|
||||
for col in 0..n {
|
||||
proj[row][col] -= vec[row] * vec[col];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
proj
|
||||
}
|
||||
|
||||
/// Merge multiple SSP projectors
|
||||
pub fn merge(projectors: &[&SspProjector]) -> NeuroResult<Self> {
|
||||
if projectors.is_empty() {
|
||||
return Err(NeuroError::Signal("No projectors to merge".to_string()));
|
||||
}
|
||||
|
||||
let n_channels = projectors[0].n_channels;
|
||||
for p in projectors {
|
||||
if p.n_channels != n_channels {
|
||||
return Err(NeuroError::Signal("Channel count mismatch".to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
let mut vectors = Vec::new();
|
||||
let mut active = Vec::new();
|
||||
let mut explained_var = Vec::new();
|
||||
let mut names = Vec::new();
|
||||
|
||||
for p in projectors {
|
||||
vectors.extend(p.vectors.clone());
|
||||
active.extend(p.active.clone());
|
||||
explained_var.extend(p.explained_var.clone());
|
||||
names.extend(p.names.clone());
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
vectors,
|
||||
active,
|
||||
explained_var,
|
||||
names,
|
||||
n_channels,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute robust threshold using median absolute deviation
|
||||
fn compute_robust_threshold(data: &[f64]) -> f64 {
|
||||
let mut sorted: Vec<f64> = data.iter().map(|&x| x.abs()).collect();
|
||||
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
|
||||
|
||||
let median = if sorted.len() % 2 == 0 {
|
||||
(sorted[sorted.len() / 2 - 1] + sorted[sorted.len() / 2]) / 2.0
|
||||
} else {
|
||||
sorted[sorted.len() / 2]
|
||||
};
|
||||
|
||||
// MAD = median(|x - median|)
|
||||
let mut deviations: Vec<f64> = data.iter().map(|&x| (x.abs() - median).abs()).collect();
|
||||
deviations.sort_by(|a, b| a.partial_cmp(b).unwrap());
|
||||
|
||||
let mad = if deviations.len() % 2 == 0 {
|
||||
(deviations[deviations.len() / 2 - 1] + deviations[deviations.len() / 2]) / 2.0
|
||||
} else {
|
||||
deviations[deviations.len() / 2]
|
||||
};
|
||||
|
||||
// 5 * MAD is a robust threshold
|
||||
5.0 * mad.max(1e-10)
|
||||
}
|
||||
|
||||
/// Find peaks above threshold with minimum separation
|
||||
fn find_peaks(data: &[f64], threshold: f64) -> Vec<usize> {
|
||||
let min_separation = 50; // Minimum samples between peaks
|
||||
let mut peaks = Vec::new();
|
||||
let mut last_peak: Option<usize> = None;
|
||||
|
||||
for i in 1..data.len().saturating_sub(1) {
|
||||
if data[i] > threshold && data[i] > data[i - 1] && data[i] > data[i + 1] {
|
||||
if let Some(last) = last_peak {
|
||||
if i - last >= min_separation {
|
||||
peaks.push(i);
|
||||
last_peak = Some(i);
|
||||
}
|
||||
} else {
|
||||
peaks.push(i);
|
||||
last_peak = Some(i);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
peaks
|
||||
}
|
||||
|
||||
/// Compute eigendecomposition of symmetric matrix using power iteration
|
||||
fn symmetric_eigen(matrix: &[Vec<f64>]) -> NeuroResult<(Vec<Vec<f64>>, Vec<f64>)> {
|
||||
let n = matrix.len();
|
||||
if n == 0 {
|
||||
return Ok((vec![], vec![]));
|
||||
}
|
||||
|
||||
let mut eigenvectors = Vec::with_capacity(n);
|
||||
let mut eigenvalues = Vec::with_capacity(n);
|
||||
let mut work_matrix = matrix.to_vec();
|
||||
|
||||
for _ in 0..n {
|
||||
// Power iteration to find largest eigenvalue/vector
|
||||
let (eigval, eigvec) = power_iteration(&work_matrix, 100)?;
|
||||
|
||||
if eigval.abs() < 1e-15 {
|
||||
break;
|
||||
}
|
||||
|
||||
eigenvalues.push(eigval);
|
||||
eigenvectors.push(eigvec.clone());
|
||||
|
||||
// Deflate: A = A - λ * v * v^T
|
||||
for i in 0..n {
|
||||
for j in 0..n {
|
||||
work_matrix[i][j] -= eigval * eigvec[i] * eigvec[j];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok((eigenvectors, eigenvalues))
|
||||
}
|
||||
|
||||
/// Power iteration to find dominant eigenvalue/vector
|
||||
fn power_iteration(matrix: &[Vec<f64>], max_iter: usize) -> NeuroResult<(f64, Vec<f64>)> {
|
||||
let n = matrix.len();
|
||||
if n == 0 {
|
||||
return Err(NeuroError::Signal("Empty matrix".to_string()));
|
||||
}
|
||||
|
||||
// Initialize with random-ish vector
|
||||
let mut v: Vec<f64> = (0..n).map(|i| ((i + 1) as f64).sin()).collect();
|
||||
|
||||
// Normalize
|
||||
let mut norm: f64 = v.iter().map(|&x| x * x).sum::<f64>().sqrt();
|
||||
for x in &mut v {
|
||||
*x /= norm;
|
||||
}
|
||||
|
||||
let mut eigenvalue = 0.0;
|
||||
|
||||
for _ in 0..max_iter {
|
||||
// w = A * v
|
||||
let w: Vec<f64> = (0..n)
|
||||
.map(|i| (0..n).map(|j| matrix[i][j] * v[j]).sum())
|
||||
.collect();
|
||||
|
||||
// Rayleigh quotient: λ = v^T * A * v
|
||||
eigenvalue = (0..n).map(|i| v[i] * w[i]).sum();
|
||||
|
||||
// Normalize w
|
||||
norm = w.iter().map(|&x| x * x).sum::<f64>().sqrt();
|
||||
if norm < 1e-15 {
|
||||
break;
|
||||
}
|
||||
|
||||
// Check convergence
|
||||
let diff: f64 = (0..n).map(|i| (w[i] / norm - v[i]).abs()).sum::<f64>();
|
||||
|
||||
v = w.iter().map(|&x| x / norm).collect();
|
||||
|
||||
if diff < 1e-10 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Ok((eigenvalue, v))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_ssp_from_epochs() {
|
||||
// Create simple artifact epochs
|
||||
let epochs = vec![
|
||||
vec![
|
||||
vec![1.0, 2.0, 3.0, 2.0, 1.0],
|
||||
vec![0.5, 1.0, 1.5, 1.0, 0.5],
|
||||
vec![0.2, 0.4, 0.6, 0.4, 0.2],
|
||||
],
|
||||
vec![
|
||||
vec![1.1, 2.1, 3.1, 2.1, 1.1],
|
||||
vec![0.55, 1.05, 1.55, 1.05, 0.55],
|
||||
vec![0.22, 0.42, 0.62, 0.42, 0.22],
|
||||
],
|
||||
];
|
||||
|
||||
let ssp = SspProjector::from_epochs(&epochs, 2, "test").unwrap();
|
||||
|
||||
assert!(ssp.n_projectors() >= 1);
|
||||
assert_eq!(ssp.n_active(), ssp.n_projectors());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ssp_apply() {
|
||||
// Create a known projector
|
||||
let ssp = SspProjector {
|
||||
vectors: vec![vec![1.0, 0.0, 0.0]], // Project out first channel
|
||||
active: vec![true],
|
||||
explained_var: vec![0.5],
|
||||
names: vec!["test-1".to_string()],
|
||||
n_channels: 3,
|
||||
};
|
||||
|
||||
let data = vec![
|
||||
vec![1.0, 2.0, 3.0],
|
||||
vec![4.0, 5.0, 6.0],
|
||||
vec![7.0, 8.0, 9.0],
|
||||
];
|
||||
|
||||
let cleaned = ssp.apply(&data).unwrap();
|
||||
|
||||
// First channel should be zeroed
|
||||
assert!(cleaned[0].iter().all(|&x| x.abs() < 1e-10));
|
||||
// Other channels unchanged
|
||||
assert_eq!(cleaned[1], vec![4.0, 5.0, 6.0]);
|
||||
assert_eq!(cleaned[2], vec![7.0, 8.0, 9.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ssp_projection_matrix() {
|
||||
let ssp = SspProjector {
|
||||
vectors: vec![vec![1.0, 0.0], vec![0.0, 1.0]],
|
||||
active: vec![true, false],
|
||||
explained_var: vec![0.5, 0.3],
|
||||
names: vec!["p1".to_string(), "p2".to_string()],
|
||||
n_channels: 2,
|
||||
};
|
||||
|
||||
let proj = ssp.compute_projection_matrix();
|
||||
|
||||
// Only first projector is active, so P = I - u1*u1^T
|
||||
// u1 = [1, 0], so P = [[0, 0], [0, 1]]
|
||||
assert!((proj[0][0] - 0.0).abs() < 1e-10);
|
||||
assert!((proj[0][1] - 0.0).abs() < 1e-10);
|
||||
assert!((proj[1][0] - 0.0).abs() < 1e-10);
|
||||
assert!((proj[1][1] - 1.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_set_active() {
|
||||
let mut ssp = SspProjector {
|
||||
vectors: vec![vec![1.0, 0.0]],
|
||||
active: vec![true],
|
||||
explained_var: vec![0.5],
|
||||
names: vec!["p1".to_string()],
|
||||
n_channels: 2,
|
||||
};
|
||||
|
||||
assert!(ssp.is_active(0));
|
||||
ssp.set_active(0, false).unwrap();
|
||||
assert!(!ssp.is_active(0));
|
||||
assert_eq!(ssp.n_active(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_merge_projectors() {
|
||||
let p1 = SspProjector {
|
||||
vectors: vec![vec![1.0, 0.0, 0.0]],
|
||||
active: vec![true],
|
||||
explained_var: vec![0.5],
|
||||
names: vec!["EOG-1".to_string()],
|
||||
n_channels: 3,
|
||||
};
|
||||
|
||||
let p2 = SspProjector {
|
||||
vectors: vec![vec![0.0, 1.0, 0.0]],
|
||||
active: vec![true],
|
||||
explained_var: vec![0.3],
|
||||
names: vec!["ECG-1".to_string()],
|
||||
n_channels: 3,
|
||||
};
|
||||
|
||||
let merged = SspProjector::merge(&[&p1, &p2]).unwrap();
|
||||
|
||||
assert_eq!(merged.n_projectors(), 2);
|
||||
assert_eq!(merged.names()[0], "EOG-1");
|
||||
assert_eq!(merged.names()[1], "ECG-1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_power_iteration() {
|
||||
// Diagonal matrix - eigenvalues are diagonal elements
|
||||
let matrix = vec![vec![3.0, 0.0], vec![0.0, 1.0]];
|
||||
|
||||
let (eigval, eigvec) = power_iteration(&matrix, 100).unwrap();
|
||||
|
||||
// Largest eigenvalue should be 3
|
||||
assert!((eigval - 3.0).abs() < 1e-6);
|
||||
// Corresponding eigenvector should be [1, 0]
|
||||
assert!((eigvec[0].abs() - 1.0).abs() < 1e-6);
|
||||
assert!(eigvec[1].abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_find_peaks() {
|
||||
let data = vec![0.0, 1.0, 5.0, 2.0, 0.0, 0.0, 0.0];
|
||||
let peaks = find_peaks(&data, 3.0);
|
||||
|
||||
assert_eq!(peaks.len(), 1);
|
||||
assert_eq!(peaks[0], 2);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,751 @@
|
||||
//! 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::error::{NeuroError, NeuroResult};
|
||||
|
||||
/// 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,
|
||||
) -> NeuroResult<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,
|
||||
) -> NeuroResult<TfrResult> {
|
||||
if data.is_empty() || data[0].is_empty() {
|
||||
return Err(NeuroError::Signal("Empty data".to_string()));
|
||||
}
|
||||
if freqs.is_empty() {
|
||||
return Err(NeuroError::Signal("No frequencies specified".to_string()));
|
||||
}
|
||||
|
||||
let n_samples = data[0].len();
|
||||
let n_freqs = freqs.len();
|
||||
let freq_min = freqs.iter().cloned().fold(f64::INFINITY, f64::min);
|
||||
let freq_max = freqs.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
|
||||
|
||||
// Decimated output length
|
||||
let n_times = (n_samples + decim - 1) / 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()
|
||||
.cloned()
|
||||
.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(|c| c.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,
|
||||
) -> NeuroResult<StftResult> {
|
||||
if data.is_empty() || data[0].is_empty() {
|
||||
return Err(NeuroError::Signal("Empty data".to_string()));
|
||||
}
|
||||
if n_fft == 0 || !n_fft.is_power_of_two() {
|
||||
return Err(NeuroError::Signal("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,
|
||||
) -> NeuroResult<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(|c| c.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>,
|
||||
) -> NeuroResult<PsdResult> {
|
||||
if data.is_empty() || data[0].is_empty() {
|
||||
return Err(NeuroError::Signal("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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user