networkit_rs/
embedding.rs

1use 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}