rustytorch/crates/specialized/rtx-geom/src/layers/gcn.rs

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])
    }
}