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

//! # rtx-neuro-connectivity
//!
//! Functional and effective connectivity analysis for MEG/EEG data.
//!
//! ## Features
//!
//! - **Spectral Connectivity**: Coherence, imaginary coherence
//! - **Phase Connectivity**: PLV, wPLI, dwPLI
//! - **Effective Connectivity**: Granger causality (coming soon)
//! - **Cross-Frequency**: Phase-amplitude coupling
//!
//! ## Example
//!
//! ```ignore
//! use rtx_neuro_connectivity::{spectral_connectivity, ConnectivityMethod};
//!
//! let conn = spectral_connectivity(&epochs, ConnectivityMethod::Coherence, sfreq)?;
//! ```

#![warn(missing_docs)]

pub mod coherence;
pub mod granger;
pub mod pac;
pub mod plv;
pub mod utils;
pub mod wpli;

// Re-export main types
pub use coherence::{coherence, coherence_pairs, imaginary_coherence};
pub use granger::{
    GrangerConfig, GrangerResult, SpectralGrangerResult, VarModel, granger_causality,
    spectral_granger,
};
pub use pac::{PacMethod, phase_amplitude_coupling};
pub use plv::{ciplv, plv, plv_epochs, wplv};
pub use wpli::{dwpli, wpli};

/// Connectivity analysis error types
#[derive(Debug, thiserror::Error)]
pub enum ConnectivityError {
    /// Invalid parameters
    #[error("Invalid parameters: {0}")]
    InvalidParameters(String),

    /// Dimension mismatch
    #[error("Dimension mismatch: {0}")]
    DimensionMismatch(String),

    /// Computation error
    #[error("Computation error: {0}")]
    ComputationError(String),

    /// Insufficient data
    #[error("Insufficient data: {0}")]
    InsufficientData(String),
}

/// Result type for connectivity operations
pub type ConnectivityResult<T> = Result<T, ConnectivityError>;

/// Connectivity method enumeration
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConnectivityMethod {
    /// Magnitude-squared coherence
    Coherence,
    /// Imaginary coherence (robust to volume conduction)
    ImaginaryCoherence,
    /// Phase locking value
    Plv,
    /// Weighted phase lag index
    Wpli,
    /// Debiased weighted phase lag index
    DwPli,
    /// Pairwise phase consistency
    Ppc,
}

/// Result of connectivity computation
#[derive(Debug, Clone)]
pub struct ConnectivityResult2D {
    /// Connectivity matrix [n_pairs x n_freqs]
    pub data: Vec<Vec<f64>>,
    /// Frequency vector
    pub freqs: Vec<f64>,
    /// Source indices for pairs
    pub sources: Vec<usize>,
    /// Target indices for pairs
    pub targets: Vec<usize>,
    /// Method used
    pub method: ConnectivityMethod,
    /// Number of epochs used
    pub n_epochs: usize,
}

impl ConnectivityResult2D {
    /// Get connectivity for a specific pair
    pub fn get_pair(&self, source: usize, target: usize) -> Option<&[f64]> {
        for (i, (&s, &t)) in self.sources.iter().zip(&self.targets).enumerate() {
            if s == source && t == target {
                return Some(&self.data[i]);
            }
        }
        None
    }

    /// Convert to symmetric matrix at a specific frequency
    pub fn to_matrix(&self, freq_idx: usize, n_channels: usize) -> Vec<Vec<f64>> {
        let mut matrix = vec![vec![0.0; n_channels]; n_channels];

        for (i, (&s, &t)) in self.sources.iter().zip(&self.targets).enumerate() {
            if freq_idx < self.data[i].len() {
                matrix[s][t] = self.data[i][freq_idx];
                matrix[t][s] = self.data[i][freq_idx]; // Symmetric
            }
        }

        // Diagonal is 1 for most methods
        for i in 0..n_channels {
            matrix[i][i] = 1.0;
        }

        matrix
    }

    /// Get mean connectivity across frequencies
    pub fn mean_connectivity(&self) -> Vec<f64> {
        self.data
            .iter()
            .map(|pair| pair.iter().sum::<f64>() / pair.len() as f64)
            .collect()
    }
}

/// Compute spectral connectivity across all channel pairs
pub fn spectral_connectivity(
    epochs: &[Vec<Vec<f64>>], // [n_epochs][n_channels][n_times]
    method: ConnectivityMethod,
    sfreq: f64,
    fmin: f64,
    fmax: f64,
    n_fft: Option<usize>,
) -> ConnectivityResult<ConnectivityResult2D> {
    if epochs.is_empty() {
        return Err(ConnectivityError::InsufficientData(
            "No epochs provided".to_string(),
        ));
    }

    let n_channels = epochs[0].len();
    let n_times = epochs[0][0].len();
    let n_fft = n_fft.unwrap_or(n_times);

    // Generate all pairs
    let mut sources = Vec::new();
    let mut targets = Vec::new();
    for i in 0..n_channels {
        for j in (i + 1)..n_channels {
            sources.push(i);
            targets.push(j);
        }
    }

    // Compute connectivity for each pair
    let data: Vec<Vec<f64>> = match method {
        ConnectivityMethod::Coherence => {
            coherence::coherence_all_pairs(epochs, sfreq, fmin, fmax, n_fft)?
        }
        ConnectivityMethod::ImaginaryCoherence => {
            coherence::imag_coherence_all_pairs(epochs, sfreq, fmin, fmax, n_fft)?
        }
        ConnectivityMethod::Plv => plv::plv_all_pairs(epochs, sfreq, fmin, fmax, n_fft)?,
        ConnectivityMethod::Wpli => wpli::wpli_all_pairs(epochs, sfreq, fmin, fmax, n_fft)?,
        ConnectivityMethod::DwPli => wpli::dwpli_all_pairs(epochs, sfreq, fmin, fmax, n_fft)?,
        ConnectivityMethod::Ppc => plv::ppc_all_pairs(epochs, sfreq, fmin, fmax, n_fft)?,
    };

    // Compute frequency vector
    let freqs = utils::fft_freqs(sfreq, n_fft, fmin, fmax);

    Ok(ConnectivityResult2D {
        data,
        freqs,
        sources,
        targets,
        method,
        n_epochs: epochs.len(),
    })
}