bixverse-methods 0.8.1

//! Graph community detection algorithms. WalkTrap lives here; Leiden and
//! Louvain live in [`crate::graph::leiden`] and are re-exported.

use rayon::prelude::*;
use rustc_hash::{FxHashMap, FxHashSet};
use std::collections::BinaryHeap;
use std::time::Instant;

use crate::prelude::*;

pub use crate::graph::leiden::{LeidenParams, leiden_sparse_graph, louvain_sparse_graph};

////////////
// Consts //
////////////

/// Neighbouring communities above which a merge scores them in parallel
const WALKTRAP_PAR_MIN_NEIGHBOURS: usize = 64;

/////////////////////
// General helpers //
/////////////////////

/// Helper function to remap communities by size
///
/// The non-determinism due to threading etc. can be a huge annoyance in
/// single cell. This function remaps the clusters by size, ensuring more (not
/// complete) reproducibility.
///
/// ### Params
///
/// * `labels` - The membership labels
///
/// ### Returns
///
/// Remapped labels with `0` -> biggest community by membership, `1` -> second
/// biggest, etc.
pub(crate) fn remap_communities_by_size(labels: &[usize]) -> Vec<usize> {
    let n = labels.len();
    let mut counts: FxHashMap<usize, usize> = FxHashMap::default();
    for &l in labels {
        *counts.entry(l).or_insert(0) += 1;
    }

    let mut sorted: Vec<(usize, usize)> = counts.into_iter().collect();
    sorted.sort_unstable_by_key(|a| std::cmp::Reverse(a.1));

    let mut remap: FxHashMap<usize, usize> = FxHashMap::default();
    for (new_label, (old_label, _)) in sorted.into_iter().enumerate() {
        remap.insert(old_label, new_label);
    }

    (0..n).map(|i| remap[&labels[i]]).collect()
}

//////////////
// Walktrap //
//////////////

/////////////////
// WalkScratch //
/////////////////

/// Per-thread dense buffers for the random-walk propagation
///
/// ### Fields
///
/// * `cur` - Current probabilities, nonzero only at `cur_nodes`
/// * `next` - Probabilities of the next step, nonzero only at `next_nodes`
/// * `cur_nodes` - Nodes holding mass in `cur`
/// * `next_nodes` - Nodes holding mass in `next`
/// * `seen` - Membership flags for `next_nodes` while a step is running
struct WalkScratch<T> {
    cur: Vec<T>,
    next: Vec<T>,
    cur_nodes: Vec<u32>,
    next_nodes: Vec<u32>,
    seen: Vec<bool>,
}

impl<T: BixverseFloat> WalkScratch<T> {
    fn new(n: usize) -> Self {
        Self {
            cur: vec![T::zero(); n],
            next: vec![T::zero(); n],
            cur_nodes: Vec::new(),
            next_nodes: Vec::new(),
            seen: vec![false; n],
        }
    }
}

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

/// Squared Euclidean distance between two index-sorted sparse vectors.
///
/// Iterates over the union of indices in a single pass; missing indices are
/// treated as zero.
///
/// ### Params
///
/// * `a` - First index-sorted sparse vector as `(index, value)` pairs
/// * `b` - Second index-sorted sparse vector as `(index, value)` pairs
///
/// ### Returns
///
/// Squared Euclidean distance between the two vectors
fn sparse_sq_dist<T>(a: &[(usize, T)], b: &[(usize, T)]) -> T
where
    T: BixverseFloat,
{
    let mut acc = T::zero();
    let mut i = 0;
    let mut j = 0;
    while i < a.len() && j < b.len() {
        let (ai, av) = a[i];
        let (bi, bv) = b[j];
        if ai == bi {
            let d = av - bv;
            acc += d * d;
            i += 1;
            j += 1;
        } else if ai < bi {
            acc += av * av;
            i += 1;
        } else {
            acc += bv * bv;
            j += 1;
        }
    }
    while i < a.len() {
        let v = a[i].1;
        acc += v * v;
        i += 1;
    }
    while j < b.len() {
        let v = b[j].1;
        acc += v * v;
        j += 1;
    }
    acc
}

