node2vec_rs/

graph.rs

use indicatif::{ProgressBar, ProgressStyle};
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use rayon::prelude::*;
use rustc_hash::{FxHashMap, FxHashSet};

/////////////
// Helpers //
/////////////

/// Compute the transition probabilities
///
/// ### Params
///
/// * `adjacency` - The adjacency of the graph stored as a HashMap
/// * `p` - p parameter in node2vec that controls probability to return to
///   origin node.
/// * `q` - q parameter in node2vec that controls probability to reach out
///   futher in the graph.
///
/// ### Returns
///
/// The transition probabilities as an `FxHashMap`.
pub fn compute_transition_prob(
    adjacency: &FxHashMap<u32, Vec<(u32, f32)>>,
    p: f32,
    q: f32,
) -> FxHashMap<(u32, u32), Vec<(u32, f32)>> {
    let neighbours_set: FxHashMap<u32, FxHashSet<u32>> = adjacency
        .iter()
        .map(|(node, edges)| (*node, edges.iter().map(|(n, _)| *n).collect()))
        .collect();

    adjacency
        .par_iter()
        .flat_map(|(curr_node, curr_edges)| {
            curr_edges.par_iter().filter_map(|(prev_node, _)| {
                adjacency.get(curr_node).map(|next_edges| {
                    let mut probs = Vec::new();
                    let mut total = 0_f32;

                    for (next_node, weight) in next_edges.iter() {
                        let unnorm_prob = if next_node == prev_node {
                            weight / p
                        } else if neighbours_set
                            .get(prev_node)
                            .map(|s| s.contains(next_node))
                            .unwrap_or(false)
                        {
                            *weight
                        } else {
                            weight / q
                        };
                        total += unnorm_prob;
                        probs.push((*next_node, unnorm_prob));
                    }

                    let mut cumulative = 0.0;
                    let normalised: Vec<(u32, f32)> = probs
                        .into_iter()
                        .map(|(node, prob)| {
                            cumulative += prob / total;
                            (node, cumulative)
                        })
                        .collect();

                    ((*prev_node, *curr_node), normalised)
                })
            })
        })
        .collect()
}

/////////////////////
// Graph structure //
/////////////////////

/// Structure to store the Node2Vec graph
///
/// ### Fields
///
/// * `adjacency` - The adjacency stored as an FxHashMap.
/// * `transition_probs` - The transition probabilities stored in a FxHashMap.
#[derive(Debug, Clone)]
pub struct Node2VecGraph {
    pub adjacency: FxHashMap<u32, Vec<(u32, f32)>>,
    pub transition_probs: FxHashMap<(u32, u32), Vec<(u32, f32)>>,
}

impl Node2VecGraph {
    /// Generates random walks from the graph
    ///
    /// ### Params
    /// * `walks_per_node` - Number of walks to generate starting from each node
    /// * `walk_length` - Length of each walk
    /// * `seed` - Random seed for reproducibility
    ///
    /// ### Returns
    ///
    /// Vector of walks, where each walk is a sequence of node IDs
    pub fn generate_walks(
        &self,
        walks_per_node: usize,
        walk_length: usize,
        seed: usize,
    ) -> Vec<Vec<u32>> {
        let total_walks = self.adjacency.len() * walks_per_node;
        let progress = ProgressBar::new(total_walks as u64);

        progress.set_style(
            ProgressStyle::default_bar()
                .template("[Random walk gen: {elapsed_precise}] {bar:40.cyan/blue} {pos}/{len}")
                .unwrap()
                .progress_chars("#>-"),
        );

        let walks = self
            .adjacency
            .par_iter()
            .flat_map(|(start_node, _)| {
                let progress = progress.clone();
                (0..walks_per_node).into_par_iter().map(move |walk_idx| {
                    let walk_seed = seed
                        .wrapping_mul(*start_node as usize)
                        .wrapping_add(walk_idx);
                    let mut rng = StdRng::seed_from_u64(walk_seed as u64);
                    let walk = self.single_walk(*start_node, walk_length, &mut rng);
                    progress.inc(1);
                    walk
                })
            })
            .collect();

        progress.finish();
        walks
    }

