use std::collections::{HashMap, HashSet};
pub type GraphEdges = HashMap<String, Vec<String>>;
const CONVERGENCE_EPSILON: f64 = 1e-8;
fn l1_delta(prev: &HashMap<String, f64>, next: &HashMap<String, f64>) -> f64 {
let mut delta = 0.0;
for (node, score) in next {
delta += (score - prev.get(node).copied().unwrap_or(0.0)).abs();
}
for (node, score) in prev {
if !next.contains_key(node) {
delta += score.abs();
}
}
delta
}
pub fn personalized_pagerank(
graph: &GraphEdges,
seeds: &[String],
damping: f32,
iterations: usize,
) -> Vec<(String, f64)> {
personalized_pagerank_holdout(graph, seeds, damping, iterations, None)
}
pub fn personalized_pagerank_holdout(
graph: &GraphEdges,
seeds: &[String],
damping: f32,
iterations: usize,
holdout: Option<&str>,
) -> Vec<(String, f64)> {
if seeds.is_empty() || graph.is_empty() {
return Vec::new();
}
let masked = |node: &str| holdout == Some(node);
let mut inlink_counts: HashMap<&str, usize> = HashMap::new();
for (node, targets) in graph {
if masked(node) {
continue;
}
for t in targets {
*inlink_counts.entry(t.as_str()).or_default() += 1;
}
}
let mut seed_weights: Vec<(&str, f64)> = seeds
.iter()
.filter(|s| graph.contains_key(s.as_str()) && !masked(s.as_str()))
.map(|s| {
let inlinks = inlink_counts.get(s.as_str()).copied().unwrap_or(0);
(s.as_str(), 1.0 / (inlinks as f64 + 1.0))
})
.collect();
let total_weight: f64 = seed_weights.iter().map(|(_, w)| w).sum();
if total_weight == 0.0 {
return Vec::new();
}
for (_, w) in &mut seed_weights {
*w /= total_weight;
}
let d = damping as f64;
let mut scores: HashMap<String, f64> = HashMap::new();
let seed_set: HashSet<&str> = seeds.iter().map(|s| s.as_str()).collect();
for &(s, w) in &seed_weights {
scores.insert(s.to_string(), w);
}
for _ in 0..iterations {
let mut new_scores: HashMap<String, f64> = HashMap::new();
for &(s, w) in &seed_weights {
*new_scores.entry(s.to_string()).or_default() += (1.0 - d) * w;
}
for (node, score) in &scores {
if masked(node) {
continue;
}
if let Some(neighbors) = graph.get(node) {
let share = d * score / neighbors.len() as f64;
for neighbor in neighbors {
*new_scores.entry(neighbor.clone()).or_default() += share;
}
}
}
let converged = l1_delta(&scores, &new_scores) < CONVERGENCE_EPSILON;
scores = new_scores;
if converged {
break;
}
}
let mut results: Vec<(String, f64)> = scores
.into_iter()
.filter(|(path, score)| *score > 1e-6 && !seed_set.contains(path.as_str()) && !masked(path))
.collect();
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
results.truncate(crate::config::TOP_K);
results
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_ppr_empty_graph() {
let graph: GraphEdges = HashMap::new();
let results = personalized_pagerank(&graph, &["a.md".to_string()], 0.5, 20);
assert!(results.is_empty());
}
#[test]
fn test_ppr_empty_seeds() {
let mut graph: GraphEdges = HashMap::new();
graph.insert("a".into(), vec!["b".into()]);
graph.insert("b".into(), vec!["a".into()]);
let results = personalized_pagerank(&graph, &[], 0.5, 20);
assert!(results.is_empty());
}
#[test]
fn test_ppr_holdout_masks_edges_and_filters_node() {
let mut graph: GraphEdges = HashMap::new();
graph.insert("a".into(), vec!["s".into(), "x".into()]);
graph.insert("s".into(), vec!["a".into(), "g".into()]);
graph.insert("g".into(), vec!["s".into()]);
graph.insert("x".into(), vec!["a".into()]);
let seeds = vec!["a".to_string()];
let unmasked = personalized_pagerank(&graph, &seeds, 0.5, 20);
let unmasked_paths: Vec<&str> = unmasked.iter().map(|r| r.0.as_str()).collect();
assert!(unmasked_paths.contains(&"s") && unmasked_paths.contains(&"g"));
let held = personalized_pagerank_holdout(&graph, &seeds, 0.5, 20, Some("s"));
let held_paths: Vec<&str> = held.iter().map(|r| r.0.as_str()).collect();
assert!(
!held_paths.contains(&"s"),
"held-out node must not be returned: {held_paths:?}"
);
assert!(
!held_paths.contains(&"g"),
"held-out node's edges must not propagate: {held_paths:?}"
);
assert!(
held_paths.contains(&"x"),
"unrelated structure must survive the holdout: {held_paths:?}"
);
}
fn dense_graph(n: usize) -> GraphEdges {
let mut graph: GraphEdges = HashMap::new();
for i in 0..n {
let targets: Vec<String> = (0..n)
.filter(|j| *j != i)
.map(|j| format!("n{j}"))
.collect();
graph.insert(format!("n{i}"), targets);
}
graph
}
fn score_map(results: &[(String, f64)]) -> HashMap<&str, f64> {
results.iter().map(|(p, s)| (p.as_str(), *s)).collect()
}
fn assert_same_scores(capped: &[(String, f64)], uncapped: &[(String, f64)], case: &str) {
let a = score_map(capped);
let b = score_map(uncapped);
assert_eq!(
a.len(),
b.len(),
"{case}: early stop changed how many nodes survive the score floor"
);
for (path, score) in &a {
let other = b
.get(path)
.unwrap_or_else(|| panic!("{case}: {path} missing from the uncapped run"));
assert!(
(score - other).abs() < 1e-6,
"{case}: early stop changed {path} score: {score} vs {other}"
);
}
}
#[test]
fn test_ppr_early_stop_matches_full_iteration() {
let mut chain: GraphEdges = HashMap::new();
for i in 0..30 {
chain.insert(format!("n{i}"), vec![format!("n{}", (i + 1) % 30)]);
}
for (case, graph) in [("dense", dense_graph(12)), ("cycle", chain)] {
let seeds = vec!["n0".to_string()];
let capped = personalized_pagerank(&graph, &seeds, 0.5, 20);
let uncapped = personalized_pagerank(&graph, &seeds, 0.5, 400);
assert_same_scores(&capped, &uncapped, case);
}
}
#[test]
fn test_ppr_early_stop_preserves_a_real_ranking() {
let mut graph: GraphEdges = HashMap::new();
graph.insert("a".into(), vec!["b".into()]);
graph.insert("b".into(), vec!["c".into()]);
graph.insert("c".into(), vec!["d".into()]);
graph.insert("d".into(), vec!["e".into()]);
let seeds = vec!["a".to_string()];
let capped = personalized_pagerank(&graph, &seeds, 0.5, 20);
let uncapped = personalized_pagerank(&graph, &seeds, 0.5, 400);
let order: Vec<&str> = capped.iter().map(|r| r.0.as_str()).collect();
assert_eq!(
order,
vec!["b", "c", "d", "e"],
"hop order is score order on a chain"
);
assert_eq!(
order,
uncapped.iter().map(|r| r.0.as_str()).collect::<Vec<_>>(),
"early stop changed rank order on a graph with no ties"
);
assert_same_scores(&capped, &uncapped, "chain");
}
#[test]
fn test_ppr_holdout_early_stop_matches_full_iteration() {
let graph = dense_graph(10);
let seeds = vec!["n0".to_string()];
let capped = personalized_pagerank_holdout(&graph, &seeds, 0.5, 20, Some("n3"));
let uncapped = personalized_pagerank_holdout(&graph, &seeds, 0.5, 400, Some("n3"));
assert!(
!score_map(&capped).contains_key("n3"),
"held-out node must stay out of the result"
);
assert_same_scores(&capped, &uncapped, "holdout");
}
#[test]
fn test_ppr_iteration_cap_is_still_honoured() {
let mut graph: GraphEdges = HashMap::new();
graph.insert("a".into(), vec!["b".into()]);
graph.insert("b".into(), vec!["c".into()]);
graph.insert("c".into(), vec!["d".into()]);
let one = personalized_pagerank(&graph, &["a".to_string()], 0.5, 1);
let paths: Vec<&str> = one.iter().map(|r| r.0.as_str()).collect();
assert!(
paths.contains(&"b"),
"one step must reach the direct neighbour: {paths:?}"
);
assert!(
!paths.contains(&"c"),
"one step must not reach two hops out: {paths:?}"
);
}
#[test]
fn test_l1_delta_counts_appearing_and_vanishing_nodes() {
let mut prev: HashMap<String, f64> = HashMap::new();
prev.insert("a".into(), 0.5);
prev.insert("gone".into(), 0.25);
let mut next: HashMap<String, f64> = HashMap::new();
next.insert("a".into(), 0.5);
next.insert("new".into(), 0.125);
assert!((l1_delta(&prev, &next) - 0.375).abs() < 1e-12);
assert_eq!(l1_delta(&prev, &prev), 0.0);
}
#[test]
fn test_ppr_chain() {
let mut graph: GraphEdges = HashMap::new();
graph.insert("a".into(), vec!["b".into()]);
graph.insert("b".into(), vec!["a".into(), "c".into()]);
graph.insert("c".into(), vec!["b".into(), "d".into()]);
graph.insert("d".into(), vec!["c".into()]);
let results = personalized_pagerank(&graph, &["a".into()], 0.5, 20);
assert!(!results.is_empty());
let paths: Vec<&str> = results.iter().map(|r| r.0.as_str()).collect();
assert!(paths.contains(&"b"));
assert!(results[0].1 >= results.last().unwrap().1);
}
}