bixverse_methods/graph/

spanning_tree.rs

//! Minimum and maximum spanning forests over CSR adjacencies.
//!
//! Kruskal with union-find. The adjacency is read as **undirected**: a stored
//! entry connects its endpoints whichever direction it sits in, so a directed
//! CSR needs no transpose. Where a pair is stored in both directions with
//! different weights, the weight that favours the objective wins (smaller for
//! the minimum forest, larger for the maximum one).
//!
//! Disconnected input yields a forest rather than an error, matching
//! `scipy.sparse.csgraph.minimum_spanning_tree`. Two deliberate divergences
//! from that routine: explicitly stored zero weights are kept as real edges
//! (scipy treats a zero as an absent edge), and ties resolve by `(weight, lo,
//! hi)` so the output does not depend on traversal order.

use rayon::prelude::*;

use crate::graph::graph_structures::UnionFind;
use crate::prelude::*;
use bixverse_core::core::math::sparse::{undirected_edges_to_csr, validate_square_csr};

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

/// Shared Kruskal core behind both public entry points.
///
/// One sort over the candidate edges, then the union-find pass. A pair stored in
/// both directions needs no separate collapse: the objective-favouring copy
/// sorts first under the Kruskal ordering, so union accepts it and every later
/// copy is rejected as a cycle. The result is identical to deduplicating first
/// and half a sort cheaper.
///
/// The sort is the whole cost, so it is the only part that runs in parallel.
/// Collecting the candidates is linear and the union-find pass is inherently
/// sequential. Rayon leaves short slices on the sequential path, so this stays
/// cheap on the tiny graphs that graph abstraction feeds it.
///
/// ### Params
///
/// * `graph` - Square CSR adjacency, read as undirected.
/// * `maximise` - Prefer the heaviest edges rather than the lightest.
///
/// ### Returns
///
/// The forest as a symmetric CSR, or an error when `graph` is not a structurally
/// sound square CSR.
fn spanning_forest<T>(
    graph: &CompressedSparseData2<T>,
    maximise: bool,
) -> Result<CompressedSparseData2<T>, MethodsErrors>
where
    T: BixverseFloat + BixverseNumeric,
{
    let n = validate_square_csr(graph)?;

    let mut edges = candidate_edges(graph, n);

    edges.par_sort_unstable_by(|a, b| {
        weight_order(a.2, b.2, maximise).then_with(|| (a.0, a.1).cmp(&(b.0, b.1)))
    });

    let mut uf = UnionFind::new(n);
    let mut kept: Vec<(u32, u32, T)> = Vec::with_capacity(n.saturating_sub(1));
    for &(lo, hi, w) in &edges {
        if uf.union(lo, hi) {
            kept.push((lo, hi, w));
        }
    }

    kept.sort_unstable_by_key(|e| (e.0, e.1));

    Ok(undirected_edges_to_csr(n, &kept))
}

/// Collect the undirected candidate edges of a CSR adjacency.
///
/// Self loops and `NaN` weights are dropped: a self loop can never belong to a
/// spanning forest, and a `NaN` would poison the comparator and silently
/// reorder the whole Kruskal sequence.
///
/// Infinities are kept. `total_cmp` orders them correctly, scipy treats them as
/// ordinary large weights, and for [maximum_spanning_forest] a `+inf` edge is
/// the *most* desirable one in the graph. Dropping it would hand the caller an
/// extra component whenever it was the only bridge between two of them.
///
/// ### Params
///
/// * `graph` - Square CSR adjacency.
/// * `n` - Node count.
///
/// ### Returns
///
/// `(lo, hi, weight)` triples with `lo < hi`, still holding both directions of
/// any pair stored twice.
fn candidate_edges<T>(graph: &CompressedSparseData2<T>, n: usize) -> Vec<(u32, u32, T)>
where
    T: BixverseFloat + BixverseNumeric,
{
    let mut edges: Vec<(u32, u32, T)> = Vec::with_capacity(graph.get_nnz());

    for i in 0..n {
        let row = i as u32;
        for idx in graph.indptr[i] as usize..graph.indptr[i + 1] as usize {
            let j = graph.indices[idx];
            let w = graph.data[idx];
            if j == row || w.is_nan() {
                continue;
            }
            edges.push((row.min(j), row.max(j), w));
        }
    }

    edges
}

