rustytorch/crates/specialized/rtx-neuro-connectivity/src/utils.rs

//! 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);
    }
}