bixverse-sc 0.9.3

//! Implementation of the BBKNN approach for single cell batch correction via
//! Polański, et al., Bioinformatics, 2019. Compared to some of the other
//! methods, this one does not generate a batch-corrected embedding space,
//! but a batch-corrected kNN graph for subsequent usage in clustering and
//! visualisations.

use ann_search_rs::utils::dist::{Dist, parse_ann_dist};
use ann_search_rs::*;
use faer::MatRef;
use rayon::prelude::*;

use crate::prelude::*;
use crate::sc_batch_correction::batch_utils::process_batch_labels;
use crate::sc_batch_correction::harmony_core::to_row_major;
use bixverse_core::core::mat_struct::MatSliceView;
use bixverse_core::core::math::sparse::*;

///////////
// BBKNN //
///////////

// Constants needed for the UMAP stuff
const SMOOTH_K_TOLERANCE: f32 = 1e-5;
const MIN_K_DIST_SCALE: f32 = 1e-3;

/// Structure to store the BBKNN pamareters
#[derive(Clone, Debug)]
pub struct BbknnParams {
    /// How many neighbours per batch to identify
    pub neighbours_within_batch: usize,
    /// Mixing ratio between union (1.0) and intersection (0.0).
    pub set_op_mix_ratio: f32,
    /// UMAP connectivity computation parameter, how many nearest neighbours of
    /// each cell are assumed to be fully connected.
    pub local_connectivity: f32,
    /// Trim the neighbours of each cell to these many to connectivities. May
    /// help with population independence and improve the tidiness of clustering.
    pub trim: Option<usize>,
    /// [KnnParams] for the various approximate nearest neighbour searches
    /// in ann-search-rs
    pub knn_params: KnnParams,
}

///////////////////
// Helpers BBKNN //
///////////////////

/// Reject a batch that cannot supply the requested number of neighbours
///
/// Shared by the CPU and GPU searches. A cell of the batch itself is excluded
/// from its own neighbour list, so the batch needs one cell more than the
/// per-batch neighbour count.
///
/// ### Params
///
/// * `batch` - The batch label, for the error message.
/// * `n_cells` - Cells in this batch.
/// * `neighbours_within_batch` - Neighbours requested per batch.
///
/// ### Returns
///
/// `Ok(())`, or `BbknnBatchTooSmall` if the batch is too small.
pub fn check_batch_size(
    batch: usize,
    n_cells: usize,
    neighbours_within_batch: usize,
) -> Result<(), ScErrors> {
    let required = neighbours_within_batch + 1;
    if n_cells < required {
        return Err(ScErrors::BbknnBatchTooSmall {
            batch,
            n_cells,
            required,
        });
    }

    Ok(())
}

