Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,573 @@
|
||||
//! Temporal GNN for dynamic connectivity analysis.
|
||||
//!
|
||||
//! Provides models for analyzing time-varying brain connectivity graphs.
|
||||
|
||||
use crate::error::{GnnError, GnnResult};
|
||||
use crate::graph::BrainGraph;
|
||||
use crate::layers::{BrainPool, BrainPoolConfig, PoolMethod};
|
||||
use crate::models::{BrainGNN, BrainNetCNN, BrainNetCNNConfig};
|
||||
use ndarray::{Array1, Array2};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Time window specification
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TimeWindow {
|
||||
/// Start time in seconds
|
||||
pub start: f64,
|
||||
/// End time in seconds
|
||||
pub end: f64,
|
||||
/// Window index
|
||||
pub index: usize,
|
||||
}
|
||||
|
||||
impl TimeWindow {
|
||||
/// Create a new time window
|
||||
pub fn new(start: f64, end: f64, index: usize) -> Self {
|
||||
Self { start, end, index }
|
||||
}
|
||||
|
||||
/// Duration of the window
|
||||
pub fn duration(&self) -> f64 {
|
||||
self.end - self.start
|
||||
}
|
||||
|
||||
/// Center time of the window
|
||||
pub fn center(&self) -> f64 {
|
||||
f64::midpoint(self.start, self.end)
|
||||
}
|
||||
}
|
||||
|
||||
/// Dynamic connectivity sequence
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DynamicConnectivity {
|
||||
/// Sequence of brain graphs (one per time window)
|
||||
pub graphs: Vec<BrainGraph>,
|
||||
/// Time windows
|
||||
pub windows: Vec<TimeWindow>,
|
||||
/// Sampling frequency
|
||||
pub sfreq: f64,
|
||||
}
|
||||
|
||||
impl DynamicConnectivity {
|
||||
/// Create from a sequence of connectivity matrices
|
||||
pub fn from_matrices(
|
||||
matrices: &[Vec<Vec<f64>>],
|
||||
channel_names: &[&str],
|
||||
window_duration: f64,
|
||||
step: f64,
|
||||
threshold: f64,
|
||||
) -> GnnResult<Self> {
|
||||
let mut graphs = Vec::with_capacity(matrices.len());
|
||||
let mut windows = Vec::with_capacity(matrices.len());
|
||||
|
||||
for (i, matrix) in matrices.iter().enumerate() {
|
||||
let graph = BrainGraph::from_connectivity_matrix(matrix, channel_names, threshold)?;
|
||||
graphs.push(graph);
|
||||
|
||||
let start = i as f64 * step;
|
||||
windows.push(TimeWindow::new(start, start + window_duration, i));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
graphs,
|
||||
windows,
|
||||
sfreq: 1.0 / step,
|
||||
})
|
||||
}
|
||||
|
||||
/// Number of time points
|
||||
pub fn n_times(&self) -> usize {
|
||||
self.graphs.len()
|
||||
}
|
||||
|
||||
/// Get graph at specific time index
|
||||
pub fn get(&self, index: usize) -> Option<&BrainGraph> {
|
||||
self.graphs.get(index)
|
||||
}
|
||||
|
||||
/// Iterator over (time_window, graph) pairs
|
||||
pub fn iter(&self) -> impl Iterator<Item = (&TimeWindow, &BrainGraph)> {
|
||||
self.windows.iter().zip(self.graphs.iter())
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration for temporal GNN
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TemporalConfig {
|
||||
/// Spatial GNN configuration
|
||||
pub spatial: BrainNetCNNConfig,
|
||||
/// Number of temporal attention heads
|
||||
pub temporal_heads: usize,
|
||||
/// Temporal hidden dimension
|
||||
pub temporal_hidden: usize,
|
||||
/// Whether to use temporal positional encoding
|
||||
pub positional_encoding: bool,
|
||||
/// Maximum sequence length
|
||||
pub max_seq_len: usize,
|
||||
}
|
||||
|
||||
impl Default for TemporalConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
spatial: BrainNetCNNConfig {
|
||||
hidden_dims: vec![32, 64],
|
||||
n_classes: 2,
|
||||
..Default::default()
|
||||
},
|
||||
temporal_heads: 4,
|
||||
temporal_hidden: 64,
|
||||
positional_encoding: true,
|
||||
max_seq_len: 100,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Temporal attention mechanism
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TemporalAttention {
|
||||
/// Number of attention heads
|
||||
n_heads: usize,
|
||||
/// Dimension per head
|
||||
head_dim: usize,
|
||||
/// Query weights per head
|
||||
wq: Vec<Array2<f64>>,
|
||||
/// Key weights per head
|
||||
wk: Vec<Array2<f64>>,
|
||||
/// Value weights per head
|
||||
wv: Vec<Array2<f64>>,
|
||||
/// Output projection
|
||||
wo: Array2<f64>,
|
||||
/// Positional encoding matrix
|
||||
positional_encoding: Option<Array2<f64>>,
|
||||
}
|
||||
|
||||
impl TemporalAttention {
|
||||
/// Create a new temporal attention layer
|
||||
pub fn new(d_model: usize, n_heads: usize, max_seq_len: usize, use_pe: bool) -> Self {
|
||||
let head_dim = d_model / n_heads;
|
||||
let scale = (1.0 / head_dim as f64).sqrt();
|
||||
|
||||
let mut wq = Vec::with_capacity(n_heads);
|
||||
let mut wk = Vec::with_capacity(n_heads);
|
||||
let mut wv = Vec::with_capacity(n_heads);
|
||||
|
||||
for _ in 0..n_heads {
|
||||
wq.push(Array2::from_shape_fn((d_model, head_dim), |_| {
|
||||
rand_simple() * scale
|
||||
}));
|
||||
wk.push(Array2::from_shape_fn((d_model, head_dim), |_| {
|
||||
rand_simple() * scale
|
||||
}));
|
||||
wv.push(Array2::from_shape_fn((d_model, head_dim), |_| {
|
||||
rand_simple() * scale
|
||||
}));
|
||||
}
|
||||
|
||||
let wo = Array2::from_shape_fn((d_model, d_model), |_| rand_simple() * scale);
|
||||
|
||||
// Sinusoidal positional encoding
|
||||
let positional_encoding = if use_pe {
|
||||
let mut pe = Array2::zeros((max_seq_len, d_model));
|
||||
for pos in 0..max_seq_len {
|
||||
for i in 0..(d_model / 2) {
|
||||
let angle = pos as f64 / (10000.0_f64).powf(2.0 * i as f64 / d_model as f64);
|
||||
pe[[pos, 2 * i]] = angle.sin();
|
||||
pe[[pos, 2 * i + 1]] = angle.cos();
|
||||
}
|
||||
}
|
||||
Some(pe)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Self {
|
||||
n_heads,
|
||||
head_dim,
|
||||
wq,
|
||||
wk,
|
||||
wv,
|
||||
wo,
|
||||
positional_encoding,
|
||||
}
|
||||
}
|
||||
|
||||
/// Forward pass on sequence of embeddings [n_times, d_model]
|
||||
pub fn forward(&self, x: &Array2<f64>) -> Array2<f64> {
|
||||
let n_times = x.nrows();
|
||||
let d_model = x.ncols();
|
||||
|
||||
// Add positional encoding
|
||||
let mut h = x.clone();
|
||||
if let Some(ref pe) = self.positional_encoding {
|
||||
for t in 0..n_times {
|
||||
for d in 0..d_model {
|
||||
h[[t, d]] += pe[[t, d]];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Multi-head attention
|
||||
let mut head_outputs = Vec::with_capacity(self.n_heads);
|
||||
|
||||
for head in 0..self.n_heads {
|
||||
let q = h.dot(&self.wq[head]); // [n_times, head_dim]
|
||||
let k = h.dot(&self.wk[head]);
|
||||
let v = h.dot(&self.wv[head]);
|
||||
|
||||
// Attention scores
|
||||
let scale = 1.0 / (self.head_dim as f64).sqrt();
|
||||
let scores = q.dot(&k.t()) * scale; // [n_times, n_times]
|
||||
|
||||
// Softmax
|
||||
let mut attn = Array2::zeros((n_times, n_times));
|
||||
for i in 0..n_times {
|
||||
let max_val = scores
|
||||
.row(i)
|
||||
.iter()
|
||||
.copied()
|
||||
.fold(f64::NEG_INFINITY, f64::max);
|
||||
let exp_sum: f64 = scores.row(i).iter().map(|&s| (s - max_val).exp()).sum();
|
||||
for j in 0..n_times {
|
||||
attn[[i, j]] = (scores[[i, j]] - max_val).exp() / exp_sum;
|
||||
}
|
||||
}
|
||||
|
||||
let out = attn.dot(&v);
|
||||
head_outputs.push(out);
|
||||
}
|
||||
|
||||
// Concatenate heads
|
||||
let mut concat = Array2::zeros((n_times, d_model));
|
||||
for (h, head_out) in head_outputs.iter().enumerate() {
|
||||
for t in 0..n_times {
|
||||
for d in 0..self.head_dim {
|
||||
concat[[t, h * self.head_dim + d]] = head_out[[t, d]];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Output projection + residual
|
||||
let output = concat.dot(&self.wo);
|
||||
x + &output
|
||||
}
|
||||
|
||||
/// Get attention weights for interpretation
|
||||
pub fn get_attention_weights(&self, x: &Array2<f64>) -> Array2<f64> {
|
||||
let n_times = x.nrows();
|
||||
|
||||
// Use first head for visualization
|
||||
let q = x.dot(&self.wq[0]);
|
||||
let k = x.dot(&self.wk[0]);
|
||||
|
||||
let scale = 1.0 / (self.head_dim as f64).sqrt();
|
||||
let scores = q.dot(&k.t()) * scale;
|
||||
|
||||
// Softmax
|
||||
let mut attn = Array2::zeros((n_times, n_times));
|
||||
for i in 0..n_times {
|
||||
let max_val = scores
|
||||
.row(i)
|
||||
.iter()
|
||||
.copied()
|
||||
.fold(f64::NEG_INFINITY, f64::max);
|
||||
let exp_sum: f64 = scores.row(i).iter().map(|&s| (s - max_val).exp()).sum();
|
||||
for j in 0..n_times {
|
||||
attn[[i, j]] = (scores[[i, j]] - max_val).exp() / exp_sum;
|
||||
}
|
||||
}
|
||||
|
||||
attn
|
||||
}
|
||||
}
|
||||
|
||||
/// Temporal Brain GNN model
|
||||
///
|
||||
/// Combines spatial GNN for each time window with temporal attention
|
||||
/// for capturing dynamic brain state changes.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TemporalBrainGNN {
|
||||
config: TemporalConfig,
|
||||
/// Spatial GNN (applied to each time window)
|
||||
spatial_gnn: BrainNetCNN,
|
||||
/// Temporal attention
|
||||
temporal_attention: TemporalAttention,
|
||||
/// Temporal convolution weights
|
||||
temporal_conv: Array2<f64>,
|
||||
/// Pooling
|
||||
pool: BrainPool,
|
||||
/// Classifier
|
||||
classifier: Array2<f64>,
|
||||
classifier_bias: Array1<f64>,
|
||||
}
|
||||
|
||||
impl TemporalBrainGNN {
|
||||
/// Create a new temporal brain GNN
|
||||
pub fn new(config: TemporalConfig) -> GnnResult<Self> {
|
||||
// Create spatial GNN
|
||||
let spatial_gnn = BrainNetCNN::new(config.spatial.clone())?;
|
||||
|
||||
// Final dimension from spatial GNN
|
||||
let spatial_dim = *config
|
||||
.spatial
|
||||
.hidden_dims
|
||||
.last()
|
||||
.unwrap_or(&config.spatial.n_features);
|
||||
|
||||
// Temporal attention
|
||||
let temporal_attention = TemporalAttention::new(
|
||||
spatial_dim,
|
||||
config.temporal_heads,
|
||||
config.max_seq_len,
|
||||
config.positional_encoding,
|
||||
);
|
||||
|
||||
// Temporal 1D convolution (simplified as matrix)
|
||||
let scale = (1.0 / config.temporal_hidden as f64).sqrt();
|
||||
let temporal_conv = Array2::from_shape_fn((spatial_dim, config.temporal_hidden), |_| {
|
||||
rand_simple() * scale
|
||||
});
|
||||
|
||||
let pool = BrainPool::new(BrainPoolConfig {
|
||||
method: PoolMethod::Attention,
|
||||
ratio: 1.0,
|
||||
});
|
||||
|
||||
// Classifier
|
||||
let classifier =
|
||||
Array2::from_shape_fn((config.temporal_hidden, config.spatial.n_classes), |_| {
|
||||
rand_simple() * scale
|
||||
});
|
||||
let classifier_bias = Array1::zeros(config.spatial.n_classes);
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
spatial_gnn,
|
||||
temporal_attention,
|
||||
temporal_conv,
|
||||
pool,
|
||||
classifier,
|
||||
classifier_bias,
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass on dynamic connectivity
|
||||
pub fn forward(&self, dyn_conn: &DynamicConnectivity) -> GnnResult<Array1<f64>> {
|
||||
let temporal_embeddings = self.forward_temporal_embeddings(dyn_conn)?;
|
||||
|
||||
// Global temporal pooling
|
||||
let pooled = self.pool.global_pool(&temporal_embeddings);
|
||||
|
||||
// Apply temporal convolution
|
||||
let h = Array1::from_iter((0..self.config.temporal_hidden).map(|j| {
|
||||
(0..pooled.len())
|
||||
.map(|i| pooled[i] * self.temporal_conv[[i, j]])
|
||||
.sum::<f64>()
|
||||
}));
|
||||
|
||||
// Classifier
|
||||
let mut logits = Array1::zeros(self.config.spatial.n_classes);
|
||||
for j in 0..self.config.spatial.n_classes {
|
||||
logits[j] = self.classifier_bias[j];
|
||||
for i in 0..h.len() {
|
||||
logits[j] += h[i] * self.classifier[[i, j]];
|
||||
}
|
||||
}
|
||||
|
||||
Ok(logits)
|
||||
}
|
||||
|
||||
/// Get temporal embeddings [n_times, spatial_dim]
|
||||
pub fn forward_temporal_embeddings(
|
||||
&self,
|
||||
dyn_conn: &DynamicConnectivity,
|
||||
) -> GnnResult<Array2<f64>> {
|
||||
let n_times = dyn_conn.n_times();
|
||||
if n_times == 0 {
|
||||
return Err(GnnError::InvalidGraph("No time windows".into()));
|
||||
}
|
||||
|
||||
let spatial_dim = *self
|
||||
.config
|
||||
.spatial
|
||||
.hidden_dims
|
||||
.last()
|
||||
.unwrap_or(&self.config.spatial.n_features);
|
||||
|
||||
// Get spatial embeddings for each time window
|
||||
let mut spatial_embeddings = Array2::zeros((n_times, spatial_dim));
|
||||
|
||||
for (t, graph) in dyn_conn.graphs.iter().enumerate() {
|
||||
let node_emb = self.spatial_gnn.forward_node_embeddings(graph)?;
|
||||
let pooled = self.pool.global_pool(&node_emb);
|
||||
|
||||
for d in 0..spatial_dim.min(pooled.len()) {
|
||||
spatial_embeddings[[t, d]] = pooled[d];
|
||||
}
|
||||
}
|
||||
|
||||
// Apply temporal attention
|
||||
let temporal_out = self.temporal_attention.forward(&spatial_embeddings);
|
||||
|
||||
Ok(temporal_out)
|
||||
}
|
||||
|
||||
/// Get temporal attention weights for interpretation
|
||||
pub fn get_temporal_attention_weights(
|
||||
&self,
|
||||
dyn_conn: &DynamicConnectivity,
|
||||
) -> GnnResult<Array2<f64>> {
|
||||
let temporal_embeddings = self.forward_temporal_embeddings(dyn_conn)?;
|
||||
Ok(self
|
||||
.temporal_attention
|
||||
.get_attention_weights(&temporal_embeddings))
|
||||
}
|
||||
|
||||
/// Predict brain state at each time window
|
||||
pub fn predict_states(&self, dyn_conn: &DynamicConnectivity) -> GnnResult<Vec<usize>> {
|
||||
let temporal_embeddings = self.forward_temporal_embeddings(dyn_conn)?;
|
||||
|
||||
let mut states = Vec::with_capacity(dyn_conn.n_times());
|
||||
|
||||
for t in 0..dyn_conn.n_times() {
|
||||
// Get embedding for this time
|
||||
let h: Vec<f64> = temporal_embeddings.row(t).to_vec();
|
||||
|
||||
// Apply temporal conv
|
||||
let conv_out: Vec<f64> = (0..self.config.temporal_hidden)
|
||||
.map(|j| {
|
||||
(0..h.len())
|
||||
.map(|i| h[i] * self.temporal_conv[[i, j]])
|
||||
.sum()
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Classify
|
||||
let mut logits = vec![0.0; self.config.spatial.n_classes];
|
||||
for j in 0..self.config.spatial.n_classes {
|
||||
logits[j] = self.classifier_bias[j];
|
||||
for i in 0..conv_out.len() {
|
||||
logits[j] += conv_out[i] * self.classifier[[i, j]];
|
||||
}
|
||||
}
|
||||
|
||||
// Argmax
|
||||
let state = logits
|
||||
.iter()
|
||||
.enumerate()
|
||||
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
|
||||
.map_or(0, |(i, _)| i);
|
||||
|
||||
states.push(state);
|
||||
}
|
||||
|
||||
Ok(states)
|
||||
}
|
||||
}
|
||||
|
||||
// Simple pseudo-random for initialization
|
||||
fn rand_simple() -> f64 {
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
static SEED: AtomicU64 = AtomicU64::new(98765);
|
||||
|
||||
let mut s = SEED.fetch_add(1, Ordering::Relaxed);
|
||||
s ^= s >> 12;
|
||||
s ^= s << 25;
|
||||
s ^= s >> 27;
|
||||
s = s.wrapping_mul(0x2545F4914F6CDD1D);
|
||||
(s as f64) / (u64::MAX as f64)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::graph::{BrainRegion, GraphBuilder, Hemisphere};
|
||||
|
||||
fn create_test_graph() -> BrainGraph {
|
||||
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(1, 3, 0.6);
|
||||
builder.add_edge(2, 3, 0.9);
|
||||
|
||||
let mut graph = builder.build().unwrap();
|
||||
for node in &mut graph.nodes {
|
||||
node.features = vec![0.5, 0.3, 0.2, 0.1];
|
||||
}
|
||||
graph
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dynamic_connectivity() {
|
||||
let channels = ["F3", "F4", "O1", "O2"];
|
||||
let matrices: Vec<Vec<Vec<f64>>> = (0..5)
|
||||
.map(|t| {
|
||||
(0..4)
|
||||
.map(|i| {
|
||||
(0..4)
|
||||
.map(|j| if i == j { 1.0 } else { 0.5 + t as f64 * 0.05 })
|
||||
.collect()
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.collect();
|
||||
|
||||
let dyn_conn =
|
||||
DynamicConnectivity::from_matrices(&matrices, &channels, 1.0, 0.5, 0.4).unwrap();
|
||||
|
||||
assert_eq!(dyn_conn.n_times(), 5);
|
||||
assert!(dyn_conn.get(2).is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_temporal_attention() {
|
||||
let attn = TemporalAttention::new(16, 4, 50, true);
|
||||
|
||||
let x = Array2::from_shape_fn((10, 16), |_| rand_simple());
|
||||
let out = attn.forward(&x);
|
||||
|
||||
assert_eq!(out.shape(), &[10, 16]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_temporal_brain_gnn() {
|
||||
// Create dynamic connectivity
|
||||
let channels = ["F3", "F4", "O1", "O2"];
|
||||
let matrices: Vec<Vec<Vec<f64>>> = (0..5)
|
||||
.map(|_| {
|
||||
(0..4)
|
||||
.map(|i| (0..4).map(|j| if i == j { 1.0 } else { 0.6 }).collect())
|
||||
.collect()
|
||||
})
|
||||
.collect();
|
||||
|
||||
let dyn_conn =
|
||||
DynamicConnectivity::from_matrices(&matrices, &channels, 1.0, 0.5, 0.4).unwrap();
|
||||
|
||||
// Create model
|
||||
let model = TemporalBrainGNN::new(TemporalConfig {
|
||||
spatial: BrainNetCNNConfig {
|
||||
n_channels: 4,
|
||||
n_features: 1, // Will use degree
|
||||
hidden_dims: vec![8, 16],
|
||||
n_classes: 2,
|
||||
..Default::default()
|
||||
},
|
||||
temporal_heads: 2,
|
||||
temporal_hidden: 8,
|
||||
max_seq_len: 10,
|
||||
..Default::default()
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let logits = model.forward(&dyn_conn).unwrap();
|
||||
assert_eq!(logits.len(), 2);
|
||||
|
||||
let states = model.predict_states(&dyn_conn).unwrap();
|
||||
assert_eq!(states.len(), 5);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user