torsh-graph 0.2.0

//! GraphSAGE (Sample and Aggregate) layer implementation
/// Crate-local result alias: the error type defaults to [`TorshError`],
/// so both `Result<T>` and `Result<T, OtherError>` stay valid.
type Result<T, E = torsh_core::error::TorshError> = std::result::Result<T, E>;

use crate::parameter::Parameter;
use crate::{GraphData, GraphLayer};
use torsh_tensor::{
    creation::{randn, zeros},
    Tensor,
};

/// GraphSAGE convolution layer
#[derive(Debug)]
pub struct SAGEConv {
    in_features: usize,
    out_features: usize,
    weight_neighbor: Parameter,
    weight_self: Parameter,
    bias: Option<Parameter>,
}

impl SAGEConv {
    /// Create a new GraphSAGE convolution layer
    pub fn new(in_features: usize, out_features: usize, bias: bool) -> Result<Self> {
        let weight_neighbor = Parameter::new(randn(&[in_features, out_features])?);
        let weight_self = Parameter::new(randn(&[in_features, out_features])?);
        let bias = if bias {
            Some(Parameter::new(zeros(&[out_features])?))
        } else {
            None
        };

        Ok(Self {
            in_features,
            out_features,
            weight_neighbor,
            weight_self,
            bias,
        })
    }

    /// Get input feature dimension
    pub fn in_features(&self) -> usize {
        self.in_features
    }

    /// Get output feature dimension
    pub fn out_features(&self) -> usize {
        self.out_features
    }

    /// Apply GraphSAGE convolution
    pub fn forward(&self, graph: &GraphData) -> Result<GraphData> {
        let num_nodes = graph.num_nodes;
        let edge_data = crate::utils::tensor_to_vec2::<f32>(&graph.edge_index)?;

        // Build adjacency list for efficient neighbor aggregation
        let mut adjacency_list: Vec<Vec<usize>> = vec![Vec::new(); num_nodes];
        for j in 0..edge_data[0].len() {
            let src = edge_data[0][j] as usize;
            let dst = edge_data[1][j] as usize;
            adjacency_list[dst].push(src);
        }

        // Aggregate neighbor features (mean aggregation)
        let mut neighbor_features = zeros(&[num_nodes, self.in_features])?;

        for node in 0..num_nodes {
            if !adjacency_list[node].is_empty() {
                let mut aggregated = zeros(&[self.in_features])?;

                for &neighbor in &adjacency_list[node] {
                    let neighbor_slice = graph.x.slice(0, neighbor, neighbor + 1)?.to_tensor()?;
                    let neighbor_feat = neighbor_slice.squeeze(0)?;
                    aggregated = aggregated.add(&neighbor_feat)?;
                }

                // Mean aggregation
                aggregated = aggregated.div_scalar(adjacency_list[node].len() as f32)?;
                // Store aggregated features for this node
                let aggregated_data = aggregated.to_vec()?;
                for (i, &value) in aggregated_data.iter().enumerate() {
                    neighbor_features.set_item(&[node, i], value)?;
                }
            }
        }

        // Transform neighbor features and self features
        let neighbor_transformed = neighbor_features.matmul(&self.weight_neighbor.clone_data())?;
        let self_transformed = graph.x.matmul(&self.weight_self.clone_data())?;

        // Combine neighbor and self representations
        let mut output_features = neighbor_transformed.add(&self_transformed)?;

        // Add bias if present
        if let Some(ref bias) = self.bias {
            output_features = output_features.add(&bias.clone_data())?;
        }

        // L2 normalize the output features (common in GraphSAGE)
        // For simplicity, using standard normalization instead of row-wise normalization
        let norm_val = output_features.norm()?;
        let epsilon = 1e-8_f32;
        let norm_scalar = norm_val.item()?.max(epsilon);
        output_features = output_features.div_scalar(norm_scalar)?;

        // Create output graph
        Ok(GraphData {
            x: output_features,
            edge_index: graph.edge_index.clone(),
            edge_attr: graph.edge_attr.clone(),
            batch: graph.batch.clone(),
            num_nodes: graph.num_nodes,
            num_edges: graph.num_edges,
        })
    }
}

impl GraphLayer for SAGEConv {
    fn forward(&self, graph: &GraphData) -> Result<GraphData> {
        self.forward(graph)
    }

    fn parameters(&self) -> Vec<Tensor> {
        let mut params = vec![
            self.weight_neighbor.clone_data(),
            self.weight_self.clone_data(),
        ];
        if let Some(ref bias) = self.bias {
            params.push(bias.clone_data());
        }
        params
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use torsh_core::device::DeviceType;
    use torsh_tensor::creation::from_vec;

    #[test]
    fn test_sage_creation() {
        let sage = SAGEConv::new(10, 20, true);
        let params = sage.expect("operation should succeed").parameters();
        assert_eq!(params.len(), 3); // weight_neighbor + weight_self + bias
    }

    #[test]
    fn test_sage_forward() {
        let sage = SAGEConv::new(4, 8, false);

        // Create test graph
        let x = from_vec(
            vec![
                1.0, 2.0, 3.0, 4.0, // node 0
                5.0, 6.0, 7.0, 8.0, // node 1
                9.0, 10.0, 11.0, 12.0, // node 2
            ],
            &[3, 4],
            DeviceType::Cpu,
        )
        .expect("operation should succeed");
        let edge_index = from_vec(vec![0.0, 1.0, 2.0, 1.0, 2.0, 0.0], &[2, 3], DeviceType::Cpu)
            .expect("from vec should succeed");
        let graph = GraphData::new(x, edge_index);

        let output = sage
            .expect("operation should succeed")
            .forward(&graph)
            .expect("operation should succeed");
        assert_eq!(output.x.shape().dims(), &[3, 8]);
        assert_eq!(output.num_nodes, 3);

        // Check that output is finite (simplified test since norm_dim doesn't exist)
        let output_values = output
            .x
            .to_vec()
            .expect("tensor to_vec conversion should succeed");
        for &val in &output_values {
            assert!(val.is_finite(), "Output should be finite");
        }
    }
}