/// Generate batch balanced kNN graph
///
/// The function generates on a per batch basis an approximate nearest neighbour
/// search index and then runs all cells against this batch-specific index.
/// The resulting indices across all cells/batches are combined subsequntly.
///
/// ### Params
///
/// * `mat` - The embedding matrix to use. Usually PCA. cells = rows, features
///   = columns.
/// * `batch_labels` - Slice indicating which cell belongs to which batch.
/// * `knn_method` - Which KnnSearch to use.
/// * `bbknn_params` - `BbknnParams` with the parameters for the BBKNN batch
///   correction.
/// * `seed` - Random seed.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// A tuple with (nearest_neighbour_indices, nearest_neighbour_distances)
fn get_batch_balanced_knn(
    mat: MatRef<f32>,
    batch_labels: &[usize],
    knn_method: &KnnSearch,
    bbknn_params: &BbknnParams,
    seed: usize,
    verbose: usize,
) -> ScKnnResults {
    let verbosity = parse_verbosity_level(verbose);

    let n_cells = mat.nrows();

    // get unique batches
    let unique_batches: Vec<usize> = {
        let mut batches: Vec<_> = batch_labels.to_vec();
        batches.sort_unstable();
        batches.dedup();
        batches
    };

    let dist_metric: Dist = parse_ann_dist(&bbknn_params.knn_params.ann_dist).unwrap_or_default();

    let n_batches = unique_batches.len();

    if n_batches == 1 {
        return Err(ScErrors::NeedAtLeastTwoBatches { n_batches });
    }

    let mut all_indices = vec![vec![0; bbknn_params.neighbours_within_batch * n_batches]; n_cells];
    let mut all_distances =
        vec![vec![0.0; bbknn_params.neighbours_within_batch * n_batches]; n_cells];
    let col_indices: Vec<usize> = (0..mat.ncols()).collect();
    // row views of a column-major matrix are strided and miss the SIMD path
    let mat_rm_data = to_row_major(mat);
    let mat_rm = MatRef::from_row_major_slice(&mat_rm_data, n_cells, mat.ncols());

    for (batch_idx, &batch) in unique_batches.iter().enumerate() {
        if verbosity.normal_verbosity() {
            println!(
                "Processing batch {} / {}: {}",
                batch_idx + 1,
                n_batches,
                batch
            );
        }

        let batch_cell_indices: Vec<usize> = batch_labels
            .iter()
            .enumerate()
            .filter(|(_, b)| **b == batch)
            .map(|(i, _)| i)
            .collect();

        // Cells of this batch have to drop themselves, so the batch must hold
        // one more cell than the per-batch neighbour count. Without the guard
        // the fill loop below walks off the end of the neighbour row.
        check_batch_size(
            batch,
            batch_cell_indices.len(),
            bbknn_params.neighbours_within_batch,
        )?;

        let sub_matrix = MatSliceView::new(mat, &batch_cell_indices, &col_indices);
        let sub_matrix = sub_matrix.to_owned();

        let (neighbour_indices, _) = match knn_method {
            KnnSearch::Annoy => {
                // annoy path with updated functions
                let index = build_annoy_index(
                    sub_matrix.as_ref(),
                    &bbknn_params.knn_params.ann_dist,
                    bbknn_params.knn_params.n_tree,
                    seed,
                )?;
                query_annoy_index(
                    mat,
                    &index,
                    bbknn_params.neighbours_within_batch + 1,
                    bbknn_params.knn_params.search_budget,
                    false,
                    verbosity.detailed_verbosity(),
                )?
            }
            KnnSearch::Hnsw => {
                // hnsw path with updated functions
                let index = build_hnsw_index(
                    sub_matrix.as_ref(),
                    bbknn_params.knn_params.m,
                    bbknn_params.knn_params.ef_construction,
                    &bbknn_params.knn_params.ann_dist,
                    seed,
                    verbosity.detailed_verbosity(),
                );
                query_hnsw_index(
                    mat,
                    &index,
                    bbknn_params.neighbours_within_batch + 1,
                    bbknn_params.knn_params.ef_search,
                    true,
                    verbosity.detailed_verbosity(),
                )?
            }
            KnnSearch::NNDescent => {
                let index = build_nndescent_index(
                    sub_matrix.as_ref(),
                    &bbknn_params.knn_params.ann_dist,
                    bbknn_params.knn_params.delta,
                    bbknn_params.knn_params.diversify_prob,
                    None, // default to 30 here as the algorithm does
                    None,
                    None,
                    None,
                    seed,
                    verbosity.detailed_verbosity(),
                )?;

                query_nndescent_index(
                    mat,
                    &index,
                    bbknn_params.neighbours_within_batch + 1,
                    bbknn_params.knn_params.ef_budget,
                    false,
                    verbosity.detailed_verbosity(),
                )?
            }
            &KnnSearch::Exhaustive => {
                let index =
                    build_exhaustive_index(sub_matrix.as_ref(), &bbknn_params.knn_params.ann_dist);

                query_exhaustive_index(
                    mat,
                    &index,
                    bbknn_params.neighbours_within_batch + 1,
                    false,
                    verbosity.detailed_verbosity(),
                )?
            }
            &KnnSearch::Ivf => {
                let index = build_ivf_index(
                    sub_matrix.as_ref(),
                    bbknn_params.knn_params.n_list,
                    None,
                    &bbknn_params.knn_params.ann_dist,
                    seed,
                    verbosity.detailed_verbosity(),
                )?;

                query_ivf_index(
                    mat,
                    &index,
                    bbknn_params.neighbours_within_batch + 1,
                    bbknn_params.knn_params.n_probe,
                    false,
                    verbosity.detailed_verbosity(),
                )?
            }
            &KnnSearch::KmKnn => {
                let index = build_kmknn_index(
                    sub_matrix.as_ref(),
                    &bbknn_params.knn_params.ann_dist,
                    bbknn_params.knn_params.n_list,
                    None,
                    seed,
                    verbosity.detailed_verbosity(),
                )?;

                query_kmknn_index(
                    mat,
                    &index,
                    bbknn_params.neighbours_within_batch + 1,
                    false,
                    verbosity.detailed_verbosity(),
                )?
            }
        };

        let col_start = batch_idx * bbknn_params.neighbours_within_batch;

        let k_within = bbknn_params.neighbours_within_batch;
        all_indices
            .par_iter_mut()
            .zip(all_distances.par_iter_mut())
            .zip(neighbour_indices.par_iter())
            .enumerate()
            .for_each(|(cell_idx, ((idx_row, dist_row), neighbours))| {
                let mut added = 0;
                let mut k_idx = 0;

                while added < k_within {
                    let global_idx = batch_cell_indices[neighbours[k_idx]];

                    if global_idx != cell_idx {
                        let dist = compute_distance_knn(
                            mat_rm.row(cell_idx),
                            mat_rm.row(global_idx),
                            &dist_metric,
                        );

                        idx_row[col_start + added] = global_idx;
                        dist_row[col_start + added] = dist;
                        added += 1;
                    }

                    k_idx += 1;
                }
            });
    }

    Ok((all_indices, all_distances))
}

