Stream predicted libraries and calculate isotope envelopes on demand

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()

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,15 +60,15 @@ pub enum UnavailableReason {
 type Counts = (i64, i64);
 type Resolution = Result<Counts, UnavailableReason>;
 
-/// 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,
     pub composition_rows: usize,
     pub total_rows: usize,
     pub unavailable: BTreeMap<UnavailableReason, usize>,
     #[serde(skip)]
-    envelopes: RowValues<[f32; 3]>,
+    composition_counts: Option<RowValues<(u16, u16)>>,
 }
 
 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<IonAnnot>) -> [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<std::num::NonZeroU16>, 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)

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

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),
     }
 }
@@ -392,7 +413,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 +429,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.
@@ -655,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);
@@ -674,31 +692,46 @@ 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]
     fn prediction_streams_directly_into_a_search_library() {
         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.provenance.is_object());
+        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();
         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!(

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.