/// Order two weights so that the objective-favouring one compares as less.
///
/// ### Params
///
/// * `a` - First weight.
/// * `b` - Second weight.
/// * `maximise` - Heaviest-first rather than lightest-first.
///
/// ### Returns
///
/// The ordering to sort by. Uses `total_cmp`, which is a total order over the
/// infinities [candidate_edges] lets through.
#[inline]
fn weight_order<T>(a: T, b: T, maximise: bool) -> std::cmp::Ordering
where
    T: BixverseFloat,
{
    if maximise {
        b.total_cmp(&a)
    } else {
        a.total_cmp(&b)
    }
}

//////////
// Main //
//////////

/// Minimum spanning forest of a CSR adjacency.
///
/// ### Params
///
/// * `graph` - Square CSR adjacency, read as undirected. Self loops and `NaN`
///   weights are skipped; infinities are ordinary weights.
///
/// ### Returns
///
/// A symmetric CSR of the same shape holding the retained edges at their
/// original weights, or an error when `graph` is not a structurally sound square
/// CSR.
pub fn minimum_spanning_forest<T>(
    graph: &CompressedSparseData2<T>,
) -> Result<CompressedSparseData2<T>, MethodsErrors>
where
    T: BixverseFloat + BixverseNumeric,
{
    spanning_forest(graph, false)
}

