legume_numeric/matrix/

knn_graph.rs

use crate::matrix::graph::WeightedGraph;
use crate::matrix::knn::all_pairs::knn_rows_l2;
use crate::matrix::knn::ivf::{knn_rows_ivf, IvfArgs, DEFAULT_N_PROBE};
use crate::matrix::knn::{EXACT_THRESHOLD, KNN_SEED};
use crate::matrix::knn_match::{ColumnDict, SearchScratch};

use indicatif::ParallelProgressIterator;
use log::info;
use nalgebra::DMatrix;
use nalgebra_sparse::CscMatrix;
use rayon::prelude::*;

const DEFAULT_BLOCK_SIZE: usize = 1000;

/// Up to this many points every row's neighbours come from the exact
/// all-pairs Gram kernel; beyond it from the inverted-file search. Both are
/// parallel and thread-count independent; the split is where `O(n²)` stops
/// being affordable.
pub const ALL_PAIRS_THRESHOLD: usize = 65_536;

pub struct KnnGraph {
    /// Symmetric CSC adjacency matrix (n_nodes x n_nodes)
    pub adjacency: CscMatrix<f32>,
    /// Sorted edge list (i < j), deduplicated
    pub edges: Vec<(usize, usize)>,
    /// Edge distances/weights, parallel to `edges`
    pub distances: Vec<f32>,
    /// Number of nodes
    pub n_nodes: usize,
}

/// Which input graph an edge of a [`KnnGraph::union_with`] came from.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EdgeSource {
    Primary,
    Secondary,
    /// Present in both inputs. Still a primary-graph edge for any consumer
    /// filtering on the primary relation, since the primary relation holds.
    Both,
}

/// How to reconcile two graphs' `distances` when merging.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DistanceMerge {
    /// Keep raw values. Correct only when both inputs measure the same thing.
    Raw,
    /// Replace each side's distances with its own within-source quantile rank
    /// in `[0, 1]` before merging. Use this whenever the inputs measure
    /// different things, so the merged column stays comparable across sources
    /// and monotone within each.
    SourceRank,
}

pub struct KnnGraphArgs {
    pub knn: usize,
    pub block_size: usize,
    /// If true, keep only reciprocal edges (i→j AND j→i).
    /// If false, keep union edges (i→j OR j→i), using min distance.
    pub reciprocal: bool,
}

impl KnnGraph {
    /// Build a KNN graph from column vectors.
    ///
    /// * `points` - transposed coordinate matrix (d x n), where each column is a point
    /// * `args` - KNN graph construction parameters
    pub fn from_columns(points: &DMatrix<f32>, args: KnnGraphArgs) -> anyhow::Result<KnnGraph> {
        Self::from_rows(&points.transpose(), args)
    }

    /// Build a KNN graph from row vectors (cells × features).
    ///
    /// * `data` - matrix (n x d), where each row is a point
    /// * `args` - KNN graph construction parameters
    pub fn from_rows(data: &DMatrix<f32>, args: KnnGraphArgs) -> anyhow::Result<KnnGraph> {
        let lists = neighbour_lists(data, &args)?;
        Self::from_neighbours(data.nrows(), &lists, args.reciprocal)
    }

    /// [`KnnGraph::from_columns_fuzzy`] over row vectors.
    pub fn from_rows_fuzzy(
        data: &DMatrix<f32>,
        args: KnnGraphArgs,
    ) -> anyhow::Result<(KnnGraph, Vec<f32>)> {
        let lists = neighbour_lists(data, &args)?;
        let graph = Self::from_neighbours(data.nrows(), &lists, args.reciprocal)?;
        let weights = umap_edge_weights(&graph.edges, &lists);
        Ok((graph, weights))
    }

