Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,307 @@
|
||||
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()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user