leiden-rs 0.6.0

//! Quality functions and graph data representation for community detection.

use gryf::core::{
    base::NeighborReference,
    id::{IdType, IntegerIdType},
    marker::Undirected,
    GraphBase, GraphRef, Neighbors, VertexSet,
};

/// Internal CSR graph representation for the Leiden algorithm.
///
/// Each undirected edge is stored twice (once per direction) so that
/// iterating over all neighbors of a node is O(degree).
pub struct GraphData {
    n: usize,
    adj_offsets: Vec<usize>,
    adj_targets: Vec<usize>,
    adj_weights: Vec<f64>,
    total_weight: f64,
    degree: Vec<f64>,
    /// Weight of each node for CPM community size tracking.
    /// For the original graph, all weights are 1.0. For aggregate
    /// graphs, each weight equals the number of original nodes
    /// represented by that aggregate node.
    node_weight: Vec<f64>,
}

impl GraphData {
    /// Extract graph data from a gryf undirected graph.
    pub fn from_gryf_graph<V, G>(
        graph: &gryf::Graph<V, f64, Undirected, G>,
    ) -> crate::error::Result<Self>
    where
        G: GraphBase<VertexId: IntegerIdType, EdgeType = Undirected>
            + Neighbors
            + VertexSet
            + GraphRef<V, f64>,
    {
        let n = graph.vertex_count();
        let mut degree = vec![0.0f64; n];
        let mut neighbor_counts = vec![0usize; n];

        for vertex_id in graph.vertices_by_id() {
            let u = vertex_id.as_usize();
            for neighbor in graph.neighbors_undirected(vertex_id) {
                neighbor_counts[u] += 1;
                let edge_ref = neighbor.edge();
                let weight = graph.edge(edge_ref.as_ref()).copied().unwrap_or(1.0);
                if !(weight.is_finite() && weight >= 0.0) {
                    return Err(crate::error::LeidenError::InvalidEdgeWeight { weight });
                }
                let v = neighbor.id().as_ref().as_usize();
                if u == v {
                    degree[u] += 2.0 * weight;
                } else {
                    degree[u] += weight;
                }
            }
        }

        let mut adj_offsets = vec![0usize; n + 1];
        for i in 0..n {
            adj_offsets[i + 1] = adj_offsets[i] + neighbor_counts[i];
        }

        let total_slots = adj_offsets[n];
        let mut adj_targets = vec![0usize; total_slots];
        let mut adj_weights = vec![0.0f64; total_slots];
        let mut insert_pos = adj_offsets.clone();

        for vertex_id in graph.vertices_by_id() {
            let u = vertex_id.as_usize();
            for neighbor in graph.neighbors_undirected(vertex_id) {
                let v = neighbor.id().as_ref().as_usize();
                let edge_ref = neighbor.edge();
                let weight = graph.edge(edge_ref.as_ref()).copied().unwrap_or(1.0);
                let pos = insert_pos[u];
                adj_targets[pos] = v;
                adj_weights[pos] = weight;
                insert_pos[u] += 1;
            }
        }

        let total_weight = degree.iter().sum::<f64>() / 2.0;

        Ok(GraphData {
            n,
            adj_offsets,
            adj_targets,
            adj_weights,
            total_weight,
            degree,
            node_weight: vec![1.0; n],
        })
    }