    /// The kNN graph of the columns of `points`, and UMAP's fuzzy membership
    /// of each edge (parallel to `edges`), as umap-learn and uwot compute it:
    ///
    /// 1. each point's weights over its OWN `knn` neighbours only:
    ///    `exp(-(d - ρ) / σ)`, ρ its nearest distance, σ set so they sum to
    ///    `log2(knn + 1)` (UMAP counts the point itself among its
    ///    `n_neighbors`, so `knn` others is `n_neighbors = knn + 1`);
    /// 2. zero toward a point it did not list;
    /// 3. the fuzzy union of the two directions, `a + b - a·b`.
    ///
    /// [`KnnGraph::fuzzy_kernel_weights`] instead calibrates each point over
    /// every edge touching it after the union, so a point many others list
    /// gets a wider kernel and a one-sided edge a weight from both ends.
    pub fn from_columns_fuzzy(
        points: &DMatrix<f32>,
        args: KnnGraphArgs,
    ) -> anyhow::Result<(KnnGraph, Vec<f32>)> {
        Self::from_rows_fuzzy(&points.transpose(), args)
    }

    /// The graph implied by every point's directed neighbour list.
    fn from_neighbours(
        nn: usize,
        lists: &[NeighbourList],
        reciprocal: bool,
    ) -> anyhow::Result<KnnGraph> {
        let n_triplets: usize = lists.iter().map(|(nb, _)| nb.len()).sum();
        info!("{n_triplets} triplets by kNN matching");
        if n_triplets == 0 {
            return Err(anyhow::anyhow!("empty triplets"));
        }

        // Filtering and the sort ran silent, which on a large pair graph is a
        // stretch of nothing between the search bar and the next log line,
        // and reads as a hang.
        let filter_spin =
            crate::matrix::progress::new_spinner("{spinner} [{elapsed_precise}] {msg}")
                .with_message("filtering edges");
        let edges = edges_from_neighbours(lists, reciprocal);
        filter_spin.finish_and_clear();
        info!(
            "{} edges after {} matching",
            edges.len(),
            if reciprocal { "reciprocal" } else { "union" }
        );

        let (edge_pairs, distances): (Vec<_>, Vec<_>) = edges.into_iter().unzip();
        let adjacency = symmetric_adjacency(nn, &edge_pairs, &distances);

        Ok(KnnGraph {
            adjacency,
            edges: edge_pairs,
            distances,
            n_nodes: nn,
        })
    }

    /// Merge two graphs over the same nodes, keeping every pair exactly once
    /// and reporting which input each came from.
    ///
    /// Borrows both inputs: a caller that unions a spatial graph with an
    /// expression one generally still needs the spatial graph afterwards, as
    /// the topology for anything that reasons about physical adjacency.
    ///
    /// Neither input is assumed sorted, nor assumed to store `i < j`. Edge
    /// order is a constructor invariant here, not a type invariant, and a
    /// hand-built `KnnGraph` can violate both.
    ///
    /// `distances` after a union are NOT a metric. Under
    /// [`DistanceMerge::SourceRank`] they are within-source quantile ranks,
    /// which keeps them comparable across sources without pretending the two
    /// measurements are the same quantity. When an edge is in both inputs the
    /// smaller value wins, matching the `reciprocal: false` convention in
    /// `build_from_dict`.
    pub fn union_with(
        &self,
        other: &KnnGraph,
        policy: DistanceMerge,
    ) -> anyhow::Result<(KnnGraph, Vec<EdgeSource>)> {
        anyhow::ensure!(
            self.n_nodes == other.n_nodes,
            "cannot union graphs over different node counts: {} vs {}",
            self.n_nodes,
            other.n_nodes
        );
        let n_nodes = self.n_nodes;

        let (a_dist, b_dist) = match policy {
            DistanceMerge::Raw => (self.distances.clone(), other.distances.clone()),
            DistanceMerge::SourceRank => (
                within_source_rank(&self.distances),
                within_source_rank(&other.distances),
            ),
        };

        // Sort then fold, NOT a keyed map. At a few million edges an ordered
        // map costs a pointer-chasing O(log n) descent and a node allocation
        // per insert, all of it serial. One flat buffer, one parallel sort and
        // one linear scan replaces that, and measured an order of magnitude
        // faster. The allocation count drops too, though only the time was
        // measured.
        //
        // The canonical key is what makes this a set operation: a pair stored
        // one way round in one input and the other way round in the other must
        // land on the same key. The dedup has to finish before the COO below,
        // which SUMS duplicate entries rather than rejecting them.
        let canonical = |&(i, j): &(usize, usize)| if i <= j { (i, j) } else { (j, i) };
        // Source as a bitmask, so folding a run is an OR rather than a case
        // analysis: 1 = primary, 2 = secondary, 3 = both.
        let mut tagged: Vec<TaggedEdge> = Vec::with_capacity(self.edges.len() + other.edges.len());
        tagged.par_extend(
            self.edges
                .par_iter()
                .zip(a_dist.par_iter())
                .map(|(e, &d)| (canonical(e), d, 1u8)),
        );
        tagged.par_extend(
            other
                .edges
                .par_iter()
                .zip(b_dist.par_iter())
                .map(|(e, &d)| (canonical(e), d, 2u8)),
        );
        let folded = fold_tagged_edges(tagged);
        let mut edges = Vec::with_capacity(folded.len());
        let mut distances = Vec::with_capacity(folded.len());
        let mut source = Vec::with_capacity(folded.len());
        for (key, dist, mask) in folded {
            edges.push(key);
            distances.push(dist);
            source.push(match mask {
                1 => EdgeSource::Primary,
                2 => EdgeSource::Secondary,
                _ => EdgeSource::Both,
            });
        }

        // Derived state, so rebuild rather than merge.
        let adjacency = symmetric_adjacency(n_nodes, &edges, &distances);

        Ok((
            KnnGraph {
                adjacency,
                edges,
                distances,
                n_nodes,
            },
            source,
        ))
    }