/// Weighted average of two index-sorted sparse vectors:
/// `(n_a · a + n_b · b) / (n_a + n_b)`, returned index-sorted.
///
/// Iterates over the union of indices in a single pass; missing indices are
/// treated as zero.
///
/// ### Params
///
/// * `a` - First index-sorted sparse vector as `(index, value)` pairs
/// * `b` - Second index-sorted sparse vector as `(index, value)` pairs
/// * `n_a` - Weight (typically community size) for vector `a`
/// * `n_b` - Weight (typically community size) for vector `b`
///
/// ### Returns
///
/// Merged index-sorted sparse vector
fn merge_sparse_vecs<T>(
    a: &[(usize, T)],
    b: &[(usize, T)],
    n_a: usize,
    n_b: usize,
) -> Vec<(usize, T)>
where
    T: BixverseFloat,
{
    let na = T::from_usize(n_a).unwrap();
    let nb = T::from_usize(n_b).unwrap();
    let nc = na + nb;
    let mut out = Vec::with_capacity(a.len() + b.len());
    let mut i = 0;
    let mut j = 0;
    while i < a.len() && j < b.len() {
        let (ai, av) = a[i];
        let (bi, bv) = b[j];
        if ai == bi {
            out.push((ai, (na * av + nb * bv) / nc));
            i += 1;
            j += 1;
        } else if ai < bi {
            out.push((ai, na * av / nc));
            i += 1;
        } else {
            out.push((bi, nb * bv / nc));
            j += 1;
        }
    }
    while i < a.len() {
        let (ai, av) = a[i];
        out.push((ai, na * av / nc));
        i += 1;
    }
    while j < b.len() {
        let (bi, bv) = b[j];
        out.push((bi, nb * bv / nc));
        j += 1;
    }
    out
}

/// Index-sort a sparse vector, optionally keeping only its `k` largest-value
/// entries first.
///
/// `None` leaves the support untruncated. `Some(k)` keeps the top-`k` entries by
/// mass (the dropped tail holds negligible walk mass), bounding stored support
/// at `k` per community vector.
fn finalise_sparse_vec<T>(mut v: Vec<(usize, T)>, max_support: Option<usize>) -> Vec<(usize, T)>
where
    T: BixverseFloat,
{
    if let Some(k) = max_support
        && v.len() > k
    {
        v.select_nth_unstable_by(k, |a, b| {
            b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
        });
        v.truncate(k);
        v.shrink_to_fit(); // important... I need to reduce the allocations here...
    }
    v.sort_unstable_by_key(|&(idx, _)| idx);
    v
}

/// Compute random-walk probability vectors for all nodes (sparse)
///
/// For each node i, simulates a `walk_length`-step random walk starting at i
/// and records the landing probabilities, scaled by the inverse square root of
/// each node's degree (Pons & Latapy 2005). Probabilities are propagated
/// through hashmaps so only reached nodes are stored.
///
/// Transition probabilities are weight-proportional: `p(i→j) = w(i,j) / deg(i)`.
/// Isolated nodes (zero degree) yield an effectively zero vector.
///
/// ### Params
///
/// * `graph` - Sparse weighted graph
/// * `walk_length` - Number of steps per random walk
/// * `max_support` - Optional usize. If provided, the only the top max support
///   entries are kept and the detail is removed. This can help with large
///   memory data sets.
///
/// ### Returns
///
/// One index-sorted sparse vector per node, where entry `(j, v)` holds
/// `v = P(walk from i lands at j) / sqrt(deg(j))`.
fn compute_walk_probabilities<T>(
    graph: &SparseGraph<T>,
    walk_length: usize,
    max_support: Option<usize>,
) -> Vec<Vec<(usize, T)>>
where
    T: BixverseFloat + std::iter::Sum + Send + Sync,
{
    let n = graph.get_node_number();
    let epsilon = T::from_f64(1e-10).unwrap();

    let inv_sqrt_deg: Vec<T> = (0..n)
        .map(|k| {
            let (_, w) = graph.get_neighbours_raw(k);
            let d: T = w.iter().copied().sum();
            if d > T::zero() {
                T::one() / d.sqrt()
            } else {
                T::zero()
            }
        })
        .collect();

    let transition_probs: Vec<Vec<(usize, T)>> = (0..n)
        .into_par_iter()
        .map(|i| {
            let (neighbours, weights) = graph.get_neighbours_raw(i);
            let degree: T = weights.iter().copied().sum();
            if degree > T::zero() {
                neighbours
                    .iter()
                    .zip(weights.iter())
                    .map(|(&j, &w)| (j as usize, w / degree))
                    .collect()
            } else {
                Vec::new()
            }
        })
        .collect();

    // Dense accumulators per thread, cleared through the touched lists
    (0..n)
        .into_par_iter()
        .map_init(
            || WalkScratch::<T>::new(n),
            |scratch, i| {
                let WalkScratch {
                    cur,
                    next,
                    cur_nodes,
                    next_nodes,
                    seen,
                } = scratch;

                cur[i] = T::one();
                cur_nodes.push(i as u32);

                for _ in 0..walk_length {
                    for &node in cur_nodes.iter() {
                        seen[node as usize] = false;
                    }
                    for &node in cur_nodes.iter() {
                        let p = cur[node as usize];
                        if p > epsilon {
                            for &(nb, tp) in &transition_probs[node as usize] {
                                if !seen[nb] {
                                    seen[nb] = true;
                                    next_nodes.push(nb as u32);
                                }
                                next[nb] += p * tp;
                            }
                        }
                    }
                    for &node in cur_nodes.iter() {
                        cur[node as usize] = T::zero();
                    }
                    cur_nodes.clear();
                    std::mem::swap(cur, next);
                    std::mem::swap(cur_nodes, next_nodes);
                }

                let row: Vec<(usize, T)> = cur_nodes
                    .iter()
                    .map(|&k| (k as usize, cur[k as usize] * inv_sqrt_deg[k as usize]))
                    .collect();
                for &node in cur_nodes.iter() {
                    cur[node as usize] = T::zero();
                    seen[node as usize] = false;
                }
                cur_nodes.clear();

                finalise_sparse_vec(row, max_support)
            },
        )
        .collect()
}

