diff --git a/imspy_connector/Cargo.toml b/imspy_connector/Cargo.toml index a2ebf3716..c0a8fe288 100644 --- a/imspy_connector/Cargo.toml +++ b/imspy_connector/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "imspy-connector" -version = "0.4.2" +version = "0.4.4" edition = "2021" [lib] diff --git a/imspy_connector/pyproject.toml b/imspy_connector/pyproject.toml index 2f8d3e26f..77b599b40 100644 --- a/imspy_connector/pyproject.toml +++ b/imspy_connector/pyproject.toml @@ -7,7 +7,7 @@ name = "imspy_connector" dependencies = [ "opentims-bruker-bridge>=1.1.0", ] -version = "0.4.2" +version = "0.4.4" requires-python = ">=3.11" classifiers = [ "Programming Language :: Rust", diff --git a/imspy_connector/src/py_dataset.rs b/imspy_connector/src/py_dataset.rs index 00e3c0b48..c40135486 100644 --- a/imspy_connector/src/py_dataset.rs +++ b/imspy_connector/src/py_dataset.rs @@ -5,12 +5,26 @@ use rustdf::data::utility::{zstd_compress, zstd_decompress, reconstruct_compress use crate::py_tims_frame::{PyTimsFrame}; use crate::py_tims_slice::PyTimsSlice; -use numpy::{IntoPyArray, PyArray1}; +use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1}; +use pyo3::exceptions::PyRuntimeError; +use rustdf::data::handle::IndexConverter; +use rustdf::data::utility::build_compressed_frames as rust_build_compressed_frames; use pyo3::types::PyList; use pyo3::{PyResult, Python}; use rustdf::data::acquisition::AcquisitionMode; use rustdf::data::handle::TimsData; +/// Borrow each numpy array as a contiguous slice, keeping the arrays' lifetime. +fn readonly_slices<'a>(arrays: &'a [PyReadonlyArray1<'_, f64>]) -> PyResult> { + arrays + .iter() + .map(|a| { + a.as_slice() + .map_err(|_| PyRuntimeError::new_err("build_compressed_frames: arrays must be contiguous")) + }) + .collect() +} + #[pyclass] pub struct PyTimsDataset { inner: TimsDataset, @@ -119,6 +133,60 @@ impl PyTimsDataset { self.inner.loader.get_index_converter().inverse_mobility_to_scan(frame_id, &inverse_mobility_values) } + /// Run the whole per-frame TDF writer pipeline for a batch of frames in parallel: + /// m/z -> TOF, 1/K0 -> scan, `(scan, tof)` dedup, interleave and zstd. + /// + /// Returns `(num_peaks, max_intensity, summed_intensity, compressed_bytes)` per frame, in input + /// order; the statistics are taken after the dedup, which is what the `Frames` row records. + /// + /// Raises if the dataset was opened with the Bruker SDK and the file carries no calibration + /// tables the SDK-free converter can use — the SDK cannot be called from several threads at + /// once, so the caller must fall back to the per-frame path in that case. + #[pyo3(signature = (frame_ids, mz, mobility, intensity, max_scans, compression_level=0, num_threads=4))] + pub fn build_compressed_frames( + &self, + py: Python<'_>, + frame_ids: Vec, + mz: Vec>, + mobility: Vec>, + intensity: Vec>, + max_scans: u32, + compression_level: i32, + num_threads: usize, + ) -> PyResult)>> { + let sdk_free = self.inner.loader.sdk_free_converter(); + let converter: Option<&(dyn IndexConverter + Sync)> = match sdk_free.as_ref() { + Some(c) => Some(c), + None => self.inner.loader.sync_index_converter(), + }; + let Some(converter) = converter else { + return Err(PyRuntimeError::new_err( + "UNAVAILABLE: build_compressed_frames needs a converter that reproduces the \ + Bruker SDK exactly; this .d either has no usable calibration tables or uses \ + MzCalibration ModelType 2, which is only accurate to ~2 ppm. Use the per-frame \ + writer path.", + )); + }; + + let mz_slices = readonly_slices(&mz)?; + let im_slices = readonly_slices(&mobility)?; + let it_slices = readonly_slices(&intensity)?; + + let frames = py + .detach(|| { + rust_build_compressed_frames( + converter, &frame_ids, &mz_slices, &im_slices, &it_slices, + max_scans, compression_level, num_threads, + ) + }) + .map_err(PyRuntimeError::new_err)?; + + Ok(frames + .into_iter() + .map(|f| (f.num_peaks, f.max_intensity, f.summed_intensity, f.data)) + .collect()) + } + #[staticmethod] pub fn compress_bytes_zstd(bytes: Vec, compression_level: i32) -> Vec { let result = zstd_compress(&bytes, compression_level).unwrap(); diff --git a/imspy_connector/src/py_dia.rs b/imspy_connector/src/py_dia.rs index 1f517fbcd..287bea0be 100644 --- a/imspy_connector/src/py_dia.rs +++ b/imspy_connector/src/py_dia.rs @@ -2226,6 +2226,34 @@ impl PyTimsDatasetDIA { PyTimsFrame::from_inner(self.inner.sample_fragment_signal(num_frames, window_group, max_intensity, take_probability)) } + /// Batched, parallel, seeded reference-noise overlay (see rustdf `overlay_reference_noise`). + /// `window_groups[i]` is the DIA window group of frame i, or None for a precursor frame. + #[pyo3(signature = (frames, window_groups, num_precursor_frames, num_fragment_frames, max_intensity_precursor, max_intensity_fragment, take_precursor, take_fragment, seed, num_threads))] + pub fn overlay_reference_noise( + &self, + frames: Vec, + window_groups: Vec>, + num_precursor_frames: usize, + num_fragment_frames: usize, + max_intensity_precursor: f64, + max_intensity_fragment: f64, + take_precursor: f64, + take_fragment: f64, + seed: u64, + num_threads: usize, + ) -> Vec { + let inner: Vec = frames.into_iter().map(|f| f.inner).collect(); + self.inner + .overlay_reference_noise( + inner, window_groups, num_precursor_frames, num_fragment_frames, + max_intensity_precursor, max_intensity_fragment, take_precursor, take_fragment, + seed, num_threads, + ) + .into_iter() + .map(PyTimsFrame::from_inner) + .collect() + } + /// Convenience helper for Python-side planner #[getter] pub fn max_global_num_scans(&self) -> usize { diff --git a/imspy_connector/src/py_peptide.rs b/imspy_connector/src/py_peptide.rs index 93a7097e7..c9253426e 100644 --- a/imspy_connector/src/py_peptide.rs +++ b/imspy_connector/src/py_peptide.rs @@ -1,5 +1,7 @@ use std::collections::{HashMap}; use pyo3::prelude::*; +use pyo3::wrap_pyfunction; +use rayon::prelude::*; use mscore::data::peptide::{FragmentType, PeptideSequence, PeptideProductIon, PeptideProductIonSeries, PeptideProductIonSeriesCollection, PeptideIon}; @@ -345,6 +347,29 @@ impl PyPeptideProductIon { } } +/// Monoisotopic mass for a batch of UniMod sequences, in parallel. +/// +/// Constructing one `PyPeptideSequence` per peptide just to read `.mono_isotopic_mass` costs about +/// 200 us each — 250 000 peptides took 48 s in a single Python loop, which was the largest remaining +/// serial cost in the simulation. The work is per-sequence independent, so do the whole batch over a +/// rayon pool and hand back one array. +#[pyfunction] +#[pyo3(signature = (sequences, num_threads=4))] +pub fn mono_isotopic_masses(py: Python<'_>, sequences: Vec, num_threads: usize) -> PyResult> { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(num_threads.max(1)) + .build() + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("thread pool: {e}")))?; + Ok(py.detach(|| { + pool.install(|| { + sequences + .par_iter() + .map(|s| PeptideSequence::new(s.clone(), None).mono_isotopic_mass()) + .collect() + }) + })) +} + #[pymodule] pub fn py_peptide(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; @@ -352,5 +377,6 @@ pub fn py_peptide(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_function(wrap_pyfunction!(mono_isotopic_masses, m)?)?; Ok(()) } \ No newline at end of file diff --git a/imspy_connector/src/py_simulation.rs b/imspy_connector/src/py_simulation.rs index 7c76af11f..c996777a0 100644 --- a/imspy_connector/src/py_simulation.rs +++ b/imspy_connector/src/py_simulation.rs @@ -202,6 +202,12 @@ pub struct PyTimsTofSyntheticsPrecursorFrameBuilder { #[pymethods] impl PyTimsTofSyntheticsPrecursorFrameBuilder { + + /// Seed the m/z jitter so a run's noise is reproducible. Without it the jitter comes from + /// `thread_rng` and the written spectra depend on thread scheduling. + pub fn set_noise_seed(&mut self, seed: u64) { + self.inner.set_noise_seed(seed); + } #[new] pub fn new(db_path: &str) -> Self { let path = std::path::Path::new(db_path); @@ -238,6 +244,12 @@ pub struct PyTimsTofSyntheticsFrameBuilderDIA { #[pymethods] impl PyTimsTofSyntheticsFrameBuilderDIA { + + /// Seed the m/z jitter so a run's noise is reproducible. Without it the jitter comes from + /// `thread_rng` and the written spectra depend on thread scheduling. + pub fn set_noise_seed(&mut self, seed: u64) { + self.inner.set_noise_seed(seed); + } #[new] #[pyo3(signature = (db_path, with_annotations, num_threads, isotope_config=None, projection_mode=None, target_p=0.999, frame_step_size=0.001, scan_step_size=0.0001, n_steps=1000, remove_epsilon=1e-4))] pub fn new( @@ -337,6 +349,12 @@ pub struct PyTimsTofSyntheticsFrameBuilderDDA { #[pymethods] impl PyTimsTofSyntheticsFrameBuilderDDA { + + /// Seed the m/z jitter so a run's noise is reproducible. Without it the jitter comes from + /// `thread_rng` and the written spectra depend on thread scheduling. + pub fn set_noise_seed(&mut self, seed: u64) { + self.inner.set_noise_seed(seed); + } #[new] #[pyo3(signature = (db_path, with_annotations, num_threads, isotope_config=None))] pub fn new( @@ -425,6 +443,12 @@ pub struct PyTimsTofLazyFrameBuilderDIA { #[pymethods] impl PyTimsTofLazyFrameBuilderDIA { + + /// Seed the m/z jitter so a run's noise is reproducible. Without it the jitter comes from + /// `thread_rng` and the written spectra depend on thread scheduling. + pub fn set_noise_seed(&mut self, seed: u64) { + self.inner.set_noise_seed(seed); + } #[new] #[pyo3(signature = (db_path, num_threads=4, projection_mode=None, target_p=0.999, frame_step_size=0.001, scan_step_size=0.0001, n_steps=1000, remove_epsilon=1e-4))] pub fn new( diff --git a/mscore/src/algorithm/peptide.rs b/mscore/src/algorithm/peptide.rs index fb42e50a3..2a0959a18 100644 --- a/mscore/src/algorithm/peptide.rs +++ b/mscore/src/algorithm/peptide.rs @@ -11,6 +11,7 @@ use rayon::ThreadPoolBuilder; use regex::Regex; use statrs::distribution::{Binomial, Discrete}; use std::collections::HashMap; +use std::sync::LazyLock; /// calculate the monoisotopic mass of a peptide sequence /// @@ -33,10 +34,20 @@ use std::collections::HashMap; /// let mass_quantized = (mass * 1e6).round() as i32; /// assert_eq!(mass_quantized, 936418877); /// ``` +/// Tables and patterns hoisted out of the per-peptide path. +/// +/// `calculate_peptide_mono_isotopic_mass` rebuilt the 2 073-entry UniMod table, rebuilt the amino +/// acid table and recompiled the regex on **every** call. That put one mass calculation at ~190 us, +/// and 250 000 peptides at ~48 s — the largest remaining serial cost in the simulation. Built once. +static AA_MASSES: LazyLock> = LazyLock::new(amino_acid_masses); +static UNIMOD_MASSES: LazyLock> = LazyLock::new(unimod_modifications_mass_numerical); +static UNIMOD_PATTERN: LazyLock = + LazyLock::new(|| Regex::new(r"\[UNIMOD:(\d+)]").unwrap()); + pub fn calculate_peptide_mono_isotopic_mass(peptide_sequence: &PeptideSequence) -> f64 { - let amino_acid_masses = amino_acid_masses(); - let modifications_mz_numerical = unimod_modifications_mass_numerical(); - let pattern = Regex::new(r"\[UNIMOD:(\d+)]").unwrap(); + let amino_acid_masses = &*AA_MASSES; + let modifications_mz_numerical = &*UNIMOD_MASSES; + let pattern = &*UNIMOD_PATTERN; let sequence = peptide_sequence.sequence.as_str(); @@ -102,7 +113,7 @@ pub fn calculate_peptide_product_ion_mono_isotopic_mass(sequence: &str, kind: Fr return 0.0; } - let amino_acid_masses = amino_acid_masses(); + let amino_acid_masses = &*AA_MASSES; // Add up raw amino acid masses and potential modifications let mass_sequence: f64 = sequence diff --git a/mscore/src/algorithm/utility.rs b/mscore/src/algorithm/utility.rs index c55c6ebe3..6c1de64c7 100644 --- a/mscore/src/algorithm/utility.rs +++ b/mscore/src/algorithm/utility.rs @@ -135,10 +135,30 @@ pub fn emg_cdf_range(lower_limit: f64, upper_limit: f64, mu: f64, sigma: f64, la pub fn calculate_bounds_emg(mu: f64, sigma: f64, lambda: f64, step_size: f64, target: f64, lower_start: f64, upper_start: f64, n_steps: Option) -> (f64, f64) { assert!(0.0 <= target && target <= 1.0, "target must be in [0, 1]"); + // Degenerate parameters must not reach the allocation below. An infinite sigma makes + // `(upper - lower) / step_size` infinite, `as usize` saturates it to usize::MAX, and the + // `collect` then aborts the whole process with a capacity overflow rather than returning. + // Hand back non-finite bounds and let the caller decide what an unusable peak means. + if !mu.is_finite() || !sigma.is_finite() || !lambda.is_finite() || !(step_size > 0.0) { + return (f64::NAN, f64::NAN); + } + let lower_initial = mu - lower_start * sigma - 2.0; let upper_initial = mu + upper_start * sigma; - let steps = ((upper_initial - lower_initial) / step_size).round() as usize; + let span = upper_initial - lower_initial; + if !span.is_finite() || span < 0.0 { + return (f64::NAN, f64::NAN); + } + // A finite but absurd span would still allocate unboundedly: 1e8 points is already 800 MB. + // Anything past that is a nonsensical elution window, not a slow one. + const MAX_STEPS: f64 = 1e8; + let steps_f = (span / step_size).round(); + if !steps_f.is_finite() || steps_f > MAX_STEPS { + return (f64::NAN, f64::NAN); + } + + let steps = steps_f as usize; let search_space: Vec = (0..=steps).map(|i| lower_initial + i as f64 * step_size).collect(); let calc_cdf = |low: usize, high: usize| -> f64 { @@ -177,17 +197,26 @@ pub fn calculate_bounds_emg(mu: f64, sigma: f64, lambda: f64, step_size: f64, ta pub fn calculate_frame_occurrence_emg(retention_times: &[f64], rt: f64, sigma: f64, lambda_: f64, target_p: f64, step_size: f64, n_steps: Option) -> Vec { let (rt_min, rt_max) = calculate_bounds_emg(rt, sigma, lambda_, step_size, target_p, 20.0, 60.0, n_steps); - // Finding the frame closest to rt_min + // A degenerate EMG parameter upstream can make these bounds non-finite. `partial_cmp` returns + // None for NaN, so the comparators below used to unwrap a None and panic a worker thread + // mid-simulation. A peptide with no finite elution window occupies no frames; say that instead + // of aborting the run. + if !rt_min.is_finite() || !rt_max.is_finite() { + return Vec::new(); + } + + // Closest-frame search. `total_cmp` rather than `partial_cmp().unwrap()`: the inputs are finite + // by the guard above, but this can never panic regardless of what reaches it. let first_frame = retention_times.iter() .enumerate() - .min_by(|(_, &a), (_, &b)| (a - rt_min).abs().partial_cmp(&(b - rt_min).abs()).unwrap()) + .min_by(|(_, &a), (_, &b)| (a - rt_min).abs().total_cmp(&(b - rt_min).abs())) .map(|(idx, _)| idx + 1) // Rust is zero-indexed, so +1 to match Python's 1-indexing .unwrap_or(0); // Fallback in case of an empty slice // Finding the frame closest to rt_max let last_frame = retention_times.iter() .enumerate() - .min_by(|(_, &a), (_, &b)| (a - rt_max).abs().partial_cmp(&(b - rt_max).abs()).unwrap()) + .min_by(|(_, &a), (_, &b)| (a - rt_max).abs().total_cmp(&(b - rt_max).abs())) .map(|(idx, _)| idx + 1) // Same adjustment for 1-indexing .unwrap_or(0); // Fallback @@ -422,12 +451,14 @@ pub fn calculate_scan_occurrence_gaussian( .map(|(i, &val)| (val, i)) .collect(); + // `total_cmp`, not `partial_cmp().unwrap()`: the same NaN hazard as the retention-time search + // above — a non-finite bound would otherwise panic a worker thread rather than degrade. // Find the closest index to ims_lower let upper_idx = indexed_values .iter() .enumerate() .min_by(|(_, (val_a, _)), (_, (val_b, _))| { - (val_a - ims_lower).abs().partial_cmp(&(val_b - ims_lower).abs()).unwrap() + (val_a - ims_lower).abs().total_cmp(&(val_b - ims_lower).abs()) }) .map(|(idx, _)| idx) .unwrap_or(0); @@ -437,7 +468,7 @@ pub fn calculate_scan_occurrence_gaussian( .iter() .enumerate() .min_by(|(_, (val_a, _)), (_, (val_b, _))| { - (val_a - ims_upper).abs().partial_cmp(&(val_b - ims_upper).abs()).unwrap() + (val_a - ims_upper).abs().total_cmp(&(val_b - ims_upper).abs()) }) .map(|(idx, _)| idx) .unwrap_or(indexed_values.len() - 1); @@ -748,3 +779,33 @@ mod tests { assert_eq!(res_abundances.len(), 2, "Should produce 2 sets of abundances"); } } + +#[cfg(test)] +mod emg_degenerate_tests { + use super::*; + + /// A degenerate EMG parameter upstream used to make the retention-time bounds non-finite, and + /// the closest-frame search then unwrapped a `None` from `partial_cmp` and panicked a worker + /// thread mid-simulation. It must degrade to "this peptide occupies no frames" instead. + #[test] + fn non_finite_bounds_return_no_frames_instead_of_panicking() { + let rts: Vec = (0..100).map(|i| i as f64 * 0.5).collect(); + for (sigma, lambda_) in [(f64::NAN, 1.0), (1.0, f64::NAN), (0.0, 0.0), (f64::INFINITY, 1.0)] { + let frames = calculate_frame_occurrence_emg(&rts, 10.0, sigma, lambda_, 0.999, 0.0001, None); + assert!( + frames.is_empty() || frames.iter().all(|&f| f >= 0), + "sigma={sigma}, lambda={lambda_} produced {frames:?}" + ); + } + } + + #[test] + fn well_behaved_parameters_still_return_a_contiguous_frame_window() { + let rts: Vec = (0..100).map(|i| i as f64 * 0.5).collect(); + let frames = calculate_frame_occurrence_emg(&rts, 10.0, 0.5, 0.5, 0.999, 0.0001, None); + assert!(!frames.is_empty(), "a valid peak must occupy at least one frame"); + for w in frames.windows(2) { + assert_eq!(w[1], w[0] + 1, "frame window must be contiguous"); + } + } +} diff --git a/mscore/src/chemistry/utility.rs b/mscore/src/chemistry/utility.rs index 9387461d6..6e41d6bd6 100644 --- a/mscore/src/chemistry/utility.rs +++ b/mscore/src/chemistry/utility.rs @@ -1,3 +1,4 @@ +use std::sync::LazyLock; use regex::Regex; use crate::chemistry::unimod::unimod_modifications_mass; @@ -23,8 +24,11 @@ use crate::chemistry::unimod::unimod_modifications_mass; /// let tokens = unimod_sequence_to_tokens(sequence, true); /// assert_eq!(tokens, vec!["P", "E", "P", "T", "I", "D", "E[UNIMOD:1]", "H"]); /// ``` +/// Compiled once: these run per peptide, and `Regex::new` on every call dominated the cost. +static UNIMOD_BRACKETED: LazyLock = LazyLock::new(|| Regex::new(r"\[UNIMOD:\d+\]").unwrap()); + pub fn unimod_sequence_to_tokens(sequence: &str, group_modifications: bool) -> Vec { - let pattern = Regex::new(r"\[UNIMOD:\d+\]").unwrap(); + let pattern = &*UNIMOD_BRACKETED; let mut tokens = Vec::new(); let mut last_index = 0; @@ -97,12 +101,12 @@ pub fn find_unimod_patterns(input_string: &str) -> (String, Vec) { } fn remove_unimod_annotation(sequence: &str) -> String { - let pattern = Regex::new(r"\[UNIMOD:\d+]").unwrap(); + let pattern = &*UNIMOD_BRACKETED; pattern.replace_all(sequence, "").to_string() } fn extract_unimod_patterns(input_string: &str) -> Vec<(usize, usize, String)> { - let pattern = Regex::new(r"\[UNIMOD:\d+]").unwrap(); + let pattern = &*UNIMOD_BRACKETED; pattern.find_iter(input_string) .map(|mat| (mat.start(), mat.end(), mat.as_str().to_string())) .collect() diff --git a/mscore/src/data/peptide.rs b/mscore/src/data/peptide.rs index d9cfb8de5..c6f437c42 100644 --- a/mscore/src/data/peptide.rs +++ b/mscore/src/data/peptide.rs @@ -1,3 +1,4 @@ +use std::sync::LazyLock; use std::collections::{HashMap}; use bincode::{Decode, Encode}; use itertools::Itertools; @@ -262,7 +263,8 @@ impl PeptideSequence { pub fn new(raw_sequence: String, peptide_id: Option) -> Self { // constructor will parse the sequence and check if it is valid - let pattern = Regex::new(r"\[UNIMOD:(\d+)]").unwrap(); + static PAT: LazyLock = LazyLock::new(|| Regex::new(r"\[UNIMOD:(\d+)]").unwrap()); + let pattern = &*PAT; // remove the modifications from the sequence let sequence = pattern.replace_all(&raw_sequence, "").to_string(); diff --git a/mscore/src/data/spectrum.rs b/mscore/src/data/spectrum.rs index f0674543f..ea91d9d9f 100644 --- a/mscore/src/data/spectrum.rs +++ b/mscore/src/data/spectrum.rs @@ -9,7 +9,7 @@ use serde::{Serialize, Deserialize}; extern crate rand; use rand::distributions::{Uniform, Distribution}; -use rand::rngs::ThreadRng; +use rand::Rng; use statrs::distribution::Normal; /// Represents a vectorized mass spectrum. @@ -296,7 +296,14 @@ impl MzSpectrum { pub fn add_mz_noise_uniform(&self, ppm: f64, right_drag: bool) -> Self { let mut rng = rand::thread_rng(); - self.add_mz_noise(ppm, &mut rng, |rng, mz, ppm| { + self.add_mz_noise_uniform_with_rng(ppm, right_drag, &mut rng) + } + + /// Same as `add_mz_noise_uniform`, but drawing from a caller-supplied RNG so the jitter is + /// reproducible. `thread_rng` makes the written spectrum depend on which thread happened to + /// run it; a shifted m/z can cross a TOF bin boundary and change the writer's dedup. + pub fn add_mz_noise_uniform_with_rng(&self, ppm: f64, right_drag: bool, rng: &mut R) -> Self { + self.add_mz_noise(ppm, rng, |rng, mz, ppm| { let ppm_mz = match right_drag { true => mz * ppm / 1e6 / 2.0, @@ -314,16 +321,21 @@ impl MzSpectrum { pub fn add_mz_noise_normal(&self, ppm: f64) -> Self { let mut rng = rand::thread_rng(); - self.add_mz_noise(ppm, &mut rng, |rng, mz, ppm| { + self.add_mz_noise_normal_with_rng(ppm, &mut rng) + } + + /// Reproducible counterpart of `add_mz_noise_normal`; see `add_mz_noise_uniform_with_rng`. + pub fn add_mz_noise_normal_with_rng(&self, ppm: f64, rng: &mut R) -> Self { + self.add_mz_noise(ppm, rng, |rng, mz, ppm| { let ppm_mz = mz * ppm / 1e6; let dist = Normal::new(mz, ppm_mz / 3.0).unwrap(); dist.sample(rng) }) } - fn add_mz_noise(&self, ppm: f64, rng: &mut ThreadRng, noise_fn: F) -> Self + fn add_mz_noise(&self, ppm: f64, rng: &mut R, noise_fn: F) -> Self where - F: Fn(&mut ThreadRng, f64, f64) -> f64, + F: Fn(&mut R, f64, f64) -> f64, { let mz: Vec = self.mz.iter().map(|&mz_value| noise_fn(rng, mz_value, ppm)).collect(); // Clone intensity Arc (O(1)) and wrap new mz in Arc @@ -704,4 +716,54 @@ impl MzSpectrumVectorized { pub struct IndexedMzSpectrumVectorized { pub index: Vec, pub mz_vector: MzSpectrumVectorized, -} \ No newline at end of file +} +#[cfg(test)] +mod mz_noise_seeded_tests { + use super::*; + use crate::simulation::noise_rng::noise_rng; + + fn spectrum() -> MzSpectrum { + MzSpectrum::new( + vec![300.1234, 450.5678, 700.9012, 1100.3456], + vec![1000.0, 2500.0, 700.0, 150.0], + ) + } + + #[test] + fn seeded_uniform_jitter_is_reproducible() { + let s = spectrum(); + let a = s.add_mz_noise_uniform_with_rng(6.5, false, &mut noise_rng(41, &[7])); + let b = s.add_mz_noise_uniform_with_rng(6.5, false, &mut noise_rng(41, &[7])); + assert_eq!(a.mz, b.mz); + assert_eq!(a.intensity, b.intensity); + } + + #[test] + fn seeded_normal_jitter_is_reproducible() { + let s = spectrum(); + let a = s.add_mz_noise_normal_with_rng(6.5, &mut noise_rng(41, &[7])); + let b = s.add_mz_noise_normal_with_rng(6.5, &mut noise_rng(41, &[7])); + assert_eq!(a.mz, b.mz); + } + + #[test] + fn a_different_frame_key_gives_different_jitter() { + let s = spectrum(); + let a = s.add_mz_noise_normal_with_rng(6.5, &mut noise_rng(41, &[7])); + let b = s.add_mz_noise_normal_with_rng(6.5, &mut noise_rng(41, &[8])); + assert_ne!(a.mz, b.mz, "frames must not share a noise stream"); + } + + #[test] + fn the_jitter_stays_within_the_requested_ppm() { + let s = spectrum(); + let ppm = 6.5; + let out = s.add_mz_noise_uniform_with_rng(ppm, false, &mut noise_rng(41, &[1])); + assert_eq!(out.mz.len(), s.mz.len()); + for (before, after) in s.mz.iter().zip(out.mz.iter()) { + // to_resolution(6) rounds to 1e-6, so allow that on top of the ppm window. + let window = before * ppm / 1e6 + 1e-6; + assert!((after - before).abs() <= window, "{before} -> {after} exceeds {ppm} ppm"); + } + } +} diff --git a/mscore/src/simulation/annotation.rs b/mscore/src/simulation/annotation.rs index f70f05b9b..4a0fe3db8 100644 --- a/mscore/src/simulation/annotation.rs +++ b/mscore/src/simulation/annotation.rs @@ -2,7 +2,7 @@ use std::collections::{BTreeMap, HashMap}; use std::fmt::Display; use itertools::{izip, multizip}; use rand::distributions::{Uniform, Distribution}; -use rand::rngs::ThreadRng; +use rand::Rng; use statrs::distribution::Normal; use crate::data::spectrum::{MsType, ToResolution, Vectorized}; use crate::data::peptide::FragmentType; @@ -148,7 +148,12 @@ impl MzSpectrumAnnotated { pub fn add_mz_noise_uniform(&self, ppm: f64, right_drag: bool) -> Self { let mut rng = rand::thread_rng(); - self.add_mz_noise(ppm, &mut rng, |rng, mz, ppm| { + self.add_mz_noise_uniform_with_rng(ppm, right_drag, &mut rng) + } + + /// Reproducible counterpart of `add_mz_noise_uniform`; see `MzSpectrum`. + pub fn add_mz_noise_uniform_with_rng(&self, ppm: f64, right_drag: bool, rng: &mut R) -> Self { + self.add_mz_noise(ppm, rng, |rng, mz, ppm| { let ppm_mz = match right_drag { true => mz * ppm / 1e6 / 2.0, @@ -166,16 +171,21 @@ impl MzSpectrumAnnotated { pub fn add_mz_noise_normal(&self, ppm: f64) -> Self { let mut rng = rand::thread_rng(); - self.add_mz_noise(ppm, &mut rng, |rng, mz, ppm| { + self.add_mz_noise_normal_with_rng(ppm, &mut rng) + } + + /// Reproducible counterpart of `add_mz_noise_normal`; see `MzSpectrum`. + pub fn add_mz_noise_normal_with_rng(&self, ppm: f64, rng: &mut R) -> Self { + self.add_mz_noise(ppm, rng, |rng, mz, ppm| { let ppm_mz = mz * ppm / 1e6; let dist = Normal::new(mz, ppm_mz / 3.0).unwrap(); // 3 sigma ? good enough? dist.sample(rng) }) } - fn add_mz_noise(&self, ppm: f64, rng: &mut ThreadRng, noise_fn: F) -> Self + fn add_mz_noise(&self, ppm: f64, rng: &mut R, noise_fn: F) -> Self where - F: Fn(&mut ThreadRng, f64, f64) -> f64, + F: Fn(&mut R, f64, f64) -> f64, { let mz: Vec = self.mz.iter().map(|&mz_value| noise_fn(rng, mz_value, ppm)).collect(); let spectrum = MzSpectrumAnnotated { mz, intensity: self.intensity.clone(), annotations: self.annotations.clone()}; diff --git a/mscore/src/simulation/mod.rs b/mscore/src/simulation/mod.rs index 45e9e3569..3ba2e0b49 100644 --- a/mscore/src/simulation/mod.rs +++ b/mscore/src/simulation/mod.rs @@ -1,2 +1,3 @@ pub mod simulation; -pub mod annotation; \ No newline at end of file +pub mod annotation; +pub mod noise_rng; diff --git a/mscore/src/simulation/noise_rng.rs b/mscore/src/simulation/noise_rng.rs new file mode 100644 index 000000000..082da5889 --- /dev/null +++ b/mscore/src/simulation/noise_rng.rs @@ -0,0 +1,84 @@ +//! Deterministic RNG streams for the simulation's noise, keyed by identity rather than by +//! scheduling order. +//! +//! The m/z jitter and the reference-noise overlay used `rand::thread_rng`, so which thread picked +//! up a spectrum decided its noise. That makes a run irreproducible, and it is not a cosmetic +//! problem: a jittered m/z can cross a TOF bin boundary, which changes the `(scan, tof)` dedup and +//! so the peaks written to the file. +//! +//! Deriving the stream from a master seed plus the stable identifiers already in scope — frame id, +//! peptide id, ion index, scan — makes the draw independent of thread count and scheduling. + +use rand::rngs::StdRng; +use rand::SeedableRng; + +/// Golden-ratio odd constant, the usual choice for cheap integer mixing. +const MIX: u64 = 0x9E37_79B9_7F4A_7C15; + +/// Mix a master seed and a list of identifiers into one 64-bit stream key. +/// +/// Order matters: `[frame, peptide]` and `[peptide, frame]` give different streams, which is what +/// you want, since the components are not interchangeable. +#[inline] +pub fn stream_key(seed: u64, keys: &[u64]) -> u64 { + let mut h = seed ^ MIX; + for (i, &k) in keys.iter().enumerate() { + // SplitMix64-style finaliser per component, position-dependent so permutations differ. + let mut x = k.wrapping_add(MIX.wrapping_mul(i as u64 + 1)); + x = (x ^ (x >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + x = (x ^ (x >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + x ^= x >> 31; + h = (h ^ x).wrapping_mul(MIX); + } + h +} + +/// A reproducible RNG for one unit of noise, keyed by `seed` and `keys`. +#[inline] +pub fn noise_rng(seed: u64, keys: &[u64]) -> StdRng { + StdRng::seed_from_u64(stream_key(seed, keys)) +} + +#[cfg(test)] +mod tests { + use super::*; + use rand::Rng; + + fn draw(seed: u64, keys: &[u64]) -> u64 { + noise_rng(seed, keys).gen() + } + + #[test] + fn same_inputs_give_the_same_stream() { + assert_eq!(draw(41, &[7, 3, 1]), draw(41, &[7, 3, 1])); + } + + #[test] + fn the_seed_matters() { + assert_ne!(draw(41, &[7, 3, 1]), draw(42, &[7, 3, 1])); + } + + #[test] + fn each_key_component_matters() { + let base = draw(41, &[7, 3, 1]); + assert_ne!(base, draw(41, &[8, 3, 1])); + assert_ne!(base, draw(41, &[7, 4, 1])); + assert_ne!(base, draw(41, &[7, 3, 2])); + } + + #[test] + fn key_order_matters() { + assert_ne!(draw(41, &[1, 2]), draw(41, &[2, 1])); + } + + #[test] + fn neighbouring_keys_do_not_collide() { + // Frame/scan ids are dense and small; adjacent ones must not share a stream. + let mut seen = std::collections::HashSet::new(); + for frame in 0..200u64 { + for scan in 0..50u64 { + assert!(seen.insert(stream_key(41, &[frame, scan])), "collision at {frame}/{scan}"); + } + } + } +} diff --git a/mscore/src/timstof/frame.rs b/mscore/src/timstof/frame.rs index 31ae18779..317b1870c 100644 --- a/mscore/src/timstof/frame.rs +++ b/mscore/src/timstof/frame.rs @@ -582,9 +582,15 @@ impl TimsFrame { } pub fn generate_random_sample(&self, take_probability: f64) -> TimsFrame { + let mut rng = rand::thread_rng(); + self.generate_random_sample_with_rng(take_probability, &mut rng) + } + + /// Same as `generate_random_sample`, but drawing from a caller-supplied RNG so that the + /// selection is reproducible (used by the seeded, parallel reference-noise overlay). + pub fn generate_random_sample_with_rng(&self, take_probability: f64, rng: &mut R) -> TimsFrame { assert!(take_probability >= 0.0 && take_probability <= 1.0); - let mut rng = rand::thread_rng(); let mut scan = Vec::new(); let mut mobility = Vec::new(); let mut tof = Vec::new(); diff --git a/packages/imspy-core/pyproject.toml b/packages/imspy-core/pyproject.toml index 679faa7fe..f59fb98b6 100644 --- a/packages/imspy-core/pyproject.toml +++ b/packages/imspy-core/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "imspy-core" -version = "0.4.2" +version = "0.4.4" description = "Core data structures and utilities for processing timsTOF ion mobility spectrometry data." authors = [ { name = "theGreatHerrLebert", email = "davidteschner@googlemail.com" } @@ -20,7 +20,7 @@ dependencies = [ "pyarrow>=13.0", "mendeleev>=0.7.0", "toml>=0.10.2", - "imspy-connector>=0.4.2", + "imspy-connector>=0.4.4", "opentims-bruker-bridge>=1.1.0", ] diff --git a/packages/imspy-core/src/imspy_core/timstof/data.py b/packages/imspy-core/src/imspy_core/timstof/data.py index 9f2170007..240ddf662 100644 --- a/packages/imspy-core/src/imspy_core/timstof/data.py +++ b/packages/imspy-core/src/imspy_core/timstof/data.py @@ -329,6 +329,42 @@ def tof_to_mz(self, frame_id: int, tof_values: NDArray[np.int32]) -> NDArray[np. """ return self.__dataset.tof_to_mz(frame_id, tof_values) + def build_compressed_frames( + self, + frame_ids: List[int], + mz: List[NDArray[np.float64]], + mobility: List[NDArray[np.float64]], + intensity: List[NDArray[np.float64]], + max_scans: int, + compression_level: int = 0, + num_threads: int = 4, + ) -> List[tuple]: + """Run the per-frame TDF writer pipeline for a batch of frames in parallel. + + Converts m/z to TOF and 1/K0 to scan, merges duplicate (scan, tof) cells, interleaves into + the Bruker layout and zstd-compresses — all in Rust over a thread pool, so the caller only + has to append the results in order. + + Args: + frame_ids: Frame id to take the calibration from, per frame. + mz / mobility / intensity: Per-frame peak arrays, contiguous float64. + max_scans: Scan count of the output layout. + compression_level: zstd level; 0 is libzstd's default, matching the per-frame writer. + num_threads: Rayon threads. + + Returns: + One ``(num_peaks, max_intensity, summed_intensity, compressed_bytes)`` per input frame, + in input order. The statistics are taken after the dedup, as the Frames table records them. + + Raises: + RuntimeError: if the dataset uses the Bruker SDK and carries no calibration tables for + the SDK-free converter, since the SDK cannot be called from several threads at once. + """ + return self.__dataset.build_compressed_frames( + [int(f) for f in frame_ids], mz, mobility, intensity, + int(max_scans), int(compression_level), int(num_threads), + ) + def mz_to_tof(self, frame_id: int, mz_values: NDArray[np.float64]) -> NDArray[np.int32]: """Convert m/z values to TOF values. diff --git a/packages/imspy-core/src/imspy_core/timstof/dia.py b/packages/imspy-core/src/imspy_core/timstof/dia.py index 40b79eda9..1b52d0b46 100644 --- a/packages/imspy-core/src/imspy_core/timstof/dia.py +++ b/packages/imspy-core/src/imspy_core/timstof/dia.py @@ -1,5 +1,5 @@ import sqlite3 -from typing import List +from typing import List, Optional from imspy_core.core.base import RustWrapperObject from imspy_core.timstof.data import TimsDataset @@ -72,6 +72,48 @@ def sample_fragment_signal(self, num_frames: int, window_group: int, max_intensi return TimsFrame.from_py_ptr(self.__dataset.sample_fragment_signal(num_frames, window_group, max_intensity, take_probability)) + def overlay_reference_noise( + self, + frames: List[TimsFrame], + window_groups: List[Optional[int]], + num_precursor_frames: int = 5, + num_fragment_frames: int = 5, + max_intensity_precursor: float = 30.0, + max_intensity_fragment: float = 30.0, + take_precursor: float = 0.2, + take_fragment: float = 0.2, + seed: int = 0, + num_threads: int = 4, + ) -> List[TimsFrame]: + """Add sampled reference-data noise to a batch of frames, in parallel and reproducibly. + + Equivalent to calling ``sample_precursor_signal`` / ``sample_fragment_signal`` per frame and + adding the result with ``frame + noise``, but done in Rust over a thread pool with a + per-frame RNG stream derived from ``seed`` and the frame id. + + Args: + frames: Simulated frames. + window_groups: Per frame, the DIA window group (fragment frame) or None (precursor frame). + num_precursor_frames / num_fragment_frames: Reference frames sampled per output frame. + max_intensity_precursor / max_intensity_fragment: Keep reference peaks up to this intensity. + take_precursor / take_fragment: Fraction of reference peaks to keep. + seed: Master seed; identical inputs and seed give identical output. + num_threads: Rayon threads. + + Returns: + List[TimsFrame]: frames with noise added, same order as the input. + """ + assert len(frames) == len(window_groups), "frames and window_groups must have the same length" + out = self.__dataset.overlay_reference_noise( + [f.get_py_ptr() for f in frames], + [None if g is None else int(g) for g in window_groups], + int(num_precursor_frames), int(num_fragment_frames), + float(max_intensity_precursor), float(max_intensity_fragment), + float(take_precursor), float(take_fragment), + int(seed) & 0xFFFFFFFFFFFFFFFF, int(num_threads), + ) + return [TimsFrame.from_py_ptr(f) for f in out] + def read_compressed_data_full(self) -> List[bytes]: """Read compressed data. diff --git a/packages/imspy-simulation/pyproject.toml b/packages/imspy-simulation/pyproject.toml index 777b65783..79c1359f8 100644 --- a/packages/imspy-simulation/pyproject.toml +++ b/packages/imspy-simulation/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "imspy-simulation" -version = "0.4.2" +version = "0.4.5" description = "TimsTOF data simulation tools for proteomics." authors = [ { name = "theGreatHerrLebert", email = "davidteschner@googlemail.com" } @@ -10,8 +10,8 @@ license = "MIT" requires-python = ">=3.11,<3.14" dependencies = [ - "imspy-connector>=0.4.2", - "imspy-core>=0.4.2", + "imspy-connector>=0.4.4", + "imspy-core>=0.4.4", "imspy-predictors[koina]>=0.5.2", # Include Koina support for intensity prediction "pandas>=2.0", "numpy>=1.24", diff --git a/packages/imspy-simulation/src/imspy_simulation/builders/dda.py b/packages/imspy-simulation/src/imspy_simulation/builders/dda.py index 326e73d0a..e5ead06c1 100644 --- a/packages/imspy-simulation/src/imspy_simulation/builders/dda.py +++ b/packages/imspy-simulation/src/imspy_simulation/builders/dda.py @@ -110,6 +110,15 @@ def __init__( ) self._with_annotations = with_annotations + def set_noise_seed(self, seed: int) -> None: + """Seed the m/z jitter so this builder's noise is reproducible. + + Without it the jitter is drawn from a thread-local RNG, so which thread builds a frame + decides its noise — and a shifted m/z can cross a TOF bin boundary, changing the peaks + that get written. + """ + self._py_ptr.set_noise_seed(int(seed) & 0xFFFFFFFFFFFFFFFF) + @property def _ptr(self): """Get the underlying PyO3 pointer.""" diff --git a/packages/imspy-simulation/src/imspy_simulation/builders/dia.py b/packages/imspy-simulation/src/imspy_simulation/builders/dia.py index 36643f406..533fd5416 100644 --- a/packages/imspy-simulation/src/imspy_simulation/builders/dia.py +++ b/packages/imspy-simulation/src/imspy_simulation/builders/dia.py @@ -131,6 +131,15 @@ def __init__( ) self._with_annotations = with_annotations + def set_noise_seed(self, seed: int) -> None: + """Seed the m/z jitter so this builder's noise is reproducible. + + Without it the jitter is drawn from a thread-local RNG, so which thread builds a frame + decides its noise — and a shifted m/z can cross a TOF bin boundary, changing the peaks + that get written. + """ + self._py_ptr.set_noise_seed(int(seed) & 0xFFFFFFFFFFFFFFFF) + @property def _ptr(self): """Get the underlying PyO3 pointer.""" diff --git a/packages/imspy-simulation/src/imspy_simulation/builders/factory.py b/packages/imspy-simulation/src/imspy_simulation/builders/factory.py index 81761eeb0..3beb0d43e 100644 --- a/packages/imspy-simulation/src/imspy_simulation/builders/factory.py +++ b/packages/imspy-simulation/src/imspy_simulation/builders/factory.py @@ -6,7 +6,7 @@ from enum import Enum, auto from pathlib import Path -from typing import Union +from typing import Union, Optional import logging from imspy_simulation.core.protocols import FrameBuilder @@ -34,7 +34,7 @@ class LoadingStrategy(Enum): LAZY = auto() -def create_frame_builder( +def _create_frame_builder_impl( db_path: Union[str, Path], acquisition_mode: AcquisitionMode, loading_strategy: LoadingStrategy = LoadingStrategy.STANDARD, @@ -159,3 +159,19 @@ def create_frame_builder( ) raise ValueError(f"Unknown acquisition mode: {acquisition_mode}") + + +def create_frame_builder(*args, noise_seed: Optional[int] = None, **kwargs) -> FrameBuilder: + """Create a frame builder, optionally seeding its m/z jitter. + + See :func:`_create_frame_builder_impl` for the full argument list. + + Args: + noise_seed: Master seed for the m/z jitter. ``None`` keeps the legacy thread-local RNG, + where which thread builds a frame decides its noise and a run cannot be reproduced. + TimSim passes the run's ``sample_seed``. + """ + builder = _create_frame_builder_impl(*args, **kwargs) + if noise_seed is not None: + builder.set_noise_seed(int(noise_seed)) + return builder diff --git a/packages/imspy-simulation/src/imspy_simulation/tdf.py b/packages/imspy-simulation/src/imspy_simulation/tdf.py index e9cb518b3..da61db8c3 100644 --- a/packages/imspy-simulation/src/imspy_simulation/tdf.py +++ b/packages/imspy-simulation/src/imspy_simulation/tdf.py @@ -6,6 +6,7 @@ import numpy as np from pathlib import Path +from typing import List from numpy._typing import NDArray @@ -133,6 +134,26 @@ def _locked_database_error(db_path, exc: Exception) -> RuntimeError: ) +def dedup_scan_tof(scan: NDArray, tof: NDArray, intensity: NDArray) -> tuple[NDArray, NDArray, NDArray]: + """Merge peaks that fall on the same (scan, tof) cell and return them sorted by scan, then tof. + + m/z -> TOF is not injective, so several simulated peaks can land on one TOF index; their + intensities are summed. The result is byte-identical to the previous + ``np.unique(np.stack((scan, tof)), axis=0)`` + ``np.lexsort`` implementation, but ~10-15x faster: + the pair is packed into one uint64 key (scan in the high 32 bits) so a single 1-D sort orders by + scan first and tof second, which is exactly the lexicographic order the TDF encoder needs. + """ + scan = np.asarray(scan, dtype=np.uint32) + tof = np.asarray(tof, dtype=np.uint32) + intensity = np.asarray(intensity) + key = (scan.astype(np.uint64) << np.uint64(32)) | tof.astype(np.uint64) + unique_key, inverse = np.unique(key, return_inverse=True) + summed = np.bincount(inverse.ravel(), weights=intensity) + out_scan = (unique_key >> np.uint64(32)).astype(np.uint32) + out_tof = (unique_key & np.uint64(0xFFFFFFFF)).astype(np.uint32) + return out_scan, out_tof, summed.astype(np.uint32) + + class TDFWriter: def __init__(self, helper_handle: TimsDataset, path: str = "./", exp_name: str = "RAW.d", offset_bytes: int = 64, verbose: bool=False, use_rust_compression: bool=False, expect_existing: bool = False) -> None: @@ -161,6 +182,9 @@ def __init__(self, helper_handle: TimsDataset, path: str = "./", exp_name: str = ) not in ("0", "", "false", "False") self.__conn_native = None + # Binary output handle; opened lazily on first write_frame() and kept open for the + # whole run instead of open/append/close per frame (34k+ syscalls, costly on NFS). + self._bin_fh = None self._setup_connections() def _setup_connections(self) -> None: @@ -328,6 +352,24 @@ def build_frame_meta_row( scan_mode: int, frame_start_pos: int, only_frame_one: bool = False + ): + """Thin wrapper: derive the Frames statistics from an intensity array, then delegate.""" + return self._frame_meta_row( + frame, scan_mode, frame_start_pos, only_frame_one, + num_peaks=len(intensity), + max_intensity=int(np.max(intensity)) if len(intensity) > 0 else 0, + summed_intensity=int(np.sum(intensity)) if len(intensity) > 0 else 0, + ) + + def _frame_meta_row( + self, + frame: TimsFrame, + scan_mode: int, + frame_start_pos: int, + only_frame_one: bool, + num_peaks: int, + max_intensity: int, + summed_intensity: int, ): """Build a row for the frame meta data table from a TimsFrame object. Arguments: @@ -355,10 +397,10 @@ def build_frame_meta_row( r.ScanMode = scan_mode r.MsMsType = frame.ms_type r.TimsId = frame_start_pos - r.MaxIntensity = int(np.max(intensity)) if len(intensity) > 0 else 0 - r.SummedIntensities = int(np.sum(intensity)) if len(intensity) > 0 else 0 + r.MaxIntensity = int(max_intensity) + r.SummedIntensities = int(summed_intensity) r.NumScans = self.helper_handle.num_scans - r.NumPeaks = len(intensity) + r.NumPeaks = int(num_peaks) return r @@ -387,27 +429,8 @@ def compress_frame(self, frame: TimsFrame, only_frame_one: bool = False) -> (NDA scan = self.inv_mobility_to_scan(i, frame.mobility).astype(np.uint32) intensity = frame.intensity.astype(np.uint32) - # Since, mz -> tof is not bijective, we need to check for duplicates - # stack scan and tof to form a 2D array for unique grouping - scan_tof = np.stack((scan, tof), axis=1) - - # get unique (scan, tof) pairs and their inverse indices - unique_pairs, inverse_indices = np.unique(scan_tof, axis=0, return_inverse=True) - - # sum intensities for each unique (scan, tof) pair - summed_intensity = np.bincount(inverse_indices, weights=intensity) - - # now split back scan and tof - unique_scan = unique_pairs[:, 0] - unique_tof = unique_pairs[:, 1] - - # sort first by scan, then by tof - sort_idx = np.lexsort((unique_tof, unique_scan)) - - # final sorted arrays - scan = unique_scan[sort_idx] - tof = unique_tof[sort_idx] - intensity = summed_intensity[sort_idx].astype(np.uint32) + # Since mz -> tof is not bijective, merge duplicate (scan, tof) cells and sort by scan, tof. + scan, tof, intensity = dedup_scan_tof(scan, tof, intensity) # get the real data as interleaved bytes (Rust encoder or NumPy/Numba) if self.use_rust_compression: @@ -440,13 +463,88 @@ def write_frame(self, frame: TimsFrame, scan_mode: int, only_frame_one: bool = F only_frame_one )) - with open(self.binary_file, "ab") as bin_file: - bin_file.write( - (len(compressed_data) + 8).to_bytes(4, "little", signed=False) - ) - bin_file.write(int(self.helper_handle.num_scans).to_bytes(4, "little", signed=False)) - bin_file.write(compressed_data) - self.position = bin_file.tell() + bin_file = self._binary_handle() + bin_file.write( + (len(compressed_data) + 8).to_bytes(4, "little", signed=False) + ) + bin_file.write(int(self.helper_handle.num_scans).to_bytes(4, "little", signed=False)) + bin_file.write(compressed_data) + self.position = bin_file.tell() + + def _conversion_frame_id(self, frame: TimsFrame, only_frame_one: bool) -> int: + """Frame id whose calibration is used for this frame, clamped to the reference's range.""" + try: + max_index = self.helper_handle.meta_data.Id.max() + except AttributeError: + max_index = self.helper_handle.meta_data.frame_id.max() + if only_frame_one: + return 1 + return int(max_index) if frame.frame_id > max_index else int(frame.frame_id) + + def write_frames(self, frames: List[TimsFrame], scan_mode: int, only_frame_one: bool = False, + num_threads: int = 4) -> None: + """Batched counterpart of :meth:`write_frame`. + + Convert, dedup, interleave and compress every frame in one parallel Rust call, then append + the payloads in order — only the append and the running byte offset have to stay sequential. + Previously this ran per frame in Python, where the m/z -> TOF and 1/K0 -> scan conversions + went through the Bruker SDK and so could not be parallelised; measured at 56 % of writer time. + + Falls back to the per-frame path if the parallel builder is unavailable (old connector, or a + reference that carries no calibration tables for the SDK-free converter). + """ + if len(frames) == 0: + return + + builder = getattr(self.helper_handle, "build_compressed_frames", None) + if builder is not None: + conv_ids = [self._conversion_frame_id(f, only_frame_one) for f in frames] + mz = [np.ascontiguousarray(f.mz, dtype=np.float64) for f in frames] + mobility = [np.ascontiguousarray(f.mobility, dtype=np.float64) for f in frames] + intensity = [np.ascontiguousarray(f.intensity, dtype=np.float64) for f in frames] + try: + built = builder(conv_ids, mz, mobility, intensity, + int(self.helper_handle.num_scans), 0, int(num_threads)) + except RuntimeError as e: + # Only the explicit capability signal falls back. Anything else — a ragged array, + # a compression or thread-pool failure — is a real error, and silently writing the + # file through the other path would hide it behind a subtly different output. + if "UNAVAILABLE" not in str(e): + raise + warnings.warn( + f"batched frame writer unavailable, using the per-frame path: {e}", + RuntimeWarning, + ) + built = None + if built is not None: + bin_file = self._binary_handle() + num_scans_bytes = int(self.helper_handle.num_scans).to_bytes(4, "little", signed=False) + for frame, (num_peaks, max_i, sum_i, data) in zip(frames, built): + self.frame_meta_data.append( + self._frame_meta_row(frame, scan_mode, self.position, only_frame_one, + num_peaks, max_i, sum_i)) + bin_file.write((len(data) + 8).to_bytes(4, "little", signed=False)) + bin_file.write(num_scans_bytes) + bin_file.write(data) + self.position = bin_file.tell() + return + + for frame in frames: + self.write_frame(frame, scan_mode, only_frame_one) + + def _binary_handle(self): + """Return the persistent append handle for analysis.tdf_bin, opening it on first use.""" + if self._bin_fh is None or self._bin_fh.closed: + self._bin_fh = open(self.binary_file, "ab") + return self._bin_fh + + def close_binary(self) -> None: + """Flush and close the analysis.tdf_bin handle (idempotent). Called once all frames are + written; a later write_frame() re-opens it transparently.""" + if self._bin_fh is not None and not self._bin_fh.closed: + self._bin_fh.flush() + self._bin_fh.close() + self._bin_fh = None def get_frame_meta_data(self) -> pd.DataFrame: return pd.DataFrame(self.frame_meta_data) @@ -466,6 +564,7 @@ def write_frame_meta_data(self) -> None: prevents the noise-pipeline from producing duplicate-Id rows in the first place. The check below is the defence-in-depth. """ + self.close_binary() meta_df = self.get_frame_meta_data() validate_frames_id_uniqueness(meta_df) self._create_table(self.conn, meta_df, "Frames") @@ -652,33 +751,3 @@ def write_dia_ms_ms_windows(self, dia_ms_ms_windows: pd.DataFrame) -> None: self._create_table(self.conn, out, "DiaFrameMsMsWindows") - # TODO: these methods needs to be debugged - """ - def compress_frames(self, frames: List[TimsFrame], only_frame_one: bool = False, num_threads: int = 4) -> List[bytes]: - # same as compress_frame but for multiple frames - tofs, scans, intensities = [], [], [] - for frame in frames: - i = 1 if only_frame_one else frame.frame_id - tofs.append(self.mz_to_tof(i, frame.mz).astype(np.uint32)) - scans.append(self.inv_mobility_to_scan(i, frame.mobility).astype(np.uint32)) - intensities.append(frame.intensity.astype(np.uint32)) - - real_data = ims.get_data_for_compression_par(tofs, scans, intensities, self.helper_handle.num_scans, num_threads) - return [zstd.ZSTD_compress(bytes(data), 1) for data in real_data] - - def write_frames(self, frames: List[TimsFrame], scan_mode: int, only_frame_one: bool = False, num_threads: int = 4) -> None: - - compressed_data = self.compress_frames(frames, only_frame_one, num_threads=num_threads) - - for i, data in enumerate(compressed_data): - - self.frame_meta_data.append(self.build_frame_meta_row(frames[i], scan_mode, self.position, only_frame_one)) - - with open(self.binary_file, "ab") as bin_file: - bin_file.write( - (len(data) + 8).to_bytes(4, "little", signed=False) - ) - bin_file.write(int(self.helper_handle.num_scans).to_bytes(4, "little", signed=False)) - bin_file.write(data) - self.position = bin_file.tell() - """ \ No newline at end of file diff --git a/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/add_noise_from_real_data.py b/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/add_noise_from_real_data.py index d45d38088..15ac04f15 100644 --- a/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/add_noise_from_real_data.py +++ b/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/add_noise_from_real_data.py @@ -44,6 +44,8 @@ def add_real_data_noise_to_frames( num_fragment_frames: int = 5, acquisition_mode: str = 'DIA', pasef_meta: Optional[pd.DataFrame] = None, + noise_seed: Optional[int] = None, + num_threads: int = 4, ) -> List[TimsFrame]: """Add noise to frame. @@ -102,6 +104,25 @@ def add_real_data_noise_to_frames( d = acquisition_builder.frames_to_window_groups window_group_dict = dict(zip(d['frame'], d['window_group'])) fragment_frames = set(window_group_dict.keys()) + handle = acquisition_builder.tdf_writer.helper_handle + + # Fast path (imspy-connector >= 0.4.3): sample + add for the whole batch in Rust, in parallel, + # with per-frame seeded RNG streams. Falls back to the serial per-frame loop below on older + # connectors, which behaves identically apart from being unseeded. + if noise_seed is not None and hasattr(handle, "overlay_reference_noise"): + groups = [int(window_group_dict[f.frame_id]) if f.frame_id in fragment_frames else None + for f in frames] + return handle.overlay_reference_noise( + frames, groups, + num_precursor_frames=num_precursor_frames, + num_fragment_frames=num_fragment_frames, + max_intensity_precursor=intensity_max_precursor, + max_intensity_fragment=intensity_max_fragment, + take_precursor=precursor_sample_fraction, + take_fragment=fragment_sample_fraction, + seed=noise_seed, + num_threads=num_threads, + ) for frame in frames: if frame.frame_id in fragment_frames: diff --git a/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/assemble_frames.py b/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/assemble_frames.py index 26476d4dc..af50927f2 100644 --- a/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/assemble_frames.py +++ b/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/assemble_frames.py @@ -1,6 +1,7 @@ from pathlib import Path from typing import Optional import logging +import time import pandas as pd from tqdm import tqdm @@ -41,6 +42,7 @@ def assemble_frames( quad_transmission_min_probability: float = 0.5, quad_transmission_max_isotopes: int = 10, superimpose_on_reference: bool = False, + noise_seed: Optional[int] = None, ) -> None: """Assemble frames from frame ids and write them to the database. @@ -76,6 +78,8 @@ def assemble_frames( consider for transmission (default 10). superimpose_on_reference: If True, superimpose simulated signals on top of the full, unmodified reference frames (DIA only). + noise_seed: Master seed for the reference-noise sampler (DIA fast path); None keeps the + legacy unseeded per-frame sampling. Returns: None, writes frames to disk and metadata to database. @@ -114,6 +118,7 @@ def assemble_frames( quad_isotope_transmission_mode=quad_isotope_transmission_mode, quad_transmission_min_probability=quad_transmission_min_probability, quad_transmission_max_isotopes=quad_transmission_max_isotopes, + noise_seed=noise_seed, ) logger.info("Signal noise settings:") @@ -123,12 +128,17 @@ def assemble_frames( logger.info(f'Fragment m/z noise: {mz_noise_fragment}') logger.info(f'Fragment noise PPM: {fragment_noise_ppm}') + # per-phase wall clock inside the batch loop (build = parallel Rust, noise + write = main thread) + t_build = t_noise = t_write = 0.0 + n_written = 0 + # go over all frames in batches for b in tqdm(range(num_batches), total=num_batches, desc='frame assembly', ncols=100): start_index = b * batch_size stop_index = (b + 1) * batch_size ids = frame_ids[start_index:stop_index] + t0 = time.perf_counter() built_frames = frame_builder.build_frames( ids, mz_noise_precursor=mz_noise_precursor, @@ -139,6 +149,7 @@ def assemble_frames( num_threads=num_threads, fragment=fragment, ) + t1 = time.perf_counter() if superimpose_on_reference: built_frames = superimpose_reference_frames( @@ -157,14 +168,32 @@ def assemble_frames( num_fragment_frames=num_fragment_frames, acquisition_mode=acquisition_builder.acquisition_mode.mode, pasef_meta=pasef_meta, + noise_seed=noise_seed, + num_threads=num_threads, ) - - for frame in built_frames: - if acquisition_builder.acquisition_mode.mode == 'DDA': - acquisition_builder.tdf_writer.write_frame(frame, scan_mode=8) - else: - acquisition_builder.tdf_writer.write_frame(frame, scan_mode=9) - + t2 = time.perf_counter() + + scan_mode = 8 if acquisition_builder.acquisition_mode.mode == 'DDA' else 9 + acquisition_builder.tdf_writer.write_frames( + built_frames, scan_mode=scan_mode, num_threads=num_threads) + t3 = time.perf_counter() + + t_build += t1 - t0 + t_noise += t2 - t1 + t_write += t3 - t2 + n_written += len(built_frames) + + timings = { + 'frames': n_written, + 'batches': num_batches, + 'build_s': round(t_build, 2), + 'noise_s': round(t_noise, 2), + 'write_s': round(t_write, 2), + } + logger.info( + f"frame assembly phases: build {t_build:.1f} s | noise/superimpose {t_noise:.1f} s | " + f"convert+compress+write {t_write:.1f} s ({n_written} frames, {num_batches} batches)" + ) logger.info('Writing frame meta data to database ...') # write frame meta data to database @@ -216,3 +245,5 @@ def assemble_frames( acquisition_builder.synthetics_handle.get_table('pasef_meta'), id_mapping=id_mapping ) + + return timings diff --git a/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/simulate_frame_distributions_emg.py b/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/simulate_frame_distributions_emg.py index 4fbf0bea0..0249cd87e 100644 --- a/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/simulate_frame_distributions_emg.py +++ b/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/simulate_frame_distributions_emg.py @@ -69,8 +69,8 @@ def sample_sigma_lambda_emg(sigma_lower:ArrayLike, # TODO use rng sigma_hat = np.random.beta(a=sigma_alpha, b=sigma_beta, size=n) lambda_hat = np.random.beta(a=lambda_alpha, b=lambda_beta, size=n) - sigmas = sigma_lower + sigma_hat * (sigma_upper - sigma_lower) - lambdas = lambda_lower + lambda_hat * (lambda_upper - lambda_lower) + sigmas = np.maximum(sigma_lower + sigma_hat * (sigma_upper - sigma_lower), _EMG_PARAM_FLOOR) + lambdas = np.maximum(lambda_lower + lambda_hat * (lambda_upper - lambda_lower), _EMG_PARAM_FLOOR) return sigmas, lambdas @@ -115,11 +115,18 @@ def sample_sigma_k_emg(sigma_lower: ArrayLike, # TODO use rng sigma_hat = np.random.beta(a=sigma_alpha, b=sigma_beta, size=n) k_hat = np.random.beta(a=k_alpha, b=k_beta, size=n) - sigmas = sigma_lower + sigma_hat * (sigma_upper - sigma_lower) - ks = k_lower + k_hat * (k_upper - k_lower) + sigmas = np.maximum(sigma_lower + sigma_hat * (sigma_upper - sigma_lower), _EMG_PARAM_FLOOR) + ks = np.maximum(k_lower + k_hat * (k_upper - k_lower), _EMG_PARAM_FLOOR) return sigmas, ks +# An EMG peak with sigma = 0 or lambda = 0 has no width or no decay, and +# `estimate_mu_from_mode_emg` divides by their product. `np.random.beta` returns exactly 0.0 often +# enough to matter (and `k_lower_rt` defaults to 0), which made y = inf, sailed past the `y > 0` +# assert, and produced a NaN mu. The NaN then reached the Rust frame-occurrence search and panicked +# a worker on a NaN comparison. Clamp both to a floor far below any physical value instead. +_EMG_PARAM_FLOOR = 1e-12 + def erfcxinv(y:ArrayLike, n:int=10)->ArrayLike: """ Calculates the inverse of the scaled complementary error function (erfcx) @@ -131,15 +138,27 @@ def erfcxinv(y:ArrayLike, n:int=10)->ArrayLike: Returns: ArrayLike: The inverse of the scaled complementary error function at y. """ - assert np.all(y > 0), "y must be positive, as erfcx only maps to positive values." - # assert that y is an array of np.float64 y = np.array(y).astype(np.float64) + assert np.all(np.isfinite(y) & (y > 0)), ( + "y must be finite and positive; erfcx maps to (0, inf), and a non-finite y means a " + "degenerate EMG parameter (sigma or lambda at zero) reached this function." + ) # start value depends on the value of y xn_start = np.where(y<2,1/(y*np.sqrt(np.pi)), -np.sqrt(np.log(y/2, out=np.zeros_like(y), where=y>=2))) xn = xn_start - # Newton-Raphson method - for i in range(n): - xn = xn - (erfcx(xn)-y)/(2*xn*erfcx(xn)-2/np.sqrt(np.pi)) + # Newton-Raphson method. + # + # d/dx erfcx(x) = 2*x*erfcx(x) - 2/sqrt(pi), which tends to 0 as x -> +inf, i.e. for small y. + # Dividing by it there produced NaN for y below ~1.1e-8. Where the derivative underflows the + # start value is already the asymptote (erfcx(x) ~ 1/(x*sqrt(pi)), so x ~ 1/(y*sqrt(pi))), so + # keeping the current iterate is both safe and accurate: over y in [1e-8, 1e2] this agrees with + # the unguarded iteration to machine precision wherever that one produced a finite answer. + for _ in range(n): + f = erfcx(xn) - y + df = 2*xn*erfcx(xn) - 2/np.sqrt(np.pi) + step = np.divide(f, df, out=np.zeros_like(xn), where=np.abs(df) > 1e-12) + nxt = xn - step + xn = np.where(np.isfinite(nxt), nxt, xn) return xn def estimate_mu_from_mode_emg(mode: ArrayLike, sigma: ArrayLike, lambda_: ArrayLike)->ArrayLike: diff --git a/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/simulate_peptides.py b/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/simulate_peptides.py index 588809a77..7cc166a3c 100644 --- a/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/simulate_peptides.py +++ b/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/simulate_peptides.py @@ -1,9 +1,19 @@ +import os + import numpy as np import pandas as pd from typing import Optional from collections import Counter from imspy_core.data.peptide import PeptideSequence +try: # imspy-connector >= 0.4.4. The submodule is an attribute of the extension module, not an + # importable path, so `from imspy_connector.py_peptide import ...` raises ModuleNotFoundError + # and would silently leave us on the slow path. + import imspy_connector as _ims_connector + _batch_masses = _ims_connector.py_peptide.mono_isotopic_masses +except (ImportError, AttributeError): # pragma: no cover - exercised only on an older connector + _batch_masses = None + from imspy_predictors.rt.predictors import DeepChromatographyApex, load_deep_retention_time_predictor from imspy_predictors.ionization.predictors import predict_peptide_flyability_with_koina from imspy_predictors.utility import load_tokenizer_from_resources @@ -72,6 +82,10 @@ def simulate_peptides( max_length: int = 30, proteome_mix: bool = False, use_koina_model: Optional[str] = None, + num_threads: int = -1, + num_sample_peptides: Optional[int] = None, + sample_seed: Optional[int] = None, + sample_margin: float = 1.3, ) -> pd.DataFrame: """ Simulate peptides from a protein table. @@ -87,6 +101,13 @@ def simulate_peptides( max_length: Maximum length of the peptides. proteome_mix: If True, simulate a proteome mix. use_koina_model: If not None, use the Koina model to predict peptide flyability, i.e. chances of being measured in the experiment, currently only supports pfly. + num_threads: Threads for the batched monoisotopic mass calculation; -1 uses every core. + num_sample_peptides: If set, thin the table to roughly this many peptides *before* the + retention-time prediction, instead of predicting for the whole digest and discarding + afterwards. None keeps the old behaviour. + sample_seed: Seed for that thinning, so the choice is reproducible. + sample_margin: Oversample factor applied to `num_sample_peptides`, to leave headroom for + the retention-time filter below, which drops part of the early-eluting population. Returns: DataFrame with simulated peptides. """ @@ -103,7 +124,6 @@ def simulate_peptides( if len(peptide) < min_length or len(peptide) > max_length: continue - masses.append(PeptideSequence(peptide).mono_isotopic_mass) missed_cleavages.append(0) decoys.append(0) n_term.append(None) @@ -115,6 +135,15 @@ def simulate_peptides( protein_id.append(row.protein_id) i += 1 + # One batched, parallel call instead of constructing a PeptideSequence per peptide. Building + # that object per peptide cost ~190 us each — 48 s for 250 000 peptides — and was the largest + # remaining serial cost in the pipeline. Results are bitwise identical, including across thread + # counts. Falls back to the per-peptide path on a connector without the batch function. + if _batch_masses is not None: + masses = _batch_masses(sequences, num_threads if num_threads > 0 else (os.cpu_count() or 4)) + else: + masses = [PeptideSequence(p).mono_isotopic_mass for p in sequences] + peptide_table = pd.DataFrame({ "protein_id": protein_id, "peptide_id": peptide_id, @@ -128,6 +157,20 @@ def simulate_peptides( "events": events} ) + # Thin the table here, before retention-time prediction, rather than after it. + # + # The caller asks for `num_sample_peptides` but the digest produces `num_peptides_total`, and + # the retention-time predictor used to run over the whole digest so that 98 % of its output + # could be thrown away downstream. A small run spent ~40 s predicting times it never used. + # + # `sample_margin` is headroom for the retention-time filter further down, which removes part of + # the early-eluting population; without it the table could fall below what the caller asked for. + # The filter itself is density-based (it thins the early bins to the median bin count), so it + # behaves the same on a uniformly thinned table as on the full one. + if num_sample_peptides is not None and len(peptide_table) > num_sample_peptides: + target = min(len(peptide_table), int(np.ceil(num_sample_peptides * max(sample_margin, 1.0)))) + peptide_table = peptide_table.sample(n=target, random_state=sample_seed).reset_index(drop=True) + # Simulate peptide efficiency using rejection sampling in log-normal space efficiency = generate_normal_efficiency( n=len(peptide_table), @@ -202,9 +245,12 @@ def simulate_peptides( random_indices = np.array([], dtype=int) rt_filter[random_indices] = True + n_excluded = int((~rt_filter).sum()) peptide_table = peptide_table[rt_filter] if verbose: - print(f"Excluded {len(peptide_table) - rt_filter.sum()} peptides with low retention times.") + # Count before the filter is applied: the old message subtracted the filtered length + # from itself and so always reported 0. + print(f"Excluded {n_excluded} peptides with low retention times.") return peptide_table diff --git a/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/simulate_proteins.py b/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/simulate_proteins.py index 0776574ef..0e16ef952 100644 --- a/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/simulate_proteins.py +++ b/packages/imspy-simulation/src/imspy_simulation/timsim/jobs/simulate_proteins.py @@ -145,13 +145,90 @@ def protein_to_peptides(fasta, peptide = indexed_db[idx] peptide_set.add(peptide.to_unimod_sequence()) - return peptide_set + # Sorted, not a set: `sample_peptides_from_proteins` draws from this with + # `np.random.choice`, and a set's iteration order depends on Python's per-process + # string hash seed. That made the draw differ between runs even with the RNG seeded. + return sorted(peptide_set) except Exception as e: # Proteins with no valid peptides after digestion will raise an exception return None +def proteins_to_peptides_batched( + names, + sequences, + missed_cleavages: int = 2, + min_len: int = 7, + max_len: int = 30, + cleave_at: str = 'KR', + restrict: str = 'P', + generate_decoys: bool = False, + c_terminal: bool = True, + variable_mods: dict = {"M": ["[UNIMOD:35]"], "[": ["[UNIMOD:1]"]}, + static_mods: dict = {"C": "[UNIMOD:4]"}, + chunk_size: int = 2000, +) -> list: + """Digest many proteins with ONE sage indexed database per chunk instead of one per protein. + + Returns, for every input protein, the same UniMod peptide strings that + ``protein_to_peptides`` produces for that protein alone, as a sorted list: sage attaches every protein a + peptide occurs in (``peptide.proteins``), so a shared peptide is credited to each of its + proteins exactly as the per-protein digest did. Proteins without a valid peptide get an + empty list (the per-protein version returned None; callers drop both). + + Motivation: ``generate_indexed_database()`` sets up sage's parallel index build; calling it + 20 000 times (once per protein) spent ~3 000 core-seconds on a 120-core machine for a + 93 s stage that ran *faster* with 4 threads. One database per chunk is ~10x cheaper. + Headers are the positional index into ``names``; sage keeps the header token up to the first + whitespace, so protein names never leak into the mapping. Decoy proteins (``rev_`` prefix) + are credited to their target protein, matching the per-protein behaviour. + """ + names = list(names) + sequences = list(sequences) + result = [set() for _ in names] + if not names: + return result + enzyme_builder = EnzymeBuilder( + missed_cleavages=missed_cleavages, + min_len=min_len, + max_len=max_len, + cleave_at=cleave_at, + restrict=restrict, + c_terminal=c_terminal, + ) + for start in range(0, len(names), chunk_size): + stop = min(start + chunk_size, len(names)) + fasta = "".join( + f">{i}\n" + "\n".join(sequences[i][j:j + 60] for j in range(0, len(sequences[i]), 60)) + "\n" + for i in range(start, stop) + ) + sage_config = SageSearchConfiguration( + fasta=fasta, + enzyme_builder=enzyme_builder, + static_mods=static_mods, + variable_mods=variable_mods, + generate_decoys=generate_decoys, + bucket_size=int(np.power(2, 6)), + ) + try: + with SuppressStderr(): # sage panics (harmlessly) on proteins that yield no peptide + indexed_db = sage_config.generate_indexed_database() + except Exception: + # whole chunk without a single valid peptide: leave the sets empty + continue + for k in range(indexed_db.num_peptides): + peptide = indexed_db[PeptideIx(idx=k)] + seq = peptide.to_unimod_sequence() + for prot in peptide.proteins: + if prot.startswith("rev_"): + prot = prot[4:] + result[int(prot)].add(seq) + # Sorted lists rather than sets: see `protein_to_peptides`. The downstream sampler draws + # from these, and set iteration order is not stable across processes. + return [sorted(s) for s in result] + + def get_tenzer_hokey(): # Tenzer constants delta_exp = 0.005 @@ -246,21 +323,37 @@ def simulate_proteins( print("Skipping digestion, returning raw protein sequences as peptides.") # Generate peptides - sample["peptides"] = sample.apply(lambda f: protein_to_peptides( - generate_single_fasta( - f.name, - f.sequence, - ), - generate_decoys=generate_decoys, - variable_mods=variable_mods, - static_mods=static_mods, - cleave_at=cleave_at, - restrict=restrict, - missed_cleavages=missed_cleavages, - min_len=min_len, - max_len=max_len, - digest=digest, - ), axis=1) + if digest: + # one sage database per chunk of proteins instead of one per protein (see + # proteins_to_peptides_batched); same per-protein peptide sets as protein_to_peptides + peptide_sets = proteins_to_peptides_batched( + sample.index, sample.sequence, + generate_decoys=generate_decoys, + variable_mods=variable_mods, + static_mods=static_mods, + cleave_at=cleave_at, + restrict=restrict, + missed_cleavages=missed_cleavages, + min_len=min_len, + max_len=max_len, + ) + sample["peptides"] = [s if len(s) > 0 else None for s in peptide_sets] + else: + sample["peptides"] = sample.apply(lambda f: protein_to_peptides( + generate_single_fasta( + f.name, + f.sequence, + ), + generate_decoys=generate_decoys, + variable_mods=variable_mods, + static_mods=static_mods, + cleave_at=cleave_at, + restrict=restrict, + missed_cleavages=missed_cleavages, + min_len=min_len, + max_len=max_len, + digest=digest, + ), axis=1) # Remove None values sample = sample[sample.peptides.notnull()] diff --git a/packages/imspy-simulation/src/imspy_simulation/timsim/simulator.py b/packages/imspy-simulation/src/imspy_simulation/timsim/simulator.py index c024b2041..193a622ad 100644 --- a/packages/imspy-simulation/src/imspy_simulation/timsim/simulator.py +++ b/packages/imspy-simulation/src/imspy_simulation/timsim/simulator.py @@ -9,11 +9,14 @@ import argparse import logging import time +import json +import resource from pathlib import Path from dataclasses import dataclass, field from typing import Dict, Optional import toml +import numpy as np import pandas as pd from tabulate import tabulate @@ -62,6 +65,94 @@ def format_duration(seconds: float) -> str: return f"{hours}h {minutes}m" +class StageTracker: + """Per-stage wall clock, CPU seconds (all threads) and peak RSS, logged at every stage + transition and dumped to ``timings.json`` next to ``synthetic_data.db``. + + Wall is ``time.perf_counter``; CPU is ``resource.getrusage(RUSAGE_SELF)`` user+sys, which + includes rayon/torch worker threads, so ``cpu_s / wall_s`` is the average number of busy + cores during the stage. ``peak_rss_gb`` is the process high-water mark reached by the end of + the stage (monotonic). ``laps`` are optional named sub-intervals inside a stage. + """ + + def __init__(self, logger: logging.Logger) -> None: + self._logger = logger + self.records: list[dict] = [] + self._current: Optional[dict] = None + self._t0 = 0.0 + self._cpu0 = 0.0 + self._lap_t = 0.0 + + @staticmethod + def _cpu_seconds() -> float: + ru = resource.getrusage(resource.RUSAGE_SELF) + return ru.ru_utime + ru.ru_stime + + @staticmethod + def _peak_rss_gb() -> float: + return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 ** 2 # KiB -> GiB on Linux + + def begin(self, name: str) -> None: + self.end() + self._current = {'stage': name, 'laps': {}} + self._t0 = self._lap_t = time.perf_counter() + self._cpu0 = self._cpu_seconds() + + def lap(self, label: str) -> None: + """Record the time since the stage start (or the previous lap) under ``label``.""" + if self._current is None: + return + now = time.perf_counter() + self._current['laps'][label] = round(now - self._lap_t, 2) + self._lap_t = now + + def annotate(self, **fields) -> None: + """Attach extra measurements (e.g. per-phase timers from a job) to the current stage.""" + if self._current is not None: + self._current.update(fields) + + def end(self) -> None: + if self._current is None: + return + wall = time.perf_counter() - self._t0 + cpu = self._cpu_seconds() - self._cpu0 + rec = self._current + rec.update(wall_s=round(wall, 2), cpu_s=round(cpu, 2), + avg_cores=round(cpu / wall, 2) if wall > 0 else 0.0, + peak_rss_gb=round(self._peak_rss_gb(), 2)) + if not rec['laps']: + rec.pop('laps') + self.records.append(rec) + self._current = None + laps = f" | laps {rec['laps']}" if 'laps' in rec else "" + self._logger.info( + f" stage done: {rec['stage']}: wall {wall:.1f} s | cpu {cpu:.1f} core-s " + f"(avg {rec['avg_cores']:.1f} cores) | peak RSS {rec['peak_rss_gb']:.2f} GB{laps}" + ) + + def summary_lines(self) -> list[str]: + total = sum(r['wall_s'] for r in self.records) or 1.0 + rows = [(r['stage'], f"{r['wall_s']:.1f}", f"{100 * r['wall_s'] / total:.0f}%", + f"{r['cpu_s']:.1f}", f"{r['avg_cores']:.1f}", f"{r['peak_rss_gb']:.2f}") + for r in self.records] + return tabulate(rows, headers=['stage', 'wall s', 'share', 'cpu s', 'avg cores', 'peak RSS GB'], + tablefmt='simple').splitlines() + + def dump(self, path) -> None: + self.end() + payload = { + 'stages': self.records, + 'total_wall_s': round(sum(r['wall_s'] for r in self.records), 2), + 'total_cpu_s': round(sum(r['cpu_s'] for r in self.records), 2), + 'peak_rss_gb': round(self._peak_rss_gb(), 2), + } + try: + with open(path, 'w') as fh: + json.dump(payload, fh, indent=1) + except OSError as e: # never fail a finished simulation over telemetry + self._logger.warning(f"could not write {path}: {e}") + + @dataclass class SimulationStats: """Track statistics for simulation summary.""" @@ -1310,9 +1401,15 @@ def main(): # Initialize timer and stats timer = SimulationTimer() stats = SimulationStats() + stages = StageTracker(logger) + + def section(title: str) -> str: + """Start timing ``title`` and return its banner (closes the previous stage).""" + stages.begin(title) + return section_header(title, use_unicode) # Configure GPU - logger.info(section_header("GPU Configuration", use_unicode)) + logger.info(section("GPU Configuration")) configure_gpu_memory(memory_limit_gb=config.gpu_memory_limit_gb, use_gpu=config.use_gpu) # Handle macOS Bruker SDK limitation @@ -1322,7 +1419,7 @@ def main(): use_bruker_sdk = False # Load modifications config - logger.info(section_header("Loading Configuration", use_unicode)) + logger.info(section("Loading Configuration")) script_dir = Path(__file__).parent modifications_path = config.modifications if not modifications_path or modifications_path == "": @@ -1348,6 +1445,21 @@ def main(): for key, value in config_dict.items(): logger.debug(f" {key}: {value}") + # Seed the global NumPy RNG for the whole run. + # + # Until 0.4.6 `sample_seed` reached exactly two places: the pandas down-sample below and the + # reference-noise overlay. Everything that decides *what* gets simulated — which peptides are + # drawn per protein and their abundances (simulate_peptides), the EMG retention-time shape, + # ion-mobility and scan variance, phosphosite choice, the decoy draw in digest_fasta — used + # bare `np.random.*` against the unseeded global RNG. Two runs of an identical config therefore + # shared only 11 of ~1690 peptides: same configuration, effectively unrelated datasets. + # + # Seeding once here makes a run reproducible from its config, which is what `sample_seed` + # always claimed to do. Note this changes which peptides a given config selects compared with + # releases before 0.4.6, so previously generated datasets are not reproduced by re-running + # their config — they were never reproducible in the first place. + np.random.seed(config.sample_seed & 0xFFFFFFFF) + # Prepare paths save_path = check_path(config.save_path) # An Astral run has no Bruker reference .d (it builds from a .raw template). @@ -1360,7 +1472,7 @@ def main(): f.write(tabulate(table_data, headers=["Argument", "Value"], tablefmt="grid")) # Build the acquisition (frames, scans, etc.) - logger.info(section_header("Building Acquisition", use_unicode)) + logger.info(section("Building Acquisition")) if not config.silent_mode: logger.info("") logger.info(f" Experiment: {name}") @@ -1486,7 +1598,7 @@ def main(): factors = {} if config.from_existing and not resume_after: - logger.info(section_header("Loading Existing Simulation", use_unicode)) + logger.info(section("Loading Existing Simulation")) # Load existing simulation data if config.acquisition_type == 'DIA': existing_sim_handle = SyntheticExperimentDataHandleDIA(database_path=config.existing_path) @@ -1610,7 +1722,7 @@ def main(): # ---------------------------------------- findings_result = None if config.from_findings and not resume_after: - logger.info(section_header("Loading Search Engine Findings", use_unicode)) + logger.info(section("Loading Search Engine Findings")) rt_lower = acquisition_builder.frame_table['time'].min() rt_upper = acquisition_builder.frame_table['time'].max() findings_result = load_findings( @@ -1638,7 +1750,7 @@ def main(): protein_list, peptide_list = [], [] if not config.from_existing and not config.from_findings and not resume_after: - logger.info(section_header("Processing FASTA Files", use_unicode)) + logger.info(section("Processing FASTA Files")) for fasta_name, fasta_path in fastas.items(): if not config.silent_mode: logger.info("") @@ -1668,6 +1780,7 @@ def main(): remove_degenerate_peptides=config.remove_degenerate_peptides, ) + stages.lap('digest') # JOB 1: Simulate peptides if not config.silent_mode: logger.info(" Creating peptides from proteins...") @@ -1683,8 +1796,12 @@ def main(): min_length=config.min_len, max_length=config.max_len, proteome_mix=config.proteome_mix, + num_threads=config.num_threads, + num_sample_peptides=config.num_sample_peptides if config.sample_peptides else None, + sample_seed=config.sample_seed, ) + stages.lap('sample_peptides+rt_filter') if config.proteome_mix: # Scale by mixture factor peptides_tmp['events'] *= mixture_factor @@ -1725,7 +1842,7 @@ def main(): logger.info(f" Total peptides to simulate: {peptides.shape[0]}") # JOB 3: Simulate retention times - logger.info(section_header("Simulating Retention Times", use_unicode)) + logger.info(section("Simulating Retention Times")) # Support both rt_model and deprecated koina_rt_model rt_model = config.rt_model or config.koina_rt_model if rt_model: @@ -1739,7 +1856,7 @@ def main(): # Simulate RT for from_findings when not provided in the input if config.from_findings and findings_result is not None and not findings_result.has_rt and not resume_after: - logger.info(section_header("Simulating Retention Times (not in findings)", use_unicode)) + logger.info(section("Simulating Retention Times (not in findings)")) rt_model = config.rt_model or config.koina_rt_model if rt_model: logger.info(f" Using RT model: {rt_model}") @@ -1764,7 +1881,7 @@ def main(): if resume_after != "ions": # JOB 4: Frame distributions - logger.info(section_header("Simulating Frame Distributions", use_unicode)) + logger.info(section("Simulating Frame Distributions")) peptides = simulate_frame_distributions_emg( peptides=peptides, frames=acquisition_builder.frame_table, @@ -1819,7 +1936,7 @@ def main(): ) if need_charge: - logger.info(section_header("Simulating Charge States", use_unicode)) + logger.info(section("Simulating Charge States")) # JOB 5: Charge states ions = simulate_charge_states( peptides=peptides, @@ -1838,7 +1955,7 @@ def main(): ions = ions.drop_duplicates(subset=['sequence', 'charge']) if need_im: - logger.info(section_header("Simulating Ion Mobilities", use_unicode)) + logger.info(section("Simulating Ion Mobilities")) # JOB 6: Ion mobilities if config.ccs_model: logger.info(f" Using CCS model: {config.ccs_model}") @@ -1855,7 +1972,7 @@ def main(): if not config.from_existing: # JOB 7: Precursor isotopic distributions (always needed) - logger.info(section_header("Simulating Precursor Isotope Patterns", use_unicode)) + logger.info(section("Simulating Precursor Isotope Patterns")) ions = simulate_precursor_spectra_sequence( ions=ions, num_threads=num_threads, @@ -1863,7 +1980,7 @@ def main(): ) # JOB 8: Scan distributions - logger.info(section_header("Simulating Scan Distributions", use_unicode)) + logger.info(section("Simulating Scan Distributions")) ions = simulate_scan_distributions_with_variance( ions=ions, scans=acquisition_builder.scan_table, @@ -1908,7 +2025,7 @@ def main(): ) if config.acquisition_type == 'DDA': - logger.info(section_header("Simulating DDA-PASEF Selection", use_unicode)) + logger.info(section("Simulating DDA-PASEF Selection")) pasef_meta, precursors = simulate_dda_pasef_selection_scheme( acquisition_builder=acquisition_builder, verbose=not config.silent_mode, @@ -1932,7 +2049,7 @@ def main(): else: # Resuming from ions checkpoint — populate DB tables - logger.info(section_header("Restoring Database from Checkpoint", use_unicode)) + logger.info(section("Restoring Database from Checkpoint")) acquisition_builder.synthetics_handle.create_table(table_name='proteins', table=proteins) acquisition_builder.synthetics_handle.create_table(table_name='peptides', table=peptides) acquisition_builder.synthetics_handle.create_table(table_name='ions', table=ions) @@ -1943,7 +2060,7 @@ def main(): logger.info(f" Restored: proteins ({len(proteins)}), peptides ({len(peptides)}), ions ({len(ions)})") # JOB 9: Simulate fragment intensities - logger.info(section_header("Simulating Fragment Intensities", use_unicode)) + logger.info(section("Simulating Fragment Intensities")) # Support both intensity_model and deprecated fragment_intensity_model intensity_model = config.intensity_model or config.fragment_intensity_model if intensity_model: @@ -2122,8 +2239,8 @@ def main(): logger=logger, ) else: - logger.info(section_header("Assembling Frames", use_unicode)) - assemble_frames( + logger.info(section("Assembling Frames")) + assembly_timings = assemble_frames( acquisition_builder=acquisition_builder, frames=acquisition_builder.frame_table, batch_size=config.batch_size, @@ -2148,7 +2265,10 @@ def main(): quad_transmission_min_probability=config.quad_transmission_min_probability, quad_transmission_max_isotopes=config.quad_transmission_max_isotopes, superimpose_on_reference=config.superimpose_on_reference, + noise_seed=config.sample_seed, ) + if assembly_timings: + stages.annotate(**assembly_timings) # mzPROV provenance: sign the authored Bruker .d (self-disclosure that this is # TimSim-simulated data). Bruker-only — mzprov v0 doesn't canonicalize vendor # .raw, so the Thermo build-from-template branch above is intentionally skipped. @@ -2175,7 +2295,7 @@ def main(): # Optional: Generate preview video for visual inspection if config.generate_preview_video: if VIDEO_GENERATION_AVAILABLE: - logger.info(section_header("Generating Preview Video", use_unicode)) + logger.info(section("Generating Preview Video")) data_path = os.path.join(save_path, name, f"{name}.d") video_path = os.path.join(save_path, f"{name}_preview.mp4") mode = 'dda' if config.acquisition_type == 'DDA' else 'dia' @@ -2200,7 +2320,12 @@ def main(): # Print completion banner and summary print(simulation_complete_banner(use_unicode)) + stages.end() total_time = timer.format_duration(timer.total_elapsed()) + logger.info(" Stage timings") + for line in stages.summary_lines(): + logger.info(f" {line}") + logger.info("") logger.info(" Simulation Summary") logger.info(f" {'─' * 40}") logger.info(f" Experiment: {stats.experiment_name}") @@ -2213,6 +2338,12 @@ def main(): logger.info(f" Total time: {total_time}") logger.info(f" Output: {stats.output_path}") logger.info("") + try: + timings_path = Path(acquisition_builder.path) / 'timings.json' + except NameError: # pragma: no cover - acquisition builder is always created in main() + timings_path = Path(save_path) / 'timings.json' + stages.dump(timings_path) + logger.info(f" Stage timings written to {timings_path}") if __name__ == '__main__': diff --git a/packages/imspy-simulation/src/imspy_simulation/utility.py b/packages/imspy-simulation/src/imspy_simulation/utility.py index 1faee5340..55fa06c28 100644 --- a/packages/imspy-simulation/src/imspy_simulation/utility.py +++ b/packages/imspy-simulation/src/imspy_simulation/utility.py @@ -182,9 +182,20 @@ def read_acquisition_config(acquisition_name: str = 'dia') -> Dict[str, Any]: # Function to convert a list (or a pandas series) to a JSON string def python_list_to_json_string(lst, as_float=True, num_decimals: int = 4) -> str: + """Serialise a 1-D sequence as a JSON array (floats rounded to ``num_decimals``, or ints). + + Vectorised: one ``np.round`` over the whole array instead of a NumPy scalar call per element + (this helper runs once per peptide/ion in four pipeline stages, so the per-element version + cost minutes on a 250k simulation). Output is identical to the previous implementation. + """ if as_float: - return json.dumps([float(np.round(x, num_decimals)) for x in lst]) - return json.dumps([int(x) for x in lst]) + arr = np.asarray(lst) + if arr.dtype.kind != 'f': # ints, bools, object lists of Python numbers + arr = arr.astype(np.float64) + # round in the input precision (float32 stays float32) and widen afterwards, exactly as + # the former per-element ``float(np.round(x, n))`` did, so the strings are byte-identical. + return json.dumps(np.round(arr, num_decimals).astype(np.float64).tolist()) + return json.dumps(np.asarray(lst).astype(np.int64).tolist()) # load peptides and ions diff --git a/packages/imspy-simulation/tests/test_batched_frame_writer.py b/packages/imspy-simulation/tests/test_batched_frame_writer.py new file mode 100644 index 000000000..6be15c31b --- /dev/null +++ b/packages/imspy-simulation/tests/test_batched_frame_writer.py @@ -0,0 +1,153 @@ +"""The batched writer must produce exactly what the per-frame writer produced. + +`TDFWriter.write_frames` moves conversion, dedup, interleave and zstd into one parallel Rust call. +The output is a vendor file format, so "equivalent" is not good enough: the compressed payloads must +be byte-identical to the per-frame path, and the Frames statistics must match. +""" +import os + +import numpy as np +import pytest +import zstd + +from imspy_simulation.tdf import dedup_scan_tof +from imspy_simulation.utility import get_compressible_data + +REF_D = "/scratch/timsim-demo/TIMSIM-HeLa-RUSTW-SMOKE/TIMSIM-HeLa-RUSTW-SMOKE.d" +pytestmark = pytest.mark.skipif(not os.path.exists(REF_D), reason="local reference .d not available") + + +def _reference(ds, frames, num_scans): + """Exactly what TDFWriter.compress_frame did, frame by frame, through the SDK.""" + out = [] + for f in frames: + tof = np.array(ds.mz_to_tof(f.frame_id, f.mz)).astype(np.uint32) + scan = np.array(ds.inverse_mobility_to_scan(f.frame_id, f.mobility)).astype(np.uint32) + intensity = f.intensity.astype(np.uint32) + scan, tof, intensity = dedup_scan_tof(scan, tof, intensity) + blob = zstd.ZSTD_compress(bytes(get_compressible_data(tof, scan, intensity, num_scans)), 0) + out.append((len(intensity), + int(intensity.max()) if len(intensity) else 0, + int(intensity.sum()), blob)) + return out + + +def _load(n=120): + from imspy_core.timstof import TimsDatasetDIA + ds = TimsDatasetDIA(REF_D, in_memory=False, use_bruker_sdk=True) + frames = [f for f in (ds.get_tims_frame(i) for i in range(1, n + 1)) if len(f.mz) > 0] + assert frames, "reference produced no non-empty frames" + return ds, frames, int(ds.num_scans) + + +@pytest.mark.parametrize("num_threads", [1, 4, 16]) +def test_batched_writer_is_byte_identical_to_per_frame(num_threads): + ds, frames, num_scans = _load() + if not hasattr(ds, "build_compressed_frames"): + pytest.skip("connector without build_compressed_frames") + ref = _reference(ds, frames, num_scans) + got = ds.build_compressed_frames( + [f.frame_id for f in frames], + [np.ascontiguousarray(f.mz, dtype=np.float64) for f in frames], + [np.ascontiguousarray(f.mobility, dtype=np.float64) for f in frames], + [np.ascontiguousarray(f.intensity, dtype=np.float64) for f in frames], + num_scans, 0, num_threads, + ) + assert len(got) == len(ref) + for i, ((n1, m1, s1, b1), (n2, m2, s2, b2)) in enumerate(zip(ref, got)): + assert (n1, m1, s1) == (n2, m2, s2), f"frame {i}: Frames statistics differ" + # A zstd frame written without the decompressed size in its header cannot be read back by + # the simple API every existing .d was written with, so the header shape matters too. + assert b1 == bytes(b2), f"frame {i}: compressed payload differs" + + +def test_batched_writer_output_still_decompresses_with_the_simple_api(): + ds, frames, num_scans = _load(n=30) + if not hasattr(ds, "build_compressed_frames"): + pytest.skip("connector without build_compressed_frames") + got = ds.build_compressed_frames( + [f.frame_id for f in frames], + [np.ascontiguousarray(f.mz, dtype=np.float64) for f in frames], + [np.ascontiguousarray(f.mobility, dtype=np.float64) for f in frames], + [np.ascontiguousarray(f.intensity, dtype=np.float64) for f in frames], + num_scans, 0, 4, + ) + for _, _, _, blob in got: + assert len(zstd.ZSTD_uncompress(bytes(blob))) > 0 + + +def test_ragged_arrays_raise_instead_of_panicking(): + """A ragged triplet used to index out of bounds inside Rust; it must surface as an error.""" + ds, frames, num_scans = _load(n=10) + if not hasattr(ds, "build_compressed_frames"): + pytest.skip("connector without build_compressed_frames") + mz = [np.ascontiguousarray(f.mz, dtype=np.float64) for f in frames] + mobility = [np.ascontiguousarray(f.mobility, dtype=np.float64) for f in frames] + intensity = [np.ascontiguousarray(f.intensity, dtype=np.float64) for f in frames] + mobility[0] = mobility[0][:-1] # one value short + with pytest.raises(RuntimeError, match="must have the same length"): + ds.build_compressed_frames([f.frame_id for f in frames], mz, mobility, intensity, + num_scans, 0, 4) + + +def test_mismatched_outer_lengths_raise(): + ds, frames, num_scans = _load(n=10) + if not hasattr(ds, "build_compressed_frames"): + pytest.skip("connector without build_compressed_frames") + mz = [np.ascontiguousarray(f.mz, dtype=np.float64) for f in frames] + mobility = [np.ascontiguousarray(f.mobility, dtype=np.float64) for f in frames] + intensity = [np.ascontiguousarray(f.intensity, dtype=np.float64) for f in frames] + with pytest.raises(RuntimeError, match="same length"): + ds.build_compressed_frames([f.frame_id for f in frames][:-1], mz, mobility, intensity, + num_scans, 0, 4) + + +def test_two_writers_produce_equivalent_d_folders(tmp_path): + """End-to-end: write the same frames through both paths and compare the .d folders. + + The per-frame check above compares payloads in isolation; this one compares what actually + lands on disk — the Frames table (including the TimsId byte offsets, which depend on every + preceding frame's compressed length) and the decompressed content of every frame. + + Note the files are equivalent, not byte-identical: the Python `zstd` module and the Rust + `zstd` crate link different libzstd builds and can pick a different, equally valid encoding + of the same input. Observed at 4 bytes in 8.6 MB, in 2 frames of 600, all decompressing to + identical payloads. The contract is the decompressed data and the table, not the encoding. + """ + import sqlite3 + import pandas as pd + from imspy_simulation.tdf import TDFWriter + + ds, frames, _ = _load(n=150) + if not hasattr(ds, "build_compressed_frames"): + pytest.skip("connector without build_compressed_frames") + + def write(name, fn): + w = TDFWriter(helper_handle=ds, path=str(tmp_path), exp_name=name) + fn(w) + w.close_binary() + w.write_frame_meta_data() + return tmp_path / name + + p_old = write("old.d", lambda w: [w.write_frame(f, scan_mode=9) for f in frames]) + p_new = write("new.d", lambda w: w.write_frames(frames, scan_mode=9, num_threads=8)) + + old = pd.read_sql("select * from Frames order by Id", sqlite3.connect(str(p_old / "analysis.tdf"))) + new = pd.read_sql("select * from Frames order by Id", sqlite3.connect(str(p_new / "analysis.tdf"))) + assert old.equals(new), "Frames table differs between the per-frame and batched writers" + + b_old = (p_old / "analysis.tdf_bin").read_bytes() + b_new = (p_new / "analysis.tdf_bin").read_bytes() + assert len(b_old) == len(b_new) + + # Walk the frames by their recorded offsets and compare decompressed payloads. + for _, row in old.iterrows(): + off = int(row.TimsId) + for buf in (b_old, b_new): + assert int.from_bytes(buf[off:off + 4], "little") > 0 + n_old = int.from_bytes(b_old[off:off + 4], "little") + n_new = int.from_bytes(b_new[off:off + 4], "little") + assert n_old == n_new, f"frame {int(row.Id)}: compressed length differs" + d_old = zstd.ZSTD_uncompress(bytes(b_old[off + 8:off + n_old])) + d_new = zstd.ZSTD_uncompress(bytes(b_new[off + 8:off + n_new])) + assert d_old == d_new, f"frame {int(row.Id)}: decompressed payload differs" diff --git a/packages/imspy-simulation/tests/test_batched_peptide_masses.py b/packages/imspy-simulation/tests/test_batched_peptide_masses.py new file mode 100644 index 000000000..7c751d272 --- /dev/null +++ b/packages/imspy-simulation/tests/test_batched_peptide_masses.py @@ -0,0 +1,51 @@ +"""The batched monoisotopic mass path must match the per-peptide one exactly. + +Constructing a `PeptideSequence` per peptide cost ~190 us each, which made 250 000 peptides take +48 s and was the largest remaining serial cost in the pipeline. Two changes: the lookup tables and +the UniMod regex are built once rather than per call (9.4x on its own, and it helps every caller), +and the connector gained a batched parallel entry point. Mass values must not move. +""" +import numpy as np +import pytest + +from imspy_core.data.peptide import PeptideSequence + +_ims = pytest.importorskip("imspy_connector") +mono_isotopic_masses = getattr(getattr(_ims, "py_peptide", None), "mono_isotopic_masses", None) +if mono_isotopic_masses is None: # pragma: no cover + pytest.skip("connector without mono_isotopic_masses", allow_module_level=True) + + +def _peptides(n, seed=0): + rng = np.random.default_rng(seed) + aa = list("ACDEFGHIKLMNPQRSTVWY") + return ["".join(rng.choice(aa, int(rng.integers(7, 31)))) for _ in range(n)] + + +@pytest.mark.parametrize("num_threads", [1, 4, 16]) +def test_batched_masses_match_the_per_peptide_path(num_threads): + peps = _peptides(2000) + ref = [PeptideSequence(p).mono_isotopic_mass for p in peps] + got = mono_isotopic_masses(peps, num_threads) + assert got == ref, "batched masses differ from the per-peptide path" + + +def test_batched_masses_are_thread_count_invariant(): + """Bitwise, not approximately: float addition is not associative, so a parallel sum that + varied with scheduling would make the simulation irreproducible.""" + peps = _peptides(5000, seed=7) + assert mono_isotopic_masses(peps, 1) == mono_isotopic_masses(peps, 16) + + +def test_modifications_are_still_counted(): + """The UniMod regex moved into a static; modified peptides must still pick up their delta.""" + plain, mod = "PEPTIDEK", "PEPTIDEM[UNIMOD:35]K" + m_plain, m_mod = mono_isotopic_masses([plain, mod], 1) + assert m_mod > m_plain + # UNIMOD:35 is oxidation, +15.994915, on top of the extra methionine residue. + assert abs((m_mod - m_plain) - (PeptideSequence("PEPTIDEMK").mono_isotopic_mass + - PeptideSequence(plain).mono_isotopic_mass) - 15.994915) < 1e-6 + + +def test_empty_batch(): + assert mono_isotopic_masses([], 4) == [] diff --git a/packages/imspy-simulation/tests/test_emg_degenerate_params.py b/packages/imspy-simulation/tests/test_emg_degenerate_params.py new file mode 100644 index 000000000..55b1e11be --- /dev/null +++ b/packages/imspy-simulation/tests/test_emg_degenerate_params.py @@ -0,0 +1,69 @@ +"""The EMG elution model must not emit NaN for degenerate parameter draws. + +`np.random.beta` returns exactly 0.0 often enough to matter, and `k_lower_rt` defaults to 0, so a +zero sigma or lambda was reachable. `estimate_mu_from_mode_emg` divides by their product, which made +the erfcx target infinite, slipped past the positivity assert, and produced a NaN mu. That NaN +reached the Rust frame-occurrence search and panicked a worker thread mid-run. +""" +import numpy as np +import pytest + +from imspy_simulation.timsim.jobs.simulate_frame_distributions_emg import ( + erfcxinv, estimate_mu_from_mode_emg, sample_sigma_k_emg, sample_sigma_lambda_emg, +) +from scipy.special import erfcx + + +@pytest.mark.parametrize("y", [1e-12, 1e-10, 1e-8, 1e-6, 1e-3, 0.1, 1.0, 2.0, 10.0, 1e3]) +def test_erfcxinv_is_finite_and_inverts_erfcx(y): + x = erfcxinv(np.array([y]))[0] + assert np.isfinite(x), f"erfcxinv({y}) is not finite" + # erfcx overflows for very negative x, so only round-trip where that is representable. + back = erfcx(x) + if np.isfinite(back) and back > 0: + assert abs(back - y) / y < 1e-6, f"erfcxinv({y}) -> {x} -> {back}" + + +def test_erfcxinv_rejects_a_degenerate_target(): + """An infinite target means a zero sigma or lambda got through; fail loudly, never silently NaN.""" + with pytest.raises(AssertionError): + erfcxinv(np.array([np.inf])) + + +@pytest.mark.parametrize("sampler,second", [(sample_sigma_k_emg, "k"), (sample_sigma_lambda_emg, "lambda")]) +def test_sampled_parameters_are_strictly_positive(sampler, second): + """The draw must never hand a zero downstream, even with a zero lower bound (k_lower_rt is 0). + + Beta parameters are chosen to make a 0.0 draw likely: a small alpha underflows to exactly zero. + """ + np.random.seed(0) + kw = dict(sigma_lower=0.0, sigma_upper=1.0, sigma_alpha=0.01, sigma_beta=10.0, n=50000) + kw.update({f"{second}_lower": 0.0, f"{second}_upper": 10.0, + f"{second}_alpha": 0.01, f"{second}_beta": 10.0}) + sigmas, others = sampler(**kw) + assert (sigmas > 0).all(), f"{(sigmas <= 0).sum()} sigma draws were non-positive" + assert (others > 0).all(), f"{(others <= 0).sum()} {second} draws were non-positive" + mu = estimate_mu_from_mode_emg(np.full_like(sigmas, 10.0), sigmas, others) + assert np.isfinite(mu).all(), f"{(~np.isfinite(mu)).sum()} non-finite mu values" + + +def test_thinning_happens_before_retention_time_prediction(): + """The expensive predictor must see only the peptides we intend to keep. + + The digest produces `num_peptides_total` regardless of the requested complexity, so predicting + retention times for the whole digest meant a small run spent most of its time on peptides it + then discarded. Thinning first must still deliver at least the requested count, since the + retention-time filter drops part of the early-eluting population afterwards. + """ + import inspect + from imspy_simulation.timsim.jobs import simulate_peptides as mod + + src = inspect.getsource(mod.simulate_peptides) + thin = src.index("peptide_table.sample(") + predict = src.index("simulate_separation_times_pandas") + assert thin < predict, "thinning must precede retention-time prediction" + + sig = inspect.signature(mod.simulate_peptides).parameters + assert "num_sample_peptides" in sig and "sample_margin" in sig + # Headroom must be real, or the retention-time filter can leave us short of the request. + assert sig["sample_margin"].default > 1.0 diff --git a/packages/imspy-simulation/tests/test_perf_batched_digest_and_overlay.py b/packages/imspy-simulation/tests/test_perf_batched_digest_and_overlay.py new file mode 100644 index 000000000..f4c46406e --- /dev/null +++ b/packages/imspy-simulation/tests/test_perf_batched_digest_and_overlay.py @@ -0,0 +1,173 @@ +"""Parity tests: batched sage digest == per-protein digest; seeded parallel noise overlay is +deterministic and equals the serial per-frame sampling in distribution.""" +import os +import numpy as np +import pandas as pd +import pytest + +from imspy_simulation.timsim.jobs.simulate_proteins import ( + protein_to_peptides, proteins_to_peptides_batched, generate_single_fasta, parse_fasta_to_dataframe, +) + +FASTA = "/scratch/timsim-demo/timsim-necro/flow/hela_subset.fasta" +REF_D = "/scratch/timsim-demo/TIMSIM-HeLa-RUSTW-SMOKE/TIMSIM-HeLa-RUSTW-SMOKE.d" + +VAR = {"M": ["[UNIMOD:35]"], "[": ["[UNIMOD:1]"]} +STAT = {"C": "[UNIMOD:4]"} + + +@pytest.mark.parametrize("decoys", [False, True]) +def test_batched_digest_matches_per_protein(decoys): + tbl = _proteins(n=400) + batched = proteins_to_peptides_batched( + tbl.index, tbl.sequence, generate_decoys=decoys, variable_mods=VAR, static_mods=STAT, + chunk_size=150, # force several chunks + ) + n_nonempty = 0 + for (idx, row), got in zip(tbl.iterrows(), batched): + ref = protein_to_peptides( + generate_single_fasta(idx, row.sequence), generate_decoys=decoys, + variable_mods=VAR, static_mods=STAT, cleave_at='KR', restrict='P', + missed_cleavages=2, min_len=7, max_len=30, digest=True, + ) + if not ref: + assert got == [], f"protein {idx}: per-protein digest empty, batched gave {len(got)}" + else: + assert got == sorted(ref), f"protein {idx}: differing peptides" + # Sorted, so the draw downstream cannot depend on set iteration order. + assert got == sorted(got), f"protein {idx}: batched digest is not sorted" + n_nonempty += 1 + assert n_nonempty >= min(len(tbl) - 1, 300) + + +def test_batched_digest_empty_input(): + assert proteins_to_peptides_batched([], []) == [] + + +# Synthetic proteins: several tryptic sites each, so every one yields multiple peptides. +# Deliberately not the local HeLa fixture — the ordering guarantee is what is under test, and a +# test that silently skips on a machine without that file proves nothing in CI. +_SYNTH = { + "sp|TEST1|SYN1": "MKVLATRSEDKPEPTIDEKSAMPLERGGVTLLKDIFFERENTKAVQWYRPLNESTGGK", + "sp|TEST2|SYN2": "MSSGQKALTTRPEPTIDEKMNQWYFRAVCDEKLLGGTTRSEQPPLKYYTTRDDFGHIK", + "sp|TEST3|SYN3": "MTTKAAQPLRSEEDNKPEPTIDERGGWWYTKLNMFCVVRDDQQTTLKSSAHHPPYYRK", + "sp|TEST4|SYN4": "MNNKEFGHLRPPQTTSKVVCCWYKDDEEFFRAALLNMSTTKGGHHYYPPQRWWEEDDK", +} + + +def _proteins(n: int) -> pd.DataFrame: + """Protein table for the digest tests: the local HeLa subset when present, else the synthetic + set above. CI has no local fixtures, and a skipped parity test would leave the batched digest + unchecked there.""" + if os.path.exists(FASTA): + return parse_fasta_to_dataframe(FASTA).sample(n=n, random_state=7) + return pd.DataFrame({"sequence": list(_SYNTH.values())}, index=list(_SYNTH)) + + +def test_digest_output_is_order_stable_not_a_set(): + """Set iteration order varies with Python's per-process hash seed; the sampler draws from + these collections, so an unordered container makes a seeded run irreproducible.""" + ids, seqs = list(_SYNTH), list(_SYNTH.values()) + batched = proteins_to_peptides_batched(ids, seqs, variable_mods=VAR, static_mods=STAT) + assert any(len(p) > 1 for p in batched), "fixture produced nothing to order" + for peptides in batched: + assert isinstance(peptides, list), "digest must return an ordered container" + assert peptides == sorted(peptides) + for idx, seq in zip(ids, seqs): + got = protein_to_peptides(generate_single_fasta(idx, seq), generate_decoys=False, + variable_mods=VAR, static_mods=STAT, cleave_at='KR', restrict='P', + missed_cleavages=2, min_len=7, max_len=30, digest=True) + if got: + assert isinstance(got, list) and got == sorted(got) + + +@pytest.mark.skipif(not os.path.exists(REF_D), reason="local reference .d not available") +def test_overlay_reference_noise_is_seeded_and_matches_serial_statistics(): + from imspy_core.timstof import TimsDatasetDIA + ds = TimsDatasetDIA(REF_D, in_memory=False) + if not hasattr(ds, "overlay_reference_noise"): + pytest.skip("imspy-core without overlay_reference_noise") + info = ds.dia_ms_ms_info + wg_of = dict(zip(info.Frame, info.WindowGroup)) + ids = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18] + frames = [ds.get_tims_frame(i) for i in ids] + groups = [int(wg_of[i]) if i in wg_of else None for i in ids] + kw = dict(num_precursor_frames=5, num_fragment_frames=5, max_intensity_precursor=150000.0, + max_intensity_fragment=150000.0, take_precursor=0.2, take_fragment=0.2, num_threads=4) + a = ds.overlay_reference_noise(frames, groups, seed=41, **kw) + b = ds.overlay_reference_noise(frames, groups, seed=41, **kw) + c = ds.overlay_reference_noise(frames, groups, seed=42, **kw) + assert len(a) == len(frames) + for fa, fb, fc, f0 in zip(a, b, c, frames): + assert fa.frame_id == f0.frame_id + np.testing.assert_array_equal(fa.mz, fb.mz) + np.testing.assert_array_equal(fa.intensity, fb.intensity) + np.testing.assert_array_equal(fa.scan, fb.scan) + assert len(fa.mz) >= len(f0.mz) # noise only adds peaks + assert not (len(fa.mz) == len(fc.mz) and np.array_equal(fa.mz, fc.mz)) # seed matters + # serial reference path (unseeded) must give the same added-peak count within sampling noise + added_par = np.array([len(x.mz) - len(f.mz) for x, f in zip(a, frames)], dtype=float) + added_ser = [] + for f, g in zip(frames, groups): + if g is None: + noise = ds.sample_precursor_signal(5, 150000.0, 0.2) + else: + noise = ds.sample_fragment_signal(5, g, 150000.0, 0.2) + added_ser.append(len((f + noise).mz) - len(f.mz)) + added_ser = np.array(added_ser, dtype=float) + assert added_par.sum() > 0 and added_ser.sum() > 0 + ratio = added_par.sum() / added_ser.sum() + assert 0.5 < ratio < 2.0, f"parallel/serial added-peak ratio {ratio:.2f}" + + +# --------------------------------------------------------------------------- +# Bruker SDK thread-safety regression guards. +# +# `tims_index_to_mz` is not safe to call concurrently on one SDK handle: before the +# fix, the same seed produced m/z differing by up to 4e-5 Th on ~8 % of peaks, in +# 3/8 runs at 4 threads and 7/8 at 32. rustdf now reads through the SDK-free formula +# converter inside parallel regions instead of serialising, so these must hold with +# the SDK enabled (the default). +# --------------------------------------------------------------------------- + +@pytest.mark.skipif(not os.path.exists(REF_D), reason="local reference .d not available") +@pytest.mark.parametrize("threads", [1, 4, 32]) +def test_overlay_is_bit_identical_across_thread_counts_with_sdk(threads): + from imspy_core.timstof import TimsDatasetDIA + ds = TimsDatasetDIA(REF_D, in_memory=False, use_bruker_sdk=True) + info = ds.dia_ms_ms_info + wg_of = dict(zip(info.Frame, info.WindowGroup)) + ids = list(range(1, 19)) + frames = [ds.get_tims_frame(i) for i in ids] + groups = [int(wg_of[i]) if i in wg_of else None for i in ids] + kw = dict(num_precursor_frames=5, num_fragment_frames=5, max_intensity_precursor=150000.0, + max_intensity_fragment=150000.0, take_precursor=0.2, take_fragment=0.2) + ref = ds.overlay_reference_noise(frames, groups, seed=41, num_threads=1, **kw) + for rep in range(5): + got = ds.overlay_reference_noise(frames, groups, seed=41, num_threads=threads, **kw) + for a, b in zip(ref, got): + assert len(a.mz) == len(b.mz), f"rep {rep}: peak count differs on frame {a.frame_id}" + np.testing.assert_array_equal(a.mz, b.mz) + np.testing.assert_array_equal(a.intensity, b.intensity) + + +@pytest.mark.skipif(not os.path.exists(REF_D), reason="local reference .d not available") +def test_sdk_free_converter_matches_the_sdk(): + """The general fix rests on this: the formula converter reproduces the SDK, so parallel + regions may use it instead of dropping to one thread.""" + from imspy_core.timstof import TimsDatasetDIA + sdk = TimsDatasetDIA(REF_D, in_memory=False, use_bruker_sdk=True) + fml = TimsDatasetDIA(REF_D, in_memory=False, use_bruker_sdk=False) + checked = 0 + for fid in [1, 2, 3, 5, 8, 13, 21, 34, 55, 89]: + a, b = sdk.get_tims_frame(fid), fml.get_tims_frame(fid) + if len(a.mz) == 0: + continue + assert len(a.mz) == len(b.mz) + np.testing.assert_array_equal(a.scan, b.scan) + np.testing.assert_array_equal(a.intensity, b.intensity) + # last-bit agreement, orders of magnitude below any search tolerance + assert np.abs(a.mz - b.mz).max() < 1e-8, "m/z drift between SDK and formula converter" + assert np.abs(a.mobility - b.mobility).max() < 1e-12 + checked += 1 + assert checked >= 5 diff --git a/packages/imspy-simulation/tests/test_perf_serial_hotspots.py b/packages/imspy-simulation/tests/test_perf_serial_hotspots.py new file mode 100644 index 000000000..1d25fca43 --- /dev/null +++ b/packages/imspy-simulation/tests/test_perf_serial_hotspots.py @@ -0,0 +1,147 @@ +"""Parity tests for the serial hot-spot fixes (packed-key (scan, tof) dedup, vectorised JSON +helper, persistent tdf_bin handle, stage tracker). + +Each fix must reproduce the previous implementation byte-for-byte; the reference +implementations are copied here verbatim from the pre-change code. +""" +import io +import json +import logging +import os +import tempfile + +import numpy as np +import pytest + +from imspy_simulation.tdf import dedup_scan_tof +from imspy_simulation.utility import python_list_to_json_string + + +# --------------------------------------------------------------------------- references +def _dedup_reference(scan, tof, intensity): + scan_tof = np.stack((scan.astype(np.uint32), tof.astype(np.uint32)), axis=1) + unique_pairs, inverse_indices = np.unique(scan_tof, axis=0, return_inverse=True) + summed_intensity = np.bincount(inverse_indices.ravel(), weights=intensity) + unique_scan = unique_pairs[:, 0] + unique_tof = unique_pairs[:, 1] + sort_idx = np.lexsort((unique_tof, unique_scan)) + return unique_scan[sort_idx], unique_tof[sort_idx], summed_intensity[sort_idx].astype(np.uint32) + + +def _json_reference(lst, as_float=True, num_decimals=4): + if as_float: + return json.dumps([float(np.round(x, num_decimals)) for x in lst]) + return json.dumps([int(x) for x in lst]) + + +# --------------------------------------------------------------------------- dedup +@pytest.mark.parametrize("seed,n,n_scans,n_tof", [(0, 10, 5, 20), (1, 5_000, 927, 400_000), + (2, 60_000, 927, 400_000), (3, 3_000, 3, 10)]) +def test_dedup_scan_tof_matches_reference(seed, n, n_scans, n_tof): + rng = np.random.default_rng(seed) + scan = rng.integers(0, n_scans, n).astype(np.uint32) + tof = rng.integers(0, n_tof, n).astype(np.uint32) + intensity = rng.integers(1, 5_000, n).astype(np.uint32) + ref = _dedup_reference(scan, tof, intensity) + got = dedup_scan_tof(scan, tof, intensity) + for r, g in zip(ref, got): + np.testing.assert_array_equal(r, g) + assert got[0].dtype == np.uint32 and got[1].dtype == np.uint32 and got[2].dtype == np.uint32 + # order contract: sorted by scan, then tof; no duplicate cells left + key = got[0].astype(np.uint64) << np.uint64(32) | got[1].astype(np.uint64) + assert np.all(np.diff(key) > 0) + + +def test_dedup_scan_tof_sums_duplicates(): + scan = np.array([2, 2, 1, 2], dtype=np.uint32) + tof = np.array([7, 7, 9, 3], dtype=np.uint32) + inten = np.array([10, 5, 1, 2], dtype=np.uint32) + s, t, i = dedup_scan_tof(scan, tof, inten) + assert s.tolist() == [1, 2, 2] + assert t.tolist() == [9, 3, 7] + assert i.tolist() == [1, 2, 15] + + +def test_dedup_scan_tof_empty(): + e = np.array([], dtype=np.uint32) + s, t, i = dedup_scan_tof(e, e, e) + assert len(s) == len(t) == len(i) == 0 + + +def test_dedup_scan_tof_large_tof_values_do_not_collide(): + # tof close to 2**32 must stay separate from scan+1, tof=0 + scan = np.array([0, 1], dtype=np.uint32) + tof = np.array([np.iinfo(np.uint32).max, 0], dtype=np.uint32) + inten = np.array([1, 2], dtype=np.uint32) + s, t, i = dedup_scan_tof(scan, tof, inten) + assert s.tolist() == [0, 1] and t.tolist() == [np.iinfo(np.uint32).max, 0] and i.tolist() == [1, 2] + + +# --------------------------------------------------------------------------- json helper +@pytest.mark.parametrize("lst", [ + [], [0.0], [1.23456789, 2.5, 1e-7, 123456.123456], list(np.random.default_rng(0).random(300)), + np.random.default_rng(1).random(50).astype(np.float32), [1, 2, 3], np.array([0.99995, 0.00005]), +]) +def test_json_float_matches_reference(lst): + assert python_list_to_json_string(lst) == _json_reference(lst) + assert python_list_to_json_string(lst, num_decimals=2) == _json_reference(lst, num_decimals=2) + + +@pytest.mark.parametrize("lst", [[], [0], [1, 2, 3], np.array([5, 6, 7], dtype=np.int32), + np.array([1.0, 2.0]), list(range(1000))]) +def test_json_int_matches_reference(lst): + assert python_list_to_json_string(lst, as_float=False) == _json_reference(lst, as_float=False) + + +def test_json_nan_roundtrip_like_reference(): + lst = [float('nan'), 1.0] + assert python_list_to_json_string(lst) == _json_reference(lst) + + +# --------------------------------------------------------------------------- persistent handle +def test_tdf_writer_binary_handle_reopens_and_closes(): + from imspy_simulation.tdf import TDFWriter + w = TDFWriter.__new__(TDFWriter) # bypass __init__ (needs a Bruker handle) + with tempfile.TemporaryDirectory() as d: + w.binary_file = os.path.join(d, "analysis.tdf_bin") + w._bin_fh = None + fh = w._binary_handle() + fh.write(b"abc") + assert w._binary_handle() is fh + w.close_binary() + assert w._bin_fh is None + w.close_binary() # idempotent + w._binary_handle().write(b"def") + w.close_binary() + with open(w.binary_file, "rb") as r: + assert r.read() == b"abcdef" + + +# --------------------------------------------------------------------------- stage tracker +def test_stage_tracker_records_and_dumps(): + from imspy_simulation.timsim.simulator import StageTracker + logger = logging.getLogger("stage-tracker-test") + stream = io.StringIO() + handler = logging.StreamHandler(stream) + logger.addHandler(handler) + logger.setLevel(logging.INFO) + st = StageTracker(logger) + st.begin("A") + st.lap("first") + st.lap("second") + st.begin("B") + st.annotate(build_s=1.5) + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "timings.json") + st.dump(path) + payload = json.load(open(path)) + names = [r['stage'] for r in payload['stages']] + assert names == ["A", "B"] + a, b = payload['stages'] + assert set(a['laps']) == {"first", "second"} + assert 'laps' not in b and b['build_s'] == 1.5 + for r in (a, b): + for k in ('wall_s', 'cpu_s', 'avg_cores', 'peak_rss_gb'): + assert k in r + assert "stage done: A" in stream.getvalue() and "stage done: B" in stream.getvalue() + assert len(st.summary_lines()) >= 3 diff --git a/rustdf/src/data/dda.rs b/rustdf/src/data/dda.rs index 8d9b212fa..5fd3ecde2 100644 --- a/rustdf/src/data/dda.rs +++ b/rustdf/src/data/dda.rs @@ -683,14 +683,24 @@ impl TimsDatasetDDA { None => pasef_info.iter().collect(), }; - // Note: The Bruker SDK is NOT thread-safe, so we must use sequential iteration - // when the SDK is being used for index conversion. - let uses_bruker_sdk = self.loader.uses_bruker_sdk(); + // The Bruker SDK's index conversion is not safe to call concurrently on one handle, but + // the SDK-free formula converter evaluates the very calibration curves the file carries + // (matched the SDK to 3e-11 Th / 1 ULP on the reference data) and *is* `Sync`. So in SDK + // mode we read through that instead of dropping to a single thread; only a file without + // usable calibration tables still forces the serial path. + let sdk_free = self.loader.sdk_free_converter(); + let conv: Option<&(dyn IndexConverter + Sync)> = + sdk_free.as_ref().map(|c| c as &(dyn IndexConverter + Sync)); + let read_frame = |frame_id: u32| match conv { + Some(c) => self.loader.get_frame_with_converter(frame_id, c), + None => self.loader.get_frame(frame_id), + }; + let force_serial = self.loader.uses_bruker_sdk() && conv.is_none(); // Helper closure to process a single PASEF fragment let process_fragment = |pasef_info: &PasefMsMsMeta| -> PASEFDDAFragment { // get the frame - let frame = self.loader.get_frame(pasef_info.frame_id as u32); + let frame = read_frame(pasef_info.frame_id as u32); // get five percent of the scan range let scan_margin = (pasef_info.scan_num_end - pasef_info.scan_num_begin) / 20; @@ -718,8 +728,8 @@ impl TimsDatasetDDA { } }; - if uses_bruker_sdk { - // Sequential processing when using Bruker SDK (not thread-safe) + if force_serial { + // No thread-safe converter available: fall back to one thread. filtered_pasef_info.iter().map(|info| process_fragment(info)).collect() } else { // Parallel processing when using simple index converter (thread-safe) @@ -783,9 +793,19 @@ impl TimsDatasetDDA { } // Step 4: Build fragment data with aggregation (merge frames for same precursor) - // Note: The Bruker SDK is NOT thread-safe, so we must use sequential iteration - // when the SDK is being used for index conversion. - let uses_bruker_sdk = self.loader.uses_bruker_sdk(); + // The Bruker SDK's index conversion is not safe to call concurrently on one handle, but + // the SDK-free formula converter evaluates the very calibration curves the file carries + // (matched the SDK to 3e-11 Th / 1 ULP on the reference data) and *is* `Sync`. So in SDK + // mode we read through that instead of dropping to a single thread; only a file without + // usable calibration tables still forces the serial path. + let sdk_free = self.loader.sdk_free_converter(); + let conv: Option<&(dyn IndexConverter + Sync)> = + sdk_free.as_ref().map(|c| c as &(dyn IndexConverter + Sync)); + let read_frame = |frame_id: u32| match conv { + Some(c) => self.loader.get_frame_with_converter(frame_id, c), + None => self.loader.get_frame(frame_id), + }; + let force_serial = self.loader.uses_bruker_sdk() && conv.is_none(); // Helper closure to process a single precursor let process_precursor = |(precursor_id, pasef_infos): (&i64, &Vec<&PasefMsMsMeta>)| -> Option { @@ -807,7 +827,7 @@ impl TimsDatasetDDA { for pasef_info in pasef_infos { // Get the frame and filter by scan range - let frame = self.loader.get_frame(pasef_info.frame_id as u32); + let frame = read_frame(pasef_info.frame_id as u32); // Get five percent of the scan range for margin let scan_margin = (pasef_info.scan_num_end - pasef_info.scan_num_begin) / 20; @@ -860,8 +880,8 @@ impl TimsDatasetDDA { }) }; - let fragment_data: Vec = if uses_bruker_sdk { - // Sequential processing when using Bruker SDK (not thread-safe) + let fragment_data: Vec = if force_serial { + // No thread-safe converter available: fall back to one thread. pasef_by_precursor .iter() .filter_map(process_precursor) diff --git a/rustdf/src/data/dia.rs b/rustdf/src/data/dia.rs index bd692e013..7fa1eeb55 100644 --- a/rustdf/src/data/dia.rs +++ b/rustdf/src/data/dia.rs @@ -634,7 +634,7 @@ impl TimsDatasetDIA { let mut sampled_frames: Vec = self.collect_typed_noise_samples(&pool, num_frames, max_intensity, take_probability, &mut rng, - |meta| meta.id as u32); + |meta| meta.id as u32, None); Self::accumulate_or_typed_empty(&mut sampled_frames, target_type) } @@ -661,7 +661,7 @@ impl TimsDatasetDIA { let mut sampled_frames: Vec = self.collect_typed_noise_samples(&pool, num_frames, max_intensity, take_probability, &mut rng, - |&fid| fid); + |&fid| fid, None); Self::accumulate_or_typed_empty(&mut sampled_frames, target_type) } @@ -676,6 +676,7 @@ impl TimsDatasetDIA { take_probability: f64, rng: &mut impl rand::Rng, key_fn: impl Fn(&T) -> u32, + conv: Option<&(dyn IndexConverter + Sync)>, ) -> Vec { let mut out: Vec = Vec::with_capacity(num_frames); if pool.is_empty() || num_frames == 0 { @@ -687,12 +688,17 @@ impl TimsDatasetDIA { while out.len() < num_frames && attempts < max_attempts { attempts += 1; let Some(candidate_meta) = pool.choose(rng) else { break }; - let candidate = self - .loader - .get_frame(key_fn(candidate_meta)) + let frame_id = key_fn(candidate_meta); + // `conv` is the SDK-free converter handed in by parallel callers (see + // `overlay_reference_noise`); the serial samplers read with the loader's own. + let raw = match conv { + Some(c) => self.loader.get_frame_with_converter(frame_id, c), + None => self.loader.get_frame(frame_id), + }; + let candidate = raw .filter_ranged(0.0, 2000.0, 0, 1000, 0.0, 5.0, 1.0, max_intensity, 0, i32::MAX) - .generate_random_sample(take_probability); + .generate_random_sample_with_rng(take_probability, rng); // Skip the loader's empty-frame Default-fallback so the // returned noise never carries MsType::Unknown into the // outer simulator add. @@ -729,6 +735,90 @@ impl TimsDatasetDIA { } /// All DIA window_group IDs present in the file (sorted unique). + /// Reference-noise overlay for a whole batch, in parallel and reproducible. + /// + /// For every simulated frame, draws `num_*_frames` random reference frames of the same kind + /// (precursor frames for MS1, frames of the same DIA window group for MS2), keeps a + /// `take_*` fraction of their peaks below `max_intensity_*`, accumulates them and adds the + /// result onto the simulated frame with `TimsFrame + TimsFrame` — exactly the per-frame + /// `sample_precursor_signal` / `sample_fragment_signal` + `frame + noise` sequence the + /// Python job used to run serially, but over a rayon pool and with a per-frame RNG stream + /// derived from `seed` and the frame id, so the result does not depend on scheduling. + /// + /// `window_groups[i]` is `Some(group)` for a fragment frame and `None` for a precursor frame. + pub fn overlay_reference_noise( + &self, + frames: Vec, + window_groups: Vec>, + num_precursor_frames: usize, + num_fragment_frames: usize, + max_intensity_precursor: f64, + max_intensity_fragment: f64, + take_precursor: f64, + take_fragment: f64, + seed: u64, + num_threads: usize, + ) -> Vec { + assert_eq!(frames.len(), window_groups.len(), "frames and window_groups must align"); + let precursor_pool: Vec = self + .meta_data + .iter() + .filter(|x| x.ms_ms_type == 0) + .map(|m| m.id as u32) + .collect(); + let mut fragment_pools: FxHashMap> = FxHashMap::default(); + for info in &self.dia_ms_ms_info { + fragment_pools.entry(info.window_group).or_default().push(info.frame_id); + } + // The Bruker SDK's `tims_index_to_mz` is not safe to call concurrently on one handle + // (same seed gave m/z that differed by ~1e-8 relative on ~8 % of peaks, run to run). + // Read the reference frames through the SDK-free formula converter instead; if the + // file carries no calibration tables, fall back to one thread so the result stays + // deterministic. + let sdk_free = self.loader.sdk_free_converter(); + let threads = if self.loader.uses_bruker_sdk() && sdk_free.is_none() { + 1 + } else { + num_threads.max(1) + }; + let conv: Option<&(dyn IndexConverter + Sync)> = + sdk_free.as_ref().map(|c| c as &(dyn IndexConverter + Sync)); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap(); + pool.install(|| { + frames + .into_par_iter() + .zip(window_groups.into_par_iter()) + .map(|(frame, wg)| { + use rand::SeedableRng; + let stream = (frame.frame_id as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15); + let mut rng = rand::rngs::StdRng::seed_from_u64(seed ^ stream); + let noise = match wg { + Some(group) => { + let ids: &[u32] = fragment_pools + .get(&group) + .map(|v| v.as_slice()) + .unwrap_or(&[]); + let mut s = self.collect_typed_noise_samples( + ids, num_fragment_frames, max_intensity_fragment, + take_fragment, &mut rng, |&fid| fid, conv); + Self::accumulate_or_typed_empty(&mut s, MsType::FragmentDia) + } + None => { + let mut s = self.collect_typed_noise_samples( + &precursor_pool, num_precursor_frames, max_intensity_precursor, + take_precursor, &mut rng, |&fid| fid, conv); + Self::accumulate_or_typed_empty(&mut s, MsType::Precursor) + } + }; + frame + noise + }) + .collect() + }) + } + pub fn dia_window_groups(&self) -> Vec { let mut gs: Vec = self .dia_ms_ms_info diff --git a/rustdf/src/data/handle.rs b/rustdf/src/data/handle.rs index 25267ce5e..8c9661b00 100644 --- a/rustdf/src/data/handle.rs +++ b/rustdf/src/data/handle.rs @@ -528,6 +528,30 @@ impl BrukerFormulaConverter { }) } + /// True only when substituting this converter for the Bruker SDK cannot change a single + /// output value, which is the only condition under which it is safe to swap in. + /// + /// Two things have to hold for **every** frame in the run: + /// * its `MzCalibration` row is `ModelType` 1 — the only m/z model this converter reproduces + /// bit-exactly. ModelType 2 shares the C0/C1/C2 curve but omits the proprietary C8..C14 + /// fine correction, landing ~2 ppm away; measured on a real ModelType-2 run that moves + /// **11.5 % of arbitrary m/z values into a neighbouring TOF bin**, which would change the + /// `(scan, tof)` dedup and therefore the bytes written to `analysis.tdf_bin`. + /// * it resolves to a real calibration row. Frames with NULL `T1`/`T2`, a missing `Frames` + /// row, or a link to an unsupported calibration row silently fall back to the default + /// calibration in `mz_for`/`im_for`, and the SDK does no such thing. + /// + /// Callers that cannot use a non-exact converter must fall back to the SDK, serially. + pub fn is_sdk_exact(&self) -> bool { + if self.frames.is_empty() { + return false; + } + self.frames.values().all(|f| { + self.mz_rows.get(&f.mz_id).map(|r| r.model_type == 1).unwrap_or(false) + && self.im_rows.contains_key(&f.tims_id) + }) + } + /// Spread (max - min, in degC) of `Frames.T1` and `Frames.T2` across the run. /// /// A near-zero spread means per-frame and fixed-frame calibration agree; a @@ -666,8 +690,11 @@ pub struct TimsLazyLoder { pub index_converter: TimsIndexConverter, } -impl TimsData for TimsLazyLoder { - fn get_frame(&self, frame_id: u32) -> TimsFrame { +impl TimsLazyLoder { + /// Same as [`TimsData::get_frame`] but converts tof/scan with an explicit converter instead of + /// `self.index_converter`. Parallel readers use this to bypass the Bruker SDK, whose + /// `tims_index_to_mz` is not safe to call concurrently on one handle. + pub fn get_frame_with_converter(&self, frame_id: u32, conv: &dyn IndexConverter) -> TimsFrame { let frame_index = (frame_id - 1) as usize; // turns out, there can be empty frames in the data, check for that, if so, return an empty frame @@ -783,10 +810,8 @@ impl TimsData for TimsLazyLoder { let intensity_dbl = intensities_.iter().map(|&x| x as f64).collect::>(); let tof_i32 = tof_indices_.iter().map(|&x| x as i32).collect::>(); - let mz = self.index_converter.tof_to_mz(frame_id, &tof_indices_); - let inv_mobility = self - .index_converter - .scan_to_inverse_mobility(frame_id, &scan); + let mz = conv.tof_to_mz(frame_id, &tof_indices_); + let inv_mobility = conv.scan_to_inverse_mobility(frame_id, &scan); let ms_type_raw = self.raw_data_layout.frame_meta_data[frame_index].ms_ms_type; let ms_type = match ms_type_raw { @@ -823,10 +848,8 @@ impl TimsData for TimsLazyLoder { let tof_i32 = tof.iter().map(|&x| x as i32).collect(); let scan = flatten_scan_values(&scan, true); - let mz = self.index_converter.tof_to_mz(frame_id, &tof); - let inv_mobility = self - .index_converter - .scan_to_inverse_mobility(frame_id, &scan); + let mz = conv.tof_to_mz(frame_id, &tof); + let inv_mobility = conv.scan_to_inverse_mobility(frame_id, &scan); let ms_type_raw = self.raw_data_layout.frame_meta_data[frame_index].ms_ms_type; @@ -856,6 +879,12 @@ impl TimsData for TimsLazyLoder { } } } +} + +impl TimsData for TimsLazyLoder { + fn get_frame(&self, frame_id: u32) -> TimsFrame { + self.get_frame_with_converter(frame_id, &self.index_converter) + } fn get_raw_frame(&self, frame_id: u32) -> RawTimsFrame { let frame_index = (frame_id - 1) as usize; @@ -956,8 +985,11 @@ pub struct TimsInMemoryLoader { compressed_data: Vec, } -impl TimsData for TimsInMemoryLoader { - fn get_frame(&self, frame_id: u32) -> TimsFrame { +impl TimsInMemoryLoader { + /// Same as [`TimsData::get_frame`] but converts tof/scan with an explicit converter instead of + /// `self.index_converter`. Parallel readers use this to bypass the Bruker SDK, whose + /// `tims_index_to_mz` is not safe to call concurrently on one handle. + pub fn get_frame_with_converter(&self, frame_id: u32, conv: &dyn IndexConverter) -> TimsFrame { let raw_frame = self.get_raw_frame(frame_id); let raw_frame = match raw_frame.ms_type { @@ -973,10 +1005,8 @@ impl TimsData for TimsInMemoryLoader { let tof_i32 = raw_frame.tof.iter().map(|&x| x as i32).collect(); let scan = flatten_scan_values(&raw_frame.scan, true); - let mz = self.index_converter.tof_to_mz(frame_id, &raw_frame.tof); - let inverse_mobility = self - .index_converter - .scan_to_inverse_mobility(frame_id, &scan); + let mz = conv.tof_to_mz(frame_id, &raw_frame.tof); + let inverse_mobility = conv.scan_to_inverse_mobility(frame_id, &scan); let ims_frame = ImsFrame::new( raw_frame.retention_time, @@ -993,6 +1023,12 @@ impl TimsData for TimsInMemoryLoader { ims_frame, } } +} + +impl TimsData for TimsInMemoryLoader { + fn get_frame(&self, frame_id: u32) -> TimsFrame { + self.get_frame_with_converter(frame_id, &self.index_converter) + } fn get_raw_frame(&self, frame_id: u32) -> RawTimsFrame { let frame_index = (frame_id - 1) as usize; @@ -1407,6 +1443,57 @@ impl TimsDataLoader { } } + /// `get_frame` with an explicit index converter (see [`TimsLazyLoder::get_frame_with_converter`]). + pub fn get_frame_with_converter(&self, frame_id: u32, conv: &dyn IndexConverter) -> TimsFrame { + match self { + TimsDataLoader::InMemory(loader) => loader.get_frame_with_converter(frame_id, conv), + TimsDataLoader::Lazy(loader) => loader.get_frame_with_converter(frame_id, conv), + } + } + + /// Path of the `.d` folder this loader reads from. + pub fn raw_data_path(&self) -> &str { + match self { + TimsDataLoader::InMemory(loader) => &loader.raw_data_layout.raw_data_path, + TimsDataLoader::Lazy(loader) => &loader.raw_data_layout.raw_data_path, + } + } + + /// The loader's own index converter, if it is safe to use from several threads — i.e. + /// everything except the Bruker SDK wrapper, whose `tims_index_to_mz` must not be called + /// concurrently on one handle. + pub fn sync_index_converter(&self) -> Option<&(dyn IndexConverter + Sync)> { + let converter = match self { + TimsDataLoader::InMemory(loader) => &loader.index_converter, + TimsDataLoader::Lazy(loader) => &loader.index_converter, + }; + match converter { + TimsIndexConverter::BrukerLib(_) => None, + TimsIndexConverter::Simple(c) => Some(c), + TimsIndexConverter::Calibrated(c) => Some(c), + TimsIndexConverter::Lookup(c) => Some(c), + TimsIndexConverter::BrukerFormula(c) => Some(c), + } + } + + /// A converter that is safe to use from many threads at once, for readers that want to + /// run in parallel while the dataset itself was opened with the Bruker SDK. + /// + /// Returns `None` when the loader's own converter is already thread-safe (anything but + /// `BrukerLib`), or when the file carries no calibration tables the SDK-free + /// [`BrukerFormulaConverter`] could evaluate — callers must then serialise their reads. + pub fn sdk_free_converter(&self) -> Option { + if !self.uses_bruker_sdk() { + return None; + } + BrukerFormulaConverter::from_d_folder(self.raw_data_path(), 1) + .ok() + // Only substitute when the result is indistinguishable from the SDK's; otherwise the + // caller has to stay on the SDK and serialise. See `is_sdk_exact`. + .filter(|c| c.is_sdk_exact()) + .map(TimsIndexConverter::BrukerFormula) + } + /// Check if the Bruker SDK is being used for index conversion. /// The Bruker SDK is NOT thread-safe, so parallel operations that call /// the index converter must be disabled when using the SDK. diff --git a/rustdf/src/data/utility.rs b/rustdf/src/data/utility.rs index 18f14e960..cb9e93a54 100644 --- a/rustdf/src/data/utility.rs +++ b/rustdf/src/data/utility.rs @@ -51,7 +51,7 @@ pub fn zstd_compress(decompressed_data: &[u8], compression_level: i32) -> io::Re /// frame data straight into the encoder, producing negative/garbage TOF deltas /// that vendor readers (e.g. DiaNN) reject. Mirroring the Python preprocessing /// here keeps the two writers byte-for-byte identical. -fn sort_dedup_scan_tof( +pub fn sort_dedup_scan_tof( scans: &[u32], tofs: &[u32], intensities: &[u32], @@ -375,3 +375,97 @@ pub fn merge_ranges(mut ranges: Vec<(usize, usize)>) -> Vec<(usize, usize)> { out.push(cur); out } + +/// One frame's worth of writer output: the statistics the TDF `Frames` row needs, plus the +/// zstd-compressed payload ready to append to `analysis.tdf_bin`. +pub struct CompressedFrame { + pub num_peaks: u32, + pub max_intensity: u32, + pub summed_intensity: u64, + pub data: Vec, +} + +/// The whole per-frame TDF writer pipeline — m/z → TOF and 1/K0 → scan conversion, `(scan, tof)` +/// dedup, interleave into the Bruker layout, zstd — for a batch of frames, over a rayon pool. +/// +/// This used to run one frame at a time in Python because the conversion went through the Bruker +/// SDK, which is not safe to call concurrently on one handle. Given a `Sync` converter (the SDK-free +/// `BrukerFormulaConverter` reproduces the SDK's integer output exactly in both of these directions) +/// the whole pipeline is per-frame independent, so only the final append has to stay ordered. +/// +/// `mz`, `mobility` and `intensity` are parallel slices, one entry per frame, matching `frame_ids`. +/// Results come back in input order. +pub fn build_compressed_frames( + converter: &(dyn crate::data::handle::IndexConverter + Sync), + frame_ids: &[u32], + mz: &[&[f64]], + mobility: &[&[f64]], + intensity: &[&[f64]], + max_scans: u32, + compression_level: i32, + num_threads: usize, +) -> Result, String> { + if frame_ids.len() != mz.len() || frame_ids.len() != mobility.len() || frame_ids.len() != intensity.len() { + return Err(format!( + "frame_ids ({}), mz ({}), mobility ({}) and intensity ({}) must have the same length", + frame_ids.len(), mz.len(), mobility.len(), intensity.len() + )); + } + // Per frame the three arrays index each other in `sort_dedup_scan_tof`, so a ragged triplet + // would be an out-of-bounds panic in Rust rather than an error the caller can handle. + for i in 0..frame_ids.len() { + if mz[i].len() != mobility[i].len() || mz[i].len() != intensity[i].len() { + return Err(format!( + "frame {} (id {}): mz ({}), mobility ({}) and intensity ({}) must have the same length", + i, frame_ids[i], mz[i].len(), mobility[i].len(), intensity[i].len() + )); + } + } + + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(num_threads.max(1)) + .build() + .map_err(|e| format!("could not build the writer thread pool: {e}"))?; + + pool.install(|| { + (0..frame_ids.len()) + .into_par_iter() + .map(|i| { + let frame_id = frame_ids[i]; + let mz_values: Vec = mz[i].to_vec(); + let im_values: Vec = mobility[i].to_vec(); + + let tof = converter.mz_to_tof(frame_id, &mz_values); + let scan = converter.inverse_mobility_to_scan(frame_id, &im_values); + let raw_intensity: Vec = intensity[i].iter().map(|&x| x as u32).collect(); + + // m/z → TOF is not injective, so several peaks can land on one cell; sum them and + // sort by (scan, tof), which is the order the encoder's delta coding needs. + let (scan, tof, intensity) = sort_dedup_scan_tof(&scan, &tof, &raw_intensity); + + let num_peaks = intensity.len() as u32; + let max_intensity = intensity.iter().copied().max().unwrap_or(0); + let summed_intensity: u64 = intensity.iter().map(|&x| x as u64).sum(); + + let mut tof_delta = tof.clone(); + modify_tofs(&mut tof_delta, &scan); + let peak_cnts = get_peak_cnts(max_scans, &scan); + let interleaved: Vec = tof_delta + .iter() + .zip(intensity.iter()) + .flat_map(|(t, i)| [*t, *i]) + .collect(); + let real_data = get_realdata(&peak_cnts, &interleaved); + + // `zstd::bulk::compress` (ZSTD_compress) writes the decompressed size into the + // frame header; `encode_all` does not, and readers that use the simple API — + // including the Python `zstd` module every existing .d was written with — refuse + // a frame without it. Keep the header shape the format already has. + let data = zstd::bulk::compress(real_data.as_slice(), compression_level) + .map_err(|e| format!("frame id {frame_id}: zstd compression failed: {e}"))?; + + Ok(CompressedFrame { num_peaks, max_intensity, summed_intensity, data }) + }) + .collect() + }) +} diff --git a/rustdf/src/sim/dda.rs b/rustdf/src/sim/dda.rs index 0b6f3bbda..8f4c07283 100644 --- a/rustdf/src/sim/dda.rs +++ b/rustdf/src/sim/dda.rs @@ -1,3 +1,4 @@ +use mscore::simulation::noise_rng::noise_rng; use mscore::algorithm::isotope::calculate_precursor_transmission_factor; use crate::sim::containers::IsotopeTransmissionMode; use mscore::data::peptide::{PeptideIon, PeptideProductIonSeriesCollection}; @@ -35,9 +36,17 @@ pub struct TimsTofSyntheticsFrameBuilderDDA { pub fragment_ions_with_complementary: Option>, /// Physical instrument capabilities (P5e). Default = Bruker timsTOF. pub capabilities: InstrumentCapabilities, + /// Master seed for the simulation's noise. 0 keeps the legacy `thread_rng` behaviour + /// for API users that never set it; TimSim sets it from the run's `sample_seed`. + pub noise_seed: u64, } impl TimsTofSyntheticsFrameBuilderDDA { + + /// Set the master seed for the m/z jitter, making a run's noise reproducible. + pub fn set_noise_seed(&mut self, seed: u64) { + self.noise_seed = seed; + } /// Create a new DDA frame builder. /// /// # Arguments @@ -107,6 +116,7 @@ impl TimsTofSyntheticsFrameBuilderDDA { num_threads, )); Self { + noise_seed: 0, path: path.to_str().unwrap().to_string(), precursor_frame_builder: synthetics, transmission_settings, @@ -124,6 +134,7 @@ impl TimsTofSyntheticsFrameBuilderDDA { num_threads, )); Self { + noise_seed: 0, path: path.to_str().unwrap().to_string(), precursor_frame_builder: synthetics, transmission_settings, @@ -184,6 +195,7 @@ impl TimsTofSyntheticsFrameBuilderDDA { }; Self { + noise_seed: 0, path: String::new(), precursor_frame_builder, transmission_settings, @@ -576,6 +588,9 @@ impl TimsTofSyntheticsFrameBuilderDDA { intensity_min: Option, right_drag: Option, ) -> TimsFrame { + // One RNG per frame, keyed by the master seed and the frame id, so the m/z jitter + // does not depend on which thread builds the frame. See mscore::simulation::noise_rng. + let mut rng = noise_rng(self.noise_seed, &[frame_id as u64]); // Cache frame-level lookups let ms_type = if self.precursor_frame_builder.precursor_frame_id_set.contains(&frame_id) { MsType::Unknown @@ -745,9 +760,9 @@ impl TimsTofSyntheticsFrameBuilderDDA { let mz_spectrum = if mz_noise_fragment { if uniform { - final_spectrum.add_mz_noise_uniform(fragment_ppm, right_drag_val) + final_spectrum.add_mz_noise_uniform_with_rng(fragment_ppm, right_drag_val, &mut rng) } else { - final_spectrum.add_mz_noise_normal(fragment_ppm) + final_spectrum.add_mz_noise_normal_with_rng(fragment_ppm, &mut rng) } } else { final_spectrum @@ -790,9 +805,9 @@ impl TimsTofSyntheticsFrameBuilderDDA { let precursor_mz_spectrum = if mz_noise_fragment { if uniform { - precursor_scaled.add_mz_noise_uniform(fragment_ppm, right_drag_val) + precursor_scaled.add_mz_noise_uniform_with_rng(fragment_ppm, right_drag_val, &mut rng) } else { - precursor_scaled.add_mz_noise_normal(fragment_ppm) + precursor_scaled.add_mz_noise_normal_with_rng(fragment_ppm, &mut rng) } } else { precursor_scaled @@ -850,6 +865,9 @@ impl TimsTofSyntheticsFrameBuilderDDA { intensity_min: Option, right_drag: Option, ) -> TimsFrameAnnotated { + // One RNG per frame, keyed by the master seed and the frame id, so the m/z jitter + // does not depend on which thread builds the frame. See mscore::simulation::noise_rng. + let mut rng = noise_rng(self.noise_seed, &[frame_id as u64]); // Cache frame-level lookups let ms_type = if self.precursor_frame_builder.precursor_frame_id_set.contains(&frame_id) { MsType::Unknown @@ -979,9 +997,9 @@ impl TimsTofSyntheticsFrameBuilderDDA { let mz_spectrum = if mz_noise_fragment { if uniform { - scaled_spec.add_mz_noise_uniform(fragment_ppm, right_drag_val) + scaled_spec.add_mz_noise_uniform_with_rng(fragment_ppm, right_drag_val, &mut rng) } else { - scaled_spec.add_mz_noise_normal(fragment_ppm) + scaled_spec.add_mz_noise_normal_with_rng(fragment_ppm, &mut rng) } } else { scaled_spec @@ -1028,9 +1046,9 @@ impl TimsTofSyntheticsFrameBuilderDDA { let precursor_final = if mz_noise_fragment { if uniform { - precursor_scaled.add_mz_noise_uniform(fragment_ppm, right_drag_val) + precursor_scaled.add_mz_noise_uniform_with_rng(fragment_ppm, right_drag_val, &mut rng) } else { - precursor_scaled.add_mz_noise_normal(fragment_ppm) + precursor_scaled.add_mz_noise_normal_with_rng(fragment_ppm, &mut rng) } } else { precursor_scaled diff --git a/rustdf/src/sim/dia.rs b/rustdf/src/sim/dia.rs index 2511fbea5..5535a7698 100644 --- a/rustdf/src/sim/dia.rs +++ b/rustdf/src/sim/dia.rs @@ -1,3 +1,4 @@ +use mscore::simulation::noise_rng::noise_rng; use mscore::algorithm::isotope::{ calculate_precursor_transmission_factor, calculate_transmission_dependent_fragment_ion_isotope_distribution, @@ -103,9 +104,17 @@ pub struct TimsTofSyntheticsFrameBuilderDIA { Option>, /// Physical instrument capabilities (P5e). Default = Bruker timsTOF. pub capabilities: InstrumentCapabilities, + /// Master seed for the simulation's noise. 0 keeps the legacy `thread_rng` behaviour + /// for API users that never set it; TimSim sets it from the run's `sample_seed`. + pub noise_seed: u64, } impl TimsTofSyntheticsFrameBuilderDIA { + + /// Set the master seed for the m/z jitter, making a run's noise reproducible. + pub fn set_noise_seed(&mut self, seed: u64) { + self.noise_seed = seed; + } pub fn new(path: &Path, with_annotations: bool, num_threads: usize) -> rusqlite::Result { Self::new_with_config(path, with_annotations, num_threads, IsotopeTransmissionConfig::default()) } @@ -176,6 +185,7 @@ impl TimsTofSyntheticsFrameBuilderDIA { num_threads, )); Ok(Self { + noise_seed: 0, path: path.to_str().unwrap().to_string(), precursor_frame_builder: synthetics, transmission_settings, @@ -195,6 +205,7 @@ impl TimsTofSyntheticsFrameBuilderDIA { num_threads, )); Ok(Self { + noise_seed: 0, path: path.to_str().unwrap().to_string(), precursor_frame_builder: synthetics, transmission_settings, @@ -587,6 +598,9 @@ impl TimsTofSyntheticsFrameBuilderDIA { intensity_min: Option, right_drag: Option, ) -> TimsFrame { + // One RNG per frame, keyed by the master seed and the frame id, so the m/z jitter + // does not depend on which thread builds the frame. See mscore::simulation::noise_rng. + let mut rng = noise_rng(self.noise_seed, &[frame_id as u64]); // Cache frame-level lookups once let ms_type = if self.precursor_frame_builder.precursor_frame_id_set.contains(&frame_id) { MsType::Unknown @@ -730,9 +744,9 @@ impl TimsTofSyntheticsFrameBuilderDIA { let mz_spectrum = if mz_noise_fragment { if uniform { - final_spectrum.add_mz_noise_uniform(fragment_ppm, right_drag_val) + final_spectrum.add_mz_noise_uniform_with_rng(fragment_ppm, right_drag_val, &mut rng) } else { - final_spectrum.add_mz_noise_normal(fragment_ppm) + final_spectrum.add_mz_noise_normal_with_rng(fragment_ppm, &mut rng) } } else { final_spectrum @@ -775,9 +789,9 @@ impl TimsTofSyntheticsFrameBuilderDIA { let precursor_mz_spectrum = if mz_noise_fragment { if uniform { - precursor_scaled.add_mz_noise_uniform(fragment_ppm, right_drag_val) + precursor_scaled.add_mz_noise_uniform_with_rng(fragment_ppm, right_drag_val, &mut rng) } else { - precursor_scaled.add_mz_noise_normal(fragment_ppm) + precursor_scaled.add_mz_noise_normal_with_rng(fragment_ppm, &mut rng) } } else { precursor_scaled @@ -956,6 +970,9 @@ impl TimsTofSyntheticsFrameBuilderDIA { intensity_min: Option, right_drag: Option, ) -> TimsFrameAnnotated { + // One RNG per frame, keyed by the master seed and the frame id, so the m/z jitter + // does not depend on which thread builds the frame. See mscore::simulation::noise_rng. + let mut rng = noise_rng(self.noise_seed, &[frame_id as u64]); // Cache frame-level lookups let ms_type = if self.precursor_frame_builder.precursor_frame_id_set.contains(&frame_id) { MsType::Unknown @@ -1049,9 +1066,9 @@ impl TimsTofSyntheticsFrameBuilderDIA { let mz_spectrum = if mz_noise_fragment { if uniform { - scaled_spec.add_mz_noise_uniform(fragment_ppm, right_drag_val) + scaled_spec.add_mz_noise_uniform_with_rng(fragment_ppm, right_drag_val, &mut rng) } else { - scaled_spec.add_mz_noise_normal(fragment_ppm) + scaled_spec.add_mz_noise_normal_with_rng(fragment_ppm, &mut rng) } } else { scaled_spec @@ -1095,9 +1112,9 @@ impl TimsTofSyntheticsFrameBuilderDIA { let precursor_final = if mz_noise_fragment { if uniform { - precursor_scaled.add_mz_noise_uniform(fragment_ppm, right_drag_val) + precursor_scaled.add_mz_noise_uniform_with_rng(fragment_ppm, right_drag_val, &mut rng) } else { - precursor_scaled.add_mz_noise_normal(fragment_ppm) + precursor_scaled.add_mz_noise_normal_with_rng(fragment_ppm, &mut rng) } } else { precursor_scaled @@ -1352,6 +1369,7 @@ mod p6a_fragment_kernel_tests { frame_to_rt.insert(2u32, 30.0f32); let precursor = TimsTofSyntheticsPrecursorFrameBuilder { + noise_seed: 0, ions: BTreeMap::new(), peptides: BTreeMap::new(), scans: Vec::new(), @@ -1375,6 +1393,7 @@ mod p6a_fragment_kernel_tests { ); TimsTofSyntheticsFrameBuilderDIA { + noise_seed: 0, path: String::new(), precursor_frame_builder: precursor, transmission_settings: TimsTransmissionDIA::new( diff --git a/rustdf/src/sim/lazy_builder.rs b/rustdf/src/sim/lazy_builder.rs index 81da79e57..1e51771d8 100644 --- a/rustdf/src/sim/lazy_builder.rs +++ b/rustdf/src/sim/lazy_builder.rs @@ -4,6 +4,7 @@ //! memory-efficient alternatives to their non-lazy counterparts that only load //! peptide/ion data for the frames being built rather than loading everything upfront. +use mscore::simulation::noise_rng::noise_rng; use mscore::data::peptide::PeptideProductIonSeriesCollection; use mscore::data::spectrum::{IndexedMzSpectrum, MsType, MzSpectrum}; use mscore::timstof::collision::{TimsTofCollisionEnergy, TimsTofCollisionEnergyDIA}; @@ -51,9 +52,17 @@ pub struct TimsTofLazyFrameBuilderDIA { pub num_threads: usize, /// Source for occurrence/abundance distributions (legacy columns or projector) pub source: DistributionSource, + /// Master seed for the simulation's noise. 0 keeps the legacy `thread_rng` behaviour + /// for API users that never set it; TimSim sets it from the run's `sample_seed`. + pub noise_seed: u64, } impl TimsTofLazyFrameBuilderDIA { + + /// Set the master seed for the m/z jitter, making a run's noise reproducible. + pub fn set_noise_seed(&mut self, seed: u64) { + self.noise_seed = seed; + } /// Create a new lazy frame builder. /// /// Only loads static metadata (frames, scans, transmission settings). @@ -96,6 +105,7 @@ impl TimsTofLazyFrameBuilderDIA { let fragmentation_settings = handle.get_collision_energy_dia(); Ok(Self { + noise_seed: 0, db_path: path.to_str().unwrap().to_string(), frames, scans, @@ -295,6 +305,9 @@ impl TimsTofLazyFrameBuilderDIA { frame_to_abundances: &BTreeMap, Vec)>, peptide_to_events: &BTreeMap, ) -> TimsFrame { + // One RNG per frame, keyed by the master seed and the frame id, so the m/z jitter + // does not depend on which thread builds the frame. See mscore::simulation::noise_rng. + let mut rng = noise_rng(self.noise_seed, &[frame_id as u64]); let ms_type = MsType::Precursor; let rt = *self.frame_to_rt.get(&frame_id).unwrap_or(&0.0) as f64; @@ -328,9 +341,9 @@ impl TimsTofLazyFrameBuilderDIA { let mz_spectrum = if mz_noise_precursor { if uniform { - scaled_spec.add_mz_noise_uniform(precursor_noise_ppm, right_drag) + scaled_spec.add_mz_noise_uniform_with_rng(precursor_noise_ppm, right_drag, &mut rng) } else { - scaled_spec.add_mz_noise_normal(precursor_noise_ppm) + scaled_spec.add_mz_noise_normal_with_rng(precursor_noise_ppm, &mut rng) } } else { scaled_spec @@ -385,6 +398,9 @@ impl TimsTofLazyFrameBuilderDIA { peptide_to_events: &BTreeMap, fragment_ions_map: &Option)>>, ) -> TimsFrame { + // One RNG per frame, keyed by the master seed and the frame id, so the m/z jitter + // does not depend on which thread builds the frame. See mscore::simulation::noise_rng. + let mut rng = noise_rng(self.noise_seed, &[frame_id as u64]); let ms_type = MsType::FragmentDia; let rt = *self.frame_to_rt.get(&frame_id).unwrap_or(&0.0) as f64; @@ -478,9 +494,9 @@ impl TimsTofLazyFrameBuilderDIA { let mz_spectrum = if mz_noise_fragment { if uniform { - scaled_spec.add_mz_noise_uniform(fragment_noise_ppm, right_drag) + scaled_spec.add_mz_noise_uniform_with_rng(fragment_noise_ppm, right_drag, &mut rng) } else { - scaled_spec.add_mz_noise_normal(fragment_noise_ppm) + scaled_spec.add_mz_noise_normal_with_rng(fragment_noise_ppm, &mut rng) } } else { scaled_spec diff --git a/rustdf/src/sim/precursor.rs b/rustdf/src/sim/precursor.rs index e33a02e58..19e9055cf 100644 --- a/rustdf/src/sim/precursor.rs +++ b/rustdf/src/sim/precursor.rs @@ -1,3 +1,4 @@ +use mscore::simulation::noise_rng::noise_rng; use mscore::data::peptide::PeptideIon; use mscore::data::spectrum::{IndexedMzSpectrum, MsType, MzSpectrum}; use mscore::simulation::annotation::{ @@ -37,9 +38,17 @@ pub struct TimsTofSyntheticsPrecursorFrameBuilder { pub peptide_to_events: BTreeMap, /// Mapping from ion_id to (peptide_id, charge) for DDA precursor lookup pub ion_id_to_peptide_charge: BTreeMap, + /// Master seed for the simulation's noise. 0 keeps the legacy `thread_rng` behaviour + /// for API users that never set it; TimSim sets it from the run's `sample_seed`. + pub noise_seed: u64, } impl TimsTofSyntheticsPrecursorFrameBuilder { + + /// Set the master seed for the m/z jitter, making a run's noise reproducible. + pub fn set_noise_seed(&mut self, seed: u64) { + self.noise_seed = seed; + } /// Create a new instance of TimsTofSynthetics /// /// # Arguments @@ -90,6 +99,7 @@ impl TimsTofSyntheticsPrecursorFrameBuilder { } Self { + noise_seed: 0, ions: TimsTofSyntheticsDataHandle::build_peptide_to_ion_map(&ions), peptides: TimsTofSyntheticsDataHandle::build_peptide_map(&peptides), scans: scans.clone(), @@ -245,6 +255,9 @@ impl TimsTofSyntheticsPrecursorFrameBuilder { precursor_noise_ppm: f64, right_drag: bool, ) -> TimsFrame { + // One RNG per frame, keyed by the master seed and the frame id, so the m/z jitter + // does not depend on which thread builds the frame. See mscore::simulation::noise_rng. + let mut rng = noise_rng(self.noise_seed, &[frame_id as u64]); // Cache frame-level lookups let ms_type = if self.precursor_frame_id_set.contains(&frame_id) { MsType::Precursor @@ -262,9 +275,9 @@ impl TimsTofSyntheticsPrecursorFrameBuilder { for (scan, scaled_spec) in contributions { let mz_spectrum = if mz_noise_precursor { if uniform { - scaled_spec.add_mz_noise_uniform(precursor_noise_ppm, right_drag) + scaled_spec.add_mz_noise_uniform_with_rng(precursor_noise_ppm, right_drag, &mut rng) } else { - scaled_spec.add_mz_noise_normal(precursor_noise_ppm) + scaled_spec.add_mz_noise_normal_with_rng(precursor_noise_ppm, &mut rng) } } else { scaled_spec @@ -353,6 +366,9 @@ impl TimsTofSyntheticsPrecursorFrameBuilder { precursor_noise_ppm: f64, right_drag: bool, ) -> TimsFrameAnnotated { + // One RNG per frame, keyed by the master seed and the frame id, so the m/z jitter + // does not depend on which thread builds the frame. See mscore::simulation::noise_rng. + let mut rng = noise_rng(self.noise_seed, &[frame_id as u64]); // Cache frame-level lookups let ms_type = if self.precursor_frame_id_set.contains(&frame_id) { MsType::Precursor @@ -402,9 +418,9 @@ impl TimsTofSyntheticsPrecursorFrameBuilder { let mz_spectrum = if mz_noise_precursor { if uniform { - scaled_spec.add_mz_noise_uniform(precursor_noise_ppm, right_drag) + scaled_spec.add_mz_noise_uniform_with_rng(precursor_noise_ppm, right_drag, &mut rng) } else { - scaled_spec.add_mz_noise_normal(precursor_noise_ppm) + scaled_spec.add_mz_noise_normal_with_rng(precursor_noise_ppm, &mut rng) } } else { scaled_spec @@ -543,6 +559,7 @@ mod tests { frame_to_rt.insert(1u32, 60.0f32); TimsTofSyntheticsPrecursorFrameBuilder { + noise_seed: 0, ions: BTreeMap::new(), peptides: BTreeMap::new(), scans: Vec::new(),