414 lines
12 KiB
Rust
414 lines
12 KiB
Rust
//! Spectral coherence estimation.
|
|
//!
|
|
//! Coherence measures the linear relationship between two signals in the frequency domain.
|
|
//!
|
|
//! ## Mathematical Background
|
|
//!
|
|
//! Magnitude-squared coherence: Cxy(f) = |Sxy(f)|² / (Sxx(f) * Syy(f))
|
|
//!
|
|
//! where Sxy is cross-spectral density and Sxx, Syy are auto-spectral densities.
|
|
|
|
use crate::{ConnectivityError, ConnectivityResult, utils};
|
|
use num_complex::Complex64;
|
|
|
|
/// Compute coherence between two signals
|
|
///
|
|
/// # Arguments
|
|
/// * `x` - First signal [n_samples]
|
|
/// * `y` - Second signal [n_samples]
|
|
/// * `sfreq` - Sampling frequency (Hz)
|
|
/// * `n_fft` - FFT length (None = signal length)
|
|
///
|
|
/// # Returns
|
|
/// (coherence, frequencies) where coherence is [n_freqs]
|
|
pub fn coherence(
|
|
x: &[f64],
|
|
y: &[f64],
|
|
sfreq: f64,
|
|
n_fft: Option<usize>,
|
|
) -> ConnectivityResult<(Vec<f64>, Vec<f64>)> {
|
|
if x.len() != y.len() {
|
|
return Err(ConnectivityError::DimensionMismatch(
|
|
"Signals must have same length".to_string(),
|
|
));
|
|
}
|
|
|
|
let n = x.len();
|
|
let n_fft = n_fft.unwrap_or(n);
|
|
|
|
// Apply window
|
|
let window = utils::hanning_window(n);
|
|
let x_win = utils::apply_window(x, &window);
|
|
let y_win = utils::apply_window(y, &window);
|
|
|
|
// Compute FFT
|
|
let x_fft = utils::rfft(&x_win, n_fft);
|
|
let y_fft = utils::rfft(&y_win, n_fft);
|
|
|
|
// Cross and auto spectral densities
|
|
let sxy = utils::cross_spectral_density(&x_fft, &y_fft);
|
|
let sxx = utils::power_spectral_density(&x_fft);
|
|
let syy = utils::power_spectral_density(&y_fft);
|
|
|
|
// Coherence: |Sxy|² / (Sxx * Syy)
|
|
// Clamp to [0, 1] due to numerical precision
|
|
let coh: Vec<f64> = sxy
|
|
.iter()
|
|
.zip(sxx.iter())
|
|
.zip(syy.iter())
|
|
.map(|((xy, &xx), &yy)| {
|
|
let denom = xx * yy;
|
|
if denom > 1e-15 {
|
|
(xy.norm_sqr() / denom).clamp(0.0, 1.0)
|
|
} else {
|
|
0.0
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
let freqs = utils::fft_freqs(sfreq, n_fft, 0.0, sfreq / 2.0);
|
|
|
|
Ok((coh, freqs))
|
|
}
|
|
|
|
/// Compute coherence between specific channel pairs
|
|
///
|
|
/// # Arguments
|
|
/// * `x` - First signal [n_samples]
|
|
/// * `y` - Second signal [n_samples]
|
|
/// * `sfreq` - Sampling frequency
|
|
/// * `fmin` - Minimum frequency
|
|
/// * `fmax` - Maximum frequency
|
|
/// * `n_fft` - FFT length
|
|
pub fn coherence_pairs(
|
|
x: &[f64],
|
|
y: &[f64],
|
|
sfreq: f64,
|
|
fmin: f64,
|
|
fmax: f64,
|
|
n_fft: usize,
|
|
) -> ConnectivityResult<Vec<f64>> {
|
|
let (coh, freqs) = coherence(x, y, sfreq, Some(n_fft))?;
|
|
|
|
// Filter to frequency range
|
|
let result: Vec<f64> = coh
|
|
.iter()
|
|
.zip(freqs.iter())
|
|
.filter(|&(_, f)| *f >= fmin && *f <= fmax)
|
|
.map(|(&c, _)| c)
|
|
.collect();
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
/// Compute imaginary coherence (robust to volume conduction)
|
|
///
|
|
/// Imaginary coherence only considers the imaginary part of the cross-spectrum,
|
|
/// which is zero for volume-conducted signals (zero phase lag).
|
|
pub fn imaginary_coherence(
|
|
x: &[f64],
|
|
y: &[f64],
|
|
sfreq: f64,
|
|
n_fft: Option<usize>,
|
|
) -> ConnectivityResult<(Vec<f64>, Vec<f64>)> {
|
|
if x.len() != y.len() {
|
|
return Err(ConnectivityError::DimensionMismatch(
|
|
"Signals must have same length".to_string(),
|
|
));
|
|
}
|
|
|
|
let n = x.len();
|
|
let n_fft = n_fft.unwrap_or(n);
|
|
|
|
let window = utils::hanning_window(n);
|
|
let x_win = utils::apply_window(x, &window);
|
|
let y_win = utils::apply_window(y, &window);
|
|
|
|
let x_fft = utils::rfft(&x_win, n_fft);
|
|
let y_fft = utils::rfft(&y_win, n_fft);
|
|
|
|
let sxy = utils::cross_spectral_density(&x_fft, &y_fft);
|
|
let sxx = utils::power_spectral_density(&x_fft);
|
|
let syy = utils::power_spectral_density(&y_fft);
|
|
|
|
// Imaginary coherence: Im(Sxy)² / (Sxx * Syy)
|
|
// Clamp to [0, 1]
|
|
let imcoh: Vec<f64> = sxy
|
|
.iter()
|
|
.zip(sxx.iter())
|
|
.zip(syy.iter())
|
|
.map(|((xy, &xx), &yy)| {
|
|
let denom = xx * yy;
|
|
if denom > 1e-15 {
|
|
((xy.im * xy.im) / denom).clamp(0.0, 1.0)
|
|
} else {
|
|
0.0
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
let freqs = utils::fft_freqs(sfreq, n_fft, 0.0, sfreq / 2.0);
|
|
|
|
Ok((imcoh, freqs))
|
|
}
|
|
|
|
/// Compute coherence across all channel pairs for epoched data
|
|
pub fn coherence_all_pairs(
|
|
epochs: &[Vec<Vec<f64>>],
|
|
sfreq: f64,
|
|
fmin: f64,
|
|
fmax: f64,
|
|
n_fft: usize,
|
|
) -> ConnectivityResult<Vec<Vec<f64>>> {
|
|
if epochs.is_empty() {
|
|
return Err(ConnectivityError::InsufficientData(
|
|
"No epochs provided".to_string(),
|
|
));
|
|
}
|
|
|
|
let n_channels = epochs[0].len();
|
|
let _n_epochs = epochs.len();
|
|
let freq_indices = utils::freq_indices(sfreq, n_fft, fmin, fmax);
|
|
let _n_freqs = freq_indices.len();
|
|
|
|
// All channel pairs
|
|
let mut pairs: Vec<(usize, usize)> = Vec::new();
|
|
for i in 0..n_channels {
|
|
for j in (i + 1)..n_channels {
|
|
pairs.push((i, j));
|
|
}
|
|
}
|
|
|
|
let window = utils::hanning_window(epochs[0][0].len());
|
|
|
|
// Accumulate cross and auto spectra across epochs
|
|
let mut sxy_sum: Vec<Vec<Complex64>> =
|
|
vec![vec![Complex64::new(0.0, 0.0); n_fft / 2 + 1]; pairs.len()];
|
|
let mut sxx_sum: Vec<Vec<f64>> = vec![vec![0.0; n_fft / 2 + 1]; n_channels];
|
|
let _syy_sum: Vec<Vec<f64>> = vec![vec![0.0; n_fft / 2 + 1]; n_channels];
|
|
|
|
for epoch in epochs {
|
|
// FFT all channels
|
|
let ffts: Vec<Vec<Complex64>> = epoch
|
|
.iter()
|
|
.map(|ch| {
|
|
let windowed = utils::apply_window(ch, &window);
|
|
utils::rfft(&windowed, n_fft)
|
|
})
|
|
.collect();
|
|
|
|
// Accumulate auto-spectra
|
|
for (ch, fft) in ffts.iter().enumerate() {
|
|
for (i, c) in fft.iter().enumerate() {
|
|
sxx_sum[ch][i] += c.norm_sqr();
|
|
}
|
|
}
|
|
|
|
// Accumulate cross-spectra
|
|
for (p, &(i, j)) in pairs.iter().enumerate() {
|
|
for (k, (a, b)) in ffts[i].iter().zip(&ffts[j]).enumerate() {
|
|
sxy_sum[p][k] += a * b.conj();
|
|
}
|
|
}
|
|
}
|
|
|
|
// Compute coherence
|
|
let mut result = Vec::with_capacity(pairs.len());
|
|
|
|
for (p, &(i, j)) in pairs.iter().enumerate() {
|
|
let coh: Vec<f64> = freq_indices
|
|
.iter()
|
|
.map(|&k| {
|
|
let denom = sxx_sum[i][k] * sxx_sum[j][k];
|
|
if denom > 1e-15 {
|
|
(sxy_sum[p][k].norm_sqr() / denom).clamp(0.0, 1.0)
|
|
} else {
|
|
0.0
|
|
}
|
|
})
|
|
.collect();
|
|
result.push(coh);
|
|
}
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
/// Compute imaginary coherence across all channel pairs
|
|
pub fn imag_coherence_all_pairs(
|
|
epochs: &[Vec<Vec<f64>>],
|
|
sfreq: f64,
|
|
fmin: f64,
|
|
fmax: f64,
|
|
n_fft: usize,
|
|
) -> ConnectivityResult<Vec<Vec<f64>>> {
|
|
if epochs.is_empty() {
|
|
return Err(ConnectivityError::InsufficientData(
|
|
"No epochs provided".to_string(),
|
|
));
|
|
}
|
|
|
|
let n_channels = epochs[0].len();
|
|
let freq_indices = utils::freq_indices(sfreq, n_fft, fmin, fmax);
|
|
|
|
let mut pairs: Vec<(usize, usize)> = Vec::new();
|
|
for i in 0..n_channels {
|
|
for j in (i + 1)..n_channels {
|
|
pairs.push((i, j));
|
|
}
|
|
}
|
|
|
|
let window = utils::hanning_window(epochs[0][0].len());
|
|
|
|
let mut sxy_sum: Vec<Vec<Complex64>> =
|
|
vec![vec![Complex64::new(0.0, 0.0); n_fft / 2 + 1]; pairs.len()];
|
|
let mut sxx_sum: Vec<Vec<f64>> = vec![vec![0.0; n_fft / 2 + 1]; n_channels];
|
|
|
|
for epoch in epochs {
|
|
let ffts: Vec<Vec<Complex64>> = epoch
|
|
.iter()
|
|
.map(|ch| {
|
|
let windowed = utils::apply_window(ch, &window);
|
|
utils::rfft(&windowed, n_fft)
|
|
})
|
|
.collect();
|
|
|
|
for (ch, fft) in ffts.iter().enumerate() {
|
|
for (i, c) in fft.iter().enumerate() {
|
|
sxx_sum[ch][i] += c.norm_sqr();
|
|
}
|
|
}
|
|
|
|
for (p, &(i, j)) in pairs.iter().enumerate() {
|
|
for (k, (a, b)) in ffts[i].iter().zip(&ffts[j]).enumerate() {
|
|
sxy_sum[p][k] += a * b.conj();
|
|
}
|
|
}
|
|
}
|
|
|
|
let mut result = Vec::with_capacity(pairs.len());
|
|
|
|
for (p, &(i, j)) in pairs.iter().enumerate() {
|
|
let imcoh: Vec<f64> = freq_indices
|
|
.iter()
|
|
.map(|&k| {
|
|
let denom = sxx_sum[i][k] * sxx_sum[j][k];
|
|
if denom > 1e-15 {
|
|
let im = sxy_sum[p][k].im;
|
|
((im * im) / denom).clamp(0.0, 1.0)
|
|
} else {
|
|
0.0
|
|
}
|
|
})
|
|
.collect();
|
|
result.push(imcoh);
|
|
}
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::f64::consts::PI;
|
|
|
|
fn create_test_signals() -> (Vec<f64>, Vec<f64>) {
|
|
let n = 500;
|
|
let sfreq = 100.0;
|
|
|
|
// Two signals with known coherence
|
|
let x: Vec<f64> = (0..n)
|
|
.map(|i| {
|
|
let t = i as f64 / sfreq;
|
|
(2.0 * PI * 10.0 * t).sin() + 0.5 * (2.0 * PI * 20.0 * t).sin()
|
|
})
|
|
.collect();
|
|
|
|
let y: Vec<f64> = (0..n)
|
|
.map(|i| {
|
|
let t = i as f64 / sfreq;
|
|
0.8 * (2.0 * PI * 10.0 * t).sin() + 0.3 * (2.0 * PI * 30.0 * t).sin()
|
|
})
|
|
.collect();
|
|
|
|
(x, y)
|
|
}
|
|
|
|
#[test]
|
|
fn test_coherence() {
|
|
let (x, y) = create_test_signals();
|
|
|
|
let (coh, freqs) = coherence(&x, &y, 100.0, None).unwrap();
|
|
|
|
// Should have positive coherence
|
|
assert!(coh.len() > 0);
|
|
assert!(coh.iter().all(|&c| c >= 0.0 && c <= 1.0));
|
|
|
|
// Coherence at 10 Hz should be higher (shared component)
|
|
let idx_10hz = freqs.iter().position(|&f| (f - 10.0).abs() < 1.0).unwrap();
|
|
assert!(coh[idx_10hz] > 0.5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_perfect_coherence() {
|
|
let n = 500;
|
|
// Use a pure sinusoid at a specific frequency
|
|
let x: Vec<f64> = (0..n)
|
|
.map(|i| (2.0 * PI * 10.0 * i as f64 / 100.0).sin())
|
|
.collect();
|
|
let y = x.clone(); // Perfect copy
|
|
|
|
let (coh, freqs) = coherence(&x, &y, 100.0, None).unwrap();
|
|
|
|
// Coherence at the signal frequency should be very high
|
|
let idx_10hz = freqs.iter().position(|&f| (f - 10.0).abs() < 1.0).unwrap();
|
|
assert!(
|
|
coh[idx_10hz] > 0.95,
|
|
"Coherence at signal frequency should be ~1.0, got {}",
|
|
coh[idx_10hz]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_imaginary_coherence() {
|
|
let (x, y) = create_test_signals();
|
|
|
|
let (imcoh, _) = imaginary_coherence(&x, &y, 100.0, None).unwrap();
|
|
|
|
// Imaginary coherence should be bounded [0, 1]
|
|
assert!(imcoh.iter().all(|&c| c >= 0.0 && c <= 1.0));
|
|
}
|
|
|
|
#[test]
|
|
fn test_coherence_all_pairs() {
|
|
let n = 200;
|
|
let sfreq = 100.0;
|
|
|
|
// Create 3-channel epochs
|
|
let epochs: Vec<Vec<Vec<f64>>> = (0..5)
|
|
.map(|_| {
|
|
(0..3)
|
|
.map(|ch| {
|
|
(0..n)
|
|
.map(|i| {
|
|
let t = i as f64 / sfreq;
|
|
(2.0 * PI * (10.0 + ch as f64) * t).sin()
|
|
})
|
|
.collect()
|
|
})
|
|
.collect()
|
|
})
|
|
.collect();
|
|
|
|
let result = coherence_all_pairs(&epochs, sfreq, 5.0, 30.0, n).unwrap();
|
|
|
|
// 3 channels = 3 pairs
|
|
assert_eq!(result.len(), 3);
|
|
// All coherence values should be valid
|
|
assert!(
|
|
result
|
|
.iter()
|
|
.all(|pair| pair.iter().all(|&c| c >= 0.0 && c <= 1.0))
|
|
);
|
|
}
|
|
}
|