1use std::collections::{BTreeMap, HashMap};
2use std::time::Instant;
3
4use serde::{Deserialize, Serialize};
5
6use crate::alignment::{AlignmentResult, RunInput};
7use crate::lfq::LfqConfig;
8
9#[path = "group_confidence.rs"]
10mod group_confidence;
11pub use group_confidence::{EvidenceAudit, PermutationAudit, VARIANT_NAMES};
12
13type GroupSpans = ((f64, f64), (f64, f64), Option<(f64, f64)>);
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
18#[serde(default, deny_unknown_fields)]
19pub struct ConsensusConfig {
20 pub mz_ppm: f64,
22 pub rt_window_pct: f64,
24 pub im_tolerance: f64,
26 pub max_group_qvalue: f64,
30}
31
32impl Default for ConsensusConfig {
33 fn default() -> Self {
34 Self {
35 mz_ppm: 20.0,
36 rt_window_pct: 0.02,
37 im_tolerance: 0.05,
38 max_group_qvalue: 0.05,
39 }
40 }
41}
42
43impl ConsensusConfig {
44 pub fn validate(&self) -> Result<(), crate::error::KothError> {
45 for (name, value) in [
46 ("mz_ppm", self.mz_ppm),
47 ("rt_window_pct", self.rt_window_pct),
48 ("im_tolerance", self.im_tolerance),
49 ] {
50 if !value.is_finite() || value < 0.0 {
51 return Err(crate::error::KothError::ConfigError(format!(
52 "lfq.consensus.{name} must be finite and nonnegative"
53 )));
54 }
55 }
56 if !self.max_group_qvalue.is_finite() || !(0.0..=1.0).contains(&self.max_group_qvalue) {
57 return Err(crate::error::KothError::ConfigError(
58 "lfq.consensus.max_group_qvalue must be between 0 and 1".into(),
59 ));
60 }
61 Ok(())
62 }
63}
64
65#[derive(Debug, Clone)]
67pub struct ConsensusFeature {
68 pub ref_mz: f64,
70 pub ref_rt: f64,
72 pub ref_im: f64,
74 pub charge: u8,
75 pub theoretical_pattern: Vec<f64>,
77 pub neutral_mass: f64,
79 pub seed_combined_score: f64,
81 pub seed_run_idx: usize,
83 pub n_contributing_runs: usize,
85 pub per_run_feature: Vec<Option<u32>>,
93 pub members: Vec<(usize, u32)>,
95 pub seed_feature_idx: u32,
99 pub group_score: f64,
101 pub group_qvalue: f64,
103}
104
105#[derive(Clone)]
108struct ProjectedFeature {
109 run_idx: usize,
110 run_rank: usize,
112 measurement_key: u64,
114 feature_idx: u32,
117 ref_mz: f64,
118 ref_rt: f64,
119 ref_im: f64,
120 charge: u8,
121 neutral_mass: f64,
122 combined_score: f64,
123 theoretical_pattern: Vec<f64>,
124}
125
126fn emit_group(projected: &[ProjectedFeature], group: &[usize], n_runs: usize) -> ConsensusFeature {
127 let mut best_per_run: HashMap<usize, usize> = HashMap::new();
129 for &i in group {
130 let run_idx = projected[i].run_idx;
131 let entry = best_per_run.entry(run_idx).or_insert(i);
132 if quality_order(&projected[i], &projected[*entry]).is_lt() {
133 *entry = i;
134 }
135 }
136
137 let seed_idx = *group
138 .iter()
139 .min_by(|&&a, &&b| quality_order(&projected[a], &projected[b]))
140 .unwrap();
141 let seed = &projected[seed_idx];
142
143 let mut per_run_feature: Vec<Option<u32>> = vec![None; n_runs];
145 for (&run_idx, &proj_i) in &best_per_run {
146 per_run_feature[run_idx] = Some(projected[proj_i].feature_idx);
147 }
148
149 ConsensusFeature {
150 ref_mz: seed.ref_mz,
151 ref_rt: seed.ref_rt,
152 ref_im: seed.ref_im,
153 charge: seed.charge,
154 theoretical_pattern: seed.theoretical_pattern.clone(),
155 neutral_mass: seed.neutral_mass,
156 seed_combined_score: seed.combined_score,
157 seed_run_idx: seed.run_idx,
158 n_contributing_runs: best_per_run.len(),
159 per_run_feature,
160 members: group
161 .iter()
162 .map(|&i| (projected[i].run_idx, projected[i].feature_idx))
163 .collect(),
164 seed_feature_idx: seed.feature_idx,
165 group_score: 0.0,
166 group_qvalue: 1.0,
167 }
168}
169
170fn project_features(runs: &[RunInput], alignment: &AlignmentResult) -> Vec<ProjectedFeature> {
177 const PROTON: f64 = 1.007_276_466_621;
178
179 let t_project = Instant::now();
180 let mut projected: Vec<ProjectedFeature> = Vec::new();
181
182 let mut run_order: Vec<usize> = (0..runs.len()).collect();
183 run_order.sort_by_key(|&i| &runs[i].name);
184 let mut run_ranks = vec![0; runs.len()];
185 for (rank, &i) in run_order.iter().enumerate() {
186 run_ranks[i] = rank;
187 }
188 for (run_idx, run) in runs.iter().enumerate() {
189 let is_reference = run_idx == alignment.reference_idx;
190
191 for (feat_idx, feat) in run.features.iter().enumerate() {
192 if feat.feature.charge == 0 || !feat.combined_score.is_finite() {
193 continue;
194 }
195
196 let native_mz = feat.monoisotopic_mz();
197 let native_rt = feat.feature.rt_apex();
198 let native_im = feat.feature.im_apex();
199
200 let (ref_mz, ref_rt, ref_im) = if is_reference {
201 (native_mz, native_rt, native_im)
202 } else {
203 let al = &alignment.alignments[&run.name];
204 (
205 al.correct_mz(native_mz, native_rt),
206 al.warp_rt(native_rt),
207 al.correct_im(native_im, native_rt),
208 )
209 };
210
211 let neutral_mass =
212 ref_mz * feat.feature.charge as f64 - feat.feature.charge as f64 * PROTON;
213
214 if ![ref_mz, ref_rt, ref_im, neutral_mass]
215 .iter()
216 .all(|x| x.is_finite())
217 || neutral_mass <= 0.0
218 {
219 continue;
220 }
221 use std::hash::{Hash, Hasher};
222 let mut key = std::collections::hash_map::DefaultHasher::new();
223 for x in [
224 feat.cosine_score,
225 feat.isotope_score,
226 feat.feature.ppm_error,
227 feat.feature.rt_start(),
228 feat.feature.rt_end(),
229 feat.feature.total_intensity(),
230 ] {
231 x.to_bits().hash(&mut key);
232 }
233 feat.neutron_offset.hash(&mut key);
234 feat.feature.n_scans_total().hash(&mut key);
235 for x in &feat.theoretical_pattern {
236 x.to_bits().hash(&mut key);
237 }
238 projected.push(ProjectedFeature {
239 run_idx,
240 run_rank: run_ranks[run_idx],
241 measurement_key: key.finish(),
242 feature_idx: feat_idx as u32,
243 ref_mz,
244 ref_rt,
245 ref_im,
246 charge: feat.feature.charge,
247 neutral_mass,
248 combined_score: feat.combined_score,
249 theoretical_pattern: feat.theoretical_pattern.clone(),
250 });
251 }
252 }
253
254 log::info!(
255 "[timing] consensus project ({} features): {:.2?}",
256 projected.len(),
257 t_project.elapsed()
258 );
259
260 projected
261}
262
263pub fn build_consensus_candidates(
264 runs: &[RunInput],
265 alignment: &AlignmentResult,
266 config: &LfqConfig,
267) -> Vec<ConsensusFeature> {
268 let projected = project_features(runs, alignment);
269 let range = runs[alignment.reference_idx].rt_range();
270 group_confidence::score_candidates(&projected, &config.consensus, range.1 - range.0, runs.len())
271}
272
273pub fn audit_consensus(
275 runs: &[RunInput],
276 alignment: &AlignmentResult,
277 config: &LfqConfig,
278 full_controls: bool,
279) -> EvidenceAudit {
280 let projected = project_features(runs, alignment);
281 let range = runs[alignment.reference_idx].rt_range();
282 group_confidence::audit_candidates(
283 &projected,
284 &config.consensus,
285 range.1 - range.0,
286 runs.len(),
287 full_controls,
288 )
289}
290
291pub fn audit_permutation_consensus(
293 runs: &[RunInput],
294 alignment: &AlignmentResult,
295 config: &LfqConfig,
296) -> PermutationAudit {
297 let projected = project_features(runs, alignment);
298 let range = runs[alignment.reference_idx].rt_range();
299 group_confidence::audit_permutations(
300 &projected,
301 &config.consensus,
302 range.1 - range.0,
303 runs.len(),
304 )
305}
306
307pub fn build_consensus(
310 runs: &[RunInput],
311 alignment: &AlignmentResult,
312 config: &LfqConfig,
313) -> Vec<ConsensusFeature> {
314 config
315 .consensus
316 .validate()
317 .expect("invalid consensus configuration");
318 let candidates = build_consensus_candidates(runs, alignment, config);
319 let count = candidates.len();
320 let retained: Vec<_> = candidates
321 .into_iter()
322 .filter(|g| {
323 g.n_contributing_runs >= 2 && g.group_qvalue <= config.consensus.max_group_qvalue
324 })
325 .collect();
326 log::info!(
327 "Consensus confidence: retained {} / {} candidates at group q <= {}",
328 retained.len(),
329 count,
330 config.consensus.max_group_qvalue
331 );
332 retained
333}
334
335fn quality_order(a: &ProjectedFeature, b: &ProjectedFeature) -> std::cmp::Ordering {
336 b.combined_score
337 .total_cmp(&a.combined_score)
338 .then(a.charge.cmp(&b.charge))
339 .then(a.neutral_mass.total_cmp(&b.neutral_mass))
340 .then(a.ref_rt.total_cmp(&b.ref_rt))
341 .then(a.ref_im.total_cmp(&b.ref_im))
342 .then(a.run_rank.cmp(&b.run_rank))
343 .then(a.measurement_key.cmp(&b.measurement_key))
344 .then(a.feature_idx.cmp(&b.feature_idx))
345}
346
347struct Group {
348 members: Vec<usize>,
349 mass: (f64, f64),
350 rt: (f64, f64),
351 im: Option<(f64, f64)>,
352}
353impl Group {
354 fn new(i: usize, p: &ProjectedFeature) -> Self {
355 Self {
356 members: vec![i],
357 mass: (p.neutral_mass, p.neutral_mass),
358 rt: (p.ref_rt, p.ref_rt),
359 im: (p.ref_im != 0.0).then_some((p.ref_im, p.ref_im)),
360 }
361 }
362 fn spans(&self, p: &ProjectedFeature) -> GroupSpans {
363 let mass = (
364 self.mass.0.min(p.neutral_mass),
365 self.mass.1.max(p.neutral_mass),
366 );
367 let rt = (self.rt.0.min(p.ref_rt), self.rt.1.max(p.ref_rt));
368 let im = if p.ref_im == 0.0 {
369 self.im
370 } else {
371 Some(self.im.map_or((p.ref_im, p.ref_im), |(lo, hi)| {
372 (lo.min(p.ref_im), hi.max(p.ref_im))
373 }))
374 };
375 (mass, rt, im)
376 }
377}
378
379fn cluster_projected(
382 projected: &[ProjectedFeature],
383 config: &ConsensusConfig,
384 rt_span: f64,
385 n_runs: usize,
386) -> Vec<ConsensusFeature> {
387 assert!(config.mz_ppm.is_finite() && config.mz_ppm >= 0.0);
388 assert!(config.rt_window_pct.is_finite() && config.rt_window_pct >= 0.0);
389 assert!(config.im_tolerance.is_finite() && config.im_tolerance >= 0.0);
390 let rt_tol = config.rt_window_pct * rt_span;
391 let ratio = 1.0 + config.mz_ppm / 1e6;
392 let mut order: Vec<usize> = (0..projected.len()).collect();
393 order.sort_by(|&a, &b| quality_order(&projected[a], &projected[b]));
394 let mut index: BTreeMap<(u8, u64), Vec<usize>> = BTreeMap::new();
395 let mut groups: Vec<Group> = Vec::new();
396 for i in order {
397 let p = &projected[i];
398 let lo = (p.neutral_mass / ratio).to_bits();
399 let hi = (p.neutral_mass * ratio).to_bits();
400 let mut best: Option<(f64, usize)> = None;
401 for ids in index
402 .range((p.charge, lo)..=(p.charge, hi))
403 .map(|(_, ids)| ids)
404 {
405 for &g in ids {
406 let (mass, rt, im) = groups[g].spans(p);
407 if (mass.1 - mass.0) / mass.0 * 1e6 > config.mz_ppm
408 || rt.1 - rt.0 > rt_tol
409 || im.is_some_and(|(lo, hi)| hi - lo > config.im_tolerance)
410 {
411 continue;
412 }
413 let seed = &projected[groups[g].members[0]];
414 let residual = ((p.neutral_mass - seed.neutral_mass) / seed.neutral_mass * 1e6
415 / config.mz_ppm.max(f64::EPSILON))
416 .powi(2)
417 + ((p.ref_rt - seed.ref_rt) / rt_tol.max(f64::EPSILON)).powi(2)
418 + if p.ref_im != 0.0 && seed.ref_im != 0.0 {
419 ((p.ref_im - seed.ref_im) / config.im_tolerance.max(f64::EPSILON)).powi(2)
420 } else {
421 0.0
422 };
423 if best.is_none_or(|(r, j)| residual < r || (residual == r && g < j)) {
424 best = Some((residual, g));
425 }
426 }
427 }
428 if let Some((_, g)) = best {
429 let (mass, rt, im) = groups[g].spans(p);
430 groups[g].mass = mass;
431 groups[g].rt = rt;
432 groups[g].im = im;
433 groups[g].members.push(i);
434 } else {
435 index
436 .entry((p.charge, p.neutral_mass.to_bits()))
437 .or_default()
438 .push(groups.len());
439 groups.push(Group::new(i, p));
440 }
441 }
442 let mut out: Vec<_> = groups
443 .iter()
444 .map(|g| emit_group(projected, &g.members, n_runs))
445 .collect();
446 out.sort_by(|a, b| {
447 a.charge
448 .cmp(&b.charge)
449 .then(a.neutral_mass.total_cmp(&b.neutral_mass))
450 .then(a.ref_im.total_cmp(&b.ref_im))
451 .then(a.ref_rt.total_cmp(&b.ref_rt))
452 });
453 out
454}
455
456#[cfg(test)]
457mod tests {
458 use super::*;
459
460 fn proj(
461 run_idx: usize,
462 feature_idx: u32,
463 combined_score: f64,
464 ref_mz: f64,
465 ) -> ProjectedFeature {
466 ProjectedFeature {
467 run_idx,
468 run_rank: run_idx,
469 measurement_key: 0,
470 feature_idx,
471 ref_mz,
472 ref_rt: 10.0,
473 ref_im: 0.0,
474 charge: 2,
475 neutral_mass: 1000.0,
476 combined_score,
477 theoretical_pattern: vec![0.6, 0.3, 0.1],
478 }
479 }
480
481 #[test]
485 fn split_peak_dedup_counts_runs_not_features() {
486 let projected = vec![
487 proj(0, 10, 0.5, 500.00), proj(0, 11, 0.9, 500.10), proj(1, 20, 0.7, 500.20), ];
491 let cf = emit_group(&projected, &[0, 1, 2], 2);
492 assert_eq!(cf.n_contributing_runs, 2, "two runs, not three features");
493 assert_eq!(cf.seed_run_idx, 0); assert!((cf.ref_mz - 500.10).abs() < 1e-9);
495 assert_eq!(cf.per_run_feature[0], Some(11)); assert_eq!(cf.per_run_feature[1], Some(20));
497 }
498
499 #[test]
505 fn seed_selection_is_deterministic_under_score_tie() {
506 let projected = vec![proj(0, 1, 0.8, 500.0), proj(1, 2, 0.8, 600.0)];
507 let first = emit_group(&projected, &[0, 1], 2);
508 for _ in 0..64 {
509 let cf = emit_group(&projected, &[0, 1], 2);
510 assert_eq!(
511 cf.ref_mz, first.ref_mz,
512 "seed m/z must be stable across calls"
513 );
514 assert_eq!(cf.seed_run_idx, first.seed_run_idx);
515 }
516 }
517 pub(super) fn at(run: usize, id: u32, mass: f64, rt: f64, score: f64) -> ProjectedFeature {
518 let mut p = proj(run, id, score, mass / 2.0 + 1.007276466621);
519 p.neutral_mass = mass;
520 p.ref_rt = rt;
521 p
522 }
523 fn cfg() -> ConsensusConfig {
524 ConsensusConfig::default()
525 }
526 fn replicated(
527 p: &[ProjectedFeature],
528 c: &ConsensusConfig,
529 span: f64,
530 runs: usize,
531 ) -> Vec<ConsensusFeature> {
532 cluster_projected(p, c, span, runs)
533 .into_iter()
534 .filter(|g| g.n_contributing_runs >= 2)
535 .collect()
536 }
537 #[test]
538 fn interleaved_low_seed_keeps_its_valid_run_support() {
539 let p = vec![
540 at(0, 0, 1000.0, 10.0, 0.7),
541 at(0, 1, 1000.001, 40.0, 0.9),
542 at(1, 0, 1000.002, 10.0, 0.9),
543 ];
544 let g = replicated(&p, &cfg(), 100.0, 2);
545 assert_eq!(g.len(), 1);
546 assert_eq!(g[0].n_contributing_runs, 2);
547 assert_eq!(g[0].members.len(), 2);
548 }
549 #[test]
550 fn weak_members_are_candidates_without_a_seed_override() {
551 let p = vec![
552 at(0, 0, 1000.0, 10.0, 0.3),
553 at(1, 0, 1000.002, 10.1, 0.2),
554 at(0, 1, 1100.0, 20.0, 0.4),
555 ];
556 let g = replicated(&p, &cfg(), 100.0, 2);
557 assert_eq!(g.len(), 1);
558 assert_eq!(g[0].seed_combined_score, 0.3);
559 }
560 #[test]
561 fn replicated_support_is_counted_after_interleaved_members_join() {
562 let p = vec![
563 at(0, 0, 1000.0, 10.0, 0.9),
564 at(0, 1, 1000.001, 40.0, 0.9),
565 at(1, 0, 1000.002, 10.0, 0.9),
566 ];
567 let g = replicated(&p, &cfg(), 100.0, 2);
568 assert_eq!(g.len(), 1);
569 assert_eq!(g[0].n_contributing_runs, 2);
570 }
571 #[test]
572 fn chain_cannot_exceed_mass_or_rt_span() {
573 let p = vec![
574 at(0, 0, 1000.0, 10.0, 0.99),
575 at(1, 0, 1000.015, 11.5, 0.9),
576 at(2, 0, 1000.030, 13.0, 0.9),
577 ];
578 let g = cluster_projected(&p, &cfg(), 100.0, 3);
579 assert_eq!(g.len(), 2);
580 assert_eq!(g.iter().map(|g| g.members.len()).sum::<usize>(), 3);
581 for g in g {
582 assert!(g.members.len() <= 2);
583 }
584 }
585 #[test]
587 fn rt_outliers_can_form_two_replicated_groups_at_the_same_mass() {
588 let p = vec![
589 at(0, 0, 1000.0, 10.0, 0.99),
590 at(1, 0, 1000.0, 10.1, 0.95),
591 at(2, 0, 1000.0, 10.4, 0.90),
592 at(3, 0, 1000.0, 10.5, 0.85),
593 ];
594 let g = cluster_projected(&p, &cfg(), 15.0, 4);
595 assert_eq!(g.len(), 2);
596 assert!(g.iter().all(|g| g.n_contributing_runs == 2));
597 let ownership: std::collections::HashSet<_> =
598 g.iter().flat_map(|g| g.members.iter().copied()).collect();
599 assert_eq!(ownership.len(), p.len());
600 assert_eq!(g[0].neutral_mass, g[1].neutral_mass);
602 }
603
604 #[test]
605 fn residual_monoisotope_error_forms_a_separate_group() {
606 let c13 = crate::lfq::grid::C13_NEUTRON;
607 let p = vec![
608 at(0, 0, 1000.0, 10.0, 0.99),
609 at(1, 0, 1000.0, 10.0, 0.95),
610 at(2, 0, 1000.0 + c13, 10.0, 0.90),
611 at(3, 0, 1000.0 + c13, 10.0, 0.85),
612 ];
613 let g = cluster_projected(&p, &cfg(), 15.0, 4);
614 assert_eq!(g.len(), 2);
615 assert!(g.iter().all(|g| g.n_contributing_runs == 2));
616 assert!((g[1].neutral_mass - g[0].neutral_mass - c13).abs() < 1e-9);
617 }
618
619 #[test]
620 fn per_run_winner_uses_its_own_score_and_keeps_alternatives() {
621 let p = vec![
622 at(0, 0, 1000.0, 10.0, 0.99),
623 at(1, 10, 1000.001, 10.0, 0.51),
624 at(1, 20, 1000.002, 10.0, 0.90),
625 ];
626 let g = replicated(&p, &cfg(), 100.0, 2);
627 assert_eq!(g.len(), 1);
628 assert_eq!(g[0].per_run_feature[1], Some(20));
629 assert_eq!(g[0].members.len(), 3);
630 assert_eq!(g[0].n_contributing_runs, 2);
631 }
632 #[test]
633 fn missing_im_does_not_bridge_incompatible_conformers() {
634 let mut p = vec![
635 at(0, 0, 1000.0, 10.0, 0.99),
636 at(1, 0, 1000.0, 10.0, 0.95),
637 at(2, 0, 1000.0, 10.0, 0.9),
638 ];
639 p[0].ref_im = 1.0;
640 p[1].ref_im = 0.0;
641 p[2].ref_im = 1.2;
642 assert_eq!(cluster_projected(&p, &cfg(), 100.0, 3).len(), 2);
643 }
644 #[test]
645 fn reordering_observations_and_run_indices_preserves_group_coordinates() {
646 let mut p = vec![
647 at(0, 0, 1000.0, 10.0, 0.9),
648 at(1, 0, 1000.015, 11.5, 0.9),
649 at(2, 0, 1000.03, 13.0, 0.9),
650 ];
651 let signature = |g: Vec<ConsensusFeature>| {
652 g.iter()
653 .map(|f| (f.neutral_mass, f.ref_rt, f.n_contributing_runs))
654 .collect::<Vec<_>>()
655 };
656 let expected = signature(cluster_projected(&p, &cfg(), 100.0, 3));
657 p.reverse();
658 for x in &mut p {
659 x.run_idx = 2 - x.run_idx;
660 }
661 assert_eq!(expected, signature(cluster_projected(&p, &cfg(), 100.0, 3)));
662 }
663}