    /// Get neighbors of a node from the CSC adjacency matrix
    pub fn neighbors(&self, node: usize) -> &[usize] {
        let offsets = self.adjacency.col_offsets();
        let start = offsets[node];
        let end = offsets[node + 1];
        &self.adjacency.row_indices()[start..end]
    }

    pub fn num_edges(&self) -> usize {
        self.edges.len()
    }

    pub fn num_nodes(&self) -> usize {
        self.n_nodes
    }

    /// Convert distances to similarity weights using an exponential kernel:
    /// `w = exp(-d / σ)` where σ = median distance.
    ///
    /// Returns weights parallel to `self.edges`, all in (0, 1].
    /// Consistent with the softmax(-d) pattern used in counterfactual
    /// inference (data_beans::alg) but with a global bandwidth.
    pub fn exp_kernel_weights(&self) -> Vec<f32> {
        if self.distances.is_empty() {
            return Vec::new();
        }
        let sigma = crate::matrix::utils::median(&self.distances);
        let sigma = if sigma <= 0.0 { 1.0 } else { sigma };
        info!("exp_kernel_weights: σ (median distance) = {:.4}", sigma);
        self.distances.iter().map(|&d| (-d / sigma).exp()).collect()
    }

    /// Adaptive-bandwidth kernel weights with local connectivity.
    ///
    /// Per-point sigma calibration (originated in t-SNE, van der Maaten
    /// & Hinton 2008) ensures every node has the same effective number
    /// of neighbors, preventing isolated singletons in sparse regions.
    /// The rho subtraction and fuzzy-union symmetrization follow UMAP
    /// (McInnes et al. 2018), matching the scanpy default for Leiden.
    ///
    /// Algorithm:
    /// 1. rho_i = distance to nearest neighbor (local connectivity)
    /// 2. sigma_i via binary search: sum_j exp(-(d_ij - rho_i)/sigma_i) = log2(k)
    /// 3. Directed weight: w(i→j) = exp(-(d_ij - rho_i) / sigma_i)
    /// 4. Symmetrize: w_sym = w(i→j) + w(j→i) - w(i→j) * w(j→i)
    ///
    /// Returns weights parallel to `self.edges`, all in (0, 1].
    pub fn fuzzy_kernel_weights(&self) -> Vec<f32> {
        if self.distances.is_empty() {
            return Vec::new();
        }

        let offsets = self.adjacency.col_offsets();
        let row_indices = self.adjacency.row_indices();
        let values = self.adjacency.values();

        // Step 1-2: compute rho and sigma per node — independent per node.
        let (rho, sigma): (Vec<f32>, Vec<f32>) = (0..self.n_nodes)
            .into_par_iter()
            .map(|i| {
                let start = offsets[i];
                let end = offsets[i + 1];
                let dists: Vec<f32> = (start..end).map(|idx| values[idx]).collect();
                if dists.is_empty() {
                    return (0.0_f32, 1.0_f32);
                }
                let rho_i = dists.iter().cloned().fold(f32::INFINITY, f32::min);
                let target = (dists.len() as f32).log2();
                let sigma_i = smooth_knn_sigma(&dists, rho_i, target);
                (rho_i, sigma_i)
            })
            .unzip();

        // Step 3-4: compute directed weights and symmetrize per edge —
        // independent per edge, only reads rho/sigma.
        self.edges
            .par_iter()
            .map(|&(i, j)| {
                let d_ij = self.edge_distance_directed(offsets, row_indices, values, i, j);
                let w_ij = directed_umap_weight(d_ij, rho[i], sigma[i]);
                let d_ji = self.edge_distance_directed(offsets, row_indices, values, j, i);
                let w_ji = directed_umap_weight(d_ji, rho[j], sigma[j]);
                // fuzzy union: P(at least one edge) = P(A) + P(B) - P(A)*P(B)
                w_ij + w_ji - w_ij * w_ji
            })
            .collect()
    }

