Skip to main content

koth_ms/
lib.rs

1//! koth_ff: A high-performance LC-MS feature finder for timsTOF and mzML data.
2//!
3//! ## Pipeline
4//!
5//! 1. **Hill detection** – groups MS1 peaks across scans into chromatographic hills
6//! 2. **Feature finding** – groups hills into isotope envelopes with charge state assignment
7//! 3. **Scoring** – scores isotope pattern quality using the averagine model
8//!
9//! ## Quick start
10//!
11//! ```no_run
12//! use std::path::Path;
13//! use koth_ms::{run_hills, run_features, run_scoring, config::KothConfig};
14//!
15//! let config = KothConfig::default();
16//! let input = Path::new("data.mzML");
17//!
18//! let spectra = koth_ms::read_spectra(input, &config.file).unwrap();
19//! let hills = run_hills(&spectra, &config.hills, &config.file);
20//! let features = run_features(&hills, &config.features, &config.file).unwrap();
21//! let scored = run_scoring(&features, &config.scoring, &config.features, config.file.polarity);
22//! ```
23
24/// Crate version with the git commit it was built from, e.g. `0.1.0 (a1b2c3d4e5f6)`,
25/// or `0.1.0 (a1b2c3d4e5f6-dirty)` when the working tree had uncommitted changes.
26///
27/// This is what `--version` reports on both binaries, and what the benchmark's
28/// pin check reads. `KOTH_GIT_SHA` is set by `build.rs`; it is `unknown` when the
29/// crate is built from a source archive with no git directory.
30pub const VERSION: &str = concat!(env!("CARGO_PKG_VERSION"), " (", env!("KOTH_GIT_SHA"), ")");
31
32pub mod alignment;
33pub mod config;
34pub mod error;
35pub mod features;
36pub mod hills;
37pub mod input;
38pub mod io;
39pub mod lfq;
40pub mod mem;
41pub mod models;
42pub mod output;
43pub mod pipeline;
44pub mod scoring;
45pub mod stats;
46
47// In-process / streaming API surface. Re-exported at the crate root so callers
48// write `koth_ms::run_pipeline` / `koth_ms::PipelineSink` without reaching into
49// the module path.
50pub use pipeline::{
51    group_ms2_hills_by_window, run_pipeline, run_pipeline_from_spectra, run_pipeline_streaming,
52    run_pipeline_streaming_from_spectra, run_pipeline_with_ms2, FeatureFindingOutput,
53    PipelineOptions, PipelineSink,
54};
55
56// Re-export the core in-memory result types at the crate root so a downstream
57// crate (koth_tracer, uno) can name them without depending on the module layout.
58pub use models::{Feature, Hill, IsolationWindow, Polarity, ScoredFeature, Spectrum};
59
60use std::collections::HashMap;
61use std::path::Path;
62
63use config::{FeaturesConfig, FileConfig, HillsConfig, ScoringConfig};
64use error::KothError;
65use rand::seq::SliceRandom;
66
67/// Read MS1 spectra from an mzML file or Bruker .d directory.
68pub fn read_spectra(path: &Path, file: &FileConfig) -> Result<Vec<Spectrum>, KothError> {
69    io::read_spectra(path, file)
70}
71
72/// Apply the decoy construction to a full spectrum set: when `file.decoy_mode`
73/// is set, shuffle the **entire** set (Fisher–Yates over `rand::thread_rng`)
74/// before hill detection; otherwise return it untouched.
75///
76/// This is the single definition of the "shuffle the whole set, then detect"
77/// step. Every entry point — in-memory ([`run_hills`]), streaming
78/// ([`hills_streaming_inner`]), and the MS2-buffered pipeline path
79/// ([`pipeline::run_pipeline_streaming_from_spectra`]) — routes its decoy shuffle
80/// through here so the RNG source and full-set semantics live in exactly one
81/// place.
82pub(crate) fn shuffle_if_decoy(mut spectra: Vec<Spectrum>, file: &FileConfig) -> Vec<Spectrum> {
83    if file.decoy_mode {
84        log::info!(
85            "Decoy mode: shuffling {} spectra before hill detection",
86            spectra.len()
87        );
88        spectra.shuffle(&mut rand::thread_rng());
89    }
90    spectra
91}
92
93/// Stage 1: Detect chromatographic hills from MS1 spectra.
94pub fn run_hills(spectra: &[Spectrum], config: &HillsConfig, file: &FileConfig) -> Vec<Hill> {
95    let spectra = shuffle_if_decoy(spectra.to_vec(), file);
96    hills::detect_hills(&spectra, config, file)
97}
98
99/// Stage 1 (streaming): Detect hills from mzML, Bruker `.d`, or Thermo `.raw`
100/// using bounded spectrum buffers. Completed hills are retained for assembly;
101/// native readers also retain scan metadata. Decoy shuffling and optional TIC
102/// normalization still collect spectra. Gzipped mzML is decompressed as it is
103/// parsed, never held whole in memory.
104pub fn run_hills_streaming(
105    path: &Path,
106    config: &HillsConfig,
107    file: &FileConfig,
108) -> Result<Vec<Hill>, KothError> {
109    hills_streaming_inner(path, config, file)
110}
111
112/// Stage 1 (MS2): Detect MS2 hills partitioned by the precursor isolation window
113/// of each MS2 spectrum (DIA channels).
114///
115/// Supported inputs: mzML (`.mzML`/`.mzML.gz`), Bruker **diaPASEF** `.d`
116/// (`--features tdf`), and Thermo **DIA** `.raw` (`--features thermo`). On a
117/// diaPASEF `.d` the fixed isolation windows are read straight from the raw
118/// frames' quadrupole settings and each window segment is IM-collapsed into an
119/// MS2 spectrum (see [`io::bruker::read_bruker_ms2`]). On a DIA `.raw` each MS2
120/// scan becomes one MS2 spectrum stamped with its precursor isolation window
121/// (see `io::thermo::read_thermo_ms2`). ddaPASEF `.d` and DDA `.raw` are
122/// unsupported for MS2 (precursor reconstruction is out of scope): they yield an
123/// empty set plus a warning, leaving MS1 unaffected.
124pub fn run_ms2_hills_streaming(
125    path: &Path,
126    config: &HillsConfig,
127    file: &FileConfig,
128) -> Result<Vec<Hill>, KothError> {
129    let fmt = io::detect_format(path);
130
131    #[cfg(feature = "tdf")]
132    if fmt == io::InputFormat::BrukerD {
133        // Bruker diaPASEF: stream one MS2 spectrum per isolation-window segment
134        // out of the raw `.d`, then group by window exactly as the mzML path
135        // does. ddaPASEF / unknown acquisitions return an empty Vec (with a
136        // warning) from `read_bruker_ms2`, so MS1 is left untouched.
137        let spectra = io::bruker::read_bruker_ms2(path, file)?;
138        if spectra.is_empty() {
139            return Ok(Vec::new());
140        }
141        log::info!(
142            "Streaming MS2 hill detection from {} ({} Bruker diaPASEF MS2 spectra)",
143            path.display(),
144            spectra.len()
145        );
146        return Ok(hills::detect_ms2_hills_from_iter(
147            spectra.into_iter(),
148            config,
149            file,
150        ));
151    }
152
153    #[cfg(feature = "thermo")]
154    if fmt == io::InputFormat::ThermoRaw {
155        // Thermo DIA: one MS2 spectrum per MS2 scan, stamped with its precursor
156        // isolation window, then grouped by window exactly as the mzML path does.
157        // A DDA `.raw` (or one that cannot be confidently classified as DIA)
158        // returns an empty Vec (with a warning) from `read_thermo_ms2`, so MS1 is
159        // left untouched.
160        let spectra = io::thermo::read_thermo_ms2(path)?;
161        if spectra.is_empty() {
162            return Ok(Vec::new());
163        }
164        log::info!(
165            "Streaming MS2 hill detection from {} ({} Thermo DIA MS2 spectra)",
166            path.display(),
167            spectra.len()
168        );
169        return Ok(hills::detect_ms2_hills_from_iter(
170            spectra.into_iter(),
171            config,
172            file,
173        ));
174    }
175
176    if !matches!(fmt, io::InputFormat::Mzml | io::InputFormat::MzmlGz) {
177        log::warn!(
178            "MS2 hill detection supports mzML, Bruker diaPASEF .d, and Thermo DIA .raw \
179             (.raw needs --features thermo) inputs only; skipping '{}'",
180            path.display()
181        );
182        return Ok(Vec::new());
183    }
184    let (iter, status) = io::mzml::stream_mzml_ms2_checked(path)?;
185    log::info!("Streaming MS2 hill detection from {}", path.display());
186    let hills = hills::detect_ms2_hills_from_iter(iter, config, file);
187    status.check()?;
188    Ok(hills)
189}
190
191fn hills_streaming_inner(
192    path: &Path,
193    config: &HillsConfig,
194    file: &FileConfig,
195) -> Result<Vec<Hill>, KothError> {
196    let iter = io::stream_spectra(path, file)?;
197    if file.decoy_mode {
198        // Full-run shuffling deliberately retains its collecting semantics.
199        let spectra = shuffle_if_decoy(iter.collect::<Result<Vec<_>, _>>()?, file);
200        return Ok(hills::detect_hills_from_iter(
201            spectra.into_iter(),
202            config,
203            file,
204        ));
205    }
206
207    log::info!("Streaming MS1 hill detection from {}", path.display());
208    // Bruker already owns a bounded producer. Thermo and mzML are decoded on
209    // this prefetch thread, overlapping I/O with hill detection.
210    let iter: io::SpectrumStream = if io::detect_format(path) == io::InputFormat::BrukerD {
211        iter
212    } else {
213        Box::new(io::prefetch::prefetch(iter, PREFETCH_CAPACITY))
214    };
215    detect_hills_from_results(iter, config, file)
216}
217
218/// Stop on the first reader error and discard partial hills. The detector's
219/// infallible iterator API must never turn an I/O failure into a successful run.
220fn detect_hills_from_results(
221    iter: io::SpectrumStream,
222    config: &HillsConfig,
223    file: &FileConfig,
224) -> Result<Vec<Hill>, KothError> {
225    let mut error = None;
226    let spectra = iter.map_while(|item| match item {
227        Ok(spectrum) => Some(spectrum),
228        Err(e) => {
229            error = Some(e);
230            None
231        }
232    });
233    let hills = hills::detect_hills_from_iter(spectra, config, file);
234    match error {
235        Some(error) => Err(error),
236        None => Ok(hills),
237    }
238}
239
240/// Bounded look-ahead in spectra for the mzML and Thermo reader thread.
241const PREFETCH_CAPACITY: usize = 64;
242
243/// Stage 2: Detect isotope features from hills.
244pub fn run_features(
245    hills: &[Hill],
246    config: &FeaturesConfig,
247    file: &FileConfig,
248) -> Result<Vec<Feature>, KothError> {
249    if has_multiple_faims_channels(hills.iter().map(|h| h.faims_cv)) {
250        let mut groups: HashMap<Option<u32>, (Option<f32>, Vec<Hill>)> = HashMap::new();
251        for hill in hills {
252            let cv = hills::canonical_faims_cv(hill.faims_cv);
253            groups
254                .entry(cv.map(f32::to_bits))
255                .or_insert_with(|| (cv, Vec::new()))
256                .1
257                .push(hill.clone());
258        }
259        let mut groups: Vec<_> = groups.into_values().collect();
260        groups.sort_by(|a, b| hills::cmp_faims_cv(a.0, b.0));
261
262        let mut out = Vec::new();
263        for (cv, group) in groups {
264            log::info!(
265                "Detecting features for FAIMS CV {} ({} hills)",
266                cv.map_or_else(|| "missing".to_string(), |v| v.to_string()),
267                group.len()
268            );
269            out.extend(run_features_one_channel(&group, config, file)?);
270        }
271        return Ok(out);
272    }
273    run_features_one_channel(hills, config, file)
274}
275
276fn run_features_one_channel(
277    hills: &[Hill],
278    config: &FeaturesConfig,
279    file: &FileConfig,
280) -> Result<Vec<Feature>, KothError> {
281    if file.mz_recalibration {
282        if file.decoy_mode {
283            log::info!(
284                "Decoy mode active — skipping m/z recalibration (shuffled spectra would skew the surface)"
285            );
286        } else {
287            match features::learn_recal_model(hills, config, file) {
288                Some(model) => {
289                    log::info!(
290                        "m/z recalibration: learned isotope-consistency surface from {} spacing samples (global offset {:.3} ppm, σ {:.3} ppm) — applying in pass 2",
291                        model.n_samples,
292                        model.global_offset(),
293                        model.global_sigma,
294                    );
295                    return Ok(features::detect_features_with_recal(
296                        hills,
297                        config,
298                        file,
299                        Some(&model),
300                    ));
301                }
302                None => {
303                    log::warn!(
304                        "m/z recalibration enabled but too few isotope-spacing samples to build a surface — using uncorrected features"
305                    );
306                }
307            }
308        }
309    }
310    Ok(features::detect_features(hills, config, file))
311}
312
313/// Stage 3: Score isotope features using the averagine model.
314///
315/// Applies three AND-ed retention filters from `FeaturesConfig`:
316///   - `min_isotope_score` (Bhattacharyya vs averagine)
317///   - `min_cosine_score`  (mean chromatographic cosine of isotope hills)
318///   - `min_combined_score` (= isotope × cosine)
319///
320/// All thresholds at 0.0 keep every scored feature.
321pub fn run_scoring(
322    features: &[Feature],
323    scoring_cfg: &ScoringConfig,
324    features_cfg: &FeaturesConfig,
325    polarity: Polarity,
326) -> Vec<ScoredFeature> {
327    if has_multiple_faims_channels(features.iter().map(Feature::faims_cv)) {
328        let mut groups: HashMap<Option<u32>, (Option<f32>, Vec<Feature>)> = HashMap::new();
329        for feature in features {
330            let cv = hills::canonical_faims_cv(feature.faims_cv());
331            debug_assert!(
332                feature
333                    .hills
334                    .iter()
335                    .all(|h| hills::canonical_faims_cv(h.faims_cv) == cv),
336                "feature contains hills from different FAIMS CV channels"
337            );
338            groups
339                .entry(cv.map(f32::to_bits))
340                .or_insert_with(|| (cv, Vec::new()))
341                .1
342                .push(feature.clone());
343        }
344        let mut groups: Vec<_> = groups.into_values().collect();
345        groups.sort_by(|a, b| hills::cmp_faims_cv(a.0, b.0));
346
347        let mut out = Vec::new();
348        for (_cv, group) in groups {
349            out.extend(run_scoring_one_channel(
350                &group,
351                scoring_cfg,
352                features_cfg,
353                polarity,
354            ));
355        }
356        return out;
357    }
358    run_scoring_one_channel(features, scoring_cfg, features_cfg, polarity)
359}
360
361fn run_scoring_one_channel(
362    features: &[Feature],
363    scoring_cfg: &ScoringConfig,
364    features_cfg: &FeaturesConfig,
365    polarity: Polarity,
366) -> Vec<ScoredFeature> {
367    let mut scored = scoring::score_features(
368        features,
369        scoring_cfg,
370        &features_cfg.sulfur_offsets,
371        &features_cfg.isotope_model.model(),
372        polarity,
373    );
374    let any_filter = features_cfg.min_isotope_score > 0.0
375        || features_cfg.min_cosine_score > 0.0
376        || features_cfg.min_combined_score > 0.0;
377    if any_filter {
378        let before = scored.len();
379        scored.retain(|sf| {
380            sf.isotope_score >= features_cfg.min_isotope_score
381                && sf.cosine_score >= features_cfg.min_cosine_score
382                && sf.combined_score >= features_cfg.min_combined_score
383        });
384        log::info!(
385            "Retained {}/{} features after score filters (iso\u{2265}{:.2}, cos\u{2265}{:.2}, comb\u{2265}{:.2})",
386            scored.len(),
387            before,
388            features_cfg.min_isotope_score,
389            features_cfg.min_cosine_score,
390            features_cfg.min_combined_score
391        );
392    }
393    scored
394}
395
396fn has_multiple_faims_channels(values: impl Iterator<Item = Option<f32>>) -> bool {
397    let mut first: Option<Option<u32>> = None;
398    for value in values {
399        let key = hills::canonical_faims_cv(value).map(f32::to_bits);
400        match first {
401            None => first = Some(key),
402            Some(first_key) if first_key != key => return true,
403            Some(_) => {}
404        }
405    }
406    false
407}
408
409#[cfg(test)]
410mod shuffle_tests {
411    use super::*;
412    use config::FileConfig;
413
414    fn spectra(n: usize) -> Vec<Spectrum> {
415        (0..n)
416            .map(|i| Spectrum {
417                scan_index: i,
418                retention_time: i as f64 * 0.1,
419                peaks: Vec::new(),
420                ms_level: 1,
421                isolation_window: None,
422                faims_cv: None,
423            })
424            .collect()
425    }
426
427    /// Decoy off: the set is returned untouched, in the exact input order (the
428    /// non-decoy path the byte-parity tests exercise).
429    #[test]
430    fn no_decoy_is_identity() {
431        let file = FileConfig {
432            decoy_mode: false,
433            ..FileConfig::default()
434        };
435        let out = shuffle_if_decoy(spectra(64), &file);
436        let idx: Vec<usize> = out.iter().map(|s| s.scan_index).collect();
437        assert_eq!(idx, (0..64).collect::<Vec<_>>());
438    }
439
440    /// Decoy on: the FULL set is permuted with no drops or duplicates (a decoy
441    /// run must shuffle every spectrum, not a subset). Order is `thread_rng`, so
442    /// we assert the preserved multiset, which is the load-bearing invariant.
443    #[test]
444    fn decoy_permutes_full_set() {
445        let file = FileConfig {
446            decoy_mode: true,
447            ..FileConfig::default()
448        };
449        let out = shuffle_if_decoy(spectra(256), &file);
450        assert_eq!(out.len(), 256);
451        let mut idx: Vec<usize> = out.iter().map(|s| s.scan_index).collect();
452        idx.sort_unstable();
453        assert_eq!(idx, (0..256).collect::<Vec<_>>());
454    }
455}
456
457#[cfg(test)]
458mod ms1_streaming_tests {
459    use super::*;
460
461    #[test]
462    fn reader_error_discards_partial_detection() {
463        let items = vec![
464            Ok(Spectrum {
465                scan_index: 0,
466                retention_time: 1.0,
467                peaks: vec![],
468                ms_level: 1,
469                isolation_window: None,
470                faims_cv: None,
471            }),
472            Err(KothError::ThermoError("decode failure".into())),
473        ];
474        let result = detect_hills_from_results(
475            Box::new(items.into_iter()),
476            &HillsConfig::default(),
477            &FileConfig::default(),
478        );
479        assert!(matches!(result, Err(KothError::ThermoError(_))));
480    }
481}