use crate::algorithms::AnalyzeAlgorithm;
pub const MAX_EMBEDDING_DIMENSIONS: usize = 4_096;
pub const MAX_HASHGNN_DIMENSIONS: usize = 8_192;
#[derive(Debug, Clone, PartialEq)]
pub enum EmbeddingOptions {
Node2Vec(Node2VecOptions),
GraphSage(GraphSageOptions),
FastRandomProjection(FastRpOptions),
HashGnn(HashGnnOptions),
}
#[derive(Debug, Clone, PartialEq)]
pub struct EmbeddingAnalyzeOptions {
pub by: AnalyzeAlgorithm,
pub via: Option<String>,
pub directed: bool,
pub weight: Option<String>,
pub options: EmbeddingOptions,
}
#[derive(Debug, Clone, PartialEq)]
pub struct Node2VecOptions {
pub dimensions: usize,
pub walk_length: usize,
pub walks_per_node: usize,
pub p: f64,
pub q: f64,
pub window_size: usize,
pub negative_samples: usize,
pub epochs: usize,
pub learning_rate: f64,
pub seed: u64,
}
impl Default for Node2VecOptions {
fn default() -> Self {
Self {
dimensions: 128,
walk_length: 80,
walks_per_node: 10,
p: 1.0,
q: 1.0,
window_size: 10,
negative_samples: 5,
epochs: 1,
learning_rate: 0.025,
seed: 0,
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum GraphSageAggregator {
#[default]
Mean,
}
#[derive(Debug, Clone, PartialEq)]
pub struct GraphSageOptions {
pub dimensions: usize,
pub hidden_dimensions: usize,
pub layers: usize,
pub sample_sizes: Vec<usize>,
pub aggregator: GraphSageAggregator,
pub epochs: usize,
pub negative_samples: usize,
pub learning_rate: f64,
pub feature_properties: Vec<String>,
pub seed: u64,
}
impl Default for GraphSageOptions {
fn default() -> Self {
Self {
dimensions: 256,
hidden_dimensions: 256,
layers: 2,
sample_sizes: vec![25, 10],
aggregator: GraphSageAggregator::Mean,
epochs: 1,
negative_samples: 20,
learning_rate: 0.000_002,
feature_properties: Vec::new(),
seed: 0,
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct FastRpOptions {
pub dimensions: usize,
pub iteration_weights: Vec<f64>,
pub normalization_strength: f64,
pub feature_weight: f64,
pub feature_properties: Vec<String>,
pub seed: u64,
}
impl Default for FastRpOptions {
fn default() -> Self {
Self {
dimensions: 128,
iteration_weights: vec![0.0, 1.0, 1.0],
normalization_strength: 0.0,
feature_weight: 0.0,
feature_properties: Vec::new(),
seed: 0,
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct HashGnnOptions {
pub dimensions: usize,
pub iterations: usize,
pub embedding_density: f64,
pub heterogeneous: bool,
pub node_type_property: Option<String>,
pub relationship_type_property: Option<String>,
pub seed: u64,
}
impl Default for HashGnnOptions {
fn default() -> Self {
Self {
dimensions: 256,
iterations: 2,
embedding_density: 0.25,
heterogeneous: false,
node_type_property: None,
relationship_type_property: None,
seed: 0,
}
}
}