/// Compute Ward's linkage criterion between two communities
///
/// Returns the increase in total within-cluster variance that would result from
/// merging communities A and B, given their squared centroid distance:
///
/// ```Ward(A, B) = (n_a · n_b) / (n_a + n_b) · d²(A, B)```
///
/// ### Params
///
/// * `d_sq` - Squared distance between the community centroids
/// * `n_a` - Size of community A
/// * `n_b` - Size of community B
///
/// ### Returns
///
/// Ward criterion value (lower means a cheaper merge)
#[inline(always)]
fn ward_criterion<T>(d_sq: T, n_a: usize, n_b: usize) -> T
where
    T: BixverseFloat,
{
    let na = T::from_usize(n_a).unwrap();
    let nb = T::from_usize(n_b).unwrap();
    (na * nb) / (na + nb) * d_sq
}

/// Resolve final community labels from a union-find parent array
///
/// Follows each node's parent chain to its root and assigns a contiguous
/// integer label `0..num_communities` in traversal order. The result is a flat
/// label per node with no gaps.
///
/// ### Params
///
/// * `n` - Number of original nodes
/// * `parent` - Union-find parent array; `parent[x] == x` denotes a root
///
/// ### Returns
///
/// Vector of community labels, one per node
fn finalise_labels(n: usize, parent: &[usize]) -> Vec<usize> {
    let mut labels = vec![0usize; n];
    let mut label_map: FxHashMap<usize, usize> = FxHashMap::default();
    let mut next = 0usize;
    for node in 0..n {
        let mut x = node;
        while parent[x] != x {
            x = parent[x];
        }
        let label = *label_map.entry(x).or_insert_with(|| {
            let l = next;
            next += 1;
            l
        });
        labels[node] = label;
    }
    labels
}

