use num_traits::{ToPrimitive, Zero};
use rand::prelude::*;
use rayon::prelude::*;
use std::cell::RefCell;
use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, Ordering::Relaxed};
use crate::graph::community_detections::remap_communities_by_size;
use crate::prelude::*;
const LEIDEN_PAR_MIN_NODES: usize = 4096;
const PAR_MOVE_MAX_ROUNDS: usize = 1000;
const PAR_MOVE_MIN_LEN: usize = 64;
const NEW_CLUSTER: u32 = u32::MAX;
const MAX_ITERATIONS: usize = 1000;
const QUALITY_CHUNK: usize = 4096;
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub enum LocalMoving {
#[default]
Queue,
Coloured,
}
pub fn parse_local_moving(s: &str) -> Option<LocalMoving> {
match s.to_lowercase().as_str() {
"queue" => Some(LocalMoving::Queue),
"coloured" | "colored" => Some(LocalMoving::Coloured),
_ => None,
}
}
#[derive(Clone, Copy, Debug)]
pub struct LeidenParams {
pub resolution: f64,
pub beta: f64,
pub n_iterations: Option<usize>,
pub min_gain: f64,
pub refine: bool,
pub seed: u64,
pub local_moving: LocalMoving,
}
impl LeidenParams {
pub fn new(
resolution: f64,
beta: f64,
n_iterations: Option<usize>,
refine: bool,
seed: u64,
) -> Self {
Self {
resolution,
beta,
n_iterations,
min_gain: 0.0,
refine,
seed,
local_moving: LocalMoving::Queue,
}
}
}
impl Default for LeidenParams {
fn default() -> Self {
Self {
resolution: 1.0,
beta: 0.01,
n_iterations: Some(5),
min_gain: 0.0,
refine: true,
seed: 42,
local_moving: LocalMoving::Queue,
}
}
}
trait LevelGraph: Sync {
type W: Copy + Zero + Send + Sync + ToPrimitive;
fn n(&self) -> usize;
fn row(&self, i: usize) -> (&[u32], &[Self::W]);
}
impl<T> LevelGraph for SparseGraph<T>
where
T: BixverseFloat + Send + Sync + std::iter::Sum,
{
type W = T;
#[inline]
fn n(&self) -> usize {
self.get_node_number()
}
#[inline]
fn row(&self, i: usize) -> (&[u32], &[T]) {
self.get_neighbours_raw(i)
}
}
struct AggGraph<W> {
indptr: Vec<usize>,
indices: Vec<u32>,
weights: Vec<W>,
}
impl<W: Copy + Zero + Send + Sync + ToPrimitive> LevelGraph for AggGraph<W> {
type W = W;
#[inline]
fn n(&self) -> usize {
self.indptr.len() - 1
}
#[inline]
fn row(&self, i: usize) -> (&[u32], &[W]) {
let (s, e) = (self.indptr[i], self.indptr[i + 1]);
(&self.indices[s..e], &self.weights[s..e])
}
}
#[inline(always)]
fn wf<W: ToPrimitive>(w: W) -> f64 {
w.to_f64().expect("float to f64 cannot fail")
}
#[inline]
fn splitmix64(mut x: u64) -> u64 {
x = x.wrapping_add(0x9E37_79B9_7F4A_7C15);
x = (x ^ (x >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
x = (x ^ (x >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
x ^ (x >> 31)
}
fn derive_seed(seed: u64, parts: &[u64]) -> u64 {
parts
.iter()
.fold(splitmix64(seed), |acc, &p| splitmix64(acc ^ p))
}
fn reindex(labels: &mut [u32]) -> usize {
let mut map = vec![u32::MAX; labels.len()];
let mut next = 0u32;
for l in labels.iter_mut() {
let m = &mut map[*l as usize];
if *m == u32::MAX {
*m = next;
next += 1;
}
*l = *m;
}
next as usize
}
fn group_by_label(labels: &[u32], k: usize) -> (Vec<usize>, Vec<u32>) {
let mut offsets = vec![0usize; k + 1];
for &l in labels {
offsets[l as usize + 1] += 1;
}
for c in 0..k {
offsets[c + 1] += offsets[c];
}
let mut cursor = offsets.clone();
let mut members = vec![0u32; labels.len()];
for (v, &l) in labels.iter().enumerate() {
members[cursor[l as usize]] = v as u32;
cursor[l as usize] += 1;
}
(offsets, members)
}
fn level_quality<G: LevelGraph>(
g: &G,
node_w: &[f64],
membership: &[u32],
res: f64,
two_m: f64,
) -> f64 {
let n = g.n();
let partials: Vec<f64> = (0..n.div_ceil(QUALITY_CHUNK))
.into_par_iter()
.map(|chunk| {
let mut s = 0.0;
for v in chunk * QUALITY_CHUNK..((chunk + 1) * QUALITY_CHUNK).min(n) {
let c = membership[v];
let (nbrs, ws) = g.row(v);
for (&u, &w) in nbrs.iter().zip(ws) {
if membership[u as usize] == c {
s += wf(w);
}
}
}
s
})
.collect();
let internal: f64 = partials.iter().sum();
let mut tot = vec![0.0f64; n];
for (&c, &k) in membership.iter().zip(node_w) {
tot[c as usize] += k;
}
let squares: f64 = tot.iter().map(|t| t * t).sum();
(internal - res * squares) / two_m
}
pub fn modularity<T>(graph: &SparseGraph<T>, labels: &[usize], resolution: f64) -> f64
where
T: BixverseFloat + Send + Sync + std::iter::Sum,
{
let n = graph.get_node_number();
let k = labels.iter().copied().max().map_or(0, |m| m + 1);
let mut internal = 0.0f64;
let mut two_m = 0.0f64;
let mut cluster_w = vec![0.0f64; k];
for v in 0..n {
let (nbrs, ws) = graph.get_neighbours_raw(v);
let mut deg = 0.0;
for (&u, &w) in nbrs.iter().zip(ws) {
let w = wf(w);
deg += w;
if labels[u as usize] == labels[v] {
internal += w;
}
}
two_m += deg;
cluster_w[labels[v]] += deg;
}
if two_m == 0.0 {
return 0.0;
}
let penalty: f64 = cluster_w.iter().map(|&c| c * c).sum::<f64>() * resolution / two_m;
(internal - penalty) / two_m
}
fn move_nodes<G: LevelGraph>(
g: &G,
node_w: &[f64],
membership: &mut [u32],
res: f64,
max_evals: usize,
shuffle: bool,
rng: &mut StdRng,
) -> (usize, bool) {
let n = g.n();
let queue_order: Vec<u32> = if shuffle {
let mut order: Vec<u32> = (0..n as u32).collect();
order.shuffle(rng);
order
} else {
let start = if n > 0 { rng.random_range(0..n) } else { 0 };
(start..n).chain(0..start).map(|v| v as u32).collect()
};
let mut changed = false;
let mut queue: VecDeque<u32> = queue_order.into();
let mut stable = vec![false; n];
let mut cluster_w = vec![0.0f64; n];
let mut cluster_n = vec![0u32; n];
for v in 0..n {
let c = membership[v] as usize;
cluster_w[c] += node_w[v];
cluster_n[c] += 1;
}
let mut empty: Vec<u32> = (0..n as u32)
.filter(|&c| cluster_n[c as usize] == 0)
.collect();
let mut w_to = vec![0.0f64; n];
let mut added = vec![false; n];
let mut nbr_clusters: Vec<u32> = Vec::with_capacity(64);
let mut evals = 0usize;
while let Some(v) = queue.pop_front() {
if evals >= max_evals {
break;
}
evals += 1;
let v = v as usize;
let cur = membership[v];
let kv = node_w[v];
cluster_w[cur as usize] -= kv;
cluster_n[cur as usize] -= 1;
if cluster_n[cur as usize] == 0 {
empty.push(cur);
}
let e = *empty
.last()
.expect("node just left a cluster, so one is free or v is alone");
nbr_clusters.push(e);
added[e as usize] = true;
let (nbrs, ws) = g.row(v);
for (&u, &w) in nbrs.iter().zip(ws) {
if u as usize == v {
continue;
}
let c = membership[u as usize];
if !added[c as usize] {
added[c as usize] = true;
nbr_clusters.push(c);
}
w_to[c as usize] += wf(w);
}
let mut best = cur;
let mut best_gain = w_to[cur as usize] - kv * cluster_w[cur as usize] * res;
for &c in &nbr_clusters {
let gain = w_to[c as usize] - kv * cluster_w[c as usize] * res;
if gain > best_gain {
best = c;
best_gain = gain;
}
w_to[c as usize] = 0.0;
added[c as usize] = false;
}
nbr_clusters.clear();
cluster_w[best as usize] += kv;
cluster_n[best as usize] += 1;
if empty.last() == Some(&best) {
empty.pop();
}
stable[v] = true;
if best != cur {
changed = true;
membership[v] = best;
for &u in nbrs {
let u = u as usize;
if stable[u] && membership[u] != best {
stable[u] = false;
queue.push_back(u as u32);
}
}
}
}
(reindex(membership), changed)
}
#[derive(Default)]
struct ClusterTable {
keys: Vec<u32>,
vals: Vec<f64>,
used: Vec<u32>,
}
impl ClusterTable {
#[inline]
fn reserve(&mut self, deg: usize) {
let cap = (2 * deg).next_power_of_two().max(16);
if self.keys.len() < cap {
self.keys = vec![u32::MAX; cap];
self.vals = vec![0.0; cap];
}
}
#[inline]
fn slot(&self, c: u32) -> usize {
let mask = self.keys.len() - 1;
let mut i = (c.wrapping_mul(0x9E37_79B1) as usize) & mask;
while self.keys[i] != c && self.keys[i] != u32::MAX {
i = (i + 1) & mask;
}
i
}
#[inline]
fn add(&mut self, c: u32, w: f64) {
let i = self.slot(c);
if self.keys[i] == u32::MAX {
self.keys[i] = c;
self.vals[i] = w;
self.used.push(i as u32);
} else {
self.vals[i] += w;
}
}
#[inline]
fn get(&self, c: u32) -> f64 {
let i = self.slot(c);
if self.keys[i] == c { self.vals[i] } else { 0.0 }
}
#[inline]
fn clear(&mut self) {
for &i in &self.used {
self.keys[i as usize] = u32::MAX;
}
self.used.clear();
}
}
thread_local! {
static MOVE_SCRATCH: RefCell<ClusterTable> = RefCell::new(ClusterTable::default());
}
#[derive(Clone, Copy, Debug)]
struct Move {
v: u32,
to: u32,
w_to: f64,
w_from: f64,
}
#[allow(clippy::too_many_arguments)]
fn local_move<G: LevelGraph>(
g: &G,
node_w: &[f64],
membership: &mut [u32],
res: f64,
max_evals: usize,
shuffle: bool,
mode: LocalMoving,
rng: &mut StdRng,
) -> (usize, bool) {
if g.n() < LEIDEN_PAR_MIN_NODES {
return move_nodes(g, node_w, membership, res, max_evals, shuffle, rng);
}
match mode {
LocalMoving::Queue => move_nodes(g, node_w, membership, res, max_evals, shuffle, rng),
LocalMoving::Coloured => move_nodes_coloured(g, node_w, membership, res, max_evals, rng),
}
}
#[allow(clippy::too_many_arguments)]
#[inline]
fn best_move<G: LevelGraph>(
g: &G,
v: u32,
membership: &[u32],
cluster_w: &[f64],
cluster_n: &[u32],
kv: f64,
res: f64,
table: &mut ClusterTable,
) -> Option<Move> {
let (nbrs, ws) = g.row(v as usize);
table.reserve(nbrs.len());
for (&u, &w) in nbrs.iter().zip(ws) {
if u != v {
table.add(membership[u as usize], wf(w));
}
}
let own = membership[v as usize];
let w_from = table.get(own);
let stay = w_from - kv * (cluster_w[own as usize] - kv) * res;
let mut best = (own, stay);
let mut w_best = 0.0;
for &i in &table.used {
let c = table.keys[i as usize];
if c == own {
continue;
}
let gain = table.vals[i as usize] - kv * cluster_w[c as usize] * res;
if gain > best.1 || (gain == best.1 && best.0 != own && c < best.0) {
best = (c, gain);
w_best = table.vals[i as usize];
}
}
if best.0 == own && cluster_n[own as usize] > 1 && 0.0 > stay {
best = (NEW_CLUSTER, 0.0);
w_best = 0.0;
}
table.clear();
(best.0 != own).then_some(Move {
v,
to: best.0,
w_to: w_best,
w_from,
})
}
fn decide_moves<G: LevelGraph>(
g: &G,
nodes: &[u32],
membership: &[u32],
cluster_w: &[f64],
cluster_n: &[u32],
node_w: &[f64],
res: f64,
) -> Vec<Move> {
let decide = |&v: &u32| {
MOVE_SCRATCH.with(|s| {
best_move(
g,
v,
membership,
cluster_w,
cluster_n,
node_w[v as usize],
res,
&mut s.borrow_mut(),
)
})
};
nodes
.par_iter()
.with_min_len(PAR_MOVE_MIN_LEN)
.filter_map(decide)
.collect()
}
fn flag_neighbours<G: LevelGraph>(
g: &G,
moved: &[u32],
membership: &[u32],
pending: &[AtomicBool],
) -> Vec<u32> {
let flag = |&v: &u32| {
let c = membership[v as usize];
g.row(v as usize).0.iter().copied().filter(move |&u| {
let p = &pending[u as usize];
membership[u as usize] != c && !p.load(Relaxed) && !p.swap(true, Relaxed)
})
};
moved
.par_iter()
.with_min_len(PAR_MOVE_MIN_LEN)
.flat_map_iter(flag)
.collect()
}
fn take_pending(list: &mut Vec<u32>, pending: &[AtomicBool]) -> Vec<u32> {
let mut nodes = std::mem::take(list);
nodes.par_sort_unstable();
nodes
.par_iter()
.with_min_len(PAR_MOVE_MIN_LEN)
.for_each(|&v| pending[v as usize].store(false, Relaxed));
nodes
}
fn cluster_totals(membership: &[u32], node_w: &[f64]) -> (Vec<f64>, Vec<u32>, Vec<u32>) {
let n = membership.len();
let mut cluster_w = vec![0.0f64; n];
let mut cluster_n = vec![0u32; n];
for (v, &c) in membership.iter().enumerate() {
cluster_w[c as usize] += node_w[v];
cluster_n[c as usize] += 1;
}
let empty = (0..n as u32)
.rev()
.filter(|&c| cluster_n[c as usize] == 0)
.collect();
(cluster_w, cluster_n, empty)
}
fn greedy_colouring<G: LevelGraph>(g: &G, order: &[u32]) -> (Vec<Vec<u32>>, Vec<u32>) {
let n = g.n();
let mut colour = vec![u32::MAX; n];
let mut used: Vec<bool> = Vec::new();
let mut classes: Vec<Vec<u32>> = Vec::new();
for &v in order {
let nbrs = g.row(v as usize).0;
let mut mask = 0u64;
for &u in nbrs {
let c = colour[u as usize];
if c < 64 {
mask |= 1 << c;
}
}
let c = if mask != u64::MAX {
mask.trailing_ones() as usize
} else {
used.resize(classes.len(), false);
for &u in nbrs {
let cu = colour[u as usize];
if cu != u32::MAX {
used[cu as usize] = true;
}
}
let c = used.iter().position(|&x| !x).unwrap_or(used.len());
used.iter_mut().for_each(|x| *x = false);
c
};
if c == classes.len() {
classes.push(Vec::new());
}
colour[v as usize] = c as u32;
classes[c].push(v);
}
classes.par_iter_mut().for_each(|c| c.sort_unstable());
(classes, colour)
}
fn seeded_colouring<G: LevelGraph>(g: &G, rng: &mut StdRng) -> (Vec<Vec<u32>>, Vec<u32>) {
let n = g.n();
let start = rng.random_range(0..n) as u32;
let order: Vec<u32> = (start..n as u32).chain(0..start).collect();
greedy_colouring(g, &order)
}
fn move_nodes_coloured<G: LevelGraph>(
g: &G,
node_w: &[f64],
membership: &mut [u32],
res: f64,
max_evals: usize,
rng: &mut StdRng,
) -> (usize, bool) {
let n = g.n();
let (classes, colour) = seeded_colouring(g, rng);
let (mut cluster_w, mut cluster_n, mut empty) = cluster_totals(membership, node_w);
let pending: Vec<AtomicBool> = (0..n).map(|_| AtomicBool::new(true)).collect();
let mut lists = classes;
let mut moved: Vec<u32> = Vec::new();
let mut changed = false;
let mut evals = 0usize;
'sweeps: for _ in 0..PAR_MOVE_MAX_ROUNDS {
if lists.iter().all(|l| l.is_empty()) {
break;
}
for c in 0..lists.len() {
if lists[c].is_empty() {
continue;
}
let nodes = take_pending(&mut lists[c], &pending);
evals += nodes.len();
let moves = decide_moves(g, &nodes, membership, &cluster_w, &cluster_n, node_w, res);
moved.clear();
for m in moves {
let (v, kv) = (m.v as usize, node_w[m.v as usize]);
let from = membership[v] as usize;
let stay = m.w_from - kv * (cluster_w[from] - kv) * res;
let go = if m.to == NEW_CLUSTER {
0.0
} else {
m.w_to - kv * cluster_w[m.to as usize] * res
};
if go <= stay || (m.to == NEW_CLUSTER && cluster_n[from] == 1) {
continue;
}
let to = if m.to == NEW_CLUSTER {
empty
.pop()
.expect("a cluster with two or more nodes leaves one id free")
} else {
m.to
} as usize;
cluster_w[from] -= kv;
cluster_n[from] -= 1;
if cluster_n[from] == 0 {
empty.push(from as u32);
}
cluster_w[to] += kv;
cluster_n[to] += 1;
membership[v] = to as u32;
moved.push(v as u32);
}
changed |= !moved.is_empty();
for u in flag_neighbours(g, &moved, membership, &pending) {
lists[colour[u as usize] as usize].push(u);
}
if evals >= max_evals {
break 'sweeps;
}
}
}
(reindex(membership), changed)
}
#[derive(Default)]
struct RefineScratch {
cluster_w: Vec<f64>,
ext: Vec<f64>,
w_to: Vec<f64>,
added: Vec<bool>,
non_singleton: Vec<bool>,
cands: Vec<u32>,
cum: Vec<f64>,
gains: Vec<f64>,
}
impl RefineScratch {
fn reset(&mut self, s: usize) {
for v in [&mut self.cluster_w, &mut self.ext, &mut self.w_to] {
v.clear();
v.resize(s, 0.0);
}
for v in [&mut self.added, &mut self.non_singleton] {
v.clear();
v.resize(s, false);
}
}
}
#[allow(clippy::too_many_arguments)]
fn refine_cluster<G: LevelGraph>(
g: &G,
node_w: &[f64],
membership: &[u32],
local_idx: &[u32],
members: &[u32],
cluster: u32,
res: f64,
beta: f64,
seed: u64,
sc: &mut RefineScratch,
) -> (Vec<u32>, usize) {
let s = members.len();
let mut label: Vec<u32> = (0..s as u32).collect();
if s == 1 {
return (label, 1);
}
sc.reset(s);
let mut total_w = 0.0f64;
for (i, &v) in members.iter().enumerate() {
let v = v as usize;
sc.cluster_w[i] = node_w[v];
total_w += node_w[v];
let (nbrs, ws) = g.row(v);
for (&u, &w) in nbrs.iter().zip(ws) {
if u as usize != v && membership[u as usize] == cluster {
sc.ext[i] += wf(w);
}
}
}
let mut rng = StdRng::seed_from_u64(seed);
let mut order: Vec<u32> = (0..s as u32).collect();
order.shuffle(&mut rng);
for &i in &order {
let i = i as usize;
let cur = label[i] as usize;
let kv = node_w[members[i] as usize];
let well_connected = |sc: &RefineScratch, c: usize| {
sc.ext[c] >= sc.cluster_w[c] * (total_w - sc.cluster_w[c]) * res
};
if sc.non_singleton[cur] || !well_connected(sc, cur) {
continue;
}
let ext_v = sc.ext[cur];
sc.cluster_w[cur] = 0.0;
sc.cands.push(cur as u32);
sc.added[cur] = true;
let (nbrs, ws) = g.row(members[i] as usize);
for (&u, &w) in nbrs.iter().zip(ws) {
let u = u as usize;
if u == members[i] as usize || membership[u] != cluster {
continue;
}
let c = label[local_idx[u] as usize] as usize;
if !sc.added[c] {
sc.added[c] = true;
sc.cands.push(c as u32);
}
sc.w_to[c] += wf(w);
}
let mut best = cur;
let mut max_gain = 0.0f64;
for &c in &sc.cands {
let c = c as usize;
let gain = if well_connected(sc, c) {
sc.w_to[c] - kv * sc.cluster_w[c] * res
} else {
f64::NAN
};
if gain > max_gain {
max_gain = gain;
best = c;
}
sc.gains.push(gain);
}
let mut total = 0.0f64;
for &gain in &sc.gains {
if gain >= 0.0 {
total += ((gain - max_gain) / beta).exp();
}
sc.cum.push(total);
}
let chosen = if total > 0.0 && total.is_finite() {
let r = rng.random::<f64>() * total;
let idx = sc.cum.partition_point(|&x| x <= r).min(sc.cands.len() - 1);
sc.cands[idx] as usize
} else {
best
};
let w_chosen = sc.w_to[chosen];
for &c in &sc.cands {
sc.w_to[c as usize] = 0.0;
sc.added[c as usize] = false;
}
sc.cands.clear();
sc.gains.clear();
sc.cum.clear();
sc.cluster_w[chosen] += kv;
if chosen != cur {
sc.ext[chosen] += ext_v - 2.0 * w_chosen;
label[i] = chosen as u32;
sc.non_singleton[chosen] = true;
}
}
let k = reindex(&mut label);
(label, k)
}
fn refine_partition<G: LevelGraph>(
g: &G,
node_w: &[f64],
membership: &[u32],
n_clusters: usize,
res: f64,
beta: f64,
seed: u64,
) -> (Vec<u32>, usize) {
let n = g.n();
let (offsets, members) = group_by_label(membership, n_clusters);
let mut local_idx = vec![0u32; n];
for c in 0..n_clusters {
for (i, &v) in members[offsets[c]..offsets[c + 1]].iter().enumerate() {
local_idx[v as usize] = i as u32;
}
}
let run = |sc: &mut RefineScratch, c: usize| {
refine_cluster(
g,
node_w,
membership,
&local_idx,
&members[offsets[c]..offsets[c + 1]],
c as u32,
res,
beta,
derive_seed(seed, &[c as u64]),
sc,
)
};
let per_cluster: Vec<(Vec<u32>, usize)> = if n >= LEIDEN_PAR_MIN_NODES {
(0..n_clusters)
.into_par_iter()
.map_init(RefineScratch::default, run)
.collect()
} else {
let mut sc = RefineScratch::default();
(0..n_clusters).map(|c| run(&mut sc, c)).collect()
};
let mut refined = vec![0u32; n];
let mut base = 0u32;
for (c, (labels, k)) in per_cluster.into_iter().enumerate() {
for (&v, &l) in members[offsets[c]..offsets[c + 1]].iter().zip(&labels) {
refined[v as usize] = base + l;
}
base += k as u32;
}
(refined, base as usize)
}
fn aggregate<G: LevelGraph>(
g: &G,
node_w: &[f64],
refined: &[u32],
n_refined: usize,
membership: &[u32],
) -> (AggGraph<f64>, Vec<f64>, Vec<u32>) {
let (offsets, members) = group_by_label(refined, n_refined);
let row = |(acc, seen, touched): &mut (Vec<f64>, Vec<bool>, Vec<u32>), c: usize| {
if acc.len() < n_refined {
acc.resize(n_refined, 0.0);
seen.resize(n_refined, false);
}
let mut w_sum = 0.0;
for &v in &members[offsets[c]..offsets[c + 1]] {
w_sum += node_w[v as usize];
let (nbrs, ws) = g.row(v as usize);
for (&u, &w) in nbrs.iter().zip(ws) {
let c2 = refined[u as usize];
if c2 as usize == c {
continue;
}
if !seen[c2 as usize] {
seen[c2 as usize] = true;
touched.push(c2);
}
acc[c2 as usize] += wf(w);
}
}
let mut idx = Vec::with_capacity(touched.len());
let mut wts = Vec::with_capacity(touched.len());
for &c2 in touched.iter() {
idx.push(c2);
wts.push(acc[c2 as usize]);
acc[c2 as usize] = 0.0;
seen[c2 as usize] = false;
}
touched.clear();
(idx, wts, w_sum)
};
let rows: Vec<(Vec<u32>, Vec<f64>, f64)> = if g.n() >= LEIDEN_PAR_MIN_NODES {
(0..n_refined)
.into_par_iter()
.map_init(|| (Vec::new(), Vec::new(), Vec::new()), row)
.collect()
} else {
let mut scratch = (Vec::new(), Vec::new(), Vec::new());
(0..n_refined).map(|c| row(&mut scratch, c)).collect()
};
let mut indptr = Vec::with_capacity(n_refined + 1);
indptr.push(0usize);
for r in &rows {
indptr.push(indptr.last().unwrap() + r.0.len());
}
let nnz = *indptr.last().unwrap();
let mut indices = Vec::with_capacity(nnz);
let mut weights = Vec::with_capacity(nnz);
let mut agg_w = Vec::with_capacity(n_refined);
for (idx, wts, w_sum) in rows {
indices.extend_from_slice(&idx);
weights.extend_from_slice(&wts);
agg_w.push(w_sum);
}
let agg_membership = (0..n_refined)
.map(|c| membership[members[offsets[c]] as usize])
.collect();
(
AggGraph {
indptr,
indices,
weights,
},
agg_w,
agg_membership,
)
}
fn bfs_order<G: LevelGraph>(g: &G) -> Vec<u32> {
let n = g.n();
let mut seen = vec![false; n];
let mut order = Vec::with_capacity(n);
for s in 0..n {
if seen[s] {
continue;
}
seen[s] = true;
let mut head = order.len();
order.push(s as u32);
while head < order.len() {
let v = order[head] as usize;
head += 1;
for &u in g.row(v).0 {
if !seen[u as usize] {
seen[u as usize] = true;
order.push(u);
}
}
}
}
order
}
fn permute_graph<G: LevelGraph>(g: &G, order: &[u32], pos: &[u32]) -> AggGraph<G::W> {
let mut indptr = Vec::with_capacity(order.len() + 1);
indptr.push(0usize);
for &old in order {
indptr.push(indptr.last().unwrap() + g.row(old as usize).0.len());
}
let nnz = *indptr.last().unwrap();
let mut indices = vec![0u32; nnz];
let mut weights = vec![G::W::zero(); nnz];
let mut rows = Vec::with_capacity(order.len());
let (mut idx_rest, mut w_rest) = (indices.as_mut_slice(), weights.as_mut_slice());
for r in 0..order.len() {
let len = indptr[r + 1] - indptr[r];
let (idx_row, idx_tail) = idx_rest.split_at_mut(len);
let (w_row, w_tail) = w_rest.split_at_mut(len);
rows.push((idx_row, w_row));
idx_rest = idx_tail;
w_rest = w_tail;
}
rows.par_iter_mut()
.zip(order.par_iter())
.for_each(|((idx_row, w_row), &old)| {
let (nbrs, ws) = g.row(old as usize);
for (d, &u) in idx_row.iter_mut().zip(nbrs) {
*d = pos[u as usize];
}
w_row.copy_from_slice(ws);
});
AggGraph {
indptr,
indices,
weights,
}
}
#[allow(clippy::too_many_arguments)]
fn leiden_iteration<G, E, F>(
graph: &G,
node_w: &[f64],
membership: &mut [u32],
res: f64,
params: &LeidenParams,
iteration: u64,
max_evals: usize,
multi_level: bool,
level0: &mut F,
) -> Result<bool, E>
where
G: LevelGraph,
F: FnMut(&G, &[f64], &mut [u32], f64, &mut StdRng) -> Result<(usize, bool), E>,
{
let n = graph.n();
let mut rng = StdRng::seed_from_u64(derive_seed(params.seed, &[iteration]));
reindex(membership);
let (n_clusters, mut changed) = level0(graph, node_w, membership, res, &mut rng)?;
if !multi_level || n_clusters == n {
return Ok(changed);
}
let (refined, n_refined) = refined_or_unrefined(
graph, node_w, membership, n_clusters, res, params, iteration, 0,
);
let mut aggregate_node = refined.clone();
let (mut g, mut w, mut agg_memb) = aggregate(graph, node_w, &refined, n_refined, membership);
let mut level = 1u64;
loop {
let gn = g.n();
let (nc, ch) = local_move(
&g,
&w,
&mut agg_memb,
res,
max_evals.saturating_mul(gn),
!params.refine,
params.local_moving,
&mut rng,
);
changed |= ch;
for (m, &a) in membership.iter_mut().zip(&aggregate_node) {
*m = agg_memb[a as usize];
}
if nc == gn {
break;
}
let (refined, n_refined) =
refined_or_unrefined(&g, &w, &agg_memb, nc, res, params, iteration, level);
for a in aggregate_node.iter_mut() {
*a = refined[*a as usize];
}
(g, w, agg_memb) = aggregate(&g, &w, &refined, n_refined, &agg_memb);
level += 1;
}
Ok(changed)
}
#[allow(clippy::too_many_arguments)]
fn refined_or_unrefined<G: LevelGraph>(
g: &G,
node_w: &[f64],
membership: &[u32],
n_clusters: usize,
res: f64,
params: &LeidenParams,
iteration: u64,
level: u64,
) -> (Vec<u32>, usize) {
if params.refine {
let seed = derive_seed(params.seed, &[iteration, level]);
let (refined, k) =
refine_partition(g, node_w, membership, n_clusters, res, params.beta, seed);
if k < g.n() {
return (refined, k);
}
}
(membership.to_vec(), n_clusters)
}
pub trait LocalMover {
type Error: From<MethodsErrors>;
fn prepare<W: ToPrimitive + Sync>(
&mut self,
indptr: &[usize],
indices: &[u32],
weights: &[W],
node_w: &[f64],
) -> Result<(), Self::Error>;
fn move_nodes(
&mut self,
membership: &mut [u32],
classes: &[Vec<u32>],
colour: &[u32],
res: f64,
) -> Result<bool, Self::Error>;
}
fn detect<T, E, F>(
graph: &SparseGraph<T>,
params: &LeidenParams,
max_evals: usize,
multi_level: bool,
mut level0: F,
) -> Result<Vec<usize>, E>
where
T: BixverseFloat + Send + Sync + std::iter::Sum,
E: From<MethodsErrors> + Send,
F: FnMut(&AggGraph<T>, &[f64], &mut [u32], f64, &mut StdRng) -> Result<(usize, bool), E> + Send,
{
if graph.is_directed() {
return Err(MethodsErrors::GraphDirectedError.into());
}
let n = graph.get_node_number();
if n == 0 {
return Ok(Vec::new());
}
assert!(n < u32::MAX as usize, "more than u32::MAX nodes");
let order = bfs_order(graph);
let mut pos = vec![0u32; n];
for (new, &old) in order.iter().enumerate() {
pos[old as usize] = new as u32;
}
let g = permute_graph(graph, &order, &pos);
let node_w: Vec<f64> = (0..n)
.into_par_iter()
.map(|v| g.row(v).1.iter().map(|&w| wf(w)).sum())
.collect();
let two_m: f64 = node_w.iter().sum();
let mut membership: Vec<u32> = (0..n as u32).collect();
if two_m > 0.0 {
let res = params.resolution / two_m;
rayon::scope(|_| -> Result<(), E> {
let until_stable = params.n_iterations.is_none();
let cap = params.n_iterations.map_or(MAX_ITERATIONS, |k| k.max(1));
let use_gain = until_stable && params.min_gain > 0.0;
let mut quality = if use_gain {
level_quality(&g, &node_w, &membership, res, two_m)
} else {
0.0
};
for it in 0..cap {
let before = until_stable.then(|| membership.clone());
leiden_iteration(
&g,
&node_w,
&mut membership,
res,
params,
it as u64,
max_evals,
multi_level,
&mut level0,
)?;
if let Some(mut before) = before {
reindex(&mut before);
let mut after = membership.clone();
reindex(&mut after);
if before == after {
break;
}
}
if use_gain {
let q = level_quality(&g, &node_w, &membership, res, two_m);
if q - quality <= params.min_gain {
break;
}
quality = q;
}
}
Ok(())
})?;
}
let labels: Vec<usize> = pos
.iter()
.map(|&new| membership[new as usize] as usize)
.collect();
Ok(remap_communities_by_size(&labels))
}
pub fn community_detection<T>(
graph: &SparseGraph<T>,
params: &LeidenParams,
max_evals: usize,
multi_level: bool,
) -> Result<Vec<usize>, MethodsErrors>
where
T: BixverseFloat + Send + Sync + std::iter::Sum,
{
let (shuffle, mode) = (!params.refine, params.local_moving);
detect(
graph,
params,
max_evals,
multi_level,
|g, node_w, membership, res, rng| {
let evals = max_evals.saturating_mul(g.n());
Ok(local_move(
g, node_w, membership, res, evals, shuffle, mode, rng,
))
},
)
}
pub fn community_detection_with<T, M>(
graph: &SparseGraph<T>,
params: &LeidenParams,
max_evals: usize,
multi_level: bool,
mover: &mut M,
) -> Result<Vec<usize>, M::Error>
where
T: BixverseFloat + Send + Sync + std::iter::Sum,
M: LocalMover + Send,
M::Error: Send,
{
let mut prepared = false;
detect(
graph,
params,
max_evals,
multi_level,
|g, node_w, membership, res, rng| {
if !prepared {
mover.prepare(&g.indptr, &g.indices, &g.weights, node_w)?;
prepared = true;
}
let (classes, colour) = seeded_colouring(g, rng);
let changed = mover.move_nodes(membership, &classes, &colour, res)?;
Ok((reindex(membership), changed))
},
)
}
pub fn leiden_sparse_graph<T>(
graph: &SparseGraph<T>,
params: Option<LeidenParams>,
) -> Result<Vec<usize>, MethodsErrors>
where
T: BixverseFloat + Send + Sync + std::iter::Sum,
{
let params = params.unwrap_or_default();
community_detection(graph, ¶ms, usize::MAX, true)
}
pub fn louvain_sparse_graph<T>(
graph: &SparseGraph<T>,
resolution: T,
max_iter: usize,
multi_level: bool,
seed: usize,
) -> Result<Vec<usize>, MethodsErrors>
where
T: BixverseFloat + Send + Sync + std::iter::Sum,
{
let iterations = if multi_level { 2 } else { 1 };
let params = LeidenParams::new(wf(resolution), 0.01, Some(iterations), false, seed as u64);
community_detection(graph, ¶ms, max_iter, multi_level)
}
#[cfg(test)]
mod tests {
use super::*;
use bixverse_core::core::math::sparse::coo_to_csr;
const KARATE: [usize; 156] = [
0, 1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0, 7, 0, 8, 0, 10, 0, 11, 0, 12, 0, 13, 0, 17, 0, 19,
0, 21, 0, 31, 1, 2, 1, 3, 1, 7, 1, 13, 1, 17, 1, 19, 1, 21, 1, 30, 2, 3, 2, 7, 2, 27, 2,
28, 2, 32, 2, 9, 2, 8, 2, 13, 3, 7, 3, 12, 3, 13, 4, 6, 4, 10, 5, 6, 5, 10, 5, 16, 6, 16,
8, 30, 8, 32, 8, 33, 9, 33, 13, 33, 14, 32, 14, 33, 15, 32, 15, 33, 18, 32, 18, 33, 19, 33,
20, 32, 20, 33, 22, 32, 22, 33, 23, 25, 23, 27, 23, 32, 23, 33, 23, 29, 24, 25, 24, 27, 24,
31, 25, 31, 26, 29, 26, 33, 27, 33, 28, 31, 28, 33, 29, 32, 29, 33, 30, 32, 30, 33, 31, 32,
31, 33, 32, 33,
];
fn build_undirected<T: BixverseFloat + BixverseNumeric + std::iter::Sum>(
n: usize,
edges: &[(usize, usize, T)],
directed: bool,
) -> SparseGraph<T> {
let mut rows = Vec::with_capacity(edges.len() * 2);
let mut cols = Vec::with_capacity(edges.len() * 2);
let mut vals = Vec::with_capacity(edges.len() * 2);
for &(u, v, w) in edges {
rows.extend([u, v]);
cols.extend([v, u]);
vals.extend([w, w]);
}
let csr = coo_to_csr(&rows.index_cast(), &cols.index_cast(), &vals, (n, n));
SparseGraph::new(n, csr, directed)
}
fn karate() -> SparseGraph<f64> {
let edges: Vec<(usize, usize, f64)> =
KARATE.chunks_exact(2).map(|e| (e[0], e[1], 1.0)).collect();
build_undirected(34, &edges, false)
}
fn planted(
k: usize,
size: usize,
d_in: usize,
d_out: usize,
seed: u64,
) -> Vec<(usize, usize, f64)> {
let n = k * size;
let mut rng = StdRng::seed_from_u64(seed);
let mut edges = Vec::new();
for v in 0..n {
let b = v / size;
for _ in 0..d_in {
let u = b * size + rng.random_range(0..size);
if u != v {
edges.push((v, u, rng.random_range(1..=3) as f64));
}
}
for _ in 0..d_out {
let u = rng.random_range(0..n);
if u / size != b {
edges.push((v, u, 1.0));
}
}
}
edges
}
fn n_communities(labels: &[usize]) -> usize {
labels.iter().max().map_or(0, |m| m + 1)
}
fn all_connected(graph: &SparseGraph<f64>, labels: &[usize]) -> bool {
let n = graph.get_node_number();
let mut seen = vec![false; n];
let mut roots = vec![false; n_communities(labels)];
for s in 0..n {
if seen[s] {
continue;
}
if roots[labels[s]] {
return false;
}
roots[labels[s]] = true;
let mut stack = vec![s];
seen[s] = true;
while let Some(v) = stack.pop() {
for &u in graph.get_neighbours_raw(v).0 {
let u = u as usize;
if !seen[u] && labels[u] == labels[v] {
seen[u] = true;
stack.push(u);
}
}
}
}
true
}
#[test]
fn test_modularity_matches_igraph() {
let g = karate();
let labels: Vec<usize> = (0..34).map(|i| i % 2).collect();
approx::assert_relative_eq!(
modularity(&g, &labels, 1.0),
-0.000082182774,
epsilon = 1e-10
);
approx::assert_relative_eq!(
modularity(&g, &labels, 0.5),
0.249958908613,
epsilon = 1e-10
);
}
#[test]
fn test_leiden_karate_modularity() {
let g = karate();
for seed in 0..10 {
let p = LeidenParams {
seed,
..Default::default()
};
let labels = leiden_sparse_graph(&g, Some(p)).unwrap();
let q = modularity(&g, &labels, 1.0);
assert!(q >= 0.418803 - 1e-6, "seed {seed}: modularity {q}");
assert!(all_connected(&g, &labels));
}
}
#[test]
fn test_louvain_karate_modularity() {
let g = karate();
for seed in 0..10 {
let labels = louvain_sparse_graph(&g, 1.0, 50, true, seed).unwrap();
let q = modularity(&g, &labels, 1.0);
assert!(q >= 0.388560 - 1e-6, "seed {seed}: modularity {q}");
}
}
#[test]
fn test_leiden_recovers_planted_blocks() {
let edges = planted(8, 200, 10, 1, 7);
let g = build_undirected(1600, &edges, false);
let labels = leiden_sparse_graph(&g, None).unwrap();
assert_eq!(n_communities(&labels), 8);
for v in 0..1600 {
assert_eq!(labels[v], labels[(v / 200) * 200]);
}
}
#[test]
fn test_leiden_thread_count_invariant() {
let edges = planted(20, 400, 8, 3, 11);
let g = build_undirected(8000, &edges, false);
assert!(g.get_node_number() >= LEIDEN_PAR_MIN_NODES);
let run = |threads: usize| {
rayon::ThreadPoolBuilder::new()
.num_threads(threads)
.build()
.unwrap()
.install(|| leiden_sparse_graph(&g, None).unwrap())
};
let one = run(1);
assert_eq!(one, run(4));
assert!(all_connected(&g, &one));
}
#[test]
fn test_leiden_until_stable() {
let edges = planted(20, 400, 8, 3, 11);
let g = build_undirected(8000, &edges, false);
let q_two = modularity(&g, &leiden_sparse_graph(&g, None).unwrap(), 1.0);
for min_gain in [0.0, 1e-4] {
let p = LeidenParams {
n_iterations: None,
min_gain,
..Default::default()
};
let run = |threads: usize| {
rayon::ThreadPoolBuilder::new()
.num_threads(threads)
.build()
.unwrap()
.install(|| leiden_sparse_graph(&g, Some(p)).unwrap())
};
let one = run(1);
assert_eq!(one, run(4), "min_gain {min_gain}");
assert!(all_connected(&g, &one));
let q = modularity(&g, &one, 1.0);
assert!(q >= q_two - 1e-3, "min_gain {min_gain}: {q} vs {q_two}");
}
}
fn run_mode(g: &SparseGraph<f64>, mode: LocalMoving, refine: bool, seed: u64) -> Vec<usize> {
let p = LeidenParams {
refine,
n_iterations: Some(if refine { 2 } else { 1 }),
seed,
local_moving: mode,
..Default::default()
};
community_detection(g, &p, usize::MAX, true).unwrap()
}
#[test]
fn test_parallel_local_moving_thread_count_invariant() {
let edges = planted(20, 400, 8, 3, 11);
let g = build_undirected(8000, &edges, false);
for refine in [true, false] {
let run = |threads: usize| {
rayon::ThreadPoolBuilder::new()
.num_threads(threads)
.build()
.unwrap()
.install(|| run_mode(&g, LocalMoving::Coloured, refine, 3))
};
let one = run(1);
assert_eq!(one, run(4), "refine {refine}");
if refine {
assert!(all_connected(&g, &one));
}
}
}
#[test]
fn test_parallel_local_moving_recovers_planted_blocks() {
let edges = planted(8, 600, 10, 1, 7);
let g = build_undirected(4800, &edges, false);
assert!(g.get_node_number() >= LEIDEN_PAR_MIN_NODES);
for refine in [true, false] {
let labels = run_mode(&g, LocalMoving::Coloured, refine, 1);
assert_eq!(n_communities(&labels), 8, "refine {refine}");
for v in 0..4800 {
assert_eq!(labels[v], labels[(v / 600) * 600], "refine {refine}");
}
}
}
#[test]
fn test_greedy_colouring_matches_first_fit() {
let mut edges = planted(4, 100, 12, 2, 5);
for u in 0..80 {
for v in u + 1..80 {
edges.push((u, v, 1.0));
}
}
let g = build_undirected(400, &edges, false);
let order: Vec<u32> = (150..400).chain(0..150).collect();
let (classes, colour) = greedy_colouring(&g, &order);
let mut want = vec![u32::MAX; 400];
for &v in &order {
let taken: Vec<u32> = g
.row(v as usize)
.0
.iter()
.map(|&u| want[u as usize])
.collect();
want[v as usize] = (0..).find(|c| !taken.contains(c)).unwrap();
}
assert_eq!(colour, want);
assert!(classes.len() >= 80);
for (c, members) in classes.iter().enumerate() {
assert!(members.is_sorted());
assert!(members.iter().all(|&v| colour[v as usize] == c as u32));
}
}
#[test]
fn test_parallel_local_moving_small_graph_uses_queue() {
let g = karate();
assert_eq!(
run_mode(&g, LocalMoving::Coloured, true, 5),
run_mode(&g, LocalMoving::Queue, true, 5)
);
}
#[test]
fn test_parallel_local_moving_modularity_close_to_queue() {
let edges = planted(20, 400, 8, 3, 11);
let g = build_undirected(8000, &edges, false);
for refine in [true, false] {
let q_queue = modularity(&g, &run_mode(&g, LocalMoving::Queue, refine, 2), 1.0);
let q = modularity(&g, &run_mode(&g, LocalMoving::Coloured, refine, 2), 1.0);
assert!(
q >= 0.99 * q_queue,
"refine {refine}: {q} vs queue {q_queue}"
);
}
}
#[test]
fn test_leiden_same_seed_same_labels() {
let edges = planted(10, 100, 6, 3, 3);
let g = build_undirected(1000, &edges, false);
let p = LeidenParams {
seed: 5,
..Default::default()
};
assert_eq!(
leiden_sparse_graph(&g, Some(p)).unwrap(),
leiden_sparse_graph(&g, Some(p)).unwrap()
);
}
#[test]
fn test_leiden_f32_matches_f64() {
let edges = planted(6, 150, 6, 3, 9);
let edges32: Vec<(usize, usize, f32)> =
edges.iter().map(|&(u, v, w)| (u, v, w as f32)).collect();
let g64 = build_undirected(900, &edges, false);
let g32 = build_undirected(900, &edges32, false);
assert_eq!(
leiden_sparse_graph(&g64, None).unwrap(),
leiden_sparse_graph(&g32, None).unwrap()
);
}
#[test]
fn test_leiden_communities_connected() {
let edges = planted(30, 60, 3, 2, 21);
let g = build_undirected(1800, &edges, false);
for seed in 0..5 {
let p = LeidenParams {
seed,
..Default::default()
};
let labels = leiden_sparse_graph(&g, Some(p)).unwrap();
assert!(all_connected(&g, &labels), "seed {seed}");
}
}
#[test]
fn test_leiden_directed_errors() {
let edges: Vec<(usize, usize, f64)> =
KARATE.chunks_exact(2).map(|e| (e[0], e[1], 1.0)).collect();
let g = build_undirected(34, &edges, true);
assert!(leiden_sparse_graph(&g, None).is_err());
}
#[test]
fn test_leiden_empty_and_edgeless() {
let g = build_undirected::<f64>(0, &[], false);
assert!(leiden_sparse_graph(&g, None).unwrap().is_empty());
let g = build_undirected::<f64>(3, &[], false);
assert_eq!(leiden_sparse_graph(&g, None).unwrap().len(), 3);
}
}