630 lines
18 KiB
Rust
630 lines
18 KiB
Rust
//! Brain-specific GNN layers.
|
|
//!
|
|
//! These layers extend rtx-geom's GNN layers with brain-specific functionality.
|
|
|
|
use crate::error::{GnnError, GnnResult};
|
|
use crate::graph::BrainGraph;
|
|
use ndarray::{Array1, Array2, Axis};
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// Configuration for BrainConv layer
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct BrainConvConfig {
|
|
/// Input feature dimension
|
|
pub in_features: usize,
|
|
/// Output feature dimension
|
|
pub out_features: usize,
|
|
/// Whether to add self-loops
|
|
pub add_self_loops: bool,
|
|
/// Whether to normalize adjacency
|
|
pub normalize: bool,
|
|
/// Activation function
|
|
pub activation: Activation,
|
|
/// Dropout probability
|
|
pub dropout: f64,
|
|
}
|
|
|
|
impl Default for BrainConvConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
in_features: 64,
|
|
out_features: 64,
|
|
add_self_loops: true,
|
|
normalize: true,
|
|
activation: Activation::ReLU,
|
|
dropout: 0.0,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Activation functions
|
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
|
|
pub enum Activation {
|
|
/// No activation
|
|
None,
|
|
/// ReLU
|
|
ReLU,
|
|
/// Leaky ReLU
|
|
LeakyReLU,
|
|
/// ELU
|
|
ELU,
|
|
/// Tanh
|
|
Tanh,
|
|
/// Sigmoid
|
|
Sigmoid,
|
|
/// GELU
|
|
GELU,
|
|
}
|
|
|
|
impl Activation {
|
|
/// Apply activation function
|
|
pub fn apply(&self, x: f64) -> f64 {
|
|
match self {
|
|
Activation::None => x,
|
|
Activation::ReLU => x.max(0.0),
|
|
Activation::LeakyReLU => {
|
|
if x > 0.0 {
|
|
x
|
|
} else {
|
|
0.01 * x
|
|
}
|
|
}
|
|
Activation::ELU => {
|
|
if x > 0.0 {
|
|
x
|
|
} else {
|
|
x.exp() - 1.0
|
|
}
|
|
}
|
|
Activation::Tanh => x.tanh(),
|
|
Activation::Sigmoid => 1.0 / (1.0 + (-x).exp()),
|
|
Activation::GELU => {
|
|
// Approximate GELU
|
|
0.5 * x * (1.0 + (0.7978845608 * (x + 0.044715 * x * x * x)).tanh())
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Apply to array
|
|
pub fn apply_array(&self, x: &Array2<f64>) -> Array2<f64> {
|
|
x.mapv(|v| self.apply(v))
|
|
}
|
|
}
|
|
|
|
/// Brain-specific graph convolution layer
|
|
///
|
|
/// Extends standard GCN with brain topology awareness.
|
|
#[derive(Debug, Clone)]
|
|
pub struct BrainConv {
|
|
config: BrainConvConfig,
|
|
/// Weight matrix [in_features, out_features]
|
|
weights: Array2<f64>,
|
|
/// Bias vector [out_features]
|
|
bias: Array1<f64>,
|
|
}
|
|
|
|
impl BrainConv {
|
|
/// Create a new BrainConv layer
|
|
pub fn new(config: BrainConvConfig) -> GnnResult<Self> {
|
|
// Xavier initialization
|
|
let scale = (2.0 / (config.in_features + config.out_features) as f64).sqrt();
|
|
|
|
let weights = Array2::from_shape_fn((config.in_features, config.out_features), |_| {
|
|
(rand_simple() * 2.0 - 1.0) * scale
|
|
});
|
|
|
|
let bias = Array1::zeros(config.out_features);
|
|
|
|
Ok(Self {
|
|
config,
|
|
weights,
|
|
bias,
|
|
})
|
|
}
|
|
|
|
/// Forward pass
|
|
pub fn forward(&self, x: &Array2<f64>, adj: &Array2<f64>) -> GnnResult<Array2<f64>> {
|
|
let n = x.nrows();
|
|
if adj.nrows() != n || adj.ncols() != n {
|
|
return Err(GnnError::DimensionMismatch(format!(
|
|
"Adjacency {}x{} doesn't match features {}",
|
|
adj.nrows(),
|
|
adj.ncols(),
|
|
n
|
|
)));
|
|
}
|
|
|
|
// Normalize adjacency
|
|
let adj_norm = if self.config.normalize {
|
|
normalize_adjacency(adj, self.config.add_self_loops)
|
|
} else if self.config.add_self_loops {
|
|
adj + &Array2::<f64>::eye(n)
|
|
} else {
|
|
adj.clone()
|
|
};
|
|
|
|
// Message passing: A * X * W + b
|
|
let h = adj_norm.dot(x).dot(&self.weights);
|
|
let mut output = h + &self.bias;
|
|
|
|
// Activation
|
|
output = self.config.activation.apply_array(&output);
|
|
|
|
Ok(output)
|
|
}
|
|
|
|
/// Forward with brain graph
|
|
pub fn forward_graph(&self, graph: &BrainGraph) -> GnnResult<Array2<f64>> {
|
|
let x = graph.node_feature_matrix()?;
|
|
let adj = graph.adjacency_matrix();
|
|
self.forward(&x, &adj)
|
|
}
|
|
}
|
|
|
|
/// Configuration for BrainAttention layer
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct BrainAttentionConfig {
|
|
/// Input feature dimension
|
|
pub in_features: usize,
|
|
/// Output feature dimension
|
|
pub out_features: usize,
|
|
/// Number of attention heads
|
|
pub n_heads: usize,
|
|
/// Dropout probability
|
|
pub dropout: f64,
|
|
/// Whether to concatenate heads (true) or average (false)
|
|
pub concat_heads: bool,
|
|
/// Negative slope for LeakyReLU
|
|
pub negative_slope: f64,
|
|
}
|
|
|
|
impl Default for BrainAttentionConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
in_features: 64,
|
|
out_features: 64,
|
|
n_heads: 4,
|
|
dropout: 0.0,
|
|
concat_heads: true,
|
|
negative_slope: 0.2,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Brain-specific graph attention layer
|
|
#[derive(Debug, Clone)]
|
|
pub struct BrainAttention {
|
|
config: BrainAttentionConfig,
|
|
/// Weight matrices per head [n_heads, in_features, out_features/n_heads]
|
|
weights: Vec<Array2<f64>>,
|
|
/// Attention weights [n_heads, 2 * out_features/n_heads]
|
|
attention: Vec<Array1<f64>>,
|
|
}
|
|
|
|
impl BrainAttention {
|
|
/// Create a new BrainAttention layer
|
|
pub fn new(config: BrainAttentionConfig) -> GnnResult<Self> {
|
|
let head_dim = if config.concat_heads {
|
|
config.out_features / config.n_heads
|
|
} else {
|
|
config.out_features
|
|
};
|
|
|
|
let scale = (2.0 / (config.in_features + head_dim) as f64).sqrt();
|
|
|
|
let mut weights = Vec::with_capacity(config.n_heads);
|
|
let mut attention = Vec::with_capacity(config.n_heads);
|
|
|
|
for _ in 0..config.n_heads {
|
|
weights.push(Array2::from_shape_fn(
|
|
(config.in_features, head_dim),
|
|
|_| (rand_simple() * 2.0 - 1.0) * scale,
|
|
));
|
|
attention.push(Array1::from_shape_fn(2 * head_dim, |_| {
|
|
(rand_simple() * 2.0 - 1.0) * scale
|
|
}));
|
|
}
|
|
|
|
Ok(Self {
|
|
config,
|
|
weights,
|
|
attention,
|
|
})
|
|
}
|
|
|
|
/// Forward pass
|
|
pub fn forward(&self, x: &Array2<f64>, adj: &Array2<f64>) -> GnnResult<Array2<f64>> {
|
|
let n = x.nrows();
|
|
let head_dim = self.weights[0].ncols();
|
|
|
|
// Process each attention head
|
|
let mut head_outputs = Vec::with_capacity(self.config.n_heads);
|
|
|
|
for head in 0..self.config.n_heads {
|
|
// Transform features
|
|
let h = x.dot(&self.weights[head]); // [n, head_dim]
|
|
|
|
// Compute attention coefficients
|
|
let mut attn_matrix = Array2::zeros((n, n));
|
|
|
|
for i in 0..n {
|
|
for j in 0..n {
|
|
if adj[[i, j]] > 0.0 || i == j {
|
|
// Concatenate hi || hj
|
|
let mut concat: Vec<f64> = Vec::with_capacity(2 * head_dim);
|
|
concat.extend(h.row(i).iter().copied());
|
|
concat.extend(h.row(j).iter().copied());
|
|
|
|
// Attention score
|
|
let score: f64 = concat
|
|
.iter()
|
|
.zip(self.attention[head].iter())
|
|
.map(|(&c, &a)| c * a)
|
|
.sum();
|
|
|
|
// LeakyReLU
|
|
attn_matrix[[i, j]] = if score > 0.0 {
|
|
score
|
|
} else {
|
|
self.config.negative_slope * score
|
|
};
|
|
}
|
|
}
|
|
}
|
|
|
|
// Softmax over neighbors
|
|
for i in 0..n {
|
|
let row = attn_matrix.row(i);
|
|
let max_val = row.iter().copied().fold(f64::NEG_INFINITY, f64::max);
|
|
let exp_sum: f64 = row.iter().map(|&v| (v - max_val).exp()).sum();
|
|
|
|
for j in 0..n {
|
|
if adj[[i, j]] > 0.0 || i == j {
|
|
attn_matrix[[i, j]] = (attn_matrix[[i, j]] - max_val).exp() / exp_sum;
|
|
} else {
|
|
attn_matrix[[i, j]] = 0.0;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Aggregate
|
|
let head_out = attn_matrix.dot(&h);
|
|
head_outputs.push(head_out);
|
|
}
|
|
|
|
// Combine heads
|
|
if self.config.concat_heads {
|
|
// Concatenate along feature dimension
|
|
let mut output = Array2::zeros((n, self.config.out_features));
|
|
for (i, head_out) in head_outputs.iter().enumerate() {
|
|
for j in 0..n {
|
|
for k in 0..head_dim {
|
|
output[[j, i * head_dim + k]] = head_out[[j, k]];
|
|
}
|
|
}
|
|
}
|
|
Ok(output)
|
|
} else {
|
|
// Average heads
|
|
let mut output = Array2::zeros((n, self.config.out_features));
|
|
for head_out in &head_outputs {
|
|
output += head_out;
|
|
}
|
|
output /= self.config.n_heads as f64;
|
|
Ok(output)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Configuration for BrainPool layer
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct BrainPoolConfig {
|
|
/// Pooling method
|
|
pub method: PoolMethod,
|
|
/// Pooling ratio (for TopK, SAGPool)
|
|
pub ratio: f64,
|
|
}
|
|
|
|
impl Default for BrainPoolConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
method: PoolMethod::Mean,
|
|
ratio: 0.5,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Pooling methods
|
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
|
|
pub enum PoolMethod {
|
|
/// Mean pooling (global)
|
|
Mean,
|
|
/// Max pooling (global)
|
|
Max,
|
|
/// Sum pooling (global)
|
|
Sum,
|
|
/// Attention-weighted pooling
|
|
Attention,
|
|
/// Region-based pooling (pools by brain region)
|
|
Region,
|
|
}
|
|
|
|
/// Brain-specific graph pooling layer
|
|
#[derive(Debug, Clone)]
|
|
pub struct BrainPool {
|
|
config: BrainPoolConfig,
|
|
/// Attention weights for attention pooling
|
|
attention_weights: Option<Array1<f64>>,
|
|
}
|
|
|
|
impl BrainPool {
|
|
/// Create a new BrainPool layer
|
|
pub fn new(config: BrainPoolConfig) -> Self {
|
|
Self {
|
|
config,
|
|
attention_weights: None,
|
|
}
|
|
}
|
|
|
|
/// Initialize attention weights
|
|
pub fn init_attention(&mut self, n_features: usize) {
|
|
self.attention_weights = Some(Array1::from_shape_fn(n_features, |_| rand_simple() * 0.1));
|
|
}
|
|
|
|
/// Global pooling (for graph-level predictions)
|
|
pub fn global_pool(&self, x: &Array2<f64>) -> Array1<f64> {
|
|
match self.config.method {
|
|
PoolMethod::Mean => x.mean_axis(Axis(0)).unwrap(),
|
|
PoolMethod::Max => {
|
|
let mut result = Array1::zeros(x.ncols());
|
|
for j in 0..x.ncols() {
|
|
result[j] = x
|
|
.column(j)
|
|
.iter()
|
|
.copied()
|
|
.fold(f64::NEG_INFINITY, f64::max);
|
|
}
|
|
result
|
|
}
|
|
PoolMethod::Sum => x.sum_axis(Axis(0)),
|
|
PoolMethod::Attention => {
|
|
if let Some(ref attn) = self.attention_weights {
|
|
// Compute attention scores
|
|
let scores: Vec<f64> = x
|
|
.rows()
|
|
.into_iter()
|
|
.map(|row| {
|
|
row.iter()
|
|
.zip(attn.iter())
|
|
.map(|(&r, &a)| r * a)
|
|
.sum::<f64>()
|
|
})
|
|
.collect();
|
|
|
|
// Softmax
|
|
let max_score = scores.iter().copied().fold(f64::NEG_INFINITY, f64::max);
|
|
let exp_scores: Vec<f64> =
|
|
scores.iter().map(|&s| (s - max_score).exp()).collect();
|
|
let sum_exp: f64 = exp_scores.iter().sum();
|
|
let attn_probs: Vec<f64> = exp_scores.iter().map(|&e| e / sum_exp).collect();
|
|
|
|
// Weighted sum
|
|
let mut result = Array1::zeros(x.ncols());
|
|
for (i, &prob) in attn_probs.iter().enumerate() {
|
|
for j in 0..x.ncols() {
|
|
result[j] += prob * x[[i, j]];
|
|
}
|
|
}
|
|
result
|
|
} else {
|
|
x.mean_axis(Axis(0)).unwrap()
|
|
}
|
|
}
|
|
PoolMethod::Region => {
|
|
// Falls back to mean for now
|
|
x.mean_axis(Axis(0)).unwrap()
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Region-based pooling using brain graph
|
|
pub fn pool_by_region(&self, graph: &BrainGraph) -> GnnResult<Array2<f64>> {
|
|
use crate::graph::BrainRegion;
|
|
|
|
let x = graph.node_feature_matrix()?;
|
|
let n_features = x.ncols();
|
|
|
|
// Pool by region
|
|
let regions = [
|
|
BrainRegion::Frontal,
|
|
BrainRegion::Central,
|
|
BrainRegion::Temporal,
|
|
BrainRegion::Parietal,
|
|
BrainRegion::Occipital,
|
|
];
|
|
|
|
let mut pooled = Array2::zeros((regions.len(), n_features));
|
|
|
|
for (r_idx, ®ion) in regions.iter().enumerate() {
|
|
let indices: Vec<usize> = graph
|
|
.nodes
|
|
.iter()
|
|
.filter(|n| n.region == region)
|
|
.map(|n| n.index)
|
|
.collect();
|
|
|
|
if indices.is_empty() {
|
|
continue;
|
|
}
|
|
|
|
// Mean pooling over region nodes
|
|
for &i in &indices {
|
|
for j in 0..n_features {
|
|
pooled[[r_idx, j]] += x[[i, j]];
|
|
}
|
|
}
|
|
for j in 0..n_features {
|
|
pooled[[r_idx, j]] /= indices.len() as f64;
|
|
}
|
|
}
|
|
|
|
Ok(pooled)
|
|
}
|
|
}
|
|
|
|
/// Edge convolution for learning edge representations
|
|
#[derive(Debug, Clone)]
|
|
pub struct EdgeConv {
|
|
/// Input edge feature dimension
|
|
in_features: usize,
|
|
/// Output edge feature dimension
|
|
out_features: usize,
|
|
/// Weight matrix
|
|
weights: Array2<f64>,
|
|
/// Bias
|
|
bias: Array1<f64>,
|
|
/// Activation
|
|
activation: Activation,
|
|
}
|
|
|
|
impl EdgeConv {
|
|
/// Create a new EdgeConv layer
|
|
pub fn new(in_features: usize, out_features: usize) -> Self {
|
|
let scale = (2.0 / (in_features + out_features) as f64).sqrt();
|
|
|
|
Self {
|
|
in_features,
|
|
out_features,
|
|
weights: Array2::from_shape_fn((in_features, out_features), |_| {
|
|
(rand_simple() * 2.0 - 1.0) * scale
|
|
}),
|
|
bias: Array1::zeros(out_features),
|
|
activation: Activation::ReLU,
|
|
}
|
|
}
|
|
|
|
/// Forward pass on edge features
|
|
pub fn forward(&self, edge_features: &Array2<f64>) -> Array2<f64> {
|
|
let h = edge_features.dot(&self.weights) + &self.bias;
|
|
self.activation.apply_array(&h)
|
|
}
|
|
}
|
|
|
|
// Simple pseudo-random for initialization (deterministic for reproducibility)
|
|
fn rand_simple() -> f64 {
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
static SEED: AtomicU64 = AtomicU64::new(12345);
|
|
|
|
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)
|
|
}
|
|
|
|
/// Normalize adjacency matrix with symmetric normalization
|
|
fn normalize_adjacency(adj: &Array2<f64>, add_self_loops: bool) -> Array2<f64> {
|
|
let n = adj.nrows();
|
|
let mut a = adj.clone();
|
|
|
|
if add_self_loops {
|
|
for i in 0..n {
|
|
a[[i, i]] = 1.0;
|
|
}
|
|
}
|
|
|
|
// Compute degree
|
|
let degrees: Vec<f64> = a.sum_axis(Axis(1)).to_vec();
|
|
|
|
// D^{-1/2} A D^{-1/2}
|
|
let mut norm = Array2::zeros((n, n));
|
|
for i in 0..n {
|
|
for j in 0..n {
|
|
if a[[i, j]] > 0.0 {
|
|
let d_i = degrees[i].max(1e-10);
|
|
let d_j = degrees[j].max(1e-10);
|
|
norm[[i, j]] = a[[i, j]] / (d_i * d_j).sqrt();
|
|
}
|
|
}
|
|
}
|
|
|
|
norm
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_brain_conv() {
|
|
let config = BrainConvConfig {
|
|
in_features: 4,
|
|
out_features: 8,
|
|
..Default::default()
|
|
};
|
|
|
|
let layer = BrainConv::new(config).unwrap();
|
|
|
|
let x = Array2::from_shape_fn((5, 4), |_| rand_simple());
|
|
let adj = Array2::from_shape_fn((5, 5), |_| if rand_simple() > 0.5 { 1.0 } else { 0.0 });
|
|
|
|
let output = layer.forward(&x, &adj).unwrap();
|
|
assert_eq!(output.shape(), &[5, 8]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_brain_attention() {
|
|
let config = BrainAttentionConfig {
|
|
in_features: 4,
|
|
out_features: 8,
|
|
n_heads: 2,
|
|
concat_heads: true,
|
|
..Default::default()
|
|
};
|
|
|
|
let layer = BrainAttention::new(config).unwrap();
|
|
|
|
let x = Array2::from_shape_fn((5, 4), |_| rand_simple());
|
|
let mut adj = Array2::zeros((5, 5));
|
|
for i in 0..5 {
|
|
adj[[i, i]] = 1.0;
|
|
if i > 0 {
|
|
adj[[i, i - 1]] = 1.0;
|
|
}
|
|
}
|
|
|
|
let output = layer.forward(&x, &adj).unwrap();
|
|
assert_eq!(output.shape(), &[5, 8]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_brain_pool() {
|
|
let x = Array2::from_shape_fn((10, 4), |_| rand_simple());
|
|
|
|
// Mean pooling
|
|
let pool = BrainPool::new(BrainPoolConfig {
|
|
method: PoolMethod::Mean,
|
|
..Default::default()
|
|
});
|
|
let pooled = pool.global_pool(&x);
|
|
assert_eq!(pooled.len(), 4);
|
|
|
|
// Max pooling
|
|
let pool_max = BrainPool::new(BrainPoolConfig {
|
|
method: PoolMethod::Max,
|
|
..Default::default()
|
|
});
|
|
let pooled_max = pool_max.global_pool(&x);
|
|
assert_eq!(pooled_max.len(), 4);
|
|
}
|
|
|
|
#[test]
|
|
fn test_activation() {
|
|
assert_eq!(Activation::ReLU.apply(-1.0), 0.0);
|
|
assert_eq!(Activation::ReLU.apply(1.0), 1.0);
|
|
assert!((Activation::Sigmoid.apply(0.0) - 0.5).abs() < 1e-10);
|
|
assert!(Activation::Tanh.apply(0.0).abs() < 1e-10);
|
|
}
|
|
}
|