    /// Look up the distance from node `from` to node `to` in the CSC adjacency.
    fn edge_distance_directed(
        &self,
        offsets: &[usize],
        row_indices: &[usize],
        values: &[f32],
        from: usize,
        to: usize,
    ) -> f32 {
        let start = offsets[from];
        let end = offsets[from + 1];
        for idx in start..end {
            if row_indices[idx] == to {
                return values[idx];
            }
        }
        f32::INFINITY
    }
}

impl WeightedGraph for KnnGraph {
    fn num_nodes(&self) -> usize {
        self.n_nodes
    }

    fn num_edges(&self) -> usize {
        self.edges.len()
    }

    fn neighbors_with_weight<'a>(
        &'a self,
        node: usize,
    ) -> Box<dyn Iterator<Item = (usize, f32)> + 'a> {
        let offsets = self.adjacency.col_offsets();
        let start = offsets[node];
        let end = offsets[node + 1];
        let rows = &self.adjacency.row_indices()[start..end];
        let vals = &self.adjacency.values()[start..end];
        Box::new(rows.iter().zip(vals.iter()).map(|(&i, &w)| (i, w)))
    }
}

/// Binary search for per-point sigma (UMAP's smooth_knn_dist).
///
/// Finds sigma such that: sum_j exp(-max(0, d_j - rho) / sigma) = target
fn smooth_knn_sigma(dists: &[f32], rho: f32, target: f32) -> f32 {
    const TOLERANCE: f32 = 1e-5;
    const MAX_ITER: usize = 64;

    let mean_dist: f32 = dists.iter().sum::<f32>() / dists.len().max(1) as f32;
    let min_sigma = 1e-3 * mean_dist;

    let mut lo = 0.0f32;
    let mut hi = f32::INFINITY;
    let mut mid = 1.0f32;

    for _ in 0..MAX_ITER {
        let mut psum = 0.0f32;
        for &d in dists {
            let gap = d - rho;
            if gap > 0.0 {
                psum += (-gap / mid).exp();
            } else {
                psum += 1.0;
            }
        }

        if (psum - target).abs() < TOLERANCE {
            break;
        }

        if psum > target {
            hi = mid;
            mid = (lo + hi) / 2.0;
        } else {
            lo = mid;
            if hi.is_infinite() {
                mid *= 2.0;
            } else {
                mid = (lo + hi) / 2.0;
            }
        }
    }

    mid.max(min_sigma)
}

/// Compute a single directed UMAP membership weight.
fn directed_umap_weight(d: f32, rho: f32, sigma: f32) -> f32 {
    if d.is_infinite() || sigma <= 0.0 {
        return 0.0;
    }
    let gap = d - rho;
    if gap <= 0.0 {
        1.0
    } else {
        (-gap / sigma).exp()
    }
}

////////////////////////
// Leiden integration //
////////////////////////

impl WeightedGraph for crate::leiden::Network {
    fn num_nodes(&self) -> usize {
        self.nodes()
    }

