623 lines
18 KiB
Rust
623 lines
18 KiB
Rust
//! Pre-built brain GNN model architectures.
|
|
//!
|
|
//! Includes BrainNetCNN, BrainGAT, and BrainTransformer models.
|
|
|
|
use crate::error::{GnnError, GnnResult};
|
|
use crate::graph::BrainGraph;
|
|
use crate::layers::{
|
|
Activation, BrainAttention, BrainAttentionConfig, BrainConv, BrainConvConfig, BrainPool,
|
|
BrainPoolConfig, PoolMethod,
|
|
};
|
|
use ndarray::{Array1, Array2};
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// Trait for brain GNN models
|
|
pub trait BrainGNN {
|
|
/// Forward pass on a brain graph
|
|
fn forward(&self, graph: &BrainGraph) -> GnnResult<Array1<f64>>;
|
|
|
|
/// Forward pass returning node embeddings
|
|
fn forward_node_embeddings(&self, graph: &BrainGraph) -> GnnResult<Array2<f64>>;
|
|
|
|
/// Get number of parameters
|
|
fn n_parameters(&self) -> usize;
|
|
}
|
|
|
|
/// BrainNetCNN configuration
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct BrainNetCNNConfig {
|
|
/// Number of input channels
|
|
pub n_channels: usize,
|
|
/// Number of input features per channel
|
|
pub n_features: usize,
|
|
/// Hidden dimensions for each layer
|
|
pub hidden_dims: Vec<usize>,
|
|
/// Number of output classes
|
|
pub n_classes: usize,
|
|
/// Dropout probability
|
|
pub dropout: f64,
|
|
/// Pooling method
|
|
pub pooling: PoolMethod,
|
|
}
|
|
|
|
impl Default for BrainNetCNNConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
n_channels: 64,
|
|
n_features: 1,
|
|
hidden_dims: vec![32, 64, 128],
|
|
n_classes: 2,
|
|
dropout: 0.5,
|
|
pooling: PoolMethod::Mean,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// BrainNetCNN model
|
|
///
|
|
/// A graph convolutional network inspired by BrainNetCNN paper.
|
|
/// Uses edge-to-edge and edge-to-node convolutions.
|
|
#[derive(Debug, Clone)]
|
|
pub struct BrainNetCNN {
|
|
config: BrainNetCNNConfig,
|
|
/// Graph convolution layers
|
|
conv_layers: Vec<BrainConv>,
|
|
/// Pooling layer
|
|
pool: BrainPool,
|
|
/// Final classifier weights
|
|
classifier: Array2<f64>,
|
|
/// Classifier bias
|
|
classifier_bias: Array1<f64>,
|
|
}
|
|
|
|
impl BrainNetCNN {
|
|
/// Create a new BrainNetCNN model
|
|
pub fn new(config: BrainNetCNNConfig) -> GnnResult<Self> {
|
|
let mut conv_layers = Vec::with_capacity(config.hidden_dims.len());
|
|
|
|
let mut prev_dim = config.n_features;
|
|
for &hidden_dim in &config.hidden_dims {
|
|
conv_layers.push(BrainConv::new(BrainConvConfig {
|
|
in_features: prev_dim,
|
|
out_features: hidden_dim,
|
|
add_self_loops: true,
|
|
normalize: true,
|
|
activation: Activation::ReLU,
|
|
dropout: config.dropout,
|
|
})?);
|
|
prev_dim = hidden_dim;
|
|
}
|
|
|
|
let pool = BrainPool::new(BrainPoolConfig {
|
|
method: config.pooling,
|
|
ratio: 0.5,
|
|
});
|
|
|
|
// Classifier
|
|
let final_dim = *config.hidden_dims.last().unwrap_or(&config.n_features);
|
|
let scale = (2.0 / (final_dim + config.n_classes) as f64).sqrt();
|
|
let classifier =
|
|
Array2::from_shape_fn((final_dim, config.n_classes), |_| rand_simple() * scale);
|
|
let classifier_bias = Array1::zeros(config.n_classes);
|
|
|
|
Ok(Self {
|
|
config,
|
|
conv_layers,
|
|
pool,
|
|
classifier,
|
|
classifier_bias,
|
|
})
|
|
}
|
|
}
|
|
|
|
impl BrainGNN for BrainNetCNN {
|
|
fn forward(&self, graph: &BrainGraph) -> GnnResult<Array1<f64>> {
|
|
let embeddings = self.forward_node_embeddings(graph)?;
|
|
|
|
// Global pooling
|
|
let pooled = self.pool.global_pool(&embeddings);
|
|
|
|
// Classification
|
|
let mut logits = Array1::zeros(self.config.n_classes);
|
|
for j in 0..self.config.n_classes {
|
|
logits[j] = self.classifier_bias[j];
|
|
for i in 0..pooled.len() {
|
|
logits[j] += pooled[i] * self.classifier[[i, j]];
|
|
}
|
|
}
|
|
|
|
Ok(logits)
|
|
}
|
|
|
|
fn forward_node_embeddings(&self, graph: &BrainGraph) -> GnnResult<Array2<f64>> {
|
|
let mut x = graph.node_feature_matrix()?;
|
|
let adj = graph.adjacency_matrix();
|
|
|
|
// Apply conv layers
|
|
for layer in &self.conv_layers {
|
|
x = layer.forward(&x, &adj)?;
|
|
}
|
|
|
|
Ok(x)
|
|
}
|
|
|
|
fn n_parameters(&self) -> usize {
|
|
let mut count = 0;
|
|
for _layer in &self.conv_layers {
|
|
// Approximate parameter count
|
|
count += 100; // Placeholder
|
|
}
|
|
count += self.classifier.len() + self.classifier_bias.len();
|
|
count
|
|
}
|
|
}
|
|
|
|
/// BrainGAT configuration
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct BrainGATConfig {
|
|
/// Number of input channels
|
|
pub n_channels: usize,
|
|
/// Number of input features per channel
|
|
pub n_features: usize,
|
|
/// Hidden dimensions
|
|
pub hidden_dims: Vec<usize>,
|
|
/// Number of attention heads per layer
|
|
pub n_heads: Vec<usize>,
|
|
/// Number of output classes
|
|
pub n_classes: usize,
|
|
/// Dropout probability
|
|
pub dropout: f64,
|
|
/// Whether to concatenate attention heads
|
|
pub concat_heads: bool,
|
|
}
|
|
|
|
impl Default for BrainGATConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
n_channels: 64,
|
|
n_features: 1,
|
|
hidden_dims: vec![64, 64],
|
|
n_heads: vec![4, 4],
|
|
n_classes: 2,
|
|
dropout: 0.5,
|
|
concat_heads: true,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Brain Graph Attention Network
|
|
#[derive(Debug, Clone)]
|
|
pub struct BrainGAT {
|
|
config: BrainGATConfig,
|
|
/// Attention layers
|
|
attn_layers: Vec<BrainAttention>,
|
|
/// Pooling layer
|
|
pool: BrainPool,
|
|
/// Final classifier
|
|
classifier: Array2<f64>,
|
|
classifier_bias: Array1<f64>,
|
|
}
|
|
|
|
impl BrainGAT {
|
|
/// Create a new BrainGAT model
|
|
pub fn new(config: BrainGATConfig) -> GnnResult<Self> {
|
|
if config.hidden_dims.len() != config.n_heads.len() {
|
|
return Err(GnnError::InvalidConfig(
|
|
"hidden_dims and n_heads must have same length".into(),
|
|
));
|
|
}
|
|
|
|
let mut attn_layers = Vec::with_capacity(config.hidden_dims.len());
|
|
|
|
let mut prev_dim = config.n_features;
|
|
for (&hidden_dim, &n_heads) in config.hidden_dims.iter().zip(&config.n_heads) {
|
|
attn_layers.push(BrainAttention::new(BrainAttentionConfig {
|
|
in_features: prev_dim,
|
|
out_features: hidden_dim,
|
|
n_heads,
|
|
dropout: config.dropout,
|
|
concat_heads: config.concat_heads,
|
|
..Default::default()
|
|
})?);
|
|
prev_dim = hidden_dim;
|
|
}
|
|
|
|
let pool = BrainPool::new(BrainPoolConfig::default());
|
|
|
|
// Classifier
|
|
let final_dim = *config.hidden_dims.last().unwrap_or(&config.n_features);
|
|
let scale = (2.0 / (final_dim + config.n_classes) as f64).sqrt();
|
|
let classifier =
|
|
Array2::from_shape_fn((final_dim, config.n_classes), |_| rand_simple() * scale);
|
|
let classifier_bias = Array1::zeros(config.n_classes);
|
|
|
|
Ok(Self {
|
|
config,
|
|
attn_layers,
|
|
pool,
|
|
classifier,
|
|
classifier_bias,
|
|
})
|
|
}
|
|
}
|
|
|
|
impl BrainGNN for BrainGAT {
|
|
fn forward(&self, graph: &BrainGraph) -> GnnResult<Array1<f64>> {
|
|
let embeddings = self.forward_node_embeddings(graph)?;
|
|
let pooled = self.pool.global_pool(&embeddings);
|
|
|
|
let mut logits = Array1::zeros(self.config.n_classes);
|
|
for j in 0..self.config.n_classes {
|
|
logits[j] = self.classifier_bias[j];
|
|
for i in 0..pooled.len() {
|
|
logits[j] += pooled[i] * self.classifier[[i, j]];
|
|
}
|
|
}
|
|
|
|
Ok(logits)
|
|
}
|
|
|
|
fn forward_node_embeddings(&self, graph: &BrainGraph) -> GnnResult<Array2<f64>> {
|
|
let mut x = graph.node_feature_matrix()?;
|
|
let adj = graph.adjacency_matrix();
|
|
|
|
for layer in &self.attn_layers {
|
|
x = layer.forward(&x, &adj)?;
|
|
}
|
|
|
|
Ok(x)
|
|
}
|
|
|
|
fn n_parameters(&self) -> usize {
|
|
self.classifier.len() + self.classifier_bias.len() + 1000 // Approximate
|
|
}
|
|
}
|
|
|
|
/// BrainTransformer configuration
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct BrainTransformerConfig {
|
|
/// Number of input channels
|
|
pub n_channels: usize,
|
|
/// Number of input features per channel
|
|
pub n_features: usize,
|
|
/// Model dimension
|
|
pub d_model: usize,
|
|
/// Number of transformer layers
|
|
pub n_layers: usize,
|
|
/// Number of attention heads
|
|
pub n_heads: usize,
|
|
/// Feed-forward dimension
|
|
pub d_ff: usize,
|
|
/// Number of output classes
|
|
pub n_classes: usize,
|
|
/// Dropout probability
|
|
pub dropout: f64,
|
|
/// Whether to use graph structure in attention
|
|
pub use_graph_mask: bool,
|
|
}
|
|
|
|
impl Default for BrainTransformerConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
n_channels: 64,
|
|
n_features: 1,
|
|
d_model: 64,
|
|
n_layers: 4,
|
|
n_heads: 4,
|
|
d_ff: 256,
|
|
n_classes: 2,
|
|
dropout: 0.1,
|
|
use_graph_mask: true,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Brain Transformer model
|
|
///
|
|
/// A transformer-based model for brain connectivity analysis.
|
|
/// Can optionally mask attention based on graph structure.
|
|
#[derive(Debug, Clone)]
|
|
pub struct BrainTransformer {
|
|
config: BrainTransformerConfig,
|
|
/// Input projection
|
|
input_proj: Array2<f64>,
|
|
/// Transformer layers
|
|
layers: Vec<TransformerLayer>,
|
|
/// Pooling
|
|
pool: BrainPool,
|
|
/// Classifier
|
|
classifier: Array2<f64>,
|
|
classifier_bias: Array1<f64>,
|
|
}
|
|
|
|
/// A single transformer layer
|
|
#[derive(Debug, Clone)]
|
|
struct TransformerLayer {
|
|
/// Query projection
|
|
wq: Array2<f64>,
|
|
/// Key projection
|
|
wk: Array2<f64>,
|
|
/// Value projection
|
|
wv: Array2<f64>,
|
|
/// Output projection
|
|
wo: Array2<f64>,
|
|
/// FFN layer 1
|
|
ff1: Array2<f64>,
|
|
/// FFN layer 2
|
|
ff2: Array2<f64>,
|
|
/// Layer norms (simulated with scale)
|
|
scale: f64,
|
|
}
|
|
|
|
impl BrainTransformer {
|
|
/// Create a new BrainTransformer model
|
|
pub fn new(config: BrainTransformerConfig) -> GnnResult<Self> {
|
|
let scale = (1.0 / config.d_model as f64).sqrt();
|
|
|
|
// Input projection
|
|
let input_proj = Array2::from_shape_fn((config.n_features, config.d_model), |_| {
|
|
rand_simple() * scale
|
|
});
|
|
|
|
// Transformer layers
|
|
let mut layers = Vec::with_capacity(config.n_layers);
|
|
for _ in 0..config.n_layers {
|
|
layers.push(TransformerLayer {
|
|
wq: Array2::from_shape_fn((config.d_model, config.d_model), |_| {
|
|
rand_simple() * scale
|
|
}),
|
|
wk: Array2::from_shape_fn((config.d_model, config.d_model), |_| {
|
|
rand_simple() * scale
|
|
}),
|
|
wv: Array2::from_shape_fn((config.d_model, config.d_model), |_| {
|
|
rand_simple() * scale
|
|
}),
|
|
wo: Array2::from_shape_fn((config.d_model, config.d_model), |_| {
|
|
rand_simple() * scale
|
|
}),
|
|
ff1: Array2::from_shape_fn((config.d_model, config.d_ff), |_| {
|
|
rand_simple() * scale
|
|
}),
|
|
ff2: Array2::from_shape_fn((config.d_ff, config.d_model), |_| {
|
|
rand_simple() * scale
|
|
}),
|
|
scale: 1.0,
|
|
});
|
|
}
|
|
|
|
let pool = BrainPool::new(BrainPoolConfig::default());
|
|
|
|
// Classifier
|
|
let classifier = Array2::from_shape_fn((config.d_model, config.n_classes), |_| {
|
|
rand_simple() * scale
|
|
});
|
|
let classifier_bias = Array1::zeros(config.n_classes);
|
|
|
|
Ok(Self {
|
|
config,
|
|
input_proj,
|
|
layers,
|
|
pool,
|
|
classifier,
|
|
classifier_bias,
|
|
})
|
|
}
|
|
|
|
/// Self-attention with optional graph masking
|
|
fn self_attention(
|
|
&self,
|
|
layer: &TransformerLayer,
|
|
x: &Array2<f64>,
|
|
mask: Option<&Array2<f64>>,
|
|
) -> Array2<f64> {
|
|
let n = x.nrows();
|
|
let d = x.ncols();
|
|
|
|
// Compute Q, K, V
|
|
let q = x.dot(&layer.wq);
|
|
let k = x.dot(&layer.wk);
|
|
let v = x.dot(&layer.wv);
|
|
|
|
// Attention scores
|
|
let scale = 1.0 / (d as f64).sqrt();
|
|
let mut scores = q.dot(&k.t()) * scale;
|
|
|
|
// Apply graph mask if provided
|
|
if let Some(m) = mask {
|
|
for i in 0..n {
|
|
for j in 0..n {
|
|
if m[[i, j]] < 0.5 && i != j {
|
|
scores[[i, j]] = f64::NEG_INFINITY;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Softmax
|
|
for i in 0..n {
|
|
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 {
|
|
scores[[i, j]] = (scores[[i, j]] - max_val).exp() / exp_sum;
|
|
}
|
|
}
|
|
|
|
// Apply attention
|
|
let attn_out = scores.dot(&v);
|
|
attn_out.dot(&layer.wo)
|
|
}
|
|
|
|
/// Feed-forward network
|
|
fn ffn(&self, layer: &TransformerLayer, x: &Array2<f64>) -> Array2<f64> {
|
|
// FFN with GELU activation
|
|
let h = x.dot(&layer.ff1);
|
|
let h_act =
|
|
h.mapv(|v| 0.5 * v * (1.0 + (0.7978845608 * (v + 0.044715 * v * v * v)).tanh()));
|
|
h_act.dot(&layer.ff2)
|
|
}
|
|
}
|
|
|
|
impl BrainGNN for BrainTransformer {
|
|
fn forward(&self, graph: &BrainGraph) -> GnnResult<Array1<f64>> {
|
|
let embeddings = self.forward_node_embeddings(graph)?;
|
|
let pooled = self.pool.global_pool(&embeddings);
|
|
|
|
let mut logits = Array1::zeros(self.config.n_classes);
|
|
for j in 0..self.config.n_classes {
|
|
logits[j] = self.classifier_bias[j];
|
|
for i in 0..pooled.len() {
|
|
logits[j] += pooled[i] * self.classifier[[i, j]];
|
|
}
|
|
}
|
|
|
|
Ok(logits)
|
|
}
|
|
|
|
fn forward_node_embeddings(&self, graph: &BrainGraph) -> GnnResult<Array2<f64>> {
|
|
let mut x = graph.node_feature_matrix()?;
|
|
|
|
// Project to model dimension
|
|
x = x.dot(&self.input_proj);
|
|
|
|
// Graph mask
|
|
let mask = if self.config.use_graph_mask {
|
|
Some(graph.adjacency_matrix())
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// Apply transformer layers
|
|
for layer in &self.layers {
|
|
// Self-attention with residual
|
|
let attn_out = self.self_attention(layer, &x, mask.as_ref());
|
|
x = x + attn_out;
|
|
|
|
// FFN with residual
|
|
let ffn_out = self.ffn(layer, &x);
|
|
x = x + ffn_out;
|
|
}
|
|
|
|
Ok(x)
|
|
}
|
|
|
|
fn n_parameters(&self) -> usize {
|
|
let mut count = self.input_proj.len();
|
|
for layer in &self.layers {
|
|
count += layer.wq.len()
|
|
+ layer.wk.len()
|
|
+ layer.wv.len()
|
|
+ layer.wo.len()
|
|
+ layer.ff1.len()
|
|
+ layer.ff2.len();
|
|
}
|
|
count += self.classifier.len() + self.classifier_bias.len();
|
|
count
|
|
}
|
|
}
|
|
|
|
// Simple pseudo-random for initialization
|
|
fn rand_simple() -> f64 {
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
static SEED: AtomicU64 = AtomicU64::new(54321);
|
|
|
|
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::GraphBuilder;
|
|
use crate::graph::{BrainRegion, Hemisphere};
|
|
|
|
fn create_test_graph() -> BrainGraph {
|
|
let mut builder = GraphBuilder::new();
|
|
|
|
// 4-node test graph
|
|
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();
|
|
|
|
// Add features
|
|
for node in &mut graph.nodes {
|
|
node.features = vec![0.5, 0.3, 0.2, 0.1];
|
|
}
|
|
|
|
graph
|
|
}
|
|
|
|
#[test]
|
|
fn test_brainnetcnn() {
|
|
let graph = create_test_graph();
|
|
|
|
let model = BrainNetCNN::new(BrainNetCNNConfig {
|
|
n_channels: 4,
|
|
n_features: 4,
|
|
hidden_dims: vec![8, 16],
|
|
n_classes: 3,
|
|
..Default::default()
|
|
})
|
|
.unwrap();
|
|
|
|
let logits = model.forward(&graph).unwrap();
|
|
assert_eq!(logits.len(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_braingat() {
|
|
let graph = create_test_graph();
|
|
|
|
let model = BrainGAT::new(BrainGATConfig {
|
|
n_channels: 4,
|
|
n_features: 4,
|
|
hidden_dims: vec![8, 8],
|
|
n_heads: vec![2, 2],
|
|
n_classes: 2,
|
|
..Default::default()
|
|
})
|
|
.unwrap();
|
|
|
|
let logits = model.forward(&graph).unwrap();
|
|
assert_eq!(logits.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_brain_transformer() {
|
|
let graph = create_test_graph();
|
|
|
|
let model = BrainTransformer::new(BrainTransformerConfig {
|
|
n_channels: 4,
|
|
n_features: 4,
|
|
d_model: 16,
|
|
n_layers: 2,
|
|
n_heads: 2,
|
|
d_ff: 32,
|
|
n_classes: 2,
|
|
..Default::default()
|
|
})
|
|
.unwrap();
|
|
|
|
let logits = model.forward(&graph).unwrap();
|
|
assert_eq!(logits.len(), 2);
|
|
|
|
let embeddings = model.forward_node_embeddings(&graph).unwrap();
|
|
assert_eq!(embeddings.shape(), &[4, 16]);
|
|
}
|
|
}
|