networkit_rs/
embedding.rs1use cxx::UniquePtr;
2use miette::IntoDiagnostic;
3
4use crate::{
5 base::Algorithm,
6 bridge::{self, *},
7};
8
9pub struct Node2Vec {
10 inner: UniquePtr<bridge::Node2Vec>,
11 n_nodes: usize,
12 dim: usize,
13}
14
15#[derive(Clone)]
16pub struct NodeFeatures {
17 data: Vec<f64>,
18 n_nodes: usize,
19 dim: usize,
20}
21
22impl NodeFeatures {
23 pub fn iter<'a>(&'a self) -> impl Iterator<Item = &'a [f64]> {
24 (0..self.n_nodes).map(|idx| self.get(idx))
25 }
26 pub fn get(&self, idx: usize) -> &[f64] {
27 &self.data[idx * self.dim..idx * self.dim + self.dim]
28 }
29}
30
31impl Node2Vec {
32 pub fn new(
33 g: &crate::Graph,
34 p: Option<f64>,
35 q: Option<f64>,
36 l: Option<u64>,
37 n: Option<u64>,
38 d: Option<u64>,
39 ) -> Self {
40 let d = d.unwrap_or(128);
41 let n_nodes = g.number_of_nodes() as usize;
42 Self {
43 inner: NewNode2Vec(
44 g,
45 p.unwrap_or(1.),
46 q.unwrap_or(1.),
47 l.unwrap_or(80),
48 n.unwrap_or(10),
49 d,
50 ),
51 n_nodes,
52 dim: d as usize,
53 }
54 }
55 pub fn get_features(&self) -> NodeFeatures {
56 let mut data = Vec::with_capacity(self.n_nodes * self.dim);
57 Node2VecGetFeatures(&self.inner, &mut data);
58 NodeFeatures {
59 data,
60 n_nodes: self.n_nodes,
61 dim: self.dim,
62 }
63 }
64}
65
66impl Algorithm for Node2Vec {
67 fn run(&mut self) -> miette::Result<()> {
68 self.inner.pin_mut().run().into_diagnostic()
69 }
70
71 fn has_finished(&self) -> bool {
72 self.inner.hasFinished()
73 }
74}