    fn num_edges(&self) -> usize {
        crate::leiden::Network::edge_count(self)
    }

    fn neighbors_with_weight<'a>(
        &'a self,
        node: usize,
    ) -> Box<dyn Iterator<Item = (usize, f32)> + 'a> {
        Box::new(self.neighbors(node).map(|(n, w)| (n, w as f32)))
    }
}

/// Convert a modularity resolution `gamma` to the CPM scale expected by the
/// Leiden crate, given the total undirected edge weight of the graph.
///
/// CPM resolution = `gamma / (2 * total_edge_weight)`. Guards against division
/// by zero for degenerate graphs by clamping the denominator to at least 1.
#[must_use]
pub fn modularity_to_cpm_resolution(modularity_gamma: f64, total_edge_weight: f64) -> f64 {
    modularity_gamma / (2.0 * total_edge_weight).max(1.0)
}

impl KnnGraph {
    /// Convert this KNN graph to a Leiden `Network` with modularity objective.
    ///
    /// Node weights = weighted degree, edge weights = fuzzy kernel weights.
    /// Returns `(network, total_edge_weight)`. Pass `total_edge_weight` to
    /// [`modularity_to_cpm_resolution`] to get a CPM-scale resolution.
    pub fn to_leiden_network(&self) -> (crate::leiden::Network, f64) {
        self.to_leiden_network_with(&self.fuzzy_kernel_weights())
    }

    /// [`KnnGraph::to_leiden_network`] with the edge weights given, parallel
    /// to `edges` (e.g. from [`KnnGraph::from_rows_fuzzy`]).
    pub fn to_leiden_network_with(&self, weights: &[f32]) -> (crate::leiden::Network, f64) {
        let n = self.n_nodes;

        let mut node_degree = vec![0.0f32; n];
        let mut n_edges = vec![0usize; n];
        let mut total_edge_weight = 0.0f64;
        for (&(i, j), &w) in self.edges.iter().zip(weights.iter()) {
            node_degree[i] += w;
            node_degree[j] += w;
            n_edges[i] += 1;
            n_edges[j] += 1;
            total_edge_weight += w as f64;
        }

        let mut network = crate::leiden::Network::with_nodes(&node_degree, &n_edges);
        for (&(i, j), &w) in self.edges.iter().zip(weights.iter()) {
            network.add_edge(i, j, w);
        }

        (network, total_edge_weight)
    }
}

/// Run Leiden clustering at a fixed (already-scaled) resolution.
///
/// Returns cluster labels as `Vec<usize>` (not necessarily contiguous).
pub fn run_leiden(
    network: &crate::leiden::Network,
    n: usize,
    resolution: f64,
    seed: Option<usize>,
) -> Vec<usize> {
    use crate::leiden::clustering::SimpleClustering;
    use crate::leiden::Clustering;
    use crate::leiden::Leiden;

    let mut leiden = Leiden::new(resolution, 0.01, seed);
    let mut clustering = SimpleClustering::init_different_clusters(n);

    for iter in 0..10 {
        let updated = leiden.iterate(network, &mut clustering);
        info!(
            "  Leiden iter {}: {} clusters{}",
            iter + 1,
            clustering.num_clusters(),
            if !updated { " (converged)" } else { "" }
        );
        if !updated {
            break;
        }
    }

    (0..n).map(|i| clustering.get(i)).collect()
}

