bixverse-rs 0.3.6

//! Scrublet doublet detection for single cell data.
//!
//! Implements the method from Wolock, et al., Cell Systems, 2019. Synthetic
//! doublets are generated by averaging random cell pairs, embedded alongside
//! observed cells via PCA, and scored by the proportion of simulated neighbours
//! in a kNN graph. Threshold selection uses Otsu's method on the observed score
//! distribution.

use faer::Mat;
use half::f16;
use rayon::prelude::*;
use rustc_hash::{FxHashMap, FxHashSet};
use std::time::Instant;

use crate::prelude::*;
use crate::single_cell::sc_processing::utils_doublets::*;

///////////
// Types //
///////////

/// Scrublet doublet scores for observed and simulated cells.
///
/// ### Fields
///
/// * `0` - Doublet scores for observed cells.
/// * `1` - Standard errors for observed cell scores.
/// * `2` - Doublet scores for simulated doublets.
/// * `3` - Standard errors for simulated doublet scores.
type ScrubletDoubletScores = (Vec<f32>, Vec<f32>, Vec<f32>, Vec<f32>);

/// Final return type from `run_scrublet`.
///
/// ### Fields
///
/// * `0` - The `ScrubletResult` with predictions and statistics.
/// * `1` - Optional combined PCA embedding (observed + simulated).
/// * `2` - Optional first parent indices for each simulated doublet.
/// * `3` - Optional second parent indices for each simulated doublet.
pub type FinalScrubletRes = (
    ScrubletResult,
    Option<Mat<f32>>,
    Option<Vec<usize>>,
    Option<Vec<usize>>,
);

////////////////////////
// Params and results //
////////////////////////

/// Parameters for the Scrublet algorithm.
///
/// Controls preprocessing, doublet simulation, PCA, kNN construction and
/// threshold selection.
#[derive(Clone, Debug)]
pub struct ScrubletParams {
    // -- processing --
    /// Whether to log-transform counts after normalisation.
    pub log_transform: bool,
    /// Whether to mean-centre genes before PCA.
    pub mean_center: bool,
    /// Whether to scale genes to unit variance before PCA.
    pub normalise_variance: bool,
    /// Optional target library size. Defaults to the mean HVG library size.
    pub target_size: Option<f32>,

    // -- hvg --
    /// Percentile threshold for HVG selection.
    pub min_gene_var_pctl: f32,
    /// HVG method: `"vst"`, `"mvb"`, or `"dispersion"`.
    pub hvg_method: String,
    /// Loess span for VST fitting.
    pub loess_span: f64,
    /// Optional clip max for variance stabilisation.
    pub clip_max: Option<f32>,
    /// Binning strategy
    pub binning_strategy: String,
    /// Number of bins for hvg
    pub n_bins: usize,

    // -- pca --
    /// Number of principal components.
    pub no_pcs: usize,
    /// Whether to use randomised SVD.
    pub random_svd: bool,

    // -- scrublet --
    /// Ratio of simulated doublets to observed cells.
    pub sim_doublet_ratio: f32,
    /// Expected doublet rate (typically 0.05-0.10).
    pub expected_doublet_rate: f32,
    /// Uncertainty in the expected doublet rate.
    pub stdev_doublet_rate: f32,
    /// Number of histogram bins for Otsu threshold detection.
    pub n_bins_hist: usize,
    /// Optional manual threshold. If `None`, Otsu's method is used.
    pub manual_threshold: Option<f32>,

    // -- knn --
    /// [KnnParams] for the various approximate nearest neighbour searches
    /// in ann-search-rs
    pub knn_params: KnnParams,
}

