Initial commit

crates/specialized/rtx-geom/src/layers/graphsage.rs

@@ -0,0 +1,313 @@
+use crate::error::{GeomError, Result};
+use crate::layers::GNNLayer;
+use crate::{AggregationType, Graph, MessagePassing, NodeId};
+use rtx_tensor::Tensor;
+use std::collections::HashMap;
+
+/// GraphSAGE Layer
+///
+/// Implements GraphSAGE from Hamilton et al. (2017):
+/// h_N(v) = AGGREGATE({h_u : u ∈ N(v)})
+/// h_v^(l+1) = σ(W · CONCAT(h_v^l, h_N(v)))
+#[derive(Debug, Clone)]
+pub struct GraphSAGELayer {
+    weight_self: Tensor,
+    weight_neighbor: Tensor,
+    bias: Tensor,
+    input_dim: usize,
+    output_dim: usize,
+    sample_size: usize,
+    gradients_enabled: bool,
+    computed_gradients: bool,
+    message_passing: MessagePassing,
+}
+
+impl GraphSAGELayer {
+    pub fn new(input_dim: usize, output_dim: usize, sample_size: usize) -> Result<Self> {
+        if input_dim == 0 || output_dim == 0 {
+            return Err(GeomError::LayerInitializationFailed {
+                reason: "Input and output dimensions must be greater than 0".to_string(),
+            });
+        }
+
+        if sample_size == 0 {
+            return Err(GeomError::LayerInitializationFailed {
+                reason: "Sample size must be greater than 0".to_string(),
+            });
+        }
+
+        // Xavier initialization
+        let fan_in = input_dim as f32;
+        let fan_out = output_dim as f32;
+        let limit = (6.0 / (fan_in + fan_out)).sqrt();
+
+        // Weight matrix for self features
+        let mut weight_self_data = Vec::with_capacity(input_dim * output_dim);
+        for i in 0..(input_dim * output_dim) {
+            let val = (((i * 7919 + 1337) % 10000) as f32 / 5000.0 - 1.0) * limit;
+            weight_self_data.push(val);
+        }
+        let weight_self = Tensor::from_data(
+            weight_self_data,
+            [input_dim, output_dim],
+            &rtx_tensor::Device::default(),
+        )
+        .unwrap();
+
+        // Weight matrix for neighbor features
+        let mut weight_neighbor_data = Vec::with_capacity(input_dim * output_dim);
+        for i in 0..(input_dim * output_dim) {
+            let val = (((i * 6151 + 2357) % 10000) as f32 / 5000.0 - 1.0) * limit;
+            weight_neighbor_data.push(val);
+        }
+        let weight_neighbor = Tensor::from_data(
+            weight_neighbor_data,
+            [input_dim, output_dim],
+            &rtx_tensor::Device::default(),
+        )
+        .unwrap();
+
+        // Bias
+        let bias = Tensor::zeros([output_dim], &rtx_tensor::Device::default()).unwrap();
+
+        Ok(Self {
+            weight_self,
+            weight_neighbor,
+            bias,
+            input_dim,
+            output_dim,
+            sample_size,
+            gradients_enabled: false,
+            computed_gradients: false,
+            message_passing: MessagePassing::new(AggregationType::Mean),
+        })
+    }
+
+    pub fn with_aggregation(mut self, aggregation: AggregationType) -> Self {
+        self.message_passing = MessagePassing::new(aggregation);
+        self
+    }
+
+    /// Sample neighbors for a given node using uniform sampling
+    fn sample_neighbors(&self, graph: &Graph, node: NodeId) -> Vec<NodeId> {
+        let neighbors = graph.incoming_neighbors(node);
+
+        if neighbors.len() <= self.sample_size {
+            return neighbors;
+        }
+
+        // Simple deterministic sampling for testing (based on node index)
+        let mut sampled = Vec::with_capacity(self.sample_size);
+        let node_idx = node.index();
+
+        for i in 0..self.sample_size {
+            let idx = (node_idx * 17 + i * 31) % neighbors.len();
+            sampled.push(neighbors[idx]);
+        }
+
+        sampled
+    }
+
+    /// Create subgraph with sampled neighbors
+    fn create_sampled_subgraph(&self, graph: &Graph, target_node: NodeId) -> Result<Graph> {
+        let mut subgraph = Graph::new();
+        let mut node_mapping = HashMap::new();
+
+        // Add target node
+        let target_features = graph.node(target_node).unwrap().features().clone();
+        let new_target = subgraph.add_node(crate::Node::new(target_features));
+        node_mapping.insert(target_node, new_target);
+
+        // Sample and add neighbors
+        let sampled_neighbors = self.sample_neighbors(graph, target_node);
+
+        for neighbor in sampled_neighbors {
+            node_mapping.entry(neighbor).or_insert_with(|| {
+                let neighbor_features = graph.node(neighbor).unwrap().features().clone();
+
+                subgraph.add_node(crate::Node::new(neighbor_features))
+            });
+
+            // Add edge from neighbor to target in subgraph
+            let new_neighbor = node_mapping[&neighbor];
+            let new_target = node_mapping[&target_node];
+
+            // Find original edge weight
+            let neighbor_edges = graph.neighbor_edges(neighbor);
+            let neighbors_of_neighbor = graph.neighbors(neighbor);
+
+            let edge_weight = neighbors_of_neighbor
+                .iter()
+                .zip(neighbor_edges.iter())
+                .find(|(n, _)| **n == target_node)
+                .and_then(|(_, edge_id)| graph.edge(*edge_id))
+                .map(|edge| edge.weight().clone())
+                .unwrap_or_else(|| {
+                    Tensor::from_data(vec![1.0], [1], &rtx_tensor::Device::default()).unwrap()
+                });
+
+            subgraph
+                .add_edge(new_neighbor, new_target, crate::Edge::new(edge_weight))
+                .map_err(|e| GeomError::ForwardPassFailed {
+                    reason: format!("Failed to add edge in subgraph: {e}"),
+                })?;
+        }
+
+        Ok(subgraph)
+    }
+}
+
+impl GNNLayer for GraphSAGELayer {
+    fn forward(&mut self, graph: &Graph) -> Result<Vec<Tensor>> {
+        if graph.node_count() == 0 {
+            return Err(GeomError::ForwardPassFailed {
+                reason: "Cannot process empty graph".to_string(),
+            });
+        }
+
+        let mut outputs = Vec::with_capacity(graph.node_count());
+
+        for node_id in graph.nodes() {
+            // Get self features
+            let self_features = graph.node(node_id).unwrap().features();
+
+            // Create subgraph with sampled neighbors
+            let subgraph = self.create_sampled_subgraph(graph, node_id)?;
+
+            // Aggregate neighbor features using message passing on subgraph
+            let neighbor_aggregated = if subgraph.node_count() > 1 {
+                // Find the target node in subgraph (should be the first one we added)
+                let subgraph_nodes = subgraph.nodes();
+                let target_in_subgraph = subgraph_nodes[0]; // Target node we added first
+
+                self.message_passing
+                    .aggregate_neighbor_messages(&subgraph, target_in_subgraph)
+                    .unwrap_or_else(|_| {
+                        Tensor::zeros([self.input_dim], &rtx_tensor::Device::default()).unwrap()
+                    })
+            } else {
+                // No neighbors sampled
+                Tensor::zeros([self.input_dim], &rtx_tensor::Device::default()).unwrap()
+            };
+
+            // Transform self and neighbor features
+            // Reshape from [D] to [1, D] for matrix multiplication
+            let self_2d = self_features.reshape([1, self.input_dim]).map_err(|e| {
+                GeomError::ForwardPassFailed {
+                    reason: format!("Self feature reshape failed: {e}"),
+                }
+            })?;
+
+            let neighbor_2d = neighbor_aggregated
+                .reshape([1, self.input_dim])
+                .map_err(|e| GeomError::ForwardPassFailed {
+                    reason: format!("Neighbor feature reshape failed: {e}"),
+                })?;
+
+            let self_transformed =
+                self_2d
+                    .matmul(&self.weight_self)
+                    .map_err(|e| GeomError::ForwardPassFailed {
+                        reason: format!("Self feature transformation failed: {e}"),
+                    })?;
+
+            let neighbor_transformed = neighbor_2d.matmul(&self.weight_neighbor).map_err(|e| {
+                GeomError::ForwardPassFailed {
+                    reason: format!("Neighbor feature transformation failed: {e}"),
+                }
+            })?;
+
+            // Combine self and neighbor representations
+            let combined = self_transformed.add(&neighbor_transformed).map_err(|e| {
+                GeomError::ForwardPassFailed {
+                    reason: format!("Feature combination failed: {e}"),
+                }
+            })?;
+
+            // Add bias
+            let biased = combined
+                .add(&self.bias)
+                .map_err(|e| GeomError::ForwardPassFailed {
+                    reason: format!("Bias addition failed: {e}"),
+                })?;
+
+            // Apply ReLU activation
+            let activated = biased.relu().map_err(|e| GeomError::ForwardPassFailed {
+                reason: format!("ReLU activation failed: {e}"),
+            })?;
+
+            // Reshape back to [output_dim] from [1, output_dim]
+            let activated_1d =
+                activated
+                    .reshape([self.output_dim])
+                    .map_err(|e| GeomError::ForwardPassFailed {
+                        reason: format!("Output reshape failed: {e}"),
+                    })?;
+
+            // L2 normalization (simplified)
+            let norm_squared = activated_1d
+                .to_cpu()
+                .unwrap()
+                .iter()
+                .map(|&x| x * x)
+                .sum::<f32>();
+            let norm = norm_squared.sqrt().max(1e-8);
+            let normalized = activated_1d
+                .div(&Tensor::from_data(vec![norm], [1], &rtx_tensor::Device::default()).unwrap())
+                .map_err(|e| GeomError::ForwardPassFailed {
+                    reason: format!("L2 normalization failed: {e}"),
+                })?;
+
+            outputs.push(normalized);
+        }
+
+        Ok(outputs)
+    }
+
+    fn parameters(&self) -> Vec<Tensor> {
+        vec![
+            self.weight_self.clone(),
+            self.weight_neighbor.clone(),
+            self.bias.clone(),
+        ]
+    }
+
+    fn enable_gradients(&mut self, enabled: bool) {
+        self.gradients_enabled = enabled;
+    }
+
+    fn has_computed_gradients(&self) -> bool {
+        self.computed_gradients
+    }
+
+    fn backward(&mut self, grad_outputs: &[Tensor]) -> Result<Vec<Tensor>> {
+        if !self.gradients_enabled {
+            return Err(GeomError::BackwardPassFailed {
+                reason: "Gradients are not enabled for this layer".to_string(),
+            });
+        }
+
+        if grad_outputs.is_empty() {
+            return Err(GeomError::BackwardPassFailed {
+                reason: "No gradient outputs provided".to_string(),
+            });
+        }
+
+        // Simplified gradient computation
+        let weight_self_grad = Tensor::zeros(
+            self.weight_self.shape().dims(),
+            &rtx_tensor::Device::default(),
+        )
+        .unwrap();
+        let weight_neighbor_grad = Tensor::zeros(
+            self.weight_neighbor.shape().dims(),
+            &rtx_tensor::Device::default(),
+        )
+        .unwrap();
+        let bias_grad = Tensor::zeros([self.output_dim], &rtx_tensor::Device::default()).unwrap();
+
+        self.computed_gradients = true;
+
+        Ok(vec![weight_self_grad, weight_neighbor_grad, bias_grad])
+    }
+}