Initial commit

crates/specialized/rtx-neuro-gnn/src/explain.rs

@@ -0,0 +1,659 @@
+//! GNN explainability for brain networks.
+//!
+//! Provides methods to understand which edges (connections) and nodes (channels)
+//! are most important for model predictions.
+
+use crate::error::{GnnError, GnnResult};
+use crate::graph::{BrainGraph, BrainRegion, Hemisphere};
+use crate::models::BrainGNN;
+use ndarray::Array2;
+use serde::{Deserialize, Serialize};
+
+/// Edge importance score
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct EdgeImportance {
+    /// Source node index
+    pub source: usize,
+    /// Target node index
+    pub target: usize,
+    /// Source node name
+    pub source_name: String,
+    /// Target node name
+    pub target_name: String,
+    /// Importance score
+    pub importance: f64,
+    /// Whether this is interhemispheric
+    pub interhemispheric: bool,
+}
+
+/// Node importance score
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct NodeImportance {
+    /// Node index
+    pub index: usize,
+    /// Node name (channel)
+    pub name: String,
+    /// Importance score
+    pub importance: f64,
+    /// Hemisphere
+    pub hemisphere: Hemisphere,
+    /// Brain region
+    pub region: BrainRegion,
+}
+
+/// Interpretation method
+#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
+pub enum InterpretationMethod {
+    /// Gradient-based saliency
+    Gradient,
+    /// Integrated gradients
+    IntegratedGradients,
+    /// Edge masking (ablation)
+    EdgeMasking,
+    /// GNNExplainer (learnable masks)
+    GnnExplainer,
+    /// Attention weights (if model uses attention)
+    Attention,
+}
+
+/// Configuration for GNN explainer
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ExplainerConfig {
+    /// Interpretation method
+    pub method: InterpretationMethod,
+    /// Number of steps for integrated gradients
+    pub n_steps: usize,
+    /// Perturbation magnitude for gradient estimation
+    pub epsilon: f64,
+    /// Top-K edges to return
+    pub top_k_edges: usize,
+    /// Top-K nodes to return
+    pub top_k_nodes: usize,
+    /// Target class for explanation (None = predicted class)
+    pub target_class: Option<usize>,
+}
+
+impl Default for ExplainerConfig {
+    fn default() -> Self {
+        Self {
+            method: InterpretationMethod::EdgeMasking,
+            n_steps: 50,
+            epsilon: 0.01,
+            top_k_edges: 20,
+            top_k_nodes: 10,
+            target_class: None,
+        }
+    }
+}
+
+/// GNN Explainer for brain networks
+pub struct GnnExplainer<M: BrainGNN> {
+    model: M,
+    config: ExplainerConfig,
+}
+
+impl<M: BrainGNN> GnnExplainer<M> {
+    /// Create a new GNN explainer
+    pub fn new(model: M, config: ExplainerConfig) -> Self {
+        Self { model, config }
+    }
+
+    /// Explain model prediction on a graph
+    pub fn explain(&self, graph: &BrainGraph) -> GnnResult<ExplanationResult> {
+        // Get baseline prediction
+        let logits = self.model.forward(graph)?;
+        let predicted_class = logits
+            .iter()
+            .enumerate()
+            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
+            .map_or(0, |(i, _)| i);
+
+        let target_class = self.config.target_class.unwrap_or(predicted_class);
+
+        // Compute edge importance based on method
+        let edge_scores = match self.config.method {
+            InterpretationMethod::EdgeMasking => self.compute_edge_masking(graph, target_class)?,
+            InterpretationMethod::Gradient | InterpretationMethod::IntegratedGradients => {
+                self.compute_gradient_importance(graph, target_class)?
+            }
+            InterpretationMethod::GnnExplainer => {
+                self.compute_gnn_explainer(graph, target_class)?
+            }
+            InterpretationMethod::Attention => {
+                // Return edge weights as importance (placeholder)
+                graph.edge_weights()
+            }
+        };
+
+        // Build edge importance list
+        let mut edge_importance: Vec<EdgeImportance> = graph
+            .edges
+            .iter()
+            .zip(edge_scores.iter())
+            .map(|(edge, &importance)| EdgeImportance {
+                source: edge.source,
+                target: edge.target,
+                source_name: graph.nodes[edge.source].name.clone(),
+                target_name: graph.nodes[edge.target].name.clone(),
+                importance,
+                interhemispheric: edge.interhemispheric,
+            })
+            .collect();
+
+        // Sort by importance (descending)
+        edge_importance.sort_by(|a, b| b.importance.partial_cmp(&a.importance).unwrap());
+        edge_importance.truncate(self.config.top_k_edges);
+
+        // Compute node importance (aggregate from edges)
+        let mut node_scores = vec![0.0; graph.n_nodes()];
+        for (edge, &score) in graph.edges.iter().zip(edge_scores.iter()) {
+            node_scores[edge.source] += score.abs();
+            node_scores[edge.target] += score.abs();
+        }
+
+        // Normalize node scores
+        let max_node = node_scores.iter().copied().fold(0.0, f64::max);
+        if max_node > 0.0 {
+            for score in &mut node_scores {
+                *score /= max_node;
+            }
+        }
+
+        // Build node importance list
+        let mut node_importance: Vec<NodeImportance> = graph
+            .nodes
+            .iter()
+            .zip(node_scores.iter())
+            .map(|(node, &importance)| NodeImportance {
+                index: node.index,
+                name: node.name.clone(),
+                importance,
+                hemisphere: node.hemisphere,
+                region: node.region,
+            })
+            .collect();
+
+        node_importance.sort_by(|a, b| b.importance.partial_cmp(&a.importance).unwrap());
+        node_importance.truncate(self.config.top_k_nodes);
+
+        Ok(ExplanationResult {
+            predicted_class,
+            target_class,
+            confidence: logits[target_class],
+            edge_importance,
+            node_importance,
+            method: self.config.method,
+        })
+    }
+
+    /// Edge masking (ablation study)
+    fn compute_edge_masking(&self, graph: &BrainGraph, target_class: usize) -> GnnResult<Vec<f64>> {
+        let baseline_logits = self.model.forward(graph)?;
+        let baseline_score = baseline_logits[target_class];
+
+        let mut importance = Vec::with_capacity(graph.n_edges());
+
+        for edge_idx in 0..graph.n_edges() {
+            // Create graph with edge removed
+            let mut masked_graph = graph.clone();
+            masked_graph.edges.remove(edge_idx);
+
+            let masked_logits = self.model.forward(&masked_graph)?;
+            let masked_score = masked_logits[target_class];
+
+            // Importance = drop in score when edge is removed
+            importance.push(baseline_score - masked_score);
+        }
+
+        // Normalize
+        let max_imp = importance.iter().copied().fold(0.0, f64::max);
+        if max_imp > 0.0 {
+            for imp in &mut importance {
+                *imp /= max_imp;
+            }
+        }
+
+        Ok(importance)
+    }
+
+    /// Gradient-based importance (finite differences approximation)
+    fn compute_gradient_importance(
+        &self,
+        graph: &BrainGraph,
+        target_class: usize,
+    ) -> GnnResult<Vec<f64>> {
+        let eps = self.config.epsilon;
+        let mut importance = Vec::with_capacity(graph.n_edges());
+
+        for edge_idx in 0..graph.n_edges() {
+            // Perturb edge weight up
+            let mut graph_plus = graph.clone();
+            graph_plus.edges[edge_idx].weight += eps;
+
+            // Perturb edge weight down
+            let mut graph_minus = graph.clone();
+            graph_minus.edges[edge_idx].weight -= eps;
+
+            let logits_plus = self.model.forward(&graph_plus)?;
+            let logits_minus = self.model.forward(&graph_minus)?;
+
+            // Finite difference gradient
+            let gradient = (logits_plus[target_class] - logits_minus[target_class]) / (2.0 * eps);
+            importance.push(gradient.abs());
+        }
+
+        // Normalize
+        let max_imp = importance.iter().copied().fold(0.0, f64::max);
+        if max_imp > 0.0 {
+            for imp in &mut importance {
+                *imp /= max_imp;
+            }
+        }
+
+        Ok(importance)
+    }
+
+    /// GNNExplainer (simplified version using iterative masking)
+    fn compute_gnn_explainer(
+        &self,
+        graph: &BrainGraph,
+        target_class: usize,
+    ) -> GnnResult<Vec<f64>> {
+        let n_edges = graph.n_edges();
+        let mut mask = vec![0.5; n_edges]; // Start with uniform mask
+
+        // Simple optimization loop
+        let lr = 0.1;
+        let n_iterations = 100;
+
+        for _ in 0..n_iterations {
+            // Compute gradient for each mask value
+            let mut gradients = vec![0.0; n_edges];
+
+            for i in 0..n_edges {
+                let eps: f64 = 0.01;
+
+                // Mask + epsilon
+                let mut mask_plus = mask.clone();
+                mask_plus[i] = (mask_plus[i] + eps).min(1.0);
+                let score_plus = self.masked_forward(graph, &mask_plus, target_class)?;
+
+                // Mask - epsilon
+                let mut mask_minus = mask.clone();
+                mask_minus[i] = (mask_minus[i] - eps).max(0.0);
+                let score_minus = self.masked_forward(graph, &mask_minus, target_class)?;
+
+                gradients[i] = (score_plus - score_minus) / (2.0 * eps);
+            }
+
+            // Update mask (maximize score)
+            for i in 0..n_edges {
+                mask[i] = (mask[i] + lr * gradients[i]).clamp(0.0, 1.0);
+            }
+        }
+
+        Ok(mask)
+    }
+
+    /// Forward pass with edge masking
+    fn masked_forward(
+        &self,
+        graph: &BrainGraph,
+        mask: &[f64],
+        target_class: usize,
+    ) -> GnnResult<f64> {
+        let mut masked_graph = graph.clone();
+
+        for (i, edge) in masked_graph.edges.iter_mut().enumerate() {
+            edge.weight *= mask[i];
+        }
+
+        let logits = self.model.forward(&masked_graph)?;
+        Ok(logits[target_class])
+    }
+}
+
+/// Result of explanation
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ExplanationResult {
+    /// Predicted class
+    pub predicted_class: usize,
+    /// Target class for explanation
+    pub target_class: usize,
+    /// Confidence (logit) for target class
+    pub confidence: f64,
+    /// Top-K important edges
+    pub edge_importance: Vec<EdgeImportance>,
+    /// Top-K important nodes
+    pub node_importance: Vec<NodeImportance>,
+    /// Method used
+    pub method: InterpretationMethod,
+}
+
+impl ExplanationResult {
+    /// Get summary statistics
+    pub fn summary(&self) -> ExplanationSummary {
+        let n_interhemi = self
+            .edge_importance
+            .iter()
+            .filter(|e| e.interhemispheric)
+            .count();
+
+        let mut region_counts: std::collections::HashMap<BrainRegion, usize> =
+            std::collections::HashMap::new();
+        for node in &self.node_importance {
+            *region_counts.entry(node.region).or_insert(0) += 1;
+        }
+
+        let top_region = region_counts
+            .into_iter()
+            .max_by_key(|&(_, count)| count)
+            .map(|(region, _)| region);
+
+        ExplanationSummary {
+            predicted_class: self.predicted_class,
+            confidence: self.confidence,
+            n_important_edges: self.edge_importance.len(),
+            n_important_nodes: self.node_importance.len(),
+            n_interhemispheric: n_interhemi,
+            top_region,
+            top_edge: self.edge_importance.first().cloned(),
+            top_node: self.node_importance.first().cloned(),
+        }
+    }
+
+    /// Get edges for specific region
+    pub fn edges_for_region(&self, _region: BrainRegion) -> Vec<&EdgeImportance> {
+        // Note: This is a simplified implementation
+        // In practice, would need to check node regions
+        self.edge_importance.iter().collect()
+    }
+
+    /// Export as connectivity matrix (importance weighted)
+    pub fn to_importance_matrix(&self, n_nodes: usize) -> Array2<f64> {
+        let mut matrix = Array2::zeros((n_nodes, n_nodes));
+
+        for edge in &self.edge_importance {
+            if edge.source < n_nodes && edge.target < n_nodes {
+                matrix[[edge.source, edge.target]] = edge.importance;
+                matrix[[edge.target, edge.source]] = edge.importance; // Symmetric
+            }
+        }
+
+        matrix
+    }
+}
+
+/// Summary of explanation
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ExplanationSummary {
+    /// Predicted class
+    pub predicted_class: usize,
+    /// Confidence score
+    pub confidence: f64,
+    /// Number of important edges identified
+    pub n_important_edges: usize,
+    /// Number of important nodes identified
+    pub n_important_nodes: usize,
+    /// Number of interhemispheric connections in top edges
+    pub n_interhemispheric: usize,
+    /// Most represented brain region
+    pub top_region: Option<BrainRegion>,
+    /// Most important edge
+    pub top_edge: Option<EdgeImportance>,
+    /// Most important node
+    pub top_node: Option<NodeImportance>,
+}
+
+/// Biomarker discovery across multiple subjects
+pub struct BiomarkerDiscovery<M: BrainGNN + Clone> {
+    model: M,
+    explainer_config: ExplainerConfig,
+}
+
+impl<M: BrainGNN + Clone> BiomarkerDiscovery<M> {
+    /// Create a new biomarker discovery instance
+    pub fn new(model: M, config: ExplainerConfig) -> Self {
+        Self {
+            model,
+            explainer_config: config,
+        }
+    }
+
+    /// Find consistent biomarkers across subjects
+    pub fn find_biomarkers(
+        &self,
+        graphs: &[BrainGraph],
+        labels: &[usize],
+    ) -> GnnResult<BiomarkerResult> {
+        if graphs.len() != labels.len() {
+            return Err(GnnError::DimensionMismatch(
+                "graphs and labels must have same length".into(),
+            ));
+        }
+
+        let explainer = GnnExplainer::new(self.model.clone(), self.explainer_config.clone());
+
+        // Collect explanations for each class
+        let n_classes = labels.iter().max().map_or(0, |&x| x + 1);
+        let mut class_edge_scores: Vec<std::collections::HashMap<(String, String), Vec<f64>>> = (0
+            ..n_classes)
+            .map(|_| std::collections::HashMap::new())
+            .collect();
+
+        for (graph, &label) in graphs.iter().zip(labels.iter()) {
+            let explanation = explainer.explain(graph)?;
+
+            for edge in &explanation.edge_importance {
+                let key = (edge.source_name.clone(), edge.target_name.clone());
+                class_edge_scores[label]
+                    .entry(key)
+                    .or_default()
+                    .push(edge.importance);
+            }
+        }
+
+        // Find consistently important edges per class
+        let mut class_biomarkers = Vec::with_capacity(n_classes);
+
+        for class_scores in class_edge_scores {
+            let mut biomarkers: Vec<(String, String, f64, f64)> = class_scores
+                .into_iter()
+                .filter(|(_, scores)| scores.len() >= 2) // At least 2 occurrences
+                .map(|((src, tgt), scores)| {
+                    let mean = scores.iter().sum::<f64>() / scores.len() as f64;
+                    let std = (scores.iter().map(|&x| (x - mean).powi(2)).sum::<f64>()
+                        / scores.len() as f64)
+                        .sqrt();
+                    (src, tgt, mean, std)
+                })
+                .collect();
+
+            // Sort by mean importance / std (consistency)
+            biomarkers.sort_by(|a, b| {
+                let score_a = a.2 / (a.3 + 0.1);
+                let score_b = b.2 / (b.3 + 0.1);
+                score_b.partial_cmp(&score_a).unwrap()
+            });
+
+            class_biomarkers.push(biomarkers);
+        }
+
+        Ok(BiomarkerResult {
+            class_biomarkers,
+            n_subjects: graphs.len(),
+            n_classes,
+        })
+    }
+}
+
+/// Result of biomarker discovery
+#[derive(Debug, Clone)]
+pub struct BiomarkerResult {
+    /// Biomarkers per class: (source, target, mean_importance, std)
+    pub class_biomarkers: Vec<Vec<(String, String, f64, f64)>>,
+    /// Number of subjects analyzed
+    pub n_subjects: usize,
+    /// Number of classes
+    pub n_classes: usize,
+}
+
+impl BiomarkerResult {
+    /// Get top biomarkers for a class
+    pub fn top_biomarkers(&self, class: usize, k: usize) -> Vec<(String, String, f64)> {
+        self.class_biomarkers
+            .get(class)
+            .map(|b| {
+                b.iter()
+                    .take(k)
+                    .map(|(src, tgt, mean, _)| (src.clone(), tgt.clone(), *mean))
+                    .collect()
+            })
+            .unwrap_or_default()
+    }
+
+    /// Get discriminative biomarkers (different between classes)
+    pub fn discriminative_biomarkers(
+        &self,
+        class_a: usize,
+        class_b: usize,
+    ) -> Vec<(String, String, f64)> {
+        let a_edges: std::collections::HashSet<_> = self
+            .class_biomarkers
+            .get(class_a)
+            .map(|b| {
+                b.iter()
+                    .take(20)
+                    .map(|(s, t, _, _)| (s.clone(), t.clone()))
+                    .collect()
+            })
+            .unwrap_or_default();
+
+        let b_edges: std::collections::HashSet<_> = self
+            .class_biomarkers
+            .get(class_b)
+            .map(|b| {
+                b.iter()
+                    .take(20)
+                    .map(|(s, t, _, _)| (s.clone(), t.clone()))
+                    .collect()
+            })
+            .unwrap_or_default();
+
+        // Find edges in A but not in B
+        a_edges
+            .difference(&b_edges)
+            .filter_map(|(s, t)| {
+                self.class_biomarkers.get(class_a).and_then(|b| {
+                    b.iter()
+                        .find(|(src, tgt, _, _)| src == s && tgt == t)
+                        .map(|(src, tgt, mean, _)| (src.clone(), tgt.clone(), *mean))
+                })
+            })
+            .collect()
+    }
+}
+
+impl<M: BrainGNN + Clone> GnnExplainer<M> {
+    /// Explain using a reference to the model
+    pub fn explain_ref(&self, graph: &BrainGraph) -> GnnResult<ExplanationResult> {
+        self.explain(graph)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::graph::{GraphBuilder, Hemisphere};
+    use crate::models::{BrainNetCNN, BrainNetCNNConfig};
+
+    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_edge_masking_explanation() {
+        let graph = create_test_graph();
+
+        let model = BrainNetCNN::new(BrainNetCNNConfig {
+            n_channels: 4,
+            n_features: 4,
+            hidden_dims: vec![8, 16],
+            n_classes: 2,
+            ..Default::default()
+        })
+        .unwrap();
+
+        let explainer = GnnExplainer::new(
+            model,
+            ExplainerConfig {
+                method: InterpretationMethod::EdgeMasking,
+                top_k_edges: 4,
+                top_k_nodes: 4,
+                ..Default::default()
+            },
+        );
+
+        let result = explainer.explain(&graph).unwrap();
+
+        assert!(result.edge_importance.len() <= 4);
+        assert!(result.node_importance.len() <= 4);
+        assert!(result.predicted_class < 2);
+    }
+
+    #[test]
+    fn test_explanation_summary() {
+        let graph = create_test_graph();
+
+        let model = BrainNetCNN::new(BrainNetCNNConfig {
+            n_channels: 4,
+            n_features: 4,
+            hidden_dims: vec![8],
+            n_classes: 2,
+            ..Default::default()
+        })
+        .unwrap();
+
+        let explainer = GnnExplainer::new(model, ExplainerConfig::default());
+        let result = explainer.explain(&graph).unwrap();
+        let summary = result.summary();
+
+        assert!(summary.n_important_edges > 0);
+        assert!(summary.predicted_class < 2);
+    }
+
+    #[test]
+    fn test_importance_matrix() {
+        let graph = create_test_graph();
+
+        let model = BrainNetCNN::new(BrainNetCNNConfig {
+            n_channels: 4,
+            n_features: 4,
+            hidden_dims: vec![8],
+            n_classes: 2,
+            ..Default::default()
+        })
+        .unwrap();
+
+        let explainer = GnnExplainer::new(model, ExplainerConfig::default());
+        let result = explainer.explain(&graph).unwrap();
+        let matrix = result.to_importance_matrix(4);
+
+        assert_eq!(matrix.shape(), &[4, 4]);
+    }
+}