Merge commit 'd803bfe2b1fe7f5e219e50ac20d6801a0a58ac75' as 'vendor/ruvector'

vendor/ruvector/crates/ruvector-attention-node/src/graph.rs

@@ -0,0 +1,430 @@
+//! NAPI-RS bindings for graph attention mechanisms
+//!
+//! Provides Node.js bindings for:
+//! - Edge-featured attention (GATv2-style)
+//! - Graph RoPE (Rotary Position Embeddings for graphs)
+//! - Dual-space attention (Euclidean + Hyperbolic)
+
+use napi::bindgen_prelude::*;
+use napi_derive::napi;
+use ruvector_attention::graph::{
+    DualSpaceAttention as RustDualSpace, DualSpaceConfig as RustDualConfig,
+    EdgeFeaturedAttention as RustEdgeFeatured, EdgeFeaturedConfig as RustEdgeConfig,
+    GraphRoPE as RustGraphRoPE, RoPEConfig as RustRoPEConfig,
+};
+use ruvector_attention::traits::Attention;
+
+// ============================================================================
+// Edge-Featured Attention
+// ============================================================================
+
+/// Configuration for edge-featured attention
+#[napi(object)]
+pub struct EdgeFeaturedConfig {
+    pub node_dim: u32,
+    pub edge_dim: u32,
+    pub num_heads: u32,
+    pub concat_heads: Option<bool>,
+    pub add_self_loops: Option<bool>,
+    pub negative_slope: Option<f64>,
+}
+
+/// Edge-featured attention (GATv2-style)
+#[napi]
+pub struct EdgeFeaturedAttention {
+    inner: RustEdgeFeatured,
+    config: EdgeFeaturedConfig,
+}
+
+#[napi]
+impl EdgeFeaturedAttention {
+    /// Create a new edge-featured attention instance
+    ///
+    /// # Arguments
+    /// * `config` - Edge-featured attention configuration
+    #[napi(constructor)]
+    pub fn new(config: EdgeFeaturedConfig) -> Self {
+        let rust_config = RustEdgeConfig {
+            node_dim: config.node_dim as usize,
+            edge_dim: config.edge_dim as usize,
+            num_heads: config.num_heads as usize,
+            concat_heads: config.concat_heads.unwrap_or(true),
+            add_self_loops: config.add_self_loops.unwrap_or(true),
+            negative_slope: config.negative_slope.unwrap_or(0.2) as f32,
+            dropout: 0.0,
+        };
+        Self {
+            inner: RustEdgeFeatured::new(rust_config),
+            config,
+        }
+    }
+
+    /// Create with simple parameters
+    #[napi(factory)]
+    pub fn simple(node_dim: u32, edge_dim: u32, num_heads: u32) -> Self {
+        Self::new(EdgeFeaturedConfig {
+            node_dim,
+            edge_dim,
+            num_heads,
+            concat_heads: Some(true),
+            add_self_loops: Some(true),
+            negative_slope: Some(0.2),
+        })
+    }
+
+    /// Compute attention without edge features (standard attention)
+    #[napi]
+    pub fn compute(
+        &self,
+        query: Float32Array,
+        keys: Vec<Float32Array>,
+        values: Vec<Float32Array>,
+    ) -> Result<Float32Array> {
+        let query_slice = query.as_ref();
+        let keys_vec: Vec<Vec<f32>> = keys.into_iter().map(|k| k.to_vec()).collect();
+        let values_vec: Vec<Vec<f32>> = values.into_iter().map(|v| v.to_vec()).collect();
+        let keys_refs: Vec<&[f32]> = keys_vec.iter().map(|k| k.as_slice()).collect();
+        let values_refs: Vec<&[f32]> = values_vec.iter().map(|v| v.as_slice()).collect();
+
+        let result = self
+            .inner
+            .compute(query_slice, &keys_refs, &values_refs)
+            .map_err(|e| Error::from_reason(e.to_string()))?;
+
+        Ok(Float32Array::new(result))
+    }
+
+    /// Compute attention with edge features
+    ///
+    /// # Arguments
+    /// * `query` - Query vector
+    /// * `keys` - Array of key vectors
+    /// * `values` - Array of value vectors
+    /// * `edge_features` - Array of edge feature vectors (same length as keys)
+    #[napi]
+    pub fn compute_with_edges(
+        &self,
+        query: Float32Array,
+        keys: Vec<Float32Array>,
+        values: Vec<Float32Array>,
+        edge_features: Vec<Float32Array>,
+    ) -> Result<Float32Array> {
+        let query_slice = query.as_ref();
+        let keys_vec: Vec<Vec<f32>> = keys.into_iter().map(|k| k.to_vec()).collect();
+        let values_vec: Vec<Vec<f32>> = values.into_iter().map(|v| v.to_vec()).collect();
+        let edge_features_vec: Vec<Vec<f32>> =
+            edge_features.into_iter().map(|e| e.to_vec()).collect();
+
+        let keys_refs: Vec<&[f32]> = keys_vec.iter().map(|k| k.as_slice()).collect();
+        let values_refs: Vec<&[f32]> = values_vec.iter().map(|v| v.as_slice()).collect();
+        let edges_refs: Vec<&[f32]> = edge_features_vec.iter().map(|e| e.as_slice()).collect();
+
+        let result = self
+            .inner
+            .compute_with_edges(query_slice, &keys_refs, &values_refs, &edges_refs)
+            .map_err(|e| Error::from_reason(e.to_string()))?;
+
+        Ok(Float32Array::new(result))
+    }
+
+    /// Get the node dimension
+    #[napi(getter)]
+    pub fn node_dim(&self) -> u32 {
+        self.config.node_dim
+    }
+
+    /// Get the edge dimension
+    #[napi(getter)]
+    pub fn edge_dim(&self) -> u32 {
+        self.config.edge_dim
+    }
+
+    /// Get the number of heads
+    #[napi(getter)]
+    pub fn num_heads(&self) -> u32 {
+        self.config.num_heads
+    }
+}
+
+// ============================================================================
+// Graph RoPE Attention
+// ============================================================================
+
+/// Configuration for Graph RoPE attention
+#[napi(object)]
+pub struct RoPEConfig {
+    pub dim: u32,
+    pub max_position: u32,
+    pub base: Option<f64>,
+    pub scaling_factor: Option<f64>,
+}
+
+/// Graph RoPE attention (Rotary Position Embeddings for graphs)
+#[napi]
+pub struct GraphRoPEAttention {
+    inner: RustGraphRoPE,
+    config: RoPEConfig,
+}
+
+#[napi]
+impl GraphRoPEAttention {
+    /// Create a new Graph RoPE attention instance
+    ///
+    /// # Arguments
+    /// * `config` - RoPE configuration
+    #[napi(constructor)]
+    pub fn new(config: RoPEConfig) -> Self {
+        let rust_config = RustRoPEConfig {
+            dim: config.dim as usize,
+            max_position: config.max_position as usize,
+            base: config.base.unwrap_or(10000.0) as f32,
+            scaling_factor: config.scaling_factor.unwrap_or(1.0) as f32,
+        };
+        Self {
+            inner: RustGraphRoPE::new(rust_config),
+            config,
+        }
+    }
+
+    /// Create with simple parameters
+    #[napi(factory)]
+    pub fn simple(dim: u32, max_position: u32) -> Self {
+        Self::new(RoPEConfig {
+            dim,
+            max_position,
+            base: Some(10000.0),
+            scaling_factor: Some(1.0),
+        })
+    }
+
+    /// Compute attention without positional encoding
+    #[napi]
+    pub fn compute(
+        &self,
+        query: Float32Array,
+        keys: Vec<Float32Array>,
+        values: Vec<Float32Array>,
+    ) -> Result<Float32Array> {
+        let query_slice = query.as_ref();
+        let keys_vec: Vec<Vec<f32>> = keys.into_iter().map(|k| k.to_vec()).collect();
+        let values_vec: Vec<Vec<f32>> = values.into_iter().map(|v| v.to_vec()).collect();
+        let keys_refs: Vec<&[f32]> = keys_vec.iter().map(|k| k.as_slice()).collect();
+        let values_refs: Vec<&[f32]> = values_vec.iter().map(|v| v.as_slice()).collect();
+
+        let result = self
+            .inner
+            .compute(query_slice, &keys_refs, &values_refs)
+            .map_err(|e| Error::from_reason(e.to_string()))?;
+
+        Ok(Float32Array::new(result))
+    }
+
+    /// Compute attention with graph positions
+    ///
+    /// # Arguments
+    /// * `query` - Query vector
+    /// * `keys` - Array of key vectors
+    /// * `values` - Array of value vectors
+    /// * `query_position` - Position of query node
+    /// * `key_positions` - Positions of key nodes (e.g., hop distances)
+    #[napi]
+    pub fn compute_with_positions(
+        &self,
+        query: Float32Array,
+        keys: Vec<Float32Array>,
+        values: Vec<Float32Array>,
+        query_position: u32,
+        key_positions: Vec<u32>,
+    ) -> Result<Float32Array> {
+        let query_slice = query.as_ref();
+        let keys_vec: Vec<Vec<f32>> = keys.into_iter().map(|k| k.to_vec()).collect();
+        let values_vec: Vec<Vec<f32>> = values.into_iter().map(|v| v.to_vec()).collect();
+        let keys_refs: Vec<&[f32]> = keys_vec.iter().map(|k| k.as_slice()).collect();
+        let values_refs: Vec<&[f32]> = values_vec.iter().map(|v| v.as_slice()).collect();
+        let positions_usize: Vec<usize> = key_positions.into_iter().map(|p| p as usize).collect();
+
+        let result = self
+            .inner
+            .compute_with_positions(
+                query_slice,
+                &keys_refs,
+                &values_refs,
+                query_position as usize,
+                &positions_usize,
+            )
+            .map_err(|e| Error::from_reason(e.to_string()))?;
+
+        Ok(Float32Array::new(result))
+    }
+
+    /// Apply rotary embedding to a vector
+    #[napi]
+    pub fn apply_rotary(&self, vector: Float32Array, position: u32) -> Float32Array {
+        let v = vector.as_ref();
+        let result = self.inner.apply_rotary(v, position as usize);
+        Float32Array::new(result)
+    }
+
+    /// Convert graph distance to position bucket
+    #[napi]
+    pub fn distance_to_position(distance: u32, max_distance: u32) -> u32 {
+        RustGraphRoPE::distance_to_position(distance as usize, max_distance as usize) as u32
+    }
+
+    /// Get the dimension
+    #[napi(getter)]
+    pub fn dim(&self) -> u32 {
+        self.config.dim
+    }
+
+    /// Get the max position
+    #[napi(getter)]
+    pub fn max_position(&self) -> u32 {
+        self.config.max_position
+    }
+}
+
+// ============================================================================
+// Dual-Space Attention
+// ============================================================================
+
+/// Configuration for dual-space attention
+#[napi(object)]
+pub struct DualSpaceConfig {
+    pub dim: u32,
+    pub curvature: f64,
+    pub euclidean_weight: f64,
+    pub hyperbolic_weight: f64,
+    pub temperature: Option<f64>,
+}
+
+/// Dual-space attention (Euclidean + Hyperbolic)
+#[napi]
+pub struct DualSpaceAttention {
+    inner: RustDualSpace,
+    config: DualSpaceConfig,
+}
+
+#[napi]
+impl DualSpaceAttention {
+    /// Create a new dual-space attention instance
+    ///
+    /// # Arguments
+    /// * `config` - Dual-space configuration
+    #[napi(constructor)]
+    pub fn new(config: DualSpaceConfig) -> Self {
+        let rust_config = RustDualConfig {
+            dim: config.dim as usize,
+            curvature: config.curvature as f32,
+            euclidean_weight: config.euclidean_weight as f32,
+            hyperbolic_weight: config.hyperbolic_weight as f32,
+            learn_weights: false,
+            temperature: config.temperature.unwrap_or(1.0) as f32,
+        };
+        Self {
+            inner: RustDualSpace::new(rust_config),
+            config,
+        }
+    }
+
+    /// Create with simple parameters (equal weights)
+    #[napi(factory)]
+    pub fn simple(dim: u32, curvature: f64) -> Self {
+        Self::new(DualSpaceConfig {
+            dim,
+            curvature,
+            euclidean_weight: 0.5,
+            hyperbolic_weight: 0.5,
+            temperature: Some(1.0),
+        })
+    }
+
+    /// Create with custom weights
+    #[napi(factory)]
+    pub fn with_weights(
+        dim: u32,
+        curvature: f64,
+        euclidean_weight: f64,
+        hyperbolic_weight: f64,
+    ) -> Self {
+        Self::new(DualSpaceConfig {
+            dim,
+            curvature,
+            euclidean_weight,
+            hyperbolic_weight,
+            temperature: Some(1.0),
+        })
+    }
+
+    /// Compute dual-space attention
+    #[napi]
+    pub fn compute(
+        &self,
+        query: Float32Array,
+        keys: Vec<Float32Array>,
+        values: Vec<Float32Array>,
+    ) -> Result<Float32Array> {
+        let query_slice = query.as_ref();
+        let keys_vec: Vec<Vec<f32>> = keys.into_iter().map(|k| k.to_vec()).collect();
+        let values_vec: Vec<Vec<f32>> = values.into_iter().map(|v| v.to_vec()).collect();
+        let keys_refs: Vec<&[f32]> = keys_vec.iter().map(|k| k.as_slice()).collect();
+        let values_refs: Vec<&[f32]> = values_vec.iter().map(|v| v.as_slice()).collect();
+
+        let result = self
+            .inner
+            .compute(query_slice, &keys_refs, &values_refs)
+            .map_err(|e| Error::from_reason(e.to_string()))?;
+
+        Ok(Float32Array::new(result))
+    }
+
+    /// Get space contributions (Euclidean and Hyperbolic scores separately)
+    #[napi]
+    pub fn get_space_contributions(
+        &self,
+        query: Float32Array,
+        keys: Vec<Float32Array>,
+    ) -> SpaceContributions {
+        let query_slice = query.as_ref();
+        let keys_vec: Vec<Vec<f32>> = keys.into_iter().map(|k| k.to_vec()).collect();
+        let keys_refs: Vec<&[f32]> = keys_vec.iter().map(|k| k.as_slice()).collect();
+
+        let (euc_scores, hyp_scores) = self.inner.get_space_contributions(query_slice, &keys_refs);
+
+        SpaceContributions {
+            euclidean_scores: Float32Array::new(euc_scores),
+            hyperbolic_scores: Float32Array::new(hyp_scores),
+        }
+    }
+
+    /// Get the dimension
+    #[napi(getter)]
+    pub fn dim(&self) -> u32 {
+        self.config.dim
+    }
+
+    /// Get the curvature
+    #[napi(getter)]
+    pub fn curvature(&self) -> f64 {
+        self.config.curvature
+    }
+
+    /// Get the Euclidean weight
+    #[napi(getter)]
+    pub fn euclidean_weight(&self) -> f64 {
+        self.config.euclidean_weight
+    }
+
+    /// Get the Hyperbolic weight
+    #[napi(getter)]
+    pub fn hyperbolic_weight(&self) -> f64 {
+        self.config.hyperbolic_weight
+    }
+}
+
+/// Space contribution scores
+#[napi(object)]
+pub struct SpaceContributions {
+    pub euclidean_scores: Float32Array,
+    pub hyperbolic_scores: Float32Array,
+}