/// Helper to calculate smooth kNN distances
///
/// ### Params
///
/// * `dist` - Slice of distance vectors from the kNN
/// * `k` - Number of neighbours.
/// * `local_connectivity` - Determines the minimum distance threshold (rho) by
///   interpolating between the nearest non-zero neighbors. If it's 1.5, you'd
///   interpolate between the 1st and 2nd neighbor distances.
/// * `smook_k_tol` - Tolerance parameter that controls:
///   - Whether to apply interpolation when computing rho
///   - When to stop the binary search (when |psum - target| < smook_k_tol)
/// * `min_k_dist_scale` - Minimum scaling factor for sigma to prevent it from
///   becoming too small
///
/// ### Returns
///
/// Tuple of (rho, sigma)
pub fn smooth_knn_dist(
    dist: &[Vec<f32>],
    k: f32,
    local_connectivity: f32,
    smook_k_tol: f32,
    min_k_dist_scale: f32,
) -> (Vec<f32>, Vec<f32>) {
    let n = dist.len();
    let n_neighbours = dist[0].len();
    let target = k.log2();

    let mean_dist = dist.iter().flat_map(|d| d.iter()).sum::<f32>() / (n * n_neighbours) as f32;

    let res: Vec<(f32, f32)> = dist
        .par_iter()
        .map(|dist_i| {
            let mut rho = 0.0_f32;

            let non_zero_dist: Vec<f32> = dist_i.iter().filter(|&&d| d > 0.0).copied().collect();

            if non_zero_dist.len() >= local_connectivity as usize {
                let index = local_connectivity.floor() as usize;
                let interpolation = local_connectivity - local_connectivity.floor();

                if index > 0 {
                    rho = non_zero_dist[index - 1];
                    if interpolation > smook_k_tol {
                        rho += interpolation * (non_zero_dist[index] - non_zero_dist[index - 1]);
                    }
                } else {
                    rho = interpolation * non_zero_dist[0];
                }
            } else if !non_zero_dist.is_empty() {
                rho = *non_zero_dist
                    .iter()
                    .max_by(|a, b| a.partial_cmp(b).unwrap())
                    .unwrap();
            }

            // binary search for sigma
            let mut lo = 0.0_f32;
            let mut hi = f32::INFINITY;
            let mut mid = 1.0_f32;

            for _ in 0..64 {
                let mut psum = 0.0_f32;

                for j in 1..n_neighbours {
                    let d = dist_i[j] - rho;
                    if d > 0.0_f32 {
                        psum += (-d / mid).exp();
                    } else {
                        psum += 1.0;
                    }
                }

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

                if psum > target {
                    hi = mid;
                    mid = (lo + hi) / 2.0;
                } else {
                    lo = mid;
                    if hi == f32::INFINITY {
                        mid *= 2.0;
                    } else {
                        mid = (lo + hi) / 2.0;
                    }
                }
            }

            let mut sigma = mid;

            if rho > 0.0 {
                let mean_i = dist_i.iter().sum::<f32>() / n_neighbours as f32;
                if sigma < min_k_dist_scale * mean_i {
                    sigma = min_k_dist_scale * mean_i;
                }
            } else if sigma < min_k_dist_scale * mean_dist {
                sigma = min_k_dist_scale * mean_dist
            }

            (sigma, rho)
        })
        .collect();

    let sigmas = res.iter().map(|r| r.0).collect();
    let rhos = res.iter().map(|r| r.1).collect();

    (sigmas, rhos)
}

