diff --git a/Cargo.toml b/Cargo.toml index 62e071a..2056ee6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,7 +9,7 @@ edition = "2021" readme = "README.md" # Own version (decoupled from workspace.version): data-beans evolves its CLI / # QC API independently of the shared utility crates. -version = "0.9.1" +version = "0.10.0" rust-version = "1.91" # Ship only the sources, README and license: nothing else in the working tree # (data stores, saved name lists, scratch files) can end up on crates.io. @@ -27,7 +27,7 @@ path = "src/sim_main.rs" required-features = ["sim"] [features] -default = ["aux", "tui", "tensor", "ndarray"] +default = ["aux", "tui", "tensor", "ndarray", "mcp", "plot"] # Read sparse data into dense candle `Tensor`s (`read_*_tensor`). Off, no # candle is linked: for crates that only need matrices. tensor = ["legume-numeric/tensor"] @@ -40,7 +40,13 @@ aux = [] # Full-screen terminal views (ratatui): the CLI's interactive sessions, and # `data_beans::interactive` for downstream crates building their own. The # views run in legume-tui sessions, which a dropped connection leaves running. -tui = ["dep:ratatui", "dep:legume-tui"] +tui = ["dep:ratatui", "dep:legume-tui", "legume-mcp?/tui"] +# `data-beans mcp`: the commands as Model Context Protocol tools for an +# agent, over stdio. +mcp = ["dep:legume-mcp"] +# `data-beans plot`: QC and statistics figures (PNG, SVG, PDF) through +# legume-plot, and the MCP `plot` tool. +plot = ["dep:legume-plot", "dep:base64"] alg = ["aux"] sim = ["alg"] # Pulls in libhdf5 — opt-in so installs on hosts without HDF5 headers still @@ -90,8 +96,15 @@ statrs = { version = "0.18", default-features = false } ratatui = { version = "0.30", default-features = false, features = ["crossterm"], optional = true } # Detachable sessions for the views; no pictures, so no graphics. legume-tui = { version = "0.3.2", optional = true, default-features = false, features = ["clap", "logs"] } +# Its `tui` view (`data-beans mcp watch`) comes on with data-beans' `tui`. +legume-mcp = { version = "0.1", optional = true } +legume-plot = { version = "0.3.2", optional = true } +base64 = { version = "0.22", optional = true } [dev-dependencies] # The crate's harness, driving the views through pseudo-terminals in # tests/detach_session.rs. legume-tui = { version = "0.3.2", default-features = false, features = ["testing"] } +# tests/mcp_server.rs follows a server through the watch pane. +legume-mcp = { version = "0.1", features = ["tui"] } +ratatui = { version = "0.30", default-features = false, features = ["crossterm"] } diff --git a/README.md b/README.md index 932209f..3012217 100644 --- a/README.md +++ b/README.md @@ -32,3 +32,128 @@ Here, we implement a wrapper for sparse matrix backends for swift data access by ├── indices (column indices) └── indptr (row pointers) ``` + +## Agents (MCP) + +`data-beans mcp` serves data-beans to an AI agent as [Model Context Protocol](https://modelcontextprotocol.io) tools over stdio, for any MCP client: Claude Code, Claude Desktop, Cursor, VS Code, Gemini CLI, Codex CLI and others. The server talks only to the client that started it, over stdin and stdout; data-beans sends nothing anywhere itself. What leaves your machine is up to the agent's own client and model. + +### Connect a client + +The server is in the default build (the `mcp` feature; `cargo install data-beans` has it). Importing `.h5ad` and 10x `.h5` files needs `--features hdf5` too. + +The server keeps its records in a project folder: its working directory, or `LEGUME_MCP_PROJECT`. Set it for clients that start servers somewhere else (Claude Desktop, say). + +Most clients (Claude Desktop, Cursor, Gemini CLI) take an `mcpServers` entry: + +```json +{ + "mcpServers": { + "data-beans": { + "command": "data-beans", + "args": ["mcp"], + "env": { "LEGUME_MCP_PROJECT": "/path/to/project" } + } + } +} +``` + +Claude Code, from the project folder: + +```sh +claude mcp add data-beans -e LEGUME_MCP_PROJECT="$PWD" -- data-beans mcp +``` + +VS Code (`.vscode/mcp.json`): + +```json +{ + "servers": { + "data-beans": { + "type": "stdio", + "command": "data-beans", + "args": ["mcp"], + "env": { "LEGUME_MCP_PROJECT": "${workspaceFolder}" } + } + } +} +``` + +Codex CLI (`~/.codex/config.toml`): + +```toml +[mcp_servers.data-beans] +command = "data-beans" +args = ["mcp"] +env = { LEGUME_MCP_PROJECT = "/path/to/project" } +``` + +### What the agent is told + +The server hands these instructions to the client when it connects. Clients that do not pass MCP instructions on to the model can paste them into an `AGENTS.md` or the system prompt: + +```text +data-beans stores sparse single-cell matrices, features (genes, peaks) as rows and cells as columns, in .zarr.zip, .zarr or .h5 backends. + +Look before you derive: +1. info gives a backend's size and metadata; list_contents shows what an outside .zarr or .h5 holds before you import it. +2. names pages through row or column names, with a regex filter, to find the exact names to pass on. +3. statistics summarizes each row or column (nnz, tot, mu, sig) with quantiles and the top entries by a chosen statistic; values reads the counts of a few named rows or columns. + +Quality control: qc runs the cell QC senna and pinto apply (a near-empty floor, then robust MAD bands on detected genes, total counts and the mito fraction). Call it without `out` first: it reports each metric's median, the band and how many cells fall outside, and writes nothing. Adjust n_mads or the caps until the drops are what the data warrant, then call it with `out` to write the kept cells as a new backend. plot shows the same settings as a figure (histograms with the band edges, counts against genes by verdict), so look at it before and after tuning. + +Derive new backends with import (mtx, zarr triplets, fragments, h5ad, 10x; say in `source` where the input came from), subset, subsample, split, convert and merge. Each writes to `out`, which must not exist yet: no tool replaces a file. Leave `out` out and the output goes in today's results folder, named after its input and the tool. Outputs are .zarr.zip archives unless `out` ends in .zarr (a directory) or .h5. Paths are resolved against the server's working directory; results give the paths written. + +A person may be watching (`data-beans mcp watch`): a "Note from the user" in a result is theirs to you; take it into account. A call they stopped fails with "stopped by the user": do not run it again unasked. +``` + +### Tools + +| Tool | What it does | Writes | +|---|---|---| +| `info` | a backend's size (rows, columns, non-zeros) and metadata | | +| `list_contents` | the groups and arrays in a `.zarr`, `.zarr.zip` or `.h5`, before importing it | | +| `names` | row or column names, paged, filtered by a case-insensitive regex | | +| `statistics` | per-row or per-column nnz, total, mean and sd: quantiles and the top entries | a `.parquet` when `out` is given | +| `values` | the counts of up to 100 named or indexed rows or columns | a table when `out` is given | +| `qc` | senna's and pinto's cell QC: the near-empty floor, MAD bands on genes, counts and mito fraction, caps | the kept cells when `out` is given; a dry run otherwise | +| `plot` | the QC figure, or a statistic's histogram, returned as an image | the figure | +| `import` | a backend from `.mtx`, zarr triplets, fragments, `.h5ad` or 10x `.h5`, its `source` noted | a backend | +| `subset` | some rows or columns: by `names`, a `pattern` (regex) or `indices` (ranges such as `0-19612`) | a backend | +| `subsample` | cells and/or genes drawn at random | a backend | +| `split` | cells split train/test, or into K folds; by region or group with `options` | the halves | +| `convert` | zarr ↔ h5, `.zarr` ↔ `.zarr.zip` | a backend | +| `merge` | backends side by side, batch-labelled; `common_rows` keeps the rows all share | a backend and its `.batch.gz` | +| `notebook` | the agent's plan, observations and conclusions (from legume-mcp) | the day's `notebook.md` | + +No tool replaces a file: an output that is an input, or anything named after the output beside it, is refused, and `squeeze`, which rewrites its input, is not offered. + +### Records + +The server keeps the project's records the way [legume-mcp](https://crates.io/crates/legume-mcp) lays them out, after Noble's guide to organizing computational biology projects [1]: + +- outputs the agent does not name go in `results//`, named after their input and the tool (`pbmc.imported.qc.subset.zarr.zip`); +- every step that writes goes into the day's `results//runall.sh`, with paths relative to the project and steps skipped when their outputs exist, so a copy of the project replays the session; +- `results/INDEX.tsv` lists each output with a one-line summary; +- `import` notes where its input came from (the agent's `source`: a URL, an accession, a version) in `data/README.md`; +- the agent keeps `results//notebook.md` through the `notebook` tool. + +`LEGUME_MCP_RECORDS=off` keeps no records. + +[1] W. S. Noble, "A quick guide to organizing computational biology projects", PLoS Comput Biol 5(7): e1000424 (2009). doi:[10.1371/journal.pcbi.1000424](https://doi.org/10.1371/journal.pcbi.1000424) + +### Watching + +```sh +data-beans mcp watch +``` + +follows a running data-beans server in a detachable session: the agent's calls as they come, each one's log and result, the notebook. `s` stops the selected call, `n` leaves the agent a note (it arrives with its next result, and goes into the notebook), `o` opens a backend the call wrote in the statistics explorer, `F` follows again after you move. `--watch-only` offers no stop or note. With several servers running it asks which; `data-beans mcp watch ` names one. `LEGUME_MCP_WATCH=off` keeps a server from being watched. + +### The same on the command line + +```sh +data-beans qc data.zarr.zip --qc-mito-pattern '(?i)^MT-' # summary only +data-beans qc data.zarr.zip --qc-mito-pattern '(?i)^MT-' -o clean # and the kept cells +data-beans plot qc data.zarr.zip --qc-mito-pattern '(?i)^MT-' -o qc.png +data-beans plot stat data.zarr.zip -s row --stat sum -o genes.png +``` diff --git a/src/handlers/analysis.rs b/src/handlers/analysis.rs index 0d3346a..df4881e 100644 --- a/src/handlers/analysis.rs +++ b/src/handlers/analysis.rs @@ -12,7 +12,6 @@ use legume_numeric::matrix::common_io::*; use legume_numeric::matrix::membership::Membership; use legume_numeric::matrix::traits::RunningStatOps; use log::{info, warn}; -use regex::Regex; #[derive(ValueEnum, Clone, Debug, PartialEq)] #[clap(rename_all = "lowercase")] @@ -263,14 +262,9 @@ fn rows_matching(cmd_args: &RunStatArgs, data: &SparseIoVec) -> anyhow::Result = row_names - .iter() - .enumerate() - .filter_map(|(i, name)| re.is_match(name).then_some(i)) - .collect(); + 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, @@ -285,13 +279,17 @@ fn rows_matching(cmd_args: &RunStatArgs, data: &SparseIoVec) -> anyhow::Result &'static str { + pub(crate) fn metric(&self) -> &'static str { match self { HistStat::Nnz => "nnz", HistStat::Sum => "sum", @@ -302,7 +300,7 @@ impl HistStat { /// Pull the chosen statistic out of a running-stat accumulator (works for /// both the row and column variants via the shared `RunningStatOps` trait). - fn values>>(&self, s: &S) -> Vec { + pub(crate) fn values>>(&self, s: &S) -> Vec { match self { HistStat::Nnz => s.count_positives(), HistStat::Sum => s.sum(), diff --git a/src/handlers/cell_qc.rs b/src/handlers/cell_qc.rs new file mode 100644 index 0000000..0a7c5ce --- /dev/null +++ b/src/handlers/cell_qc.rs @@ -0,0 +1,300 @@ +//! `data-beans qc`: senna's and pinto's cell QC (`data_beans::qc_lib`, the +//! same flags) run on one backend, its outcome summarized on stdout, and the +//! cells (and features) it keeps written to a new backend when asked. +//! +//! It works in the library's types, since `qc_lib` does: the binary's own +//! copies of the modules are other types. + +use clap::Args; +use data_beans::hdf5_io::resolve_backend_file; +use data_beans::qc_lib::{ + compute_qc, robust_band, write_qc_report, QcArgs, QcConfig, QcReport, RobustBand, +}; +use data_beans::sparse_io::{open_sparse_matrix, SparseIo, SparseIoBackend}; +use data_beans::sparse_io_vector::SparseIoVec; +use data_beans::zarr_io::{finalize_output, output_target, prepare_output}; +use legume_numeric::matrix::common_io::write_lines; +use log::info; +use std::sync::Arc; + +#[derive(Args, Debug)] +pub struct CellQcArgs { + /// data file -- `.zarr`, `.zarr.zip`, or `.h5` + pub data_file: Box, + + #[command(flatten)] + pub qc: QcArgs, + + /// output: the cells (and features) QC keeps, as a new backend; without + /// it, only the summary is printed + #[arg(short, long)] + pub output: Option>, + + /// backend for the output file + #[arg(long, value_enum, default_value = "zarr")] + pub backend: SparseIoBackend, + + /// keep a `.zarr` directory instead of producing a `.zarr.zip` archive + #[arg(long = "no-zip", default_value_t = true, action = clap::ArgAction::SetFalse)] + pub zip: bool, + + /// cells per rayon job (omit for auto-scaling by feature count) + #[arg(long)] + pub block_size: Option, +} + +/// The mito pattern `qc` and `plot qc` start from, where senna's and +/// pinto's QC has none; an empty `--qc-mito-pattern=` turns it off. +pub const MITO: &str = "(?i)^MT-"; + +/// A backend's QC: what [`compute_qc`] made of it, under which config. +pub struct CellQc { + pub matrix: Arc>>, + pub cfg: QcConfig, + pub report: QcReport, +} + +impl CellQc { + pub fn compute( + data_file: &str, + qc: &QcArgs, + block_size: Option, + ) -> anyhow::Result { + let mut qc = qc.clone(); + if qc.qc_mito_pattern.as_deref() == Some("") { + qc.qc_mito_pattern = None; + } + let cfg = qc + .to_config() + .ok_or_else(|| anyhow::anyhow!("--no-qc leaves nothing to do"))?; + let (backend, file) = resolve_backend_file(data_file, None)?; + let matrix: Arc>> = + Arc::from(open_sparse_matrix(&file, &backend)?); + let mut data = SparseIoVec::new(); + data.push(matrix.clone(), None)?; + let report = compute_qc(&data, &cfg, block_size)?; + Ok(Self { + matrix, + cfg, + report, + }) + } + + /// The per-cell metrics QC measured, in a fixed order. + pub fn metrics(&self) -> Vec> { + let (r, cfg) = (&self.report, &self.cfg); + let mut v = vec![ + Metric::counts("n_genes", "detected genes", &r.n_genes, cfg.mad_on_n_genes), + Metric::counts( + "total_counts", + "total counts", + &r.total_counts, + cfg.mad_on_counts, + ), + ]; + if let Some(f) = &r.mito_frac { + v.push(Metric::fraction( + ("mito_frac", "mitochondrial fraction"), + f, + cfg.mad_on_mito, + cfg.mito_max_frac, + )); + } + if let Some(f) = &r.ribo_frac { + v.push(Metric::fraction( + ("ribo_frac", "ribosomal fraction"), + f, + false, + cfg.ribo_max_frac, + )); + } + v + } + + /// The robust band of `m`, on its own scale, as qc_lib fits it: over + /// the cells above the near-empty floor. + pub fn band(&self, m: &Metric) -> Option { + let consider: Vec = self.report.near_empty.iter().map(|&e| !e).collect(); + robust_band(m.values, self.cfg.n_mads, m.count, Some(&consider)).map(|b| b.raw(m.count)) + } + + /// The edge of `m`'s band QC drops cells past, when it does: its floor + /// for a count, its ceiling for a fraction. + pub fn edge(&self, m: &Metric) -> Option { + let band = self.band(m).filter(|_| m.banded)?; + Some(if m.count { band.lo } else { band.hi }) + } +} + +/// A per-cell metric: a count, banded on the log1p scale and held from +/// below, or a fraction, banded as is and held from above (and capped). +pub struct Metric<'a> { + pub name: &'static str, + /// What a figure calls it. + #[cfg_attr(not(feature = "plot"), allow(dead_code))] + pub label: &'static str, + pub values: &'a [f32], + pub count: bool, + /// Whether QC drops cells outside its band. + pub banded: bool, + pub cap: Option, +} + +impl<'a> Metric<'a> { + fn counts(name: &'static str, label: &'static str, values: &'a [f32], banded: bool) -> Self { + Self { + name, + label, + values, + count: true, + banded, + cap: None, + } + } + + fn fraction( + (name, label): (&'static str, &'static str), + values: &'a [f32], + banded: bool, + cap: Option, + ) -> Self { + Self { + name, + label, + values, + count: false, + banded, + cap, + } + } + + /// What its edge is, and the side of it a dropped cell is on. + pub fn bound(&self) -> (&'static str, &'static str) { + if self.count { + ("min", "below") + } else { + ("max", "above") + } + } + + /// Whether `v` is on the kept side of `edge`. + pub fn within(&self, v: f32, edge: f32) -> bool { + if self.count { + v >= edge + } else { + v <= edge + } + } + + /// `v` shown past the f32 noise of the log1p round trip: counts to a + /// tenth, fractions to four places. + fn shown(&self, v: f32) -> String { + if self.count { + format!("{v:.1}") + } else { + format!("{v:.4}") + } + } +} + +pub fn run_cell_qc(args: &CellQcArgs) -> anyhow::Result<()> { + let qc = CellQc::compute(&args.data_file, &args.qc, args.block_size)?; + let (matrix, report) = (&qc.matrix, &qc.report); + // The cell names, read only when something is written with them. + let cells = match args.qc.qc_report.is_some() || args.output.is_some() { + true => matrix.column_names()?, + false => vec![], + }; + if let Some(path) = &args.qc.qc_report { + write_qc_report(path, &cells, report)?; + info!("wrote the per-cell report: {path}"); + } + + let keep_cells = report.emit_idx_unmasked(); + let keep_rows: Vec = (0..report.feature_keep.len()) + .filter(|&r| report.feature_keep[r]) + .collect(); + let summary = summary(&qc, keep_cells.len(), keep_rows.len()); + write_lines(&summary, "stdout")?; + + let Some(output) = &args.output else { + return Ok(()); + }; + anyhow::ensure!(!keep_cells.is_empty(), "QC keeps no cells; nothing written"); + // The output is cleared before the input is read: it must not be it. + let target = output_target(output, args.backend.clone(), args.zip)?; + for p in [&target.effective, &target.working, &target.found] { + anyhow::ensure!( + !crate::handlers::merging::same_file(&args.data_file, p), + "the output {p} is the input {}: give -o another name", + args.data_file + ); + } + let rows = matrix.row_names()?; + let out_rows: Vec> = keep_rows.iter().map(|&r| rows[r].clone()).collect(); + let out_cells: Vec> = keep_cells.iter().map(|&c| cells[c].clone()).collect(); + let (effective, backend_out, file_out) = + prepare_output(output, args.backend.clone(), args.zip)?; + let row_filter = (keep_rows.len() < rows.len()).then_some(keep_rows.as_slice()); + let (_, _, nnz) = data_beans::column_subset::stream_column_selection( + &**matrix, + &keep_cells, + row_filter, + &out_rows, + &out_cells, + &file_out, + &backend_out, + )?; + let written = finalize_output(&file_out, &effective)?; + info!( + "done: {written} ({} features x {} cells, {nnz} non-zeros)", + out_rows.len(), + out_cells.len() + ); + Ok(()) +} + +/// The outcome, as `key:\tvalue` lines: the counts kept and dropped, and for +/// each metric its median, the edge QC held it to, and how many cells fell +/// past it (a cell can fall past several). +fn summary(qc: &CellQc, kept_cells: usize, kept_features: usize) -> Vec> { + let report = &qc.report; + let mut lines: Vec> = vec![]; + let mut say = |key: &str, value: String| lines.push(format!("{key}:\t{value}").into()); + let near_empty = report.near_empty.iter().filter(|&&e| e).count(); + say("cells", report.train_keep.len().to_string()); + say("cells_kept", kept_cells.to_string()); + say("near_empty", near_empty.to_string()); + // Dropped as outliers, apart from the near-empty: what the figure shows. + let outliers = (report.near_empty.iter().zip(&report.train_keep)) + .filter(|&(&e, &k)| !e && !k) + .count(); + say("outliers", outliers.to_string()); + say("features", report.feature_keep.len().to_string()); + say("features_kept", kept_features.to_string()); + + // Cells above the near-empty floor past `edge`. + let past = |m: &Metric, edge: f32| { + m.values + .iter() + .zip(&report.near_empty) + .filter(|&(&v, &e)| !e && !m.within(v, edge)) + .count() + }; + for m in qc.metrics() { + let name = m.name; + if let Some(band) = qc.band(&m) { + say(&format!("{name}.median"), m.shown(band.median)); + } + if let Some(edge) = qc.edge(&m) { + let (bound, side) = m.bound(); + say(&format!("{name}.{bound}"), m.shown(edge)); + say(&format!("{name}.{side}"), past(&m, edge).to_string()); + } + if let Some(cap) = m.cap { + say(&format!("{name}.cap"), cap.to_string()); + say(&format!("{name}.above_cap"), past(&m, cap).to_string()); + } + } + lines +} diff --git a/src/handlers/inspection.rs b/src/handlers/inspection.rs index 00d1c64..3d45209 100644 --- a/src/handlers/inspection.rs +++ b/src/handlers/inspection.rs @@ -4,7 +4,7 @@ use crate::interactive::stat_tui::Side; use crate::interactive::ui::Tui; use crate::sparse_io::*; use crate::utilities::io_helpers::{ - parse_index_spec, read_col_names, MAX_COLUMN_NAME_IDX, MAX_ROW_NAME_IDX, + parse_index_spec_within, read_col_names, MAX_COLUMN_NAME_IDX, MAX_ROW_NAME_IDX, }; use crate::utilities::name_matching::match_by_substring; use clap::Args; @@ -225,9 +225,7 @@ pub fn take_columns(args: &TakeColumnsArgs) -> anyhow::Result<()> { let column_names: Vec> = columns.iter().map(|&i| names[i].clone()).collect(); (data.read_columns_dmatrix(columns)?, column_names) } else if let Some(columns) = columns { - let columns = parse_index_spec(&columns)?; - let n_columns = data.num_columns().unwrap_or(0); - let columns: Vec = columns.into_iter().filter(|&i| i < n_columns).collect(); + let columns = parse_index_spec_within(&columns, data.num_columns().unwrap_or(0))?; if columns.is_empty() { return Err(anyhow::anyhow!("invalid indexes")); @@ -308,9 +306,7 @@ pub fn take_rows(args: &TakeRowsArgs) -> anyhow::Result<()> { let row_names: Vec> = rows.iter().map(|&i| names[i].clone()).collect(); (data_backend.read_rows_dmatrix(rows)?, row_names) } else if let Some(rows) = rows { - let rows = parse_index_spec(&rows)?; - let n_rows = data_backend.num_rows().unwrap_or(0); - let rows: Vec = rows.into_iter().filter(|&i| i < n_rows).collect(); + let rows = parse_index_spec_within(&rows, data_backend.num_rows().unwrap_or(0))?; if rows.is_empty() { return Err(anyhow::anyhow!("invalid indexes")); diff --git a/src/handlers/merging.rs b/src/handlers/merging.rs index 6f889b7..326ce2f 100644 --- a/src/handlers/merging.rs +++ b/src/handlers/merging.rs @@ -68,6 +68,16 @@ pub struct MergeBackendArgs { #[arg(long = "no-zip", default_value_t = true, action = clap::ArgAction::SetFalse)] pub zip: bool, + #[arg( + long, + default_value_t = false, + help = "Keep only the rows every backend has", + long_help = "Keep only the rows (features) every backend has, matched by name,\n\ + in the first backend's order. Without it, the backends must have\n\ + the same rows in the same order." + )] + pub common_rows: bool, + #[arg( long, default_value_t = false, @@ -285,16 +295,21 @@ pub fn run_merge_backend(args: &MergeBackendArgs) -> anyhow::Result<()> { path: Box, backend: SparseIoBackend, ncol: usize, + nnz: u64, col_offset: u64, - nnz_offset: u64, + /// Under `--common-rows`, each row's place among the common rows. + row_map: Option>>, } let mut batches: Vec = Vec::with_capacity(num_batches); let mut row_names: Vec> = vec![]; let mut column_names: Vec> = vec![]; let mut column_batch_names: Vec> = vec![]; + // Under `--common-rows`: each batch's row names, and its row index to + // count the common rows' entries. + let mut batch_rows: Vec>> = vec![]; + let mut batch_indptr: Vec> = vec![]; let mut col_offset: u64 = 0; - let mut nnz_offset: u64 = 0; let mut metas: Vec = Vec::with_capacity(num_batches); for (batch_idx, data_file) in args.data_files.iter().enumerate() { @@ -315,14 +330,17 @@ pub fn run_merge_backend(args: &MergeBackendArgs) -> anyhow::Result<()> { .num_non_zeros() .ok_or_else(|| anyhow::anyhow!("missing nnz in {}", data_file))?; - if batch_idx == 0 { + if args.common_rows { + batch_rows.push(data.row_names()?); + batch_indptr.push(data.row_indptr().to_vec()); + } else if batch_idx == 0 { row_names = data.row_names()?; } else { info!("checking if the row names are consistent"); anyhow::ensure!( row_names == data.row_names()?, "{} has other row names than {}: merge-backend needs the same rows, in order; \ - subset them to common rows first", + subset them to common rows first, or pass --common-rows", data_file, args.data_files[0] ); @@ -341,16 +359,53 @@ pub fn run_merge_backend(args: &MergeBackendArgs) -> anyhow::Result<()> { path: data_file.clone(), backend, ncol, + nnz: nnz as u64, col_offset, - nnz_offset, + row_map: None, }); col_offset += ncol as u64; - nnz_offset += nnz as u64; } + if args.common_rows { + row_names = common_names(&batch_rows); + anyhow::ensure!(!row_names.is_empty(), "the backends share no row names"); + info!("{} rows common to all backends", row_names.len()); + let place: HashMap<&str, u64> = row_names + .iter() + .enumerate() + .map(|(i, n)| (n.as_ref(), i as u64)) + .collect(); + for ((h, rows), indptr) in batches.iter_mut().zip(&batch_rows).zip(&batch_indptr) { + let mut seen = vec![false; row_names.len()]; + // A name a backend repeats keeps its first row only. + let map: Vec> = rows + .iter() + .map(|n| { + let i = *place.get(n.as_ref())?; + (!std::mem::replace(&mut seen[i as usize], true)).then_some(i) + }) + .collect(); + if map.iter().enumerate().all(|(r, m)| *m == Some(r as u64)) { + continue; // these are the common rows, in order + } + // The rows kept hold this many entries: the row index says. + anyhow::ensure!( + indptr.len() == rows.len() + 1, + "{} has no by-row index to count its common rows' entries", + h.path + ); + h.nnz = map + .iter() + .enumerate() + .filter(|(_, m)| m.is_some()) + .map(|(r, _)| indptr[r + 1] - indptr[r]) + .sum(); + h.row_map = Some(map); + } + } let total_ncol = col_offset as usize; - let total_nnz = nnz_offset as usize; + let total_nnz = batches.iter().map(|h| h.nnz).sum::() as usize; let total_nrow = row_names.len(); info!( "Found {} columns/barcodes across {} batches (total nnz = {})", @@ -369,6 +424,7 @@ pub fn run_merge_backend(args: &MergeBackendArgs) -> anyhow::Result<()> { out.begin_streaming_csc((total_nrow, total_ncol, total_nnz))?; let prog_bar = styled_progress_bar(num_batches as u64, "batches"); + let mut nnz_offset: u64 = 0; for h in &batches { let src = open_sparse_matrix(&h.path, &h.backend)?; let jobs = create_jobs(h.ncol, total_nrow, args.block_size); @@ -382,8 +438,15 @@ pub fn run_merge_backend(args: &MergeBackendArgs) -> anyhow::Result<()> { .map(|(lb, ub)| -> anyhow::Result> { let (_, _, t) = src.read_triplets_by_columns((*lb..*ub).collect())?; let lb_u64 = *lb as u64; + let row_map = h.row_map.as_deref(); Ok(t.into_iter() - .map(move |(i, j_local, x)| (i, j_local + lb_u64, x)) + .filter_map(move |(i, j_local, x)| { + let i = match row_map { + Some(m) => m[i as usize]?, + None => i, + }; + Some((i, j_local + lb_u64, x)) + }) .collect()) }) .collect::>>()? @@ -394,11 +457,12 @@ pub fn run_merge_backend(args: &MergeBackendArgs) -> anyhow::Result<()> { let (local_colptr, row_indices, values) = build_csc_slab(triplets_batch, h.ncol); out.append_csc_slab( h.col_offset, - h.nnz_offset, + nnz_offset, &local_colptr, &row_indices, &values, )?; + nnz_offset += h.nnz; prog_bar.inc(1); } @@ -453,9 +517,15 @@ pub fn run_merge_backend(args: &MergeBackendArgs) -> anyhow::Result<()> { Ok(()) } +/// A merge's batch labels: `{output}.batch.gz`, the output without its +/// backend suffix. +pub(crate) fn batch_file(effective_output: &str) -> String { + format!("{}.batch.gz", strip_backend_suffix(effective_output)) +} + /// Whether paths `a` and `b` name the same file: compared resolved when /// they exist, as written (made absolute) when not. -fn same_file(a: &str, b: &str) -> bool { +pub(crate) fn same_file(a: &str, b: &str) -> bool { let resolve = |p: &str| std::fs::canonicalize(p).or_else(|_| std::path::absolute(p)); matches!((resolve(a), resolve(b)), (Ok(x), Ok(y)) if x == y) } @@ -472,7 +542,8 @@ pub(crate) fn write_batches_and_finalize( batch_map: &HashMap, Box>, ) -> anyhow::Result<()> { let stem = strip_backend_suffix(effective_output); - // The name itself, dots and all (`sample.v2`), not cut at its last dot. + // The output's name itself, dots and all (`sample.v2`), not cut at its + // last dot. let default_batch: Box = std::path::Path::new(stem) .file_name() .map_or_else(|| stem.into(), |n| n.to_string_lossy().into()); @@ -481,10 +552,27 @@ pub(crate) fn write_batches_and_finalize( .iter() .map(|k| batch_map.get(k).unwrap_or(&default_batch).clone()) .collect(); - write_lines(&labels, &format!("{stem}.batch.gz"))?; + write_lines(&labels, &batch_file(effective_output))?; finalize_zarr_output(backend_file, effective_output) } +/// The names in every list, in the first list's order, each once. +fn common_names(lists: &[Vec>]) -> Vec> { + let Some((first, rest)) = lists.split_first() else { + return vec![]; + }; + let others: Vec> = rest + .iter() + .map(|l| l.iter().map(|n| n.as_ref()).collect()) + .collect(); + let mut seen = rustc_hash::FxHashSet::default(); + first + .iter() + .filter(|&n| others.iter().all(|o| o.contains(n.as_ref())) && seen.insert(n.as_ref())) + .cloned() + .collect() +} + pub fn run_merge_mtx(args: &MergeMtxArgs) -> anyhow::Result<()> { let directories = args.data_directories.clone(); diff --git a/src/handlers/merging_tests.rs b/src/handlers/merging_tests.rs index 0579bd6..e41d126 100644 --- a/src/handlers/merging_tests.rs +++ b/src/handlers/merging_tests.rs @@ -47,6 +47,7 @@ fn merge_args(files: &[&str], output: &str) -> MergeBackendArgs { backend: SparseIoBackend::Zarr, output: output.into(), zip: false, + common_rows: false, do_squeeze: false, row_nnz_cutoff: 0, column_nnz_cutoff: 0, diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs index e6c7450..8bb7c84 100644 --- a/src/handlers/mod.rs +++ b/src/handlers/mod.rs @@ -1,8 +1,11 @@ pub mod analysis; pub mod builders; +pub mod cell_qc; pub mod explore; pub mod exporters; pub mod inspection; pub mod listing; pub mod merging; +#[cfg(feature = "plot")] +pub mod plotting; pub mod transformation; diff --git a/src/handlers/plotting.rs b/src/handlers/plotting.rs new file mode 100644 index 0000000..16a93bc --- /dev/null +++ b/src/handlers/plotting.rs @@ -0,0 +1,605 @@ +//! `data-beans plot`: figures of a backend's QC and statistics, drawn as an +//! SVG with axes and rendered by legume-plot (PNG through resvg, PDF through +//! svg2pdf). Scatter points are rasterized, as legume-plot does, so a +//! million cells make an image, not a million SVG elements. + +use crate::handlers::analysis::{HistStat, StatDim}; +use crate::handlers::cell_qc::CellQc; +use crate::hdf5_io::resolve_backend_file; +use crate::qc::{collect_column_stat, collect_row_stat}; +use crate::sparse_io::open_sparse_matrix; +use base64::Engine; +use clap::{Args, Subcommand}; +use data_beans::qc_lib::QcArgs; +use legume_plot::rasterize::{rasterize_group_png, Extent, PointShape, RadiusSpec}; +use legume_plot::{escape_xml as escape, FigureFormats, Rgb}; +use log::info; +use std::fmt::Write as _; +use std::path::Path; + +#[derive(Args, Debug)] +pub struct PlotArgs { + #[command(subcommand)] + pub kind: PlotKind, +} + +#[derive(Subcommand, Debug)] +pub enum PlotKind { + /// Cell QC: each metric's histogram with the edge QC held it to, and + /// total counts against detected genes, cells coloured by verdict + #[command( + mut_args(|a| a.hide(false)), + mut_arg("qc_mito_pattern", |a| a.default_value(crate::handlers::cell_qc::MITO)) + )] + Qc(PlotQcArgs), + /// The histogram of one statistic over the rows or the columns + Stat(PlotStatArgs), +} + +#[derive(Args, Debug)] +pub struct PlotQcArgs { + /// data file -- `.zarr`, `.zarr.zip`, or `.h5` + pub data_file: Box, + + #[command(flatten)] + pub qc: QcArgs, + + /// figure file: `.png`, `.svg` or `.pdf` + #[arg(short, long)] + pub output: Box, + + /// cells per rayon job (omit for auto-scaling by feature count) + #[arg(long)] + pub block_size: Option, +} + +#[derive(Args, Debug)] +pub struct PlotStatArgs { + /// data file -- `.zarr`, `.zarr.zip`, or `.h5` + pub data_file: Box, + + /// over rows (features) or columns (cells) + #[arg(short = 's', long, value_enum, default_value = "column")] + pub stat_dim: StatDim, + + /// the statistic, on a log10(x+1) axis + #[arg(long, value_enum, default_value = "nnz")] + pub stat: HistStat, + + /// figure file: `.png`, `.svg` or `.pdf` + #[arg(short, long)] + pub output: Box, + + /// cells per rayon job (omit for auto-scaling by feature count) + #[arg(long)] + pub block_size: Option, +} + +const KEPT: Rgb = (59, 111, 182); +const OUTLIER: Rgb = (209, 73, 91); +const NEAR_EMPTY: Rgb = (154, 154, 154); +const EDGE: Rgb = (209, 73, 91); +const CAP: Rgb = (224, 138, 30); +const INK: &str = "#333"; +const FONT: &str = "Helvetica, Arial, sans-serif"; + +const PANEL_W: f32 = 460.0; +const PANEL_H: f32 = 340.0; +const TITLE_H: f32 = 44.0; + +pub fn run_plot(args: &PlotArgs) -> anyhow::Result<()> { + let (fig, output) = match &args.kind { + PlotKind::Qc(a) => { + let qc = CellQc::compute(&a.data_file, &a.qc, a.block_size)?; + (qc_figure(&qc, &a.data_file)?, &a.output) + } + PlotKind::Stat(a) => (stat_figure(a)?, &a.output), + }; + write_figure(fig, output)?; + info!("wrote {output}"); + Ok(()) +} + +/// Write `fig` as the format `output`'s extension names. +fn write_figure(fig: Figure, output: &str) -> anyhow::Result<()> { + let (w, h) = (fig.w, fig.h); + let path = Path::new(output); + let ext = path + .extension() + .and_then(|e| e.to_str()) + .unwrap_or_default(); + let want = FigureFormats { + svg: ext == "svg", + png: ext == "png", + pdf: ext == "pdf", + }; + anyhow::ensure!( + want.svg || want.png || want.pdf, + "{output}: name a .png, .svg or .pdf file" + ); + let base = path.with_extension(""); + legume_plot::write_figure(&fig.finish(), w, h, &base.to_string_lossy(), want)?; + Ok(()) +} + +/// The histogram of one row or column statistic. +fn stat_figure(a: &PlotStatArgs) -> anyhow::Result
{ + let (backend, file) = resolve_backend_file(&a.data_file, None)?; + let data = open_sparse_matrix(&file, &backend)?; + let (what, values) = match a.stat_dim { + StatDim::Row => ( + "row", + a.stat + .values(&collect_row_stat(data.as_ref(), a.block_size)?), + ), + StatDim::Column => ( + "column", + a.stat + .values(&collect_column_stat(data.as_ref(), a.block_size)?), + ), + }; + let label = format!("{} per {what}", a.stat.metric()); + let mut fig = Figure::new(1, 1, &format!("{}: {label}", a.data_file)); + let all = Series { + label: format!("{} {what}s", values.len()), + color: KEPT, + values, + }; + fig.histogram(0, (&label, &format!("{what}s")), &[all], true, &[]); + Ok(fig) +} + +/// The QC figure: a histogram per metric, stacked by verdict, with its +/// band edge and cap; then total counts against detected genes. +fn qc_figure(qc: &CellQc, title: &str) -> anyhow::Result
{ + let r = &qc.report; + let n = r.train_keep.len(); + let near_empty: Vec = (0..n).filter(|&c| r.near_empty[c]).collect(); + let outlier: Vec = (0..n) + .filter(|&c| !r.near_empty[c] && !r.train_keep[c]) + .collect(); + let kept = r.emit_idx_unmasked(); + let verdicts: [(String, Rgb, &Vec); 3] = [ + (format!("kept {}", kept.len()), KEPT, &kept), + (format!("outlier {}", outlier.len()), OUTLIER, &outlier), + ( + format!("near-empty {}", near_empty.len()), + NEAR_EMPTY, + &near_empty, + ), + ]; + // Each verdict's cells, as what `of` makes of a cell. + fn by_verdict( + verdicts: &[(String, Rgb, &Vec)], + of: impl Fn(usize) -> T, + ) -> Vec> { + verdicts + .iter() + .map(|(label, color, idx)| Series { + label: label.clone(), + color: *color, + values: idx.iter().map(|&c| of(c)).collect(), + }) + .collect() + } + + let metrics = qc.metrics(); + let panels = metrics.len() + 1; + let cols = if panels > 4 { 3 } else { 2 }; + let mut fig = Figure::new( + cols, + panels.div_ceil(cols), + &format!( + "{title}: {n} cells, {} kept, {} outliers, {} near-empty (bands {} MADs)", + kept.len(), + outlier.len(), + near_empty.len(), + qc.cfg.n_mads + ), + ); + for (i, m) in metrics.iter().enumerate() { + let mut marks = vec![]; + if let Some(edge) = qc.edge(m) { + marks.push(Mark { + at: edge, + color: EDGE, + label: m.bound().0, + }); + } + if let Some(cap) = m.cap { + marks.push(Mark { + at: cap, + color: CAP, + label: "cap", + }); + } + let groups = by_verdict(&verdicts, |c| m.values[c]); + fig.histogram(i, (m.label, "cells"), &groups, m.count, &marks); + } + let layers = by_verdict(&verdicts, |c| (r.total_counts[c], r.n_genes[c])); + fig.scatter(metrics.len(), "total counts", "detected genes", &layers)?; + Ok(fig) +} + +/// Values (or points) drawn in one colour. +struct Series { + label: String, + color: Rgb, + values: Vec, +} + +/// A vertical line at a value, labelled. +struct Mark { + at: f32, + color: Rgb, + label: &'static str, +} + +/// `v` on a log10(x+1) axis when `log`, else as is. +fn scale(v: f32, log: bool) -> f32 { + if log { + (v.max(0.0) + 1.0).log10() + } else { + v + } +} + +/// An axis: the data range it shows, on a log10(x+1) scale or as is. +#[derive(Clone, Copy)] +struct Axis { + lo: f32, + hi: f32, + log: bool, +} + +impl Axis { + fn fit(values: impl Iterator, log: bool) -> Self { + let t = |v: f32| scale(v, log); + let (mut lo, mut hi) = (f32::INFINITY, f32::NEG_INFINITY); + for v in values.filter(|v| v.is_finite()) { + lo = lo.min(t(v)); + hi = hi.max(t(v)); + } + if !lo.is_finite() { + (lo, hi) = (0.0, 1.0); + } + if hi - lo < 1e-6 { + hi = lo + 1.0; + } + let pad = (hi - lo) * 0.03; + Self { + lo: lo - pad, + hi: hi + pad, + log, + } + } + + fn t(self, v: f32) -> f32 { + scale(v, self.log) + } + + /// Where `v` falls, 0 at `lo` to 1 at `hi`. + fn frac(self, v: f32) -> f32 { + (self.t(v) - self.lo) / (self.hi - self.lo) + } + + /// Tick values (on the data's own scale) and their labels: powers of + /// ten on a log axis, steps of 1, 2 or 5 times a power of ten otherwise. + fn ticks(self) -> Vec<(f32, String)> { + if self.log { + // 0, then 1, 2, 5 times each power of ten; powers alone when + // those crowd. + let within = |x: f32| (self.lo..=self.hi).contains(&self.t(x)); + let pick = |steps: &[f32]| -> Vec { + let mut v = vec![0.0]; + let mut p = 1.0f32; + while self.t(p) <= self.hi { + v.extend(steps.iter().map(|m| m * p)); + p *= 10.0; + } + v.into_iter().filter(|&x| within(x)).collect() + }; + let mut v = pick(&[1.0, 2.0, 5.0]); + if v.len() > 8 { + v = pick(&[1.0]); + } + return v.into_iter().map(|x| (x, short(x))).collect(); + } + let span = self.hi - self.lo; + let raw = span / 5.0; + let mag = 10f32.powf(raw.log10().floor()); + let step = [1.0, 2.0, 5.0, 10.0] + .iter() + .map(|m| m * mag) + .find(|s| span / s <= 6.0) + .unwrap_or(10.0 * mag); + let mut v = vec![]; + let mut x = (self.lo / step).ceil() * step; + while x <= self.hi { + // Snap the float drift around zero (-0.0, 1e-9). + let shown = if x.abs() < step * 1e-3 { 0.0 } else { x }; + v.push((shown, short(shown))); + x += step; + } + v + } +} + +/// A number for a tick label: 1k, 2.5M, 0.05. +fn short(x: f32) -> String { + let a = x.abs(); + let trim = |s: String| { + if s.contains('.') { + s.trim_end_matches('0').trim_end_matches('.').to_string() + } else { + s + } + }; + if a >= 1e6 { + format!("{}M", trim(format!("{:.1}", x / 1e6))) + } else if a >= 1e3 { + format!("{}k", trim(format!("{:.1}", x / 1e3))) + } else if a >= 1.0 || a == 0.0 { + trim(format!("{x:.1}")) + } else { + trim(format!("{x:.3}")) + } +} + +fn rgb((r, g, b): Rgb) -> String { + format!("rgb({r},{g},{b})") +} + +/// A grid of panels, written as SVG. +struct Figure { + svg: String, + cols: usize, + w: u32, + h: u32, +} + +/// A panel's plot area, in figure pixels. +#[derive(Clone, Copy)] +struct Area { + x: f32, + y: f32, + w: f32, + h: f32, +} + +impl Figure { + fn new(cols: usize, rows: usize, title: &str) -> Self { + let w = (cols as f32 * PANEL_W) as u32; + let h = (rows as f32 * PANEL_H + TITLE_H) as u32; + let mut svg = format!( + "\n\ + \n" + ); + let _ = writeln!( + svg, + "{}", + escape(title) + ); + Self { svg, cols, w, h } + } + + /// The plot area of panel `i`, inside its margins. + fn area(&self, i: usize) -> Area { + let (col, row) = (i % self.cols, i / self.cols); + let (left, right, top, bottom) = (64.0, 18.0, 30.0, 48.0); + Area { + x: col as f32 * PANEL_W + left, + y: TITLE_H + row as f32 * PANEL_H + top, + w: PANEL_W - left - right, + h: PANEL_H - top - bottom, + } + } + + /// Axes, ticks and labels around `a`. + fn axes(&mut self, a: Area, x: Axis, y: Axis, xlabel: &str, ylabel: &str) { + let s = &mut self.svg; + let (bx, by) = (a.x, a.y + a.h); + let _ = writeln!( + s, + "", + a.x, a.y, a.w, a.h + ); + for (v, label) in x.ticks() { + let px = a.x + x.frac(v) * a.w; + let _ = writeln!( + s, + "\ + {label}", + by + 4.0, + by + 16.0 + ); + } + for (v, label) in y.ticks() { + let py = by - y.frac(v) * a.h; + let _ = writeln!( + s, + "\ + {label}", + bx - 4.0, + bx - 6.0, + py + 4.0 + ); + } + let _ = writeln!( + s, + "{}", + a.x + a.w / 2.0, + by + 34.0, + escape(xlabel) + ); + let (lx, ly) = (a.x - 46.0, a.y + a.h / 2.0); + let _ = writeln!( + s, + "{}", + escape(ylabel) + ); + } + + /// A legend in a row above `a`. + fn legend(&mut self, a: Area, entries: &[(&str, Rgb)]) { + let mut x = a.x; + let y = a.y - 20.0; + for (label, color) in entries { + let _ = writeln!( + self.svg, + "\ + {}", + rgb(*color), + x + 14.0, + y + 9.0, + escape(label) + ); + x += 24.0 + 6.5 * label.chars().count() as f32; + } + } + + /// Panel `i`: the histogram of `groups`' values, stacked, on a log10(x+1) + /// axis when `log`, with `marks` drawn across it. + fn histogram( + &mut self, + i: usize, + (xlabel, ylabel): (&str, &str), + groups: &[Series], + log: bool, + marks: &[Mark], + ) { + const BINS: usize = 50; + let a = self.area(i); + let all = groups.iter().flat_map(|g| g.values.iter().copied()); + let x = Axis::fit(all.chain(marks.iter().map(|m| m.at)), log); + let bin = |v: f32| -> Option { + let f = x.frac(v); + (f.is_finite() && (0.0..=1.0).contains(&f)) + .then(|| ((f * BINS as f32) as usize).min(BINS - 1)) + }; + let counts: Vec<[usize; BINS]> = groups + .iter() + .map(|g| { + let mut c = [0usize; BINS]; + for b in g.values.iter().filter_map(|&v| bin(v)) { + c[b] += 1; + } + c + }) + .collect(); + let top = (0..BINS) + .map(|b| counts.iter().map(|c| c[b]).sum::()) + .max() + .unwrap_or(0) + .max(1); + let y = Axis { + lo: 0.0, + hi: top as f32 * 1.05, + log: false, + }; + let bw = a.w / BINS as f32; + for b in 0..BINS { + let mut base = 0usize; + for (g, c) in groups.iter().zip(&counts) { + if c[b] == 0 { + continue; + } + let y0 = a.y + a.h - y.frac(base as f32) * a.h; + let y1 = a.y + a.h - y.frac((base + c[b]) as f32) * a.h; + let _ = writeln!( + self.svg, + "", + a.x + b as f32 * bw, + bw.max(0.5) - 0.3, + y0 - y1, + rgb(g.color) + ); + base += c[b]; + } + } + for (k, m) in marks.iter().enumerate() { + let px = a.x + x.frac(m.at) * a.w; + let _ = writeln!( + self.svg, + "\ + {} {}", + a.y, + a.y + a.h, + rgb(m.color), + px + 3.0, + a.y + 12.0 + 13.0 * k as f32, + rgb(m.color), + m.label, + short(m.at) + ); + } + self.axes(a, x, y, xlabel, ylabel); + if groups.len() > 1 { + let entries: Vec<(&str, Rgb)> = + groups.iter().map(|g| (g.label.as_str(), g.color)).collect(); + self.legend(a, &entries); + } + } + + /// Panel `i`: each layer's points, on log10(x+1) axes, rasterized. + fn scatter( + &mut self, + i: usize, + xlabel: &str, + ylabel: &str, + layers: &[Series<(f32, f32)>], + ) -> anyhow::Result<()> { + let a = self.area(i); + let all = || layers.iter().flat_map(|l| l.values.iter()); + let x = Axis::fit(all().map(|p| p.0), true); + let y = Axis::fit(all().map(|p| p.1), true); + let ext = Extent { + w: a.w.round() as u32, + h: a.h.round() as u32, + }; + let n: usize = layers.iter().map(|l| l.values.len()).sum(); + let radius = if n > 100_000 { 0.8 } else { 1.8 }; + for Series { + color, + values: points, + .. + } in layers + { + if points.is_empty() { + continue; + } + let px: Vec<(f32, f32)> = points + .iter() + .map(|&(u, v)| (x.frac(u) * a.w, a.h - y.frac(v) * a.h)) + .collect(); + let png = rasterize_group_png( + &px, + ext, + RadiusSpec::Scalar(radius), + *color, + 0.6, + PointShape::Circle, + )?; + let _ = writeln!( + self.svg, + "", + a.x, + a.y, + ext.w, + ext.h, + base64::engine::general_purpose::STANDARD.encode(png) + ); + } + self.axes(a, x, y, xlabel, ylabel); + let entries: Vec<(&str, Rgb)> = + layers.iter().map(|l| (l.label.as_str(), l.color)).collect(); + self.legend(a, &entries); + Ok(()) + } + + fn finish(mut self) -> String { + self.svg.push_str("\n"); + self.svg + } +} diff --git a/src/handlers/transformation.rs b/src/handlers/transformation.rs index 3249e95..d9f61d0 100644 --- a/src/handlers/transformation.rs +++ b/src/handlers/transformation.rs @@ -49,14 +49,20 @@ pub struct SubsetColumnsArgs { /// data file -- either `.zarr` or `.h5` pub data_file: Box, - /// column indices to take: e.g., `0,1,2,3` - #[arg(short = 'i', long, value_delimiter = ',')] - pub column_indices: Option>, + /// column indices to take: singles and inclusive ranges, e.g. `0,1,2,3`, + /// `1-10`, or `1-20,50-55` (past the last column, cut) + #[arg(short = 'i', long, conflicts_with_all = ["name_file", "pattern", "interactive"])] + pub column_indices: Option>, /// column name file where each line is a column name - #[arg(short = 'f', long)] + #[arg(short = 'f', long, conflicts_with_all = ["pattern", "interactive"])] pub name_file: Option>, + /// the columns whose names this regex finds (case-insensitive), e.g. + /// `/count/spliced$` or `^MT-` + #[arg(long, conflicts_with = "interactive")] + pub pattern: Option>, + /// pick the columns in a full-screen table: mark with Space, Enter to finish #[arg(short = 'I', long, default_value_t = false)] pub interactive: bool, @@ -68,8 +74,16 @@ pub struct SubsetColumnsArgs { #[arg(short = 'd', long, default_value = "@")] pub delimiter: char, - /// enable prefix matching (stored name is prefix of query or vice versa) - #[arg(long, default_value_t = true)] + /// enable prefix matching (stored name is prefix of query or vice + /// versa); `--allow-prefix=false` turns it off, as `--exact-names` does + #[arg( + long, + action = clap::ArgAction::Set, + num_args = 0..=1, + require_equals = true, + default_missing_value = "true", + default_value_t = true + )] pub allow_prefix: bool, /// Match names exactly (plus the delimiter base key): no prefix matching. @@ -106,14 +120,20 @@ pub struct SubsetRowsArgs { /// data file -- either `.zarr` or `.h5` pub data_file: Box, - /// row indices to take: e.g., `0,1,2,3` - #[arg(short = 'i', long, value_delimiter = ',')] - pub row_indices: Option>, + /// row indices to take: singles and inclusive ranges, e.g. `0,1,2,3`, + /// `1-10`, or `1-20,50-55` (past the last row, cut) + #[arg(short = 'i', long, conflicts_with_all = ["name_file", "pattern", "interactive"])] + pub row_indices: Option>, /// row name file where each line is a row name - #[arg(short = 'f', long)] + #[arg(short = 'f', long, conflicts_with_all = ["pattern", "interactive"])] pub name_file: Option>, + /// the rows whose names this regex finds (case-insensitive), e.g. + /// `/count/spliced$` or `^MT-` + #[arg(long, conflicts_with = "interactive")] + pub pattern: Option>, + /// pick the rows in a full-screen table: mark with Space, Enter to finish #[arg(short = 'I', long, default_value_t = false)] pub interactive: bool, @@ -125,8 +145,16 @@ pub struct SubsetRowsArgs { #[arg(short = 'd', long, default_value = "@")] pub delimiter: char, - /// enable prefix matching (stored name is prefix of query or vice versa) - #[arg(long, default_value_t = true)] + /// enable prefix matching (stored name is prefix of query or vice + /// versa); `--allow-prefix=false` turns it off, as `--exact-names` does + #[arg( + long, + action = clap::ArgAction::Set, + num_args = 0..=1, + require_equals = true, + default_missing_value = "true", + default_value_t = true + )] pub allow_prefix: bool, /// Match names exactly (plus the delimiter base key): no prefix matching. diff --git a/src/handlers/transformation/squeeze.rs b/src/handlers/transformation/squeeze.rs index 502b7b0..8c19c36 100644 --- a/src/handlers/transformation/squeeze.rs +++ b/src/handlers/transformation/squeeze.rs @@ -488,6 +488,7 @@ fn run_squeeze_then_merge( // The merge writes what the user asked squeeze for, so it takes // squeeze's own zip decision rather than always leaving a directory. zip: cmd_args.zip, + common_rows: false, do_squeeze: false, row_nnz_cutoff: 0, column_nnz_cutoff: 0, diff --git a/src/handlers/transformation/subset.rs b/src/handlers/transformation/subset.rs index bcaf439..9d339c0 100644 --- a/src/handlers/transformation/subset.rs +++ b/src/handlers/transformation/subset.rs @@ -6,8 +6,9 @@ use crate::interactive::stat_tui::Side; use crate::interactive::ui::Tui; use crate::sparse_io::*; use crate::utilities::io_helpers::{ - read_col_names, read_row_names, MAX_COLUMN_NAME_IDX, MAX_ROW_NAME_IDX, + parse_index_spec_within, read_col_names, read_row_names, MAX_COLUMN_NAME_IDX, MAX_ROW_NAME_IDX, }; +use crate::utilities::name_matching::names_matching; use crate::zarr_io::{apply_zip_flag, finalize_zarr_output, materialize_writable_backend}; use legume_numeric::matrix::common_io::*; @@ -97,7 +98,13 @@ pub fn subset_columns(args: &SubsetColumnsArgs) -> anyhow::Result<()> { Some(picked) => Some(picked), None => return Ok(()), }, - None => args.column_indices.clone(), + // By index or pattern: resolved on the input before staging, so a + // selection that fails copies nothing. + None => early_columns( + args.column_indices.as_deref(), + args.pattern.as_deref(), + &args.data_file, + )?, }; // Written off the screen: a `kill` meanwhile lets it finish. @@ -169,7 +176,7 @@ pub fn subset_columns(args: &SubsetColumnsArgs) -> anyhow::Result<()> { idx } else { return Err(anyhow::anyhow!( - "either `column-indices`, `name-file`, or `--interactive` must be provided" + "either `column-indices`, `name-file`, `pattern`, or `--interactive` must be provided" )); }; @@ -212,7 +219,13 @@ pub fn subset_rows(args: &SubsetRowsArgs) -> anyhow::Result<()> { Some(picked) => Some(picked), None => return Ok(()), }, - None => args.row_indices.clone(), + // By index or pattern: resolved on the input before staging, so a + // selection that fails copies nothing. + None => early_rows( + args.row_indices.as_deref(), + args.pattern.as_deref(), + &args.data_file, + )?, }; // Written off the screen: a `kill` meanwhile lets it finish. @@ -280,7 +293,7 @@ pub fn subset_rows(args: &SubsetRowsArgs) -> anyhow::Result<()> { idx } else { return Err(anyhow::anyhow!( - "either `row-indices`, `name-file`, or `--interactive` must be provided" + "either `row-indices`, `name-file`, `pattern`, or `--interactive` must be provided" )); }; @@ -364,3 +377,52 @@ pub fn reorder_rows(args: &ReorderRowsArgs) -> anyhow::Result<()> { #[cfg(test)] #[path = "subset/tests.rs"] mod tests; + +/// The columns `-i` or `--pattern` selects in `data_file`, read without +/// staging anything; none when neither is given. +fn early_columns( + spec: Option<&str>, + pattern: Option<&str>, + data_file: &str, +) -> anyhow::Result>> { + early(spec, pattern, data_file, Side::Columns) +} + +/// [`early_columns`] for rows. +fn early_rows( + spec: Option<&str>, + pattern: Option<&str>, + data_file: &str, +) -> anyhow::Result>> { + early(spec, pattern, data_file, Side::Rows) +} + +fn early( + spec: Option<&str>, + pattern: Option<&str>, + data_file: &str, + side: Side, +) -> anyhow::Result>> { + if spec.is_none() && pattern.is_none() { + return Ok(None); + } + let data = open_sparse_matrix_by_path(data_file)?; + let (what, n) = match side { + Side::Columns => ("columns", data.num_columns().unwrap_or(0)), + Side::Rows => ("rows", data.num_rows().unwrap_or(0)), + }; + let idx = if let Some(spec) = spec { + parse_index_spec_within(spec, n)? + } else { + let pattern = pattern.unwrap_or_default(); + let names = match side { + Side::Columns => data.column_names()?, + Side::Rows => data.row_names()?, + }; + let idx = names_matching(&names, pattern)?; + anyhow::ensure!(!idx.is_empty(), "no {what} match `{pattern}`"); + idx + }; + info!("selected {} {what}", idx.len()); + Ok(Some(idx)) +} diff --git a/src/main.rs b/src/main.rs index a9ee0a6..a94d426 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,6 +5,8 @@ mod column_subset; mod convert; mod handlers; mod hdf5_io; +#[cfg(feature = "mcp")] +mod mcp; // The interactive views live in the library (feature "tui") so downstream // crates can build on them; the CLI uses the same module. use data_beans::interactive; @@ -28,6 +30,7 @@ use crate::handlers::builders::{ run_build_from_fragments, run_build_from_mtx, run_build_from_zarr_triplets, FromFragmentsArgs, FromMtxArgs, FromZarrArgs, }; +use crate::handlers::cell_qc::{run_cell_qc, CellQcArgs}; #[cfg(feature = "hdf5")] use crate::handlers::exporters::{run_export_to_h5ad, ToH5adArgs}; use crate::handlers::exporters::{run_export_to_mtx, ToMtxArgs}; @@ -164,6 +167,13 @@ fn run(cli: &Cli) -> anyhow::Result<()> { Commands::Statistics(args) => { run_stat(args)?; } + #[cfg(feature = "plot")] + Commands::Plot(args) => { + crate::handlers::plotting::run_plot(args)?; + } + Commands::Qc(args) => { + run_cell_qc(args)?; + } Commands::Histogram(args) => { run_histogram(args)?; } @@ -203,6 +213,8 @@ fn run(cli: &Cli) -> anyhow::Result<()> { Commands::MergeBackend(args) => { run_merge_backend(args)?; } + #[cfg(feature = "mcp")] + Commands::Mcp(args) => mcp::run_mcp(args)?, Commands::Attach(args) => legume_tui::cli::attach(&TOOL, args)?, Commands::Sessions => legume_tui::cli::sessions(&TOOL)?, Commands::Kill(args) => legume_tui::cli::kill(&TOOL, args)?, @@ -644,6 +656,34 @@ enum Commands { )] Statistics(RunStatArgs), + #[command( + about = "Cell QC: near-empty floor and MAD outliers, as senna and pinto apply it", + long_about = "Run the cell QC senna and pinto apply, with the same --qc-* flags, on one backend.\n\ + Cells below the near-empty floor (--qc-min-cell-nnz) and cells outside the\n\ + robust band `median +/- --qc-mads * MAD * 1.4826` on detected features and total\n\ + counts (and the mito fraction, given --qc-mito-pattern) are dropped.\n\ + \n\ + Prints the outcome: cells and features kept, each metric's median, the band\n\ + it was held to and how many cells fell outside. With -o, writes the kept\n\ + cells (and features, under --qc-feature-min-cells) to a new backend;\n\ + --qc-report writes the per-cell metrics and verdicts.", + mut_args(|a| a.hide(false)), + mut_arg("qc_mito_pattern", |a| a.default_value(crate::handlers::cell_qc::MITO)) + )] + Qc(CellQcArgs), + + #[cfg(feature = "plot")] + #[command( + about = "Figures of a backend's QC or statistics (.png, .svg, .pdf)", + long_about = "Draw a backend's figures:\n\ + - `plot qc`: each QC metric's histogram, stacked by verdict, with the edge\n\ + QC held it to (the same --qc-* flags as `qc`), and total counts against\n\ + detected genes, cells coloured kept, outlier or near-empty;\n\ + - `plot stat`: the histogram of one row or column statistic.\n\ + The output's extension picks the format." + )] + Plot(crate::handlers::plotting::PlotArgs), + #[command( about = "ASCII log-scale histogram of a row/column statistic", long_about = "Print an ASCII log10(x+1) histogram of a statistic.\n\ @@ -658,6 +698,17 @@ enum Commands { )] Histogram(RunHistogramArgs), + #[cfg(feature = "mcp")] + #[command( + about = "Serve the commands to an agent over MCP (stdio)", + long_about = "Serve data-beans to an agent as Model Context Protocol tools over stdio\n\ + (`claude mcp add data-beans -- data-beans mcp`). Its tools read a\n\ + backend's size, names, statistics and values, and build new backends\n\ + by import, subset, subsample, split, convert and merge. No tool\n\ + replaces an existing file; squeeze, which rewrites its input, is left out." + )] + Mcp(mcp::McpArgs), + #[command( about = "Return to a detached data-beans view", long_about = legume_tui::cli::about(&TOOL, legume_tui::cli::QUIT_KEY).attach @@ -703,3 +754,34 @@ fn interactive_session_help(more: &str) -> String { const STDOUT_IN_A_SESSION: &str = "Written to stdout, the result goes to the session's \ terminal, and the part written while detached is lost; \ give -o a file to keep it. "; + +#[cfg(test)] +mod cli_tests { + use super::Cli; + use clap::{ArgAction, CommandFactory}; + + /// A flag that is on by default must have a way off: a bare `SetTrue` + /// flag defaulting to true cannot be turned off from the command line + /// (`--qc-mad-on-genes` and `--allow-prefix` were), so such a flag + /// takes a value (`--flag=false`) instead. + #[test] + fn no_flag_is_stuck_on() { + fn walk(cmd: &clap::Command, path: &str, stuck: &mut Vec) { + for arg in cmd.get_arguments() { + let on = arg.get_default_values().iter().any(|v| v == "true"); + if on && matches!(arg.get_action(), ArgAction::SetTrue) { + stuck.push(format!("{path} --{}", arg.get_id())); + } + } + for sub in cmd.get_subcommands() { + walk(sub, &format!("{path} {}", sub.get_name()), stuck); + } + } + let mut stuck = vec![]; + walk(&Cli::command(), "data-beans", &mut stuck); + assert!( + stuck.is_empty(), + "flags that cannot be turned off: {stuck:?}" + ); + } +} diff --git a/src/mcp.rs b/src/mcp.rs new file mode 100644 index 0000000..13af0be --- /dev/null +++ b/src/mcp.rs @@ -0,0 +1,1289 @@ +//! `data-beans mcp`: data-beans' inspection and backend-building commands as +//! Model Context Protocol tools, served over stdio with `legume-mcp`, so an +//! agent can look inside a backend and derive new ones from it. +//! +//! Every tool runs as a child `data-beans` process, as from the shell: +//! stdout carries the protocol, a child's log lines come back as progress, +//! and a cancelled call kills its child. Reads return what the child printed, +//! parsed and cut to a page. Writes take an `out` that must not exist yet: +//! no tool replaces a file, and `squeeze`, which rewrites its input, is not +//! offered. +//! +//! The project's records are legume-mcp's, laid out after Noble (2009), "A +//! quick guide to organizing computational biology projects", PLoS Comput +//! Biol 5(7):e1000424: each writing step goes into the day's replayable +//! `results//runall.sh` with its outputs (skipped on replay when they +//! exist) and a line in `results/INDEX.tsv`; an output the agent does not +//! name goes in that day's folder; `import` notes its input's source in +//! `data/README.md`. + +use crate::handlers::merging::{batch_file, same_file}; +use crate::hdf5_io::strip_backend_suffix; +use crate::sparse_io::SparseIoBackend; +use crate::utilities::io_helpers::{index_count, parse_index_spec}; +use crate::utilities::name_matching::{match_by_substring, names_matching}; +use crate::zarr_io::{output_target, OutputTarget}; +use anyhow::{Context, Result}; +use clap::{Args, Subcommand}; +use legume_mcp::{Call, Content, Handler, ServerInfo, ToolSpec}; +use serde_json::{json, Map, Value}; +use std::path::{Path, PathBuf}; +use std::process::Command; + +#[derive(Args, Debug)] +pub struct McpArgs { + #[command(subcommand)] + pub then: Option, +} + +#[derive(Subcommand, Debug)] +pub enum McpCommand { + /// Watch a data-beans MCP server while an agent uses it: its calls as + /// they come, each one's log and result, the notebook; s stops the + /// selected call, n leaves the agent a note, o opens what a call wrote + /// in the statistics explorer + Watch(WatchArgs), +} + +#[derive(Args, Debug)] +pub struct WatchArgs { + /// the server's watch folder (as `ls $XDG_RUNTIME_DIR/legume-mcp` + /// lists them); without it, the one live server, or a pick among them + pub server: Option, + + /// watch only: no stop, no notes + #[arg(long)] + pub watch_only: bool, +} + +const INSTRUCTIONS: &str = "\ +data-beans stores sparse single-cell matrices, features (genes, peaks) as \ +rows and cells as columns, in .zarr.zip, .zarr or .h5 backends. + +Look before you derive: +1. info gives a backend's size and metadata; list_contents shows what an \ +outside .zarr or .h5 holds before you import it. +2. names pages through row or column names, with a regex filter, to find \ +the exact names to pass on. +3. statistics summarizes each row or column (nnz, tot, mu, sig) with \ +quantiles and the top entries by a chosen statistic; values reads the \ +counts of a few named rows or columns. + +Quality control: qc runs the cell QC senna and pinto apply (a near-empty \ +floor, then robust MAD bands on detected genes, total counts and the mito \ +fraction). Call it without `out` first: it reports each metric's median, the \ +band and how many cells fall outside, and writes nothing. Adjust n_mads or \ +the caps until the drops are what the data warrant, then call it with `out` \ +to write the kept cells as a new backend. plot shows the same settings as a \ +figure (histograms with the band edges, counts against genes by verdict), \ +so look at it before and after tuning. + +Derive new backends with import (mtx, zarr triplets, fragments, h5ad, 10x; \ +say in `source` where the input came from), subset, subsample, split, \ +convert and merge. Each writes to `out`, which must not exist yet: no tool \ +replaces a file. Leave `out` out and the output goes in today's results \ +folder, named after its input and the tool. Outputs are .zarr.zip archives \ +unless `out` ends in .zarr (a directory) or .h5. Paths are resolved against \ +the server's working directory; results give the paths written. + +A person may be watching (`data-beans mcp watch`): a \"Note from the user\" \ +in a result is theirs to you; take it into account. A call they stopped \ +fails with \"stopped by the user\": do not run it again unasked."; + +/// The most entries a read returns, and the most a `values` call takes. +const MAX_LIMIT: usize = 1000; +const MAX_TAKEN: usize = 100; + +fn schema(properties: Value, required: &[&str]) -> Value { + json!({ "type": "object", "properties": properties, "required": required }) +} + +/// The `import` sources this build can read. +fn import_commands() -> Vec<&'static str> { + let mut v = vec!["from-mtx", "from-zarr", "from-fragments"]; + if cfg!(feature = "hdf5") { + v.extend(["from-h5ad", "from-10x-matrix", "from-10x-molecule"]); + } + v +} + +/// The cell-QC settings `qc` and `plot` share, and `data`. +fn qc_properties(data: &Value) -> Value { + json!({ + "data": data, + "mito_pattern": { "type": "string", "description": "Regex over feature names for mitochondrial genes (default (?i)^MT-; empty turns the mito metric off)" }, + "n_mads": { "type": "number", "exclusiveMinimum": 0, "description": "Band half-width in MADs (default 5; smaller drops more)" }, + "min_cell_nnz": { "type": "integer", "minimum": 0, "description": "Near-empty floor on detected features (default 2)" }, + "mad_on_genes": { "type": "boolean", "description": "Band on detected genes (default true)" }, + "mad_on_counts": { "type": "boolean", "description": "Band on total counts (default true)" }, + "mad_on_mito": { "type": "boolean", "description": "Band on the mito fraction (default true); false leaves only mito_max_frac, if given" }, + "mito_max_frac": { "type": "number", "description": "Hard cap on the mito fraction, 0..1" }, + "ribo_pattern": { "type": "string", "description": "Regex for ribosomal genes, e.g. (?i)^RP[SL]" }, + "ribo_max_frac": { "type": "number", "description": "Hard cap on the ribo fraction, 0..1" }, + "feature_min_cells": { "type": "integer", "minimum": 0, "description": "Also drop features detected in fewer cells (default 0, off)" }, + "auto_cutoff": { "type": "boolean", "description": "Also drop cells below the trough between ambient and real cells' detected genes" }, + }) +} + +fn tools() -> Vec { + let data = json!({ "type": "string", "description": "Backend: .zarr.zip, .zarr or .h5" }); + let axis = json!({ "type": "string", "enum": ["rows", "columns"] }); + let out = json!({ + "type": "string", + "description": "Output path; must not exist. A .zarr.zip archive unless it ends in .zarr or .h5. Omitted, a new name in today's results folder" + }); + let options = json!({ + "type": "array", "items": { "type": "string" }, + "description": "Further command-line options, as `data-beans --help` lists them" + }); + let limit = json!({ "type": "integer", "minimum": 1, "maximum": MAX_LIMIT, "description": "Entries to return (default 100)" }); + vec![ + ToolSpec::new( + "info", + "A backend's number of rows, columns and non-zeros, and its metadata.", + schema(json!({ "data": data }), &["data"]), + ), + ToolSpec::new( + "list_contents", + "The groups and arrays inside a .zarr, .zarr.zip or .h5 file, to find the fields `import` needs.", + schema(json!({ "file": { "type": "string" } }), &["file"]), + ), + ToolSpec::new( + "names", + "A page of a backend's row or column names, those matching `pattern` (a case-insensitive regex) when given. Returns the total, the number matched and the page with each name's index.", + schema( + json!({ + "data": data, "axis": axis, + "pattern": { "type": "string" }, + "offset": { "type": "integer", "minimum": 0 }, + "limit": limit, + }), + &["data", "axis"], + ), + ), + ToolSpec::new( + "statistics", + "Per-row or per-column statistics (nnz, tot, mu, sig) over one or more backends: quantiles of each, and the top entries by `sort_by`. With `out` (.parquet), the full table is written there instead.", + schema( + json!({ + "data": { "type": "array", "items": { "type": "string" }, "minItems": 1 }, + "axis": axis, + "row_pattern": { "type": "string", "description": "Column statistics count only rows matching this case-insensitive regex, e.g. ^MT-" }, + "sort_by": { "type": "string", "enum": ["nnz", "tot", "mu", "sig"] }, + "ascending": { "type": "boolean" }, + "limit": limit, + "out": { "type": "string", "description": "A .parquet file for the full table; must not exist" }, + }), + &["data", "axis"], + ), + ), + ToolSpec::new( + "values", + "The counts of at most 100 named (or indexed) rows or columns, as a dense table: one line per entry of the other axis, all-zero lines left out unless `keep_zeros`, at most `limit` lines (default 100, up to 1000) returned. With `out` (.parquet or .tsv), the full table is written there instead.", + schema( + json!({ + "data": data, "axis": axis, + "names": { "type": "array", "items": { "type": "string" }, "description": "Names to take (substring matching)" }, + "indices": { "type": "string", "description": "0-based indices and inclusive ranges, e.g. 0-9,20" }, + "keep_zeros": { "type": "boolean" }, + "limit": limit, + "out": { "type": "string" }, + }), + &["data", "axis"], + ), + ), + ToolSpec::new( + "qc", + "Cell QC as senna and pinto apply it: cells below the near-empty floor, or outside `median ± n_mads·MAD` on log detected genes and log total counts (from below) or on the mito fraction (from above), are dropped. Returns the counts kept and dropped and, per metric, its median, band and the cells outside it. Without `out` nothing is written; with it, the kept cells go to a new backend.", + schema( + { + let mut p = qc_properties(&data); + p["report"] = json!({ "type": "string", "description": "A .tsv(.gz) for the per-cell metrics and verdicts; must not exist" }); + p["out"] = json!({ "type": "string", "description": "Output for the kept cells; must not exist. Omit for a dry run" }); + p + }, + &["data"], + ), + ), + #[cfg(feature = "plot")] + ToolSpec::new( + "plot", + "A figure, returned as an image (unless `out` is an .svg or .pdf). kind qc (default): each QC metric's histogram, stacked kept/outlier/near-empty, with the edge QC held it to, and total counts against detected genes by verdict; it takes the qc settings, so plot before and after tuning them. kind stat: the histogram of one statistic over rows or columns, on a log10(x+1) axis. With `out` (.png, .svg or .pdf) the figure is kept there too.", + schema( + { + let mut p = qc_properties(&data); + p["kind"] = json!({ "type": "string", "enum": ["qc", "stat"] }); + p["axis"] = json!({ "type": "string", "enum": ["rows", "columns"], "description": "For kind stat (default columns)" }); + p["stat"] = json!({ "type": "string", "enum": ["nnz", "tot", "mu", "sig"], "description": "For kind stat (default nnz)" }); + p["out"] = json!({ "type": "string", "description": "A .png, .svg or .pdf to keep; must not exist" }); + p + }, + &["data"], + ), + ), + ToolSpec::new( + "import", + "Build a backend from another format with `data-beans `; from-mtx takes the feature and barcode files as options (-r, -c).", + schema( + json!({ + "command": { "type": "string", "enum": import_commands() }, + "input": { "type": "string" }, + "source": { "type": "string", "description": "Where the input came from (a URL or accession, and its version), kept in the project's data/README.md" }, + "out": out, "options": options, + }), + &["command", "input"], + ), + ), + ToolSpec::new( + "subset", + "A new backend with only some rows or columns, chosen by one of: `names` (exact unless `prefix_match`), `pattern` (a case-insensitive regex over the names, e.g. /count/spliced$), or `indices` (0-based, as a list or as text with ranges, e.g. \"0-19612,20000\").", + schema( + json!({ + "data": data, "axis": axis, + "names": { "type": "array", "items": { "type": "string" } }, + "pattern": { "type": "string" }, + "indices": { "description": "0-based indices: a list, or text with ranges (0-9,20)", "anyOf": [ + { "type": "array", "items": { "type": "integer", "minimum": 0 } }, + { "type": "string" } + ] }, + "prefix_match": { "type": "boolean" }, + "out": out, "options": options, + }), + &["data", "axis"], + ), + ), + ToolSpec::new( + "subsample", + "A new backend of randomly drawn cells (columns) and/or genes (rows), by count or fraction; an unset axis is kept whole.", + schema( + json!({ + "data": data, + "cells": { "type": "integer", "minimum": 1 }, + "cell_frac": { "type": "number" }, + "genes": { "type": "integer", "minimum": 1 }, + "gene_frac": { "type": "number" }, + "seed": { "type": "integer", "minimum": 0 }, + "out": out, + }), + &["data"], + ), + ), + ToolSpec::new( + "split", + "Split cells into train/test halves ({out}.train, {out}.test) or K folds ({out}.fold{k}.train/test). On spatial data pass --coord/--grid in `options`; --groups keeps a donor or sample on one side. `out` is a prefix here.", + schema( + json!({ + "data": data, + "test_frac": { "type": "number" }, + "folds": { "type": "integer", "minimum": 2 }, + "seed": { "type": "integer", "minimum": 0 }, + "out": { "type": "string", "description": "Output prefix; nothing named {out}.train*, {out}.test* or {out}.fold* may exist. Omitted, a new prefix in today's results folder" }, + "options": options, + }), + &["data"], + ), + ), + ToolSpec::new( + "convert", + "Copy a backend into another format: zarr <-> h5, .zarr <-> .zarr.zip; the format follows `out`.", + schema(json!({ "data": data, "out": out }), &["data"]), + ), + ToolSpec::new( + "merge", + "Merge backends into one: their columns side by side, each tagged with its batch, and the batch labels written beside it (`data-beans merge-backend`). The backends must have the same rows in the same order, unless `common_rows`, which keeps the rows all of them have, in the first one's order.", + schema( + json!({ + "data": { "type": "array", "items": { "type": "string" }, "minItems": 2 }, + "common_rows": { "type": "boolean" }, + "out": out, "options": options, + }), + &["data"], + ), + ), + ] +} + +/// What o opens on a call: a backend it wrote, in the statistics explorer. +fn open(call: &legume_mcp::watch::view::CallView) -> Option { + let written = written_backend(&call.text)?; + let mut cmd = Command::new(own_exe().ok()?); + cmd.args(["statistics", &written, "-s", "column", "--interactive"]); + Some(cmd) +} + +/// The backend a result says it wrote (`written`, or the first of a list), +/// when it is one. The feed cuts a result at 4 KB: past that, nothing opens. +fn written_backend(text: &str) -> Option { + // The first value: a note the person left may follow it. + let v: Value = serde_json::Deserializer::from_str(text) + .into_iter() + .next()? + .ok()?; + let w = &v["written"]; + let path = w.as_str().or_else(|| w.get(0)?.as_str())?; + (strip_backend_suffix(path) != path).then(|| path.to_string()) +} + +pub fn run_mcp(args: &McpArgs) -> Result<()> { + if let Some(McpCommand::Watch(w)) = &args.then { + let tool = &crate::interactive::TOOL; + return legume_mcp::watch::view::run(tool, w.server.clone(), !w.watch_only, open); + } + let info = ServerInfo { + name: "data-beans".into(), + version: env!("CARGO_PKG_VERSION").into(), + instructions: Some(INSTRUCTIONS.into()), + }; + legume_mcp::serve(info, DataBeans) +} + +struct DataBeans; + +impl Handler for DataBeans { + fn tools(&self) -> Vec { + tools() + } + + fn call(&mut self, name: &str, args: Value, cx: &Call) -> Result> { + let a = Arguments(args); + let v = match name { + "info" => summary(cx, &a.path("data")?)?, + "list_contents" => return list_contents(&a, cx), + "names" => names(&a, cx)?, + "statistics" => statistics(&a, cx)?, + "values" => values(&a, cx)?, + "qc" => qc(&a, cx)?, + #[cfg(feature = "plot")] + "plot" => return plot(&a, cx), + "import" => import(&a, cx)?, + "subset" => subset(&a, cx)?, + "subsample" => { + let data = a.path("data")?; + let out = a.backend_out(cx, &named(&data, "subsample"))?; + let mut argv = argv(["subsample", &data]); + argv.extend(a.numbers(&["cells", "cell_frac", "genes", "gene_frac", "seed"])?); + let what = format!("cells and genes drawn at random from {data}"); + written(cx, argv, &out, &what, &[data])? + } + "split" => split(&a, cx)?, + "convert" => { + let data = a.path("data")?; + let out = a.backend_out(cx, &named(&data, "converted"))?; + let what = format!("{data} in another format"); + written(cx, argv(["convert", &data]), &out, &what, &[data])? + } + "merge" => merge(&a, cx)?, + _ => anyhow::bail!("no tool {name}"), + }; + Ok(vec![Content::Json(v)]) + } +} + +/// A tool's arguments. +struct Arguments(Value); + +impl Arguments { + fn str(&self, key: &str) -> Option<&str> { + self.0.get(key).and_then(Value::as_str) + } + + fn req(&self, key: &str) -> Result { + self.str(key) + .map(str::to_string) + .with_context(|| format!("`{key}` is required")) + } + + /// A path, which goes on the command line where a leading `-` would + /// read as an option. + fn path(&self, key: &str) -> Result { + let p = self.req(key)?; + anyhow::ensure!(!p.starts_with('-'), "`{key}` {p}: write it as ./{p}"); + // Commands that take several split them at commas. + anyhow::ensure!( + !p.contains(','), + "`{key}` {p}: a path with a comma; rename it" + ); + Ok(p) + } + + /// Paths for a command that takes a list of them, which it also splits + /// at commas. + fn paths(&self, key: &str) -> Result> { + let v = self.strings(key)?; + if let Some(p) = v.iter().find(|p| p.starts_with('-')) { + anyhow::bail!("`{key}` {p}: write it as ./{p}"); + } + if let Some(p) = v.iter().find(|p| p.contains(',')) { + anyhow::bail!("`{key}` {p}: a path with a comma cannot go in a list; rename it"); + } + Ok(v) + } + + /// `key`'s value, when it is given and not null. + fn given(&self, key: &str) -> Option<&Value> { + self.0.get(key).filter(|v| !v.is_null()) + } + + fn parse(&self, key: &str) -> Result> { + self.given(key) + .map(|v| serde_json::from_value(v.clone())) + .transpose() + .with_context(|| format!("reading `{key}`")) + } + + fn flag(&self, key: &str) -> Result { + Ok(self.parse(key)?.unwrap_or(false)) + } + + fn strings(&self, key: &str) -> Result> { + Ok(self.parse(key)?.unwrap_or_default()) + } + + /// The numbers given among `keys`, as `--key-name value`. + fn numbers(&self, keys: &[&str]) -> Result> { + let mut v = vec![]; + for k in keys { + match self.given(k) { + None => {} + Some(n @ Value::Number(_)) => { + v.extend([format!("--{}", k.replace('_', "-")), n.to_string()]) + } + Some(other) => anyhow::bail!("`{k}` is a number, not {other}"), + } + } + Ok(v) + } + + fn rows(&self) -> Result { + match self.req("axis")?.as_str() { + "rows" | "row" => Ok(true), + "columns" | "column" => Ok(false), + other => anyhow::bail!("`axis` is rows or columns, not {other}"), + } + } + + fn limit(&self) -> Result { + Ok(self.parse("limit")?.unwrap_or(100).clamp(1, MAX_LIMIT)) + } + + /// `out`, or, when the agent names none, `{name}{ext}` in today's + /// results folder, `-2`, `-3`… added to `name` (before `ext`) while + /// `taken` says a name is in use. + fn out_or( + &self, + cx: &Call, + (name, ext): (&str, &str), + taken: impl Fn(&str) -> bool, + ) -> Result { + if self.given("out").is_some() { + return self.path("out"); + } + let dir = cx + .results_dir() + .context("`out` is required: this server keeps no results folder")?; + (1..10_000) + .map(|k| match k { + 1 => dir.join(format!("{name}{ext}")), + k => dir.join(format!("{name}-{k}{ext}")), + }) + .map(|p| p.to_string_lossy().into_owned()) + .find(|p| !taken(p)) + .context("no free name in the results folder") + } + + /// [`Self::out_or`] for a backend: a name is taken when what the command + /// would write for it exists. + fn backend_out(&self, cx: &Call, name: &str) -> Result { + self.out_or(cx, (name, ""), |c| { + backend_target(c).map_or(true, |t| { + outputs_beside(strip_backend_suffix(&t.effective)).map_or(true, |v| !v.is_empty()) + }) + }) + } + + /// `options`, refusing those that would name the output, change its + /// format or open a view. Short options stand alone (`-r X`, not + /// `-rX`), so none hides inside a cluster. + fn options(&self) -> Result> { + let v = self.strings("options")?; + const TAKEN: &[&str] = &[ + "-o", + "--output", + "--output-directory", + "--qc-report", + "--save-histogram", + "--backend", + "--no-zip", + "-I", + "--interactive", + "--", + ]; + for o in &v { + let name = o.split('=').next().unwrap_or(o); + let cluster = o.starts_with('-') + && !o.starts_with("--") + && o.len() > 2 + && o.chars().nth(1).is_some_and(|c| c.is_ascii_alphabetic()); + anyhow::ensure!( + !cluster, + "{o}: give a short option apart from its value (`-r X`)" + ); + anyhow::ensure!( + !TAKEN.contains(&name), + "{o} is not taken in `options`: outputs go to `out`, their format follows it, and no view opens here" + ); + } + Ok(v) + } +} + +fn argv(words: [&str; N]) -> Vec { + words.map(String::from).into() +} + +/// `data-beans` itself, with `argv`. A binary replaced under the running +/// server (`cargo install`) is run as the process's own image. +fn command(argv: &[String]) -> Result { + let mut cmd = Command::new(own_exe()?); + // `-v` so its log lines come back as progress. + cmd.arg("-v").args(argv); + Ok(cmd) +} + +/// The data-beans binary this is: the process's own image when the file was +/// replaced under it (`cargo install`). +fn own_exe() -> Result { + let exe = std::env::current_exe().context("finding the data-beans binary")?; + let own = Path::new("/proc/self/exe"); + Ok(if !exe.exists() && own.exists() { + own.to_path_buf() + } else { + exe + }) +} + +/// `data-beans` run with `argv`, a step that writes nothing, to its end; +/// its stdout. The records keep the call, not the command. +fn data_beans(cx: &Call, argv: &[String]) -> Result { + cx.run_reading(command(argv)?) +} + +/// [`data_beans`] for a step that writes `outputs`: the project's records +/// list each with `what` (what was made, from what), and a replay skips the +/// step when they all exist. +fn data_beans_making( + cx: &Call, + argv: &[String], + outputs: &[impl AsRef], + what: &str, +) -> Result { + let outputs: Vec = outputs.iter().map(|p| p.as_ref().to_path_buf()).collect(); + cx.run_making(command(argv)?, &outputs, what) +} + +/// `{stem}.{tool}`: `data`'s name without its folder and backend suffix, +/// then the tool, so a chain of unnamed outputs spells its lineage +/// (`pbmc.imported.qc.subset`). The name an output takes when the agent +/// gives none. +fn named(data: &str, tool: &str) -> String { + let stem = Path::new(strip_backend_suffix(data)) + .file_name() + .map_or_else(|| data.to_string(), |n| n.to_string_lossy().into_owned()); + format!("{stem}.{tool}") +} + +/// Where a backend-writing command given `-o out` writes, as it resolves +/// it under its default `--backend zarr`; zipped unless `out` names a +/// `.zarr` directory. +fn backend_target(out: &str) -> Result { + output_target(out, SparseIoBackend::Zarr, !out.ends_with(".zarr")) +} + +fn absolute(p: &str) -> String { + std::path::absolute(p) + .map(|p| p.to_string_lossy().into_owned()) + .unwrap_or_else(|_| p.to_string()) +} + +/// A file a tool is about to write: refused if it exists. +fn fresh(path: &str) -> Result<()> { + anyhow::ensure!( + !Path::new(path).exists(), + "{path} exists; not replacing it: choose another `out`" + ); + Ok(()) +} + +/// Run `argv`, a command that writes the backend `out` (made of `what`) +/// from `inputs`, with `-o`; say where the backend is and how big. +fn written( + cx: &Call, + argv: Vec, + out: &str, + what: &str, + inputs: &[String], +) -> Result { + Ok(Value::Object( + written_with(cx, argv, out, what, inputs, &[])?.0, + )) +} + +/// [`written`], with `also` written beside it; and what the command +/// printed. The output must not be an input, and nothing else named after +/// it (`{stem}.*`) may exist beside it: data-beans clears the backend and +/// its working copy before writing, and writes its side files (batch +/// labels, an importer's metadata) over whatever is there. The side files +/// it wrote come back as `also`. +fn written_with( + cx: &Call, + mut argv: Vec, + out: &str, + what: &str, + inputs: &[String], + also: &[&str], +) -> Result<(Map, String)> { + let target = backend_target(out)?; + let is_input = |p: &str| inputs.iter().any(|i| same_file(i, p)); + for p in [&target.effective, &target.working, &target.found] { + anyhow::ensure!(!is_input(p), "{p} is an input: choose another `out`"); + } + let stem = strip_backend_suffix(&target.effective).to_string(); + let before: Vec = outputs_beside(&stem)? + .into_iter() + .filter(|p| !is_input(p)) + .collect(); + anyhow::ensure!( + before.is_empty(), + "{} exist beside {out}; not replacing them: choose another `out`", + before.join(", ") + ); + // The full name, so the command cannot resolve it otherwise (`subset` + // keeps an extensionless output in its input's format). + argv.extend(["-o".into(), target.effective.to_string()]); + // A `.zarr` directory asked for, as `backend_target` resolved it. + if out.ends_with(".zarr") { + argv.push("--no-zip".into()); + } + let path: &str = &target.found; + let mut outputs = vec![path]; + outputs.extend(also); + let printed = data_beans_making(cx, &argv, &outputs, what)?; + let written = absolute(path); + // Nothing was there before, so the rest there now was written too. + let side: Vec = outputs_beside(&stem)? + .into_iter() + .filter(|p| *p != written && !is_input(p)) + .collect(); + let mut v = Map::new(); + v.insert("written".into(), json!(written)); + if !side.is_empty() { + v.insert("also".into(), json!(side)); + } + v.insert("info".into(), summary(cx, path)?); + Ok((v, printed)) +} + +/// Run `argv`, a command that writes the table `out` (made of `what`). +fn to_file(cx: &Call, mut argv: Vec, out: &str, what: &str) -> Result { + fresh(out)?; + argv.extend(["-o".into(), out.into()]); + data_beans_making(cx, &argv, &[out], what)?; + Ok(json!({ "written": absolute(out) })) +} + +fn summary(cx: &Call, data: &str) -> Result { + let text = data_beans(cx, &argv(["info", data]))?; + Ok(Value::Object(key_values(&text.replace("number_of_", "")))) +} + +/// `key:\tvalue` lines as an object, numbers as numbers; a dotted key +/// (`meta.sample`, `n_genes.median`) nests. +fn key_values(text: &str) -> Map { + let mut v = Map::new(); + for line in text.lines() { + let Some((key, value)) = line.split_once(":\t") else { + continue; + }; + let value = match (value.parse::(), value.parse::()) { + (Ok(n), _) => json!(n), + (_, Ok(x)) if x.is_finite() => json!(x), + _ => json!(value), + }; + let mut at = &mut v; + let mut parts = key.split('.').peekable(); + while let Some(part) = parts.next() { + if parts.peek().is_none() { + at.insert(part.into(), value); + break; + } + let next = at.entry(part).or_insert_with(|| json!({})); + if !next.is_object() { + *next = json!({}); + } + at = next.as_object_mut().unwrap(); + } + } + v +} + +fn list_contents(a: &Arguments, cx: &Call) -> Result> { + let file = a.path("file")?; + let command = if file.ends_with(".h5") || file.ends_with(".h5ad") || file.ends_with(".hdf5") { + "list-h5" + } else { + "list-zarr" + }; + Ok(vec![Content::Text(data_beans( + cx, + &argv([command, &file]), + )?)]) +} + +/// The command listing a backend's row or column names. +fn names_command(rows: bool) -> &'static str { + if rows { + "row-names" + } else { + "column-names" + } +} + +fn names(a: &Arguments, cx: &Call) -> Result { + let text = data_beans(cx, &argv([names_command(a.rows()?), &a.path("data")?]))?; + let lines: Vec<&str> = text.lines().collect(); + // Which names the pattern finds, as `subset --pattern` would. + let found: Option> = match a.str("pattern") { + Some(p) => { + let names: Vec> = lines.iter().map(|&n| n.into()).collect(); + Some(names_matching(&names, p)?.into_iter().collect()) + } + None => None, + }; + let offset: usize = a.parse("offset")?.unwrap_or(0); + let limit = a.limit()?; + let (mut total, mut matched, mut page) = (0, 0, vec![]); + for (i, n) in lines.iter().enumerate() { + total += 1; + if found.as_ref().is_none_or(|f| f.contains(&i)) { + if (offset..offset.saturating_add(limit)).contains(&matched) { + page.push(json!({ "index": i, "name": n })); + } + matched += 1; + } + } + Ok(json!({ "total": total, "matched": matched, "offset": offset, "names": page })) +} + +/// One line of a tab-separated table as an object: its first field, a +/// name, as text; the rest numbers where they parse. +fn record<'a>(header: &[&str], line: impl Iterator) -> Value { + let m: Map = header + .iter() + .zip(line) + .enumerate() + .map(|(j, (h, x))| { + let v = match x.parse::() { + Ok(n) if j > 0 && n.is_finite() => json!(n), + _ => json!(x), + }; + (h.to_string(), v) + }) + .collect(); + Value::Object(m) +} + +/// A table's header; `statistics` heads its table `#name`. +fn header(line: Option<&str>) -> Vec<&str> { + line.unwrap_or_default() + .trim_start_matches('#') + .split('\t') + .collect() +} + +fn statistics(a: &Arguments, cx: &Call) -> Result { + let data = a.paths("data")?; + anyhow::ensure!(!data.is_empty(), "`data` is required"); + let dim = if a.rows()? { "row" } else { "column" }; + let mut argv = argv(["statistics", &data.join(","), "-s", dim]); + if let Some(p) = a.str("row_pattern") { + argv.push(format!("--row-name-pattern={p}")); + } + if let Some(out) = a.str("out") { + anyhow::ensure!(out.ends_with(".parquet"), "`out` is a .parquet file"); + let what = format!("{dim} statistics of {}", data.join(", ")); + return to_file(cx, argv, out, &what); + } + let text = data_beans(cx, &argv)?; + let mut lines = text.lines(); + let header = header(lines.next()); + let stats = header.get(1..).unwrap_or_default(); + let sort_by = a.str("sort_by").unwrap_or("tot"); + let col = stats + .iter() + .position(|h| *h == sort_by) + .with_context(|| format!("no `{sort_by}` among {stats:?}"))?; + // Each line once: itself, and its statistics. + let mut rows: Vec<(&str, Vec)> = lines + .map(|l| { + let v = l.split('\t').skip(1); + (l, v.map(|x| x.parse().unwrap_or(f64::NAN)).collect()) + }) + .collect(); + let mut quantiles = Map::new(); + for (j, h) in stats.iter().enumerate() { + let mut xs: Vec = rows + .iter() + .filter_map(|(_, v)| v.get(j).copied()) + .filter(|x| x.is_finite()) + .collect(); + if xs.is_empty() { + continue; + } + xs.sort_by(f64::total_cmp); + let q = |p: f64| xs[((xs.len() - 1) as f64 * p).round() as usize]; + quantiles.insert( + h.to_string(), + json!({ "min": q(0.0), "q25": q(0.25), "median": q(0.5), "q75": q(0.75), "max": q(1.0) }), + ); + } + let count = rows.len(); + let ascending = a.flag("ascending")?; + let key = |r: &(&str, Vec)| r.1.get(col).copied().unwrap_or(f64::NAN); + // A statistic that is not a number goes last either way. + let order = |x: &(&str, Vec), y: &(&str, Vec)| { + let (kx, ky) = (key(x), key(y)); + let by = if ascending { + kx.total_cmp(&ky) + } else { + ky.total_cmp(&kx) + }; + kx.is_nan().cmp(&ky.is_nan()).then(by) + }; + // The top `limit` picked, then only those sorted. + let k = a.limit()?.min(count); + if k > 0 && k < count { + rows.select_nth_unstable_by(k - 1, order); + rows.truncate(k); + } + rows.sort_by(order); + let top: Vec = rows + .iter() + .map(|(l, _)| record(&header, l.split('\t'))) + .collect(); + Ok(json!({ "count": count, "quantiles": quantiles, "sort_by": sort_by, "top": top })) +} + +fn values(a: &Arguments, cx: &Call) -> Result { + let rows = a.rows()?; + let data = a.path("data")?; + // Resolved to indices first, so the cap holds before anything is read. + let names = a.strings("names")?; + let taken: Vec = match (names.is_empty(), a.str("indices")) { + (false, None) => { + let all: Vec> = data_beans(cx, &argv([names_command(rows), &data]))? + .lines() + .map(Box::from) + .collect(); + let queries: Vec> = names.iter().map(|n| n.as_str().into()).collect(); + match_by_substring(&all, &queries, if rows { "row" } else { "column" })?.0 + } + (true, Some(spec)) => { + // Counted before it is expanded: a range can be any size. + let count = index_count(spec)?; + anyhow::ensure!( + count <= MAX_TAKEN, + "{count} entries taken; at most {MAX_TAKEN}" + ); + parse_index_spec(spec)? + } + _ => anyhow::bail!("give `names` or `indices` (as text, e.g. 0-9,20), one of them"), + }; + anyhow::ensure!( + taken.len() <= MAX_TAKEN, + "{} entries taken; at most {MAX_TAKEN}: name them more exactly", + taken.len() + ); + let spec: Vec = taken.iter().map(usize::to_string).collect(); + let command = if rows { "rows" } else { "columns" }; + let argv = argv([command, &data, "-i", &spec.join(",")]); + if let Some(out) = a.str("out") { + let what = format!("the counts of {} {command} of {data}", taken.len()); + return to_file(cx, argv, out, &what); + } + let text = data_beans(cx, &argv)?; + let mut lines = text.lines(); + let header = header(lines.next()); + let keep_zeros = a.flag("keep_zeros")?; + let limit = a.limit()?; + let nonzero = |x: &str| x != "0" && x.parse::().is_ok_and(|v| v != 0.0); + let (mut total, mut kept, mut page) = (0, 0, vec![]); + for l in lines { + total += 1; + if keep_zeros || l.split('\t').skip(1).any(nonzero) { + if page.len() < limit { + page.push(record(&header, l.split('\t'))); + } + kept += 1; + } + } + Ok(json!({ + "taken": header.get(1..).unwrap_or_default(), + "lines": total, + "kept": kept, + "values": page, + })) +} + +/// The `--qc-*` flags for the cell-QC settings among `a`. +fn qc_flags(a: &Arguments) -> Result> { + let mut v = vec![]; + // The command starts from (?i)^MT-; an empty pattern turns it off. + if let Some(mito) = a.str("mito_pattern") { + v.push(format!("--qc-mito-pattern={mito}")); + } + // `--flag=value`, so a value cannot read as an option. + if let Some(ribo) = a.str("ribo_pattern") { + v.push(format!("--qc-ribo-pattern={ribo}")); + } + let numbers = [ + ("n_mads", "--qc-mads"), + ("min_cell_nnz", "--qc-min-cell-nnz"), + ("mito_max_frac", "--qc-mito-max-frac"), + ("ribo_max_frac", "--qc-ribo-max-frac"), + ("feature_min_cells", "--qc-feature-min-cells"), + ]; + for (key, flag) in numbers { + match a.given(key) { + None => {} + Some(n @ Value::Number(_)) => v.push(format!("{flag}={n}")), + Some(other) => anyhow::bail!("`{key}` is a number, not {other}"), + } + } + for (key, flag) in [ + ("mad_on_genes", "--qc-mad-on-genes"), + ("mad_on_counts", "--qc-mad-on-counts"), + ("mad_on_mito", "--qc-mad-on-mito"), + ] { + if let Some(on) = a.parse::(key)? { + v.push(format!("{flag}={on}")); + } + } + if a.flag("auto_cutoff")? { + v.push("--qc-auto-cutoff".into()); + } + Ok(v) +} + +fn qc(a: &Arguments, cx: &Call) -> Result { + let data = a.path("data")?; + let mut argv = argv(["qc", &data]); + argv.extend(qc_flags(a)?); + let report = a.str("report"); + if let Some(report) = report { + fresh(report)?; + argv.push(format!("--qc-report={report}")); + } + let what = format!("the cells of {data} that cell QC keeps"); + // Without `out` a dry run: only the summary, and the report if asked. + let (mut v, printed) = match (a.given("out"), report) { + (Some(_), _) => { + let out = a.path("out")?; + written_with( + cx, + argv, + &out, + &what, + std::slice::from_ref(&data), + report.as_slice(), + )? + } + (None, Some(r)) => { + let what = format!("per-cell QC metrics and verdicts of {data}"); + (Map::new(), data_beans_making(cx, &argv, &[r], &what)?) + } + (None, None) => (Map::new(), data_beans(cx, &argv)?), + }; + v.extend(key_values(&printed)); + if let Some(report) = report { + v.insert("report".into(), json!(absolute(report))); + } + Ok(Value::Object(v)) +} + +#[cfg(feature = "plot")] +fn plot(a: &Arguments, cx: &Call) -> Result> { + let data = a.path("data")?; + let kind = a.str("kind").unwrap_or("qc"); + let mut argv = match kind { + "qc" => { + let mut v = argv(["plot", "qc", &data]); + v.extend(qc_flags(a)?); + v + } + "stat" => { + // Columns unless `axis` says rows. + let rows = a.given("axis").is_some() && a.rows()?; + let dim = if rows { "row" } else { "column" }; + // The names `statistics` reports them by, which `--stat` takes too. + let stat = a.str("stat").unwrap_or("nnz"); + anyhow::ensure!( + ["nnz", "tot", "mu", "sig"].contains(&stat), + "`stat` is nnz, tot, mu or sig, not {stat}" + ); + argv(["plot", "stat", &data, "-s", dim, "--stat", stat]) + } + other => anyhow::bail!("`kind` is qc or stat, not {other}"), + }; + // A figure in `out`'s format, else a PNG in the results folder; a + // server keeping no records shows it to the agent only. + let temp; + let kept = a.given("out").is_some() || cx.results_dir().is_some(); + let file = if kept { + let file = a.out_or(cx, (&named(&data, kind), ".png"), |c| Path::new(c).exists())?; + fresh(&file)?; + file + } else { + temp = tempfile::Builder::new().suffix(".png").tempfile()?; + temp.path().to_string_lossy().into_owned() + }; + argv.extend(["-o".into(), file.clone()]); + if kept { + data_beans_making(cx, &argv, &[&file], &format!("the {kind} figure of {data}"))?; + } else { + data_beans(cx, &argv)?; + } + let mut v = vec![]; + if file.ends_with(".png") { + v.push(Content::Png(std::fs::read(&file)?)); + } + if kept { + v.push(Content::Json(json!({ "written": absolute(&file) }))); + } + Ok(v) +} + +fn import(a: &Arguments, cx: &Call) -> Result { + let command = a.req("command")?; + anyhow::ensure!( + import_commands().contains(&command.as_str()), + "`command` is one of {:?}", + import_commands() + ); + let input = a.path("input")?; + // The input's own suffixes (`.mtx.gz`, `.h5ad`) go. + let file = Path::new(&input).file_name().map(|n| n.to_string_lossy()); + let stem = file + .as_deref() + .and_then(|n| n.split('.').next()) + .unwrap_or("data"); + let out = a.backend_out(cx, &format!("{stem}.imported"))?; + let what = format!("{input} as a backend ({command})"); + let mut argv = argv([&command, &input]); + argv.extend(a.options()?); + let v = written(cx, argv, &out, &what, std::slice::from_ref(&input))?; + // Where the input came from, in the project's data/README.md, once it + // is in. + let source = a.str("source").unwrap_or("a local file, source not given"); + cx.provenance(Path::new(&input), source); + Ok(v) +} + +/// `indices` as the CLI takes them: text (`0-9,20`) as given, a list with +/// its ascending runs folded into ranges. +fn index_spec(v: &Value) -> Result { + if let Value::String(spec) = v { + return Ok(spec.clone()); + } + let ix: Vec = serde_json::from_value(v.clone()).context("reading `indices`")?; + // Runs folded as they come: the order asked for is the order written. + let mut runs: Vec<(usize, usize)> = vec![]; + for i in ix { + match runs.last_mut() { + Some((_, end)) if *end + 1 == i => *end = i, + _ => runs.push((i, i)), + } + } + let spec: Vec = runs + .iter() + .map(|&(a, b)| { + if a == b { + a.to_string() + } else { + format!("{a}-{b}") + } + }) + .collect(); + Ok(spec.join(",")) +} + +fn subset(a: &Arguments, cx: &Call) -> Result { + let rows = a.rows()?; + let (command, side) = if rows { + ("subset-rows", "rows") + } else { + ("subset-columns", "columns") + }; + let data = a.path("data")?; + let out = a.backend_out(cx, &named(&data, "subset"))?; + let mut argv = argv([command, &data]); + // An empty list of names is no choice. + let names = a.strings("names")?; + anyhow::ensure!( + a.given("pattern").is_none_or(Value::is_string), + "`pattern` is a regex, as text" + ); + let given = [ + !names.is_empty(), + a.given("pattern").is_some(), + a.given("indices").is_some(), + ]; + anyhow::ensure!( + given.iter().filter(|&&g| g).count() == 1, + "give one of `names`, `pattern` or `indices`" + ); + // The names go to the command in a file, kept until it has run. + let file; + let what = if let Some(p) = a.str("pattern") { + // `--flag=value`, so a pattern cannot read as an option. + argv.push(format!("--pattern={p}")); + format!("the {side} of {data} matching {p}") + } else if let Some(v) = a.given("indices") { + let spec = index_spec(v)?; + argv.push(format!("--{}-indices={spec}", &side[..side.len() - 1])); + format!("{} {side} of {data}", index_count(&spec)?) + } else { + file = tempfile::NamedTempFile::new()?; + // Through the call, so the session script carries the names too. + cx.write_file(file.path(), &(names.join("\n") + "\n"))?; + argv.extend(["-f".into(), file.path().to_string_lossy().into_owned()]); + if !a.flag("prefix_match")? { + argv.push("--exact-names".into()); + } + format!("{} named {side} of {data}", names.len()) + }; + argv.extend(a.options()?); + written(cx, argv, &out, &what, &[data]) +} + +fn merge(a: &Arguments, cx: &Call) -> Result { + let data = a.paths("data")?; + anyhow::ensure!(data.len() >= 2, "`data` takes two or more backends"); + let out = a.backend_out(cx, &named(&data[0], "merged"))?; + // merge-backend's batch labels, beside the backend. + let batches = batch_file(&backend_target(&out)?.effective); + let what = format!("{} merged", data.join(" + ")); + let mut argv = argv(["merge-backend"]); + argv.extend(data.iter().cloned()); + if a.flag("common_rows")? { + argv.push("--common-rows".into()); + } + argv.extend(a.options()?); + let mut v = written_with(cx, argv, &out, &what, &data, &[&batches])?.0; + v.insert("batches".into(), json!(absolute(&batches))); + Ok(Value::Object(v)) +} + +/// The entries beside `stem` named after it (`{stem}.*`): what a command +/// writing under that name leaves there (a backend and its working copy, +/// batch labels, an importer's metadata, split's halves). +fn outputs_beside(stem: &str) -> Result> { + let name = Path::new(stem) + .file_name() + .context("`out` names no file")? + .to_string_lossy() + .into_owned(); + let parent = Path::new(stem) + .parent() + .filter(|d| !d.as_os_str().is_empty()); + let dir = std::path::absolute(parent.unwrap_or(Path::new(".")))?; + let entries = match std::fs::read_dir(&dir) { + Ok(e) => e, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(vec![]), + Err(e) => return Err(e.into()), + }; + let mut v: Vec = entries + .filter_map(|e| e.ok()) + .map(|e| e.file_name().to_string_lossy().into_owned()) + .filter(|n| { + n.strip_prefix(name.as_str()) + .is_some_and(|r| r.starts_with('.')) + }) + .map(|n| dir.join(n).to_string_lossy().into_owned()) + .collect(); + v.sort(); + Ok(v) +} + +fn split(a: &Arguments, cx: &Call) -> Result { + let data = a.path("data")?; + let prefix = a.out_or(cx, (&named(&data, "split"), ""), |c| { + outputs_beside(c).map_or(true, |v| !v.is_empty()) + })?; + let before = outputs_beside(&prefix)?; + anyhow::ensure!( + before.is_empty(), + "{} exist beside {prefix}; not replacing them: choose another `out`", + before.join(", ") + ); + let folds: Option = a.parse("folds")?; + // What split will write: two halves, or two per fold, all zipped. + let halves: Vec = match folds { + Some(k) => (0..k) + .flat_map(|f| ["train", "test"].map(|h| format!("{prefix}.fold{f}.{h}.zarr.zip"))) + .collect(), + None => ["train", "test"] + .map(|h| format!("{prefix}.{h}.zarr.zip")) + .into(), + }; + let mut argv = argv(["split", &data, "-o", &prefix]); + argv.extend(a.numbers(&["test_frac", "folds", "seed"])?); + argv.extend(a.options()?); + data_beans_making( + cx, + &argv, + &halves, + &format!("cells of {data} split for testing"), + )?; + // Nothing was there before, so all there is now was written. + Ok(json!({ "written": outputs_beside(&prefix)? })) +} + +#[cfg(test)] +mod tests { + use super::written_backend; + + /// The README gives the instructions for clients that do not pass them + /// on: they must be the server's, word for word. + #[test] + fn the_readme_says_what_the_server_tells_the_agent() { + let readme = include_str!("../README.md"); + assert!( + readme.contains(super::INSTRUCTIONS), + "README.md's \"What the agent is told\" block differs from mcp.rs INSTRUCTIONS" + ); + } + + #[test] + fn a_result_opens_the_backend_it_wrote() { + let one = r#"{ "written": "/p/results/x.subset.zarr.zip", "info": {} }"#; + assert_eq!( + written_backend(one).as_deref(), + Some("/p/results/x.subset.zarr.zip") + ); + let halves = r#"{ "written": ["/p/cv.test.zarr.zip", "/p/cv.train.zarr.zip"] }"#; + assert_eq!( + written_backend(halves).as_deref(), + Some("/p/cv.test.zarr.zip") + ); + // A note the person left follows the result. + let noted = format!("{one}\nNote from the user: use exact names"); + assert_eq!( + written_backend(¬ed).as_deref(), + Some("/p/results/x.subset.zarr.zip") + ); + // A figure, a table, a read: nothing to open in the explorer. + assert_eq!(written_backend(r#"{ "written": "/p/qc.png" }"#), None); + assert_eq!(written_backend(r#"{ "rows": 3 }"#), None); + assert_eq!(written_backend("[image/png, 3 bytes]"), None); + } +} diff --git a/src/qc_lib.rs b/src/qc_lib.rs index 46fe79b..ccd42c2 100644 --- a/src/qc_lib.rs +++ b/src/qc_lib.rs @@ -114,26 +114,45 @@ fn median_in_place(xs: &mut [f32]) -> f32 { xs[mid] } -/// Robust outlier keep-mask: `keep[i] = false` when `values[i]` falls -/// outside `median ± n_mads · MAD · 1.4826` on the requested `tail`. -/// Count-like metrics should pass `log1p = true` so the band is symmetric -/// on the multiplicative scale. Consolidates the median+MAD idiom used in -/// `cnv::hmm`. -/// `consider` (when `Some`) restricts the median/MAD band to the cells it -/// marks `true` — so e.g. near-empty cells don't contaminate the robust -/// center — while the keep decision is still returned for every cell. -pub fn robust_outlier_keep( +/// The robust band [`robust_outlier_keep`] tests against: `median ± n_mads · +/// MAD · 1.4826`, on the `log1p` scale when asked. +#[derive(Clone, Copy, Debug)] +pub struct RobustBand { + pub median: f32, + pub lo: f32, + pub hi: f32, +} + +impl RobustBand { + /// The band on the values' own scale, undoing `log1p`. + pub fn raw(self, log1p: bool) -> Self { + let back = |v: f32| if log1p { v.exp_m1() } else { v }; + Self { + median: back(self.median), + lo: back(self.lo), + hi: back(self.hi), + } + } +} + +/// A value on the scale its band is fitted on: `ln(1 + v)` for counts. +fn band_scale(v: f32, log1p: bool) -> f32 { + if log1p { + v.max(0.0).ln_1p() + } else { + v + } +} + +/// [`RobustBand`] of `values` (the `consider`ed ones, when given); `None` +/// when none is finite, or when they have no spread (a MAD of 0). +pub fn robust_band( values: &[f32], n_mads: f32, - tail: Tail, log1p: bool, consider: Option<&[bool]>, -) -> Vec { - let n = values.len(); - if n == 0 { - return vec![]; - } - let xform = |v: f32| if log1p { v.max(0.0).ln_1p() } else { v }; +) -> Option { + let xform = |v: f32| band_scale(v, log1p); // Band statistics computed over the considered subset only. let mut xs: Vec = values @@ -144,16 +163,48 @@ pub fn robust_outlier_keep( .filter(|v| v.is_finite()) .collect(); if xs.is_empty() { - return vec![true; n]; + return None; } // O(n) median via quickselect (no full sort). let median = median_in_place(&mut xs); let mut dev: Vec = xs.iter().map(|&v| (v - median).abs()).collect(); - let mad = (median_in_place(&mut dev) * 1.4826).max(1e-8); + let mad = median_in_place(&mut dev) * 1.4826; + // No spread: more than half the values are one value (a mito fraction + // of 0 in most nuclei, say). A band of width 0 would call every other + // value an outlier, so there is none. + if mad <= 0.0 { + return None; + } + + Some(RobustBand { + median, + lo: median - n_mads * mad, + hi: median + n_mads * mad, + }) +} - let lo = median - n_mads * mad; - let hi = median + n_mads * mad; +/// Robust outlier keep-mask: `keep[i] = false` when `values[i]` falls +/// outside `median ± n_mads · MAD · 1.4826` on the requested `tail`. +/// Count-like metrics should pass `log1p = true` so the band is symmetric +/// on the multiplicative scale. Consolidates the median+MAD idiom used in +/// `cnv::hmm`. +/// `consider` (when `Some`) restricts the median/MAD band to the cells it +/// marks `true` — so e.g. near-empty cells don't contaminate the robust +/// center — while the keep decision is still returned for every cell. +pub fn robust_outlier_keep( + values: &[f32], + n_mads: f32, + tail: Tail, + log1p: bool, + consider: Option<&[bool]>, +) -> Vec { + let n = values.len(); + let Some(band) = robust_band(values, n_mads, log1p, consider) else { + return vec![true; n]; + }; + let xform = |v: f32| band_scale(v, log1p); + let (lo, hi) = (band.lo, band.hi); values .iter() .map(|&raw| { @@ -701,6 +752,11 @@ pub struct QcArgs { #[arg( long = "qc-mad-on-genes", hide = true, + // `--qc-mad-on-genes=false` turns it off; a bare flag leaves it on. + action = clap::ArgAction::Set, + num_args = 0..=1, + require_equals = true, + default_missing_value = "true", default_value_t = true, help = "MAD-outlier drop on the per-cell detected-feature count", long_help = "Drop cells whose detected-feature count falls outside `median +/- --qc-mads * MAD * 1.4826`.\n\ @@ -717,6 +773,11 @@ pub struct QcArgs { #[arg( long = "qc-mad-on-counts", hide = true, + // `--qc-mad-on-counts=false` turns it off; a bare flag leaves it on. + action = clap::ArgAction::Set, + num_args = 0..=1, + require_equals = true, + default_missing_value = "true", default_value_t = true, help = "MAD-outlier drop on per-cell total counts", long_help = "Drop cells whose total count falls outside `median +/- --qc-mads * MAD * 1.4826`.\n\ @@ -724,6 +785,22 @@ pub struct QcArgs { )] pub qc_mad_on_counts: bool, + #[arg( + long = "qc-mad-on-mito", + hide = true, + // `--qc-mad-on-mito=false` turns it off; a bare flag leaves it on. + action = clap::ArgAction::Set, + num_args = 0..=1, + require_equals = true, + default_missing_value = "true", + default_value_t = true, + help = "MAD-outlier drop on the mito fraction, given --qc-mito-pattern", + long_help = "Drop cells whose mito fraction is above `median + --qc-mads * MAD * 1.4826`.\n\ + On whenever --qc-mito-pattern is given. `--qc-mad-on-mito=false` leaves\n\ + only the hard cap, --qc-mito-max-frac, when one is given." + )] + pub qc_mad_on_mito: bool, + #[arg( long = "qc-auto-cutoff", hide = true, @@ -758,7 +835,7 @@ impl QcArgs { // selects the rows to measure. mad_on_n_genes: self.qc_mad_on_genes, mad_on_counts: self.qc_mad_on_counts, - mad_on_mito: self.qc_mito_pattern.is_some(), + mad_on_mito: self.qc_mito_pattern.is_some() && self.qc_mad_on_mito, feature_min_cells: self.qc_feature_min_cells, drop_outliers: true, // Off by default: the near-empty floor plus the model's own @@ -785,6 +862,17 @@ mod qc_tests { assert!(keep_up[5], "lower outlier kept under Tail::Upper"); } + #[test] + fn no_spread_flags_no_one() { + // Most cells read no mito at all: the MAD is 0, and a cell with one + // mito read is no outlier for it. + let mut v = vec![0.0f32; 60]; + v.extend([0.01, 0.02, 0.05, 0.3]); + assert!(robust_band(&v, 5.0, false, None).is_none()); + let keep = robust_outlier_keep(&v, 5.0, Tail::Upper, false, None); + assert!(keep.iter().all(|&k| k)); + } + #[test] fn robust_uniform_keeps_all() { let v = vec![7.0; 20]; diff --git a/src/qc_lib_tests.rs b/src/qc_lib_tests.rs index b73ef03..d004be9 100644 --- a/src/qc_lib_tests.rs +++ b/src/qc_lib_tests.rs @@ -17,6 +17,7 @@ fn default_args() -> QcArgs { qc_histogram: false, qc_mad_on_genes: true, qc_mad_on_counts: true, + qc_mad_on_mito: true, qc_auto_cutoff: false, } } diff --git a/src/utilities/io_helpers.rs b/src/utilities/io_helpers.rs index f7f3255..677cd68 100644 --- a/src/utilities/io_helpers.rs +++ b/src/utilities/io_helpers.rs @@ -62,6 +62,19 @@ pub fn read_col_names( /// `0,2,5`, `1-10`, or `1-20,50-55`. Whitespace around tokens is ignored. /// Ranges are inclusive on both ends (`1-10` -> 1,2,...,10). pub fn parse_index_spec(spec: &str) -> anyhow::Result> { + let out: Vec = index_ranges(spec)? + .into_iter() + .flat_map(|(a, b)| a..=b) + .collect(); + if out.is_empty() { + anyhow::bail!("no valid indices parsed from `{}`", spec); + } + Ok(out) +} + +/// The inclusive ranges an index spec names, as [`parse_index_spec`] reads +/// it, without listing their indices. +pub fn index_ranges(spec: &str) -> anyhow::Result> { let mut out = Vec::new(); for token in spec.split(',') { let token = token.trim(); @@ -81,22 +94,50 @@ pub fn parse_index_spec(spec: &str) -> anyhow::Result> { if start > end { anyhow::bail!("range start {} exceeds end {} in `{}`", start, end, token); } - out.extend(start..=end); + out.push((start, end)); } None => { let idx: usize = token .parse() .map_err(|_| anyhow::anyhow!("invalid index `{}`", token))?; - out.push(idx); + out.push((idx, idx)); } } } - if out.is_empty() { - anyhow::bail!("no valid indices parsed from `{}`", spec); - } Ok(out) } +/// How many indices a spec names, without listing them: a range can be any +/// size. +pub fn index_count(spec: &str) -> anyhow::Result { + Ok(index_ranges(spec)? + .iter() + // b >= a; a range to usize::MAX would wrap `b - a + 1` to 0. + .fold(0usize, |n, (a, b)| { + n.saturating_add((b - a).saturating_add(1)) + })) +} + +/// [`parse_index_spec`] on a side of `n` entries: ranges end at its end and +/// indices past it are dropped. More than `n` left means one repeats, which +/// is refused before anything is listed. +pub fn parse_index_spec_within(spec: &str, n: usize) -> anyhow::Result> { + let ranges: Vec<(usize, usize)> = index_ranges(spec)? + .into_iter() + .filter(|&(a, _)| a < n) + .map(|(a, b)| (a, b.min(n - 1))) + .collect(); + let count = ranges + .iter() + .fold(0usize, |c, (a, b)| c.saturating_add(b - a + 1)); + anyhow::ensure!(count > 0, "no index in `{spec}` is below {n}"); + anyhow::ensure!( + count <= n, + "`{spec}` repeats an index: it names {count} of {n}" + ); + Ok(ranges.into_iter().flat_map(|(a, b)| a..=b).collect()) +} + // Re-export constants for convenience pub use crate::sparse_io::{MAX_COLUMN_NAME_IDX, MAX_ROW_NAME_IDX}; @@ -137,3 +178,26 @@ mod tests { assert!(parse_index_spec("").is_err()); } } + +#[cfg(test)] +mod index_spec_tests { + use super::{index_count, parse_index_spec_within}; + + #[test] + fn ranges_are_counted_and_cut_at_the_end() { + assert_eq!(index_count("0-9,20").unwrap(), 11); + assert_eq!(index_count("0-1000000000000").unwrap(), 1_000_000_000_001); + let to_max = format!("0-{}", usize::MAX); + assert_eq!(index_count(&to_max).unwrap(), usize::MAX); + // A huge range on a side of 5 lists 5, not a trillion. + assert_eq!( + parse_index_spec_within("0-1000000000000", 5).unwrap(), + [0, 1, 2, 3, 4] + ); + assert_eq!(parse_index_spec_within("3,7,1-2", 5).unwrap(), [3, 1, 2]); + // Repeats are refused before they are listed. + assert!(parse_index_spec_within(&"0-4,".repeat(1000), 5).is_err()); + assert!(parse_index_spec_within("9-12", 5).is_err()); + assert!(index_count("5-2").is_err()); + } +} diff --git a/src/utilities/name_matching.rs b/src/utilities/name_matching.rs index 621677b..751b712 100644 --- a/src/utilities/name_matching.rs +++ b/src/utilities/name_matching.rs @@ -575,3 +575,17 @@ pub fn match_by_substring( #[cfg(test)] mod tests; + +/// The places of the names `pattern` (a regex, matched case-insensitively) +/// finds in, in order. +pub fn names_matching(names: &[Box], pattern: &str) -> anyhow::Result> { + let re = regex::RegexBuilder::new(pattern) + .case_insensitive(true) + .build() + .map_err(|e| anyhow::anyhow!("invalid pattern `{pattern}`: {e}"))?; + Ok(names + .iter() + .enumerate() + .filter_map(|(i, n)| re.is_match(n).then_some(i)) + .collect()) +} diff --git a/src/zarr_io.rs b/src/zarr_io.rs index 40ab4f7..4c9cdd0 100644 --- a/src/zarr_io.rs +++ b/src/zarr_io.rs @@ -132,6 +132,40 @@ pub fn apply_zip_flag(output: &str, zip: bool, backend: &SparseIoBackend) -> Box } } +/// Where a handler given `-o output` writes, before it writes anything. +pub struct OutputTarget { + /// The output as the handler takes it: `.zarr.zip` added under zip. + pub effective: Box, + pub backend: SparseIoBackend, + /// The file the backend is written to (the un-zipped `.zarr` directory + /// for a `.zarr.zip` target). + pub working: Box, + /// What the user finds once [`finalize_output`] has run: the archive for + /// a `.zarr.zip` target, else the working file. + pub found: Box, +} + +/// [`OutputTarget`] for `output`, touching nothing. +pub fn output_target( + output: &str, + backend: SparseIoBackend, + zip: bool, +) -> anyhow::Result { + let effective = apply_zip_flag(output, zip, &backend); + let (backend, working) = crate::hdf5_io::resolve_backend_file(&effective, Some(backend))?; + let found = if effective.ends_with(".zarr.zip") { + effective.clone() + } else { + working.clone() + }; + Ok(OutputTarget { + effective, + backend, + working, + found, + }) +} + /// Resolve the output target for a handler that writes a fresh backend, and /// clear any pre-existing file at that path. Returns /// `(effective_output, backend, working_file)`, where `working_file` is the @@ -142,15 +176,13 @@ pub fn prepare_output( backend: SparseIoBackend, zip: bool, ) -> anyhow::Result<(Box, SparseIoBackend, Box)> { - use crate::hdf5_io::resolve_backend_file; use legume_numeric::matrix::common_io::remove_file; - let effective_output = apply_zip_flag(output, zip, &backend); - let (backend, working_file) = resolve_backend_file(&effective_output, Some(backend))?; - if std::path::Path::new(working_file.as_ref()).exists() { - remove_file(&working_file)?; + let t = output_target(output, backend, zip)?; + if std::path::Path::new(t.working.as_ref()).exists() { + remove_file(&t.working)?; } - Ok((effective_output, backend, working_file)) + Ok((t.effective, t.backend, t.working)) } /// Finalize a handler's output: re-zip the working `.zarr` directory when the diff --git a/tests/mcp_server.rs b/tests/mcp_server.rs new file mode 100644 index 0000000..f0d6b5c --- /dev/null +++ b/tests/mcp_server.rs @@ -0,0 +1,624 @@ +//! `data-beans mcp` driven over stdio, as an agent's client drives it: a +//! small backend is built, read through the tools, and derived from, and a +//! write onto an existing file is refused. + +// The binary needs `tui`. +#![cfg(all(feature = "mcp", feature = "tui"))] + +use serde_json::{json, Value}; +use std::io::{BufRead, BufReader, Write}; +use std::path::Path; +use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio}; + +const EXE: &str = env!("CARGO_BIN_EXE_data-beans"); + +struct Client { + child: Child, + stdin: ChildStdin, + stdout: BufReader, + next: u64, +} + +impl Client { + /// A server keeping no watch folder (none in the real runtime dir). + fn start(dir: &Path) -> Self { + Self::start_with(dir, &[("LEGUME_MCP_WATCH", Path::new("off"))]) + } + + fn start_with(dir: &Path, env: &[(&str, &Path)]) -> Self { + let mut child = Command::new(EXE) + .arg("mcp") + .current_dir(dir) + .envs(env.iter().copied()) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .spawn() + .unwrap(); + let stdin = child.stdin.take().unwrap(); + let stdout = BufReader::new(child.stdout.take().unwrap()); + Self { + child, + stdin, + stdout, + next: 0, + } + } + + fn request(&mut self, method: &str, params: Value) -> Value { + self.next += 1; + let id = self.next; + let msg = json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }); + writeln!(self.stdin, "{msg}").unwrap(); + loop { + let mut line = String::new(); + assert!( + self.stdout.read_line(&mut line).unwrap() > 0, + "server closed" + ); + let v: Value = serde_json::from_str(&line).unwrap(); + if v["id"] == json!(id) { + return v; + } + } + } + + /// A tool's result: its structured content, or its error text. + fn call(&mut self, name: &str, args: Value) -> Result { + let r = self.request("tools/call", json!({ "name": name, "arguments": args })); + let r = &r["result"]; + if r["isError"] == json!(true) { + Err(r["content"][0]["text"].as_str().unwrap().to_string()) + } else { + Ok(r["structuredContent"].clone()) + } + } +} + +impl Drop for Client { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +/// A 6 x 4 backend `base.zarr.zip` in `dir`: rows MT-CO1, MT-ND1, G2..G5; +/// row i is non-zero in column (i mod 4), with count i + 1. +fn backend(dir: &Path) { + let mut mtx = String::from("%%MatrixMarket matrix coordinate integer general\n6 4 6\n"); + for i in 0..6 { + mtx += &format!("{} {} {}\n", i + 1, i % 4 + 1, i + 1); + } + std::fs::write(dir.join("m.mtx"), mtx).unwrap(); + std::fs::write(dir.join("rows.tsv"), "MT-CO1\nMT-ND1\nG2\nG3\nG4\nG5\n").unwrap(); + std::fs::write(dir.join("cols.tsv"), "C0\nC1\nC2\nC3\n").unwrap(); + let status = Command::new(EXE) + .args([ + "from-mtx", "m.mtx", "-r", "rows.tsv", "-c", "cols.tsv", "-o", "base", + ]) + .current_dir(dir) + .stderr(Stdio::null()) + .status() + .unwrap(); + assert!(status.success()); +} + +/// A server in a fresh directory holding the `backend` fixture, past its +/// handshake. +fn setup() -> (tempfile::TempDir, Client) { + let dir = tempfile::tempdir().unwrap(); + backend(dir.path()); + let mut c = Client::start(dir.path()); + c.request("initialize", json!({ "protocolVersion": "2025-06-18" })); + (dir, c) +} + +#[test] +fn tools_read_derive_and_refuse_to_replace() { + let dir = tempfile::tempdir().unwrap(); + backend(dir.path()); + let mut c = Client::start(dir.path()); + + let init = c.request("initialize", json!({ "protocolVersion": "2025-06-18" })); + assert_eq!(init["result"]["serverInfo"]["name"], "data-beans"); + let tools = c.request("tools/list", json!({})); + let names: Vec<&str> = tools["result"]["tools"] + .as_array() + .unwrap() + .iter() + .map(|t| t["name"].as_str().unwrap()) + .collect(); + assert!(names.contains(&"statistics") && !names.contains(&"squeeze")); + + let info = c.call("info", json!({ "data": "base.zarr.zip" })).unwrap(); + assert_eq!(info["rows"], 6); + assert_eq!(info["columns"], 4); + assert_eq!(info["nonzeros"], 6); + + let mt = c + .call( + "names", + json!({ "data": "base.zarr.zip", "axis": "rows", "pattern": "^mt-" }), + ) + .unwrap(); + assert_eq!(mt["matched"], 2); + assert_eq!(mt["names"][1], json!({ "index": 1, "name": "MT-ND1" })); + + // Row tot is i + 1: G5 (6) tops it, MT-CO1 (1) is last. + let stat = c + .call( + "statistics", + json!({ "data": ["base.zarr.zip"], "axis": "rows", "sort_by": "tot", "limit": 1 }), + ) + .unwrap(); + assert_eq!(stat["count"], 6); + assert_eq!(stat["top"][0]["name"], "G5"); + assert_eq!(stat["quantiles"]["tot"]["min"], 1.0); + + let vals = c + .call( + "values", + json!({ "data": "base.zarr.zip", "axis": "rows", "names": ["G4"] }), + ) + .unwrap(); + // G4 is row 4: count 5 in column C0, the only line kept. + assert_eq!(vals["kept"], 1); + assert_eq!(vals["values"][0], json!({ "column_name": "C0", "G4": 5.0 })); + + let sub = c + .call( + "subset", + json!({ "data": "base.zarr.zip", "axis": "rows", "names": ["MT-CO1", "G2"], "out": "sub" }), + ) + .unwrap(); + assert!(sub["written"].as_str().unwrap().ends_with("sub.zarr.zip")); + assert_eq!(sub["info"]["rows"], 2); + + let again = c.call( + "subset", + json!({ "data": "base.zarr.zip", "axis": "rows", "names": ["G3"], "out": "sub" }), + ); + assert!(again.unwrap_err().contains("not replacing")); + + let dirout = c + .call( + "convert", + json!({ "data": "base.zarr.zip", "out": "copy.zarr" }), + ) + .unwrap(); + assert!(dir.path().join("copy.zarr").is_dir()); + assert_eq!(dirout["info"]["nonzeros"], 6); + + let merged = c + .call( + "merge", + json!({ "data": ["base.zarr.zip", "copy.zarr"], "out": "both" }), + ) + .unwrap(); + assert!(dir.path().join("both.zarr.zip").exists()); + assert_eq!(merged["info"]["columns"], 8); + let mismatched = c.call( + "merge", + json!({ "data": ["base.zarr.zip", "sub.zarr.zip"], "out": "x" }), + ); + assert!(mismatched.unwrap_err().contains("other row names")); + + // `out` is the only way to name an output. + let refused = c.call( + "import", + json!({ "command": "from-mtx", "input": "m.mtx", "out": "y", "options": ["--output=z"] }), + ); + assert!(refused.unwrap_err().contains("not taken")); + // Nor hidden in a cluster of short options. + let refused = c.call( + "import", + json!({ "command": "from-mtx", "input": "m.mtx", "out": "y", "options": ["-Io"] }), + ); + assert!(refused.unwrap_err().contains("apart from its value")); + assert!(!dir.path().join("y.zarr.zip").exists()); +} + +#[test] +fn qc_plot_and_merge_on_common_rows() { + let (dir, mut c) = setup(); + + // C2 and C3 hold one entry each, under the near-empty floor of 2; C0 + // and C1 hold two, with MT-CO1 1 of 6 counts and MT-ND1 2 of 8. + let dry = c.call("qc", json!({ "data": "base.zarr.zip" })).unwrap(); + assert_eq!(dry["cells"], 4); + assert_eq!(dry["near_empty"], 2); + assert_eq!(dry["cells_kept"], 2); + assert!(dry.get("written").is_none()); + assert!(dry["mito_frac"]["median"].is_number()); + let kept = c + .call( + "qc", + json!({ "data": "base.zarr.zip", "out": "kept", "report": "qc.tsv" }), + ) + .unwrap(); + assert_eq!(kept["info"]["columns"], 2); + assert!(dir.path().join("qc.tsv").exists()); + + if cfg!(feature = "plot") { + let r = c.request( + "tools/call", + json!({ "name": "plot", "arguments": { "data": "base.zarr.zip", "out": "qc.png" } }), + ); + let content = &r["result"]["content"]; + assert_eq!(content[0]["type"], "image"); + assert_eq!(content[0]["mimeType"], "image/png"); + assert!(dir.path().join("qc.png").exists()); + } + + let sub = c + .call( + "subset", + json!({ "data": "base.zarr.zip", "axis": "rows", "names": ["G2", "MT-CO1"], "out": "sub" }), + ) + .unwrap(); + assert_eq!(sub["info"]["rows"], 2); + let merged = c + .call( + "merge", + json!({ "data": ["base.zarr.zip", "sub.zarr.zip"], "out": "both", "common_rows": true }), + ) + .unwrap(); + assert_eq!(merged["info"]["rows"], 2); + assert_eq!(merged["info"]["columns"], 8); + // Each batch keeps its two rows' entries: MT-CO1 (1) and G2 (3). + assert_eq!(merged["info"]["nonzeros"], 4); + assert!(merged["batches"] + .as_str() + .unwrap() + .ends_with("both.batch.gz")); +} + +#[test] +fn unnamed_outputs_go_to_the_results_folder_and_the_records() { + let (dir, mut c) = setup(); + + let imported = c + .call( + "import", + json!({ "command": "from-mtx", "input": "m.mtx", "options": ["-r", "rows.tsv", "-c", "cols.tsv"], "source": "the test's fixture" }), + ) + .unwrap(); + let path = imported["written"].as_str().unwrap(); + assert!(path.contains("/results/") && path.ends_with("m.imported.zarr.zip")); + // A second unnamed output of the same kind takes the next free name. + let first = c + .call("subsample", json!({ "data": path, "cells": 2 })) + .unwrap(); + let second = c + .call("subsample", json!({ "data": path, "cells": 3 })) + .unwrap(); + assert!(first["written"] + .as_str() + .unwrap() + .ends_with("m.imported.subsample.zarr.zip")); + assert!(second["written"] + .as_str() + .unwrap() + .ends_with("m.imported.subsample-2.zarr.zip")); + + // A figure's number goes before its extension. + if cfg!(feature = "plot") { + for want in ["m.imported.qc.png", "m.imported.qc-2.png"] { + let r = c.request( + "tools/call", + json!({ "name": "plot", "arguments": { "data": path } }), + ); + let text = r["result"]["content"][1]["text"] + .as_str() + .unwrap_or_default(); + assert!(text.contains(want), "{r}"); + } + } + + let index = std::fs::read_to_string(dir.path().join("results/INDEX.tsv")).unwrap(); + assert!(index.contains("m.imported.subsample-2.zarr.zip")); + let readme = std::fs::read_to_string(dir.path().join("data/README.md")).unwrap(); + assert!(readme.contains("the test's fixture")); +} + +/// What a call should come to. +enum Want { + /// Not an error, and the check holds on its structured content. + Ok(fn(&Value) -> bool), + /// A tool error (not a protocol one) whose text has this in it. + Refused(&'static str), +} + +#[test] +fn every_tool_answers_or_refuses_as_it_should() { + use Want::{Ok, Refused}; + let (_dir, mut c) = setup(); + + // Every schema is an object whose required keys are among its properties. + let tools = c.request("tools/list", json!({})); + for t in tools["result"]["tools"].as_array().unwrap() { + let s = &t["inputSchema"]; + assert_eq!(s["type"], "object", "{}", t["name"]); + for k in s["required"].as_array().into_iter().flatten() { + assert!( + s["properties"].get(k.as_str().unwrap()).is_some(), + "{}: {k}", + t["name"] + ); + } + } + + let cases: Vec<(&str, Value, Want)> = vec![ + ( + "info", + json!({ "data": "missing.zarr.zip" }), + Refused("failed"), + ), + ("info", json!({ "data": "-x" }), Refused("write it as ./-x")), + ( + "list_contents", + json!({ "file": "base.zarr.zip" }), + Ok(|_| true), + ), + ( + "names", + json!({ "data": "base.zarr.zip", "axis": "columns", "offset": 3, "limit": 5 }), + Ok(|v| v["names"].as_array().unwrap().len() == 1), + ), + ( + "names", + json!({ "data": "base.zarr.zip", "axis": "sideways" }), + Refused("rows or columns"), + ), + ( + "names", + json!({ "data": "base.zarr.zip", "axis": "rows", "pattern": "(" }), + Refused("pattern"), + ), + ( + "statistics", + json!({ "data": ["base.zarr.zip"], "axis": "columns", "limit": 0 }), + Ok(|v| v["top"].as_array().unwrap().len() == 1), + ), + ( + "statistics", + json!({ "data": ["base.zarr.zip"], "axis": "rows", "sort_by": "median" }), + Refused("median"), + ), + ( + "statistics", + json!({ "data": ["base.zarr.zip"], "axis": "rows", "out": "s.parquet" }), + Ok(|v| v["written"].as_str().unwrap().ends_with("s.parquet")), + ), + ( + "statistics", + json!({ "data": ["base.zarr.zip"], "axis": "rows", "out": "s.parquet" }), + Refused("not replacing"), + ), + ( + "values", + json!({ "data": "base.zarr.zip", "axis": "rows", "indices": "0-100000000000" }), + Refused("at most 100"), + ), + ( + "values", + json!({ "data": "base.zarr.zip", "axis": "rows" }), + Refused("one of them"), + ), + ( + "qc", + json!({ "data": "base.zarr.zip", "n_mads": "three" }), + Refused("`n_mads` is a number"), + ), + ( + "import", + json!({ "command": "rm-rf", "input": "m.mtx" }), + Refused("one of"), + ), + // m.mtx sits beside `m`, but it is the input, and stays. + ( + "import", + json!({ "command": "from-mtx", "input": "m.mtx", "out": "m" }), + Ok(|v| v["written"].as_str().unwrap().ends_with("m.zarr.zip")), + ), + // The same again: m.zarr.zip is there now, and is no input. + ( + "import", + json!({ "command": "from-mtx", "input": "m.mtx", "out": "m" }), + Refused("not replacing"), + ), + // Rows by pattern, and by a range of indices. + ( + "subset", + json!({ "data": "base.zarr.zip", "axis": "rows", "pattern": "^g[23]$", "out": "g23" }), + Ok(|v| v["info"]["rows"] == 2), + ), + ( + "subset", + json!({ "data": "base.zarr.zip", "axis": "rows", "indices": "0-2", "out": "first3" }), + Ok(|v| v["info"]["rows"] == 3), + ), + ( + "subset", + json!({ "data": "base.zarr.zip", "axis": "rows", "pattern": "^nope$", "out": "none" }), + Refused("no rows match"), + ), + // An index list keeps its order; an empty names list is no choice. + ( + "subset", + json!({ "data": "base.zarr.zip", "axis": "rows", "indices": [3, 1, 2], "names": [], "out": "in_order" }), + Ok(|v| v["info"]["rows"] == 3), + ), + // A selection that fails copies nothing beside `out`. + ( + "subset", + json!({ "data": "base.zarr.zip", "axis": "rows", "indices": "0-5,0-5", "out": "repeats" }), + Refused("repeats an index"), + ), + ( + "subset", + json!({ "data": "base.zarr.zip", "axis": "rows", "indices": "0-2", "out": "repeats" }), + Ok(|v| v["info"]["rows"] == 3), + ), + // A selection flag of its own: the command refuses a second one. + ( + "subset", + json!({ "data": "base.zarr.zip", "axis": "rows", "indices": [0], "out": "x1", "options": ["-i", "1,2"] }), + Refused("cannot be used multiple times"), + ), + // An output that is the input itself. + ( + "convert", + json!({ "data": "base.zarr.zip", "out": "base.zarr.zip" }), + Refused("is an input"), + ), + // Beside its input under the same name, in another format. + ( + "convert", + json!({ "data": "base.zarr.zip", "out": "base.zarr" }), + Ok(|v| v["written"].as_str().unwrap().ends_with("base.zarr")), + ), + ( + "import", + json!({ "command": "from-mtx", "input": "m.mtx", "out": "y", "options": ["--qc-report=z"] }), + Refused("not taken"), + ), + ( + "subset", + json!({ "data": "base.zarr.zip", "axis": "columns", "indices": [0, 1, 2], "out": "three" }), + Ok(|v| v["info"]["columns"] == 3), + ), + ( + "subsample", + json!({ "data": "base.zarr.zip", "cells": "2", "out": "s2" }), + Refused("is a number"), + ), + ( + "subsample", + json!({ "data": "base.zarr.zip", "cells": 2, "out": "s2" }), + Ok(|v| v["info"]["columns"] == 2), + ), + ( + "split", + json!({ "data": "base.zarr.zip", "folds": 2, "out": "cv" }), + Ok(|v| v["written"].as_array().unwrap().len() == 4), + ), + ( + "split", + json!({ "data": "base.zarr.zip", "folds": 2, "out": "cv" }), + Refused("not replacing"), + ), + ( + "merge", + json!({ "data": ["base.zarr.zip"] }), + Refused("two or more"), + ), + ( + "notebook", + json!({ "kind": "observation", "text": "every tool checked" }), + Ok(|_| true), + ), + ]; + for (tool, args, want) in cases { + let r = c.request("tools/call", json!({ "name": tool, "arguments": args })); + // A tool's failure is a result, never a protocol error. + assert!(r.get("error").is_none(), "{tool} {args}: {r}"); + let r = &r["result"]; + let text = r["content"][0]["text"].as_str().unwrap_or_default(); + match want { + Ok(check) => { + assert_eq!(r["isError"], false, "{tool} {args}: {text}"); + assert!(check(&r["structuredContent"]), "{tool} {args}: {r}"); + } + Refused(why) => { + assert_eq!(r["isError"], true, "{tool} {args}: {r}"); + assert!(text.contains(why), "{tool} {args}: {text}"); + } + } + } + + if cfg!(feature = "plot") { + let r = c.call("plot", json!({ "data": "base.zarr.zip", "kind": "violin" })); + assert!(r.unwrap_err().contains("qc or stat")); + } + + // An unknown tool is a protocol error. + let r = c.request( + "tools/call", + json!({ "name": "no_such_tool", "arguments": {} }), + ); + assert!(r.get("error").is_some(), "{r}"); +} + +#[test] +fn a_cancelled_call_gets_no_answer_and_the_server_goes_on() { + let (_dir, mut c) = setup(); + let id = 1000; + let call = json!({ "jsonrpc": "2.0", "id": id, "method": "tools/call", + "params": { "name": "info", "arguments": { "data": "base.zarr.zip" } } }); + let cancel = json!({ "jsonrpc": "2.0", "method": "notifications/cancelled", + "params": { "requestId": id } }); + writeln!(c.stdin, "{call}\n{cancel}").unwrap(); + // The ping is answered; whatever comes before it is not the call's + // answer unless the call finished before the cancel arrived. + let pong = c.request("ping", json!({})); + assert!(pong.get("result").is_some(), "{pong}"); +} + +#[test] +fn a_person_follows_the_server_and_leaves_the_agent_a_note() { + use legume_mcp::watch::{self, view::AgentPane, view::Handled}; + use ratatui::{backend::TestBackend, crossterm::event::KeyCode, Terminal}; + + let dir = tempfile::tempdir().unwrap(); + backend(dir.path()); + let base = dir.path().join("watch"); + let mut c = Client::start_with(dir.path(), &[("LEGUME_MCP_WATCH", &base)]); + c.request("initialize", json!({})); + c.call("info", json!({ "data": "base.zarr.zip" })).unwrap(); + + let server = watch::live_in(&base) + .pop() + .expect("the server's watch folder"); + assert_eq!(server.meta.server, "data-beans"); + let mut pane = AgentPane::new(&server.dir, true).unwrap(); + let screen = |pane: &mut AgentPane| { + let mut t = Terminal::new(TestBackend::new(160, 30)).unwrap(); + t.draw(|f| pane.render(f.area(), f.buffer_mut())).unwrap(); + t.backend().to_string() + }; + let shows = |pane: &mut AgentPane, want: &str| { + for _ in 0..200 { + pane.poll(); + if screen(pane).contains(want) { + return; + } + std::thread::sleep(std::time::Duration::from_millis(50)); + } + panic!("the pane never showed {want}:\n{}", screen(pane)); + }; + shows(&mut pane, "info"); + + // The person leaves a note; the agent's next result carries it. + pane.key(KeyCode::Char('n').into()); + for ch in "use exact names".chars() { + pane.key(KeyCode::Char(ch).into()); + } + pane.key(KeyCode::Enter.into()); + let r = c.request( + "tools/call", + json!({ "name": "subset", "arguments": { "data": "base.zarr.zip", "axis": "rows", "names": ["G2"], "out": "g2" } }), + ); + assert!( + r.to_string() + .contains("Note from the user: use exact names"), + "{r}" + ); + shows(&mut pane, "subset"); + + // Opened, the subset hands over what it wrote. + match pane.key(KeyCode::Char('o').into()) { + Handled::Open(call) => assert!(call.text.contains("g2.zarr.zip"), "{}", call.text), + _ => panic!("o opens the subset"), + } +}