use rayon::prelude::*;
use rustc_hash::FxHashMap;
use std::time::Instant;
use thousands::Separable;
use crate::core::base::info::{BinningStrategy, parse_bin_strategy_type};
use crate::core::base::loess::*;
use crate::prelude::*;
#[derive(Clone, Debug)]
pub struct HvgRes {
pub mean: Vec<f64>,
pub var: Vec<f64>,
pub var_exp: Vec<f64>,
pub var_std: Vec<f64>,
}
#[derive(Clone, Debug)]
pub struct HvgDispersionRes {
pub mean: Vec<f64>,
pub dispersion: Vec<f64>,
pub dispersion_scaled: Vec<f64>,
pub bin: Vec<i32>,
}
#[derive(Debug, Clone, Copy, Default)]
pub enum HvgMethod {
#[default]
Vst,
MeanVarBin,
Dispersion,
}
pub fn parse_hvg_method(s: &str) -> Option<HvgMethod> {
match s.to_lowercase().as_str() {
"vst" => Some(HvgMethod::Vst),
"meanvarbin" => Some(HvgMethod::MeanVarBin),
"dispersion" => Some(HvgMethod::Dispersion),
_ => None,
}
}
#[inline]
pub fn calculate_disp_stats_filtered(
gene: &CscGeneChunk,
cell_idx_map: &FxHashMap<u32, u32>,
no_cells: usize,
) -> (f32, f32) {
let n = no_cells as f64;
let mut sum = 0f64;
let mut sum_sq = 0f64;
for i in 0..gene.indices.len() {
if cell_idx_map.contains_key(&gene.indices[i]) {
let v = (gene.data_norm[i].to_f32() as f64).exp_m1();
sum += v;
sum_sq += v * v;
}
}
let mean = sum / n;
let var = if n > 1.0 {
((sum_sq - n * mean * mean) / (n - 1.0)).max(0.0)
} else {
0.0
};
let exp_mean = mean.ln_1p() as f32;
let log_vmr = if mean > 0.0 && var > 0.0 {
(var / mean).ln() as f32
} else {
f32::NAN
};
(exp_mean, log_vmr)
}
fn bin_features(means: &[f32], method: BinningStrategy, n_bins: usize) -> Vec<i32> {
let n = means.len();
let valid: Vec<bool> = means.iter().map(|v| v.is_finite()).collect();
match method {
BinningStrategy::EqualWidth => {
let (mut vmin, mut vmax) = (f32::INFINITY, f32::NEG_INFINITY);
for i in 0..n {
if valid[i] {
if means[i] < vmin {
vmin = means[i];
}
if means[i] > vmax {
vmax = means[i];
}
}
}
if !vmin.is_finite() || !vmax.is_finite() || vmax <= vmin {
return vec![-1; n];
}
let width = (vmax - vmin) / n_bins as f32;
(0..n)
.map(|i| {
if !valid[i] {
-1
} else {
let idx = ((means[i] - vmin) / width) as i32;
idx.clamp(0, n_bins as i32 - 1)
}
})
.collect()
}
BinningStrategy::EqualFrequency => {
let mut nonzero: Vec<f32> = means
.iter()
.zip(valid.iter())
.filter_map(|(&v, &ok)| if ok && v > 0.0 { Some(v) } else { None })
.collect();
if nonzero.len() < 2 {
return vec![-1; n];
}
nonzero.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let last = nonzero.len() - 1;
let breaks: Vec<f32> = (0..=n_bins)
.map(|i| {
let p = i as f32 / n_bins as f32;
let idx = (p * last as f32).round() as usize;
nonzero[idx.min(last)]
})
.collect();
(0..n)
.map(|i| {
if !valid[i] {
return -1;
}
let v = means[i];
if v < breaks[0] {
return 0;
}
for b in 0..n_bins {
if v <= breaks[b + 1] {
return b as i32;
}
}
(n_bins - 1) as i32
})
.collect()
}
}
}
fn scale_within_bins(dispersion: &[f32], bins: &[i32], n_bins: usize) -> Vec<f32> {
let mut sums = vec![0f64; n_bins];
let mut sums_sq = vec![0f64; n_bins];
let mut counts = vec![0usize; n_bins];
for (&d, &b) in dispersion.iter().zip(bins.iter()) {
if b >= 0 && d.is_finite() {
let bi = b as usize;
let df = d as f64;
sums[bi] += df;
sums_sq[bi] += df * df;
counts[bi] += 1;
}
}
let stats: Vec<(f64, f64)> = (0..n_bins)
.map(|i| {
if counts[i] < 2 {
(0.0, 0.0)
} else {
let c = counts[i] as f64;
let m = sums[i] / c;
let var = ((sums_sq[i] - c * m * m) / (c - 1.0)).max(0.0);
(m, var.sqrt())
}
})
.collect();
dispersion
.iter()
.zip(bins.iter())
.map(|(&d, &b)| {
if b < 0 || !d.is_finite() {
return 0.0;
}
let (m, sd) = stats[b as usize];
if sd > 0.0 {
((d as f64 - m) / sd) as f32
} else {
0.0
}
})
.collect()
}
pub fn build_disp_result(
means: Vec<f32>,
dispersions: Vec<f32>,
binning: BinningStrategy,
n_bins: usize,
) -> HvgDispersionRes {
let bins = bin_features(&means, binning, n_bins);
let scaled = scale_within_bins(&dispersions, &bins, n_bins);
let means_clean: Vec<f32> = means
.iter()
.map(|&v| if v.is_finite() { v } else { 0.0 })
.collect();
let disp_clean: Vec<f32> = dispersions
.iter()
.map(|&v| if v.is_finite() { v } else { 0.0 })
.collect();
HvgDispersionRes {
mean: means_clean.r_float_convert(),
dispersion: disp_clean.r_float_convert(),
dispersion_scaled: scaled.r_float_convert(),
bin: bins,
}
}
#[inline]
pub fn calculate_mean_var_filtered(
gene: &CscGeneChunk,
cell_idx_map: &FxHashMap<u32, u32>,
no_cells: usize,
) -> (f32, f32) {
let no_cells = no_cells as f32;
let mut sum = 0f32;
let mut nnz = 0usize;
for i in 0..gene.indices.len() {
if cell_idx_map.contains_key(&gene.indices[i]) {
sum += gene.data_raw.get(i) as f32;
nnz += 1;
}
}
let n_zero = no_cells - nnz as f32;
let mean = sum / no_cells;
let mut sum_sq_diff = 0f32;
for i in 0..gene.indices.len() {
if cell_idx_map.contains_key(&gene.indices[i]) {
let val = gene.data_raw.get(i) as f32;
let diff = val - mean;
sum_sq_diff += diff * diff;
}
}
let var = (sum_sq_diff + n_zero * mean * mean) / no_cells;
(mean, var)
}
#[inline]
pub fn calculate_std_variance_filtered(
gene: &CscGeneChunk,
cell_idx_map: &FxHashMap<u32, u32>,
mean: f32,
expected_var: f32,
clip_max: f32,
no_cells: usize,
) -> f32 {
let no_cells_f32 = no_cells as f32;
let expected_sd = expected_var.sqrt();
let mut sum_standardised = 0f32;
let mut sum_sq_standardised = 0f32;
let mut nnz = 0usize;
for i in 0..gene.indices.len() {
if cell_idx_map.contains_key(&gene.indices[i]) {
let val_f32 = gene.data_raw.get(i) as f32;
let norm = ((val_f32 - mean) / expected_sd)
.min(clip_max)
.max(-clip_max);
sum_standardised += norm;
sum_sq_standardised += norm * norm;
nnz += 1;
}
}
let n_zeros = no_cells - nnz;
if n_zeros > 0 {
let standardised_zero = ((-mean) / expected_sd).min(clip_max).max(-clip_max);
sum_standardised += n_zeros as f32 * standardised_zero;
sum_sq_standardised += n_zeros as f32 * standardised_zero * standardised_zero;
}
let standardised_mean = sum_standardised / no_cells_f32;
(sum_sq_standardised / no_cells_f32) - (standardised_mean * standardised_mean)
}
pub fn get_hvg_vst(
f_path: &str,
cell_indices: &[usize],
loess_span: f32,
clip_max: Option<f32>,
verbose: usize,
) -> Result<HvgRes, BixverseErrors> {
let verbosity = parse_verbosity_level(verbose);
let start_total = Instant::now();
let start_read = Instant::now();
let reader = ParallelSparseReader::new(f_path)?;
let mut gene_chunks: Vec<CscGeneChunk> = reader.get_all_genes()?;
let no_cells = cell_indices.len();
let cell_idx_map: FxHashMap<u32, u32> = cell_indices
.iter()
.enumerate()
.map(|(new_idx, &old_idx)| (old_idx as u32, new_idx as u32))
.collect();
let end_read = start_read.elapsed();
if verbosity.normal_verbosity() {
println!("HVG (VST): Load in data in {:.2?}", end_read);
}
let start_gene_stats = Instant::now();
let results: Vec<(f32, f32)> = gene_chunks
.par_iter_mut()
.map(|chunk| calculate_mean_var_filtered(chunk, &cell_idx_map, no_cells))
.collect();
let end_gene_stats = start_gene_stats.elapsed();
if verbosity.normal_verbosity() {
println!(
"HVG (VST): Calculated gene statistics in {:.2?}",
end_gene_stats
);
}
let start_loess = Instant::now();
let (means, vars): (Vec<f32>, Vec<f32>) = results.into_iter().unzip();
let clip_max = clip_max.unwrap_or((no_cells as f32).sqrt());
let means_log10: Vec<f32> = means.iter().map(|x| x.log10()).collect();
let vars_log10: Vec<f32> = vars.iter().map(|x| x.log10()).collect();
let loess = LoessRegression::new(loess_span, 2);
let loess_res = loess.fit(&means_log10, &vars_log10);
let end_loess = start_loess.elapsed();
if verbosity.normal_verbosity() {
println!("HVG (VST): Fitted Loess in {:.2?}", end_loess);
}
let start_standard = Instant::now();
let var_standardised: Vec<f32> = gene_chunks
.par_iter()
.zip(loess_res.fitted_vals.par_iter())
.zip(means.par_iter())
.map(|((chunk_i, var_i), mean_i)| {
let expected_var = 10_f32.powf(*var_i);
calculate_std_variance_filtered(
chunk_i,
&cell_idx_map,
*mean_i,
expected_var,
clip_max,
no_cells,
)
})
.collect();
let end_standard = start_standard.elapsed();
if verbosity.normal_verbosity() {
println!("HVG (VST): Standardised variance in {:.2?}", end_standard);
}
let total = start_total.elapsed();
if verbosity.normal_verbosity() {
println!("HVG (VST): Total run time -> {:.2?}", total);
}
Ok(HvgRes {
mean: means.r_float_convert(),
var: vars.r_float_convert(),
var_exp: loess_res.fitted_vals.r_float_convert(),
var_std: var_standardised.r_float_convert(),
})
}
pub fn get_hvg_vst_streaming(
f_path: &str,
cell_indices: &[usize],
loess_span: f32,
clip_max: Option<f32>,
verbose: usize,
) -> Result<HvgRes, BixverseErrors> {
let verbosity = parse_verbosity_level(verbose);
let start_total = Instant::now();
let reader = ParallelSparseReader::new(f_path)?;
let header = reader.get_header();
let no_genes = header.total_genes;
let no_cells = cell_indices.len();
let cell_idx_map: FxHashMap<u32, u32> = cell_indices
.iter()
.enumerate()
.map(|(new_idx, &old_idx)| (old_idx as u32, new_idx as u32))
.collect();
if verbosity.normal_verbosity() {
println!(
"HVG (VST, streaming): Pass 1/2 -> Calculating mean and variance for {} genes...",
no_genes.separate_with_underscores()
);
}
let start_pass1 = Instant::now();
const GENE_BATCH_SIZE: usize = 1000;
let num_batches = no_genes.div_ceil(GENE_BATCH_SIZE);
let mut means = Vec::with_capacity(no_genes);
let mut vars = Vec::with_capacity(no_genes);
for batch_idx in 0..num_batches {
if verbosity.detailed_verbosity() && batch_idx % 5 == 0 {
let progress = (batch_idx + 1) as f32 / num_batches as f32 * 100.0;
println!(" Progress: {:.1}%", progress);
}
let start_gene = batch_idx * GENE_BATCH_SIZE;
let end_gene = ((batch_idx + 1) * GENE_BATCH_SIZE).min(no_genes);
let gene_indices: Vec<usize> = (start_gene..end_gene).collect();
let start_loading = Instant::now();
let mut genes = reader.read_gene_parallel(&gene_indices)?;
let end_loading = start_loading.elapsed();
if verbosity.detailed_verbosity() {
println!(" Loaded batch in: {:.2?}.", end_loading);
}
let start_batch = Instant::now();
let batch_results: Vec<(f32, f32)> = genes
.par_iter_mut()
.map(|gene| calculate_mean_var_filtered(gene, &cell_idx_map, no_cells))
.collect();
let end_batch = start_batch.elapsed();
if verbosity.detailed_verbosity() {
println!(" Finished calculations in: {:.2?}.", end_batch);
}
for (mean, var) in batch_results {
means.push(mean);
vars.push(var);
}
drop(genes)
}
let end_pass1 = start_pass1.elapsed();
if verbosity.normal_verbosity() {
println!(" Calculated gene statistics: {:.2?}", end_pass1);
}
let start_loess = Instant::now();
let clip_max = clip_max.unwrap_or((no_cells as f32).sqrt());
let means_log10: Vec<f32> = means.iter().map(|x| x.log10()).collect();
let vars_log10: Vec<f32> = vars.iter().map(|x| x.log10()).collect();
let loess = LoessRegression::new(loess_span, 2);
let loess_res = loess.fit(&means_log10, &vars_log10);
let end_loess = start_loess.elapsed();
if verbosity.normal_verbosity() {
println!(" Fitted Loess: {:.2?}", end_loess);
println!("HVG (VST, streaming): Pass 2/2 -> Calculating standardised variance...");
}
let start_pass2 = Instant::now();
let mut var_standardised = Vec::with_capacity(no_genes);
for batch_idx in 0..num_batches {
if verbosity.detailed_verbosity() && batch_idx % 5 == 0 {
let progress = (batch_idx + 1) as f32 / num_batches as f32 * 100.0;
println!(" Progress: {:.1}%", progress);
}
let start_gene = batch_idx * GENE_BATCH_SIZE;
let end_gene = ((batch_idx + 1) * GENE_BATCH_SIZE).min(no_genes);
let gene_indices: Vec<usize> = (start_gene..end_gene).collect();
let mut genes = reader.read_gene_parallel(&gene_indices)?;
let start_batch = Instant::now();
let batch_std_vars: Vec<f32> = genes
.par_iter_mut()
.enumerate()
.map(|(local_idx, gene)| {
let gene_idx = start_gene + local_idx;
let expected_var = 10_f32.powf(loess_res.fitted_vals[gene_idx]);
calculate_std_variance_filtered(
gene,
&cell_idx_map,
means[gene_idx],
expected_var,
clip_max,
no_cells,
)
})
.collect();
let end_batch = start_batch.elapsed();
if verbosity.detailed_verbosity() {
println!(
" Finished calculating standardised variance in: {:.2?}.",
end_batch
);
}
drop(genes);
var_standardised.extend(batch_std_vars);
}
let end_pass2 = start_pass2.elapsed();
if verbosity.normal_verbosity() {
println!(
" Calculated standardised variance total: {:.2?}",
end_pass2
);
}
let total = start_total.elapsed();
if verbosity.normal_verbosity() {
println!("HVG (VST, streaming): Total run time -> {:.2?}", total);
}
Ok(HvgRes {
mean: means.r_float_convert(),
var: vars.r_float_convert(),
var_exp: loess_res.fitted_vals.r_float_convert(),
var_std: var_standardised.r_float_convert(),
})
}
pub fn get_hvg_dispersion(
f_path: &str,
cell_indices: &[usize],
binning: &str,
n_bins: usize,
verbose: usize,
) -> Result<HvgDispersionRes, BixverseErrors> {
let verbosity = parse_verbosity_level(verbose);
let start_total = Instant::now();
let binning = parse_bin_strategy_type(binning).unwrap_or_default();
let start_read = Instant::now();
let reader = ParallelSparseReader::new(f_path)?;
let mut gene_chunks: Vec<CscGeneChunk> = reader.get_all_genes()?;
let no_cells = cell_indices.len();
let cell_idx_map: FxHashMap<u32, u32> = cell_indices
.iter()
.enumerate()
.map(|(new_idx, &old_idx)| (old_idx as u32, new_idx as u32))
.collect();
if verbosity.normal_verbosity() {
println!(
"HVG (dispersion): Load in data in {:.2?}",
start_read.elapsed()
);
}
let start_stats = Instant::now();
let results: Vec<(f32, f32)> = gene_chunks
.par_iter_mut()
.map(|chunk| calculate_disp_stats_filtered(chunk, &cell_idx_map, no_cells))
.collect();
if verbosity.normal_verbosity() {
println!(
"HVG (dispersion): Calculated gene statistics in {:.2?}",
start_stats.elapsed()
);
}
let (means, dispersions): (Vec<f32>, Vec<f32>) = results.into_iter().unzip();
let start_bin = Instant::now();
let res = build_disp_result(means, dispersions, binning, n_bins);
if verbosity.normal_verbosity() {
println!(
"HVG (dispersion): Binning and scaling in {:.2?}",
start_bin.elapsed()
);
println!(
"HVG (dispersion): Total run time -> {:.2?}",
start_total.elapsed()
);
}
Ok(res)
}
pub fn get_hvg_dispersion_streaming(
f_path: &str,
cell_indices: &[usize],
binning: &str,
n_bins: usize,
verbose: usize,
) -> Result<HvgDispersionRes, BixverseErrors> {
let verbosity = parse_verbosity_level(verbose);
let start_total = Instant::now();
let binning = parse_bin_strategy_type(binning).unwrap_or_default();
let reader = ParallelSparseReader::new(f_path)?;
let header = reader.get_header();
let no_genes = header.total_genes;
let no_cells = cell_indices.len();
let cell_idx_map: FxHashMap<u32, u32> = cell_indices
.iter()
.enumerate()
.map(|(new_idx, &old_idx)| (old_idx as u32, new_idx as u32))
.collect();
if verbosity.normal_verbosity() {
println!(
"HVG (dispersion, streaming): Calculating dispersion stats for {} genes...",
no_genes.separate_with_underscores()
);
}
const GENE_BATCH_SIZE: usize = 1000;
let num_batches = no_genes.div_ceil(GENE_BATCH_SIZE);
let mut means = Vec::with_capacity(no_genes);
let mut dispersions = Vec::with_capacity(no_genes);
for batch_idx in 0..num_batches {
if verbosity.detailed_verbosity() && batch_idx % 5 == 0 {
let progress = (batch_idx + 1) as f32 / num_batches as f32 * 100.0;
println!(" Progress: {:.1}%", progress);
}
let start_gene = batch_idx * GENE_BATCH_SIZE;
let end_gene = ((batch_idx + 1) * GENE_BATCH_SIZE).min(no_genes);
let gene_indices: Vec<usize> = (start_gene..end_gene).collect();
let mut genes = reader.read_gene_parallel(&gene_indices)?;
let batch_results: Vec<(f32, f32)> = genes
.par_iter_mut()
.map(|gene| calculate_disp_stats_filtered(gene, &cell_idx_map, no_cells))
.collect();
for (m, d) in batch_results {
means.push(m);
dispersions.push(d);
}
drop(genes);
}
let start_bin = Instant::now();
let res = build_disp_result(means, dispersions, binning, n_bins);
if verbosity.normal_verbosity() {
println!(
"HVG (dispersion, streaming): Binning and scaling in {:.2?}",
start_bin.elapsed()
);
println!(
"HVG (dispersion, streaming): Total run time -> {:.2?}",
start_total.elapsed()
);
}
Ok(res)
}
pub fn get_hvg_mvb(
f_path: &str,
cell_indices: &[usize],
binning: &str,
n_bins: usize,
verbose: usize,
) -> Result<HvgDispersionRes, BixverseErrors> {
get_hvg_dispersion(f_path, cell_indices, binning, n_bins, verbose)
}
pub fn get_hvg_mvb_streaming(
f_path: &str,
cell_indices: &[usize],
binning: &str,
n_bins: usize,
verbose: usize,
) -> Result<HvgDispersionRes, BixverseErrors> {
get_hvg_dispersion_streaming(f_path, cell_indices, binning, n_bins, verbose)
}
pub fn get_hvg_vst_batch_aware(
f_path: &str,
cell_indices: &[usize],
batch_labels: &[usize],
loess_span: f32,
clip_max: Option<f32>,
verbose: usize,
) -> Result<Vec<HvgRes>, BixverseErrors> {
let verbosity = parse_verbosity_level(verbose);
let start_total = Instant::now();
let start_setup = Instant::now();
let n_batches = *batch_labels.iter().max().unwrap() + 1;
let mut batch_cell_maps: Vec<FxHashMap<u32, u32>> = vec![FxHashMap::default(); n_batches];
let mut batch_sizes = vec![0usize; n_batches];
for (&old_idx, &batch) in cell_indices.iter().zip(batch_labels.iter()) {
batch_cell_maps[batch].insert(old_idx as u32, batch_sizes[batch] as u32);
batch_sizes[batch] += 1;
}
let end_setup = start_setup.elapsed();
if verbosity.normal_verbosity() {
println!(
"HVG (VST, batch-aware): Setup batch maps in {:.2?}",
end_setup
);
println!(" Processing {} batches", n_batches);
}
let start_read = Instant::now();
let reader = ParallelSparseReader::new(f_path)?;
let mut gene_chunks: Vec<CscGeneChunk> = reader.get_all_genes()?;
let end_read = start_read.elapsed();
if verbosity.normal_verbosity() {
println!("HVG (VST, batch-aware): Loaded data in {:.2?}", end_read);
}
let start_pass_1 = Instant::now();
let mut batch_means: Vec<Vec<f32>> = vec![Vec::new(); n_batches];
let mut batch_vars: Vec<Vec<f32>> = vec![Vec::new(); n_batches];
for batch_idx in 0..n_batches {
let results: Vec<(f32, f32)> = gene_chunks
.par_iter_mut()
.map(|chunk| {
calculate_mean_var_filtered(
chunk,
&batch_cell_maps[batch_idx],
batch_sizes[batch_idx],
)
})
.collect();
let (means, vars): (Vec<f32>, Vec<f32>) = results.into_iter().unzip();
batch_means[batch_idx] = means;
batch_vars[batch_idx] = vars;
}
let end_pass_1 = start_pass_1.elapsed();
if verbosity.normal_verbosity() {
println!(
"HVG (VST, batch-aware): Calculated gene statistics per batch in {:.2?}",
end_pass_1
);
}
let start_loess = Instant::now();
let mut batch_loess_results = Vec::with_capacity(n_batches);
let mut batch_clip_max = Vec::with_capacity(n_batches);
for batch_idx in 0..n_batches {
let clip = clip_max.unwrap_or((batch_sizes[batch_idx] as f32).sqrt());
batch_clip_max.push(clip);
let means_log10: Vec<f32> = batch_means[batch_idx].iter().map(|x| x.log10()).collect();
let vars_log10: Vec<f32> = batch_vars[batch_idx].iter().map(|x| x.log10()).collect();
let loess = LoessRegression::new(loess_span, 2);
let loess_res = loess.fit(&means_log10, &vars_log10);
batch_loess_results.push(loess_res);
}
let end_loess = start_loess.elapsed();
if verbosity.normal_verbosity() {
println!(
"HVG (VST, batch-aware): Fitted loess per batch in {:.2?}",
end_loess
);
}
let start_pass_2 = Instant::now();
let mut batch_std_vars: Vec<Vec<f32>> = vec![Vec::new(); n_batches];
for batch_idx in 0..n_batches {
let var_standardised: Vec<f32> = gene_chunks
.par_iter()
.enumerate()
.map(|(gene_idx, chunk)| {
let expected_var =
10_f32.powf(batch_loess_results[batch_idx].fitted_vals[gene_idx]);
calculate_std_variance_filtered(
chunk,
&batch_cell_maps[batch_idx],
batch_means[batch_idx][gene_idx],
expected_var,
batch_clip_max[batch_idx],
batch_sizes[batch_idx],
)
})
.collect();
batch_std_vars[batch_idx] = var_standardised;
}
let end_pass_2 = start_pass_2.elapsed();
if verbosity.normal_verbosity() {
println!(
"HVG (VST, batch-aware): Calculated standardised variance per batch in {:.2?}",
end_pass_2
);
}
let total = start_total.elapsed();
if verbosity.normal_verbosity() {
println!("HVG (VST, batch-aware): Total runtime -> {:.2?}", total);
}
let res = batch_means
.into_iter()
.zip(batch_vars)
.zip(batch_loess_results)
.zip(batch_std_vars)
.map(|(((means, vars), loess_res), std_vars)| HvgRes {
mean: means.r_float_convert(),
var: vars.r_float_convert(),
var_exp: loess_res.fitted_vals.r_float_convert(),
var_std: std_vars.r_float_convert(),
})
.collect();
Ok(res)
}
pub fn get_hvg_vst_batch_aware_streaming(
f_path: &str,
cell_indices: &[usize],
batch_labels: &[usize],
loess_span: f32,
clip_max: Option<f32>,
verbose: usize,
) -> Result<Vec<HvgRes>, BixverseErrors> {
let verbosity = parse_verbosity_level(verbose);
let start_total = Instant::now();
let reader = ParallelSparseReader::new(f_path)?;
let header = reader.get_header();
let no_genes = header.total_genes;
let start_setup = Instant::now();
let n_batches = *batch_labels.iter().max().unwrap() + 1;
let mut batch_cell_maps: Vec<FxHashMap<u32, u32>> = vec![FxHashMap::default(); n_batches];
let mut batch_sizes = vec![0usize; n_batches];
for (&old_idx, &batch) in cell_indices.iter().zip(batch_labels.iter()) {
batch_cell_maps[batch].insert(old_idx as u32, batch_sizes[batch] as u32);
batch_sizes[batch] += 1;
}
let end_setup = start_setup.elapsed();
if verbosity.normal_verbosity() {
println!(
"HVG (VST, streaming, batch-aware): Setup batch maps in {:.2?}",
end_setup
);
println!(" Processing {} batches", n_batches);
}
if verbosity.normal_verbosity() {
println!(
"HVG (VST, streaming, batch-aware): Pass 1/2 -> Calculating mean and variance for {} genes...",
no_genes.separate_with_underscores()
);
}
let start_pass1 = Instant::now();
const GENE_BATCH_SIZE: usize = 1000;
let num_batches = no_genes.div_ceil(GENE_BATCH_SIZE);
let mut batch_means: Vec<Vec<f32>> = vec![Vec::with_capacity(no_genes); n_batches];
let mut batch_vars: Vec<Vec<f32>> = vec![Vec::with_capacity(no_genes); n_batches];
for chunk_idx in 0..num_batches {
if verbosity.detailed_verbosity() && chunk_idx % 5 == 0 {
let progress = (chunk_idx + 1) as f32 / num_batches as f32 * 100.0;
println!(" Progress: {:.1}%", progress);
}
let start_gene = chunk_idx * GENE_BATCH_SIZE;
let end_gene = ((chunk_idx + 1) * GENE_BATCH_SIZE).min(no_genes);
let gene_indices: Vec<usize> = (start_gene..end_gene).collect();
let mut genes = reader.read_gene_parallel(&gene_indices)?;
for batch_idx in 0..n_batches {
let results: Vec<(f32, f32)> = genes
.par_iter_mut()
.map(|gene| {
calculate_mean_var_filtered(
gene,
&batch_cell_maps[batch_idx],
batch_sizes[batch_idx],
)
})
.collect();
for (mean, var) in results {
batch_means[batch_idx].push(mean);
batch_vars[batch_idx].push(var);
}
}
}
let end_pass1 = start_pass1.elapsed();
if verbosity.normal_verbosity() {
println!(" Calculated gene statistics per batch: {:.2?}", end_pass1);
}
let start_loess = Instant::now();
let mut batch_loess_results = Vec::with_capacity(n_batches);
let mut batch_clip_max = Vec::with_capacity(n_batches);
for batch_idx in 0..n_batches {
let clip = clip_max.unwrap_or((batch_sizes[batch_idx] as f32).sqrt());
batch_clip_max.push(clip);
let means_log10: Vec<f32> = batch_means[batch_idx].iter().map(|x| x.log10()).collect();
let vars_log10: Vec<f32> = batch_vars[batch_idx].iter().map(|x| x.log10()).collect();
let loess = LoessRegression::new(loess_span, 2);
let loess_res = loess.fit(&means_log10, &vars_log10);
batch_loess_results.push(loess_res);
}
let end_loess = start_loess.elapsed();
if verbosity.normal_verbosity() {
println!(" Fitted loess per batch: {:.2?}", end_loess);
println!(
"HVG (VST, streaming, batch-aware): Pass 2/2 -> Calculating standardised variance..."
);
}
let start_pass2 = Instant::now();
let mut batch_std_vars: Vec<Vec<f32>> = vec![Vec::with_capacity(no_genes); n_batches];
for chunk_idx in 0..num_batches {
if verbosity.detailed_verbosity() && chunk_idx % 5 == 0 {
let progress = (chunk_idx + 1) as f32 / num_batches as f32 * 100.0;
println!(" Progress: {:.1}%", progress);
}
let start_gene = chunk_idx * GENE_BATCH_SIZE;
let end_gene = ((chunk_idx + 1) * GENE_BATCH_SIZE).min(no_genes);
let gene_indices: Vec<usize> = (start_gene..end_gene).collect();
let mut genes = reader.read_gene_parallel(&gene_indices)?;
for batch_idx in 0..n_batches {
let std_vars: Vec<f32> = genes
.par_iter_mut()
.enumerate()
.map(|(local_idx, gene)| {
let gene_idx = start_gene + local_idx;
let expected_var =
10_f32.powf(batch_loess_results[batch_idx].fitted_vals[gene_idx]);
calculate_std_variance_filtered(
gene,
&batch_cell_maps[batch_idx],
batch_means[batch_idx][gene_idx],
expected_var,
batch_clip_max[batch_idx],
batch_sizes[batch_idx],
)
})
.collect();
batch_std_vars[batch_idx].extend(std_vars);
}
}
let end_pass2 = start_pass2.elapsed();
if verbosity.normal_verbosity() {
println!(
" Calculated standardised variance per batch: {:.2?}",
end_pass2
);
}
let total = start_total.elapsed();
if verbosity.normal_verbosity() {
println!(
"HVG (VST, streaming, batch-aware): Total runtime -> {:.2?}",
total
);
}
let res = batch_means
.into_iter()
.zip(batch_vars)
.zip(batch_loess_results)
.zip(batch_std_vars)
.map(|(((means, vars), loess_res), std_vars)| HvgRes {
mean: means.r_float_convert(),
var: vars.r_float_convert(),
var_exp: loess_res.fitted_vals.r_float_convert(),
var_std: std_vars.r_float_convert(),
})
.collect();
Ok(res)
}
pub fn get_hvg_dispersion_batch_aware(
f_path: &str,
cell_indices: &[usize],
batch_labels: &[usize],
binning: &str,
n_bins: usize,
verbose: usize,
) -> Result<Vec<HvgDispersionRes>, BixverseErrors> {
let start_total = Instant::now();
let binning = parse_bin_strategy_type(binning).unwrap_or_default();
let verbosity = parse_verbosity_level(verbose);
let n_batches = *batch_labels.iter().max().unwrap() + 1;
let mut batch_cell_maps: Vec<FxHashMap<u32, u32>> = vec![FxHashMap::default(); n_batches];
let mut batch_sizes = vec![0usize; n_batches];
for (&old_idx, &batch) in cell_indices.iter().zip(batch_labels.iter()) {
batch_cell_maps[batch].insert(old_idx as u32, batch_sizes[batch] as u32);
batch_sizes[batch] += 1;
}
if verbosity.normal_verbosity() {
println!(
"HVG (Dispersion, batch-aware): Processing {} batches",
n_batches
);
}
let start_read = Instant::now();
let reader = ParallelSparseReader::new(f_path)?;
let mut gene_chunks: Vec<CscGeneChunk> = reader.get_all_genes()?;
if verbosity.normal_verbosity() {
println!(
"HVG (Dispersion, batch-aware): Loaded data in {:.2?}",
start_read.elapsed()
);
}
let mut out = Vec::with_capacity(n_batches);
for batch_idx in 0..n_batches {
let start_batch = Instant::now();
let results: Vec<(f32, f32)> = gene_chunks
.par_iter_mut()
.map(|chunk| {
calculate_disp_stats_filtered(
chunk,
&batch_cell_maps[batch_idx],
batch_sizes[batch_idx],
)
})
.collect();
let (means, dispersions): (Vec<f32>, Vec<f32>) = results.into_iter().unzip();
let res = build_disp_result(means, dispersions, binning, n_bins);
out.push(res);
if verbosity.detailed_verbosity() {
println!(
" Batch {}/{}: {:.2?}",
batch_idx + 1,
n_batches,
start_batch.elapsed()
);
}
}
if verbosity.normal_verbosity() {
println!(
"HVG (Dispersion, batch-aware): Total runtime -> {:.2?}",
start_total.elapsed()
);
}
Ok(out)
}
pub fn get_hvg_dispersion_batch_aware_streaming(
f_path: &str,
cell_indices: &[usize],
batch_labels: &[usize],
binning: &str,
n_bins: usize,
verbose: usize,
) -> Result<Vec<HvgDispersionRes>, BixverseErrors> {
let start_total = Instant::now();
let verbosity = parse_verbosity_level(verbose);
let binning = parse_bin_strategy_type(binning).unwrap_or_default();
let reader = ParallelSparseReader::new(f_path)?;
let header = reader.get_header();
let no_genes = header.total_genes;
let n_batches = *batch_labels.iter().max().unwrap() + 1;
let mut batch_cell_maps: Vec<FxHashMap<u32, u32>> = vec![FxHashMap::default(); n_batches];
let mut batch_sizes = vec![0usize; n_batches];
for (&old_idx, &batch) in cell_indices.iter().zip(batch_labels.iter()) {
batch_cell_maps[batch].insert(old_idx as u32, batch_sizes[batch] as u32);
batch_sizes[batch] += 1;
}
if verbosity.normal_verbosity() {
println!(
"HVG (Dispersion, batch-aware, streaming): Processing {} batches",
n_batches
);
println!(
" Calculating dispersion stats for {} genes...",
no_genes.separate_with_underscores()
);
}
const GENE_BATCH_SIZE: usize = 1000;
let num_batches = no_genes.div_ceil(GENE_BATCH_SIZE);
let mut batch_means: Vec<Vec<f32>> = vec![Vec::with_capacity(no_genes); n_batches];
let mut batch_dispersions: Vec<Vec<f32>> = vec![Vec::with_capacity(no_genes); n_batches];
for chunk_idx in 0..num_batches {
if verbosity.detailed_verbosity() && chunk_idx % 5 == 0 {
let progress = (chunk_idx + 1) as f32 / num_batches as f32 * 100.0;
println!(" Progress: {:.1}%", progress);
}
let start_gene = chunk_idx * GENE_BATCH_SIZE;
let end_gene = ((chunk_idx + 1) * GENE_BATCH_SIZE).min(no_genes);
let gene_indices: Vec<usize> = (start_gene..end_gene).collect();
let mut genes = reader.read_gene_parallel(&gene_indices)?;
for batch_idx in 0..n_batches {
let results: Vec<(f32, f32)> = genes
.par_iter_mut()
.map(|gene| {
calculate_disp_stats_filtered(
gene,
&batch_cell_maps[batch_idx],
batch_sizes[batch_idx],
)
})
.collect();
for (m, d) in results {
batch_means[batch_idx].push(m);
batch_dispersions[batch_idx].push(d);
}
}
drop(genes);
}
let out: Vec<HvgDispersionRes> = batch_means
.into_iter()
.zip(batch_dispersions)
.map(|(means, dispersions)| build_disp_result(means, dispersions, binning, n_bins))
.collect();
if verbosity.normal_verbosity() {
println!(
"HVG (Dispersion, batch-aware, streaming): Total runtime -> {:.2?}",
start_total.elapsed()
);
}
Ok(out)
}
pub fn get_hvg_mvb_batch_aware(
f_path: &str,
cell_indices: &[usize],
batch_labels: &[usize],
binning: &str,
n_bins: usize,
verbose: usize,
) -> Result<Vec<HvgDispersionRes>, BixverseErrors> {
get_hvg_dispersion_batch_aware(f_path, cell_indices, batch_labels, binning, n_bins, verbose)
}
pub fn get_hvg_mvb_batch_aware_streaming(
f_path: &str,
cell_indices: &[usize],
batch_labels: &[usize],
binning: &str,
n_bins: usize,
verbose: usize,
) -> Result<Vec<HvgDispersionRes>, BixverseErrors> {
get_hvg_dispersion_batch_aware_streaming(
f_path,
cell_indices,
batch_labels,
binning,
n_bins,
verbose,
)
}