/// Sorts the kNN data by distance
///
/// ### Params
///
/// * `knn_indices` - The indices of the k-nearest neighbours.
/// * `knn_dists` - The distances of the k-nearest neighbours.
///
/// ### Side effect
///
/// Resorts the indices and distances of the mutable inputs.
fn sort_knn_by_distance(knn_indices: &mut [Vec<usize>], knn_dists: &mut [Vec<f32>]) {
    knn_indices
        .par_iter_mut()
        .zip(knn_dists.par_iter_mut())
        .for_each_init(Vec::new, |pairs: &mut Vec<(f32, usize)>, (idx, dist)| {
            pairs.clear();
            pairs.extend(dist.iter().copied().zip(idx.iter().copied()));

            pairs.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());

            for (j, &(d, i)) in pairs.iter().enumerate() {
                dist[j] = d;
                idx[j] = i;
            }
        });
}

/// Compute membership strengths
///
/// Converts kNN distances into fuzzy membership weights using UMAP-style
/// exponential decay. Distances within rho get weight 1.0, distances beyond
/// decay as exp(-(dist - rho) / sigma).
///
/// ### Params
///
/// * `knn_indices` - Indices of k-nearest neighbours for each sample.
/// * `knn_dists` - Distances to k-nearest neighbours for each sample.
/// * `sigmas` - Bandwidth parameters for exponential decay per sample.
/// * `rhos` - Distance thresholds for hard-core similarity per sample.
///
/// ### Returns
///
/// Tuple of (rows, cols, vals) in COO format representing the weighted
/// connectivity graph
fn compute_membership_strengths(
    knn_indices: &[Vec<usize>],
    knn_dists: &[Vec<f32>],
    sigmas: &[f32],
    rhos: &[f32],
) -> (Vec<usize>, Vec<usize>, Vec<f32>) {
    let n_samples = knn_indices.len();
    let n_neighbours = knn_indices[0].len();

    let total = n_samples * n_neighbours;
    let mut rows = vec![0usize; total];
    let mut cols = vec![0usize; total];
    let mut vals = vec![0.0f32; total];

    rows.par_chunks_mut(n_neighbours)
        .zip(cols.par_chunks_mut(n_neighbours))
        .zip(vals.par_chunks_mut(n_neighbours))
        .enumerate()
        .for_each(|(i, ((row_r, row_c), row_v))| {
            for j in 0..n_neighbours {
                let neighbor = knn_indices[i][j];

                row_v[j] = if neighbor == i {
                    0.0
                } else if knn_dists[i][j] - rhos[i] <= 0.0 || sigmas[i] == 0.0 {
                    1.0
                } else {
                    (-(knn_dists[i][j] - rhos[i]) / sigmas[i]).exp()
                };
                row_r[j] = i;
                row_c[j] = neighbor;
            }
        });

    (rows, cols, vals)
}