    /// Performs a single biased random walk
    ///
    /// ### Params
    ///
    /// * `start_node` - Node to start the walk from
    /// * `walk_length` - Maximum length of the walk
    /// * `rng` - Random number generator
    ///
    /// ### Returns
    ///
    /// The vector of node IDs for this walk
    fn single_walk(&self, start_node: u32, walk_length: usize, rng: &mut StdRng) -> Vec<u32> {
        let mut walk: Vec<u32> = Vec::with_capacity(walk_length);
        walk.push(start_node);

        if walk_length == 1 {
            return walk;
        }

        let mut curr = if let Some(neighbours) = self.adjacency.get(&start_node) {
            self.sample_neighbor(neighbours, rng)
        } else {
            return walk;
        };

        walk.push(curr);

        for _ in 2..walk_length {
            let prev = walk[walk.len() - 2];

            if let Some(probs) = self.transition_probs.get(&(prev, curr)) {
                curr = self.sample_from_cumulative(probs, rng);
                walk.push(curr);
            } else {
                break;
            }
        }

        walk
    }

    /// Samples a neighbour based on edge weights
    ///
    /// ### Params
    ///
    /// * `neighbours` - Slice of (node, weight) tuples
    /// * `rng` - Random number generator
    ///
    /// ### Returns
    ///
    /// Node ID based on the neighbours
    fn sample_neighbor(&self, neighbours: &[(u32, f32)], rng: &mut impl Rng) -> u32 {
        let total: f32 = neighbours.iter().map(|(_, w)| w).sum();
        let mut rand_val = rng.random::<f32>() * total;

        for (node, weight) in neighbours {
            rand_val -= weight;
            if rand_val <= 0.0 {
                return *node;
            }
        }

        neighbours[0].0
    }

    /// Samples from a cumulative probability distribution
    ///
    /// ### Params
    ///
    /// * `cumulative` - Cumulative probabilities as (node, cumulative_prob)
    ///   pairs
    /// * `rng` - Random number generator
    ///
    /// ### Returns
    ///
    /// The node ID
    fn sample_from_cumulative(&self, cumulative: &[(u32, f32)], rng: &mut impl Rng) -> u32 {
        let rand_val = rng.random::<f32>();

        match cumulative.binary_search_by(|(_, cum_prob)| cum_prob.partial_cmp(&rand_val).unwrap())
        {
            Ok(idx) => cumulative[idx].0,
            Err(idx) => {
                if idx < cumulative.len() {
                    cumulative[idx].0
                } else {
                    cumulative[cumulative.len() - 1].0
                }
            }
        }
    }
}

#[cfg(test)]
mod graph_tests {
    use crate::prelude::*;
    use rustc_hash::FxHashMap;

    #[test]
    fn test_transition_probs_sum_to_one() {
        let mut adjacency = FxHashMap::default();
        adjacency.insert(1, vec![(2, 1.0), (3, 1.0)]);
        adjacency.insert(2, vec![(1, 1.0), (3, 1.0)]);
        adjacency.insert(3, vec![(1, 1.0), (2, 1.0)]);

        let probs = compute_transition_prob(&adjacency, 1.0, 1.0);

        // Check all cumulative probabilities end at ~1.0
        for (_, cumulative_probs) in probs.iter() {
            let last_prob = cumulative_probs.last().unwrap().1;
            assert!((last_prob - 1.0).abs() < 1e-6);
        }
    }

