use crate::handlers::explore::{off_screen, open_data, open_session, side_stats, values_reader};
use crate::hdf5_io::*;
use crate::interactive::stat_tui::{explore_in, Purpose, Side, StatExplorer};
use crate::interactive::ui::Tui;
use crate::qc::*;
use crate::sparse_io::*;
use crate::sparse_io_vector::*;
use data_beans::figure::save as save_figure;
use data_beans::figure::stat::{figure as stat_figure, StatData};
use clap::{Args, ValueEnum};
use legume_numeric::matrix::common_io::*;
use legume_numeric::matrix::membership::Membership;
use legume_numeric::matrix::traits::RunningStatOps;
use log::{info, warn};
#[derive(ValueEnum, Clone, Debug, PartialEq)]
#[clap(rename_all = "lowercase")]
pub enum StatDim {
Row,
Column,
}
#[derive(Args, Debug)]
pub struct RunStatArgs {
#[arg(
required = true,
value_delimiter = ',',
help = "Input data files in '.zarr' or '.h5' format",
long_help = "Provide data files in either '.zarr' or '.h5' format.\n\
Convert '.mtx' files with the 'data-beans from-mtx' command."
)]
pub data_files: Vec<Box<str>>,
#[arg(
short,
long,
value_enum,
help = "Statistics dimension (row or column; default column)",
long_help = "Choose whether to compute statistics over rows or columns (default\n\
column). The explorer starts on this side; Tab switches."
)]
pub stat_dim: Option<StatDim>,
#[arg(
short,
long,
help = "Row name regex pattern for column statistics",
long_help = "Regex pattern selecting row names.\n\
It applies when accumulating statistics over columns.\n\
Only rows matching this pattern will be included. Examples:\n\
'^MT-' (starts with MT-), 'GAPDH$' (ends with GAPDH),\n\
'^(MT|RPL|RPS)-' (mitochondrial or ribosomal genes).\n\
Matching is case-insensitive."
)]
pub row_name_pattern: Option<Box<str>>,
#[arg(
short = 'g',
long,
help = "Column group membership file for row statistics",
long_help = "File defining column group membership.\n\
It applies when accumulating statistics over rows.\n\
The statistics are then computed group-wise."
)]
pub column_group_file: Option<Box<str>>,
#[arg(
short = 'd',
long,
default_value = "@",
help = "Delimiter for extracting base barcode from column names",
long_help = "Delimiter used to extract the base barcode for matching.\n\
With delimiter '@',\n\
column 'ACGT-1@batch1' matches membership key 'ACGT-1@batch2',\n\
via base key 'ACGT-1'."
)]
pub delimiter: char,
#[arg(
long,
alias = "preload-data",
default_value_t = false,
help = "Preload data into memory for faster processing",
long_help = "Preload all column data into memory before computing statistics.\n\
This can significantly speed up processing but requires more memory."
)]
pub preload: bool,
#[arg(
long,
help = "Cells per rayon job (omit for auto-scaling by feature count)"
)]
pub block_size: Option<usize>,
#[arg(
short,
long,
default_value = "stdout",
help = "Output statistics file",
long_help = "Specify the output file for statistics.\n\
You can provide a '.parquet' file for efficient storage,\n\
or use 'stdout' to print results to the console."
)]
pub output: Box<str>,
#[arg(
long,
default_value_t = false,
help = "Run straight through: no explorer, even on a terminal",
long_help = "On a terminal, the statistics open in the explorer once computed: a\n\
sortable, filterable table of every row or column beside a histogram of\n\
nnz, sum, mean or sd. A file --output is still written; printing to stdout\n\
is skipped. --batch-process (or no terminal) prints or writes them and\n\
stops. --column-group-file always runs straight through."
)]
pub batch_process: bool,
#[arg(
long,
long_help = "Write the statistics figure under this name: the histograms of nnz, sum,\n\
mean and sd over the rows or columns, and mean against sd with the\n\
Poisson line. NAME.pdf and NAME.png, or the one file NAME names when it\n\
ends in .pdf, .png or .svg. Not available with --column-group-file."
)]
pub figure: Option<Box<str>>,
#[command(flatten)]
pub session: legume_tui::cli::SessionArgs,
}
pub fn run_stat(cmd_args: &RunStatArgs) -> anyhow::Result<()> {
let grouped = cmd_args.stat_dim == Some(StatDim::Row) && cmd_args.column_group_file.is_some();
let explore_after = !grouped && wants_explorer(cmd_args);
let session = explore_after
.then(|| open_session("stat", &cmd_args.session))
.transpose()?;
let mut tui = session.as_ref().map(Tui::new);
let output = cmd_args.output.clone();
dirname(&output).as_deref().map(mkdir).transpose()?;
let data = open_data(&cmd_args.data_files, cmd_args.preload)?;
let dim = cmd_args.stat_dim.clone().unwrap_or(StatDim::Column);
if dim == StatDim::Row {
if let Some(column_group_file) = &cmd_args.column_group_file {
let cols = data.column_names()?;
let membership = Membership::from_file(column_group_file, 0, 1, true)?
.with_delimiter(cmd_args.delimiter);
let (column_membership, stats) = membership.match_keys(&cols);
info!(
"Column matching: {} exact + {} base_key + {} prefix = {}/{} matched",
stats.exact,
stats.base_key,
stats.prefix,
stats.total_matched(),
stats.total()
);
if column_membership.is_empty() {
let data_sample: Vec<_> = cols.iter().take(3).collect();
let memb_sample = membership.sample_keys(3);
info!("Data columns sample: {:?}", data_sample);
info!("Membership keys sample: {:?}", memb_sample);
}
let unique_groups = membership.unique_groups();
info!(
"Will collect stats for {} groups: {:?}",
unique_groups.len(),
unique_groups
);
let (group_names, group_stats) = collect_stratified_row_stat_across_vec(
&data,
&column_membership,
cmd_args.block_size,
)?;
info!(
"Collected {} group stats: {:?}",
group_names.len(),
group_names
);
if cmd_args.figure.is_some() {
warn!("--figure is not available with --column-group-file; skipping it");
}
if cmd_args.output.eq_ignore_ascii_case("stdout") {
for (g, row_stat) in group_names.iter().zip(group_stats.iter()) {
let out: Vec<Box<str>> = row_stat
.to_string_vec(&data.row_names()?, "\t")?
.into_iter()
.map(|s| format!("{}\t{}", g, s).into_boxed_str())
.collect();
write_lines(&out, &cmd_args.output)?;
}
} else {
use legume_numeric::matrix::sparse_stat::save_grouped_stats_parquet;
info!("writing out: {}", cmd_args.output);
save_grouped_stats_parquet(
&cmd_args.output,
&data.row_names()?,
&group_names,
&group_stats,
)?;
}
return Ok(());
}
}
let side = match dim {
StatDim::Row => Side::Rows,
StatDim::Column => Side::Columns,
};
let save = !(explore_after && cmd_args.output.eq_ignore_ascii_case("stdout"));
let select_rows = rows_matching(cmd_args, &data)?;
let select_rows = select_rows.as_deref();
let output = save.then_some(&*cmd_args.output);
let first = off_screen(tui.as_mut().filter(|_| output.is_some()), || {
side_stats(&data, side, select_rows, cmd_args.block_size, output)
})?;
if let Some(name) = &cmd_args.figure {
let title = cmd_args
.data_files
.iter()
.map(|f| basename(f).unwrap_or_else(|_| f.clone()).to_string())
.collect::<Vec<_>>()
.join(", ");
let v = &first.values;
let svg = stat_figure(&StatData {
title: &title,
side: match side {
Side::Rows => "rows",
Side::Columns => "columns",
},
values: [&v[0], &v[1], &v[2], &v[3]],
marked: None,
cursor: None,
});
let paths = off_screen(tui.as_mut(), || save_figure(&svg, name))?;
info!("wrote the statistics figure: {}", paths.join(", "));
}
if let Some(tui) = tui.as_mut() {
let title = cmd_args.data_files.join(", ");
let data = &data;
let loader =
Box::new(move |side| side_stats(data, side, select_rows, cmd_args.block_size, None));
explore_in(
tui,
StatExplorer::new(
&title,
side,
first,
Some(loader),
Some(values_reader(data)),
Purpose::Explore,
),
)?;
}
Ok(())
}
fn wants_explorer(cmd_args: &RunStatArgs) -> bool {
crate::interactive::opens_view("stat", cmd_args.batch_process)
}
fn rows_matching(cmd_args: &RunStatArgs, data: &SparseIoVec) -> anyhow::Result<Option<Vec<usize>>> {
let Some(pattern) = &cmd_args.row_name_pattern else {
return Ok(None);
};
let row_names = data.row_names()?;
let selected = crate::utilities::name_matching::names_matching(&row_names, pattern)
.map_err(|e| anyhow::anyhow!("--row-name-pattern: {e}"))?;
info!(
"Row pattern '{}' matched {}/{} rows",
pattern,
selected.len(),
row_names.len()
);
Ok(Some(selected))
}
#[derive(ValueEnum, Clone, Debug, PartialEq)]
#[clap(rename_all = "lowercase")]
pub enum HistStat {
Nnz,
#[value(alias = "tot")]
Sum,
#[value(alias = "mu")]
Mean,
#[value(alias = "sig")]
Sd,
}
impl HistStat {
fn index(&self) -> usize {
match self {
HistStat::Nnz => 0,
HistStat::Sum => 1,
HistStat::Mean => 2,
HistStat::Sd => 3,
}
}
pub(crate) fn metric(&self) -> &'static str {
match self {
HistStat::Nnz => "nnz",
HistStat::Sum => "sum",
HistStat::Mean => "mean",
HistStat::Sd => "sd",
}
}
pub(crate) fn values<S: RunningStatOps<f32, Output = Vec<f32>>>(&self, s: &S) -> Vec<f32> {
match self {
HistStat::Nnz => s.count_positives(),
HistStat::Sum => s.sum(),
HistStat::Mean => s.mean(),
HistStat::Sd => s.std(),
}
}
}
#[derive(ValueEnum, Clone, Debug, PartialEq)]
#[clap(rename_all = "lowercase")]
pub enum HistDim {
Row,
Column,
Both,
}
#[derive(Args, Debug)]
pub struct RunHistogramArgs {
pub data_file: Box<str>,
#[arg(short, long, value_enum, default_value = "nnz")]
pub stat: HistStat,
#[arg(short, long, value_enum, default_value = "both")]
pub dim: HistDim,
#[arg(long, alias = "preload-data", default_value_t = false)]
pub preload: bool,
#[arg(long)]
pub block_size: Option<usize>,
#[arg(short, long)]
pub output: Option<Box<str>>,
#[arg(
long,
default_value_t = false,
help = "Print the histograms and stop: no explorer, even on a terminal",
long_help = "On a terminal, the statistics open in the explorer instead of being printed:\n\
a sortable, filterable table of every row or column beside a histogram of\n\
nnz, sum, mean or sd, on the --stat and --dim asked (Tab switches sides).\n\
--output values are still written. --batch-process (or no terminal)\n\
prints the histograms and stops."
)]
pub batch_process: bool,
#[command(flatten)]
pub session: legume_tui::cli::SessionArgs,
}
pub fn run_histogram(cmd_args: &RunHistogramArgs) -> anyhow::Result<()> {
if crate::interactive::opens_view("histogram", cmd_args.batch_process) {
return explore_histogram(cmd_args);
}
let (backend, data_file) = resolve_backend_file(&cmd_args.data_file, None)?;
let mut data = open_sparse_matrix(&data_file, &backend)?;
if cmd_args.preload {
info!("Preloading data from {} ...", data_file);
data.preload_columns()?;
}
let metric = cmd_args.stat.metric();
if matches!(cmd_args.dim, HistDim::Row | HistDim::Both) {
let row_stat = collect_row_stat(data.as_ref(), cmd_args.block_size)?;
let vals = cmd_args.stat.values(&row_stat);
print_nnz_summary("Row (feature)", metric, &vals, 0, None);
write_stat_values(cmd_args.output.as_deref(), "row", metric, &vals)?;
}
if matches!(cmd_args.dim, HistDim::Column | HistDim::Both) {
if matches!(cmd_args.dim, HistDim::Both) {
println!();
}
let col_stat = collect_column_stat(data.as_ref(), cmd_args.block_size)?;
let vals = cmd_args.stat.values(&col_stat);
print_nnz_summary("Column (cell)", metric, &vals, 0, None);
write_stat_values(cmd_args.output.as_deref(), "col", metric, &vals)?;
}
Ok(())
}
fn explore_histogram(cmd_args: &RunHistogramArgs) -> anyhow::Result<()> {
let session = open_session("histogram", &cmd_args.session)?;
let mut tui = Tui::new(&session);
let data = open_data(std::slice::from_ref(&cmd_args.data_file), cmd_args.preload)?;
let block = cmd_args.block_size;
let metric = cmd_args.stat.metric();
let (side, both) = match cmd_args.dim {
HistDim::Row => (Side::Rows, false),
HistDim::Column => (Side::Columns, false),
HistDim::Both => (Side::Columns, true),
};
let (first, mut other) = off_screen(Some(&mut tui), || {
let first = side_stats(&data, side, None, block, None)?;
let other = both
.then(|| side_stats(&data, Side::Rows, None, block, None))
.transpose()?;
let k = cmd_args.stat.index();
let prefix = cmd_args.output.as_deref();
if let Some(rows) = &other {
write_stat_values(prefix, "row", metric, &rows.values[k])?;
}
let dim = if side == Side::Rows { "row" } else { "col" };
write_stat_values(prefix, dim, metric, &first.values[k])?;
Ok((first, other))
})?;
let data = &data;
let loader = Box::new(move |s| match other.take() {
Some(rows) => Ok(rows),
None => side_stats(data, s, None, block, None),
});
let explorer = StatExplorer::new(
&cmd_args.data_file,
side,
first,
Some(loader),
Some(values_reader(data)),
Purpose::Explore,
)
.on_stat(cmd_args.stat.index());
explore_in(&mut tui, explorer)?;
Ok(())
}
fn write_stat_values(
prefix: Option<&str>,
dim: &str,
metric: &str,
vals: &[f32],
) -> anyhow::Result<()> {
if let Some(prefix) = prefix {
let path = format!("{}.{}_{}.txt", prefix, dim, metric);
write_types(&vals.to_vec(), &path)?;
info!("wrote {}", path);
}
Ok(())
}