use annembed::fromhnsw::kgraph::KGraph;
use anyhow::Result;
use log::debug;
use num_traits::{FromPrimitive, Float};
use rayon::prelude::*;
pub fn mutual_knn<F: FromPrimitive + Float + std::fmt::UpperExp + Sync + Send + std::iter::Sum>(mut knn_graph: KGraph<F>, keep_n_edges: usize) -> Result<KGraph<F>> {
let mutual_nodes = knn_graph.get_neighbours()
.into_par_iter()
.enumerate()
.map(|(node, neighbours)| {
let mut mutual_neighbours = Vec::new();
for (neighbour_index, neighbour) in neighbours.iter().enumerate() {
if neighbour_index < keep_n_edges {
mutual_neighbours.push(neighbour.clone());
continue;
}
if knn_graph.get_neighbours()[neighbour.node].iter().any(|edge| edge.node == node) {
mutual_neighbours.push(neighbour.clone());
}
}
Ok(mutual_neighbours)
})
.collect::<Result<Vec<_>>>()?;
Ok(knn_graph)
}
pub fn intersect<F: FromPrimitive + Float + std::fmt::UpperExp + Sync + Send + std::iter::Sum>(mut base_knn_graph: KGraph<F>, intersecting_graph: KGraph<F>, keep_n_edges: usize) -> Result<KGraph<F>> {
let intersecting_nodes = intersecting_graph.get_neighbours();
let new_nodes = base_knn_graph.get_neighbours()
.into_par_iter()
.enumerate()
.map(|(node, neighbours)| {
let mut new_neighbours = Vec::new();
let node_id_in_intersecting_graph = intersecting_graph.get_idx_from_dataid(base_knn_graph.get_data_id_from_idx(node).unwrap()).unwrap();
for (neighbour_index, neighbour) in neighbours.iter().enumerate() {
if neighbour_index < keep_n_edges {
new_neighbours.push(neighbour.clone());
continue;
}
match intersecting_graph.get_idx_from_dataid(base_knn_graph.get_data_id_from_idx(neighbour.node).unwrap()) {
Some(intersect_node_index) => {
if intersecting_nodes[intersect_node_index].iter().any(|edge| edge.node == node_id_in_intersecting_graph) {
new_neighbours.push(neighbour.clone());
}
},
None => {
debug!("Node {} in base graph has a neighbour {} that is not in the intersecting graph", node, neighbour.node);
}
};
}
Ok(new_neighbours)
})
.collect::<Result<Vec<_>>>()?;
Ok(base_knn_graph)
}