69 lines
2.0 KiB
Rust
69 lines
2.0 KiB
Rust
//! # RustyTorch++ Geometry (rtx-geom)
|
|
//!
|
|
//! Graph Neural Network library with GPU acceleration for RustyTorch++
|
|
//!
|
|
//! This crate provides:
|
|
//! - Graph data structures for storing nodes and edges with features
|
|
//! - Message passing framework for GNN operations
|
|
//! - GNN layers: GCN, GAT, GraphSAGE
|
|
//! - GPU acceleration through rtx-runtime integration
|
|
//!
|
|
//! ## Example
|
|
//!
|
|
//! ```rust,no_run
|
|
//! use rtx_geom::{Graph, Node, Edge};
|
|
//! use rtx_geom::layers::GCNLayer;
|
|
//! use rtx_tensor::{Tensor, Device};
|
|
//!
|
|
//! // Create a graph
|
|
//! let mut graph = Graph::new();
|
|
//! let device = Device::default();
|
|
//!
|
|
//! // Add nodes with features
|
|
//! let node1 = graph.add_node(Node::new(Tensor::zeros(&[3], &device).unwrap()));
|
|
//! let node2 = graph.add_node(Node::new(Tensor::ones(&[3], &device).unwrap()));
|
|
//!
|
|
//! // Add edges
|
|
//! graph.add_edge(node1, node2, Edge::new(Tensor::from_data(vec![1.0], [1], &device).unwrap())).unwrap();
|
|
//!
|
|
//! // Create and use a GCN layer
|
|
//! let mut gcn = GCNLayer::new(3, 5).unwrap();
|
|
//! let output = gcn.forward(&graph).unwrap();
|
|
//! ```
|
|
|
|
pub mod error;
|
|
pub mod graph;
|
|
pub mod layers;
|
|
pub mod message;
|
|
pub mod sparse_message;
|
|
|
|
// Re-export main types for convenience
|
|
pub use error::{GeomError, Result};
|
|
pub use graph::{Edge, EdgeId, Graph, Node, NodeId};
|
|
pub use message::{AggregationType, MessagePassing};
|
|
pub use sparse_message::{
|
|
MessagePassingFactory, MessagePassingStrategy, SparseAdjacencyMatrix, SparseMessagePassing,
|
|
};
|
|
|
|
// Re-export layer types
|
|
pub use layers::{
|
|
GATLayer, GCNLayer, GNNLayer, GraphSAGELayer, GraphType, SparseGCNFactory, SparseGCNLayer,
|
|
};
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use rtx_tensor::Tensor;
|
|
|
|
#[test]
|
|
fn test_basic_functionality() {
|
|
let mut graph = Graph::new();
|
|
let node_features = Tensor::zeros([2], &rtx_tensor::Device::default()).unwrap();
|
|
let node_id = graph.add_node(Node::new(node_features));
|
|
|
|
assert_eq!(graph.node_count(), 1);
|
|
assert_eq!(graph.edge_count(), 0);
|
|
assert!(graph.node(node_id).is_some());
|
|
}
|
|
}
|