From 494ae02e102e0f457441bb37ad1903f4cf1fba23 Mon Sep 17 00:00:00 2001 From: "J. Sebastian Paez" Date: Fri, 25 Sep 2026 16:54:02 -0700 Subject: [PATCH 1/3] Stream predicted rows and calculate isotope envelopes on demand --- .../src/data_sources/reference_library.rs | 2 +- .../src/fragment_mass/isotope_plan.rs | 71 ++++-- rust/timsseek_cli/src/build_library.rs | 5 +- rust/timsseek_cli/src/predicted_library.rs | 228 ++++++++++-------- 4 files changed, 177 insertions(+), 129 deletions(-) diff --git a/rust/timsseek/src/data_sources/reference_library.rs b/rust/timsseek/src/data_sources/reference_library.rs index 45da8286..78b3b029 100644 --- a/rust/timsseek/src/data_sources/reference_library.rs +++ b/rust/timsseek/src/data_sources/reference_library.rs @@ -298,7 +298,7 @@ impl<'a> ExpectedIntensity for RefQuery<'a> { fn expected_precursor_envelope(&self) -> SmallVec<[(i8, f32); 3]> { let tgt = self.geom.row(); let IsotopeStrategy::FromComposition { n_isotopes } = self.lib.geom.capabilities().isotopes; - let env = self.lib.plan.isotopes().envelope(tgt); + let env = self.lib.plan.isotopes().envelope(tgt, &self.lib.geom); (0..n_isotopes as usize) .map(|i| (i as i8, env[i])) .collect() diff --git a/rust/timsseek/src/fragment_mass/isotope_plan.rs b/rust/timsseek/src/fragment_mass/isotope_plan.rs index 8ba7368e..ee875e7a 100644 --- a/rust/timsseek/src/fragment_mass/isotope_plan.rs +++ b/rust/timsseek/src/fragment_mass/isotope_plan.rs @@ -32,7 +32,10 @@ use timsquery::models::{ }; use timsquery::utils::constants::PROTON_MASS; -use super::averagine::isotope_dist_from_mass; +use super::averagine::{ + averagine_cs_from_mass, + isotope_dist_from_mass, +}; use crate::isotopes::peptide_isotopes; #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] @@ -57,7 +60,7 @@ pub enum UnavailableReason { type Counts = (i64, i64); type Resolution = Result; -/// Reported with the scoring plan; cached envelopes never enter feature metadata. +/// Reported with the scoring plan; envelopes are calculated only when queried. #[derive(Debug, Clone, Serialize)] pub struct IsotopePlan { pub method: IsotopeMethod, @@ -65,7 +68,7 @@ pub struct IsotopePlan { pub total_rows: usize, pub unavailable: BTreeMap, #[serde(skip)] - envelopes: RowValues<[f32; 3]>, + composition_counts: Option>, } impl IsotopePlan { @@ -79,7 +82,7 @@ impl IsotopePlan { match result { Ok(Ok(cs)) => { composition_rows += 1; - Some(cs) + Some((cs.0 as u16, cs.1 as u16)) } Ok(Err(reason)) => { *unavailable.entry(reason).or_insert(0) += 1; @@ -100,39 +103,42 @@ impl IsotopePlan { } else { IsotopeMethod::MassEstimatedCs }; - let mut invalid_mass = None; - let envelopes = geom.map_rows(|row| match method { - IsotopeMethod::CompositionCs => { - let (c, s) = counts[row].expect("library plan promised C/S counts"); - peptide_isotopes(c as u16, s as u16) - } - IsotopeMethod::MassEstimatedCs => { - // Stored geometry, never the synthetic variant's shifted mass. + if method == IsotopeMethod::MassEstimatedCs { + // Validate the model range before scoring, without making envelopes. + for row in geom.rows() { let z = f64::from(geom.charge(row)); - let envelope = isotope_dist_from_mass(geom.precursor_mz(row) * z - z * PROTON_MASS); - if !envelope.iter().all(|v| v.is_finite()) { - invalid_mass = Some(format!( + let (c, s) = averagine_cs_from_mass(geom.precursor_mz(row) * z - z * PROTON_MASS); + if !model_finite(c, s) { + return Err(format!( "entry {}: precursor mass exceeds the C/S isotope model's numerical range", geom.output_id(row) )); } - envelope } - }); - if let Some(message) = invalid_mass { - return Err(message); } + let composition_counts = (method == IsotopeMethod::CompositionCs) + .then(|| geom.map_rows(|row| counts[row].expect("library plan promised C/S counts"))); Ok(Self { method, composition_rows, total_rows, unavailable, - envelopes, + composition_counts, }) } - pub(crate) fn envelope(&self, row: RowIdx) -> &[f32; 3] { - &self.envelopes[row] + pub(crate) fn envelope(&self, row: RowIdx, geom: &TargetColumns) -> [f32; 3] { + match &self.composition_counts { + Some(counts) => { + let (c, s) = counts[row]; + peptide_isotopes(c, s) + } + None => { + // Stored geometry, never the synthetic variant's shifted mass. + let z = f64::from(geom.charge(row)); + isotope_dist_from_mass(geom.precursor_mz(row) * z - z * PROTON_MASS) + } + } } } @@ -289,16 +295,20 @@ fn elements_cs(elements: &[(Element, Option, i32)]) -> Res fn valid_counts(cs: Counts) -> Resolution { if !(0..=i64::from(u16::MAX)).contains(&cs.0) || !(0..=i64::from(u16::MAX)).contains(&cs.1) { Err(UnavailableReason::InvalidCounts) - } else if !peptide_isotopes(cs.0 as u16, cs.1 as u16) - .iter() - .all(|v| v.is_finite()) - { + } else if !model_finite(cs.0 as u16, cs.1 as u16) { Err(UnavailableReason::ModelRange) } else { Ok(cs) } } +/// Below this rate, the zero-isotope product stays positive in f32. Check the +/// exact model only for extreme counts near its underflow boundary. +fn model_finite(c: u16, s: u16) -> bool { + let rate = c as f32 * 0.011 + s as f32 * (0.0076 + 0.044); + rate <= 80.0 || peptide_isotopes(c, s).iter().all(|v| v.is_finite()) +} + #[cfg(test)] mod tests { use super::*; @@ -479,6 +489,15 @@ mod tests { #[test] fn numerical_model_limits_are_checked() { + for c in [0, 1, 100, 1000, 5000, 7000, 9000, 12000, u16::MAX] { + for s in [0, 1, 10, 100, 500, 1000, 2000, u16::MAX] { + assert_eq!( + model_finite(c, s), + peptide_isotopes(c, s).iter().all(|v| v.is_finite()), + "C={c} S={s}" + ); + } + } assert_eq!( valid_counts((65536, 0)), Err(UnavailableReason::InvalidCounts) diff --git a/rust/timsseek_cli/src/build_library.rs b/rust/timsseek_cli/src/build_library.rs index cceb6ccf..df3a3136 100644 --- a/rust/timsseek_cli/src/build_library.rs +++ b/rust/timsseek_cli/src/build_library.rs @@ -392,7 +392,7 @@ pub(crate) fn predict_in_memory( // dropped before `stream_library` is called. let progress = BuildProgress::new(); let report = progress.callback(); - let (handle, sink) = predicted_library::sink(); + let (handle, sink) = predicted_library::sink(decoys); let stats = stream_library(&stream_options(prediction, model, &report), sink).map_err(|e| { CliError::LibraryBuild { source: format!("predicting from {}: {e:#}", prediction.fasta.display()), @@ -408,7 +408,8 @@ pub(crate) fn predict_in_memory( "{} proteins -> {} peptides -> {} precursors ({} decoys) -> {} fragments", stats.proteins, stats.peptides, stats.precursors, stats.decoys, stats.fragments, ); - handle.into_library(&stats, decoys) + info!("Finalizing predicted library"); + handle.into_library(&stats) } /// Predict a library and write it, with no network and no server. diff --git a/rust/timsseek_cli/src/predicted_library.rs b/rust/timsseek_cli/src/predicted_library.rs index e026feb7..960887ed 100644 --- a/rust/timsseek_cli/src/predicted_library.rs +++ b/rust/timsseek_cli/src/predicted_library.rs @@ -67,16 +67,13 @@ pub(crate) struct PredictedLibraryHandle { /// The [`LibrarySink`] half of [`sink`], for `stream_library` to consume. pub(crate) struct PredictedLibrarySink { shared: Arc>, - /// Every row so far, owned. A [`SpectrumRow`] borrows from the prediction it - /// came out of, so nothing kept past `spectrum` can borrow; and the rows are - /// held rather than pushed as they arrive because [`build_arena`] needs the - /// whole set to order it. - rows: Vec, + arena: Option, + normalized_axis: timsquery::RtAxis, } /// What crosses from the writer thread back to the caller. /// -/// `rows` is `Some` only once `finish` ran, which is not the same as the stream +/// `arena` is `Some` only once `finish` ran, which is not the same as the stream /// having succeeded: msspeculator's `run_library` joins its inference workers /// before its writer, so a worker that panics drops the result channel, the /// writer's loop over it ends cleanly, and `finish` publishes whatever arrived @@ -86,11 +83,11 @@ pub(crate) struct PredictedLibrarySink { #[derive(Default)] struct Handoff { provenance: Option, - rows: Option>, + arena: Option, } /// Build a sink and the handle that collects from it. -pub(crate) fn sink() -> (PredictedLibraryHandle, PredictedLibrarySink) { +pub(crate) fn sink(decoys: DecoyPolicy) -> (PredictedLibraryHandle, PredictedLibrarySink) { let shared = Arc::new(Mutex::new(Handoff::default())); ( PredictedLibraryHandle { @@ -98,7 +95,8 @@ pub(crate) fn sink() -> (PredictedLibraryHandle, PredictedLibrarySink) { }, PredictedLibrarySink { shared, - rows: Vec::new(), + arena: Some(ArenaRows::new(decoys)), + normalized_axis: timsquery::RtAxis::NormalizedIndex { scale: None }, }, ) } @@ -116,29 +114,20 @@ impl PredictedLibraryHandle { /// msspeculator increments it once per `spectrum` call, immediately before /// making it, so the two are the same count of the same thing: what the sink /// was handed, ahead of the decoy policy dropping any of it. - /// - /// The decoy policy arrives here rather than at [`sink`] because it decides - /// which rows reach the arena, and no row reaches it until every row has - /// arrived. - pub(crate) fn into_library( - self, - stats: &LibraryStats, - decoys: DecoyPolicy, - ) -> Result { + pub(crate) fn into_library(self, stats: &LibraryStats) -> Result { let handoff = std::mem::take(&mut *self.shared.lock().expect("handoff mutex poisoned")); - let Some(rows) = handoff.rows else { + let Some(arena_rows) = handoff.arena else { return Err(CliError::LibraryBuild { source: "the prediction stream handed over no rows, so it did not finish" .to_string(), }); }; - if rows.len() != stats.precursors { + if arena_rows.received != stats.precursors { return Err(CliError::LibraryBuild { source: format!( "the prediction reported {} precursors but the sink received {}, so what it \ handed over is a prefix of the library that was asked for", - stats.precursors, - rows.len(), + stats.precursors, arena_rows.received, ), }); } @@ -149,16 +138,9 @@ impl PredictedLibraryHandle { .to_string(), }); }; - let normalized_axis = timsquery::RtAxis::NormalizedIndex { - scale: provenance - .pointer("/retention/normalized/scale") - .and_then(serde_json::Value::as_str) - .map(str::to_owned), - }; - let arena = - build_arena(rows, decoys, &normalized_axis).map_err(|e| CliError::LibraryBuild { - source: format!("assembling the predicted library: {e:?}"), - })?; + let arena = arena_rows.seal().map_err(|e| CliError::LibraryBuild { + source: format!("assembling the predicted library: {e:?}"), + })?; let library = ReferenceLibrary::try_from(arena).map_err(|e| CliError::LibraryBuild { source: format!("finalizing the predicted library: {e:?}"), })?; @@ -171,6 +153,12 @@ impl PredictedLibraryHandle { impl PredictedLibrarySink { fn record_provenance(&mut self, provenance: serde_json::Value) { + self.normalized_axis = timsquery::RtAxis::NormalizedIndex { + scale: provenance + .pointer("/retention/normalized/scale") + .and_then(serde_json::Value::as_str) + .map(str::to_owned), + }; self.shared .lock() .expect("handoff mutex poisoned") @@ -195,13 +183,15 @@ impl LibrarySink for PredictedLibrarySink { } fn spectrum(&mut self, row: &SpectrumRow<'_>) -> Result<()> { - self.rows.push(PredictedRow::from_spectrum(row)?); + self.arena + .as_mut() + .expect("spectrum after finish") + .push(PredictedRow::from_spectrum(row)?, &self.normalized_axis); Ok(()) } fn finish(&mut self) -> Result<()> { - self.shared.lock().expect("handoff mutex poisoned").rows = - Some(std::mem::take(&mut self.rows)); + self.shared.lock().expect("handoff mutex poisoned").arena = self.arena.take(); Ok(()) } } @@ -303,41 +293,33 @@ fn ion_annot(peak: &Peak<'_>) -> Result { .map_err(|e| anyhow::anyhow!("packing {series}^{charge}: {e}")) } -/// Push every kept row into a sealed arena, with the intensity sidecar beside it. -fn build_arena( - rows: Vec, +/// Writer-owned columns. Each row enters the builder during `spectrum`. +struct ArenaRows { + geom: TargetColumnsBuilder, + frag_intens: Vec, + received: usize, decoys: DecoyPolicy, - normalized_axis: &timsquery::RtAxis, -) -> Result { - // A zero-row arena seals: ids and groups are both vacuously consistent, the - // parse gate has nothing to reject, and the result searches to zero results - // without ever reporting why. msspeculator only refuses an empty digest, so - // a FASTA whose peptides all fall outside the windows gets this far. - if rows.is_empty() { - return Err(TargetReadingError::SpeclibParse( - "the prediction produced no precursors; check the charge range and the length range \ - against the peptides the FASTA actually digests to" - .to_string(), - )); +} + +impl ArenaRows { + fn new(decoys: DecoyPolicy) -> Self { + Self { + geom: TargetColumnsBuilder::with_capabilities(TargetCapabilities::default_diann()), + frag_intens: Vec::new(), + received: 0, + decoys, + } } - let mut geom = TargetColumnsBuilder::with_capabilities(TargetCapabilities::default_diann()); - let mut frag_intens: Vec = Vec::new(); - - for row in &rows { - // Dropped before the push, which is what makes `Force` mean anything: - // neither the row nor its intensities reach the arena, so the seal sees - // an all-target library and resolves to the derived decoys the policy - // asked for. - if !decoys.accepts(row.is_decoy) { - continue; + + fn push(&mut self, row: PredictedRow, normalized_axis: &timsquery::RtAxis) { + self.received += 1; + // Count every prediction, including decoys excluded by the policy. + if !self.decoys.accepts(row.is_decoy) { + return; } - // Ids and group labels are one namespace, because a row that names no - // group falls back to its own id. A peptide with no pair id is therefore - // its own singleton group, which is what the seal drops the column for, - // and text on both sides so no row can mix an id shape with a group's. let group = row.group.clone().unwrap_or_else(|| row.id.clone()); - frag_intens.extend_from_slice(&row.intensities); - geom.push_row(Row { + self.frag_intens.extend_from_slice(&row.intensities); + self.geom.push_row(Row { precursor_mz: row.precursor_mz, charge: row.charge, rt: Some(timsquery::RtCoordinate { @@ -350,7 +332,6 @@ fn build_arena( }), mobility: row.mobility, frags: &row.frags, - analyte: row.analyte.as_input(), entry_name: Some(&row.id), is_decoy: row.is_decoy, @@ -359,30 +340,49 @@ fn build_arena( }); } - // The same empty arena the guard above refuses, reached the other way round: - // every row that arrived was a shipped decoy and the policy dropped all of - // them, leaving nothing to derive decoys against. - if geom.n_rows() == 0 { - return Err(TargetReadingError::SpeclibParse(format!( - "all {} predicted rows were shipped decoys, which {decoys:?} drops, so the library \ + fn seal(self) -> Result { + let Self { + geom, + frag_intens, + received, + decoys, + } = self; + // A zero-row arena seals: ids and groups are both vacuously consistent, the + // parse gate has nothing to reject, and the result searches to zero results + // without ever reporting why. msspeculator only refuses an empty digest, so + // a FASTA whose peptides all fall outside the windows gets this far. + if received == 0 { + return Err(TargetReadingError::SpeclibParse( + "the prediction produced no precursors; check the charge range and the length range \ + against the peptides the FASTA actually digests to" + .to_string(), + )); + } + // The same empty arena the guard above refuses, reached the other way round: + // every row that arrived was a shipped decoy and the policy dropped all of + // them, leaving nothing to derive decoys against. + if geom.n_rows() == 0 { + return Err(TargetReadingError::SpeclibParse(format!( + "all {} predicted rows were shipped decoys, which {decoys:?} drops, so the library \ holds no targets", - rows.len(), - ))); - } + received, + ))); + } - let geom = geom.seal(decoys)?; - if frag_intens.len() != geom.n_fragments() { - return Err(TargetReadingError::SpeclibParse(format!( - "reference-intensity sidecar ({}) must stay parallel to the fragment-label arena ({})", - frag_intens.len(), - geom.n_fragments(), - ))); - } + let geom = geom.seal(decoys)?; + if frag_intens.len() != geom.n_fragments() { + return Err(TargetReadingError::SpeclibParse(format!( + "reference-intensity sidecar ({}) must stay parallel to the fragment-label arena ({})", + frag_intens.len(), + geom.n_fragments(), + ))); + } - Ok(TargetTable::Mzpaf { - geom, - frag_intens: Some(frag_intens), - }) + Ok(TargetTable::Mzpaf { + geom, + frag_intens: Some(frag_intens), + }) + } } #[cfg(test)] @@ -500,14 +500,14 @@ mod tests { decoys: DecoyPolicy, scale: Option<&str>, ) -> PredictedLibrary { - let (handle, mut collector) = sink(); + let (handle, mut collector) = sink(decoys); collector.record_provenance(serde_json::json!({"generator":{"tool":"test"}, "retention":{"normalized":{"scale":scale}}})); for row in rows { collector.spectrum(row).expect("row converts"); } collector.finish().expect("stream finishes"); handle - .into_library(&stats(rows.len()), decoys) + .into_library(&stats(rows.len())) .expect("library seals") } @@ -641,6 +641,34 @@ mod tests { )); } + #[test] + fn spectrum_pushes_into_columns_before_finish() { + let target = Fixture::new("PEPTIDEK", "PEPTIDEK"); + let decoy = Fixture::new("PDITPEEK", "PDITPEEK"); + let (handle, mut collector) = sink(DecoyPolicy::Force); + collector.record_test_provenance(); + collector + .spectrum(&target.row(2, false, Some(1), peaks(3))) + .unwrap(); + collector + .spectrum(&decoy.row(2, true, Some(1), peaks(5))) + .unwrap(); + let arena = collector.arena.as_ref().unwrap(); + assert_eq!(arena.received, 2); + assert_eq!(arena.geom.n_rows(), 1); + assert_eq!(arena.frag_intens.len(), 3); + collector.finish().unwrap(); + assert_eq!( + handle + .into_library(&stats(2)) + .unwrap() + .library + .geometry() + .n_rows(), + 1 + ); + } + #[test] fn the_intensity_sidecar_is_as_long_as_the_fragment_label_arena() { let fixture = Fixture::new("PEPTIDEK", "PEPTIDEK"); @@ -795,12 +823,12 @@ mod tests { #[test] fn a_stream_that_never_finished_hands_over_no_library() { let fixture = Fixture::new("PEPTIDEK", "PEPTIDEK"); - let (handle, mut collector) = sink(); + let (handle, mut collector) = sink(DecoyPolicy::Never); collector .spectrum(&fixture.row(2, false, None, peaks(3))) .expect("row converts"); - let Err(error) = handle.into_library(&stats(1), DecoyPolicy::Never) else { + let Err(error) = handle.into_library(&stats(1)) else { panic!("a library that was never finished is not a library"); }; assert!( @@ -1084,7 +1112,7 @@ mod tests { #[test] fn a_stream_that_published_fewer_rows_than_it_counted_hands_over_no_library() { let fixture = Fixture::new("PEPTIDEK", "PEPTIDEK"); - let (handle, mut collector) = sink(); + let (handle, mut collector) = sink(DecoyPolicy::Never); collector.record_test_provenance(); for charge in 2..=3 { collector @@ -1093,7 +1121,7 @@ mod tests { } collector.finish().expect("stream finishes"); - let Err(error) = handle.into_library(&stats(5), DecoyPolicy::Never) else { + let Err(error) = handle.into_library(&stats(5)) else { panic!("a prefix of a library is not a library"); }; let error = format!("{error}"); @@ -1107,11 +1135,11 @@ mod tests { #[test] fn a_prediction_that_produced_no_precursors_at_all_is_an_error_and_not_an_empty_library() { - let (handle, mut collector) = sink(); + let (handle, mut collector) = sink(DecoyPolicy::Never); collector.record_test_provenance(); collector.finish().expect("stream finishes"); - let Err(error) = handle.into_library(&stats(0), DecoyPolicy::Never) else { + let Err(error) = handle.into_library(&stats(0)) else { panic!("a library with nothing in it searches to nothing and reports no reason"); }; assert!( @@ -1123,13 +1151,13 @@ mod tests { #[test] fn a_stream_that_published_rows_without_a_header_hands_over_no_library() { let fixture = Fixture::new("PEPTIDEK", "PEPTIDEK"); - let (handle, mut collector) = sink(); + let (handle, mut collector) = sink(DecoyPolicy::Never); collector .spectrum(&fixture.row(2, false, None, peaks(3))) .expect("row converts"); collector.finish().expect("stream finishes"); - let Err(error) = handle.into_library(&stats(1), DecoyPolicy::Never) else { + let Err(error) = handle.into_library(&stats(1)) else { panic!("a library whose provenance is unknown is not one this can record"); }; assert!( @@ -1141,7 +1169,7 @@ mod tests { #[test] fn a_peak_whose_series_is_not_an_mzpaf_letter_fails_the_row_rather_than_losing_the_peak() { let fixture = Fixture::new("PEPTIDEK", "PEPTIDEK"); - let (handle, mut collector) = sink(); + let (handle, mut collector) = sink(DecoyPolicy::Never); collector.record_test_provenance(); let error = collector .spectrum(&fixture.row(2, false, None, vec![peak("Q", 3, 1, 300.0, 1.0)])) @@ -1154,7 +1182,7 @@ mod tests { // Nothing partial survived the rejected row: the peak it could not label // took the whole row with it, so there is no library left to seal. collector.finish().expect("stream finishes"); - let Err(error) = handle.into_library(&stats(0), DecoyPolicy::Never) else { + let Err(error) = handle.into_library(&stats(0)) else { panic!("a rejected row left a library behind"); }; assert!( From a85ee237423ddd1fce66c2ef83ed996e7f946e79 Mon Sep 17 00:00:00 2001 From: "J. Sebastian Paez" Date: Fri, 25 Sep 2026 17:16:26 -0700 Subject: [PATCH 2/3] Track prediction sink lifecycle and keep typed provenance --- rust/timsseek_cli/src/build_library.rs | 4 +- rust/timsseek_cli/src/predicted_library.rs | 205 ++++++++++++++------- rust/timsseek_cli/src/search.rs | 6 +- 3 files changed, 146 insertions(+), 69 deletions(-) diff --git a/rust/timsseek_cli/src/build_library.rs b/rust/timsseek_cli/src/build_library.rs index df3a3136..ff072512 100644 --- a/rust/timsseek_cli/src/build_library.rs +++ b/rust/timsseek_cli/src/build_library.rs @@ -692,14 +692,14 @@ mod tests { let predicted = predict_in_memory(&prediction, DecoyPolicy::IfMissing) .expect("builtin model predicts directly into the arena"); assert!(!predicted.library.is_empty()); - assert!(predicted.provenance.is_object()); + assert!(predicted.provenance.to_json().is_object()); let dir = tempfile::tempdir().unwrap(); let library = dir.path().join("library.tsv"); std::fs::write(&library, "PrecursorMz\tProductMz\n100.0\t200.0\n").unwrap(); std::fs::write( default_sidecar(&library), - serde_json::to_vec(&predicted.provenance).unwrap(), + serde_json::to_vec(&predicted.provenance.to_json()).unwrap(), ) .unwrap(); assert_eq!( diff --git a/rust/timsseek_cli/src/predicted_library.rs b/rust/timsseek_cli/src/predicted_library.rs index 960887ed..def2e99b 100644 --- a/rust/timsseek_cli/src/predicted_library.rs +++ b/rust/timsseek_cli/src/predicted_library.rs @@ -54,9 +54,7 @@ const SECONDS_PER_MINUTE: f32 = 60.0; /// A library that was predicted rather than read, plus what produced it. pub(crate) struct PredictedLibrary { pub library: ReferenceLibrary, - /// msspeculator's provenance, as the JSON its own sidecar carries, for a - /// caller recording what a run used. - pub provenance: serde_json::Value, + pub provenance: LibraryProvenance, } /// The caller's half of [`sink`]: what the prediction produced, once it has. @@ -67,13 +65,24 @@ pub(crate) struct PredictedLibraryHandle { /// The [`LibrarySink`] half of [`sink`], for `stream_library` to consume. pub(crate) struct PredictedLibrarySink { shared: Arc>, - arena: Option, - normalized_axis: timsquery::RtAxis, + state: SinkState, +} + +enum SinkState { + AwaitingHeader { + arena: ArenaRows, + }, + Building { + arena: ArenaRows, + provenance: Box, + normalized_axis: timsquery::RtAxis, + }, + Finished, } /// What crosses from the writer thread back to the caller. /// -/// `arena` is `Some` only once `finish` ran, which is not the same as the stream +/// `completed` is `Some` only once `finish` ran, which is not the same as the stream /// having succeeded: msspeculator's `run_library` joins its inference workers /// before its writer, so a worker that panics drops the result channel, the /// writer's loop over it ends cleanly, and `finish` publishes whatever arrived @@ -82,8 +91,12 @@ pub(crate) struct PredictedLibrarySink { /// successful stream produces and checks its count against what arrived. #[derive(Default)] struct Handoff { - provenance: Option, - arena: Option, + completed: Option, +} + +struct Completed { + provenance: Box, + arena: ArenaRows, } /// Build a sink and the handle that collects from it. @@ -95,8 +108,9 @@ pub(crate) fn sink(decoys: DecoyPolicy) -> (PredictedLibraryHandle, PredictedLib }, PredictedLibrarySink { shared, - arena: Some(ArenaRows::new(decoys)), - normalized_axis: timsquery::RtAxis::NormalizedIndex { scale: None }, + state: SinkState::AwaitingHeader { + arena: ArenaRows::new(decoys), + }, }, ) } @@ -115,8 +129,17 @@ impl PredictedLibraryHandle { /// making it, so the two are the same count of the same thing: what the sink /// was handed, ahead of the decoy policy dropping any of it. pub(crate) fn into_library(self, stats: &LibraryStats) -> Result { - let handoff = std::mem::take(&mut *self.shared.lock().expect("handoff mutex poisoned")); - let Some(arena_rows) = handoff.arena else { + let completed = self + .shared + .lock() + .expect("handoff mutex poisoned") + .completed + .take(); + let Some(Completed { + provenance, + arena: arena_rows, + }) = completed + else { return Err(CliError::LibraryBuild { source: "the prediction stream handed over no rows, so it did not finish" .to_string(), @@ -131,13 +154,6 @@ impl PredictedLibraryHandle { ), }); } - let Some(provenance) = handoff.provenance else { - return Err(CliError::LibraryBuild { - source: "the prediction stream handed over rows without a header, so its \ - provenance is unknown" - .to_string(), - }); - }; let arena = arena_rows.seal().map_err(|e| CliError::LibraryBuild { source: format!("assembling the predicted library: {e:?}"), })?; @@ -146,53 +162,57 @@ impl PredictedLibraryHandle { })?; Ok(PredictedLibrary { library, - provenance, + provenance: *provenance, }) } } -impl PredictedLibrarySink { - fn record_provenance(&mut self, provenance: serde_json::Value) { - self.normalized_axis = timsquery::RtAxis::NormalizedIndex { - scale: provenance - .pointer("/retention/normalized/scale") - .and_then(serde_json::Value::as_str) - .map(str::to_owned), - }; - self.shared - .lock() - .expect("handoff mutex poisoned") - .provenance = Some(provenance); - } - - /// Stand in for [`LibrarySink::header`], which no test can call: - /// `LibraryProvenance` is `#[non_exhaustive]` and has no constructor, so only - /// msspeculator can build one. - #[cfg(test)] - fn record_test_provenance(&mut self) { - self.record_provenance(serde_json::json!({ "generator": { "tool": "test" } })); - } -} - impl LibrarySink for PredictedLibrarySink { fn header(&mut self, provenance: &LibraryProvenance) -> Result<()> { - // Kept as the JSON `to_json` builds, unflattened: this is a record of - // what produced the library, including its normalized RT scale. - self.record_provenance(provenance.to_json()); - Ok(()) + match std::mem::replace(&mut self.state, SinkState::Finished) { + SinkState::AwaitingHeader { arena } => { + self.state = SinkState::Building { + arena, + provenance: Box::new(provenance.clone()), + normalized_axis: timsquery::RtAxis::NormalizedIndex { + scale: Some(provenance.retention.normalized.scale.to_owned()), + }, + }; + Ok(()) + } + _ => bail!("prediction header called more than once or after finish"), + } } fn spectrum(&mut self, row: &SpectrumRow<'_>) -> Result<()> { - self.arena - .as_mut() - .expect("spectrum after finish") - .push(PredictedRow::from_spectrum(row)?, &self.normalized_axis); - Ok(()) + match &mut self.state { + SinkState::Building { + arena, + normalized_axis, + .. + } => { + arena.push(PredictedRow::from_spectrum(row)?, normalized_axis); + Ok(()) + } + SinkState::AwaitingHeader { .. } => bail!("prediction spectrum before header"), + SinkState::Finished => bail!("prediction spectrum after finish"), + } } fn finish(&mut self) -> Result<()> { - self.shared.lock().expect("handoff mutex poisoned").arena = self.arena.take(); - Ok(()) + match std::mem::replace(&mut self.state, SinkState::Finished) { + SinkState::Building { + arena, provenance, .. + } => { + self.shared + .lock() + .expect("handoff mutex poisoned") + .completed = Some(Completed { provenance, arena }); + Ok(()) + } + SinkState::AwaitingHeader { .. } => bail!("prediction finished before header"), + SinkState::Finished => bail!("prediction finished more than once"), + } } } @@ -388,6 +408,7 @@ impl ArenaRows { #[cfg(test)] mod tests { use std::collections::BTreeMap; + use std::sync::LazyLock; use msspeculator_core::peptide::Peptide; use msspeculator_inference::{ @@ -402,6 +423,34 @@ mod tests { use super::*; + static TEST_PROVENANCE: LazyLock = LazyLock::new(|| { + let fasta = + std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/test_data/tiny.fasta"); + let mut prediction = crate::build_library::resolve_prediction( + fasta, + &crate::config::LibraryConfig::default(), + ); + prediction.decoys = false; + prediction.max_fragments = Some(4); + crate::build_library::predict_in_memory(&prediction, DecoyPolicy::Never) + .expect("fixture provenance") + .provenance + }); + + fn test_provenance(scale: Option<&'static str>) -> LibraryProvenance { + let mut provenance = (*TEST_PROVENANCE).clone(); + if let Some(scale) = scale { + provenance.retention.normalized.scale = scale; + } + provenance + } + + impl PredictedLibrarySink { + fn record_test_provenance(&mut self) { + self.header(&test_provenance(None)).unwrap(); + } + } + /// A prediction row assembled by hand. /// /// `SpectrumRow`'s fields are public and `Residues`/`ProteinGroup` have @@ -486,8 +535,7 @@ mod tests { } } - /// Drive the sink the way `stream_library` does, minus the header, whose - /// `LibraryProvenance` cannot be built here. + /// Drive the sink with typed provenance from one real tiny prediction. /// /// Borrows the rows rather than taking them so one set can be put through /// this route and the file route both. @@ -498,10 +546,12 @@ mod tests { fn build_with_scale( rows: &[SpectrumRow<'_>], decoys: DecoyPolicy, - scale: Option<&str>, + scale: Option<&'static str>, ) -> PredictedLibrary { let (handle, mut collector) = sink(decoys); - collector.record_provenance(serde_json::json!({"generator":{"tool":"test"}, "retention":{"normalized":{"scale":scale}}})); + collector + .header(&test_provenance(scale)) + .expect("header converts"); for row in rows { collector.spectrum(row).expect("row converts"); } @@ -653,7 +703,9 @@ mod tests { collector .spectrum(&decoy.row(2, true, Some(1), peaks(5))) .unwrap(); - let arena = collector.arena.as_ref().unwrap(); + let SinkState::Building { arena, .. } = &collector.state else { + panic!("header should start building"); + }; assert_eq!(arena.received, 2); assert_eq!(arena.geom.n_rows(), 1); assert_eq!(arena.frag_intens.len(), 3); @@ -743,6 +795,7 @@ mod tests { fn declared_scale_survives_prediction_and_file_loading_only_for_normalized_rows() { let fixture = Fixture::new("PEPTIDEK", "PEPTIDEK"); for scale in [None, Some("reference anchors")] { + let declared_scale = scale.unwrap_or(TEST_PROVENANCE.retention.normalized.scale); for irt in [None, Some(37.75)] { let rows = [SpectrumRow { irt, @@ -752,7 +805,7 @@ mod tests { let predicted = build_with_scale(&rows, DecoyPolicy::Never, scale); let TargetTable::Mzpaf { geom: from_file, .. - } = via_file_table_with_scale(&rows, scale) + } = via_file_table_with_scale(&rows, Some(declared_scale)) else { panic!("expected mzpaf") }; @@ -760,7 +813,7 @@ mod tests { timsquery::RtAxis::Seconds } else { timsquery::RtAxis::NormalizedIndex { - scale: scale.map(str::to_owned), + scale: Some(declared_scale.to_owned()), } }; let sunk = predicted.library.geometry(); @@ -824,6 +877,7 @@ mod tests { fn a_stream_that_never_finished_hands_over_no_library() { let fixture = Fixture::new("PEPTIDEK", "PEPTIDEK"); let (handle, mut collector) = sink(DecoyPolicy::Never); + collector.record_test_provenance(); collector .spectrum(&fixture.row(2, false, None, peaks(3))) .expect("row converts"); @@ -877,7 +931,7 @@ mod tests { /// project's reader back off it, which is what a `build-library` followed by /// a `search` does. fn via_file_table(rows: &[SpectrumRow<'_>]) -> TargetTable { - via_file_table_with_scale(rows, None) + via_file_table_with_scale(rows, Some(TEST_PROVENANCE.retention.normalized.scale)) } fn via_file_table_with_scale(rows: &[SpectrumRow<'_>], scale: Option<&str>) -> TargetTable { @@ -1152,20 +1206,39 @@ mod tests { fn a_stream_that_published_rows_without_a_header_hands_over_no_library() { let fixture = Fixture::new("PEPTIDEK", "PEPTIDEK"); let (handle, mut collector) = sink(DecoyPolicy::Never); - collector + let error = collector .spectrum(&fixture.row(2, false, None, peaks(3))) - .expect("row converts"); - collector.finish().expect("stream finishes"); + .expect_err("spectrum before header must fail"); + assert!(format!("{error}").contains("before header")); + assert!(format!("{}", collector.finish().unwrap_err()).contains("before header")); let Err(error) = handle.into_library(&stats(1)) else { panic!("a library whose provenance is unknown is not one this can record"); }; assert!( - format!("{error}").contains("provenance"), + format!("{error}").contains("did not finish"), "unexpected error: {error}", ); } + #[test] + fn callbacks_after_finish_are_rejected() { + let fixture = Fixture::new("PEPTIDEK", "PEPTIDEK"); + let (_handle, mut collector) = sink(DecoyPolicy::Never); + collector.record_test_provenance(); + collector.finish().unwrap(); + assert!(format!("{}", collector.finish().unwrap_err()).contains("more than once")); + assert!( + format!( + "{}", + collector + .spectrum(&fixture.row(2, false, None, peaks(3))) + .unwrap_err() + ) + .contains("after finish") + ); + } + #[test] fn a_peak_whose_series_is_not_an_mzpaf_letter_fails_the_row_rather_than_losing_the_peak() { let fixture = Fixture::new("PEPTIDEK", "PEPTIDEK"); diff --git a/rust/timsseek_cli/src/search.rs b/rust/timsseek_cli/src/search.rs index b15663ae..38306a32 100644 --- a/rust/timsseek_cli/src/search.rs +++ b/rust/timsseek_cli/src/search.rs @@ -506,7 +506,11 @@ pub(crate) fn search(args: &SearchArgs) -> std::result::Result<(), errors::CliEr build_library::resolve_search_prediction(fasta.clone(), config.library.as_ref()); let predicted = build_library::predict_in_memory(&prediction, config.analysis.decoy_strategy)?; - (predicted.library, None, Some(predicted.provenance)) + ( + predicted.library, + None, + Some(predicted.provenance.to_json()), + ) } // Written and then read back rather than kept in the arena, so the // library the next run opens is the one this run searched. From ae51edce2e49c3b6f630a2328ce2d3bda9f58c79 Mon Sep 17 00:00:00 2001 From: "J. Sebastian Paez" Date: Fri, 25 Sep 2026 17:35:02 -0700 Subject: [PATCH 3/3] Cap predicted fragments by default and log resolved settings --- rust/timsseek_cli/assets/default_config.toml | 8 +-- rust/timsseek_cli/src/build_library.rs | 60 +++++++++++++++----- rust/timsseek_cli/src/config.rs | 2 +- rust/timsseek_cli/src/search.rs | 22 +++++-- 4 files changed, 68 insertions(+), 24 deletions(-) diff --git a/rust/timsseek_cli/assets/default_config.toml b/rust/timsseek_cli/assets/default_config.toml index 58665550..6c9322d6 100644 --- a/rust/timsseek_cli/assets/default_config.toml +++ b/rust/timsseek_cli/assets/default_config.toml @@ -112,9 +112,9 @@ calibration_query_rt_window_minutes = 0.083333336 # 5 seconds (5 / 60) # uri = "./results" -## How to predict a library, consulted by `timsseek build-library`. Every field -## is optional and falls through to msspeculator's own default, so the section -## can be omitted entirely; a build flag beats whatever is set here. +## How to predict a library, for FASTA searches and `timsseek build-library`. +## Every field is optional; the section can be omitted. A build flag beats a +## value here. timsseek defaults to 12 fragments per precursor and decoys on. # [library] # model = "builtin:small-v0" # missed_cleavages = 2 @@ -126,7 +126,7 @@ calibration_query_rt_window_minutes = 0.083333336 # 5 seconds (5 / 60) # variable_mods = ["M[UNIMOD:35]"] # max_variable_mods = 1 # min_intensity = 0.01 -# max_fragments = 20 +# max_fragments = 12 # decoys = true ## Acquisition and chromatography context are not settable: a build uses the ## model artifact's own defaults. Choosing a different one is a decision about diff --git a/rust/timsseek_cli/src/build_library.rs b/rust/timsseek_cli/src/build_library.rs index ff072512..be7f5fe3 100644 --- a/rust/timsseek_cli/src/build_library.rs +++ b/rust/timsseek_cli/src/build_library.rs @@ -42,11 +42,10 @@ use crate::predicted_library::{ PredictedLibrary, }; -/// msspeculator's defaults, restated here only so a partially-specified -/// `[library]` section does not have to name every field to change one. +/// Prediction defaults. Most match msspeculator; a partial `[library]` section +/// can change one setting without naming the rest. const DEFAULT_MODEL: &str = "builtin:small-v0"; -/// The one default that is this project's rather than msspeculator's, whose -/// `--decoys` is off. +/// This project's decoy default; msspeculator's `--decoys` is off. /// /// Every library built here exists to be searched, and a search needs a decoy /// for every target to put an FDR on. Predicting them costs twice the @@ -63,6 +62,7 @@ const DEFAULT_MIN_CHARGE: i64 = 2; const DEFAULT_MAX_CHARGE: i64 = 4; const DEFAULT_MAX_VARIABLE_MODS: usize = 1; const DEFAULT_MIN_INTENSITY: f64 = 0.01; +const DEFAULT_MAX_FRAGMENTS: usize = 12; const DEFAULT_FIXED_MOD: &str = "C[UNIMOD:4]"; const DEFAULT_VARIABLE_MOD: &str = "M[UNIMOD:35]"; @@ -87,6 +87,27 @@ pub struct ResolvedPrediction { pub decoys: bool, } +impl ResolvedPrediction { + /// Spell every effective prediction setting into `config_used.json` and + /// the configuration log before prediction starts. + pub(crate) fn as_library_config(&self) -> LibraryConfig { + LibraryConfig { + model: Some(self.model.clone()), + missed_cleavages: Some(self.missed_cleavages), + min_length: Some(self.min_length), + max_length: Some(self.max_length), + min_charge: Some(self.min_charge), + max_charge: Some(self.max_charge), + fixed_mods: Some(self.fixed_mods.clone()), + variable_mods: Some(self.variable_mods.clone()), + max_variable_mods: Some(self.max_variable_mods), + min_intensity: Some(self.min_intensity), + max_fragments: self.max_fragments, + decoys: Some(self.decoys), + } + } +} + /// What to predict, plus where `build-library` puts it. /// /// Two types rather than one because the output half means nothing to a caller @@ -161,7 +182,7 @@ pub fn resolve_prediction(fasta: PathBuf, library: &LibraryConfig) -> ResolvedPr .max_variable_mods .unwrap_or(DEFAULT_MAX_VARIABLE_MODS), min_intensity: library.min_intensity.unwrap_or(DEFAULT_MIN_INTENSITY), - max_fragments: library.max_fragments, + max_fragments: Some(library.max_fragments.unwrap_or(DEFAULT_MAX_FRAGMENTS)), decoys: library.decoys.unwrap_or(DEFAULT_DECOYS), } } @@ -656,14 +677,10 @@ mod tests { toml::from_str(toml).expect("a library-only configuration must parse") } - /// The literals, not the constants. `x.unwrap_or(K) == K` holds whatever `K` - /// says, so a default drifting from msspeculator's would not show; spelled - /// out, it shows in the diff. - /// - /// `decoys` is the exception and is asserted the other way round, because it - /// is deliberately not msspeculator's answer. + /// Pin the effective defaults as literals. Fragment count and decoys are + /// project choices; the other values follow msspeculator's defaults. #[test] - fn an_absent_setting_falls_through_to_msspeculators_default() { + fn an_absent_setting_uses_prediction_defaults() { let resolved = resolve_build(&args(&[]), &BuildConfig::default()).prediction; assert_eq!(resolved.model, "builtin:small-v0"); assert_eq!(resolved.missed_cleavages, 2); @@ -675,11 +692,19 @@ mod tests { assert_eq!(resolved.min_intensity, 0.01); assert_eq!(resolved.fixed_mods, ["C[UNIMOD:4]"]); assert_eq!(resolved.variable_mods, ["M[UNIMOD:35]"]); - assert_eq!(resolved.max_fragments, None); + assert_eq!(resolved.max_fragments, Some(12)); assert!( resolved.decoys, "predicted decoys beat the mass-shift ones a search would derive" ); + let recorded: LibraryConfig = + serde_json::from_value(serde_json::to_value(resolved.as_library_config()).unwrap()) + .unwrap(); + assert_eq!( + resolve_prediction(resolved.fasta.clone(), &recorded), + resolved, + "the recorded effective config must reproduce the prediction settings" + ); } #[test] @@ -687,11 +712,18 @@ mod tests { let fasta = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/test_data/tiny.fasta"); let mut prediction = resolve_prediction(fasta, &LibraryConfig::default()); prediction.decoys = false; - prediction.max_fragments = Some(4); let predicted = predict_in_memory(&prediction, DecoyPolicy::IfMissing) .expect("builtin model predicts directly into the arena"); assert!(!predicted.library.is_empty()); + assert!( + predicted.library.geometry().rows().all(|row| predicted + .library + .geometry() + .frag_range(row) + .len() + <= 12) + ); assert!(predicted.provenance.to_json().is_object()); let dir = tempfile::tempdir().unwrap(); diff --git a/rust/timsseek_cli/src/config.rs b/rust/timsseek_cli/src/config.rs index 5d242bac..0f4e4dd2 100644 --- a/rust/timsseek_cli/src/config.rs +++ b/rust/timsseek_cli/src/config.rs @@ -136,7 +136,7 @@ impl IndexingConfig { /// How to predict a library. /// -/// Every field is optional and falls through to msspeculator's own default, so +/// Every field is optional and falls through to timsseek's prediction default, so /// the section can be omitted entirely and a build flag can beat any single /// field without the others having to be spelled. This project owns the flag /// surface and none of the digestion, modification or prediction logic. diff --git a/rust/timsseek_cli/src/search.rs b/rust/timsseek_cli/src/search.rs index 38306a32..76785031 100644 --- a/rust/timsseek_cli/src/search.rs +++ b/rust/timsseek_cli/src/search.rs @@ -424,12 +424,23 @@ pub(crate) fn search(args: &SearchArgs) -> std::result::Result<(), errors::CliEr let config = load_config(args.config.as_deref())?; let (mut config, validated) = resolve_run_inputs(args, config)?; + // Record effective prediction defaults before the first log and config_used + // write, including on a run that fails before prediction completes. + let resolved_prediction = match &validated.library { + LibrarySource::Fasta(fasta) | LibrarySource::Build { fasta, .. } => Some( + build_library::resolve_search_prediction(fasta.clone(), config.library.as_ref()), + ), + LibrarySource::File(_) => None, + }; + if let Some(prediction) = &resolved_prediction { + config.library = Some(prediction.as_library_config()); + } // Held in `search()`'s scope so the instrumentation flush guard drops after // all work completes. let _tracing = init_tracing(args, &validated); - info!("Parsed configuration: {:#?}", config.clone()); + info!("Resolved configuration: {config:#?}"); alloc_track::snap!("start"); validate_inputs(&validated)?; @@ -501,11 +512,12 @@ pub(crate) fn search(args: &SearchArgs) -> std::result::Result<(), errors::CliEr } (loaded.library, loaded.tempdir, None) } - LibrarySource::Fasta(fasta) => { - let prediction = - build_library::resolve_search_prediction(fasta.clone(), config.library.as_ref()); + LibrarySource::Fasta(_) => { + let prediction = resolved_prediction + .as_ref() + .expect("FASTA source has resolved prediction settings"); let predicted = - build_library::predict_in_memory(&prediction, config.analysis.decoy_strategy)?; + build_library::predict_in_memory(prediction, config.analysis.decoy_strategy)?; ( predicted.library, None,