1use rayon::prelude::*;
23use std::cmp::Ordering;
24use std::collections::{HashMap, HashSet};
25use std::sync::Arc;
26
27const H2O_F32: f32 = 18.010565;
29
30const MONOISOTOPIC_F32: [f32; 26] = [
32 71.03711, 0.0, 103.00919, 115.02694, 129.04259, 147.0684, 57.02146, 137.05891, 113.08406, 0.0,
33 128.09496, 113.08406, 131.0405, 114.04293, 237.14774, 97.05276, 128.05858, 156.1011, 87.03203,
34 101.04768, 150.95363, 99.06841, 186.07932, 0.0, 163.06332, 0.0,
35];
36
37const VALID_AA: &[u8] = b"ACDEFGHIKLMNPQRSTVWYUO";
38
39fn monoisotopic_f32(aa: u8) -> f32 {
40 if aa.is_ascii_uppercase() {
41 MONOISOTOPIC_F32[(aa - b'A') as usize]
42 } else {
43 0.0
44 }
45}
46
47#[derive(Clone)]
49pub struct Fasta {
50 pub targets: Vec<(Arc<str>, String)>,
52 decoy_tag: String,
53 generate_decoys: bool,
54}
55
56impl Fasta {
57 pub fn parse(contents: &str, decoy_tag: &str, generate_decoys: bool) -> anyhow::Result<Fasta> {
60 let mut targets = Vec::new();
61 let mut last_id = "";
62 let mut s = String::new();
63 let push =
64 |id: &str, seq: String, targets: &mut Vec<(Arc<str>, String)>| -> anyhow::Result<()> {
65 let acc = id
66 .split_ascii_whitespace()
67 .next()
68 .ok_or_else(|| anyhow::anyhow!("FASTA entry with an empty header"))?;
69 let acc: Arc<str> = Arc::from(acc);
70 if !acc.contains(decoy_tag) || !generate_decoys {
71 targets.push((acc, seq));
72 }
73 Ok(())
74 };
75 for line in contents.lines() {
76 if line.is_empty() {
77 continue;
78 }
79 let line = line.trim();
80 if let Some(id) = line.strip_prefix('>') {
81 if !s.is_empty() {
82 push(last_id, std::mem::take(&mut s), &mut targets)?;
83 }
84 last_id = id;
85 } else {
86 s.push_str(line);
87 }
88 }
89 if !s.is_empty() {
90 push(last_id, s, &mut targets)?;
91 }
92 Ok(Fasta {
93 targets,
94 decoy_tag: decoy_tag.to_string(),
95 generate_decoys,
96 })
97 }
98
99 fn digest(&self, enzyme: &EnzymeParameters) -> Vec<Digest> {
100 self.targets
101 .par_iter()
102 .flat_map_iter(|(protein, sequence)| {
103 let is_tagged = protein.contains(self.decoy_tag.as_str());
104 let generate = self.generate_decoys;
105 enzyme
106 .digest(sequence, protein.clone())
107 .into_iter()
108 .filter_map(move |mut d| {
109 if is_tagged {
110 if generate {
111 None
112 } else {
113 d.decoy = true;
114 Some(d)
115 }
116 } else {
117 Some(d)
118 }
119 })
120 })
121 .collect()
122 }
123}
124
125#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash, Default)]
127pub enum Position {
128 Nterm,
129 Cterm,
130 Full,
131 #[default]
132 Internal,
133}
134
135#[derive(Clone, Debug)]
136struct Digest {
137 decoy: bool,
138 sequence: String,
139 protein: Arc<str>,
140 missed_cleavages: u8,
141 position: Position,
142}
143
144struct DigestGroup {
145 reference: Digest,
146 proteins: Vec<Arc<str>>,
147}
148
149fn group_digests(mut digests: Vec<Digest>) -> Vec<DigestGroup> {
150 let mut groups = Vec::new();
151 if digests.is_empty() {
152 return groups;
153 }
154 digests.sort_unstable_by(|a, b| {
155 a.position
156 .cmp(&b.position)
157 .then(a.decoy.cmp(&b.decoy))
158 .then(a.sequence.cmp(&b.sequence))
159 });
160 let mut curr = DigestGroup {
161 reference: digests[0].clone(),
162 proteins: Vec::new(),
163 };
164 for d in digests {
165 if d.decoy == curr.reference.decoy
166 && d.position == curr.reference.position
167 && d.sequence == curr.reference.sequence
168 {
169 curr.proteins.push(d.protein);
170 } else {
171 curr.proteins.sort_unstable();
172 groups.push(curr);
173 curr = DigestGroup {
174 proteins: vec![d.protein.clone()],
175 reference: d,
176 };
177 }
178 }
179 groups.push(curr);
180 groups
181}
182
183struct Enzyme {
184 cleave: Option<[bool; 256]>,
186 skip_suffix: [bool; 256],
187 c_terminal: bool,
188}
189
190struct EnzymeParameters {
191 missed_cleavages: u8,
192 min_len: usize,
193 max_len: usize,
194 enzyme: Option<Enzyme>,
196}
197
198#[derive(Clone)]
199struct Site {
200 range: std::ops::Range<usize>,
201 missed_cleavages: u8,
202}
203
204impl Enzyme {
205 fn new(cleave: &str, skip_suffix: &str, c_terminal: bool) -> anyhow::Result<Option<Self>> {
206 anyhow::ensure!(
207 cleave == "$" || cleave.bytes().all(|b| VALID_AA.contains(&b) || b == b'?'),
208 "enzyme cleavage sequence contains non-amino-acid characters: {cleave}"
209 );
210 anyhow::ensure!(
211 skip_suffix.bytes().all(|b| VALID_AA.contains(&b)),
212 "enzyme cleavage restriction contains non-amino-acid characters: {skip_suffix}"
213 );
214 let mut skip = [false; 256];
215 for b in skip_suffix.bytes() {
216 skip[b as usize] = true;
217 }
218 Ok(match cleave.replace('?', "").as_str() {
219 "" if cleave.is_empty() => None,
220 "$" => Some(Enzyme {
221 cleave: None,
222 skip_suffix: [false; 256],
223 c_terminal: true,
224 }),
225 set => {
226 let mut arr = [false; 256];
227 for b in set.bytes() {
228 arr[b as usize] = true;
229 }
230 Some(Enzyme {
231 cleave: Some(arr),
232 skip_suffix: skip,
233 c_terminal,
234 })
235 }
236 })
237 }
238
239 fn cleavage_sites(&self, sequence: &str) -> Vec<Site> {
240 let bytes = sequence.as_bytes();
241 let mut sites = Vec::new();
242 let mut left = 0;
243 let matches: Vec<(usize, usize)> = match &self.cleave {
244 None => vec![(bytes.len(), bytes.len())],
245 Some(set) => bytes
246 .iter()
247 .enumerate()
248 .filter(|(_, b)| set[**b as usize])
249 .map(|(i, _)| (i, i + 1))
250 .collect(),
251 };
252 for (start, end) in matches {
253 let right = if self.c_terminal { end } else { start };
254 if bytes
255 .get(right)
256 .is_some_and(|b| self.skip_suffix[*b as usize])
257 {
258 continue;
259 }
260 sites.push(Site {
261 range: left..right,
262 missed_cleavages: 0,
263 });
264 left = right;
265 }
266 sites.push(Site {
267 range: left..bytes.len(),
268 missed_cleavages: 0,
269 });
270 sites
271 }
272}
273
274impl EnzymeParameters {
275 fn cleavage_sites(&self, sequence: &str) -> Vec<Site> {
276 match &self.enzyme {
277 Some(e) => e.cleavage_sites(sequence),
278 None => {
279 let mut v = Vec::new();
280 for len in self.min_len..=self.max_len {
281 for i in 0..=sequence.len().saturating_sub(len) {
282 v.push(Site {
283 range: i..i + len,
284 missed_cleavages: 0,
285 });
286 }
287 }
288 v
289 }
290 }
291 }
292
293 fn digest(&self, sequence: &str, protein: Arc<str>) -> Vec<Digest> {
294 let n = sequence.len();
295 let mut sites = self.cleavage_sites(sequence);
296 let missed = if self.enzyme.is_none() {
297 0
298 } else {
299 self.missed_cleavages
300 };
301 if missed > 0 {
302 let mut extra = Vec::new();
303 for cleavage in 1..=(1 + missed) {
304 for win in sites.windows(cleavage as usize) {
305 extra.push(Site {
306 range: win[0].range.start..win[cleavage as usize - 1].range.end,
307 missed_cleavages: cleavage - 1,
308 });
309 }
310 }
311 sites.append(&mut extra);
312 }
313 let mut seen: HashSet<&str> = HashSet::new();
314 let mut digests = Vec::new();
315 for site in &sites {
316 let (start, end) = (site.range.start, site.range.end);
317 let Some(pep) = sequence.get(start..end) else {
318 continue;
319 };
320 let len = pep.len();
321 let position = match (start == 0, end == n) {
322 (true, true) => Position::Full,
323 (true, false) => Position::Nterm,
324 (false, true) => Position::Cterm,
325 (false, false) => Position::Internal,
326 };
327 if len >= self.min_len && len <= self.max_len && len > 0 && seen.insert(pep) {
328 digests.push(Digest {
329 decoy: false,
330 sequence: pep.to_string(),
331 protein: protein.clone(),
332 missed_cleavages: site.missed_cleavages,
333 position,
334 });
335 }
336 }
337 digests
338 }
339}
340
341#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
343pub enum ModSpecificity {
344 PeptideN(Option<u8>),
345 PeptideC(Option<u8>),
346 ProteinN(Option<u8>),
347 ProteinC(Option<u8>),
348 Residue(u8),
349}
350
351impl ModSpecificity {
352 pub fn parse(s: &str) -> Option<Self> {
355 if s.len() > 2 {
356 return None;
357 }
358 let rest = |r: &str| r.chars().next().map(|c| c as u8);
359 if let Some(r) = s.strip_prefix('^') {
360 return Some(Self::PeptideN(rest(r)));
361 }
362 if let Some(r) = s.strip_prefix('$') {
363 return Some(Self::PeptideC(rest(r)));
364 }
365 if let Some(r) = s.strip_prefix('[') {
366 return Some(Self::ProteinN(rest(r)));
367 }
368 if let Some(r) = s.strip_prefix(']') {
369 return Some(Self::ProteinC(rest(r)));
370 }
371 let c = s.chars().next()?;
372 VALID_AA
373 .contains(&(c as u8))
374 .then_some(Self::Residue(c as u8))
375 }
376}
377
378fn parse_mods<V>(input: &HashMap<String, V>) -> Vec<(ModSpecificity, V)>
379where
380 V: Clone,
381{
382 let mut out: Vec<(ModSpecificity, V)> = Vec::new();
383 for (s, v) in input {
384 match ModSpecificity::parse(s) {
385 Some(m) => out.push((m, v.clone())),
386 None => log::error!("skipping invalid modification string: {s:?}"),
387 }
388 }
389 out.sort_by_key(|a| a.0);
391 out.dedup_by(|later, earlier| {
392 if later.0 == earlier.0 {
393 std::mem::swap(later, earlier);
394 true
395 } else {
396 false
397 }
398 });
399 out
400}
401
402#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
403enum ModSite {
404 Nterm,
405 Cterm,
406 Sequence(u32),
407}
408
409#[derive(Clone, Debug, PartialEq)]
411pub struct Peptide {
412 pub decoy: bool,
413 pub sequence: Arc<[u8]>,
414 pub modifications: Vec<f32>,
416 pub nterm: Option<f32>,
417 pub cterm: Option<f32>,
418 pub monoisotopic: f32,
420 pub missed_cleavages: u8,
421 pub position: Position,
422 pub proteins: Vec<Arc<str>>,
424}
425
426impl Peptide {
427 fn from_group(g: DigestGroup) -> Option<Peptide> {
428 let seq = g.reference.sequence;
429 if !seq.is_ascii() {
430 return None;
431 }
432 let mut mass = H2O_F32;
433 for &c in seq.as_bytes() {
434 let m = monoisotopic_f32(c);
435 if m == 0.0 {
436 return None;
437 }
438 mass += m;
439 }
440 Some(Peptide {
441 decoy: g.reference.decoy,
442 modifications: vec![0.0; seq.len()],
443 sequence: Arc::from(seq.into_bytes().into_boxed_slice()),
444 nterm: None,
445 cterm: None,
446 monoisotopic: mass,
447 missed_cleavages: g.reference.missed_cleavages,
448 position: g.reference.position,
449 proteins: g.proteins,
450 })
451 }
452
453 fn modification_mass(&self) -> f32 {
454 self.modifications.iter().sum::<f32>()
455 + self.nterm.unwrap_or(0.0)
456 + self.cterm.unwrap_or(0.0)
457 }
458
459 fn apply_site(&mut self, site: ModSite, mass: f32) {
460 match site {
461 ModSite::Nterm => {
462 if self.nterm.is_none() {
463 self.nterm = Some(0.0 + mass);
464 }
465 }
466 ModSite::Cterm => {
467 if self.cterm.is_none() {
468 self.cterm = Some(0.0 + mass);
469 }
470 }
471 ModSite::Sequence(i) => {
472 if self.modifications[i as usize] == 0.0 {
473 self.modifications[i as usize] += mass;
474 }
475 }
476 }
477 }
478
479 fn first(&self) -> u8 {
480 *self.sequence.first().unwrap_or(&0)
481 }
482 fn last_idx(&self) -> u32 {
483 self.sequence.len().saturating_sub(1) as u32
484 }
485 fn last(&self) -> u8 {
486 *self.sequence.last().unwrap_or(&0)
487 }
488
489 fn terminal_site(&self, target: ModSpecificity) -> Option<ModSite> {
492 use ModSpecificity::*;
493 let prot_n = matches!(self.position, Position::Nterm | Position::Full);
494 let prot_c = matches!(self.position, Position::Cterm | Position::Full);
495 match target {
496 PeptideN(None) => Some(ModSite::Nterm),
497 PeptideN(Some(r)) if r == self.first() => Some(ModSite::Sequence(0)),
498 PeptideC(None) => Some(ModSite::Cterm),
499 PeptideC(Some(r)) if r == self.last() => Some(ModSite::Sequence(self.last_idx())),
500 ProteinN(None) if prot_n => Some(ModSite::Nterm),
501 ProteinN(Some(r)) if prot_n && r == self.first() => Some(ModSite::Sequence(0)),
502 ProteinC(None) if prot_c => Some(ModSite::Cterm),
503 ProteinC(Some(r)) if prot_c && r == self.last() => {
504 Some(ModSite::Sequence(self.last_idx()))
505 }
506 _ => None,
507 }
508 }
509
510 fn push_sites(&self, acc: &mut Vec<(ModSite, f32)>, target: ModSpecificity, mass: f32) {
511 if let ModSpecificity::Residue(r) = target {
512 acc.extend(
513 self.sequence
514 .iter()
515 .enumerate()
516 .filter(|(_, &aa)| aa == r)
517 .map(|(i, _)| (ModSite::Sequence(i as u32), mass)),
518 );
519 } else if let Some(site) = self.terminal_site(target) {
520 acc.push((site, mass));
521 }
522 }
523
524 fn apply_static(&mut self, target: ModSpecificity, mass: f32) {
525 if let ModSpecificity::Residue(r) = target {
526 for i in 0..self.sequence.len() {
527 if self.sequence[i] == r && self.modifications[i] == 0.0 {
528 self.modifications[i] = mass;
529 }
530 }
531 } else if let Some(site) = self.terminal_site(target) {
532 self.apply_site(site, mass);
533 }
534 }
535
536 fn apply(
539 self,
540 variable: &[(ModSpecificity, f32)],
541 fixed: &[(ModSpecificity, f32)],
542 max_mods: usize,
543 ) -> Vec<Peptide> {
544 let mut modified = vec![self.clone()];
545 if !variable.is_empty() {
546 let mut sites = Vec::new();
547 for &(target, mass) in variable {
548 self.push_sites(&mut sites, target, mass);
549 }
550 for n in 1..=max_mods {
551 for_each_combination(sites.len(), n, |idx| {
552 let mut used = HashSet::new();
553 if !idx.iter().all(|&i| used.insert(sites[i].0)) {
554 return;
555 }
556 let mut p = self.clone();
557 for &i in idx {
558 p.apply_site(sites[i].0, sites[i].1);
559 }
560 modified.push(p);
561 });
562 }
563 }
564 for p in modified.iter_mut() {
565 for &(target, mass) in fixed {
566 p.apply_static(target, mass);
567 }
568 p.monoisotopic += p.modification_mass();
569 }
570 modified
571 }
572
573 fn reverse(&self) -> Peptide {
575 let mut p = self.clone();
576 p.decoy = !self.decoy;
577 let n = p.sequence.len().saturating_sub(1);
578 if n > 1 {
579 let mut s = p.sequence.to_vec();
580 s[1..n].reverse();
581 p.sequence = Arc::from(s.into_boxed_slice());
582 p.modifications[1..n].reverse();
583 }
584 p
585 }
586
587 fn sort_key_cmp(&self, other: &Self) -> Ordering {
588 self.monoisotopic
589 .total_cmp(&other.monoisotopic)
590 .then_with(|| self.sequence.cmp(&other.sequence))
591 .then_with(|| {
592 self.modifications
593 .partial_cmp(&other.modifications)
594 .unwrap_or(Ordering::Equal)
595 })
596 .then_with(|| {
597 self.nterm
598 .partial_cmp(&other.nterm)
599 .unwrap_or(Ordering::Equal)
600 })
601 .then_with(|| {
602 self.cterm
603 .partial_cmp(&other.cterm)
604 .unwrap_or(Ordering::Equal)
605 })
606 }
607}
608
609fn for_each_combination(n: usize, k: usize, mut f: impl FnMut(&[usize])) {
611 if k == 0 || k > n {
612 return;
613 }
614 let mut idx: Vec<usize> = (0..k).collect();
615 loop {
616 f(&idx);
617 let mut i = k;
618 while i > 0 && idx[i - 1] == n - k + i - 1 {
619 i -= 1;
620 }
621 if i == 0 {
622 return;
623 }
624 idx[i - 1] += 1;
625 for j in i..k {
626 idx[j] = idx[j - 1] + 1;
627 }
628 }
629}
630
631#[derive(Debug, Clone)]
633pub struct DigestParams {
634 pub missed_cleavages: u8,
635 pub min_len: usize,
636 pub max_len: usize,
637 pub cleave_at: String,
639 pub restrict: String,
641 pub c_terminal: bool,
642 pub peptide_min_mass: f32,
643 pub peptide_max_mass: f32,
644 pub static_mods: HashMap<String, f32>,
645 pub variable_mods: HashMap<String, Vec<f32>>,
646 pub max_variable_mods: usize,
647 pub generate_decoys: bool,
648}
649
650pub fn digest(fasta: &Fasta, params: &DigestParams) -> anyhow::Result<Vec<Peptide>> {
653 let enzyme = EnzymeParameters {
654 missed_cleavages: params.missed_cleavages,
655 min_len: params.min_len,
656 max_len: params.max_len,
657 enzyme: Enzyme::new(¶ms.cleave_at, ¶ms.restrict, params.c_terminal)?,
658 };
659 let fixed = parse_mods(¶ms.static_mods);
660 let variable: Vec<(ModSpecificity, f32)> = parse_mods(¶ms.variable_mods)
661 .into_iter()
662 .flat_map(|(t, ms)| ms.into_iter().map(move |m| (t, m)))
663 .collect();
664 let max_mods = params.max_variable_mods.max(1);
665
666 let groups = group_digests(fasta.digest(&enzyme));
667 let targets: HashSet<Vec<u8>> = groups
668 .iter()
669 .filter(|g| !g.reference.decoy)
670 .map(|g| g.reference.sequence.as_bytes().to_vec())
671 .collect();
672
673 let mut peptides: Vec<Peptide> = groups
674 .into_par_iter()
675 .filter_map(Peptide::from_group)
676 .flat_map_iter(|p| {
677 let mut out = Vec::new();
678 for p in p.apply(&variable, &fixed, max_mods) {
679 if p.monoisotopic < params.peptide_min_mass
680 || p.monoisotopic > params.peptide_max_mass
681 {
682 continue;
683 }
684 if params.generate_decoys {
685 let d = p.reverse();
686 if !d.decoy || !targets.contains(&d.sequence[..]) {
687 out.push(d);
688 }
689 }
690 if !p.decoy || !targets.contains(&p.sequence[..]) {
691 out.push(p);
692 }
693 }
694 out
695 })
696 .collect();
697
698 peptides.par_sort_unstable_by(|a, b| a.sort_key_cmp(b));
699 peptides.dedup_by(|remove, keep| {
700 if remove.monoisotopic == keep.monoisotopic
701 && remove.sequence == keep.sequence
702 && remove.modifications == keep.modifications
703 && remove.nterm == keep.nterm
704 && remove.cterm == keep.cterm
705 {
706 keep.proteins.extend(remove.proteins.iter().cloned());
707 keep.decoy &= remove.decoy;
708 true
709 } else {
710 false
711 }
712 });
713 peptides
714 .par_iter_mut()
715 .for_each(|p| p.proteins.sort_unstable());
716 Ok(peptides)
717}
718
719#[cfg(test)]
720mod tests {
721 use super::*;
722
723 fn params() -> DigestParams {
724 DigestParams {
725 missed_cleavages: 0,
726 min_len: 2,
727 max_len: 50,
728 cleave_at: "KR".into(),
729 restrict: "P".into(),
730 c_terminal: true,
731 peptide_min_mass: 0.0,
732 peptide_max_mass: 1e6,
733 static_mods: HashMap::new(),
734 variable_mods: HashMap::new(),
735 max_variable_mods: 2,
736 generate_decoys: false,
737 }
738 }
739
740 #[test]
741 fn trypsin_sites() {
742 let fasta = Fasta::parse(
743 ">P1 desc\nMADEEKLPPGWEKRMSRSSGRVYYFNHITNASQWERPSGN\n",
744 "rev_",
745 false,
746 )
747 .unwrap();
748 let peps = digest(&fasta, ¶ms()).unwrap();
749 let mut seqs: Vec<String> = peps
750 .iter()
751 .map(|p| String::from_utf8(p.sequence.to_vec()).unwrap())
752 .collect();
753 seqs.sort();
754 assert_eq!(
755 seqs,
756 ["LPPGWEK", "MADEEK", "MSR", "SSGR", "VYYFNHITNASQWERPSGN"]
757 );
758 assert!(peps.iter().all(|p| &*p.proteins[0] == "P1"));
759 }
760
761 #[test]
762 fn variable_mods_and_decoys() {
763 let fasta = Fasta::parse(">P1\nAMCMGK\n", "rev_", true).unwrap();
764 let mut p = params();
765 p.generate_decoys = true;
766 p.static_mods.insert("C".into(), 57.02146);
767 p.variable_mods.insert("M".into(), vec![15.9949]);
768 let peps = digest(&fasta, &p).unwrap();
769 assert_eq!(peps.iter().filter(|p| !p.decoy).count(), 4);
771 assert_eq!(peps.iter().filter(|p| p.decoy).count(), 4);
772 assert!(peps
773 .iter()
774 .all(|p| p.modifications[2] == 57.02146 || p.decoy));
775 }
776
777 #[test]
778 fn combinations_are_lexicographic() {
779 let mut v = Vec::new();
780 for_each_combination(4, 2, |c| v.push(c.to_vec()));
781 assert_eq!(
782 v,
783 vec![
784 vec![0, 1],
785 vec![0, 2],
786 vec![0, 3],
787 vec![1, 2],
788 vec![1, 3],
789 vec![2, 3]
790 ]
791 );
792 }
793}