754 lines
23 KiB
Rust
754 lines
23 KiB
Rust
//! Brain graph construction from connectivity matrices.
|
|
//!
|
|
//! This module provides the bridge between connectivity analysis (coherence, PLV, etc.)
|
|
//! and graph neural networks.
|
|
|
|
use crate::error::{GnnError, GnnResult};
|
|
use ndarray::Array2;
|
|
use rtx_neuro_connectivity::ConnectivityResult2D;
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// Brain hemisphere
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
|
pub enum Hemisphere {
|
|
/// Left hemisphere
|
|
Left,
|
|
/// Right hemisphere
|
|
Right,
|
|
/// Midline (Fz, Cz, Pz, etc.)
|
|
Midline,
|
|
/// Unknown or unspecified
|
|
Unknown,
|
|
}
|
|
|
|
impl Hemisphere {
|
|
/// Parse from channel name (e.g., "Fp1" -> Left, "Fp2" -> Right)
|
|
pub fn from_channel_name(name: &str) -> Self {
|
|
let name = name.trim();
|
|
if name.ends_with('z') || name.ends_with('Z') {
|
|
Hemisphere::Midline
|
|
} else if let Some(last) = name.chars().last() {
|
|
if last.is_ascii_digit() {
|
|
let digit = last.to_digit(10).unwrap_or(0);
|
|
if digit % 2 == 1 {
|
|
Hemisphere::Left
|
|
} else {
|
|
Hemisphere::Right
|
|
}
|
|
} else {
|
|
Hemisphere::Unknown
|
|
}
|
|
} else {
|
|
Hemisphere::Unknown
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Brain region classification
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
|
pub enum BrainRegion {
|
|
/// Frontal lobe (Fp, F)
|
|
Frontal,
|
|
/// Central (C)
|
|
Central,
|
|
/// Temporal (T)
|
|
Temporal,
|
|
/// Parietal (P)
|
|
Parietal,
|
|
/// Occipital (O)
|
|
Occipital,
|
|
/// Unknown region
|
|
Unknown,
|
|
}
|
|
|
|
impl BrainRegion {
|
|
/// Parse from channel name (e.g., "Fp1" -> Frontal, "O2" -> Occipital)
|
|
pub fn from_channel_name(name: &str) -> Self {
|
|
let name = name.to_uppercase();
|
|
if name.starts_with("FP") || name.starts_with("AF") {
|
|
BrainRegion::Frontal
|
|
} else if name.starts_with('F') {
|
|
BrainRegion::Frontal
|
|
} else if name.starts_with('C') {
|
|
BrainRegion::Central
|
|
} else if name.starts_with('T') {
|
|
BrainRegion::Temporal
|
|
} else if name.starts_with('P') && !name.starts_with("PO") {
|
|
BrainRegion::Parietal
|
|
} else if name.starts_with('O') || name.starts_with("PO") {
|
|
BrainRegion::Occipital
|
|
} else {
|
|
BrainRegion::Unknown
|
|
}
|
|
}
|
|
|
|
/// Get approximate 3D coordinates for visualization (normalized)
|
|
pub fn approx_position(&self) -> [f32; 3] {
|
|
match self {
|
|
BrainRegion::Frontal => [0.0, 0.8, 0.3],
|
|
BrainRegion::Central => [0.0, 0.0, 0.5],
|
|
BrainRegion::Temporal => [0.9, 0.0, 0.0],
|
|
BrainRegion::Parietal => [0.0, -0.5, 0.5],
|
|
BrainRegion::Occipital => [0.0, -0.9, 0.1],
|
|
BrainRegion::Unknown => [0.0, 0.0, 0.0],
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A node in the brain graph (represents a channel/electrode)
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct BrainNode {
|
|
/// Node index
|
|
pub index: usize,
|
|
/// Channel name (e.g., "Fp1", "O2")
|
|
pub name: String,
|
|
/// Hemisphere
|
|
pub hemisphere: Hemisphere,
|
|
/// Brain region
|
|
pub region: BrainRegion,
|
|
/// 3D position [x, y, z] in normalized coordinates
|
|
pub position: [f32; 3],
|
|
/// Node features (e.g., power spectrum, time series features)
|
|
pub features: Vec<f64>,
|
|
}
|
|
|
|
impl BrainNode {
|
|
/// Create a new brain node
|
|
pub fn new(index: usize, name: &str) -> Self {
|
|
Self {
|
|
index,
|
|
name: name.to_string(),
|
|
hemisphere: Hemisphere::from_channel_name(name),
|
|
region: BrainRegion::from_channel_name(name),
|
|
position: [0.0, 0.0, 0.0],
|
|
features: Vec::new(),
|
|
}
|
|
}
|
|
|
|
/// Create with explicit hemisphere and region
|
|
pub fn with_location(
|
|
index: usize,
|
|
name: &str,
|
|
hemisphere: Hemisphere,
|
|
region: BrainRegion,
|
|
) -> Self {
|
|
let mut node = Self::new(index, name);
|
|
node.hemisphere = hemisphere;
|
|
node.region = region;
|
|
|
|
// Compute approximate position from region and hemisphere
|
|
let base_pos = region.approx_position();
|
|
node.position = match hemisphere {
|
|
Hemisphere::Left => [-0.5 + base_pos[0], base_pos[1], base_pos[2]],
|
|
Hemisphere::Right => [0.5 + base_pos[0], base_pos[1], base_pos[2]],
|
|
Hemisphere::Midline => [0.0, base_pos[1], base_pos[2]],
|
|
Hemisphere::Unknown => base_pos,
|
|
};
|
|
|
|
node
|
|
}
|
|
|
|
/// Set node features
|
|
pub fn with_features(mut self, features: Vec<f64>) -> Self {
|
|
self.features = features;
|
|
self
|
|
}
|
|
}
|
|
|
|
/// An edge in the brain graph (represents connectivity between channels)
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct BrainEdge {
|
|
/// Source node index
|
|
pub source: usize,
|
|
/// Target node index
|
|
pub target: usize,
|
|
/// Connectivity weight (e.g., coherence, PLV)
|
|
pub weight: f64,
|
|
/// Edge features (e.g., multi-frequency connectivity)
|
|
pub features: Vec<f64>,
|
|
/// Whether this is an interhemispheric connection
|
|
pub interhemispheric: bool,
|
|
}
|
|
|
|
impl BrainEdge {
|
|
/// Create a new brain edge
|
|
pub fn new(source: usize, target: usize, weight: f64) -> Self {
|
|
Self {
|
|
source,
|
|
target,
|
|
weight,
|
|
features: vec![weight],
|
|
interhemispheric: false,
|
|
}
|
|
}
|
|
|
|
/// Add multi-frequency features
|
|
pub fn with_frequency_features(mut self, features: Vec<f64>) -> Self {
|
|
self.features = features;
|
|
self
|
|
}
|
|
|
|
/// Mark as interhemispheric
|
|
pub fn with_interhemispheric(mut self, flag: bool) -> Self {
|
|
self.interhemispheric = flag;
|
|
self
|
|
}
|
|
}
|
|
|
|
/// A brain connectivity graph
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct BrainGraph {
|
|
/// Nodes (channels/electrodes)
|
|
pub nodes: Vec<BrainNode>,
|
|
/// Edges (connectivity)
|
|
pub edges: Vec<BrainEdge>,
|
|
/// Adjacency matrix (dense)
|
|
adjacency: Option<Array2<f64>>,
|
|
/// Metadata
|
|
pub metadata: GraphMetadata,
|
|
}
|
|
|
|
/// Graph metadata
|
|
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
|
pub struct GraphMetadata {
|
|
/// Subject ID
|
|
pub subject_id: Option<String>,
|
|
/// Condition/task
|
|
pub condition: Option<String>,
|
|
/// Time window [start, end] in seconds
|
|
pub time_window: Option<[f64; 2]>,
|
|
/// Frequency band [fmin, fmax] in Hz
|
|
pub freq_band: Option<[f64; 2]>,
|
|
/// Connectivity method used
|
|
pub connectivity_method: Option<String>,
|
|
/// Label for classification
|
|
pub label: Option<i64>,
|
|
/// Label name (e.g., "Control", "Patient")
|
|
pub label_name: Option<String>,
|
|
}
|
|
|
|
impl BrainGraph {
|
|
/// Create an empty brain graph
|
|
pub fn new() -> Self {
|
|
Self {
|
|
nodes: Vec::new(),
|
|
edges: Vec::new(),
|
|
adjacency: None,
|
|
metadata: GraphMetadata::default(),
|
|
}
|
|
}
|
|
|
|
/// Number of nodes
|
|
pub fn n_nodes(&self) -> usize {
|
|
self.nodes.len()
|
|
}
|
|
|
|
/// Number of edges
|
|
pub fn n_edges(&self) -> usize {
|
|
self.edges.len()
|
|
}
|
|
|
|
/// Create from a connectivity matrix
|
|
pub fn from_connectivity_matrix(
|
|
matrix: &[Vec<f64>],
|
|
channel_names: &[&str],
|
|
threshold: f64,
|
|
) -> GnnResult<Self> {
|
|
let n = matrix.len();
|
|
if n == 0 {
|
|
return Err(GnnError::InvalidGraph("Empty connectivity matrix".into()));
|
|
}
|
|
if channel_names.len() != n {
|
|
return Err(GnnError::DimensionMismatch(format!(
|
|
"Matrix size {} != channel names {}",
|
|
n,
|
|
channel_names.len()
|
|
)));
|
|
}
|
|
|
|
let mut graph = BrainGraph::new();
|
|
|
|
// Add nodes
|
|
for (i, name) in channel_names.iter().enumerate() {
|
|
graph.nodes.push(BrainNode::new(i, name));
|
|
}
|
|
|
|
// Add edges (above threshold)
|
|
for i in 0..n {
|
|
for j in 0..n {
|
|
if i != j && matrix[i][j] >= threshold {
|
|
let is_inter = graph.nodes[i].hemisphere != graph.nodes[j].hemisphere
|
|
&& graph.nodes[i].hemisphere != Hemisphere::Midline
|
|
&& graph.nodes[j].hemisphere != Hemisphere::Midline;
|
|
|
|
graph
|
|
.edges
|
|
.push(BrainEdge::new(i, j, matrix[i][j]).with_interhemispheric(is_inter));
|
|
}
|
|
}
|
|
}
|
|
|
|
// Store adjacency
|
|
let mut adj = Array2::zeros((n, n));
|
|
for edge in &graph.edges {
|
|
adj[[edge.source, edge.target]] = edge.weight;
|
|
}
|
|
graph.adjacency = Some(adj);
|
|
|
|
Ok(graph)
|
|
}
|
|
|
|
/// Create from ConnectivityResult2D at a specific frequency
|
|
pub fn from_connectivity_result(
|
|
result: &ConnectivityResult2D,
|
|
channel_names: &[&str],
|
|
freq_idx: usize,
|
|
threshold: f64,
|
|
) -> GnnResult<Self> {
|
|
let n_channels = channel_names.len();
|
|
let matrix = result.to_matrix(freq_idx, n_channels);
|
|
|
|
let mut graph = Self::from_connectivity_matrix(&matrix, channel_names, threshold)?;
|
|
|
|
// Add metadata
|
|
if let Some(&freq) = result.freqs.get(freq_idx) {
|
|
graph.metadata.freq_band = Some([freq, freq]);
|
|
}
|
|
graph.metadata.connectivity_method = Some(format!("{:?}", result.method));
|
|
|
|
Ok(graph)
|
|
}
|
|
|
|
/// Create multi-frequency graph with edge features
|
|
pub fn from_connectivity_multifreq(
|
|
result: &ConnectivityResult2D,
|
|
channel_names: &[&str],
|
|
freq_indices: &[usize],
|
|
threshold: f64,
|
|
) -> GnnResult<Self> {
|
|
let n_channels = channel_names.len();
|
|
|
|
// Start with first frequency
|
|
let matrix0 = result.to_matrix(freq_indices[0], n_channels);
|
|
let mut graph = Self::from_connectivity_matrix(&matrix0, channel_names, threshold)?;
|
|
|
|
// Add multi-frequency features to edges
|
|
for edge in &mut graph.edges {
|
|
let mut features = Vec::with_capacity(freq_indices.len());
|
|
for &freq_idx in freq_indices {
|
|
let matrix = result.to_matrix(freq_idx, n_channels);
|
|
features.push(matrix[edge.source][edge.target]);
|
|
}
|
|
edge.features = features;
|
|
}
|
|
|
|
// Update metadata
|
|
if !freq_indices.is_empty() {
|
|
let fmin = result.freqs.get(freq_indices[0]).copied().unwrap_or(0.0);
|
|
let fmax = result
|
|
.freqs
|
|
.get(*freq_indices.last().unwrap())
|
|
.copied()
|
|
.unwrap_or(0.0);
|
|
graph.metadata.freq_band = Some([fmin, fmax]);
|
|
}
|
|
|
|
Ok(graph)
|
|
}
|
|
|
|
/// Get adjacency matrix
|
|
pub fn adjacency_matrix(&self) -> Array2<f64> {
|
|
if let Some(ref adj) = self.adjacency {
|
|
adj.clone()
|
|
} else {
|
|
let n = self.n_nodes();
|
|
let mut adj = Array2::zeros((n, n));
|
|
for edge in &self.edges {
|
|
adj[[edge.source, edge.target]] = edge.weight;
|
|
}
|
|
adj
|
|
}
|
|
}
|
|
|
|
/// Get degree of each node
|
|
pub fn degrees(&self) -> Vec<f64> {
|
|
let adj = self.adjacency_matrix();
|
|
adj.sum_axis(ndarray::Axis(1)).to_vec()
|
|
}
|
|
|
|
/// Get node feature matrix [n_nodes, n_features]
|
|
pub fn node_feature_matrix(&self) -> GnnResult<Array2<f64>> {
|
|
if self.nodes.is_empty() {
|
|
return Err(GnnError::InvalidGraph("No nodes".into()));
|
|
}
|
|
|
|
let n_features = self.nodes[0].features.len();
|
|
if n_features == 0 {
|
|
// Use degree as default feature
|
|
let degrees = self.degrees();
|
|
let mut features = Array2::zeros((self.n_nodes(), 1));
|
|
for (i, &d) in degrees.iter().enumerate() {
|
|
features[[i, 0]] = d;
|
|
}
|
|
return Ok(features);
|
|
}
|
|
|
|
let mut features = Array2::zeros((self.n_nodes(), n_features));
|
|
for (i, node) in self.nodes.iter().enumerate() {
|
|
if node.features.len() != n_features {
|
|
return Err(GnnError::DimensionMismatch(format!(
|
|
"Node {} has {} features, expected {}",
|
|
i,
|
|
node.features.len(),
|
|
n_features
|
|
)));
|
|
}
|
|
for (j, &f) in node.features.iter().enumerate() {
|
|
features[[i, j]] = f;
|
|
}
|
|
}
|
|
|
|
Ok(features)
|
|
}
|
|
|
|
/// Get edge index tensor [2, n_edges] for sparse GNN
|
|
pub fn edge_index(&self) -> (Vec<usize>, Vec<usize>) {
|
|
let sources: Vec<usize> = self.edges.iter().map(|e| e.source).collect();
|
|
let targets: Vec<usize> = self.edges.iter().map(|e| e.target).collect();
|
|
(sources, targets)
|
|
}
|
|
|
|
/// Get edge weights
|
|
pub fn edge_weights(&self) -> Vec<f64> {
|
|
self.edges.iter().map(|e| e.weight).collect()
|
|
}
|
|
|
|
/// Get edge feature matrix [n_edges, n_features]
|
|
pub fn edge_feature_matrix(&self) -> Array2<f64> {
|
|
if self.edges.is_empty() {
|
|
return Array2::zeros((0, 0));
|
|
}
|
|
|
|
let n_features = self.edges[0].features.len();
|
|
let mut features = Array2::zeros((self.n_edges(), n_features));
|
|
for (i, edge) in self.edges.iter().enumerate() {
|
|
for (j, &f) in edge.features.iter().enumerate() {
|
|
features[[i, j]] = f;
|
|
}
|
|
}
|
|
features
|
|
}
|
|
|
|
/// Set node features from power spectrum
|
|
pub fn set_node_features_from_psd(&mut self, psd: &[Vec<f64>]) -> GnnResult<()> {
|
|
if psd.len() != self.n_nodes() {
|
|
return Err(GnnError::DimensionMismatch(format!(
|
|
"PSD has {} channels, graph has {} nodes",
|
|
psd.len(),
|
|
self.n_nodes()
|
|
)));
|
|
}
|
|
|
|
for (node, features) in self.nodes.iter_mut().zip(psd.iter()) {
|
|
node.features = features.clone();
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Get subgraph for a specific brain region
|
|
pub fn subgraph_by_region(&self, region: BrainRegion) -> BrainGraph {
|
|
let node_indices: Vec<usize> = self
|
|
.nodes
|
|
.iter()
|
|
.filter(|n| n.region == region)
|
|
.map(|n| n.index)
|
|
.collect();
|
|
|
|
self.subgraph(&node_indices)
|
|
}
|
|
|
|
/// Get subgraph for a specific hemisphere
|
|
pub fn subgraph_by_hemisphere(&self, hemisphere: Hemisphere) -> BrainGraph {
|
|
let node_indices: Vec<usize> = self
|
|
.nodes
|
|
.iter()
|
|
.filter(|n| n.hemisphere == hemisphere)
|
|
.map(|n| n.index)
|
|
.collect();
|
|
|
|
self.subgraph(&node_indices)
|
|
}
|
|
|
|
/// Create a subgraph with only specified nodes
|
|
pub fn subgraph(&self, node_indices: &[usize]) -> BrainGraph {
|
|
let node_set: std::collections::HashSet<usize> = node_indices.iter().copied().collect();
|
|
let index_map: std::collections::HashMap<usize, usize> = node_indices
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(new, &old)| (old, new))
|
|
.collect();
|
|
|
|
let nodes: Vec<BrainNode> = node_indices
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(new_idx, &old_idx)| {
|
|
let mut node = self.nodes[old_idx].clone();
|
|
node.index = new_idx;
|
|
node
|
|
})
|
|
.collect();
|
|
|
|
let edges: Vec<BrainEdge> = self
|
|
.edges
|
|
.iter()
|
|
.filter(|e| node_set.contains(&e.source) && node_set.contains(&e.target))
|
|
.map(|e| BrainEdge {
|
|
source: *index_map.get(&e.source).unwrap(),
|
|
target: *index_map.get(&e.target).unwrap(),
|
|
..e.clone()
|
|
})
|
|
.collect();
|
|
|
|
BrainGraph {
|
|
nodes,
|
|
edges,
|
|
adjacency: None,
|
|
metadata: self.metadata.clone(),
|
|
}
|
|
}
|
|
|
|
/// Convert to rtx-geom Graph for GNN processing
|
|
pub fn to_geom_graph(&self) -> GnnResult<rtx_geom::Graph> {
|
|
use rtx_tensor::{Device, Tensor};
|
|
|
|
let device = Device::default();
|
|
let mut graph = rtx_geom::Graph::new();
|
|
|
|
// Add nodes and track their IDs
|
|
let mut node_ids = Vec::with_capacity(self.nodes.len());
|
|
for node in &self.nodes {
|
|
let features = if node.features.is_empty() {
|
|
// Default: one-hot region encoding
|
|
let mut feat = vec![0.0f32; 5];
|
|
match node.region {
|
|
BrainRegion::Frontal => feat[0] = 1.0,
|
|
BrainRegion::Central => feat[1] = 1.0,
|
|
BrainRegion::Temporal => feat[2] = 1.0,
|
|
BrainRegion::Parietal => feat[3] = 1.0,
|
|
BrainRegion::Occipital => feat[4] = 1.0,
|
|
BrainRegion::Unknown => {}
|
|
}
|
|
feat
|
|
} else {
|
|
node.features.iter().map(|&x| x as f32).collect()
|
|
};
|
|
|
|
let tensor = Tensor::from_data(features.clone(), [features.len()], &device)
|
|
.map_err(|e| GnnError::TensorError(e.to_string()))?;
|
|
|
|
let node_id = graph.add_node(rtx_geom::Node::new(tensor));
|
|
node_ids.push(node_id);
|
|
}
|
|
|
|
// Add edges using tracked node IDs
|
|
for edge in &self.edges {
|
|
let features: Vec<f32> = edge.features.iter().map(|&x| x as f32).collect();
|
|
let tensor = Tensor::from_data(features.clone(), [features.len()], &device)
|
|
.map_err(|e| GnnError::TensorError(e.to_string()))?;
|
|
|
|
let source_id = node_ids[edge.source];
|
|
let target_id = node_ids[edge.target];
|
|
|
|
graph
|
|
.add_edge(source_id, target_id, rtx_geom::Edge::new(tensor))
|
|
.map_err(|e| GnnError::GeomError(e.to_string()))?;
|
|
}
|
|
|
|
Ok(graph)
|
|
}
|
|
}
|
|
|
|
impl Default for BrainGraph {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
/// Builder for constructing brain graphs programmatically
|
|
pub struct GraphBuilder {
|
|
nodes: Vec<BrainNode>,
|
|
edges: Vec<BrainEdge>,
|
|
metadata: GraphMetadata,
|
|
}
|
|
|
|
impl GraphBuilder {
|
|
/// Create a new graph builder
|
|
pub fn new() -> Self {
|
|
Self {
|
|
nodes: Vec::new(),
|
|
edges: Vec::new(),
|
|
metadata: GraphMetadata::default(),
|
|
}
|
|
}
|
|
|
|
/// Add a node
|
|
pub fn add_node(&mut self, name: &str, hemisphere: Hemisphere, region: BrainRegion) -> usize {
|
|
let index = self.nodes.len();
|
|
self.nodes
|
|
.push(BrainNode::with_location(index, name, hemisphere, region));
|
|
index
|
|
}
|
|
|
|
/// Add a node with auto-detected hemisphere and region
|
|
pub fn add_channel(&mut self, name: &str) -> usize {
|
|
let index = self.nodes.len();
|
|
self.nodes.push(BrainNode::new(index, name));
|
|
index
|
|
}
|
|
|
|
/// Add an edge
|
|
pub fn add_edge(&mut self, source: usize, target: usize, weight: f64) {
|
|
self.edges.push(BrainEdge::new(source, target, weight));
|
|
}
|
|
|
|
/// Set metadata
|
|
pub fn with_metadata(mut self, metadata: GraphMetadata) -> Self {
|
|
self.metadata = metadata;
|
|
self
|
|
}
|
|
|
|
/// Build the graph
|
|
pub fn build(self) -> GnnResult<BrainGraph> {
|
|
// Validate edges
|
|
let n = self.nodes.len();
|
|
for edge in &self.edges {
|
|
if edge.source >= n || edge.target >= n {
|
|
return Err(GnnError::InvalidGraph(format!(
|
|
"Edge {}->{} references invalid node (n_nodes={})",
|
|
edge.source, edge.target, n
|
|
)));
|
|
}
|
|
}
|
|
|
|
// Mark interhemispheric edges
|
|
let mut edges = self.edges;
|
|
for edge in &mut edges {
|
|
edge.interhemispheric = self.nodes[edge.source].hemisphere
|
|
!= self.nodes[edge.target].hemisphere
|
|
&& self.nodes[edge.source].hemisphere != Hemisphere::Midline
|
|
&& self.nodes[edge.target].hemisphere != Hemisphere::Midline;
|
|
}
|
|
|
|
Ok(BrainGraph {
|
|
nodes: self.nodes,
|
|
edges,
|
|
adjacency: None,
|
|
metadata: self.metadata,
|
|
})
|
|
}
|
|
}
|
|
|
|
impl Default for GraphBuilder {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
/// Create a standard 10-20 system graph
|
|
pub fn standard_10_20_graph() -> BrainGraph {
|
|
let mut builder = GraphBuilder::new();
|
|
|
|
// Add standard 10-20 channels
|
|
let channels = [
|
|
"Fp1", "Fp2", "F7", "F3", "Fz", "F4", "F8", "T3", "C3", "Cz", "C4", "T4", "T5", "P3", "Pz",
|
|
"P4", "T6", "O1", "O2",
|
|
];
|
|
|
|
for name in &channels {
|
|
builder.add_channel(name);
|
|
}
|
|
|
|
// Add full connectivity (will be thresholded later)
|
|
let n = channels.len();
|
|
for i in 0..n {
|
|
for j in (i + 1)..n {
|
|
// Default weight based on distance
|
|
builder.add_edge(i, j, 0.5);
|
|
builder.add_edge(j, i, 0.5);
|
|
}
|
|
}
|
|
|
|
builder.build().unwrap()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_hemisphere_detection() {
|
|
assert_eq!(Hemisphere::from_channel_name("Fp1"), Hemisphere::Left);
|
|
assert_eq!(Hemisphere::from_channel_name("Fp2"), Hemisphere::Right);
|
|
assert_eq!(Hemisphere::from_channel_name("Fz"), Hemisphere::Midline);
|
|
assert_eq!(Hemisphere::from_channel_name("Cz"), Hemisphere::Midline);
|
|
assert_eq!(Hemisphere::from_channel_name("O1"), Hemisphere::Left);
|
|
assert_eq!(Hemisphere::from_channel_name("O2"), Hemisphere::Right);
|
|
}
|
|
|
|
#[test]
|
|
fn test_region_detection() {
|
|
assert_eq!(BrainRegion::from_channel_name("Fp1"), BrainRegion::Frontal);
|
|
assert_eq!(BrainRegion::from_channel_name("F3"), BrainRegion::Frontal);
|
|
assert_eq!(BrainRegion::from_channel_name("C3"), BrainRegion::Central);
|
|
assert_eq!(BrainRegion::from_channel_name("T3"), BrainRegion::Temporal);
|
|
assert_eq!(BrainRegion::from_channel_name("P3"), BrainRegion::Parietal);
|
|
assert_eq!(BrainRegion::from_channel_name("O1"), BrainRegion::Occipital);
|
|
}
|
|
|
|
#[test]
|
|
fn test_graph_builder() {
|
|
let mut builder = GraphBuilder::new();
|
|
builder.add_channel("Fp1");
|
|
builder.add_channel("Fp2");
|
|
builder.add_edge(0, 1, 0.8);
|
|
|
|
let graph = builder.build().unwrap();
|
|
assert_eq!(graph.n_nodes(), 2);
|
|
assert_eq!(graph.n_edges(), 1);
|
|
assert!(graph.edges[0].interhemispheric);
|
|
}
|
|
|
|
#[test]
|
|
fn test_subgraph() {
|
|
let mut builder = GraphBuilder::new();
|
|
builder.add_node("F3", Hemisphere::Left, BrainRegion::Frontal);
|
|
builder.add_node("F4", Hemisphere::Right, BrainRegion::Frontal);
|
|
builder.add_node("O1", Hemisphere::Left, BrainRegion::Occipital);
|
|
builder.add_node("O2", Hemisphere::Right, BrainRegion::Occipital);
|
|
|
|
builder.add_edge(0, 1, 0.8);
|
|
builder.add_edge(0, 2, 0.5);
|
|
builder.add_edge(2, 3, 0.9);
|
|
|
|
let graph = builder.build().unwrap();
|
|
|
|
// Left hemisphere subgraph
|
|
let left = graph.subgraph_by_hemisphere(Hemisphere::Left);
|
|
assert_eq!(left.n_nodes(), 2);
|
|
assert_eq!(left.n_edges(), 1); // F3->O1
|
|
|
|
// Frontal subgraph
|
|
let frontal = graph.subgraph_by_region(BrainRegion::Frontal);
|
|
assert_eq!(frontal.n_nodes(), 2);
|
|
assert_eq!(frontal.n_edges(), 1); // F3->F4
|
|
}
|
|
|
|
#[test]
|
|
fn test_standard_10_20() {
|
|
let graph = standard_10_20_graph();
|
|
assert_eq!(graph.n_nodes(), 19);
|
|
assert!(graph.n_edges() > 0);
|
|
}
|
|
}
|