diff --git a/crates/escapepod-cli/src/commands/demux/detect.rs b/crates/escapepod-cli/src/commands/demux/detect.rs index a3982e9..71ed803 100644 --- a/crates/escapepod-cli/src/commands/demux/detect.rs +++ b/crates/escapepod-cli/src/commands/demux/detect.rs @@ -63,13 +63,13 @@ pub struct DetectArgs { /// `scripts/export_adapter_cnn_to_onnx.py` (those weights are CC BY-NC 4.0 /// and not bundled). Runs batched on the CPU by default; pass `--gpu` (with /// a `--features cnn-gpu` build) for onnxruntime CUDA inference. - #[arg( - long, - default_value = "llr", - value_name = "{llr,cnn}", - help_heading = "Advanced Options" - )] - pub method: String, + /// + /// **No default** — LLR is opt-in, never inferred. It costs 17.2 points of + /// downstream barcode recall against the same classifier (0.9928 -> 0.8196, + /// escapepod-models#16) and the failure is silent: it runs and produces + /// plausible-looking boundaries. + #[arg(long, value_name = "{cnn,llr}", help_heading = "Advanced Options")] + pub method: Option, /// Path to the boundary-CNN ONNX model (only used with `--method cnn`). #[cfg(feature = "cnn-detect")] @@ -160,7 +160,16 @@ fn llr_boundaries( /// Run the detect subcommand. pub fn run(args: DetectArgs) -> anyhow::Result<()> { - match args.method.as_str() { + let Some(method) = args.method.clone() else { + anyhow::bail!( + "--method {{cnn,llr}} is required: LLR is never chosen for you. Use \ + `--method cnn --cnn-model ` for the accuracy the shipped barcode \ + models were measured at, or `--method llr` to opt into the classical \ + detector (17.2 points worse on downstream barcode recall — \ + escapepod-models#16)." + ); + }; + match method.as_str() { "llr" => run_llr(args), "cnn" => { #[cfg(feature = "cnn-detect")] diff --git a/crates/escapepod-cli/src/commands/demux/info.rs b/crates/escapepod-cli/src/commands/demux/info.rs new file mode 100644 index 0000000..08cd056 --- /dev/null +++ b/crates/escapepod-cli/src/commands/demux/info.rs @@ -0,0 +1,295 @@ +//! `escpod demux --model --info`: what is this model, and what will it do? +//! +//! Answers the questions you would otherwise answer by reading a sidecar and a +//! provenance file side by side: what the model is, what it needs from you, +//! what ships inside it, and what it scored. Prints without touching any POD5, +//! so it is safe to run against a model you are about to trust. +//! +//! Supports the CTC-CRF bundle (rich: geometry, references, pinned detector, +//! metrics) and the DTW-SVM/GBM classifier JSONs (barcode set and shape). The +//! CRF is the one that benefits, because it is the one carrying a bundle. + +use std::path::Path; + +use crate::style; +use escapepod_demux::{AnyModel, load_any_model}; + +/// Print everything the model can tell us about itself. +pub fn run(model_path: &Path) -> anyhow::Result<()> { + #[cfg(feature = "crf-decode")] + if let Some(dir) = super::run::crf_bundle_dir(model_path) { + return crf_info(&dir); + } + classifier_info(model_path) +} + +/// Human-readable byte size; bundles are small enough that MB is the ceiling. +fn human(n: u64) -> String { + const K: u64 = 1024; + match n { + 0..K => format!("{n} B"), + K..0x100000 => format!("{:.1} KB", n as f64 / K as f64), + _ => format!("{:.1} MB", n as f64 / (K * K) as f64), + } +} + +fn heading(s: &str) { + println!("\n{}", style::action(s)); +} + +fn field(k: &str, v: impl std::fmt::Display) { + println!(" {:<22} {}", style::label(k), v); +} + +/// Flatten nested metric JSON into `a.b.c = value` lines. +/// +/// Metrics are untyped by design (see `CrfMetadata::metrics`), so rather than +/// guessing a shape, walk whatever is there. Arrays of scalars collapse onto +/// one line; anything deeper keeps its path, which is what makes a +/// per-recovery-level table readable. +fn walk_metrics(prefix: &str, v: &serde_json::Value, out: &mut Vec<(String, String)>) { + match v { + serde_json::Value::Object(map) => { + for (k, val) in map { + let p = if prefix.is_empty() { + k.clone() + } else { + format!("{prefix}.{k}") + }; + walk_metrics(&p, val, out); + } + } + serde_json::Value::Array(items) => { + if items.iter().all(|i| !i.is_object() && !i.is_array()) { + let joined: Vec = items.iter().map(render_scalar).collect(); + out.push((prefix.to_string(), joined.join(", "))); + } else { + for (i, item) in items.iter().enumerate() { + walk_metrics(&format!("{prefix}[{i}]"), item, out); + } + } + } + other => out.push((prefix.to_string(), render_scalar(other))), + } +} + +fn render_scalar(v: &serde_json::Value) -> String { + match v { + serde_json::Value::String(s) => s.clone(), + // Metrics arrive as f64, so a recall prints as 0.9877713334625322. + // Four decimals is the precision these are quoted at everywhere else; + // integral floats (counts) lose the trailing `.0`. + serde_json::Value::Number(n) => match n.as_f64() { + Some(f) if n.is_f64() && f.fract() == 0.0 => format!("{f:.0}"), + Some(f) if n.is_f64() => format!("{f:.4}"), + _ => n.to_string(), + }, + other => other.to_string(), + } +} + +#[cfg(feature = "crf-decode")] +fn crf_info(dir: &Path) -> anyhow::Result<()> { + use escapepod_demux::crf::{BarcodeRefs, CrfMetadata}; + + let meta = CrfMetadata::load(dir.join("metadata.json"))?; + + heading("Model"); + field("kind", "CTC-CRF barcode basecaller"); + field("bundle", dir.display()); + if let Some(m) = &meta.model { + field("id", &m.id); + if let Some(v) = &m.version { + field("version", v); + } + if let Some(c) = &m.chemistry { + field("chemistry", c); + } + if let Some(n) = &m.notes { + field("notes", n); + } + } + + heading("Signal geometry"); + field( + "window", + format!( + "[adapter_end - {}, adapter_end] ({} samples)", + meta.signal.chunk, meta.signal.chunk + ), + ); + field("stride", meta.signal.stride); + field("timesteps", meta.signal.chunk / meta.signal.stride); + field( + "standardisation", + format!( + "mean {:.3}, stdev {:.3}", + meta.standardisation.mean, meta.standardisation.stdev + ), + ); + + heading("Decoder"); + field("state_len", meta.crf.state_len); + field("n_base", meta.crf.n_base); + field("alphabet", meta.crf.alphabet.join("")); + field( + "states", + meta.crf.n_base.pow(meta.crf.state_len as u32).to_string(), + ); + // The single most misread property of this model: it cannot emit the first + // `state_len` bases of its target (escapepod-models#36). + field( + "emits", + format!( + "target[{}:] — the first {} target bases fix the initial state and are never emitted", + meta.crf.state_len, meta.crf.state_len + ), + ); + + heading("Barcode references"); + match &meta.barcodes { + Some(entries) => { + let refs = BarcodeRefs::from_pairs( + entries.iter().map(|e| (e.name.clone(), e.sequence.clone())), + )?; + field("source", "bundled (no --barcodes needed)"); + field("count", refs.len()); + let lens: Vec = entries.iter().map(|e| e.sequence.len()).collect(); + let (lo, hi) = ( + lens.iter().min().copied().unwrap_or(0), + lens.iter().max().copied().unwrap_or(0), + ); + field( + "length", + if lo == hi { + format!("{lo} nt") + } else { + format!("{lo}-{hi} nt") + }, + ); + field( + "min pairwise distance", + refs.min_pairwise_distance() + .map_or_else(|| "n/a".into(), |d| d.to_string()), + ); + println!(); + for e in entries { + println!(" {:<10} {}", style::label(&e.name), e.sequence); + } + } + None => { + field("source", "NOT bundled — --barcodes is required"); + } + } + + heading("Boundary detector"); + match &meta.boundary { + Some(b) => { + field("method", &b.method); + if let Some(id) = &b.model_id { + field("model", id); + } + match &b.onnx { + Some(o) => field("weights", format!("{o} (bundled)")), + None => field("weights", "not bundled — supply --cnn-model"), + } + field("pinned", "yes — this model was calibrated against it"); + } + None => field("pinned", "no — you must pass --method"), + } + + if let Some(metrics) = &meta.metrics { + heading("Published metrics"); + let mut rows = Vec::new(); + walk_metrics("", metrics, &mut rows); + let w = rows.iter().map(|(k, _)| k.len()).max().unwrap_or(0).min(48); + for (k, v) in rows { + println!(" {: = std::fs::read_dir(dir)? + .filter_map(Result::ok) + .filter(|e| e.path().is_file()) + .collect(); + files.sort_by_key(std::fs::DirEntry::file_name); + for f in files { + let size = f.metadata().map(|m| m.len()).unwrap_or(0); + println!( + " {:<28} {:>10}", + f.file_name().to_string_lossy(), + human(size) + ); + } + + heading("Run it"); + let needs_barcodes = meta.barcodes.is_none(); + let needs_method = meta.boundary.is_none(); + let mut cmd = format!(" escpod demux --model {} -d out/", dir.display()); + if needs_barcodes { + cmd.push_str(" \\\n --barcodes "); + } + if needs_method { + cmd.push_str(" \\\n --method cnn --cnn-model "); + } + println!("{cmd}"); + println!(); + Ok(()) +} + +fn classifier_info(path: &Path) -> anyhow::Result<()> { + let model = load_any_model(path)?; + heading("Model"); + field("bundle", path.display()); + let mapper = match &model { + AnyModel::Svm(m) => { + field("kind", "DTW-SVM fingerprint classifier"); + Some(&m.label_mapper) + } + AnyModel::Gbm(m) => { + field("kind", "GBM fingerprint classifier"); + Some(&m.label_mapper) + } + AnyModel::WarpDemux(_) => { + field( + "kind", + "WarpDemuX reference bank (demux classify --reference)", + ); + None + } + }; + if let Some(mapper) = mapper { + let mut ids: Vec = mapper.values().copied().filter(|&i| i >= 0).collect(); + ids.sort_unstable(); + ids.dedup(); + field("barcodes", ids.len()); + field( + "labels", + ids.iter() + .map(|i| format!("BC{i:02}")) + .collect::>() + .join(" "), + ); + } + heading("Boundary detector"); + field("pinned", "no — you must pass --method"); + heading("Bundled files"); + let size = std::fs::metadata(path).map(|m| m.len()).unwrap_or(0); + println!( + " {:<28} {:>10}", + path.file_name().unwrap_or_default().to_string_lossy(), + human(size) + ); + println!(); + Ok(()) +} diff --git a/crates/escapepod-cli/src/commands/demux/mod.rs b/crates/escapepod-cli/src/commands/demux/mod.rs index 582a723..f2190e6 100644 --- a/crates/escapepod-cli/src/commands/demux/mod.rs +++ b/crates/escapepod-cli/src/commands/demux/mod.rs @@ -19,6 +19,7 @@ mod classify; mod detect; mod fingerprint; mod fp_io; +mod info; #[cfg(feature = "demux-models")] pub mod models; mod run; diff --git a/crates/escapepod-cli/src/commands/demux/run.rs b/crates/escapepod-cli/src/commands/demux/run.rs index a27fd1e..b3b2f7a 100644 --- a/crates/escapepod-cli/src/commands/demux/run.rs +++ b/crates/escapepod-cli/src/commands/demux/run.rs @@ -21,6 +21,8 @@ use std::sync::mpsc::SyncSender; use crate::progress::create_progress_bar; use crate::style; +#[cfg(feature = "crf-decode")] +use escapepod_demux::crf::{BarcodeRefs, CrfEncoder, CrfScratch}; use escapepod_demux::{ AnyModel, DtwSvmModel, GbmModel, GbmPredictor, SvmPredictor, SvmWorkspace, extract_fingerprint_from_signal, load_any_model, @@ -44,11 +46,46 @@ pub struct RunArgs { #[arg(value_name = "FILES")] pub input: Vec, - /// Trained classifier JSON — DTW-SVM (`demux train-svm` / converted - /// WarpDemuX) or native GBM tree ensemble. Auto-detected by JSON shape. - #[arg(long, value_name = "FILE")] + /// Trained classifier — a DTW-SVM / GBM JSON (auto-detected by JSON + /// shape), or a CTC-CRF encoder bundle directory (`metadata.json` + the + /// ONNX graph it names). A CRF bundle also needs `--barcodes`. + #[arg(long, value_name = "FILE|DIR")] pub model: Option, + /// Barcode reference CSV (`name,sequence`) for the CTC-CRF head. Required + /// with a CRF bundle, ignored otherwise: the fingerprint heads carry their + /// own barcode set in the model JSON, whereas the CRF emits sequence and + /// has to be told what to match it against. + /// + /// These must be the sequences the model actually EMITS, which is not the + /// training target: `state_len` leading bases only fix the initial CRF + /// state and are never produced, so a 40-nt target emits 36 nt. Matching + /// against full-length targets still calls the same barcode, but inflates + /// every distance and compresses the confidence margin that `--min-margin` + /// gates on (escapepod-models#36). + #[cfg(feature = "crf-decode")] + #[arg(long, value_name = "FILE")] + pub barcodes: Option, + + /// Call a read `unclassified` when its edit-distance margin to the + /// second-best reference is below this (CRF head only). 0 keeps every + /// call, including outright ties. + #[cfg(feature = "crf-decode")] + #[arg( + long, + default_value = "0", + value_name = "N", + help_heading = "Advanced Options" + )] + pub min_margin: u32, + + /// Describe the model and exit: identity, signal geometry, bundled + /// references, pinned boundary detector, published metrics, and the exact + /// command line it needs. Reads no POD5, so it is safe to run against a + /// model before trusting it. + #[arg(long)] + pub info: bool, + /// Output directory for the per-barcode demultiplexed POD5 files #[arg(short = 'd', long, value_name = "DIR")] pub output_dir: Option, @@ -62,14 +99,17 @@ pub struct RunArgs { #[arg(long, default_value = "barcode", help_heading = "Advanced Options")] pub prefix: String, - /// Adapter detection method: `llr` (default) or `cnn`. - #[arg( - long, - default_value = "llr", - value_name = "{llr,cnn}", - help_heading = "Advanced Options" - )] - pub method: String, + /// Adapter detection method: `cnn` or `llr`. **No default** — LLR is + /// opt-in, never inferred. + /// + /// LLR boundaries cost 17.2 points of barcode recall against the same + /// classifier (0.9928 -> 0.8196, escapepod-models#16) and the failure is + /// silent: it runs and produces plausible output. So a model bundle that + /// pins its detector supplies this, and otherwise you have to say which + /// you want. Passing it explicitly overrides a bundle's choice, except + /// that a bundle pinning `cnn` refuses to be downgraded to `llr`. + #[arg(long, value_name = "{cnn,llr}", help_heading = "Advanced Options")] + pub method: Option, /// Path to the ADAPTed CNN ONNX model (only with `--method cnn`). #[cfg(feature = "cnn-detect")] @@ -332,25 +372,88 @@ fn route( }); } -/// Either classifier head the fused pipeline can drive. Detect + fingerprint are -/// model-agnostic; only the per-read classify differs — DTW-SVM (with an optional -/// GPU DTW path) or the CPU-only GBM tree walk. +/// Which classifier head the fused pipeline drives. +/// +/// The two fingerprint heads (DTW-SVM, with an optional GPU DTW path, and the +/// CPU-only GBM tree walk) share everything up to classify: detect, then a +/// fingerprint of the adapter region, then a model that maps features to a +/// class index. +/// +/// The CRF head is a different shape. It does not fingerprint at all — it +/// basecalls the barcode out of the raw pA window `[adapter_end - chunk, +/// adapter_end]` and matches the decoded sequence to a reference set by edit +/// distance. So its barcode set comes from the reference CSV rather than a +/// `label_mapper`, and its confidence is an edit-distance margin rather than a +/// probability. enum ClassifyModel { Svm(DtwSvmModel), Gbm(GbmModel), + #[cfg(feature = "crf-decode")] + Crf(Box), +} + +/// The CTC-CRF head: encoder bundle plus the references its decodes are matched +/// against. +#[cfg(feature = "crf-decode")] +struct CrfHead { + encoder: CrfEncoder, + refs: BarcodeRefs, + min_margin: u32, } impl ClassifyModel { - /// Class-index → barcode-id map (same shape on both heads), for the output - /// barcode set. - fn label_mapper(&self) -> &HashMap { + /// The set of output barcode labels, before `unclassified` is added. + /// + /// The fingerprint heads name barcodes positionally from the model's + /// `label_mapper` (`BC00`, `BC01`, ...); the CRF head uses the reference + /// names, so its output files are `barcode_nbc01.pod5` rather than + /// `barcode_BC00.pod5`. + fn barcode_names(&self) -> Vec { match self { - ClassifyModel::Svm(m) => &m.label_mapper, - ClassifyModel::Gbm(m) => &m.label_mapper, + ClassifyModel::Svm(m) => barcode_set(&m.label_mapper), + ClassifyModel::Gbm(m) => barcode_set(&m.label_mapper), + #[cfg(feature = "crf-decode")] + ClassifyModel::Crf(h) => { + // Reference order is the CSV's, but every label here becomes a + // router key and an output file, so a duplicate name (or one + // literally called `unclassified`) would leave a writer thread + // with no sender. Dedup rather than trusting the CSV. + let mut v: Vec = Vec::with_capacity(h.refs.len() + 1); + for n in h.refs.names() { + if n != UNCLASSIFIED && !v.iter().any(|s| s == n) { + v.push(n.clone()); + } + } + v.push(UNCLASSIFIED.to_string()); + v + } } } } +/// Is this `--model` a CTC-CRF encoder bundle rather than a classifier JSON? +/// +/// A bundle is a directory holding `metadata.json`, or that `metadata.json` +/// itself. Sniff the `format` key rather than trusting the extension, so a +/// stray `.json` cannot be mistaken for either kind: the CRF sidecar declares +/// `"format": "escapepod-crf-encoder/N"`, which no classifier JSON carries. +#[cfg(feature = "crf-decode")] +pub(super) fn crf_bundle_dir(path: &Path) -> Option { + let (dir, meta) = if path.is_dir() { + (path.to_path_buf(), path.join("metadata.json")) + } else if path.file_name().is_some_and(|n| n == "metadata.json") { + (path.parent()?.to_path_buf(), path.to_path_buf()) + } else { + return None; + }; + let text = std::fs::read_to_string(&meta).ok()?; + let json: serde_json::Value = serde_json::from_str(&text).ok()?; + json.get("format")? + .as_str()? + .starts_with("escapepod-crf-encoder/") + .then_some(dir) +} + pub fn run(args: RunArgs) -> anyhow::Result<()> { use crate::commands::profile::PhaseTimer; let mut timer = PhaseTimer::new(); @@ -359,30 +462,115 @@ pub fn run(args: RunArgs) -> anyhow::Result<()> { // Validate the fused-pipeline args here (not via clap `required`) so the // advanced subcommands aren't forced to supply them. - if args.input.is_empty() { - anyhow::bail!("no input POD5 file(s) given"); - } let model_path = args .model .clone() - .ok_or_else(|| anyhow::anyhow!("--model is required"))?; + .ok_or_else(|| anyhow::anyhow!("--model is required"))?; + // `--info` describes the model and exits: no input, no output dir, no POD5 + // touched. Checked before the input/output validation below so you can + // interrogate a model without inventing arguments for a run you are not + // making. + if args.info { + return super::info::run(&model_path); + } + if args.input.is_empty() { + anyhow::bail!("no input POD5 file(s) given"); + } let output_dir = args .output_dir .clone() .ok_or_else(|| anyhow::anyhow!("-d/--output-dir is required"))?; - // The fused pipeline supports both classifier heads: DTW-SVM (with an - // optional GPU DTW path) and the native GBM tree ensemble (CPU-only). Only - // the legacy reference-bank WarpDemux JSON is rejected here. - let model = match load_any_model(&model_path)? { - AnyModel::Svm(m) => ClassifyModel::Svm(m), - AnyModel::Gbm(m) => ClassifyModel::Gbm(m), - AnyModel::WarpDemux(_) => anyhow::bail!( - "`demux` needs an SVM or GBM model (DtwSvmModel / converted WarpDemuX \ - / native GBM). The reference-CSV path is only on `demux classify --reference`." - ), + // Three heads: DTW-SVM (with an optional GPU DTW path), the native GBM tree + // ensemble (CPU-only), and the CTC-CRF basecaller. A CRF bundle is a + // directory, so check for that before trying to parse `--model` as JSON. + // Only the legacy reference-bank WarpDemux JSON is rejected here. + #[cfg(feature = "crf-decode")] + let crf_dir = crf_bundle_dir(&model_path); + #[cfg(not(feature = "crf-decode"))] + let crf_dir: Option = None; + // Kept because a pinned boundary model's ONNX path is relative to the + // bundle directory, and `crf_dir` is consumed building the head below. + #[cfg(feature = "crf-decode")] + let crf_dir_for_pin = crf_dir.clone(); + + let model = match crf_dir { + #[cfg(feature = "crf-decode")] + Some(dir) => { + let encoder = CrfEncoder::load_bundle(&dir)?; + // References come from the bundle unless the caller overrides them. + // Carrying them in the bundle is what makes the plain + // `--model -d out/` form work, and it fixes the + // emitted-vs-target trimming once at export instead of at every + // call site (escapepod-models#36). + let refs = match (&args.barcodes, &encoder.metadata().barcodes) { + (Some(csv), _) => { + let r = BarcodeRefs::from_csv(csv)?; + if encoder.metadata().barcodes.is_some() { + info!( + "{} overriding the {} references in the bundle", + style::label("Barcodes:"), + style::count(encoder.metadata().barcodes.as_ref().unwrap().len()) + ); + } + r + } + (None, Some(entries)) => BarcodeRefs::from_pairs( + entries.iter().map(|e| (e.name.clone(), e.sequence.clone())), + )?, + (None, None) => anyhow::bail!( + "this CTC-CRF bundle carries no barcode references, so --barcodes \ + is required. The CRF emits sequence rather than a class \ + index and has to be told what to match it against — and those must \ + be the sequences the model EMITS (target[state_len:]), not the \ + full-length training targets." + ), + }; + info!( + "{} {} references, minimum pairwise edit distance {}", + style::label("Barcodes:"), + style::count(refs.len()), + refs.min_pairwise_distance() + .map_or_else(|| "n/a".to_string(), |d| d.to_string()), + ); + ClassifyModel::Crf(Box::new(CrfHead { + encoder, + refs, + min_margin: args.min_margin, + })) + } + #[cfg(not(feature = "crf-decode"))] + Some(_) => unreachable!("crf_dir is always None without the crf-decode feature"), + None => match load_any_model(&model_path)? { + AnyModel::Svm(m) => ClassifyModel::Svm(m), + AnyModel::Gbm(m) => ClassifyModel::Gbm(m), + AnyModel::WarpDemux(_) => anyhow::bail!( + "`demux` needs an SVM, GBM or CTC-CRF model (DtwSvmModel / converted \ + WarpDemuX / native GBM / CRF bundle directory). The reference-bank path \ + is only on `demux classify --reference`." + ), + }, }; - let detector = build_detector(&args)?; + // A CRF bundle may pin the boundary detector it was trained against; the + // ONNX path in the sidecar is relative to the bundle directory. + #[cfg(feature = "crf-decode")] + let boundary_pin = match (&model, &crf_dir_for_pin) { + (ClassifyModel::Crf(h), Some(dir)) => h.encoder.metadata().boundary.as_ref().map(|b| { + if let Some(id) = &b.model_id { + info!( + "{} {} (pinned by the model bundle)", + style::label("Boundary model:"), + style::value(id) + ); + } + (b.method.as_str(), b.onnx.as_ref().map(|o| dir.join(o))) + }), + _ => None, + }; + #[cfg(not(feature = "crf-decode"))] + let boundary_pin: Option<(&str, Option)> = None; + + let detector = build_detector(&args, boundary_pin)?; let fp = FpParams::default(); std::fs::create_dir_all(&output_dir)?; @@ -428,7 +616,7 @@ pub fn run(args: RunArgs) -> anyhow::Result<()> { // and when a channel fills, rayon workers block inside `for_each` on // `send` — and blocked rayon workers cannot be stolen from, so one // saturated writer stalls the whole pool. - let barcodes = barcode_set(model.label_mapper()); + let barcodes = model.barcode_names(); let router_depth = (ROUTER_TOTAL_SLOTS / barcodes.len().max(1)).clamp(256, 4096); let mut routers: Routers = HashMap::new(); @@ -468,16 +656,33 @@ pub fn run(args: RunArgs) -> anyhow::Result<()> { // accelerates adapter detection, so only warn when `--gpu` can do // nothing (CPU classify + CPU detect). #[cfg(any(feature = "gpu", feature = "cnn-gpu"))] - if args.gpu && args.method != "cnn" { + if args.gpu && args.method.as_deref() != Some("cnn") { tracing::warn!( "--gpu has no effect here: GBM classify is CPU-only and \ `--method {}` detection is CPU-only (use `--method cnn` for \ GPU adapter detection).", - args.method + args.method.as_deref().unwrap_or("") ); } produce_cpu_gbm(&args, &detector, gbm, fp, &routers, class_tx.as_ref(), &pb) } + #[cfg(feature = "crf-decode")] + ClassifyModel::Crf(head) => { + // Encoder inference here is tract on the CPU, one read per rayon + // worker. `demux basecall --gpu` has an onnxruntime path; it is not + // wired through the fused pipeline, so say so rather than silently + // ignoring the flag. + #[cfg(any(feature = "gpu", feature = "cnn-gpu"))] + if args.gpu && args.method.as_deref() != Some("cnn") { + tracing::warn!( + "--gpu has no effect here: fused CRF basecalling is CPU-only and \ + `--method {}` detection is CPU-only (use `--method cnn` for GPU \ + adapter detection, or `demux basecall --gpu` for a GPU encoder).", + args.method.as_deref().unwrap_or("") + ); + } + produce_cpu_crf(&args, &detector, head, &routers, class_tx.as_ref(), &pb) + } }; // Drop all senders so the writer threads see EOF. The producers only @@ -880,6 +1085,105 @@ fn produce_cpu_gbm( ) } +/// CRF producer: detect → prep the raw-pA window → CTC-CRF basecall → match the +/// decoded sequence to the references by edit distance. +/// +/// Unlike the fingerprint heads this needs **calibrated pA**, not ADC counts, +/// and it needs the window `[adapter_end - chunk, adapter_end]` to lie inside +/// the decoded prefix. It always does: the detector bounds its decode by +/// `max_obs_trace` and can only report an `adapter_end` inside what it saw, so +/// any read whose adapter was detected at all has its window available. +/// `meta.prep` still returns `None` when `adapter_end < chunk` (the adapter sits +/// too close to the read start), and those route as unclassified. +/// +/// Batching mirrors `produce_cpu_gbm`: detect the whole block at once, then +/// chunk so each rayon task preps, encodes and matches a run of reads. The +/// encode is the dominant cost and tract has no batched LSTM, so parallelism is +/// per read within the chunk. +#[cfg(feature = "crf-decode")] +fn produce_cpu_crf( + args: &RunArgs, + detector: &Detector, + head: &CrfHead, + routers: &Routers, + class_tx: Option<&SyncSender<(Uuid, String, f64)>>, + pb: &indicatif::ProgressBar, +) -> anyhow::Result<()> { + let meta = head.encoder.metadata(); + drive_blocks( + &args.input, + detector.signal_decode_bound(), + |sigs, items| { + const CRF_CHUNK: usize = 256; + let bounds = detector.detect_batch(&sigs); + let rows: Vec<_> = sigs + .into_iter() + .zip(bounds) + .zip(items) + .map(|((sig, b), item)| (sig, b, item)) + .collect(); + + rows.into_par_iter().chunks(CRF_CHUNK).for_each(|chunk| { + let n = chunk.len(); + let mut scratch = CrfScratch::new(); + for (signal, (_s, adapter_end), (read, chunks, run_infos)) in chunk { + let (barcode, conf) = (|| { + let adc = signal.as_ref()?; + // The detector reports `adapter_end` as an index into + // the decoded prefix, which is what `prep` wants. + let pa = adc_to_pa(adc, read.calibration_offset, read.calibration_scale); + let window = meta.prep(&pa, adapter_end)?; + let seq = head + .encoder + .basecall_prepped(&window, &mut scratch) + .inspect_err(|e| tracing::warn!("encoder: {e}")) + .ok()?; + let m = head.refs.match_sequence(seq.as_bytes())?; + // Same gate as `demux basecall --min-margin`, including + // its treatment of a single reference: with no runner-up + // there is no margin to test, so the call stands. + if !m.margin.is_none_or(|v| v >= head.min_margin) { + return None; + } + // Confidence is the margin, matching `demux basecall` + // and `eval_recovery.py`. A lone reference reports 0 + // rather than a fabricated distance. + Some(( + head.refs.name(m.index).to_string(), + f64::from(m.margin.unwrap_or(0)), + )) + })() + .unwrap_or_else(|| (UNCLASSIFIED.to_string(), 0.0)); + + route( + routers, + class_tx, + read.for_writing(read.run_info_index), + barcode, + chunks, + run_infos, + conf, + ); + } + pb.inc(n as u64); + }); + }, + ) +} + +/// ADC counts to picoamps, fused so there is one rounding step. +/// +/// Matches `escapepod_python::adc_to_pa` and `demux basecall`'s own conversion; +/// the reference `pod5` package computes it unfused, which differs by ~1 ulp — +/// thousands of times below the standardisation scale and irrelevant to the +/// decode. +#[cfg(feature = "crf-decode")] +fn adc_to_pa(raw: &[i16], offset: f32, scale: f32) -> Vec { + raw.iter() + .map(|&v| (f32::from(v) + offset) * scale) + .collect() +} + /// GBM counterpart to [`classify_one_cpu`]: fingerprint → GBM tree walk from a /// decoded signal and precomputed boundaries. Returns `(barcode, confidence)`; /// unfingerprintable reads route to `unclassified` (matching the SVM path). @@ -1147,8 +1451,43 @@ fn spawn_class_writer( Ok((Some(tx), Some(handle))) } -fn build_detector(args: &RunArgs) -> anyhow::Result { - match args.method.as_str() { +/// Build the adapter detector: the model bundle's pinned choice, or an explicit +/// `--method`, and an error rather than a silent guess when neither says. +/// +/// `pin` is the detector a CRF bundle declares itself calibrated against +/// (`(method, Some(onnx_path))`). The training window is defined relative to +/// that detector's `adapter_end`, so a pin is a hard requirement rather than a +/// preference. +/// +/// LLR is never inferred. It costs 17.2 points of barcode recall against the +/// same classifier and fails silently (escapepod-models#16), so it has to be +/// asked for by name. Explicit `--method` overrides a pin — that has to stay +/// possible to evaluate a new boundary model — except that a bundle pinning +/// `cnn` refuses the downgrade, which is #16's runtime guard. +fn build_detector( + args: &RunArgs, + pin: Option<(&str, Option)>, +) -> anyhow::Result { + let pinned_method = pin.as_ref().map(|(m, _)| *m); + let pinned_onnx = pin.and_then(|(_, p)| p); + let method = match (args.method.as_deref(), pinned_method) { + (Some("llr"), Some("cnn")) => anyhow::bail!( + "this model is calibrated against CNN adapter boundaries and refuses \ + `--method llr`: LLR costs 17.2 points of barcode recall on the same \ + classifier (0.9928 -> 0.8196) and fails silently. Drop `--method llr`, \ + or use a model that does not pin a detector." + ), + (Some(m), _) => m, + (None, Some(m)) => m, + (None, None) => anyhow::bail!( + "--method {{cnn,llr}} is required: this model does not pin a boundary \ + detector, and LLR is never chosen for you. Use `--method cnn --cnn-model \ + ` for the accuracy the shipped barcode models were measured at, or \ + `--method llr` to opt into the classical detector (17.2 points worse on \ + barcode recall — escapepod-models#16)." + ), + }; + match method { "llr" => Ok(Detector::Llr { min_adapter: args.min_adapter, border_trim: args.border_trim, @@ -1160,7 +1499,13 @@ fn build_detector(args: &RunArgs) -> anyhow::Result { let path = args .cnn_model .as_ref() - .ok_or_else(|| anyhow::anyhow!("--method cnn requires --cnn-model "))?; + .or(pinned_onnx.as_ref()) + .ok_or_else(|| { + anyhow::anyhow!( + "--method cnn requires --cnn-model (this model bundle \ + does not ship a boundary model)" + ) + })?; // `--gpu` with `--method cnn` runs detection on the GPU (one // batched onnxruntime call per block) when built with cnn-gpu. #[cfg(feature = "cnn-gpu")] @@ -1180,6 +1525,7 @@ fn build_detector(args: &RunArgs) -> anyhow::Result { } #[cfg(not(feature = "cnn-detect"))] { + let _ = pinned_onnx; anyhow::bail!("--method cnn requires a build with `--features cnn-detect`") } } @@ -1208,3 +1554,43 @@ fn print_summary(summary: &DemuxSummary) { ); } } + +#[cfg(all(test, feature = "crf-decode"))] +mod tests { + use super::crf_bundle_dir; + + /// `--model` sniffing must not depend on the extension or the file name + /// alone: a CRF bundle is identified by its sidecar's `format` key, so a + /// classifier JSON living next to one, or a directory without a sidecar, + /// still routes to `load_any_model`. + #[test] + fn crf_bundle_detected_by_format_key_not_by_name() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + + // Not a bundle: no metadata.json at all. + assert!(crf_bundle_dir(root).is_none()); + + // Not a bundle: metadata.json exists but declares something else. + let meta = root.join("metadata.json"); + std::fs::write(&meta, r#"{"format":"something-else/1"}"#).unwrap(); + assert!(crf_bundle_dir(root).is_none()); + assert!(crf_bundle_dir(&meta).is_none()); + + // A bundle: recognised via the directory and via the sidecar itself, + // and both resolve to the directory the ONNX is loaded from. + std::fs::write(&meta, r#"{"format":"escapepod-crf-encoder/1"}"#).unwrap(); + assert_eq!(crf_bundle_dir(root).as_deref(), Some(root)); + assert_eq!(crf_bundle_dir(&meta).as_deref(), Some(root)); + + // A classifier JSON is never mistaken for a bundle, even beside one. + let svm = root.join("model.json"); + std::fs::write(&svm, r#"{"label_mapper":{}}"#).unwrap(); + assert!(crf_bundle_dir(&svm).is_none()); + + // Malformed sidecar: fall through to the JSON loader rather than + // failing here, so the error the user sees comes from the real parse. + std::fs::write(&meta, "not json").unwrap(); + assert!(crf_bundle_dir(root).is_none()); + } +} diff --git a/crates/escapepod-cli/src/main.rs b/crates/escapepod-cli/src/main.rs index 438bd1d..acb5916 100644 --- a/crates/escapepod-cli/src/main.rs +++ b/crates/escapepod-cli/src/main.rs @@ -397,9 +397,18 @@ Examples: #[command(after_help = "\ Examples: escpod demux input.pod5 --model model.json -d out/ Fused pipeline (recommended) + escpod demux input.pod5 --model crf_bundle/ --barcodes refs.csv -d out/ \\ + --method cnn --cnn-model adapter.onnx Fused, CTC-CRF head escpod demux detect input.pod5 -o boundaries.csv escpod demux fingerprint input.pod5 --boundaries boundaries.csv -o fingerprints.csv escpod demux classify fingerprints.csv --reference barcodes.csv -o classifications.csv + escpod demux basecall input.pod5 --boundaries boundaries.csv --model crf_bundle/ \\ + --barcodes refs.csv -o classifications.csv + +The fused pipeline drives any of the three classifier heads -- DTW-SVM, GBM or +CTC-CRF -- and decodes each read's signal once. The detect/fingerprint/classify/ +basecall subcommands are those same stages run separately, which re-reads the +POD5 per stage; prefer the fused form unless you want the intermediate files. ")] Demux(commands::demux::DemuxArgs), diff --git a/crates/escapepod-demux/src/crf/barcode.rs b/crates/escapepod-demux/src/crf/barcode.rs index f164ebc..85fc2f5 100644 --- a/crates/escapepod-demux/src/crf/barcode.rs +++ b/crates/escapepod-demux/src/crf/barcode.rs @@ -153,6 +153,41 @@ impl BarcodeRefs { Ok(out) } + /// Build a reference set from `(name, sequence)` pairs. + /// + /// For references that travel inside the model bundle rather than in a + /// separate CSV. Applies the same validation as [`Self::from_csv`] — + /// non-empty, ACGT-only, no duplicate names — because a bundle is no more + /// trustworthy than a file the user pointed at. + pub fn from_pairs(pairs: impl IntoIterator) -> Result + where + N: Into, + S: AsRef, + { + let path = || std::path::PathBuf::from(""); + let mut out = Self::default(); + for (i, (name, seq)) in pairs.into_iter().enumerate() { + let seq = seq.as_ref().to_ascii_uppercase(); + if seq.is_empty() || !seq.bytes().all(|b| b"ACGT".contains(&b)) { + return Err(BarcodeError::BadSequence { + path: path(), + line: i + 1, + seq, + }); + } + let name = name.into(); + if out.names.contains(&name) { + return Err(BarcodeError::DuplicateName { path: path(), name }); + } + out.names.push(name); + out.seqs.push(seq.into_bytes()); + } + if out.names.is_empty() { + return Err(BarcodeError::Empty { path: path() }); + } + Ok(out) + } + pub fn len(&self) -> usize { self.names.len() } diff --git a/crates/escapepod-demux/src/crf/encoder.rs b/crates/escapepod-demux/src/crf/encoder.rs index 074b177..722702a 100644 --- a/crates/escapepod-demux/src/crf/encoder.rs +++ b/crates/escapepod-demux/src/crf/encoder.rs @@ -88,6 +88,78 @@ pub struct CrfMetadata { pub standardisation: Standardisation, pub signal: SignalSpec, pub crf: CrfSpec, + /// References the decoded sequence is matched against, if the bundle + /// carries them. + /// + /// A CRF emits sequence, not a class index, so it is useless without a + /// reference set — but unlike the fingerprint heads it has nowhere natural + /// to keep one, so callers were passing it separately every time. Shipping + /// it here also fixes the trimming at export rather than at each call site: + /// the model emits `target[state_len:]`, so anyone writing their own CSV + /// can silently supply full-length targets, which inflates every distance + /// and compresses the confidence margin (escapepod-models#36). A + /// caller-supplied list still overrides this. + #[serde(default)] + pub barcodes: Option>, + /// The boundary detector this model is calibrated against, if the bundle + /// pins one. + /// + /// The training window is defined relative to that detector's + /// `adapter_end`, so pairing the model with a different detector silently + /// degrades it. That coupling is a property of the model and belongs with + /// it, not in the user's shell history. + #[serde(default)] + pub boundary: Option, + /// Registry identity, for `--info` and for logging what actually ran. + #[serde(default)] + pub model: Option, + /// Published performance, carried verbatim from the model's provenance. + /// + /// Deliberately untyped: metric names differ per model kind and per + /// evaluation, and a schema here would either lag the provenance or force + /// every producer through this crate. `--info` pretty-prints whatever is + /// present rather than interpreting it. + #[serde(default)] + pub metrics: Option, +} + +/// Who this model is, for provenance in logs and `--info`. +#[derive(Debug, Clone, Deserialize)] +pub struct ModelIdent { + pub id: String, + #[serde(default)] + pub version: Option, + #[serde(default)] + pub chemistry: Option, + /// What the model does and anything a user needs to know before trusting + /// its output. + #[serde(default)] + pub notes: Option, + /// Caveats worth surfacing every time, e.g. a confounded pilot. + #[serde(default)] + pub caveats: Vec, +} + +/// One reference a decoded sequence can be matched to. +#[derive(Debug, Clone, Deserialize)] +pub struct BarcodeEntry { + pub name: String, + /// The sequence the model EMITS, not the training target. + pub sequence: String, +} + +/// The boundary detector a CRF bundle is calibrated against. +#[derive(Debug, Clone, Deserialize)] +pub struct BoundarySpec { + /// Detection method the model expects (`cnn` or `llr`). + pub method: String, + /// ONNX graph for `cnn`, relative to the sidecar. Absent means the bundle + /// names a method but does not ship the weights. + #[serde(default)] + pub onnx: Option, + /// Registry id of that model, for provenance in logs. + #[serde(default)] + pub model_id: Option, } fn default_onnx_name() -> String { diff --git a/crates/escapepod-demux/src/crf/mod.rs b/crates/escapepod-demux/src/crf/mod.rs index 7be0f73..fe27405 100644 --- a/crates/escapepod-demux/src/crf/mod.rs +++ b/crates/escapepod-demux/src/crf/mod.rs @@ -40,7 +40,7 @@ pub use lattice::{Backend, CrfDecodeError, CrfLayout, CrfScratch, decode, decode pub use barcode::{BarcodeError, BarcodeMatch, BarcodeRefs}; #[cfg(feature = "crf-decode")] -pub use encoder::{CrfEncoder, CrfError, CrfMetadata}; +pub use encoder::{BarcodeEntry, BoundarySpec, CrfEncoder, CrfError, CrfMetadata, ModelIdent}; #[cfg(feature = "crf-gpu")] pub use encoder_gpu::CrfEncoderGpu;