Initial commit

2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
@@ -0,0 +1,201 @@
+//! Utility functions for connectivity analysis.
+
+use num_complex::Complex64;
+use rustfft::{FftDirection, FftPlanner};
+
+/// Compute FFT frequencies for a given sample rate
+pub fn fft_freqs(sfreq: f64, n_fft: usize, fmin: f64, fmax: f64) -> Vec<f64> {
+    let freq_resolution = sfreq / n_fft as f64;
+    let n_freqs = n_fft / 2 + 1;
+
+    (0..n_freqs)
+        .map(|i| i as f64 * freq_resolution)
+        .filter(|&f| f >= fmin && f <= fmax)
+        .collect()
+}
+
+/// Get frequency indices for a given range
+pub fn freq_indices(sfreq: f64, n_fft: usize, fmin: f64, fmax: f64) -> Vec<usize> {
+    let freq_resolution = sfreq / n_fft as f64;
+    let n_freqs = n_fft / 2 + 1;
+
+    (0..n_freqs)
+        .filter(|&i| {
+            let f = i as f64 * freq_resolution;
+            f >= fmin && f <= fmax
+        })
+        .collect()
+}
+
+/// Compute FFT of a real signal
+pub fn rfft(signal: &[f64], n_fft: usize) -> Vec<Complex64> {
+    let mut planner = FftPlanner::new();
+    let fft = planner.plan_fft(n_fft, FftDirection::Forward);
+
+    // Zero-pad signal
+    let mut input: Vec<Complex64> = signal
+        .iter()
+        .map(|&x| Complex64::new(x, 0.0))
+        .chain(std::iter::repeat(Complex64::new(0.0, 0.0)))
+        .take(n_fft)
+        .collect();
+
+    fft.process(&mut input);
+
+    // Return only positive frequencies
+    input.truncate(n_fft / 2 + 1);
+    input
+}
+
+/// Compute cross-spectral density
+pub fn cross_spectral_density(x: &[Complex64], y: &[Complex64]) -> Vec<Complex64> {
+    x.iter().zip(y).map(|(a, b)| a * b.conj()).collect()
+}
+
+/// Compute power spectral density
+pub fn power_spectral_density(x: &[Complex64]) -> Vec<f64> {
+    x.iter().map(num_complex::Complex::norm_sqr).collect()
+}
+
+/// Apply Hanning window
+pub fn hanning_window(n: usize) -> Vec<f64> {
+    use std::f64::consts::PI;
+    (0..n)
+        .map(|i| 0.5 * (1.0 - (2.0 * PI * i as f64 / (n - 1) as f64).cos()))
+        .collect()
+}
+
+/// Apply window to signal
+pub fn apply_window(signal: &[f64], window: &[f64]) -> Vec<f64> {
+    signal.iter().zip(window).map(|(&s, &w)| s * w).collect()
+}
+
+/// Hilbert transform to get analytic signal
+pub fn hilbert(signal: &[f64]) -> Vec<Complex64> {
+    let n = signal.len();
+    let mut planner = FftPlanner::new();
+    let fft = planner.plan_fft(n, FftDirection::Forward);
+    let ifft = planner.plan_fft(n, FftDirection::Inverse);
+
+    // Forward FFT
+    let mut spectrum: Vec<Complex64> = signal.iter().map(|&x| Complex64::new(x, 0.0)).collect();
+    fft.process(&mut spectrum);
+
+    // Create analytic signal in frequency domain
+    // H(f) = 2 for f > 0, H(0) = 1, H(f) = 0 for f < 0
+    let half = (n + 1) / 2;
+    for i in 1..half {
+        spectrum[i] *= 2.0;
+    }
+    for i in half..n {
+        spectrum[i] = Complex64::new(0.0, 0.0);
+    }
+
+    // Inverse FFT
+    ifft.process(&mut spectrum);
+
+    // Normalize
+    let scale = 1.0 / n as f64;
+    spectrum.iter_mut().for_each(|c| *c *= scale);
+
+    spectrum
+}
+
+/// Extract instantaneous phase from analytic signal
+pub fn instantaneous_phase(analytic: &[Complex64]) -> Vec<f64> {
+    analytic.iter().map(|c| c.arg()).collect()
+}
+
+/// Bandpass filter using FFT
+pub fn bandpass_fft(signal: &[f64], sfreq: f64, fmin: f64, fmax: f64) -> Vec<f64> {
+    let n = signal.len();
+    let mut planner = FftPlanner::new();
+    let fft = planner.plan_fft(n, FftDirection::Forward);
+    let ifft = planner.plan_fft(n, FftDirection::Inverse);
+
+    // Forward FFT
+    let mut spectrum: Vec<Complex64> = signal.iter().map(|&x| Complex64::new(x, 0.0)).collect();
+    fft.process(&mut spectrum);
+
+    // Apply bandpass
+    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
+        };
+
+        if freq < fmin || freq > fmax {
+            spectrum[i] = Complex64::new(0.0, 0.0);
+        }
+    }
+
+    // Inverse FFT
+    ifft.process(&mut spectrum);
+
+    // Return real part, normalized
+    let scale = 1.0 / n as f64;
+    spectrum.iter().map(|c| c.re * scale).collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::f64::consts::PI;
+
+    #[test]
+    fn test_fft_freqs() {
+        let freqs = fft_freqs(100.0, 100, 0.0, 50.0);
+        assert_eq!(freqs[0], 0.0);
+        assert_eq!(freqs[1], 1.0);
+        assert_eq!(freqs.len(), 51);
+    }
+
+    #[test]
+    fn test_hanning_window() {
+        let window = hanning_window(5);
+        assert!((window[0] - 0.0).abs() < 1e-10);
+        assert!((window[2] - 1.0).abs() < 1e-10);
+        assert!((window[4] - 0.0).abs() < 1e-10);
+    }
+
+    #[test]
+    fn test_rfft() {
+        // Test with simple sinusoid
+        let n = 100;
+        let signal: Vec<f64> = (0..n)
+            .map(|i| (2.0 * PI * 10.0 * i as f64 / n as f64).sin())
+            .collect();
+
+        let spectrum = rfft(&signal, n);
+
+        // Peak should be at index 10
+        let magnitudes: Vec<f64> = spectrum.iter().map(|c| c.norm()).collect();
+        let peak_idx = magnitudes
+            .iter()
+            .enumerate()
+            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
+            .map(|(i, _)| i)
+            .unwrap();
+
+        assert_eq!(peak_idx, 10);
+    }
+
+    #[test]
+    fn test_hilbert() {
+        let n = 100;
+        let signal: Vec<f64> = (0..n)
+            .map(|i| (2.0 * PI * 5.0 * i as f64 / n as f64).sin())
+            .collect();
+
+        let analytic = hilbert(&signal);
+
+        // Magnitude should be approximately constant (envelope of pure sinusoid)
+        let magnitudes: Vec<f64> = analytic.iter().map(|c| c.norm()).collect();
+        let mean_mag = magnitudes.iter().sum::<f64>() / n as f64;
+
+        // Should be close to 1
+        assert!((mean_mag - 1.0).abs() < 0.2);
+    }
+}