/// Maximum spanning forest of a CSR adjacency.
///
/// The same routine with the weight ordering flipped. Preferring the heaviest
/// edges is what PAGA's `connectivities_tree` wants, and it avoids the
/// reference implementation's detour through the reciprocals of the weights.
///
/// ### Params
///
/// * `graph` - Square CSR adjacency, read as undirected. Self loops and `NaN`
///   weights are skipped; infinities are ordinary weights.
///
/// ### Returns
///
/// A symmetric CSR of the same shape holding the retained edges at their
/// original weights, or an error when `graph` is not a structurally sound square
/// CSR.
pub fn maximum_spanning_forest<T>(
    graph: &CompressedSparseData2<T>,
) -> Result<CompressedSparseData2<T>, MethodsErrors>
where
    T: BixverseFloat + BixverseNumeric,
{
    spanning_forest(graph, true)
}

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

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

    /// Build a symmetric CSR over `n` nodes from `(a, b, weight)` edges, storing
    /// both directions with each row sorted by column.
    fn undirected_csr(n: usize, edges: &[(usize, usize, f64)]) -> CompressedSparseData2<f64> {
        let mut rows: Vec<Vec<(u32, f64)>> = vec![Vec::new(); n];
        for &(a, b, w) in edges {
            rows[a].push((b as u32, w));
            rows[b].push((a as u32, w));
        }

        let mut indptr: Vec<u32> = vec![0];
        let mut indices: Vec<u32> = Vec::new();
        let mut data: Vec<f64> = Vec::new();
        for row in rows.iter_mut() {
            row.sort_unstable_by_key(|e| e.0);
            for &(j, w) in row.iter() {
                indices.push(j);
                data.push(w);
            }
            indptr.push(indices.len() as u32);
        }

        CompressedSparseData2::new_csr(&data, &indices, &indptr, None, (n, n))
    }

    /// Build a symmetric CSR whose rows keep the insertion order of the edge
    /// list rather than being sorted by column.
    ///
    /// A valid CSR for this crate wants ascending rows, so this is only for
    /// probing order dependence in a routine that reads the rows as an unordered
    /// bag of edges.
    fn undirected_csr_unsorted(
        n: usize,
        edges: &[(usize, usize, f64)],
    ) -> CompressedSparseData2<f64> {
        let mut rows: Vec<Vec<(u32, f64)>> = vec![Vec::new(); n];
        for &(a, b, w) in edges {
            rows[a].push((b as u32, w));
            rows[b].push((a as u32, w));
        }

        let mut indptr: Vec<u32> = vec![0];
        let mut indices: Vec<u32> = Vec::new();
        let mut data: Vec<f64> = Vec::new();
        for row in rows.iter() {
            for &(j, w) in row.iter() {
                indices.push(j);
                data.push(w);
            }
            indptr.push(indices.len() as u32);
        }

        CompressedSparseData2::new_csr(&data, &indices, &indptr, None, (n, n))
    }

    /// Read a CSR back as deduplicated undirected `(lo, hi, weight)` edges,
    /// ascending by `(lo, hi)`.
    fn edge_list(csr: &CompressedSparseData2<f64>) -> Vec<(usize, usize, f64)> {
        let n = csr.nrows();
        let mut out = Vec::new();
        for i in 0..n {
            for idx in csr.indptr[i] as usize..csr.indptr[i + 1] as usize {
                let j = csr.indices[idx] as usize;
                if i < j {
                    out.push((i, j, csr.data[idx]));
                }
            }
        }
        out.sort_unstable_by_key(|e| (e.0, e.1));
        out
    }

    /// The two objectives really do run Kruskal in opposite orders and land on
    /// different edge sets.
    #[test]
    fn test_minimum_and_maximum_forests_pick_different_edges() {
        // A 4-cycle with one chord. Cheapest three edges are 0-1, 1-2, 2-3.
        // Going the other way, 0-2 and 0-3 are taken first, which leaves 2-3
        // closing a cycle, so node 1 gets attached through 1-2 instead.
        let graph = undirected_csr(
            4,
            &[
                (0, 1, 1.0),
                (1, 2, 2.0),
                (2, 3, 3.0),
                (0, 3, 9.0),
                (0, 2, 10.0),
            ],
        );

        let min_tree = edge_list(&minimum_spanning_forest(&graph).unwrap());
        assert_eq!(min_tree, vec![(0, 1, 1.0), (1, 2, 2.0), (2, 3, 3.0)]);

        let max_tree = edge_list(&maximum_spanning_forest(&graph).unwrap());
        assert_eq!(max_tree, vec![(0, 2, 10.0), (0, 3, 9.0), (1, 2, 2.0)]);
    }

    /// A disconnected graph gives a forest, not an attempt at a single tree.
    #[test]
    fn test_forest_spans_each_component_separately() {
        // Two triangles with no edge between them: 4 kept edges, not 5.
        let graph = undirected_csr(
            6,
            &[
                (0, 1, 1.0),
                (1, 2, 2.0),
                (0, 2, 3.0),
                (3, 4, 1.0),
                (4, 5, 2.0),
                (3, 5, 3.0),
            ],
        );

        let forest = edge_list(&minimum_spanning_forest(&graph).unwrap());

        assert_eq!(forest.len(), 4);
        assert_eq!(
            forest,
            vec![(0, 1, 1.0), (1, 2, 2.0), (3, 4, 1.0), (4, 5, 2.0)]
        );
    }

    /// An isolated node keeps its row in the output shape while adding no edge.
    #[test]
    fn test_isolated_node_contributes_no_edges() {
        let graph = undirected_csr(3, &[(0, 1, 1.0)]);

        let forest = minimum_spanning_forest(&graph).unwrap();

        assert_eq!(edge_list(&forest), vec![(0, 1, 1.0)]);
        assert_eq!(forest.shape, (3, 3));
        assert_eq!(forest.indptr[3], 2);
    }

    /// When the two stored directions disagree the objective picks the weight,
    /// rather than whichever direction happens to be read first.
    #[test]
    fn test_asymmetric_pair_resolves_towards_the_objective() {
        // 0->1 stored at 5.0, 1->0 at 1.0. The minimum forest must take 1.0
        // and the maximum forest 5.0.
        let data = vec![5.0f64, 1.0];
        let indices = vec![1u32, 0];
        let indptr = vec![0u32, 1, 2];
        let graph = CompressedSparseData2::new_csr(&data, &indices, &indptr, None, (2, 2));

        assert_eq!(
            edge_list(&minimum_spanning_forest(&graph).unwrap()),
            vec![(0, 1, 1.0)]
        );
        assert_eq!(
            edge_list(&maximum_spanning_forest(&graph).unwrap()),
            vec![(0, 1, 5.0)]
        );
    }

    /// Ties resolve on `(lo, hi)`, so storage order cannot change the answer.
    #[test]
    fn test_all_equal_weights_break_ties_deterministically() {
        // Every edge of K4 at the same weight: the tie-break is (lo, hi), so
        // the star at node 0 wins. Building the same graph twice in a *permuted
        // storage order* is the point. Running the same pure function on the
        // same input twice would pass for an order-dependent implementation too.
        let mut edges = Vec::new();
        for i in 0..4 {
            for j in (i + 1)..4 {
                edges.push((i, j, 1.0));
            }
        }
        let graph = undirected_csr(4, &edges);

        let mut permuted = edges.clone();
        permuted.reverse();
        let permuted = undirected_csr_unsorted(4, &permuted);

        let first = edge_list(&minimum_spanning_forest(&graph).unwrap());
        let second = edge_list(&minimum_spanning_forest(&permuted).unwrap());

        assert_eq!(first, second);
        assert_eq!(first, vec![(0, 1, 1.0), (0, 2, 1.0), (0, 3, 1.0)]);
    }

    /// Self loops never enter the forest, and a NaN weight is filtered out.
    #[test]
    fn test_self_loops_and_nan_weights_are_dropped() {
        let data = vec![f64::NAN, 7.0, 1.0, 1.0];
        let indices = vec![0u32, 1, 0, 0];
        let indptr = vec![0u32, 2, 3, 4];
        // Row 0: self loop at NaN plus 0-1 at 7.0. Row 1: 1-0 at 1.0. Row 2:
        // 2-0 at 1.0. The NaN self loop must not survive, and 0-1 resolves to
        // the smaller stored weight.
        let graph = CompressedSparseData2::new_csr(&data, &indices, &indptr, None, (3, 3));

        assert_eq!(
            edge_list(&minimum_spanning_forest(&graph).unwrap()),
            vec![(0, 1, 1.0), (0, 2, 1.0)]
        );
    }

    /// Regression: a NaN on a genuine off-diagonal edge reached `total_cmp` and
    /// reordered the whole Kruskal sequence.
    #[test]
    fn test_a_nan_on_a_real_edge_is_dropped() {
        // The self-loop case above short-circuits on `j == row`, so it passes
        // with the NaN filter deleted. This one cannot: a NaN reaching
        // `total_cmp` sorts above every finite weight for the maximum forest and
        // silently reorders the whole Kruskal sequence.
        let graph = undirected_csr(3, &[(0, 1, f64::NAN), (0, 2, 5.0), (1, 2, 6.0)]);

        assert_eq!(
            edge_list(&maximum_spanning_forest(&graph).unwrap()),
            vec![(0, 2, 5.0), (1, 2, 6.0)]
        );
        assert_eq!(
            edge_list(&minimum_spanning_forest(&graph).unwrap()),
            vec![(0, 2, 5.0), (1, 2, 6.0)]
        );
    }

    /// Infinities are weights, not sentinels: the NaN filter must not take them.
    #[test]
    fn test_infinite_weights_are_ordinary_edges() {
        // scipy treats infinities as very large weights and so does this.
        // Filtering them out alongside NaN discards the single most desirable
        // edge of a maximum forest, and hands the caller an extra component
        // whenever it was the only bridge between two of them.
        let graph = undirected_csr(4, &[(0, 1, 1.0), (2, 3, 1.0), (1, 2, f64::INFINITY)]);

        let max_forest = edge_list(&maximum_spanning_forest(&graph).unwrap());
        assert_eq!(max_forest.len(), 3);
        assert!(max_forest.iter().any(|e| (e.0, e.1) == (1, 2)));

        // The minimum forest takes it last but still has to take it: nothing
        // else connects the two halves.
        let min_forest = edge_list(&minimum_spanning_forest(&graph).unwrap());
        assert_eq!(min_forest.len(), 3);
        assert!(min_forest.iter().any(|e| (e.0, e.1) == (1, 2)));
    }

    /// A stored zero is a real edge here, not an absent one.
    #[test]
    fn test_zero_weight_edges_are_kept() {
        // scipy's csgraph would treat the stored zero as an absent edge. Here
        // it is the cheapest edge in the graph and must win.
        let graph = undirected_csr(3, &[(0, 1, 0.0), (1, 2, 1.0), (0, 2, 2.0)]);

        assert_eq!(
            edge_list(&minimum_spanning_forest(&graph).unwrap()),
            vec![(0, 1, 0.0), (1, 2, 1.0)]
        );
    }

    /// A zero-node graph returns an empty forest rather than panicking.
    #[test]
    fn test_empty_graph_returns_empty_forest() {
        let graph = CompressedSparseData2::<f64>::new_csr(&[], &[], &[0], None, (0, 0));

        let forest = minimum_spanning_forest(&graph).unwrap();

        assert_eq!(forest.shape, (0, 0));
        assert_eq!(forest.get_nnz(), 0);
    }

    /// A CSC matrix is refused rather than read as if it were CSR.
    #[test]
    fn test_rejects_csc_input() {
        let graph = undirected_csr(3, &[(0, 1, 1.0)]).transform();

        assert!(minimum_spanning_forest(&graph).is_err());
    }

    /// A non-square adjacency is not a graph and must not be walked.
    #[test]
    fn test_rejects_non_square_input() {
        let data = vec![1.0f64];
        let indices = vec![1u32];
        let indptr = vec![0u32, 1, 1];
        let graph = CompressedSparseData2::new_csr(&data, &indices, &indptr, None, (2, 3));

        assert!(maximum_spanning_forest(&graph).is_err());
    }
}