/// Binary search on Leiden resolution to approximate `target_k` clusters.
///
/// `initial_resolution` should already be on the CPM scale
/// (i.e., `modularity_gamma / (2 * total_edge_weight)`).
/// Returns cluster labels (not necessarily contiguous).
pub fn tune_leiden_resolution(
    network: &crate::leiden::Network,
    n: usize,
    target_k: usize,
    initial_resolution: f64,
    seed: Option<usize>,
) -> Vec<usize> {
    let mut lo = 1e-6_f64;
    let mut hi = 10.0_f64;
    let mut best = run_leiden(network, n, initial_resolution, seed);
    let best_k = count_distinct(&best);

    info!(
        "  resolution={:.6e} → {} clusters (target {})",
        initial_resolution, best_k, target_k
    );

    if best_k == target_k {
        return best;
    }
    if best_k > target_k {
        hi = initial_resolution;
    } else {
        lo = initial_resolution;
    }

    let mut best_diff = best_k.abs_diff(target_k);

    for _ in 0..20 {
        let mid = (lo + hi) / 2.0;
        let result = run_leiden(network, n, mid, seed);
        let k = count_distinct(&result);
        info!("  resolution={:.6e} → {} clusters", mid, k);

        if k > target_k {
            hi = mid;
        } else {
            lo = mid;
        }

        let diff = k.abs_diff(target_k);
        if diff < best_diff {
            best = result;
            best_diff = diff;
        }

        if k == target_k || (hi - lo) / hi.max(1e-10) < 1e-4 {
            break;
        }
    }

    best
}

/// Count distinct values in a label vector.
fn count_distinct(labels: &[usize]) -> usize {
    let max = labels.iter().copied().max().unwrap_or(0);
    let mut seen = vec![false; max + 1];
    for &l in labels {
        seen[l] = true;
    }
    seen.iter().filter(|&&s| s).count()
}

/// Remap labels to contiguous 0..k.
pub fn compact_labels(labels: &mut [usize]) {
    let max = labels.iter().copied().max().unwrap_or(0);
    let mut mapping = vec![usize::MAX; max + 1];
    let mut next = 0usize;
    for l in labels.iter_mut() {
        if mapping[*l] == usize::MAX {
            mapping[*l] = next;
            next += 1;
        }
        *l = mapping[*l];
    }
}

fn create_jobs(ntot: usize, block_size: usize) -> Vec<(usize, usize)> {
    let block_size = if block_size == 0 {
        DEFAULT_BLOCK_SIZE
    } else {
        block_size
    };
    let nblock = ntot.div_ceil(block_size);
    (0..nblock)
        .map(|block| {
            let lb = block * block_size;
            let ub = ((block + 1) * block_size).min(ntot);
            (lb, ub)
        })
        .collect()
}

#[cfg(test)]
mod tests;

/// One point's neighbours: `(indices, distances)`, nearest first.
/// Every point's own neighbours (self excluded), exact up to
/// [`EXACT_THRESHOLD`] points.
fn neighbour_lists(data: &DMatrix<f32>, args: &KnnGraphArgs) -> anyhow::Result<Vec<NeighbourList>> {
    let nn = data.nrows();
    let n_neighbours = neighbours_per_point(args.knn, nn);
    Ok(if nn <= EXACT_THRESHOLD {
        let transposed = data.transpose();
        let points_vec = transposed.column_iter().collect::<Vec<_>>();
        let names = (0..nn).collect::<Vec<_>>();
        let dict = ColumnDict::from_dvector_views(points_vec, names);
        search_dict(&dict, nn, n_neighbours, args.block_size)?
    } else {
        search_rows(data, n_neighbours)
    })
}

/// UMAP's membership of a point in each of its neighbours, from its
/// distances to them (itself excluded): `exp(-(d - ρ) / σ)`, ρ the nearest
/// distance and σ set so they sum to `log2(len + 1)`, as UMAP counts the
/// point itself among its `n_neighbors`.
pub fn umap_memberships(dists: &[f32]) -> Vec<f32> {
    let Some(rho) = dists.iter().copied().reduce(f32::min) else {
        return Vec::new();
    };
    let target = ((dists.len() + 1) as f32).log2();
    let sigma = smooth_knn_sigma(dists, rho, target);
    dists
        .iter()
        .map(|&d| directed_umap_weight(d, rho, sigma))
        .collect()
}

/// UMAP's fuzzy membership of each of `edges` (canonical `i < j`) from the
/// points' own neighbour lists; see [`KnnGraph::from_columns_fuzzy`].
fn umap_edge_weights(edges: &[(usize, usize)], lists: &[NeighbourList]) -> Vec<f32> {
    // Each point's directed weights, parallel to its own list.
    let directed: Vec<Vec<f32>> = lists
        .par_iter()
        .map(|(_, dists)| umap_memberships(dists))
        .collect();
    let toward = |from: usize, to: usize| -> f32 {
        let (nb, _) = &lists[from];
        nb.iter()
            .position(|&j| j == to)
            .map_or(0.0, |at| directed[from][at])
    };
    edges
        .par_iter()
        .map(|&(i, j)| {
            let (a, b) = (toward(i, j), toward(j, i));
            a + b - a * b
        })
        .collect()
}

