Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,629 @@
|
||||
//! 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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user