/// Apply set operations to symmetrize connectivity graph
///
/// Applies fuzzy set union/intersection operations to create a symmetrized
/// connectivity matrix. Computes:
///
/// `ratio * (A + A^T - A.*A^T) + (1-ratio) * (A.*A^T)`
///
/// where ratio = 1.0 is pure union and ratio = 0.0 is pure intersection.
///
/// ### Params
///
/// * `connectivities` - Asymmetric connectivity matrix in CSR format
/// * `set_op_mix_ratio` - Mixing ratio between union (1.0) and intersection (0.0)
///
/// ### Returns
///
/// Symmetrised connectivity matrix in CSR format
fn apply_set_operations(
    connectivities: CompressedSparseData2<f32>,
    set_op_mix_ratio: f32,
) -> Result<CompressedSparseData2<f32>, ScErrors> {
    let (nrow, ncol) = connectivities.shape;
    let transpose = {
        let csc_transpose = connectivities.transform(); // A^T in CSC
        // Manually reinterpret CSC of A^T as CSR
        let mut coo_rows = Vec::new();
        let mut coo_cols = Vec::new();
        let mut coo_vals = Vec::new();

        for col in 0..csc_transpose.indptr.len() - 1 {
            for idx in csc_transpose.indptr[col]..csc_transpose.indptr[col + 1] {
                let idx = idx as usize;

                let row = csc_transpose.indices[idx];
                // For A^T: swap row/col to get proper CSR representation
                coo_rows.push(col);
                coo_cols.push(row);
                coo_vals.push(csc_transpose.data[idx]);
            }
        }
        coo_to_csr(&coo_rows.index_cast(), &coo_cols, &coo_vals, (ncol, nrow))
    };

    // Element-wise multiply: A .* A^T
    let prod = sparse_multiply_elementwise_csr(&connectivities, &transpose)?;

    // set_op_mix_ratio * (A + A^T - A.*A^T) + (1 - set_op_mix_ratio) * (A.*A^T)
    let union_part = sparse_add_csr(&connectivities, &transpose)?;
    let union_part = sparse_subtract_csr(&union_part, &prod)?;
    let union_part = sparse_scalar_multiply_csr(&union_part, set_op_mix_ratio);

    let intersect_part = sparse_scalar_multiply_csr(&prod, 1.0 - set_op_mix_ratio);

    let res = sparse_add_csr(&union_part, &intersect_part)?;

    Ok(eliminate_zeros_csr(res))
}

/// Trim weak edges from connectivity graph
///
/// For each node, keeps only the top-k strongest connections by zeroing out
/// weights below the k-th strongest edge. Applies trimming in both row and
/// column directions to ensure symmetry.
///
/// ### Params
///
/// * `connectivities` - Connectivity matrix in CSR format
/// * `trim` - Number of strongest connections to keep per node
///
/// ### Returns
///
/// Trimmed connectivity matrix in CSR format
fn trim_graph(
    mut connectivities: CompressedSparseData2<f32>,
    trim: usize,
) -> CompressedSparseData2<f32> {
    let n = connectivities.shape.0;
    let mut thresholds = vec![0.0f32; n];

    // compute thresholds

    let (indptr, data) = (&connectivities.indptr, &connectivities.data);
    thresholds.par_iter_mut().enumerate().for_each_init(
        Vec::new,
        |scratch: &mut Vec<f32>, (i, thr)| {
            let row_data = &data[indptr[i] as usize..indptr[i + 1] as usize];

            if row_data.len() <= trim {
                return;
            }

            scratch.clear();
            scratch.extend_from_slice(row_data);
            let (_, nth, _) =
                scratch.select_nth_unstable_by(trim - 1, |a, b| b.partial_cmp(a).unwrap());
            *thr = *nth;
        },
    );

    // Apply trimming twice (row then column)
    for _ in 0..2 {
        for i in 0..n {
            let row_start = connectivities.indptr[i] as usize;
            let row_end = connectivities.indptr[i + 1] as usize;

            for j in row_start..row_end {
                if connectivities.data[j] < thresholds[i] {
                    connectivities.data[j] = 0.0;
                }
            }
        }

        connectivities = eliminate_zeros_csr(connectivities);
        connectivities = connectivities.transform();
    }

    connectivities
}

////////////////
// Main BBKNN //
////////////////

