308 lines
8.3 KiB
Rust
308 lines
8.3 KiB
Rust
use crate::error::{GeomError, Result};
|
|
use rtx_tensor::Tensor;
|
|
use slotmap::{DefaultKey, Key, SlotMap};
|
|
use std::collections::HashMap;
|
|
|
|
/// Unique identifier for nodes in the graph
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
|
pub struct NodeId(DefaultKey);
|
|
|
|
impl NodeId {
|
|
pub fn from_raw(idx: usize) -> Self {
|
|
// This is unsafe but needed for testing invalid node IDs
|
|
unsafe { std::mem::transmute(idx as u64) }
|
|
}
|
|
|
|
pub fn index(&self) -> usize {
|
|
self.0.data().as_ffi() as usize
|
|
}
|
|
}
|
|
|
|
/// Unique identifier for edges in the graph
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
|
pub struct EdgeId(DefaultKey);
|
|
|
|
impl EdgeId {
|
|
pub fn index(&self) -> usize {
|
|
self.0.data().as_ffi() as usize
|
|
}
|
|
}
|
|
|
|
/// A node in the graph with associated features
|
|
#[derive(Debug, Clone)]
|
|
pub struct Node {
|
|
features: Tensor,
|
|
metadata: HashMap<String, String>,
|
|
}
|
|
|
|
impl Node {
|
|
pub fn new(features: Tensor) -> Self {
|
|
Self {
|
|
features,
|
|
metadata: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
pub fn with_metadata(features: Tensor, metadata: HashMap<String, String>) -> Self {
|
|
Self { features, metadata }
|
|
}
|
|
|
|
pub fn features(&self) -> &Tensor {
|
|
&self.features
|
|
}
|
|
|
|
pub fn features_mut(&mut self) -> &mut Tensor {
|
|
&mut self.features
|
|
}
|
|
|
|
pub fn metadata(&self) -> &HashMap<String, String> {
|
|
&self.metadata
|
|
}
|
|
|
|
pub fn set_metadata(&mut self, key: String, value: String) {
|
|
self.metadata.insert(key, value);
|
|
}
|
|
}
|
|
|
|
/// An edge in the graph with weight and optional features
|
|
#[derive(Debug, Clone)]
|
|
pub struct Edge {
|
|
weight: Tensor,
|
|
features: Option<Tensor>,
|
|
metadata: HashMap<String, String>,
|
|
}
|
|
|
|
impl Edge {
|
|
pub fn new(weight: Tensor) -> Self {
|
|
Self {
|
|
weight,
|
|
features: None,
|
|
metadata: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
pub fn with_features(weight: Tensor, features: Tensor) -> Self {
|
|
Self {
|
|
weight,
|
|
features: Some(features),
|
|
metadata: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
pub fn weight(&self) -> &Tensor {
|
|
&self.weight
|
|
}
|
|
|
|
pub fn weight_mut(&mut self) -> &mut Tensor {
|
|
&mut self.weight
|
|
}
|
|
|
|
pub fn features(&self) -> Option<&Tensor> {
|
|
self.features.as_ref()
|
|
}
|
|
|
|
pub fn features_mut(&mut self) -> Option<&mut Tensor> {
|
|
self.features.as_mut()
|
|
}
|
|
|
|
pub fn metadata(&self) -> &HashMap<String, String> {
|
|
&self.metadata
|
|
}
|
|
|
|
pub fn set_metadata(&mut self, key: String, value: String) {
|
|
self.metadata.insert(key, value);
|
|
}
|
|
}
|
|
|
|
/// Graph data structure for storing nodes and edges
|
|
#[derive(Debug, Clone)]
|
|
pub struct Graph {
|
|
nodes: SlotMap<DefaultKey, Node>,
|
|
edges: SlotMap<DefaultKey, (NodeId, NodeId, Edge)>,
|
|
incoming_adj: HashMap<NodeId, Vec<(NodeId, EdgeId)>>, // incoming edges: to <- [(from, edge_id)]
|
|
outgoing_adj: HashMap<NodeId, Vec<(NodeId, EdgeId)>>, // outgoing edges: from -> [(to, edge_id)]
|
|
}
|
|
|
|
impl Graph {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
nodes: SlotMap::new(),
|
|
edges: SlotMap::new(),
|
|
incoming_adj: HashMap::new(),
|
|
outgoing_adj: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
pub fn add_node(&mut self, node: Node) -> NodeId {
|
|
let key = self.nodes.insert(node);
|
|
|
|
NodeId(key)
|
|
}
|
|
|
|
pub fn add_edge(&mut self, from: NodeId, to: NodeId, edge: Edge) -> Result<EdgeId> {
|
|
// Validate that both nodes exist
|
|
if !self.nodes.contains_key(from.0) {
|
|
return Err(GeomError::InvalidNodeId(from.index()));
|
|
}
|
|
if !self.nodes.contains_key(to.0) {
|
|
return Err(GeomError::InvalidNodeId(to.index()));
|
|
}
|
|
|
|
let key = self.edges.insert((from, to, edge));
|
|
let edge_id = EdgeId(key);
|
|
|
|
// Update both adjacency lists
|
|
self.incoming_adj
|
|
.entry(to)
|
|
.or_default()
|
|
.push((from, edge_id));
|
|
self.outgoing_adj
|
|
.entry(from)
|
|
.or_default()
|
|
.push((to, edge_id));
|
|
|
|
Ok(edge_id)
|
|
}
|
|
|
|
pub fn node(&self, id: NodeId) -> Option<&Node> {
|
|
self.nodes.get(id.0)
|
|
}
|
|
|
|
pub fn node_mut(&mut self, id: NodeId) -> Option<&mut Node> {
|
|
self.nodes.get_mut(id.0)
|
|
}
|
|
|
|
pub fn edge(&self, id: EdgeId) -> Option<&Edge> {
|
|
self.edges.get(id.0).map(|(_, _, edge)| edge)
|
|
}
|
|
|
|
pub fn edge_mut(&mut self, id: EdgeId) -> Option<&mut Edge> {
|
|
self.edges.get_mut(id.0).map(|(_, _, edge)| edge)
|
|
}
|
|
|
|
pub fn edge_endpoints(&self, id: EdgeId) -> Option<(NodeId, NodeId)> {
|
|
self.edges.get(id.0).map(|(from, to, _)| (*from, *to))
|
|
}
|
|
|
|
pub fn node_count(&self) -> usize {
|
|
self.nodes.len()
|
|
}
|
|
|
|
pub fn edge_count(&self) -> usize {
|
|
self.edges.len()
|
|
}
|
|
|
|
pub fn nodes(&self) -> Vec<NodeId> {
|
|
self.nodes.keys().map(NodeId).collect()
|
|
}
|
|
|
|
pub fn edges(&self) -> Vec<EdgeId> {
|
|
self.edges.keys().map(EdgeId).collect()
|
|
}
|
|
|
|
/// Get outgoing neighbors (nodes that this node has edges TO)
|
|
pub fn neighbors(&self, node: NodeId) -> Vec<NodeId> {
|
|
self.outgoing_adj
|
|
.get(&node)
|
|
.map(|neighbors| {
|
|
neighbors
|
|
.iter()
|
|
.map(|(neighbor_id, _)| *neighbor_id)
|
|
.collect()
|
|
})
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
/// Get incoming neighbors (nodes that have edges TO this node) - used for message passing
|
|
pub fn incoming_neighbors(&self, node: NodeId) -> Vec<NodeId> {
|
|
self.incoming_adj
|
|
.get(&node)
|
|
.map(|neighbors| {
|
|
neighbors
|
|
.iter()
|
|
.map(|(neighbor_id, _)| *neighbor_id)
|
|
.collect()
|
|
})
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
pub fn neighbor_edges(&self, node: NodeId) -> Vec<EdgeId> {
|
|
self.outgoing_adj
|
|
.get(&node)
|
|
.map(|neighbors| neighbors.iter().map(|(_, edge_id)| *edge_id).collect())
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
pub fn incoming_neighbor_edges(&self, node: NodeId) -> Vec<EdgeId> {
|
|
self.incoming_adj
|
|
.get(&node)
|
|
.map(|neighbors| neighbors.iter().map(|(_, edge_id)| *edge_id).collect())
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
/// Generate adjacency matrix as a tensor
|
|
pub fn adjacency_matrix(&self) -> Tensor {
|
|
let n = self.node_count();
|
|
let mut matrix_data = vec![0.0f32; n * n];
|
|
|
|
// Create mapping from NodeId to matrix index
|
|
let node_to_index: HashMap<NodeId, usize> = self
|
|
.nodes()
|
|
.into_iter()
|
|
.enumerate()
|
|
.map(|(idx, node_id)| (node_id, idx))
|
|
.collect();
|
|
|
|
// Fill the adjacency matrix
|
|
for (from_id, neighbors) in &self.outgoing_adj {
|
|
if let Some(&from_idx) = node_to_index.get(from_id) {
|
|
for &(to_id, edge_id) in neighbors {
|
|
if let Some(&to_idx) = node_to_index.get(&to_id)
|
|
&& let Some(edge) = self.edge(edge_id)
|
|
{
|
|
let weight = edge.weight().to_cpu().unwrap()[0];
|
|
matrix_data[from_idx * n + to_idx] = weight;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
Tensor::from_data(matrix_data, [n, n], &rtx_tensor::Device::default()).unwrap()
|
|
}
|
|
|
|
/// Get node features as a batch tensor
|
|
pub fn node_features(&self) -> Option<Tensor> {
|
|
if self.nodes.is_empty() {
|
|
return None;
|
|
}
|
|
|
|
let first_node = self.nodes.values().next().unwrap();
|
|
let feature_dim = first_node.features().shape().dims()[0];
|
|
let num_nodes = self.node_count();
|
|
|
|
let mut features_data = Vec::with_capacity(num_nodes * feature_dim);
|
|
|
|
// Collect features in consistent order
|
|
for node_id in self.nodes() {
|
|
let node = self.node(node_id).unwrap();
|
|
let node_features = node.features().to_cpu().unwrap();
|
|
features_data.extend(node_features);
|
|
}
|
|
|
|
Some(
|
|
Tensor::from_data(
|
|
features_data,
|
|
[num_nodes, feature_dim],
|
|
&rtx_tensor::Device::default(),
|
|
)
|
|
.unwrap(),
|
|
)
|
|
}
|
|
}
|
|
|
|
impl Default for Graph {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|