/// Results from the Scrublet algorithm.
///
/// Contains doublet predictions, scores, and diagnostic statistics.
#[derive(Clone, Debug)]
pub struct ScrubletResult {
    /// Per-cell doublet predictions (`true` = doublet).
    pub predicted_doublets: Vec<bool>,
    /// Doublet score per observed cell.
    pub doublet_scores_obs: Vec<f32>,
    /// Doublet score per simulated doublet.
    pub doublet_scores_sim: Vec<f32>,
    /// Standard error of each observed cell's score.
    pub doublet_errors_obs: Vec<f32>,
    /// Z-score per observed cell: `(score - threshold) / error`.
    pub z_scores: Vec<f32>,
    /// Threshold used for doublet calling.
    pub threshold: f32,
    /// Fraction of observed cells called as doublets.
    pub detected_doublet_rate: f32,
    /// Fraction of simulated doublets above the threshold.
    pub detectable_doublet_fraction: f32,
    /// Estimated overall doublet rate.
    pub overall_doublet_rate: f32,
}

/////////////////////////
// Cell chunk addition //
/////////////////////////

impl CsrCellChunk {
    /// Combine two cells into a synthetic doublet for doublet simulation.
    ///
    /// Filters both cells to HVG genes, sums their raw counts, normalises
    /// to `target_size` using the combined HVG library size, and optionally
    /// log-transforms.
    ///
    /// ### Params
    ///
    /// * `cell1` - First parent cell.
    /// * `cell2` - Second parent cell.
    /// * `hvg_indices` - Set of HVG gene indices to retain.
    /// * `hvg_lib_size_combined` - Sum of parent HVG library sizes.
    /// * `target_size` - Normalisation target (e.g. mean HVG library size).
    /// * `log_transform` - Whether to apply `ln(1 + x)` after normalisation.
    /// * `doublet_index` - Index to assign to the new chunk.
    ///
    /// ### Returns
    ///
    /// A new `CsrCellChunk` representing the synthetic doublet.
    #[allow(clippy::too_many_arguments)]
    pub fn add_cells_scrublet(
        cell1: &CsrCellChunk,
        cell2: &CsrCellChunk,
        hvg_indices: &FxHashSet<usize>,
        hvg_lib_size_combined: usize,
        target_size: f32,
        log_transform: bool,
        doublet_index: usize,
    ) -> Self {
        let mut gene_counts: FxHashMap<u32, u32> = FxHashMap::default();

        for i in 0..cell1.indices.len() {
            if hvg_indices.contains(&(cell1.indices[i] as usize)) {
                *gene_counts.entry(cell1.indices[i]).or_insert(0) += cell1.data_raw.get(i);
            }
        }
        for i in 0..cell2.indices.len() {
            if hvg_indices.contains(&(cell2.indices[i] as usize)) {
                *gene_counts.entry(cell2.indices[i]).or_insert(0) += cell2.data_raw.get(i);
            }
        }

        let mut gene_vec: Vec<(u32, u32)> = gene_counts.into_iter().collect();
        gene_vec.sort_unstable_by_key(|&(gene, _)| gene);

        let mut data_raw = Vec::with_capacity(gene_vec.len());
        let mut data_norm = Vec::with_capacity(gene_vec.len());
        let mut indices = Vec::with_capacity(gene_vec.len());

        let norm_factor = target_size / hvg_lib_size_combined as f32;

        for (gene, count) in gene_vec {
            data_raw.push(count.min(u16::MAX as u32) as u16);

            let normalised = if log_transform {
                (count as f32 * norm_factor).ln_1p()
            } else {
                count as f32 * norm_factor
            };
            data_norm.push(F16::from(f16::from_f32(
                normalised.clamp(-65504.0, 65504.0),
            )));
            indices.push(gene);
        }

        Self {
            data_raw: RawCounts::U16(data_raw),
            data_norm,
            library_size: hvg_lib_size_combined,
            indices,
            original_index: doublet_index,
            to_keep: true,
        }
    }
}

///////////////////////
// Scrublet-specific //
///////////////////////

////////////////////
// Main structure //
////////////////////