type NeighbourList = (Vec<usize>, Vec<f32>);

/// A canonical `(i, j)` key with a distance and a bitmask saying which
/// inputs listed it.
type TaggedEdge = ((usize, usize), f32, u8);

/// One flat buffer of canonical keys, one parallel sort and one linear fold:
/// a run of equal keys keeps the smallest distance and the OR of its masks.
/// What a keyed map over tens of millions of triplets would do with an
/// allocation and a contended insert per triplet and a lookup per edge, as a
/// set operation done in place. The result is sorted by key, and the dedup
/// has to finish before any CSC build, which SUMS duplicates.
fn fold_tagged_edges(mut tagged: Vec<TaggedEdge>) -> Vec<TaggedEdge> {
    tagged.par_sort_unstable_by_key(|&(key, _, _)| key);
    tagged.dedup_by(|cur, prev| {
        if cur.0 == prev.0 {
            prev.1 = prev.1.min(cur.1);
            prev.2 |= cur.2;
            true
        } else {
            false
        }
    });
    tagged
}

/// `search_others` returns exactly this many *other* neighbours (self
/// excluded): the request clamped to the available others, floored at 1.
fn neighbours_per_point(knn: usize, nn: usize) -> usize {
    knn.min(nn.saturating_sub(1)).max(1)
}

/// Every point's neighbours from the exact per-query scan of a
/// [`ColumnDict`], in parallel blocks — the arm for small point sets.
fn search_dict(
    dict: &ColumnDict<usize>,
    nn: usize,
    n_neighbours: usize,
    block_size: usize,
) -> anyhow::Result<Vec<NeighbourList>> {
    let jobs = create_jobs(nn, block_size);
    // Every bar draws through `crate::matrix::progress` so it shares the one
    // `MultiProgress` the log bridge writes above; a bar built straight
    // from indicatif registers with neither and corrupts the log.
    let search_bar =
        crate::matrix::progress::new_progress_bar(jobs.len() as u64).with_message("kNN blocks");
    let result: anyhow::Result<Vec<Vec<NeighbourList>>> = jobs
        .into_par_iter()
        .progress_with(search_bar.clone())
        .map(|(lb, ub)| {
            // One scratch per block, reused across the block's queries.
            let mut scratch = SearchScratch::default();
            (lb..ub)
                .map(|i| dict.search_others_reuse(&i, n_neighbours, &mut scratch))
                .collect()
        })
        .collect();
    // Clear BEFORE propagating: an error would otherwise leave the bar
    // ticking over the caller's error output.
    search_bar.finish_and_clear();
    Ok(result?.into_iter().flatten().collect())
}

/// Every row's neighbours among the other rows, exactly by the all-pairs Gram
/// kernel up to [`ALL_PAIRS_THRESHOLD`] rows and by the inverted-file search
/// beyond it.
fn search_rows(rows: &DMatrix<f32>, n_neighbours: usize) -> Vec<NeighbourList> {
    let nn = rows.nrows();
    let (indices, distances) = if nn <= ALL_PAIRS_THRESHOLD {
        info!("kNN by the exact all-pairs kernel over {nn} points");
        knn_rows_l2(rows, n_neighbours)
    } else {
        info!("kNN by the inverted-file search over {nn} points");
        knn_rows_ivf(
            rows,
            &IvfArgs {
                k: n_neighbours,
                n_lists: 0,
                n_probe: DEFAULT_N_PROBE,
                seed: KNN_SEED,
            },
        )
    };
    indices.into_iter().zip(distances).collect()
}