    /// Create a `GraphData` from a list of weighted edges.
    ///
    /// Each tuple is `(source, target, weight)`. Node IDs must be in `0..node_count`.
    /// Self-loops are allowed. All weights must be finite and non-negative.
    pub fn from_edgelist(
        edges: &[(usize, usize, f64)],
        node_count: usize,
    ) -> crate::error::Result<Self> {
        let n = node_count;

        let mut degree: Vec<f64> = vec![0.0; n];
        for &(u, v, w) in edges {
            if !(w.is_finite() && w >= 0.0) {
                return Err(crate::error::LeidenError::InvalidEdgeWeight { weight: w });
            }
            if u >= n || v >= n {
                return Err(crate::error::LeidenError::InconsistentStructure {
                    message: format!("node ID {} exceeds node_count {}", u.max(v), n),
                });
            }
            if u == v {
                degree[u] += 2.0 * w;
            } else {
                degree[u] += w;
                degree[v] += w;
            }
        }

        let mut neighbor_count: Vec<usize> = vec![0; n];
        for &(u, v, _) in edges {
            neighbor_count[u] += 1;
            if u != v {
                neighbor_count[v] += 1;
            }
        }

        let mut adj_offsets: Vec<usize> = Vec::with_capacity(n + 1);
        adj_offsets.push(0);
        let mut total = 0;
        for &count in &neighbor_count {
            total += count;
            adj_offsets.push(total);
        }

        let mut adj_targets: Vec<usize> = vec![0; total];
        let mut adj_weights: Vec<f64> = vec![0.0; total];
        let mut cursor: Vec<usize> = adj_offsets[..n].to_vec();

        for &(u, v, w) in edges {
            adj_targets[cursor[u]] = v;
            adj_weights[cursor[u]] = w;
            cursor[u] += 1;
            if u != v {
                adj_targets[cursor[v]] = u;
                adj_weights[cursor[v]] = w;
                cursor[v] += 1;
            }
        }

        let node_weight: Vec<f64> = vec![1.0; n];

        Self::from_parts(
            n,
            adj_offsets,
            adj_targets,
            adj_weights,
            degree,
            node_weight,
        )
    }

    /// Construct a `GraphData` from pre-built CSR components.
    ///
    /// Returns an error if the input slices are inconsistent:
    /// - `adj_offsets.len() != n + 1`
    /// - `adj_targets.len() != adj_weights.len()`
    /// - `degree.len() != n`
    /// - `node_weight.len() != n`
    pub(crate) fn from_parts(
        n: usize,
        adj_offsets: Vec<usize>,
        adj_targets: Vec<usize>,
        adj_weights: Vec<f64>,
        degree: Vec<f64>,
        node_weight: Vec<f64>,
    ) -> crate::error::Result<Self> {
        if adj_offsets.len() != n + 1 {
            return Err(crate::error::LeidenError::InconsistentStructure {
                message: format!(
                    "adj_offsets length {} != n + 1 ({})",
                    adj_offsets.len(),
                    n + 1
                ),
            });
        }
        if adj_targets.len() != adj_weights.len() {
            return Err(crate::error::LeidenError::InconsistentStructure {
                message: format!(
                    "adj_targets length {} != adj_weights length {}",
                    adj_targets.len(),
                    adj_weights.len()
                ),
            });
        }
        if degree.len() != n {
            return Err(crate::error::LeidenError::InconsistentStructure {
                message: format!("degree length {} != n ({})", degree.len(), n),
            });
        }
        if node_weight.len() != n {
            return Err(crate::error::LeidenError::InconsistentStructure {
                message: format!("node_weight length {} != n ({})", node_weight.len(), n),
            });
        }
        if adj_offsets[0] != 0 {
            return Err(crate::error::LeidenError::InconsistentStructure {
                message: format!("adj_offsets[0] must be 0, got {}", adj_offsets[0]),
            });
        }
        if adj_offsets[n] != adj_targets.len() {
            return Err(crate::error::LeidenError::InconsistentStructure {
                message: format!(
                    "adj_offsets[n] ({}) != adj_targets.len() ({})",
                    adj_offsets[n],
                    adj_targets.len()
                ),
            });
        }
        let total_weight = degree.iter().sum::<f64>() / 2.0;
        Ok(GraphData {
            n,
            adj_offsets,
            adj_targets,
            adj_weights,
            total_weight,
            degree,
            node_weight,
        })
    }

    /// Number of nodes in the graph.
    pub fn node_count(&self) -> usize {
        self.n
    }

    /// Sum of all edge weights (each undirected edge counted once).
    pub fn total_weight(&self) -> f64 {
        self.total_weight
    }