/// Scrublet doublet detection.
///
/// Simulates doublets via random cell pairing, embeds observed and simulated
/// cells together via sparse PCA, builds a kNN graph, and scores each
/// observed cell by the Bayesian proportion of simulated neighbours.
#[derive(Clone, Debug)]
pub struct Scrublet {
    /// Path to the gene-based binary file (CSC format).
    f_path_gene: String,
    /// Path to the cell-based binary file (CSR format).
    f_path_cell: String,
    /// Algorithm parameters.
    params: ScrubletParams,
    /// Number of observed cells.
    n_cells: usize,
    /// Number of simulated doublets.
    n_cells_sim: usize,
    /// Cell indices included in this analysis.
    cells_to_keep: Vec<usize>,
    /// Per-cell library sizes computed over HVG genes only.
    hvg_library_sizes: Vec<usize>,
}

impl Scrublet {
    /// Create a new Scrublet instance.
    ///
    /// ### Params
    ///
    /// * `f_path_gene` - Path to the gene-based binary file (CSC).
    /// * `f_path_cells` - Path to the cell-based binary file (CSR).
    /// * `params` - Scrublet parameters.
    /// * `cell_indices` - Cell indices to include in the analysis.
    ///
    /// ### Returns
    ///
    /// Initialised `Scrublet`.
    pub fn new(
        f_path_gene: &str,
        f_path_cells: &str,
        params: ScrubletParams,
        cell_indices: &[usize],
    ) -> Self {
        Scrublet {
            f_path_gene: f_path_gene.to_string(),
            f_path_cell: f_path_cells.to_string(),
            params,
            n_cells: cell_indices.len(),
            n_cells_sim: 0,
            cells_to_keep: cell_indices.to_vec(),
            hvg_library_sizes: Vec::new(),
        }
    }