/// Walktrap community detection (Pons & Latapy 2005)
///
/// Identifies communities by agglomerative clustering of random-walk probability
/// vectors. Nodes that tend to co-occur on short random walks are grouped
/// together.
///
/// Algorithm outline:
///
/// 1. Compute scaled walk-probability vectors for all nodes (sparse)
/// 2. Initialise each node as its own community
/// 3. Greedily merge the adjacent pair with the smallest Ward criterion
///    until `num_clusters` communities remain
/// 4. Resolve labels via union-find
///
/// Merged community vectors are updated as the weighted average:
/// `q[A∪B] = (n_a · q[A] + n_b · q[B]) / (n_a + n_b)`
///
/// ### Params
///
/// * `graph` - Sparse weighted graph to cluster
/// * `walk_length` - Number of steps for the random walks
/// * `num_clusters` - Target number of communities to return
/// * `max_support` - Optional usize. If provided, the only the top max support
///   entries are kept and the detail is removed. This can help with large
///   memory data sets and the tail tends to hold neglible mass. This becomes
///   an approximation in this case.
/// * `verbose` - Print timing information for each stage
///
/// ### Returns
///
/// Vector of community labels (0-indexed, contiguous), one per node
pub fn walktrap_sparse_graph<T>(
    graph: &SparseGraph<T>,
    walk_length: usize,
    num_clusters: usize,
    max_support: Option<usize>,
    verbose: bool,
) -> Vec<usize>
where
    T: BixverseFloat + std::iter::Sum + Send + Sync,
{
    let n = graph.get_node_number();
    if n == 0 {
        return Vec::new();
    }
    if n <= num_clusters {
        return (0..n).collect();
    }

    let walktrap_start = Instant::now();

    if verbose {
        match max_support {
            Some(ms) => println!(
                "Calculating walk probabilities with max support of {:?}.",
                ms
            ),
            None => println!("Calculating walk probabilities."),
        }
    }

    let t0 = Instant::now();
    let q = compute_walk_probabilities(graph, walk_length, max_support);
    if verbose {
        println!(" ... done in {:.2?}", t0.elapsed());
    }

    let mut comm_q: Vec<Option<Vec<(usize, T)>>> = q.into_iter().map(Some).collect();

    let mut comm_size: Vec<usize> = vec![1; n];
    let mut active: Vec<bool> = vec![true; n];
    let mut comm_parent: Vec<usize> = (0..n).collect();
    let mut adj: Vec<FxHashSet<usize>> = vec![FxHashSet::default(); n];

    let mut edges: Vec<(usize, usize)> = Vec::new();
    for i in 0..n {
        let (neighbours, _) = graph.get_neighbours_raw(i);
        for &j in neighbours {
            let j = j as usize;
            if i != j {
                adj[i].insert(j);
                if i < j {
                    edges.push((i, j));
                }
            }
        }
    }

    let t1 = Instant::now();
    let initial: Vec<(RevOrderedFloat<T>, usize, usize)> = edges
        .par_iter()
        .map(|&(i, j)| {
            let qi = comm_q[i].as_ref().unwrap();
            let qj = comm_q[j].as_ref().unwrap();
            let d_sq = sparse_sq_dist(qi, qj);
            let crit = ward_criterion(d_sq, 1, 1);
            (RevOrderedFloat(crit), i, j)
        })
        .collect();

    if verbose {
        println!("Computing initial distances.")
    }

    let mut heap: BinaryHeap<(RevOrderedFloat<T>, usize, usize)> =
        BinaryHeap::with_capacity(initial.len() * 2);
    for entry in initial {
        heap.push(entry);
    }
    if verbose {
        println!(" ... done in {:.2?}", t1.elapsed());
    }

    let t2 = Instant::now();
    let mut num_active = n;

    while num_active > num_clusters {
        let (a, b) = loop {
            match heap.pop() {
                Some((_, i, j)) => {
                    if active[i] && active[j] {
                        break (i, j);
                    }
                }
                None => {
                    if verbose {
                        println!(
                            "Heap exhausted with {} communities active (graph likely disconnected)",
                            num_active
                        );
                        println!("Finished merging communities: {:.2?}", t2.elapsed());
                        println!(
                            "Finished WalkTrap community detection: {:.2?}",
                            walktrap_start.elapsed()
                        );
                    }
                    return finalise_labels(n, &comm_parent);
                }
            }
        };

        let n_a = comm_size[a];
        let n_b = comm_size[b];
        let n_c = n_a + n_b;

        let new_q = {
            let qa = comm_q[a].as_ref().unwrap();
            let qb = comm_q[b].as_ref().unwrap();
            merge_sparse_vecs(qa, qb, n_a, n_b)
        };

        let mut new_adj: FxHashSet<usize> = FxHashSet::default();
        for &k in &adj[a] {
            if k != b {
                new_adj.insert(k);
            }
        }
        for &k in &adj[b] {
            if k != a {
                new_adj.insert(k);
            }
        }

        active[a] = false;
        active[b] = false;
        comm_q[a] = None;
        comm_q[b] = None;
        adj[a] = FxHashSet::default();
        adj[b] = FxHashSet::default();

        let c = comm_q.len();
        comm_q.push(Some(new_q));
        comm_size.push(n_c);
        active.push(true);
        comm_parent.push(c);
        comm_parent[a] = c;
        comm_parent[b] = c;
        let new_adj_vec: Vec<usize> = new_adj.iter().copied().collect();
        adj.push(new_adj);

        for &k in &new_adj_vec {
            adj[k].remove(&a);
            adj[k].remove(&b);
            adj[k].insert(c);
        }

        let criterion = |&k: &usize| {
            let qc = comm_q[c].as_ref().unwrap();
            let qk = comm_q[k].as_ref().unwrap();
            let d_sq = sparse_sq_dist(qc, qk);
            let crit = ward_criterion(d_sq, n_c, comm_size[k]);
            (RevOrderedFloat(crit), c, k)
        };
        let new_criteria: Vec<(RevOrderedFloat<T>, usize, usize)> =
            if new_adj_vec.len() >= WALKTRAP_PAR_MIN_NEIGHBOURS {
                new_adj_vec.par_iter().map(criterion).collect()
            } else {
                new_adj_vec.iter().map(criterion).collect()
            };

        for entry in new_criteria {
            heap.push(entry);
        }

        num_active -= 1;
    }

    if verbose {
        println!("Finished merging communities: {:.2?}", t2.elapsed());
        println!(
            "Finished WalkTrap community detection: {:.2?}",
            walktrap_start.elapsed()
        );
    }

    let labels = finalise_labels(n, &comm_parent);

    remap_communities_by_size(&labels)
}