    /// Iterate over all `(neighbor, weight)` pairs for a node.
    #[inline]
    pub fn neighbors(&self, node: usize) -> impl Iterator<Item = (usize, f64)> + '_ {
        let start = self.adj_offsets[node];
        let end = self.adj_offsets[node + 1];
        self.adj_targets[start..end]
            .iter()
            .zip(self.adj_weights[start..end].iter())
            .map(|(&t, &w)| (t, w))
    }

    /// Get raw slices of neighbor targets and weights for a node.
    #[inline]
    pub fn neighbor_slices(&self, node: usize) -> (&[usize], &[f64]) {
        let start = self.adj_offsets[node];
        let end = self.adj_offsets[node + 1];
        (&self.adj_targets[start..end], &self.adj_weights[start..end])
    }

    /// Get the weighted degree of a node.
    #[inline]
    pub fn degree_of(&self, node: usize) -> f64 {
        self.degree[node]
    }

    /// Get the weight of a node (1.0 for original nodes, aggregated for super-nodes).
    #[inline]
    pub fn node_weight(&self, node: usize) -> f64 {
        self.node_weight[node]
    }

    /// Total weight of all nodes (sum of `node_weight`).
    ///
    /// Equals `n` for the original graph; equals the number of original nodes
    /// represented by all aggregate super-nodes after aggregation.
    pub fn total_node_weight(&self) -> f64 {
        self.node_weight.iter().sum()
    }
}

/// Parameters for computing the quality delta of moving a node between communities.
pub struct MoveComponents {
    /// Weighted degree of the node being moved.
    pub k_v: f64,
    /// Edge weight from the node to the target community.
    pub k_v_to_target: f64,
    /// Edge weight from the node to its current community.
    pub k_v_to_current: f64,
    /// Total weighted degree of the target community.
    pub sigma_tot_target: f64,
    /// Total weighted degree of the current community.
    pub sigma_tot_current: f64,
    /// Twice the total edge weight of the graph.
    pub two_m: f64,
    /// Total node weight in the target community.
    pub n_target: f64,
    /// Total node weight in the current community.
    pub n_current: f64,
    /// Weight of the node being moved (1.0 for original nodes, sum of originals for aggregate nodes).
    pub node_weight: f64,
    /// Total node weight across all communities in the graph.
    pub total_node_weight: f64,
}

/// Trait for quality functions used by the Leiden algorithm.
pub trait QualityFunction {
    /// Compute the quality delta of moving a node, given precomputed components.
    fn delta_move_from_components(&self, c: &MoveComponents) -> f64;

    /// Compute the total quality of a partition.
    fn total_quality(&self, data: &GraphData, partition: &crate::partition::Partition) -> f64;
}

/// Modularity: Q = Σ_c [e_c/m - γ*(Σ_c/(2m))²]
pub struct Modularity {
    /// Resolution parameter γ.
    pub resolution: f64,
}

impl Modularity {
    /// Create a new Modularity with default resolution (1.0).
    pub fn new() -> Self {
        Self { resolution: 1.0 }
    }

    /// Create a new Modularity with a custom resolution parameter.
    pub fn with_resolution(resolution: f64) -> Self {
        Self { resolution }
    }
}

impl Default for Modularity {
    fn default() -> Self {
        Self::new()
    }
}

impl QualityFunction for Modularity {
    fn delta_move_from_components(&self, c: &MoveComponents) -> f64 {
        if c.two_m == 0.0 {
            return 0.0;
        }
        (c.k_v_to_target - c.k_v_to_current) * 2.0 / c.two_m
            - self.resolution * c.k_v * (c.sigma_tot_target - c.sigma_tot_current + c.k_v) * 2.0
                / (c.two_m * c.two_m)
    }