    #[test]
    fn test_transition_probs_p_parameter() {
        // Simple chain: 1 - 2 - 3
        let mut adjacency = FxHashMap::default();
        adjacency.insert(1, vec![(2, 1.0)]);
        adjacency.insert(2, vec![(1, 1.0), (3, 1.0)]);
        adjacency.insert(3, vec![(2, 1.0)]);

        // With p=0.5, returning to origin should be more likely
        let probs_low_p = compute_transition_prob(&adjacency, 0.5, 1.0);
        // With p=2.0, returning to origin should be less likely
        let probs_high_p = compute_transition_prob(&adjacency, 2.0, 1.0);

        // When at node 2 coming from node 1, check probability of returning to 1
        let low_p_return = &probs_low_p[&(1, 2)];
        let high_p_return = &probs_high_p[&(1, 2)];

        // Find the cumulative probability for node 1 (return)
        let low_p_val = low_p_return.iter().find(|(n, _)| *n == 1).unwrap().1;
        let high_p_val = high_p_return.iter().find(|(n, _)| *n == 1).unwrap().1;

        // Lower p should give higher probability of return
        assert!(low_p_val > high_p_val);
    }

    #[test]
    fn test_transition_probs_q_parameter() {
        // Triangle: 1 - 2 - 3, with 1 - 3 edge
        let mut adjacency = FxHashMap::default();
        adjacency.insert(1, vec![(2, 1.0), (3, 1.0)]);
        adjacency.insert(2, vec![(1, 1.0), (3, 1.0)]);
        adjacency.insert(3, vec![(1, 1.0), (2, 1.0)]);

        // With q=0.5, exploring further should be more likely
        let probs_low_q = compute_transition_prob(&adjacency, 1.0, 0.5);
        // With q=2.0, exploring further should be less likely
        let probs_high_q = compute_transition_prob(&adjacency, 1.0, 2.0);

        // Both should normalise to 1.0
        for probs in [&probs_low_q, &probs_high_q] {
            for cumulative_probs in probs.values() {
                let last = cumulative_probs.last().unwrap().1;
                assert!((last - 1.0).abs() < 1e-6);
            }
        }
    }

    #[test]
    fn test_transition_probs_uniform_with_p_q_one() {
        let mut adjacency = FxHashMap::default();
        adjacency.insert(1, vec![(2, 1.0), (3, 1.0)]);
        adjacency.insert(2, vec![(1, 1.0), (3, 1.0)]);
        adjacency.insert(3, vec![(1, 1.0), (2, 1.0)]);

        let probs = compute_transition_prob(&adjacency, 1.0, 1.0);

        // With p=q=1 and equal weights, all transitions should be uniform
        for cumulative_probs in probs.values() {
            if cumulative_probs.len() == 2 {
                // Each neighbour should have ~0.5 probability
                let first_prob = cumulative_probs[0].1;
                assert!((first_prob - 0.5).abs() < 1e-6);
            }
        }
    }

    #[test]
    fn test_isolated_node() {
        let mut adjacency = FxHashMap::default();
        adjacency.insert(1, vec![]);

        let probs = compute_transition_prob(&adjacency, 1.0, 1.0);

        // Isolated node should have no transition probabilities
        assert!(probs.is_empty());
    }

    #[test]
    fn test_weighted_transitions() {
        let mut adjacency = FxHashMap::default();
        adjacency.insert(1, vec![(2, 3.0)]);
        adjacency.insert(2, vec![(1, 3.0), (3, 1.0)]);
        adjacency.insert(3, vec![(2, 1.0)]);

        let probs = compute_transition_prob(&adjacency, 1.0, 1.0);

        // At node 2 coming from 1, going back to 1 should be more likely (weight 3.0 vs 1.0)
        let probs_at_2 = &probs[&(1, 2)];
        let prob_to_1 = probs_at_2.iter().find(|(n, _)| *n == 1).unwrap().1;
        let prob_to_3 = probs_at_2.iter().find(|(n, _)| *n == 3).unwrap().1;

        assert!(prob_to_1 < prob_to_3); // cumulative, so first is less than second
        assert!((prob_to_3 - 1.0).abs() < 1e-6); // last should sum to 1
    }
}