///////////
// Tests //
///////////

#[cfg(test)]
mod tests {
    use super::*;
    use bixverse_core::core::math::sparse::CompressedSparseData2;

    fn build_barbell_graph() -> SparseGraph<f64> {
        let indptr = vec![0, 2, 4, 7, 10, 12, 14];
        let indices = vec![1, 2, 0, 2, 0, 1, 3, 2, 4, 5, 3, 5, 3, 4];
        let data = vec![1.0; 14];
        let csr =
            CompressedSparseData2::<f64, f64>::new_csr(&data, &indices, &indptr, None, (6, 6));
        SparseGraph::new(6, csr, false)
    }

    fn build_two_cliques_graph() -> SparseGraph<f64> {
        // K4 {0..3} + K4 {4..7} bridged by edge 3-4
        let mut nbrs: Vec<Vec<usize>> = vec![Vec::new(); 8];
        for i in 0..4 {
            for j in 0..4 {
                if i != j {
                    nbrs[i].push(j);
                }
            }
        }
        for i in 4..8 {
            for j in 4..8 {
                if i != j {
                    nbrs[i].push(j);
                }
            }
        }
        nbrs[3].push(4);
        nbrs[4].push(3);

        let mut indptr = vec![0usize];
        let mut indices = Vec::new();
        let mut data = Vec::new();
        for list in &mut nbrs {
            list.sort();
            for &j in list.iter() {
                indices.push(j);
                data.push(1.0);
            }
            indptr.push(indices.len());
        }

        let csr = CompressedSparseData2::<f64, f64>::new_csr(
            &data,
            &indices.index_cast(),
            &indptr.index_cast(),
            None,
            (8, 8),
        );
        SparseGraph::new(8, csr, false)
    }

    /// Louvain splits a barbell along the bridge, one community per bell.
    #[test]
    fn test_louvain_barbell() {
        let graph = build_barbell_graph();
        let comms = louvain_sparse_graph(&graph, 1.0, 10, true, 42).unwrap();
        assert_eq!(comms[0], comms[1]);
        assert_eq!(comms[1], comms[2]);
        assert_eq!(comms[3], comms[4]);
        assert_eq!(comms[4], comms[5]);
        assert_ne!(comms[2], comms[3]);
    }

    /// Walktrap finds the same barbell split as Louvain.
    #[test]
    fn test_walktrap_barbell() {
        let graph = build_barbell_graph();
        let comms = walktrap_sparse_graph(&graph, 3, 2, None, false);
        assert_eq!(comms[0], comms[1]);
        assert_eq!(comms[1], comms[2]);
        assert_eq!(comms[3], comms[4]);
        assert_eq!(comms[4], comms[5]);
        assert_ne!(comms[2], comms[3]);
    }

    /// A single bridge between two cliques does not merge them.
    #[test]
    fn test_walktrap_two_cliques() {
        let graph = build_two_cliques_graph();
        let comms = walktrap_sparse_graph(&graph, 4, 2, None, false);
        for i in 1..4 {
            assert_eq!(comms[i], comms[0]);
        }
        for i in 5..8 {
            assert_eq!(comms[i], comms[4]);
        }
        assert_ne!(comms[0], comms[4]);
    }
}