#![allow(dead_code)]
use faer::{Mat, MatRef};
use petgraph::Graph;
use rayon::prelude::*;
use rustc_hash::FxHashMap;
use std::clone::Clone;
use std::collections::BinaryHeap;
use crate::core::math::sparse::coo_to_csr;
use crate::prelude::*;
pub(crate) struct UnionFind {
parent: Vec<u32>,
size: Vec<u32>,
}
impl UnionFind {
pub(crate) fn new(n: usize) -> Self {
Self {
parent: (0..n as u32).collect(),
size: vec![1; n],
}
}
pub(crate) fn find(&mut self, mut x: u32) -> u32 {
while self.parent[x as usize] != x {
let grandparent = self.parent[self.parent[x as usize] as usize];
self.parent[x as usize] = grandparent;
x = grandparent;
}
x
}
pub(crate) fn union(&mut self, a: u32, b: u32) -> bool {
let (mut ra, mut rb) = (self.find(a), self.find(b));
if ra == rb {
return false;
}
if self.size[ra as usize] < self.size[rb as usize] {
std::mem::swap(&mut ra, &mut rb);
}
self.parent[rb as usize] = ra;
self.size[ra as usize] += self.size[rb as usize];
true
}
}
#[derive(Clone, Debug)]
pub struct SparseGraph<T: Clone> {
adjacency: CompressedSparseData2<T>,
num_nodes: usize,
directed: bool,
}
impl<T> SparseGraph<T>
where
T: Clone + BixverseFloat + std::iter::Sum,
{
pub fn new(num_nodes: usize, adjacency: CompressedSparseData2<T>, directed: bool) -> Self {
Self {
adjacency,
num_nodes,
directed,
}
}
#[inline]
pub fn get_neighbours(&self, node: usize) -> (Vec<usize>, &[T]) {
let start = self.adjacency.indptr[node] as usize;
let end = self.adjacency.indptr[node + 1] as usize;
(
self.adjacency.indices[start..end].index_cast(),
&self.adjacency.data[start..end],
)
}
#[inline]
pub fn get_node_degree(&self, node: usize) -> usize {
(self.adjacency.indptr[node + 1] - self.adjacency.indptr[node]) as usize
}
pub fn total_weight(&self) -> T {
let mut total = T::zero();
for i in 0..self.num_nodes {
let (_, weights) = self.get_neighbours(i);
total += weights.iter().copied().sum();
}
let half = T::from_f64(0.5).unwrap();
if !self.directed { total * half } else { total }
}
pub fn get_node_number(&self) -> usize {
self.num_nodes
}
pub fn is_directed(&self) -> bool {
self.directed
}
}
pub fn graph_from_strings<'a, T>(
nodes: &'a [String],
from: &[String],
to: &[String],
weights: Option<&[T]>,
undirected: bool,
) -> Graph<&'a str, T>
where
T: BixverseFloat,
{
assert_same_len!(from, to);
if let Some(weights) = weights {
assert_same_len!(from, weights);
}
let mut graph: Graph<&'a str, T> = Graph::new();
let mut term_to_idx = FxHashMap::default();
for term in nodes {
let idx = graph.add_node(term);
term_to_idx.insert(term, idx);
}
for (i, (from, to)) in from.iter().zip(to.iter()).enumerate() {
let from = *term_to_idx.get(from).unwrap();
let to = *term_to_idx.get(to).unwrap();
let weight = weights.map(|w| w[i]).unwrap_or(T::one());
graph.add_edge(from, to, weight);
if undirected {
graph.add_edge(to, from, weight);
}
}
graph
}
#[allow(dead_code)]
pub struct CoarseLevel<T>
where
T: Clone,
{
graph: SparseGraph<T>,
node_map: Vec<usize>,
coarse_weights: Vec<usize>,
}
pub struct GraphHierarchy<T>
where
T: Clone,
{
finest: SparseGraph<T>,
levels: Vec<CoarseLevel<T>>,
}
impl<T> GraphHierarchy<T>
where
T: Clone + BixverseFloat + std::iter::Sum + std::default::Default,
{
pub fn build(graph: SparseGraph<T>, reduction_ratio: T, max_levels: usize) -> Self {
assert!(
!graph.is_directed(),
"This is implemented for undirected graphs!"
);
let mut levels = Vec::new();
let mut current_graph = graph.clone();
let min_nodes = T::from_usize(100).unwrap();
for _ in 0..max_levels {
let coarse_level = Self::coarsen_level(¤t_graph);
let current_nodes = T::from_usize(current_graph.num_nodes).unwrap();
let coarse_nodes = T::from_usize(coarse_level.graph.num_nodes).unwrap();
let reduction = coarse_nodes / current_nodes;
if reduction > reduction_ratio || coarse_nodes < min_nodes {
break;
}
current_graph = coarse_level.graph.clone();
levels.push(coarse_level);
}
Self {
finest: graph,
levels,
}
}
fn coarsen_level(graph: &SparseGraph<T>) -> CoarseLevel<T> {
let matching = Self::heavy_edge_matching(graph);
let (coarse_graph, node_map, coarse_weights) = Self::contract_graph(graph, &matching);
CoarseLevel {
graph: coarse_graph,
node_map,
coarse_weights,
}
}
fn heavy_edge_matching(graph: &SparseGraph<T>) -> Vec<Option<usize>> {
let n = graph.num_nodes;
let mut matching = vec![None; n];
let mut matched = vec![false; n];
for i in 0..n {
if matched[i] {
continue;
}
let (neighbour, weights) = graph.get_neighbours(i);
let mut best_neighbour = None;
let mut best_weight = T::zero();
for (&j, &w) in neighbour.iter().zip(weights.iter()) {
if !matched[j] && j != i && w > best_weight {
best_weight = w;
best_neighbour = Some(j)
}
}
if let Some(j) = best_neighbour {
matching[i] = Some(j);
matching[j] = Some(i);
matched[i] = true;
matched[j] = true;
}
}
matching
}
fn contract_graph(
graph: &SparseGraph<T>,
matching: &[Option<usize>],
) -> (SparseGraph<T>, Vec<usize>, Vec<usize>) {
let n = graph.num_nodes;
let mut node_map = vec![0; n];
let mut coarse_id = 0;
for i in 0..n {
if node_map[i] == 0 || i == 0 {
match matching[i] {
Some(j) if j > i => {
node_map[i] = coarse_id;
node_map[j] = coarse_id;
coarse_id += 1;
}
None => {
node_map[i] = coarse_id;
coarse_id += 1;
}
_ => {} }
}
}
let num_coarse = coarse_id;
let mut coarse_weights = vec![0; num_coarse];
for i in 0..n {
coarse_weights[node_map[i]] += 1;
}
let mut rows = Vec::new();
let mut cols = Vec::new();
let mut vals = Vec::new();
for i in 0..n {
let ci = node_map[i];
let (neighbours, weights) = graph.get_neighbours(i);
for (&j, &w) in neighbours.iter().zip(weights.iter()) {
let cj = node_map[j];
if ci != cj {
rows.push(ci);
cols.push(cj);
vals.push(w);
}
}
}
let mut edges: Vec<(usize, usize, T)> = rows
.into_iter()
.zip(cols)
.zip(vals)
.map(|((r, c), v)| (r, c, v))
.collect();
edges.sort_unstable_by_key(|(r, c, _)| (*r, *c));
let mut deduped_rows = Vec::new();
let mut deduped_cols = Vec::new();
let mut deduped_vals = Vec::new();
if !edges.is_empty() {
let (mut curr_r, mut curr_c, mut curr_v) = edges[0];
for &(r, c, v) in &edges[1..] {
if r == curr_r && c == curr_c {
curr_v += v;
} else {
deduped_rows.push(curr_r);
deduped_cols.push(curr_c);
deduped_vals.push(curr_v);
(curr_r, curr_c, curr_v) = (r, c, v);
}
}
deduped_rows.push(curr_r);
deduped_cols.push(curr_c);
deduped_vals.push(curr_v);
}
let coarse_csr = coo_to_csr(
&deduped_rows.index_cast(),
&deduped_cols.index_cast(),
&deduped_vals,
(num_coarse, num_coarse),
);
let coarse_graph = SparseGraph::new(num_coarse, coarse_csr, false);
(coarse_graph, node_map, coarse_weights)
}
pub fn prolong(&self, level: usize, coarse_labels: &[usize]) -> Vec<usize> {
if level >= self.levels.len() {
return coarse_labels.to_vec();
}
let node_map = &self.levels[level].node_map;
node_map.iter().map(|&ci| coarse_labels[ci]).collect()
}
pub fn coarsest(&self) -> &SparseGraph<T> {
self.levels.last().map(|l| &l.graph).unwrap_or(&self.finest)
}
}
#[derive(Debug, Clone)]
#[allow(dead_code)] pub struct NodeData<'a> {
pub name: &'a str,
pub node_type: &'a str,
}
#[derive(Debug, Clone)]
pub struct EdgeData<'a, T> {
pub edge_type: &'a str,
pub weight: &'a T,
}
pub fn graph_from_strings_with_attributes<'a, T>(
nodes: &'a [String],
node_types: &'a [String],
from: &'a [String],
to: &'a [String],
edge_types: &'a [String],
edge_weights: &'a [T],
) -> Graph<NodeData<'a>, EdgeData<'a, T>>
where
T: BixverseFloat,
{
assert_same_len!(nodes, node_types);
assert_same_len!(from, to, edge_types, edge_weights);
let mut graph: Graph<NodeData<'_>, EdgeData<'_, T>> = Graph::new();
let mut name_to_idx = FxHashMap::default();
for (name, node_type) in nodes.iter().zip(node_types.iter()) {
let node_data = NodeData { name, node_type };
let idx = graph.add_node(node_data);
name_to_idx.insert(name, idx);
}
for (((from_name, to_name), edge_type), weight) in from
.iter()
.zip(to.iter())
.zip(edge_types.iter())
.zip(edge_weights.iter())
{
let from_idx = *name_to_idx
.get(from_name)
.unwrap_or_else(|| panic!("From node '{}' not found in nodes list", from_name));
let to_idx = *name_to_idx
.get(to_name)
.unwrap_or_else(|| panic!("From node '{}' not found in nodes list", from_name));
let edge_data = EdgeData { edge_type, weight };
graph.add_edge(from_idx, to_idx, edge_data);
}
graph
}
pub fn get_knn_graph_adj<T>(similarities: &MatRef<T>, k: usize) -> Mat<T>
where
T: BixverseFloat,
{
assert_symmetric_mat!(similarities);
let n = similarities.nrows();
let mut adjacency: Mat<T> = Mat::zeros(n, n);
let rows: Vec<Vec<(usize, T)>> = (0..n)
.into_par_iter()
.map(|i| {
let mut heap = BinaryHeap::with_capacity(k + 1);
for j in 0..n {
if i != j {
let sim = similarities[(i, j)];
heap.push((RevOrderedFloat(sim), j));
if heap.len() > k {
heap.pop(); }
}
}
heap.into_iter()
.map(|(ordered_sim, index)| (index, ordered_sim.0))
.collect()
})
.collect();
for (i, neighbors) in rows.iter().enumerate() {
for &(j, sim) in neighbors {
adjacency[(i, j)] = sim;
}
}
let two = T::from_f64(2.0).unwrap();
for i in 0..n {
for j in i + 1..n {
let val = (adjacency[(i, j)] + adjacency[(j, i)]) / two;
adjacency[(i, j)] = val;
adjacency[(j, i)] = val;
}
}
adjacency
}
pub fn adjacency_to_laplacian<T>(adjacency: &MatRef<T>, normalise: bool) -> Mat<T>
where
T: BixverseFloat,
{
assert_symmetric_mat!(adjacency);
let n = adjacency.nrows();
let degrees: Vec<T> = (0..n)
.map(|i| {
adjacency
.row(i)
.iter()
.copied()
.fold(T::zero(), |acc, x| acc + x)
})
.collect();
if !normalise {
let mut laplacian = adjacency.cloned();
for i in 0..n {
laplacian[(i, i)] = degrees[i] - adjacency[(i, i)];
for j in 0..n {
if i != j {
laplacian[(i, j)] = -adjacency[(i, j)];
}
}
}
laplacian
} else {
let inv_sqrt_d: Vec<T> = degrees
.iter()
.map(|&d| {
if d > T::from_f64(1e-10).unwrap() {
T::one() / d.sqrt()
} else {
T::zero()
}
})
.collect();
let mut laplacian = Mat::zeros(n, n);
for i in 0..n {
laplacian[(i, i)] = T::one();
for j in 0..n {
laplacian[(i, j)] -= inv_sqrt_d[i] * adjacency[(i, j)] * inv_sqrt_d[j];
}
}
laplacian
}
}
pub fn knn_to_sparse_graph<T>(knn: &[Vec<usize>], same_weight: bool) -> SparseGraph<T>
where
T: BixverseFloat + std::iter::Sum + BixverseNumeric,
{
let n_nodes = knn.len();
let mut edges = Vec::new();
for (i, neighbours) in knn.iter().enumerate() {
for &j in neighbours {
edges.push((i, j));
edges.push((j, i)); }
}
edges.sort_unstable();
let mut rows = Vec::new();
let mut cols = Vec::new();
let mut vals = Vec::new();
if !edges.is_empty() {
let (mut curr_r, mut curr_c) = edges[0];
let mut weight = T::one();
for &(r, c) in &edges[1..] {
if r == curr_r && c == curr_c {
weight += T::one();
} else {
rows.push(curr_r);
cols.push(curr_c);
vals.push(weight);
(curr_r, curr_c) = (r, c);
weight = T::one();
}
}
rows.push(curr_r);
cols.push(curr_c);
vals.push(weight);
}
let adjacency = coo_to_csr(
&rows.index_cast(),
&cols.index_cast(),
&vals,
(n_nodes, n_nodes),
);
let mut graph = SparseGraph::new(n_nodes, adjacency, false);
if same_weight {
for elem in graph.adjacency.data.iter_mut() {
*elem = T::one()
}
}
graph
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::math::sparse::CompressedSparseData2;
use faer::Mat;
#[test]
fn test_sparse_graph_properties() {
let data = vec![1.0, 1.0, 1.0, 1.0];
let indices = vec![1, 0, 2, 1]; let indptr = vec![0, 1, 3, 4]; let csr =
CompressedSparseData2::<f64, f64>::new_csr(&data, &indices, &indptr, None, (3, 3));
let graph = SparseGraph::new(3, csr, false);
assert_eq!(graph.get_node_degree(0), 1);
assert_eq!(graph.get_node_degree(1), 2);
assert_eq!(graph.total_weight(), 2.0);
let (n0_neighbors, n0_weights) = graph.get_neighbours(0);
assert_eq!(n0_neighbors, &[1]);
assert_eq!(n0_weights, &[1.0]);
}
#[test]
fn test_adjacency_to_laplacian() {
let mut adj: Mat<f64> = Mat::zeros(2, 2);
adj[(0, 1)] = 1.0;
adj[(1, 0)] = 1.0;
let lap = adjacency_to_laplacian(&adj.as_ref(), false);
assert!((lap[(0, 0)] - 1.0).abs() < 1e-6);
assert!((lap[(1, 1)] - 1.0).abs() < 1e-6);
assert!((lap[(0, 1)] - (-1.0)).abs() < 1e-6);
assert!((lap[(1, 0)] - (-1.0)).abs() < 1e-6);
}
#[test]
fn test_graph_from_strings() {
let nodes = vec!["A".to_string(), "B".to_string(), "C".to_string()];
let from = vec!["A".to_string(), "B".to_string()];
let to = vec!["B".to_string(), "C".to_string()];
let weights = vec![1.5, 2.5];
let graph = graph_from_strings(&nodes, &from, &to, Some(&weights), false);
assert_eq!(graph.node_count(), 3);
assert_eq!(graph.edge_count(), 2);
}
}