/// Undirected edges from directed neighbour lists, as `((i, j), distance)`
/// with `i < j`, sorted. `reciprocal` keeps a pair only when each point
/// listed the other; otherwise either direction suffices and the smaller of
/// the two distances is kept.
pub(crate) fn edges_from_neighbours(
    lists: &[NeighbourList],
    reciprocal: bool,
) -> Vec<((usize, usize), f32)> {
    // Direction as a bit so a run folds by OR: 1 = listed by the smaller
    // index, 2 = by the larger.
    let tagged: Vec<TaggedEdge> = lists
        .par_iter()
        .enumerate()
        .flat_map_iter(|(i, (nb, ds))| {
            nb.iter().zip(ds).map(move |(&j, &d)| {
                if i < j {
                    ((i, j), d, 1u8)
                } else {
                    ((j, i), d, 2u8)
                }
            })
        })
        .collect();
    fold_tagged_edges(tagged)
        .into_iter()
        .filter(|&(_, _, mask)| !reciprocal || mask == 3)
        .map(|(key, dist, _)| (key, dist))
        .collect()
}

/// The `n x n` symmetric adjacency implied by an undirected edge list, which
/// must be canonical (`i < j`), sorted and free of duplicates — what every
/// constructor here produces.
///
/// Each edge lands in BOTH endpoints' columns, exactly once per direction.
/// Built straight from the edge list — degrees, offsets, one fill — rather
/// than through a `CooMatrix`, whose conversion sorts every entry of the
/// whole matrix serially. Sorted input is what makes the fill enough: column
/// `c` receives its partners below `c` in ascending order as the edges
/// `(i, c)` pass, then its partners above `c` in ascending order from the run
/// of edges `(c, j)`.
pub fn symmetric_adjacency(
    n_nodes: usize,
    edges: &[(usize, usize)],
    distances: &[f32],
) -> CscMatrix<f32> {
    debug_assert!(
        edges.iter().all(|&(i, j)| i < j) && edges.windows(2).all(|w| w[0] < w[1]),
        "symmetric adjacency: edges must be canonical, sorted and unique"
    );
    let mut offsets = vec![0usize; n_nodes + 1];
    for &(i, j) in edges {
        offsets[i + 1] += 1;
        offsets[j + 1] += 1;
    }
    for c in 0..n_nodes {
        offsets[c + 1] += offsets[c];
    }
    let nnz = offsets[n_nodes];
    let mut row_indices = vec![0usize; nnz];
    let mut values = vec![0f32; nnz];
    let mut cursor = offsets[..n_nodes].to_vec();
    for (&(i, j), &v) in edges.iter().zip(distances) {
        row_indices[cursor[i]] = j;
        values[cursor[i]] = v;
        cursor[i] += 1;
        row_indices[cursor[j]] = i;
        values[cursor[j]] = v;
        cursor[j] += 1;
    }
    CscMatrix::try_from_csc_data(n_nodes, n_nodes, offsets, row_indices, values)
        .expect("symmetric adjacency: canonical sorted edges fill every column in order")
}

/// Each value replaced by its rank among the others, scaled to `[0, 1]`.
///
/// Ties take distinct adjacent ranks, which is harmless: the point is to put
/// two incomparable distance scales on one axis, not to be a faithful
/// empirical CDF.
fn within_source_rank(d: &[f32]) -> Vec<f32> {
    if d.len() <= 1 {
        return vec![0.0; d.len()];
    }
    let mut order: Vec<usize> = (0..d.len()).collect();
    // Ties break on the index, which is what makes this safe to sort in
    // parallel: an unstable parallel sort would otherwise put equal distances
    // in a run-dependent order, and these ranks are written out, so the file
    // would stop being reproducible. Grid-spaced coordinates tie constantly,
    // so this is the common case rather than a corner.
    //
    // `total_cmp` rather than `partial_cmp().unwrap_or(Equal)`: it is a total
    // order, so it needs no per-comparison branch and it gives NaN a definite
    // position instead of making it a wildcard that breaks transitivity.
    order.par_sort_unstable_by(|&a, &b| d[a].total_cmp(&d[b]).then(a.cmp(&b)));
    let mut out = vec![0.0f32; d.len()];
    let denom = (d.len() - 1) as f32;
    for (rank, &idx) in order.iter().enumerate() {
        out[idx] = rank as f32 / denom;
    }
    out
}