    /// Run the full Scrublet pipeline.
    ///
    /// ### Params
    ///
    /// * `streaming` - Stream HVG computation to reduce memory pressure.
    /// * `seed` - Seed for reproducibility.
    /// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for
    ///   detailed verbosity.
    /// * `return_combined_pca` - Whether to return the combined PCA embedding.
    /// * `return_pairs` - Whether to return the parent indices of each doublet.
    ///
    /// ### Returns
    ///
    /// `FinalScrubletRes` containing predictions, optional PCA and pair data.
    pub fn run_scrublet(
        &mut self,
        streaming: bool,
        seed: usize,
        verbose: usize,
        return_combined_pca: bool,
        return_pairs: bool,
    ) -> Result<FinalScrubletRes, BixverseErrors> {
        let verbosity = parse_verbosity_level(verbose);

        let hvg_opts = HvgOpts {
            method: self.params.hvg_method.clone(),
            loess_span: self.params.loess_span as f32,
            clip_max: self.params.clip_max,
            min_gene_var_pctl: self.params.min_gene_var_pctl,
            binning_strategy: self.params.binning_strategy.clone(),
            n_bins: self.params.n_bins,
        };
        let pca_opts = PcaOpts {
            log_transform: self.params.log_transform,
            mean_center: self.params.mean_center,
            normalise_variance: self.params.normalise_variance,
            no_pcs: self.params.no_pcs,
            random_svd: self.params.random_svd,
        };

        if verbosity.normal_verbosity() {
            println!("Identifying highly variable genes...");
        }
        let start_all = Instant::now();
        let start_hvg = Instant::now();

        let hvg_genes = select_hvg(
            &self.f_path_gene,
            &self.cells_to_keep,
            &hvg_opts,
            streaming,
            verbose,
        )?;

        if verbosity.normal_verbosity() {
            println!(
                "Using {} highly variable genes. Done in {:.2?}",
                hvg_genes.len(),
                start_hvg.elapsed()
            );
        }

        self.hvg_library_sizes =
            compute_hvg_library_sizes(&self.f_path_cell, &self.cells_to_keep, &hvg_genes)?;
        let target_size = resolve_target_size(self.params.target_size, &self.hvg_library_sizes);

        // Simulate doublets
        if verbosity.normal_verbosity() {
            println!("Simulating doublets...");
        }
        let start_sim = Instant::now();

        self.n_cells_sim = (self.n_cells as f32 * self.params.sim_doublet_ratio) as usize;
        let pairs = random_pairs(self.cells_to_keep.len(), self.n_cells_sim, seed);
        let (pair_1, pair_2): (Vec<usize>, Vec<usize>) = pairs.iter().copied().unzip();

        let sim_chunks = simulate_from_pairs(
            &pairs,
            &self.cells_to_keep,
            &self.hvg_library_sizes,
            &hvg_genes,
            &self.f_path_cell,
            target_size,
            self.params.log_transform,
        )?;

        if verbosity.normal_verbosity() {
            println!(
                "Simulated {} doublets. Done in {:.2?}",
                sim_chunks.len(),
                start_sim.elapsed()
            );
        }

        // PCA + projection
        if verbosity.normal_verbosity() {
            println!("Running PCA...");
        }

        let (combined_pca, _) = pca_and_project(
            &self.f_path_gene,
            &self.cells_to_keep,
            &hvg_genes,
            &self.hvg_library_sizes,
            target_size,
            &sim_chunks,
            &pca_opts,
            seed,
            verbose,
        )?;

        // kNN
        if verbosity.normal_verbosity() {
            println!("Building kNN graph...");
        }
        let start_knn = Instant::now();

        let k_adj = adjusted_k(self.params.knn_params.k, self.n_cells, self.n_cells_sim);
        if verbosity.normal_verbosity() {
            println!("Using {} neighbours in the kNN generation.", k_adj);
        }
        let knn_indices = dispatch_knn(
            combined_pca.as_ref(),
            k_adj,
            &self.params.knn_params,
            seed,
            verbosity.detailed_verbosity(),
        )?;

        if verbosity.normal_verbosity() {
            println!("Done with kNN generation in {:.2?}", start_knn.elapsed());
        }

        // Scoring and calling
        if verbosity.normal_verbosity() {
            println!("Calculating doublet scores...");
        }
        let start_score = Instant::now();

        let doublet_scores = self.calculate_doublet_scores(&knn_indices);
        let res = self.call_doublets(
            doublet_scores,
            self.params.manual_threshold,
            self.params.n_bins_hist,
            verbose,
        );

        if verbosity.normal_verbosity() {
            println!(
                "Done with doublet scoring and calling {:.2?}",
                start_score.elapsed()
            );
            println!("Finished Scrublet {:.2?}", start_all.elapsed());
        }

        let pca_out = if return_combined_pca {
            Some(combined_pca)
        } else {
            None
        };
        let pair_1_out = if return_pairs { Some(pair_1) } else { None };
        let pair_2_out = if return_pairs { Some(pair_2) } else { None };

        Ok((res, pca_out, pair_1_out, pair_2_out))
    }

    /// Calculate Bayesian doublet scores from the kNN graph.
    ///
    /// For each cell (observed and simulated), computes the doublet probability
    /// based on the fraction of simulated neighbours, adjusted for the expected
    /// doublet rate and simulation ratio.
    ///
    /// ### Params
    ///
    /// * `knn_indices` - kNN graph (one entry per cell, containing neighbour
    ///   indices).
    ///
    /// ### Returns
    ///
    /// `ScrubletDoubletScores`: (obs_scores, obs_errors, sim_scores, sim_errors).
    fn calculate_doublet_scores(&self, knn_indices: &[Vec<usize>]) -> ScrubletDoubletScores {
        let n_obs = self.n_cells;
        let n_sim = self.n_cells_sim;

        let r = n_sim as f32 / n_obs as f32;
        let rho = self.params.expected_doublet_rate;
        let se_rho = self.params.stdev_doublet_rate;

        let k_adj = adjusted_k(self.params.knn_params.k, self.n_cells, self.n_cells_sim);
        let n_adj = k_adj as f32;

        let scores_errors: Vec<(f32, f32)> = knn_indices
            .par_iter()
            .map(|neighbours| {
                let n_sim_neigh = neighbours.iter().filter(|&&idx| idx >= n_obs).count() as f32;
                let q = (n_sim_neigh + 1.0) / (n_adj + 2.0);
                let denominator = 1.0 - rho - q * (1.0 - rho - rho / r);
                let score = if denominator.abs() > 1e-10 {
                    (q * rho / r) / denominator
                } else {
                    0.0
                };

                let se_q = (q * (1.0 - q) / (n_adj + 3.0)).sqrt();
                let factor = q * rho / r / (denominator * denominator);
                let se_score = factor
                    * ((se_q / q * (1.0 - rho)).powi(2) + (se_rho / rho * (1.0 - q)).powi(2))
                        .sqrt();

                (score.max(0.0), se_score.max(1e-10))
            })
            .collect();

        let (scores_obs, errors_obs): (Vec<f32>, Vec<f32>) =
            scores_errors[..n_obs].iter().copied().unzip();
        let (scores_sim, errors_sim): (Vec<f32>, Vec<f32>) =
            scores_errors[n_obs..].iter().copied().unzip();

        (scores_obs, errors_obs, scores_sim, errors_sim)
    }