/// Turn batch-balanced kNN results into the BBKNN output graphs
///
/// Everything after the neighbour search: sorting, the UMAP connectivity
/// estimate, the fuzzy set operations and the optional trimming. Shared by the
/// CPU [`bbknn`] and the GPU `bbknn_gpu`, so the two cannot drift apart.
///
/// ### Params
///
/// * `knn_indices` - Indices of the batch-balanced neighbours. Resorted by
///   distance in place.
/// * `knn_dists` - Distances of the batch-balanced neighbours. Resorted in
///   place alongside the indices.
/// * `n_obs` - Number of cells.
/// * `set_op_mix_ratio` - Mixing ratio between union (1.0) and intersection
///   (0.0).
/// * `local_connectivity` - How many nearest neighbours of each cell are
///   assumed to be fully connected.
/// * `trim` - Trim the neighbours of each cell to these many connectivities.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// A tuple of two CompressedSparseData2 with `(distances, connectivities)`.
pub fn bbknn_graph_from_knn(
    knn_indices: &mut [Vec<usize>],
    knn_dists: &mut [Vec<f32>],
    n_obs: usize,
    set_op_mix_ratio: f32,
    local_connectivity: f32,
    trim: Option<usize>,
    verbose: usize,
) -> Result<(CompressedSparseData2<f32>, CompressedSparseData2<f32>), ScErrors> {
    let verbosity = parse_verbosity_level(verbose);

    sort_knn_by_distance(knn_indices, knn_dists);

    if verbosity.normal_verbosity() {
        println!("BBKNN: Calculating UMAP-based connectivities and removing weak connections.")
    }

    let n_neighbours = knn_indices[0].len();
    let (sigmas, rhos) = smooth_knn_dist(
        knn_dists,
        n_neighbours as f32,
        local_connectivity,
        SMOOTH_K_TOLERANCE,
        MIN_K_DIST_SCALE,
    );

    let (rows, cols, vals) = compute_membership_strengths(knn_indices, knn_dists, &sigmas, &rhos);

    let mut connectivities = coo_to_csr(
        &rows.index_cast(),
        &cols.index_cast(),
        &vals,
        (n_obs, n_obs),
    );

    connectivities = apply_set_operations(connectivities, set_op_mix_ratio)?;

    if verbosity.normal_verbosity() {
        println!("BBKNN: Finalising data.")
    }

    let dist = knn_to_sparse_dist(knn_indices, knn_dists, n_obs);

    if let Some(trim_val) = trim
        && trim_val > 0
    {
        connectivities = trim_graph(connectivities, trim_val);
    }

    Ok((dist, connectivities))
}

/// Run batch-balanced KNN
///
/// This is the main function that implements the BBKNN logic into Rust.
///
/// ### Params
///
/// * `mat` - The embedding matrix to use. Usually PCA. cells = rows, features
///   = columns.
/// * `batch_labels` - Slice indicating which cell belongs to which batch.
/// * `bbknn_params` - `BbknnParams` with the parameters for the BBKNN batch
///   correction.
/// * `seed` - Random seed.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// A tuple of two CompressedSparseData2 with `(indices, connectivities)`.
pub fn bbknn(
    mat: MatRef<f32>,
    batch_labels: &[usize],
    bbknn_params: &BbknnParams,
    seed: usize,
    verbose: usize,
) -> Result<(CompressedSparseData2<f32>, CompressedSparseData2<f32>), ScErrors> {
    let verbosity = parse_verbosity_level(verbose);
    let (_, n_batches) = process_batch_labels(batch_labels);

    if n_batches == 1 {
        return Err(ScErrors::NeedAtLeastTwoBatches { n_batches });
    }

    // parse and worst case KmkNnn
    let knn_method = parse_knn_method(&bbknn_params.knn_params.knn_method).unwrap_or_default();

    if verbosity.normal_verbosity() {
        println!("BBKNN: generating the batch balanced kNN values.")
    }

    let (mut knn_indices, mut knn_dists) =
        get_batch_balanced_knn(mat, batch_labels, &knn_method, bbknn_params, seed, verbose)?;

    bbknn_graph_from_knn(
        &mut knn_indices,
        &mut knn_dists,
        mat.nrows(),
        bbknn_params.set_op_mix_ratio,
        bbknn_params.local_connectivity,
        bbknn_params.trim,
        verbose,
    )
}