Initial commit

2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
@@ -0,0 +1,193 @@
+use crate::error::{GeomError, Result};
+use crate::layers::GNNLayer;
+use crate::{AggregationType, Graph, MessagePassing};
+use rtx_tensor::Tensor;
+
+/// Graph Convolutional Network (GCN) Layer
+///
+/// Implements the GCN layer from Kipf & Welling (2017):
+/// H^(l+1) = σ(D^(-1/2) A D^(-1/2) H^(l) W^(l))
+#[derive(Debug, Clone)]
+pub struct GCNLayer {
+    weight: Tensor,
+    bias: Tensor,
+    input_dim: usize,
+    output_dim: usize,
+    gradients_enabled: bool,
+    computed_gradients: bool,
+    message_passing: MessagePassing,
+}
+
+impl GCNLayer {
+    pub fn new(input_dim: usize, output_dim: 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(),
+            });
+        }
+
+        // Xavier/Glorot initialization
+        let fan_in = input_dim as f32;
+        let fan_out = output_dim as f32;
+        let limit = (6.0 / (fan_in + fan_out)).sqrt();
+
+        // Initialize weight matrix with random values in [-limit, limit]
+        let mut weight_data = Vec::with_capacity(input_dim * output_dim);
+        for i in 0..(input_dim * output_dim) {
+            // Simple pseudo-random initialization for deterministic testing
+            let val = (((i * 7919 + 1337) % 10000) as f32 / 5000.0 - 1.0) * limit;
+            weight_data.push(val);
+        }
+        let weight = Tensor::from_data(
+            weight_data,
+            [input_dim, output_dim],
+            &rtx_tensor::Device::default(),
+        )
+        .unwrap();
+
+        // Initialize bias to zeros
+        let bias = Tensor::zeros([output_dim], &rtx_tensor::Device::default()).unwrap();
+
+        Ok(Self {
+            weight,
+            bias,
+            input_dim,
+            output_dim,
+            gradients_enabled: false,
+            computed_gradients: false,
+            message_passing: MessagePassing::new(AggregationType::Mean),
+        })
+    }
+
+    pub fn with_message_passing(mut self, aggregation: AggregationType) -> Self {
+        self.message_passing = MessagePassing::new(aggregation);
+        self
+    }
+}
+
+impl GNNLayer for GCNLayer {
+    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(),
+            });
+        }
+
+        // Get node features
+        let node_features = graph
+            .node_features()
+            .ok_or_else(|| GeomError::ForwardPassFailed {
+                reason: "Failed to extract node features".to_string(),
+            })?;
+
+        // Validate input dimensions
+        let actual_input_dim = node_features.shape().dims()[1];
+        if actual_input_dim != self.input_dim {
+            return Err(GeomError::FeatureShapeMismatch {
+                expected: vec![graph.node_count(), self.input_dim],
+                actual: node_features.shape().dims().to_vec(),
+            });
+        }
+
+        // Compute messages via aggregation
+        let aggregated_features = self.message_passing.compute_messages(graph);
+
+        // Transform aggregated features: H * W + b
+        let mut outputs = Vec::with_capacity(aggregated_features.len());
+
+        for (_node_id, feature) in aggregated_features {
+            // Reshape feature from [D] to [1, D] for matrix multiplication
+            let feature_2d =
+                feature
+                    .reshape([1, self.input_dim])
+                    .map_err(|e| GeomError::ForwardPassFailed {
+                        reason: format!("Feature reshape failed: {e}"),
+                    })?;
+
+            // Linear transformation: [1, input_dim] * [input_dim, output_dim] = [1, output_dim]
+            let transformed =
+                feature_2d
+                    .matmul(&self.weight)
+                    .map_err(|e| GeomError::ForwardPassFailed {
+                        reason: format!("Matrix multiplication failed: {e}"),
+                    })?;
+
+            let output = transformed
+                .add(&self.bias)
+                .map_err(|e| GeomError::ForwardPassFailed {
+                    reason: format!("Bias addition failed: {e}"),
+                })?;
+
+            // Apply ReLU activation
+            let activated = output.relu().map_err(|e| GeomError::ForwardPassFailed {
+                reason: format!("ReLU activation failed: {e}"),
+            })?;
+
+            // Squeeze back to [output_dim] from [1, output_dim]
+            let squeezed =
+                activated
+                    .reshape([self.output_dim])
+                    .map_err(|e| GeomError::ForwardPassFailed {
+                        reason: format!("Output reshape failed: {e}"),
+                    })?;
+
+            outputs.push(squeezed);
+        }
+
+        Ok(outputs)
+    }
+
+    fn parameters(&self) -> Vec<Tensor> {
+        vec![self.weight.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(),
+            });
+        }
+
+        // Compute parameter gradients
+        // This is a simplified implementation for testing
+
+        // Weight gradients: sum over batch of input^T * grad_output
+        let weight_grad =
+            Tensor::zeros(self.weight.shape().dims(), &rtx_tensor::Device::default()).unwrap();
+
+        // Bias gradients: sum of grad_outputs
+        let mut bias_grad_data = vec![0.0f32; self.output_dim];
+        for grad_output in grad_outputs {
+            let grad_data = grad_output.to_cpu().unwrap();
+            for (i, &grad_val) in grad_data.iter().enumerate() {
+                if i < bias_grad_data.len() {
+                    bias_grad_data[i] += grad_val;
+                }
+            }
+        }
+        let bias_grad = Tensor::from_data(
+            bias_grad_data,
+            [self.output_dim],
+            &rtx_tensor::Device::default(),
+        )
+        .unwrap();
+
+        self.computed_gradients = true;
+
+        Ok(vec![weight_grad, bias_grad])
+    }
+}