    fn total_quality(&self, data: &GraphData, partition: &crate::partition::Partition) -> f64 {
        let n = data.node_count();
        let m = data.total_weight();
        if m == 0.0 {
            return 0.0;
        }

        let num_comms = partition.num_communities();
        let mut sigma_tot: Vec<f64> = vec![0.0; num_comms];
        let mut e_c: Vec<f64> = vec![0.0; num_comms];

        for node in 0..n {
            let comm = partition.community_of(node);
            if comm >= num_comms {
                continue;
            }
            sigma_tot[comm] += data.degree_of(node);
            for (neighbor, weight) in data.neighbors(node) {
                if neighbor >= node && partition.community_of(neighbor) == comm {
                    e_c[comm] += weight;
                }
            }
        }

        let two_m = 2.0 * m;
        let mut q = 0.0;
        for c in 0..num_comms {
            q += e_c[c] / m - self.resolution * (sigma_tot[c] / two_m).powi(2);
        }
        q
    }
}

/// CPM (Constant Potts Model): H = Σ_c [e_c - γ * n_c * (n_c - 1) / 2]
pub struct CPM {
    /// Resolution parameter γ.
    pub resolution: f64,
}

impl CPM {
    /// Create a new CPM with the given resolution parameter.
    pub fn new(resolution: f64) -> Self {
        Self { resolution }
    }
}

impl QualityFunction for CPM {
    fn delta_move_from_components(&self, c: &MoveComponents) -> f64 {
        c.k_v_to_target
            - c.k_v_to_current
            - self.resolution * c.node_weight * (c.n_target - c.n_current + c.node_weight)
    }

    fn total_quality(&self, data: &GraphData, partition: &crate::partition::Partition) -> f64 {
        let n = data.node_count();
        let num_comms = partition.num_communities();
        let mut e_c: Vec<f64> = vec![0.0; num_comms];
        let mut n_c: Vec<f64> = vec![0.0; num_comms];

        for node in 0..n {
            let comm = partition.community_of(node);
            if comm >= num_comms {
                continue;
            }
            n_c[comm] += data.node_weight(node);
            for (neighbor, weight) in data.neighbors(node) {
                if neighbor >= node && partition.community_of(neighbor) == comm {
                    e_c[comm] += weight;
                }
            }
        }

        let mut h = 0.0;
        for c in 0..num_comms {
            h += e_c[c] - self.resolution * n_c[c] * (n_c[c] - 1.0) / 2.0;
        }
        h
    }
}

/// RBConfiguration: Reichardt-Bornholdt with configuration model null.
///
/// Q = Σ_c [e_c - γ * K_c² / (4m)]
///
/// Mathematically equivalent to `Modularity::with_resolution(γ)`.
/// Provided for API compatibility with the leidenalg Python library.
pub struct RBConfiguration {
    /// Resolution parameter γ.
    pub resolution: f64,
}

impl RBConfiguration {
    /// Create a new RBConfiguration with default resolution (1.0).
    pub fn new() -> Self {
        Self { resolution: 1.0 }
    }

    /// Create a new RBConfiguration with a custom resolution parameter.
    pub fn with_resolution(resolution: f64) -> Self {
        Self { resolution }
    }
}

impl Default for RBConfiguration {
    fn default() -> Self {
        Self::new()
    }
}

impl QualityFunction for RBConfiguration {
    fn delta_move_from_components(&self, c: &MoveComponents) -> f64 {
        if c.two_m == 0.0 {
            return 0.0;
        }
        (c.k_v_to_target - c.k_v_to_current) * 2.0 / c.two_m
            - self.resolution * c.k_v * (c.sigma_tot_target - c.sigma_tot_current + c.k_v) * 2.0
                / (c.two_m * c.two_m)
    }

    fn total_quality(&self, data: &GraphData, partition: &crate::partition::Partition) -> f64 {
        let n = data.node_count();
        let m = data.total_weight();
        if m == 0.0 {
            return 0.0;
        }

        let num_comms = partition.num_communities();
        let mut sigma_tot: Vec<f64> = vec![0.0; num_comms];
        let mut e_c: Vec<f64> = vec![0.0; num_comms];

        for node in 0..n {
            let comm = partition.community_of(node);
            if comm >= num_comms {
                continue;
            }
            sigma_tot[comm] += data.degree_of(node);
            for (neighbor, weight) in data.neighbors(node) {
                if neighbor >= node && partition.community_of(neighbor) == comm {
                    e_c[comm] += weight;
                }
            }
        }

        let two_m = 2.0 * m;
        let mut q = 0.0;
        for c in 0..num_comms {
            q += e_c[c] / m - self.resolution * (sigma_tot[c] / two_m).powi(2);
        }
        q
    }
}