    /// Threshold doublet scores and produce final predictions.
    ///
    /// Uses `find_threshold_otsu` if no manual threshold is provided.
    ///
    /// ### Params
    ///
    /// * `doublet_scores` - Output from `calculate_doublet_scores`.
    /// * `manual_threshold` - Optional override for the automatic threshold.
    /// * `n_bins` - Number of bins for Otsu histogram.
    /// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for
    ///   detailed verbosity.
    ///
    /// ### Returns
    ///
    /// `ScrubletResult` with predictions, scores and diagnostics.
    fn call_doublets(
        &self,
        doublet_scores: ScrubletDoubletScores,
        manual_threshold: Option<f32>,
        n_bins: usize,
        verbose: usize,
    ) -> ScrubletResult {
        let verbosity = parse_verbosity_level(verbose);

        let threshold = manual_threshold.unwrap_or_else(|| {
            let t = find_threshold_otsu(&doublet_scores.0, n_bins);
            if verbosity.normal_verbosity() {
                println!("Automatically set threshold at doublet score = {:.4}", t);
            }
            t
        });

        let predicted_doublets: Vec<bool> = doublet_scores
            .0
            .iter()
            .map(|&score| score > threshold)
            .collect();

        let z_scores: Vec<f32> = doublet_scores
            .0
            .iter()
            .zip(doublet_scores.1.iter())
            .map(|(&score, &error)| (score - threshold) / error)
            .collect();

        let n_detected = predicted_doublets.iter().filter(|&&x| x).count();
        let detected_doublet_rate = n_detected as f32 / doublet_scores.0.len() as f32;

        let n_detectable = doublet_scores.2.iter().filter(|&&s| s > threshold).count();
        let detectable_doublet_fraction = n_detectable as f32 / doublet_scores.2.len() as f32;

        let overall_doublet_rate = if detectable_doublet_fraction > 0.01 {
            detected_doublet_rate / detectable_doublet_fraction
        } else {
            0.0
        };

        if verbosity.normal_verbosity() {
            println!(
                "Detected doublet rate = {:.1}%",
                100.0 * detected_doublet_rate
            );
            println!(
                "Estimated detectable doublet fraction = {:.1}%",
                100.0 * detectable_doublet_fraction
            );
            println!("Overall doublet rate:");
            println!(
                "  Expected  = {:.1}%",
                100.0 * self.params.expected_doublet_rate
            );
            println!("  Estimated = {:.1}%", 100.0 * overall_doublet_rate);
        }

        ScrubletResult {
            predicted_doublets,
            doublet_scores_obs: doublet_scores.0,
            doublet_scores_sim: doublet_scores.2,
            doublet_errors_obs: doublet_scores.1,
            z_scores,
            threshold,
            detected_doublet_rate,
            detectable_doublet_fraction,
            overall_doublet_rate,
        }
    }
}