Initial commit

crates/specialized/rtx-neuro-connectivity/src/coherence.rs

@@ -0,0 +1,413 @@
+//! 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))
+        );
+    }
+}