/// RBER: Reichardt-Bornholdt with Erdős-Rényi null model.
///
/// Q = Σ_c [e_c - γ * p * n_c * (n_c - 1) / 2]
///
/// Where p = 2m / (N*(N-1)) is the graph density and N is the total node weight.
pub struct RBER {
    /// Resolution parameter γ.
    pub resolution: f64,
}

impl RBER {
    /// Create a new RBER with the given resolution parameter.
    pub fn new(resolution: f64) -> Self {
        Self { resolution }
    }
}

impl QualityFunction for RBER {
    fn delta_move_from_components(&self, c: &MoveComponents) -> f64 {
        let total_n = c.total_node_weight;
        if total_n <= 1.0 || c.two_m == 0.0 {
            return 0.0;
        }
        let p = c.two_m / (total_n * (total_n - 1.0));
        (c.k_v_to_target - c.k_v_to_current)
            - self.resolution * p * c.node_weight * (c.n_target - c.n_current + c.node_weight)
    }

    fn total_quality(&self, data: &GraphData, partition: &crate::partition::Partition) -> f64 {
        let n = data.node_count();
        let m = data.total_weight();
        if n <= 1 || m == 0.0 {
            return 0.0;
        }

        let total_n = data.total_node_weight();
        if total_n <= 1.0 {
            return 0.0;
        }
        let p = 2.0 * m / (total_n * (total_n - 1.0));

        let num_comms = partition.num_communities();
        let mut e_c: Vec<f64> = vec![0.0; num_comms];
        let mut n_c: Vec<f64> = vec![0.0; num_comms];

        for node in 0..n {
            let comm = partition.community_of(node);
            if comm >= num_comms {
                continue;
            }
            n_c[comm] += data.node_weight(node);
            for (neighbor, weight) in data.neighbors(node) {
                if neighbor >= node && partition.community_of(neighbor) == comm {
                    e_c[comm] += weight;
                }
            }
        }

        let mut q = 0.0;
        for c in 0..num_comms {
            q += e_c[c] - self.resolution * p * n_c[c] * (n_c[c] - 1.0) / 2.0;
        }
        q
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use gryf::Graph;

    #[test]
    fn test_graph_data_extraction() {
        let mut graph: Graph<i32, f64, _> = Graph::new_undirected();
        let a = graph.add_vertex(1);
        let b = graph.add_vertex(2);
        let c = graph.add_vertex(3);
        graph.add_edge(&a, &b, 1.0);
        graph.add_edge(&b, &c, 2.0);

        let data = GraphData::from_gryf_graph(&graph).unwrap();
        assert_eq!(data.node_count(), 3);
        assert!((data.total_weight() - 3.0).abs() < 1e-10);
        assert!((data.degree_of(a.as_usize()) - 1.0).abs() < 1e-10);
        assert!((data.degree_of(b.as_usize()) - 3.0).abs() < 1e-10);
        assert!((data.degree_of(c.as_usize()) - 2.0).abs() < 1e-10);
    }

    #[test]
    fn test_modularity_delta_positive() {
        let m = Modularity::new();
        let delta = m.delta_move_from_components(&MoveComponents {
            k_v: 3.0,
            k_v_to_target: 2.0,
            k_v_to_current: 0.0,
            sigma_tot_target: 10.0,
            sigma_tot_current: 3.0,
            two_m: 20.0,
            n_target: 1.0,
            n_current: 1.0,
            node_weight: 1.0,
            total_node_weight: 10.0,
        });
        assert!(delta > 0.0);
    }

    #[test]
    fn test_cpm_delta_positive() {
        let cpm = CPM::new(0.1);
        let delta = cpm.delta_move_from_components(&MoveComponents {
            k_v: 3.0,
            k_v_to_target: 2.0,
            k_v_to_current: 0.0,
            sigma_tot_target: 10.0,
            sigma_tot_current: 3.0,
            two_m: 20.0,
            n_target: 5.0,
            n_current: 1.0,
            node_weight: 1.0,
            total_node_weight: 10.0,
        });
        // delta = 2.0 - 0.0 - 0.1 * 1.0 * (5 - 1 + 1) = 2.0 - 0.5 = 1.5
        assert!((delta - 1.5).abs() < 1e-10);
    }

    #[test]
    fn test_rbconfiguration_matches_modularity() {
        let rb = RBConfiguration::new();
        let m = Modularity::new();
        let c = MoveComponents {
            k_v: 3.0,
            k_v_to_target: 2.0,
            k_v_to_current: 0.0,
            sigma_tot_target: 10.0,
            sigma_tot_current: 3.0,
            two_m: 20.0,
            n_target: 1.0,
            n_current: 1.0,
            node_weight: 1.0,
            total_node_weight: 10.0,
        };
        assert!(
            (rb.delta_move_from_components(&c) - m.delta_move_from_components(&c)).abs() < 1e-10
        );
    }

    #[test]
    fn test_rbconfiguration_with_resolution() {
        let rb = RBConfiguration::with_resolution(2.0);
        let m = Modularity::with_resolution(2.0);
        let c = MoveComponents {
            k_v: 5.0,
            k_v_to_target: 3.0,
            k_v_to_current: 1.0,
            sigma_tot_target: 15.0,
            sigma_tot_current: 8.0,
            two_m: 30.0,
            n_target: 3.0,
            n_current: 2.0,
            node_weight: 1.0,
            total_node_weight: 20.0,
        };
        assert!(
            (rb.delta_move_from_components(&c) - m.delta_move_from_components(&c)).abs() < 1e-10
        );
    }

    #[test]
    fn test_rber_delta_positive() {
        let rber = RBER::new(1.0);
        let c = MoveComponents {
            k_v: 5.0,
            k_v_to_target: 4.0,
            k_v_to_current: 0.0,
            sigma_tot_target: 10.0,
            sigma_tot_current: 5.0,
            two_m: 20.0,
            n_target: 5.0,
            n_current: 1.0,
            node_weight: 1.0,
            total_node_weight: 10.0,
        };
        let delta = rber.delta_move_from_components(&c);
        assert!(delta > 0.0, "RBER delta should be positive, got {delta}");
    }

    #[test]
    fn test_rber_delta_calculation() {
        let rber = RBER::new(1.0);
        // p = 20 / (10 * 9) = 0.2222...
        // delta = (4 - 0) - 1.0 * 0.2222 * 1.0 * (5 - 1 + 1) = 4 - 1.111 = 2.889
        let c = MoveComponents {
            k_v: 5.0,
            k_v_to_target: 4.0,
            k_v_to_current: 0.0,
            sigma_tot_target: 10.0,
            sigma_tot_current: 5.0,
            two_m: 20.0,
            n_target: 5.0,
            n_current: 1.0,
            node_weight: 1.0,
            total_node_weight: 10.0,
        };
        let delta = rber.delta_move_from_components(&c);
        let p = 20.0 / (10.0 * 9.0);
        let expected = 4.0 - 1.0 * p * 1.0 * (5.0 - 1.0 + 1.0);
        assert!(
            (delta - expected).abs() < 1e-10,
            "expected {expected}, got {delta}"
        );
    }

    #[test]
    fn test_rber_zero_two_m() {
        let rber = RBER::new(1.0);
        let c = MoveComponents {
            k_v: 0.0,
            k_v_to_target: 0.0,
            k_v_to_current: 0.0,
            sigma_tot_target: 0.0,
            sigma_tot_current: 0.0,
            two_m: 0.0,
            n_target: 1.0,
            n_current: 1.0,
            node_weight: 1.0,
            total_node_weight: 10.0,
        };
        assert!((rber.delta_move_from_components(&c)).abs() < 1e-10);
    }
}