rustdf/sim/
library.rs

1//! Fast spectral-library writer: decode Prosit intensity vectors into a DiaNN/
2//! Spectronaut-style transition `.tsv` in parallel.
3//!
4//! The per-peptide decode (`PeptideSequence::associate_with_predicted_intensities`)
5//! is already a Rust kernel in `mscore`; the slow part of a Python library builder is
6//! calling it millions of times across the PyO3 boundary and assembling a giant
7//! DataFrame. This does the whole batch in Rust with `rayon` and streams the rows
8//! straight to disk — turning a multi-hour build into seconds.
9//!
10//! The caller (Python) supplies the cheap string columns it already has
11//! (`ModifiedPeptide` in DiaNN `(UniMod:n)` form, the stripped `PeptideSequence`) so
12//! this module needs no mod-notation regex; ordinals come from the fragment ion's own
13//! `amino_acid_count()`.
14
15use std::fs::OpenOptions;
16use std::io::{self, BufWriter, Write};
17use std::path::Path;
18
19use mscore::data::peptide::{FragmentType, PeptideSequence};
20use rayon::prelude::*;
21
22/// DiaNN/Spectronaut transition columns, in output order.
23pub const LIBRARY_HEADER: &str = "ProteinId\tGenes\tModifiedPeptide\tPeptideSequence\t\
24PrecursorCharge\tPrecursorMz\tTr_recalibrated\tFragmentType\tFragmentCharge\t\
25FragmentSeriesNumber\tProductMz\tLibraryIntensity";
26
27/// Prosit's fragment array holds 29 positions = a 30-residue peptide's b/y ions; a
28/// longer peptide overruns it and the decode kernel panics.
29const PROSIT_MAX_RESIDUES: usize = 30;
30
31/// Count backbone residues in a UNIMOD-bracket modified sequence, ignoring the
32/// `[UNIMOD:n]` tokens (and any other bracketed annotation). Cheap, allocation-free —
33/// used to skip over-length peptides before the panic-prone decode kernel.
34fn stripped_residue_count(modified_sequence: &str) -> usize {
35    let mut count = 0usize;
36    let mut in_bracket = false;
37    for c in modified_sequence.bytes() {
38        match c {
39            b'[' => in_bracket = true,
40            b']' => in_bracket = false,
41            b'A'..=b'Z' if !in_bracket => count += 1,
42            _ => {}
43        }
44    }
45    count
46}
47
48/// Decode one batch of peptides (parallel) and write/append their transitions to
49/// `out_path` as a DiaNN-style `.tsv`.
50///
51/// Per precursor: b/y ions are read off the reshaped Prosit array, intensities are
52/// renormalised to max=1, non-finite / non-positive intensities and m/z are dropped,
53/// and a precursor with fewer than `min_fragments` surviving fragments is skipped
54/// entirely (it emits no rows). Returns `(transitions_written, precursors_written)`.
55///
56/// `prosit_flat` is row-major: peptide `i`'s vector is `prosit_flat[i*n_cols..(i+1)*n_cols]`.
57/// All metadata slices must have the same length (the peptide count); `n_cols` is the
58/// Prosit vector width (174). With `append=false` the file is truncated and the header
59/// written first; with `append=true` rows are appended (no header) — for chunked builds.
60#[allow(clippy::too_many_arguments)]
61pub fn build_spectral_library_tsv(
62    out_path: &Path,
63    modified_sequences: &[String],
64    modpep_diann: &[String],
65    stripped: &[String],
66    charges: &[i32],
67    prosit_flat: &[f64],
68    n_cols: usize,
69    precursor_mz: &[f64],
70    retention_time: &[f64],
71    protein_ids: &[String],
72    genes: &[String],
73    min_fragments: usize,
74    append: bool,
75) -> io::Result<(usize, usize)> {
76    let n = modified_sequences.len();
77    for (name, len) in [
78        ("modpep_diann", modpep_diann.len()),
79        ("stripped", stripped.len()),
80        ("charges", charges.len()),
81        ("precursor_mz", precursor_mz.len()),
82        ("retention_time", retention_time.len()),
83        ("protein_ids", protein_ids.len()),
84        ("genes", genes.len()),
85    ] {
86        if len != n {
87            return Err(io::Error::new(
88                io::ErrorKind::InvalidInput,
89                format!("column '{name}' length {len} != peptide count {n}"),
90            ));
91        }
92    }
93    // The Prosit intensity vector reshapes to 29 fragment positions x 2 termini x 3
94    // charges = 174; any other width makes `reshape_prosit_array` panic.
95    const PROSIT_VEC_LEN: usize = 174;
96    if n_cols != PROSIT_VEC_LEN {
97        return Err(io::Error::new(
98            io::ErrorKind::InvalidInput,
99            format!("n_cols must be {PROSIT_VEC_LEN} (Prosit vector width), got {n_cols}"),
100        ));
101    }
102    if prosit_flat.len() != n * n_cols {
103        return Err(io::Error::new(
104            io::ErrorKind::InvalidInput,
105            format!(
106                "prosit_flat length {} != peptides {} * n_cols {}",
107                prosit_flat.len(),
108                n,
109                n_cols
110            ),
111        ));
112    }
113
114    // Decode each peptide independently; emit a String block of tsv rows (empty if the
115    // precursor is skipped). Pure per-item work -> embarrassingly parallel. The decode
116    // kernel PANICS on peptides >30 residues (they overrun Prosit's fixed 29-position
117    // fragment array). The workspace builds with panic="abort", so catch_unwind cannot
118    // recover — we PREVENT the panic by skipping over-length (and empty) peptides up
119    // front. (Inputs are validated UNIMOD sequences from the DB/digest, so malformed
120    // sequences are not expected here.)
121    let blocks: Vec<String> = (0..n)
122        .into_par_iter()
123        .map(|i| {
124            let n_res = stripped_residue_count(&modified_sequences[i]);
125            if n_res == 0 || n_res > PROSIT_MAX_RESIDUES {
126                return String::new();
127            }
128            let intensities = prosit_flat[i * n_cols..(i + 1) * n_cols].to_vec();
129            let pep = PeptideSequence::new(modified_sequences[i].clone(), None);
130            let coll = pep.associate_with_predicted_intensities(
131                charges[i],
132                FragmentType::B,
133                intensities,
134                true,
135                true,
136            );
137            // Gather surviving fragments and the per-precursor max for renormalisation.
138            let mut frags: Vec<(String, i32, usize, f64, f64)> = Vec::new();
139            let mut max_i = 0.0_f64;
140            for series in &coll.peptide_ions {
141                for ion in series.n_ions.iter().chain(series.c_ions.iter()) {
142                    let inten = ion.ion.intensity;
143                    let mz = ion.mz();
144                    if !inten.is_finite() || inten <= 0.0 || !mz.is_finite() {
145                        continue;
146                    }
147                    let ordinal = ion.ion.sequence.amino_acid_count();
148                    if inten > max_i {
149                        max_i = inten;
150                    }
151                    frags.push((ion.kind.to_string(), ion.ion.charge, ordinal, mz, inten));
152                }
153            }
154            if frags.len() < min_fragments || max_i <= 0.0 {
155                return String::new();
156            }
157            let mut block = String::with_capacity(frags.len() * 80);
158            for (kind, fch, ord, mz, inten) in &frags {
159                // ProteinId Genes ModifiedPeptide PeptideSequence PrecursorCharge
160                // PrecursorMz Tr FragmentType FragmentCharge FragmentSeriesNumber
161                // ProductMz LibraryIntensity
162                block.push_str(&format!(
163                    "{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\n",
164                    protein_ids[i],
165                    genes[i],
166                    modpep_diann[i],
167                    stripped[i],
168                    charges[i],
169                    precursor_mz[i],
170                    retention_time[i],
171                    kind,
172                    fch,
173                    ord,
174                    mz,
175                    inten / max_i,
176                ));
177            }
178            block
179        })
180        .collect();
181
182    let file = OpenOptions::new()
183        .write(true)
184        .create(true)
185        .append(append)
186        .truncate(!append)
187        .open(out_path)?;
188    let mut w = BufWriter::new(file);
189    if !append {
190        writeln!(w, "{LIBRARY_HEADER}")?;
191    }
192    let (mut n_trans, mut n_prec) = (0usize, 0usize);
193    for block in &blocks {
194        if block.is_empty() {
195            continue;
196        }
197        n_prec += 1;
198        n_trans += block.as_bytes().iter().filter(|&&b| b == b'\n').count();
199        w.write_all(block.as_bytes())?;
200    }
201    w.flush()?;
202    Ok((n_trans, n_prec))
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    fn one(s: &str) -> Vec<String> {
210        vec![s.to_string()]
211    }
212
213    #[test]
214    fn stripped_residue_count_ignores_mods() {
215        assert_eq!(stripped_residue_count("PEPTIDEK"), 8);
216        assert_eq!(stripped_residue_count("M[UNIMOD:35]PEPK"), 5);
217        assert_eq!(stripped_residue_count("[UNIMOD:1]ACDEK"), 5);
218        assert_eq!(stripped_residue_count("C[UNIMOD:4]C[UNIMOD:4]K"), 3);
219    }
220
221    #[test]
222    fn over_length_peptide_is_skipped_not_panicked() {
223        // A 35-residue peptide overruns Prosit's 29-position array; the decode kernel
224        // would panic (and panic=abort kills the process), so it must be skipped BEFORE
225        // the kernel. Build with one over-length peptide -> header only, no crash.
226        let out = std::env::temp_dir().join("rustdf_lib_overlen.tsv");
227        let long = "ACDEFGHIKLMNPQRSTVWYACDEFGHIKLMNPQR"; // 35 residues
228        assert_eq!(stripped_residue_count(long), 35);
229        let (n_trans, n_prec) = build_spectral_library_tsv(
230            &out, &one(long), &one(long), &one(long), &[2],
231            &vec![0.1; 174], 174, &[700.0], &[10.0], &one("P"), &one("G"), 3, false,
232        )
233        .expect("must not panic on over-length peptide");
234        assert_eq!((n_trans, n_prec), (0, 0), "over-length peptide skipped");
235    }
236
237    #[test]
238    fn rejects_mismatched_column_lengths() {
239        let dir = std::env::temp_dir().join("rustdf_lib_badcols.tsv");
240        // charges has 2 entries but everything else has 1 -> error, no file contract.
241        let r = build_spectral_library_tsv(
242            &dir,
243            &one("PEPTIDEK"),
244            &one("PEPTIDEK"),
245            &one("PEPTIDEK"),
246            &[2, 3],
247            &vec![0.0; 174],
248            174,
249            &[500.0],
250            &[10.0],
251            &one("P"),
252            &one("G"),
253            3,
254            false,
255        );
256        assert!(r.is_err(), "mismatched column lengths must error");
257    }
258
259    #[test]
260    fn rejects_wrong_prosit_flat_size() {
261        let dir = std::env::temp_dir().join("rustdf_lib_badflat.tsv");
262        let r = build_spectral_library_tsv(
263            &dir, &one("PEPK"), &one("PEPK"), &one("PEPK"), &[2],
264            &vec![0.0; 100], 174, &[500.0], &[10.0], &one("P"), &one("G"), 3, false,
265        );
266        assert!(r.is_err(), "prosit_flat length != n*n_cols must error");
267    }
268
269    #[test]
270    fn all_zero_intensities_emit_header_only() {
271        // A zero Prosit vector -> every fragment intensity 0 -> precursor skipped
272        // (< min_fragments). Replace mode still writes the header and an empty body.
273        let out = std::env::temp_dir().join("rustdf_lib_zero.tsv");
274        let (n_trans, n_prec) = build_spectral_library_tsv(
275            &out, &one("PEPTIDEK"), &one("PEPTIDEK"), &one("PEPTIDEK"), &[2],
276            &vec![0.0; 174], 174, &[500.0], &[10.0], &one("P"), &one("G"), 3, false,
277        )
278        .expect("write");
279        assert_eq!((n_trans, n_prec), (0, 0));
280        let body = std::fs::read_to_string(&out).expect("read");
281        assert!(body.starts_with("ProteinId\t"), "header present");
282        assert_eq!(body.lines().count(), 1, "header only, no transition rows");
283    }
284}
rustdf/sim/library.rs

rustdf/sim/
library.rs
