1use anyhow::{Context, Result};
19use rustc_hash::FxHashMap;
20use std::fs::File;
21use std::io::{BufRead, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
22use std::os::unix::fs::FileExt;
23use std::path::Path;
24use std::process::Command;
25
26use crate::snp::{self, SnpStatus};
27
28const FDB_MAGIC: &[u8; 8] = b"FLANKDB\0";
29
30const EMB_GENUS_DIST: &[u8] = include_bytes!("embedded/genus_dist.tsv.zst");
40const EMB_GENUS_LINEAGE: &[u8] = include_bytes!("embedded/genus_lineage.tsv.zst");
41const EMB_CONFORMAL: &[u8] = include_bytes!("embedded/conformal.tsv.zst");
42
43fn context_table_text(
48 db_dir: Option<&Path>,
49 filename: &str,
50 embedded_zst: &[u8],
51) -> Result<(String, String)> {
52 if let Some(dir) = db_dir {
53 let p = dir.join(filename);
54 if p.exists() {
55 let s = std::fs::read_to_string(&p).with_context(|| format!("read {:?}", p))?;
56 return Ok((s, format!("{:?}", p)));
57 }
58 }
59 let bytes =
60 zstd::decode_all(embedded_zst).with_context(|| format!("decode embedded {}", filename))?;
61 let s = String::from_utf8(bytes).with_context(|| format!("embedded {} utf8", filename))?;
62 Ok((s, format!("embedded:{}", filename)))
63}
64
65#[derive(Debug, Clone)]
73pub struct ArgPosition {
74 pub arg_name: String,
76 pub contig_name: String,
78 pub contig_seq: String,
80 pub contig_len: usize,
82 pub arg_start: usize,
84 pub arg_end: usize,
86 pub strand: char,
88 pub members: Vec<String>,
91}
92
93pub const GENUS_TIE_PCT: f64 = 1.0;
97
98#[derive(Debug, Clone)]
99pub struct GenusResult {
100 pub arg_name: String,
102 pub contig_name: String,
104 pub genus: Option<String>,
106 pub confidence: f64,
108 pub specificity: f64,
110 pub upstream_len: usize,
112 pub downstream_len: usize,
114 pub top_matches: Vec<(String, f64)>,
116 pub snp_status: SnpStatus,
118 pub upstream_seq: String,
121 pub downstream_seq: String,
123 pub context: String,
126 pub n_genera_tied: usize,
129 pub species: Option<String>,
132 pub species_top_matches: Vec<(String, f64)>,
134 pub n_species_tied: usize,
136 #[allow(dead_code)]
141 pub credible_set: Vec<(String, f64)>,
142 pub support: f64,
145 pub resolution_distance: f64,
148 pub limited_by: String,
151 pub resolution_rank: String,
154 pub resolution_taxon: String,
156 pub credible_set_grouped: String,
160}
161
162impl Default for GenusResult {
163 fn default() -> Self {
164 Self {
165 arg_name: String::new(),
166 contig_name: String::new(),
167 genus: None,
168 confidence: 0.0,
169 specificity: 0.0,
170 upstream_len: 0,
171 downstream_len: 0,
172 top_matches: vec![],
173 snp_status: SnpStatus::NotApplicable,
174 upstream_seq: String::new(),
175 downstream_seq: String::new(),
176 context: "NA".to_string(),
177 n_genera_tied: 0,
178 species: None,
179 species_top_matches: vec![],
180 n_species_tied: 0,
181 credible_set: vec![],
182 support: 0.0,
183 resolution_distance: 0.0,
184 limited_by: "none".to_string(),
185 resolution_rank: "NA".to_string(),
186 resolution_taxon: "NA".to_string(),
187 credible_set_grouped: String::new(),
188 }
189 }
190}
191
192#[derive(Debug, Clone)]
196pub struct FlankingRecord {
197 pub contig: String,
199 pub genus: String,
201 pub upstream: String,
203 pub downstream: String,
205}
206
207#[derive(Debug, Clone)]
213struct FdbIndexEntry {
214 offset: u64,
215 compressed_len: u32,
216 record_count: u32,
217}
218
219pub struct FlankingDatabase {
235 file: File,
236 index: FxHashMap<String, FdbIndexEntry>,
237 gene_name_to_key: FxHashMap<String, String>,
240}
241
242impl FlankingDatabase {
243 pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
250 let mut file = File::open(path.as_ref())
251 .with_context(|| format!("Failed to open fdb: {}", path.as_ref().display()))?;
252
253 let mut magic = [0u8; 8];
255 file.read_exact(&mut magic)?;
256 if &magic != FDB_MAGIC {
257 anyhow::bail!("Invalid fdb magic");
258 }
259
260 let mut buf4 = [0u8; 4];
261 let mut buf8 = [0u8; 8];
262
263 file.read_exact(&mut buf4)?;
264 let _version = u32::from_le_bytes(buf4);
265
266 file.read_exact(&mut buf4)?;
267 let gene_count = u32::from_le_bytes(buf4);
268
269 file.read_exact(&mut buf8)?;
270 let index_offset = u64::from_le_bytes(buf8);
271
272 file.seek(SeekFrom::Start(index_offset))?;
274 let mut index = FxHashMap::default();
275
276 for _ in 0..gene_count {
277 let mut buf2 = [0u8; 2];
278 file.read_exact(&mut buf2)?;
279 let name_len = u16::from_le_bytes(buf2) as usize;
280
281 let mut name_buf = vec![0u8; name_len];
282 file.read_exact(&mut name_buf)?;
283 let gene = String::from_utf8(name_buf)?;
284
285 file.read_exact(&mut buf8)?;
286 let offset = u64::from_le_bytes(buf8);
287
288 file.read_exact(&mut buf4)?;
289 let compressed_len = u32::from_le_bytes(buf4);
290
291 file.read_exact(&mut buf4)?;
292 let record_count = u32::from_le_bytes(buf4);
293
294 index.insert(gene, FdbIndexEntry {
295 offset,
296 compressed_len,
297 record_count,
298 });
299 }
300
301 let mut gene_name_to_key = FxHashMap::default();
304 for full_key in index.keys() {
305 let gene_name = full_key.split('|').next().unwrap_or(full_key);
307 if !gene_name_to_key.contains_key(gene_name) {
309 gene_name_to_key.insert(gene_name.to_string(), full_key.clone());
310 }
311 }
312
313 Ok(Self { file, index, gene_name_to_key })
314 }
315
316 pub fn has_gene(&self, gene: &str) -> bool {
319 if self.index.contains_key(gene) {
321 return true;
322 }
323 self.gene_name_to_key.contains_key(gene)
325 }
326
327 pub fn get_gene_records(&self, gene: &str) -> Result<Vec<FlankingRecord>> {
332 let lookup_key = if self.index.contains_key(gene) {
334 gene.to_string()
335 } else if let Some(full_key) = self.gene_name_to_key.get(gene) {
336 full_key.clone()
337 } else {
338 anyhow::bail!("Gene not found: {}", gene);
339 };
340
341 let entry = self.index.get(&lookup_key)
342 .ok_or_else(|| anyhow::anyhow!("Gene not found in index: {}", lookup_key))?;
343
344 let mut compressed = vec![0u8; entry.compressed_len as usize];
347 self.file.read_exact_at(&mut compressed, entry.offset)?;
348
349 let decompressed = zstd::decode_all(&compressed[..])?;
351 let content = String::from_utf8(decompressed)?;
352
353 let mut records = Vec::with_capacity(entry.record_count as usize);
355 let mut lines = content.lines();
356
357 let _header = lines.next();
359
360 for line in lines {
361 if line.is_empty() {
362 continue;
363 }
364 let fields: Vec<&str> = line.split('\t').collect();
365 if fields.len() < 7 {
367 continue;
368 }
369
370 records.push(FlankingRecord {
371 contig: fields[1].to_string(),
372 genus: fields[2].to_string(),
373 upstream: fields[5].to_string(),
374 downstream: fields[6].to_string(),
375 });
376 }
377
378 Ok(records)
379 }
380
381}
382
383fn ou_lookup(neighbors: &FxHashMap<String, FxHashMap<String, f64>>, a: &str, b: &str) -> f64 {
390 if a == b {
391 return 0.0;
392 }
393 if let Some(m) = neighbors.get(a) {
394 if let Some(&d) = m.get(b) {
395 return d;
396 }
397 }
398 if let Some(m) = neighbors.get(b) {
399 if let Some(&d) = m.get(a) {
400 return d;
401 }
402 }
403 ABSENT_PATRISTIC
404}
405
406const ABSENT_PATRISTIC: f64 = 3.0;
411
412fn ou_posterior(
422 likelihoods: &FxHashMap<String, f64>,
423 neighbors: &FxHashMap<String, FxHashMap<String, f64>>,
424 lambda: f64,
425) -> Vec<(String, f64)> {
426 let genera: Vec<&String> = likelihoods.keys().collect();
427 let mut post: Vec<(String, f64)> = Vec::with_capacity(genera.len());
428 for &ta in &genera {
429 let (mut num, mut den) = (0.0f64, 0.0f64);
430 for &tb in &genera {
431 let w = (-ou_lookup(neighbors, ta, tb) / lambda).exp();
432 num += w * likelihoods[tb];
433 den += w;
434 }
435 post.push((ta.clone(), if den > 0.0 { num / den } else { 0.0 }));
436 }
437 let z: f64 = post.iter().map(|(_, p)| *p).sum();
438 if z > 0.0 {
439 for p in post.iter_mut() {
440 p.1 /= z;
441 }
442 }
443 post.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
444 post
445}
446
447const GENUS_COHERENCE_RADIUS: f64 = 0.5;
455
456fn render_resolution(members: &[&str], lca_rank: &str, lca_taxon: &str, radius: f64)
467 -> (String, String) {
468 if members.len() == 1 {
469 ("genus".to_string(), members[0].to_string())
470 } else if radius <= GENUS_COHERENCE_RADIUS {
471 ("genus".to_string(), members.join("|"))
472 } else {
473 (lca_rank.to_string(), lca_taxon.to_string())
474 }
475}
476
477fn group_credible_set(
485 members: &[(String, f64)],
486 neighbors: &FxHashMap<String, FxHashMap<String, f64>>,
487 threshold: f64,
488) -> String {
489 let n = members.len();
490 if n == 0 {
491 return "NA".to_string();
492 }
493 if n == 1 {
494 return format!("{}({:.2})", members[0].0, members[0].1);
495 }
496 let mut parent: Vec<usize> = (0..n).collect();
498 fn find(parent: &mut [usize], mut x: usize) -> usize {
499 while parent[x] != x {
500 parent[x] = parent[parent[x]];
501 x = parent[x];
502 }
503 x
504 }
505 for i in 0..n {
506 for j in (i + 1)..n {
507 if ou_lookup(neighbors, &members[i].0, &members[j].0) <= threshold {
508 let (ri, rj) = (find(&mut parent, i), find(&mut parent, j));
509 if ri != rj {
510 parent[ri] = rj;
511 }
512 }
513 }
514 }
515 let mut clusters: FxHashMap<usize, Vec<usize>> = FxHashMap::default();
516 for i in 0..n {
517 let r = find(&mut parent, i);
518 clusters.entry(r).or_default().push(i);
519 }
520 let mut rendered: Vec<(f64, String)> = clusters
521 .values()
522 .map(|idxs| {
523 let mut mem = idxs.clone();
524 mem.sort_by(|&a, &b| {
525 members[b].1.partial_cmp(&members[a].1).unwrap_or(std::cmp::Ordering::Equal)
526 });
527 let mass: f64 = mem.iter().map(|&k| members[k].1).sum();
528 let s = mem.iter()
529 .map(|&k| format!("{}({:.2})", members[k].0, members[k].1))
530 .collect::<Vec<_>>()
531 .join(",");
532 (mass, s)
533 })
534 .collect();
535 rendered.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
536 rendered.into_iter().map(|(_, s)| s).collect::<Vec<_>>().join(" | ")
537}
538
539fn lca_of(genera: &[&str], lineage: &FxHashMap<String, [String; 7]>) -> (String, String) {
545 let lins: Vec<&[String; 7]> = genera.iter().filter_map(|g| lineage.get(*g)).collect();
546 if lins.is_empty() {
547 return if genera.len() == 1 {
549 ("genus".to_string(), genera[0].to_string())
550 } else {
551 ("root".to_string(), "unclassified".to_string())
552 };
553 }
554 for r in (0..6).rev() {
558 let first = &lins[0][r];
559 if !first.is_empty() && lins.iter().all(|l| &l[r] == first) {
560 return (LINEAGE_RANKS[r].to_string(), first.clone());
561 }
562 }
563 ("root".to_string(), "unclassified".to_string())
564}
565
566pub struct GenusClassifier {
571 db: FlankingDatabase,
572 minimap2_path: String,
573 min_identity: f64,
574 min_align_len: usize,
575 max_flanking: usize,
576 plasmid_contigs: rustc_hash::FxHashSet<String>,
579 species_identity: f64,
581 species_map: FxHashMap<String, String>,
583 plasmid_hi: f64,
586 plasmid_lo: f64,
587 max_ref_flanks: usize,
590 genus_neighbors: FxHashMap<String, FxHashMap<String, f64>>,
594 kernel_lambda: f64,
597 kernel_tau: f64,
599 genus_lineage: FxHashMap<String, [String; 7]>,
603 conformal_theta: f64,
608}
609
610const DEFAULT_CREDIBLE_MASS: f64 = 0.9;
612
613const LINEAGE_RANKS: [&str; 7] =
615 ["superkingdom", "phylum", "class", "order", "family", "genus", "species"];
616
617impl GenusClassifier {
618 pub fn new<P: AsRef<Path>>(
627 db_path: P,
628 minimap2_path: &str,
629 min_identity: f64,
630 min_align_len: usize,
631 max_flanking: usize,
632 plasmid_contigs_path: Option<&Path>,
633 species_identity: f64,
634 species_map_path: Option<&Path>,
635 plasmid_hi: f64,
636 plasmid_lo: f64,
637 max_ref_flanks: usize,
638 target_coverage: f64,
641 ) -> Result<Self> {
642 let db_dir = db_path.as_ref().parent().map(|p| p.to_path_buf());
643 let db = FlankingDatabase::open(db_path)?;
644 let mut plasmid_contigs = rustc_hash::FxHashSet::default();
645 if let Some(p) = plasmid_contigs_path {
646 let f = File::open(p).with_context(|| format!("open plasmid list {:?}", p))?;
647 for line in BufReader::new(f).lines() {
648 let acc = line?.trim().to_string();
649 if !acc.is_empty() { plasmid_contigs.insert(acc); }
650 }
651 }
652 let mut species_map = FxHashMap::default();
653 if let Some(p) = species_map_path {
654 let f = File::open(p).with_context(|| format!("open species map {:?}", p))?;
655 for line in BufReader::new(f).lines() {
656 let line = line?;
657 if let Some((c, sp)) = line.split_once('\t') {
658 let (c, sp) = (c.trim(), sp.trim());
659 if !c.is_empty() && !sp.is_empty() { species_map.insert(c.to_string(), sp.to_string()); }
660 }
661 }
662 }
663 let mut genus_neighbors: FxHashMap<String, FxHashMap<String, f64>> = FxHashMap::default();
667 {
668 let (gd_text, gd_src) =
669 context_table_text(db_dir.as_deref(), "genus_dist.tsv", EMB_GENUS_DIST)?;
670 let mut npair = 0usize;
671 for line in gd_text.lines() {
672 let mut it = line.split('\t');
673 if let (Some(a), Some(b), Some(d)) = (it.next(), it.next(), it.next()) {
674 if let Ok(d) = d.trim().parse::<f64>() {
675 genus_neighbors.entry(a.to_string()).or_default()
676 .insert(b.to_string(), d);
677 npair += 1;
678 }
679 }
680 }
681 eprintln!("[kernel] loaded {} genus-distance pairs from {}", npair, gd_src);
682 }
683 let mut genus_lineage: FxHashMap<String, [String; 7]> = FxHashMap::default();
686 {
687 let (gl_text, gl_src) =
688 context_table_text(db_dir.as_deref(), "genus_lineage.tsv", EMB_GENUS_LINEAGE)?;
689 for (i, line) in gl_text.lines().enumerate() {
690 if i == 0 {
691 continue; }
693 let cols: Vec<&str> = line.split('\t').collect();
694 if cols.len() >= 8 {
695 let mut lin: [String; 7] = Default::default();
696 for r in 0..7 {
697 lin[r] = cols[r + 1].to_string();
698 }
699 genus_lineage.insert(cols[0].to_string(), lin);
700 }
701 }
702 eprintln!("[kernel] loaded {} genus lineages from {}", genus_lineage.len(), gl_src);
703 }
704 let mut conformal_theta = DEFAULT_CREDIBLE_MASS;
707 let mut conformal_target = 0.0;
708 {
709 let (cf_text, cf_src) =
710 context_table_text(db_dir.as_deref(), "conformal.tsv", EMB_CONFORMAL)?;
711 let mut best = f64::MAX;
712 for line in cf_text.lines() {
713 if line.starts_with('#') || line.trim().is_empty() {
714 continue;
715 }
716 let c: Vec<&str> = line.split('\t').collect();
717 if c.len() >= 2 {
718 if let (Ok(t), Ok(th)) = (c[0].parse::<f64>(), c[1].parse::<f64>()) {
719 let d = (t - target_coverage).abs();
720 if d < best {
721 best = d;
722 conformal_theta = th;
723 conformal_target = t;
724 }
725 }
726 }
727 }
728 if conformal_target > 0.0 {
729 eprintln!("[kernel] conformal θ={:.3} for target coverage {:.2} (from {})",
730 conformal_theta, conformal_target, cf_src);
731 }
732 }
733 Ok(Self {
734 db,
735 minimap2_path: minimap2_path.to_string(),
736 min_identity,
737 min_align_len,
738 max_flanking,
739 plasmid_contigs,
740 species_identity,
741 species_map,
742 plasmid_hi,
743 plasmid_lo,
744 max_ref_flanks,
745 genus_neighbors,
746 kernel_lambda: 0.3,
747 kernel_tau: 0.01,
748 genus_lineage,
749 conformal_theta,
750 })
751 }
752
753 fn kernel_posterior(&self, likelihoods: &FxHashMap<String, f64>) -> Vec<(String, f64)> {
757 ou_posterior(likelihoods, &self.genus_neighbors, self.kernel_lambda)
758 }
759
760 pub fn classify_batch(&mut self, positions: &[ArgPosition], threads: usize) -> Result<Vec<GenusResult>> {
768 use rayon::prelude::*;
769 use std::collections::HashSet;
770
771 let this: &GenusClassifier = self;
772 let pool = rayon::ThreadPoolBuilder::new().num_threads(threads.max(1)).build()?;
773
774 let mut seen: HashSet<&str> = HashSet::new();
777 let unique_genes: Vec<&str> = positions.iter()
778 .flat_map(|p| p.members.iter().map(|s| s.as_str()))
779 .filter(|&g| this.db.has_gene(g) && seen.insert(g))
780 .collect();
781
782 pool.install(|| -> Result<Vec<GenusResult>> {
783 let loaded: Vec<(String, Vec<FlankingRecord>, FxHashMap<String, usize>)> = unique_genes
787 .par_iter()
788 .map(|&g| {
789 let recs = this.db.get_gene_records(g)?;
790 let mut dist: FxHashMap<String, usize> = FxHashMap::default();
791 for rec in &recs { *dist.entry(rec.genus.clone()).or_default() += 1; }
792 Ok((g.to_string(), recs, dist))
793 })
794 .collect::<Result<Vec<_>>>()?;
795
796 let mut recs_by_gene: FxHashMap<String, Vec<FlankingRecord>> = FxHashMap::default();
797 let mut dist_by_gene: FxHashMap<String, FxHashMap<String, usize>> = FxHashMap::default();
798 for (g, recs, dist) in loaded {
799 recs_by_gene.insert(g.clone(), recs);
800 dist_by_gene.insert(g, dist);
801 }
802
803 positions
805 .par_iter()
806 .enumerate()
807 .map(|(idx, pos)| this.classify_prepared(pos, idx, &recs_by_gene, &dist_by_gene))
808 .collect()
809 })
810 }
811
812 fn classify_prepared(
822 &self,
823 pos: &ArgPosition,
824 idx: usize,
825 recs_by_gene: &FxHashMap<String, Vec<FlankingRecord>>,
826 dist_by_gene: &FxHashMap<String, FxHashMap<String, usize>>,
827 ) -> Result<GenusResult> {
828 let (upstream, downstream) = self.extract_flanking_regions(pos);
830
831 let upstream_len = upstream.len();
832 let downstream_len = downstream.len();
833
834 let snp_status = snp::verify_snp(
836 &pos.contig_seq,
837 &pos.arg_name,
838 0,
839 pos.arg_end - pos.arg_start,
840 pos.arg_start,
841 pos.arg_end,
842 pos.strand,
843 );
844
845 if upstream_len < 50 && downstream_len < 50 {
847 return Ok(GenusResult {
848 arg_name: pos.arg_name.clone(),
849 contig_name: pos.contig_name.clone(),
850 genus: None,
851 confidence: 0.0,
852 specificity: 0.0,
853 upstream_len,
854 downstream_len,
855 top_matches: vec![],
856 snp_status,
857 upstream_seq: upstream.clone(),
858 downstream_seq: downstream.clone(),
859 context: "NA".to_string(),
860 n_genera_tied: 0,
861 species: None,
862 species_top_matches: vec![],
863 n_species_tied: 0,
864 credible_set: vec![],
865 support: 0.0,
866 resolution_distance: 0.0,
867 limited_by: "none".to_string(),
868 resolution_rank: "NA".to_string(),
869 resolution_taxon: "NA".to_string(),
870 credible_set_grouped: String::new(),
871 });
872 }
873
874 if !pos.members.iter().any(|m| recs_by_gene.contains_key(m)) {
876 return Ok(GenusResult {
877 arg_name: pos.arg_name.clone(),
878 contig_name: pos.contig_name.clone(),
879 genus: None,
880 confidence: 0.0,
881 specificity: 0.0,
882 upstream_len,
883 downstream_len,
884 top_matches: vec![("gene_not_in_db".to_string(), 0.0)],
885 snp_status,
886 upstream_seq: upstream.clone(),
887 downstream_seq: downstream.clone(),
888 context: "NA".to_string(),
889 n_genera_tied: 0,
890 species: None,
891 species_top_matches: vec![],
892 n_species_tied: 0,
893 credible_set: vec![],
894 support: 0.0,
895 resolution_distance: 0.0,
896 limited_by: "none".to_string(),
897 resolution_rank: "NA".to_string(),
898 resolution_taxon: "NA".to_string(),
899 credible_set_grouped: String::new(),
900 });
901 }
902
903 let mut ref_records: Vec<&FlankingRecord> = pos.members.iter()
906 .filter_map(|m| recs_by_gene.get(m))
907 .flatten()
908 .collect();
909 if self.max_ref_flanks > 0 && ref_records.len() > self.max_ref_flanks {
915 let step = ref_records.len() / self.max_ref_flanks;
916 ref_records = ref_records.iter().step_by(step).copied().take(self.max_ref_flanks).collect();
917 }
918 if ref_records.is_empty() {
919 return Ok(GenusResult {
920 arg_name: pos.arg_name.clone(),
921 contig_name: pos.contig_name.clone(),
922 genus: None,
923 confidence: 0.0,
924 specificity: 0.0,
925 upstream_len,
926 downstream_len,
927 top_matches: vec![("no_ref_records".to_string(), 0.0)],
928 snp_status,
929 upstream_seq: upstream.clone(),
930 downstream_seq: downstream.clone(),
931 context: "NA".to_string(),
932 n_genera_tied: 0,
933 species: None,
934 species_top_matches: vec![],
935 n_species_tied: 0,
936 credible_set: vec![],
937 support: 0.0,
938 resolution_distance: 0.0,
939 limited_by: "none".to_string(),
940 resolution_rank: "NA".to_string(),
941 resolution_taxon: "NA".to_string(),
942 credible_set_grouped: String::new(),
943 });
944 }
945
946 let temp_dir = if Path::new("/dev/shm").is_dir() {
949 std::path::PathBuf::from("/dev/shm")
950 } else {
951 std::env::temp_dir()
952 };
953 let pid = std::process::id();
954 let query_path = temp_dir.join(format!("argenus_query_{}_{}.fas", pid, idx));
955 let ref_path = temp_dir.join(format!("argenus_ref_{}_{}.fas", pid, idx));
956 let paf_path = temp_dir.join(format!("argenus_align_{}_{}.paf", pid, idx));
957
958 {
960 let mut query_file = BufWriter::new(File::create(&query_path)?);
961 if !upstream.is_empty() {
962 writeln!(query_file, ">upstream")?;
963 writeln!(query_file, "{}", upstream)?;
964 }
965 if !downstream.is_empty() {
966 writeln!(query_file, ">downstream")?;
967 writeln!(query_file, "{}", downstream)?;
968 }
969 }
970
971 {
973 let mut ref_file = BufWriter::new(File::create(&ref_path)?);
974 for (i, rec) in ref_records.iter().enumerate() {
975 if !rec.upstream.is_empty() {
976 writeln!(ref_file, ">{}|{}|up_{}", rec.genus, rec.contig, i)?;
977 writeln!(ref_file, "{}", rec.upstream)?;
978 }
979 if !rec.downstream.is_empty() {
980 writeln!(ref_file, ">{}|{}|down_{}", rec.genus, rec.contig, i)?;
981 writeln!(ref_file, "{}", rec.downstream)?;
982 }
983 }
984 }
985
986 let output = Command::new(&self.minimap2_path)
988 .args(["-x", "sr", "-t", "1", "-c", "--secondary=yes", "-N", "100", "-k", "15", "-w", "5"])
989 .arg(&ref_path)
990 .arg(&query_path)
991 .arg("-o").arg(&paf_path)
992 .stderr(std::process::Stdio::null())
993 .output()
994 .context("Failed to run minimap2")?;
995
996 if !output.status.success() {
997 let _ = std::fs::remove_file(&query_path);
999 let _ = std::fs::remove_file(&ref_path);
1000 let _ = std::fs::remove_file(&paf_path);
1001
1002 return Ok(GenusResult {
1003 arg_name: pos.arg_name.clone(),
1004 contig_name: pos.contig_name.clone(),
1005 genus: None,
1006 confidence: 0.0,
1007 specificity: 0.0,
1008 upstream_len,
1009 downstream_len,
1010 top_matches: vec![("minimap2_failed".to_string(), 0.0)],
1011 snp_status,
1012 upstream_seq: upstream.clone(),
1013 downstream_seq: downstream.clone(),
1014 context: "NA".to_string(),
1015 n_genera_tied: 0,
1016 species: None,
1017 species_top_matches: vec![],
1018 n_species_tied: 0,
1019 credible_set: vec![],
1020 support: 0.0,
1021 resolution_distance: 0.0,
1022 limited_by: "none".to_string(),
1023 resolution_rank: "NA".to_string(),
1024 resolution_taxon: "NA".to_string(),
1025 credible_set_grouped: String::new(),
1026 });
1027 }
1028
1029 let (genus_likelihood, genus_confidence, plasmid_frac, species_scores) =
1031 self.calculate_genus_scores(&paf_path)?;
1032
1033 let _ = std::fs::remove_file(&query_path);
1035 let _ = std::fs::remove_file(&ref_path);
1036 let _ = std::fs::remove_file(&paf_path);
1037
1038 let context = if self.plasmid_contigs.is_empty() || genus_likelihood.is_empty() {
1043 "NA".to_string()
1044 } else if plasmid_frac >= self.plasmid_hi {
1045 "plasmid".to_string()
1046 } else if plasmid_frac <= self.plasmid_lo {
1047 "chromosome".to_string()
1048 } else {
1049 "ambiguous".to_string()
1050 };
1051
1052 let mut genus_dist: FxHashMap<String, usize> = FxHashMap::default();
1054 for m in &pos.members {
1055 if let Some(d) = dist_by_gene.get(m) {
1056 for (k, v) in d { *genus_dist.entry(k.clone()).or_default() += *v; }
1057 }
1058 }
1059 let total_in_db: usize = genus_dist.values().sum();
1060
1061 let sorted_scores: Vec<(String, f64)> = self.kernel_posterior(&genus_likelihood);
1065
1066 let (genus, confidence, specificity) = if let Some((best_genus, _post)) = sorted_scores.first() {
1067 let genus_count = genus_dist.get(best_genus).copied().unwrap_or(0);
1068 let specificity = if total_in_db > 0 {
1069 (genus_count as f64 / total_in_db as f64) * 100.0
1070 } else {
1071 0.0
1072 };
1073 let conf = genus_confidence.get(best_genus).copied().unwrap_or(0.0);
1076 (Some(best_genus.clone()), conf, specificity)
1077 } else {
1078 (None, 0.0, 0.0)
1079 };
1080
1081 let mut credible_set: Vec<(String, f64)> = Vec::new();
1087 let mut support = 0.0f64;
1088 for (g, p) in &sorted_scores {
1089 credible_set.push((g.clone(), *p));
1090 support += *p;
1091 if support >= self.conformal_theta {
1092 break;
1093 }
1094 }
1095 let n_genera_tied = credible_set.len();
1096
1097 let resolution_distance = if credible_set.len() < 2 {
1100 0.0
1101 } else {
1102 let members: Vec<&str> = credible_set.iter().map(|(g, _)| g.as_str()).collect();
1103 let medoid = members.iter().min_by(|&&a, &&b| {
1104 let sa: f64 = credible_set.iter()
1105 .map(|(g, p)| p * ou_lookup(&self.genus_neighbors, a, g)).sum();
1106 let sb: f64 = credible_set.iter()
1107 .map(|(g, p)| p * ou_lookup(&self.genus_neighbors, b, g)).sum();
1108 sa.partial_cmp(&sb).unwrap_or(std::cmp::Ordering::Equal)
1109 }).copied().unwrap_or(members[0]);
1110 members.iter()
1111 .map(|&g| ou_lookup(&self.genus_neighbors, medoid, g))
1112 .fold(0.0f64, f64::max)
1113 };
1114
1115 let reach_up = pos.arg_start.min(self.max_flanking);
1118 let reach_dn = pos.contig_len.saturating_sub(pos.arg_end).min(self.max_flanking);
1119 let contig_limited = reach_up < self.max_flanking || reach_dn < self.max_flanking;
1120 let limited_by = if credible_set.len() < 2 {
1121 "none".to_string()
1122 } else if contig_limited {
1123 "query".to_string()
1124 } else {
1125 "biology".to_string()
1126 };
1127
1128 let (resolution_rank, resolution_taxon) = if credible_set.is_empty() {
1132 ("NA".to_string(), "NA".to_string())
1133 } else {
1134 let members: Vec<&str> = credible_set.iter().map(|(g, _)| g.as_str()).collect();
1135 let (lca_rank, lca_taxon) = lca_of(&members, &self.genus_lineage);
1136 render_resolution(&members, &lca_rank, &lca_taxon, resolution_distance)
1137 };
1138 let credible_set_grouped = if credible_set.is_empty() {
1139 String::new()
1140 } else {
1141 group_credible_set(&credible_set, &self.genus_neighbors, GENUS_COHERENCE_RADIUS)
1142 };
1143
1144 let top_matches: Vec<(String, f64)> = sorted_scores.into_iter().take(5).collect();
1145
1146 let mut sorted_species: Vec<(String, f64)> = species_scores.into_iter().collect();
1148 sorted_species.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
1149 let species = sorted_species.first().map(|(s, _)| s.clone());
1150 let n_species_tied = match sorted_species.first() {
1151 Some((_, best)) => sorted_species.iter()
1152 .filter(|(s, sc)| !s.is_empty() && *sc >= best - GENUS_TIE_PCT)
1153 .count(),
1154 None => 0,
1155 };
1156 let species_top_matches: Vec<(String, f64)> = sorted_species.into_iter().take(5).collect();
1157
1158 Ok(GenusResult {
1159 arg_name: pos.arg_name.clone(),
1160 contig_name: pos.contig_name.clone(),
1161 genus,
1162 confidence,
1163 specificity,
1164 upstream_len,
1165 downstream_len,
1166 top_matches,
1167 snp_status,
1168 upstream_seq: upstream.clone(),
1169 downstream_seq: downstream.clone(),
1170 context,
1171 n_genera_tied,
1172 species,
1173 species_top_matches,
1174 n_species_tied,
1175 credible_set,
1176 support,
1177 resolution_distance,
1178 limited_by,
1179 resolution_rank,
1180 resolution_taxon,
1181 credible_set_grouped,
1182 })
1183 }
1184
1185 fn extract_flanking_regions(&self, pos: &ArgPosition) -> (String, String) {
1189 let seq = &pos.contig_seq;
1190
1191 let upstream_end = pos.arg_start;
1193 let upstream_start = upstream_end.saturating_sub(self.max_flanking);
1194 let upstream = if upstream_end > upstream_start {
1195 seq[upstream_start..upstream_end].to_string()
1196 } else {
1197 String::new()
1198 };
1199
1200 let downstream_start = pos.arg_end;
1202 let downstream_end = (downstream_start + self.max_flanking).min(seq.len());
1203 let downstream = if downstream_end > downstream_start {
1204 seq[downstream_start..downstream_end].to_string()
1205 } else {
1206 String::new()
1207 };
1208
1209 if pos.strand == '-' {
1211 (reverse_complement(&downstream), reverse_complement(&upstream))
1212 } else {
1213 (upstream, downstream)
1214 }
1215 }
1216
1217 fn calculate_genus_scores(&self, paf_path: &Path)
1222 -> Result<(FxHashMap<String, f64>, FxHashMap<String, f64>, f64, FxHashMap<String, f64>)> {
1223 let file = File::open(paf_path)?;
1224 let reader = BufReader::new(file);
1225
1226 let mut genus_matches: FxHashMap<String, Vec<f64>> = FxHashMap::default();
1227 let min_identity_pct = self.min_identity * 100.0;
1228 let (mut plasmid_hits, mut total_hits) = (0usize, 0usize);
1230 let mut species_matches: FxHashMap<String, Vec<f64>> = FxHashMap::default();
1233 let species_pct = self.species_identity * 100.0;
1234 let do_species = self.species_identity > 0.0 && !self.species_map.is_empty();
1235
1236 for line in reader.lines() {
1237 let line = line?;
1238 let fields: Vec<&str> = line.split('\t').collect();
1239 if fields.len() < 12 {
1240 continue;
1241 }
1242
1243 let block_len: usize = fields[10].parse().unwrap_or(0);
1244 let matches: usize = fields[9].parse().unwrap_or(0);
1245
1246 if block_len < self.min_align_len {
1247 continue;
1248 }
1249
1250 let identity = if block_len > 0 {
1251 (matches as f64 / block_len as f64) * 100.0
1252 } else {
1253 0.0
1254 };
1255
1256 if identity < min_identity_pct {
1257 continue;
1258 }
1259
1260 let target_name = fields[5];
1262 let mut toks = target_name.split('|');
1263 let genus = toks.next().unwrap_or("");
1264 let contig = toks.next().unwrap_or("");
1265 if genus.is_empty() {
1268 continue;
1269 }
1270 total_hits += 1;
1271 if !contig.is_empty() && self.plasmid_contigs.contains(contig) {
1272 plasmid_hits += 1;
1273 }
1274 genus_matches.entry(genus.to_string()).or_default().push(identity);
1275
1276 if do_species && identity >= species_pct {
1278 if let Some(sp) = self.species_map.get(contig) {
1279 species_matches.entry(sp.clone()).or_default().push(identity);
1280 }
1281 }
1282 }
1283
1284 let tau = self.kernel_tau;
1288 let mut genus_likelihood: FxHashMap<String, f64> = FxHashMap::default();
1289 let mut genus_confidence: FxHashMap<String, f64> = FxHashMap::default();
1290 for (genus, scores) in genus_matches {
1291 if scores.is_empty() {
1292 continue;
1293 }
1294 let n = scores.len() as f64;
1295 let mean_id = scores.iter().sum::<f64>() / n;
1296 let mean_lik = scores.iter()
1297 .map(|id| (-(1.0 - id / 100.0) / tau).exp())
1298 .sum::<f64>() / n;
1299 genus_likelihood.insert(genus.clone(), mean_lik);
1300 genus_confidence.insert(genus, mean_id);
1301 }
1302
1303 let mut species_scores: FxHashMap<String, f64> = FxHashMap::default();
1304 for (sp, scores) in species_matches {
1305 if scores.is_empty() { continue; }
1306 let avg = scores.iter().sum::<f64>() / scores.len() as f64;
1307 species_scores.insert(sp, avg);
1308 }
1309
1310 let plasmid_frac = if total_hits > 0 { plasmid_hits as f64 / total_hits as f64 } else { 0.0 };
1311 Ok((genus_likelihood, genus_confidence, plasmid_frac, species_scores))
1312 }
1313
1314}
1315
1316fn reverse_complement(seq: &str) -> String {
1322 seq.chars()
1323 .rev()
1324 .map(|c| match c.to_ascii_uppercase() {
1325 'A' => 'T',
1326 'T' => 'A',
1327 'G' => 'C',
1328 'C' => 'G',
1329 _ => 'N',
1330 })
1331 .collect()
1332}
1333
1334#[cfg(test)]
1339mod tests {
1340 use super::*;
1341
1342 #[test]
1343 fn test_reverse_complement() {
1344 assert_eq!(reverse_complement("ATGC"), "GCAT");
1345 assert_eq!(reverse_complement("AAAA"), "TTTT");
1346 assert_eq!(reverse_complement(""), "");
1347 }
1348
1349 #[test]
1350 fn test_genus_result_default() {
1351 let result = GenusResult::default();
1352 assert!(result.genus.is_none());
1353 assert_eq!(result.confidence, 0.0);
1354 }
1355
1356 fn lik(pairs: &[(&str, f64)]) -> FxHashMap<String, f64> {
1357 pairs.iter().map(|(g, l)| (g.to_string(), *l)).collect()
1358 }
1359
1360 #[test]
1361 fn test_ou_posterior_no_neighbors_is_normalized_likelihood() {
1362 let neighbors = FxHashMap::default();
1365 let l = lik(&[("Escherichia", 3.0), ("Salmonella", 1.0)]);
1366 let post = ou_posterior(&l, &neighbors, 0.1);
1367 assert_eq!(post[0].0, "Escherichia");
1368 assert!((post[0].1 - 0.75).abs() < 1e-3, "got {}", post[0].1);
1371 assert!((post[1].1 - 0.25).abs() < 1e-3, "got {}", post[1].1);
1372 }
1373
1374 #[test]
1375 fn test_ou_posterior_borrows_from_close_relative() {
1376 let mut neighbors: FxHashMap<String, FxHashMap<String, f64>> = FxHashMap::default();
1381 neighbors.entry("Escherichia".into()).or_default().insert("Salmonella".into(), 0.0406);
1382 neighbors.entry("Salmonella".into()).or_default().insert("Escherichia".into(), 0.0406);
1383 let l = lik(&[("Escherichia", 1.0), ("Salmonella", 0.0), ("Bacillus", 0.0)]);
1384 let post = ou_posterior(&l, &neighbors, 0.3);
1385 let p: FxHashMap<_, _> = post.into_iter().collect();
1386 assert!(p["Salmonella"] > p["Bacillus"], "close relative must outrank far one");
1387 assert!(p["Salmonella"] > 0.05, "Salmonella should borrow real mass, got {}", p["Salmonella"]);
1388 assert!(p["Escherichia"] > p["Salmonella"], "the observed genus still leads");
1389 }
1390
1391 fn ent(g: &str, fam: &str, gen: &str) -> (String, [String; 7]) {
1392 (g.to_string(), [
1394 "Bacteria".into(), "Pseudomonadota".into(), "Gammaproteobacteria".into(),
1395 "Enterobacterales".into(), fam.into(), gen.into(), format!("{} sp.", gen),
1396 ])
1397 }
1398
1399 #[test]
1400 fn test_lca_ragged_rank() {
1401 let mut lin: FxHashMap<String, [String; 7]> = FxHashMap::default();
1402 for (k, v) in [ent("Escherichia", "Enterobacteriaceae", "Escherichia"),
1403 ent("Salmonella", "Enterobacteriaceae", "Salmonella")] {
1404 lin.insert(k, v);
1405 }
1406 lin.insert("Bacillus".into(), [
1407 "Bacteria".into(), "Bacillota".into(), "Bacilli".into(), "Bacillales".into(),
1408 "Bacillaceae".into(), "Bacillus".into(), "Bacillus subtilis".into(),
1409 ]);
1410 assert_eq!(lca_of(&["Escherichia"], &lin), ("genus".into(), "Escherichia".into()));
1412 assert_eq!(lca_of(&["Escherichia", "Salmonella"], &lin),
1414 ("family".into(), "Enterobacteriaceae".into()));
1415 assert_eq!(lca_of(&["Escherichia", "Bacillus"], &lin),
1417 ("superkingdom".into(), "Bacteria".into()));
1418 }
1419
1420 #[test]
1421 fn test_render_resolution_notation() {
1422 assert_eq!(render_resolution(&["Escherichia"], "genus", "Escherichia", 0.0),
1424 ("genus".into(), "Escherichia".into()));
1425 assert_eq!(render_resolution(&["Escherichia", "Shigella"], "family", "Enterobacteriaceae", 0.03),
1427 ("genus".into(), "Escherichia|Shigella".into()));
1428 let tight_many = ["A", "B", "C", "D", "E", "F"];
1430 assert_eq!(render_resolution(&tight_many, "family", "Enterobacteriaceae", 0.10),
1431 ("genus".into(), "A|B|C|D|E|F".into()));
1432 assert_eq!(render_resolution(&["Escherichia", "Klebsiella"], "family", "Enterobacteriaceae", 0.6),
1434 ("family".into(), "Enterobacteriaceae".into()));
1435 assert_eq!(render_resolution(&["Escherichia", "Bacillus"], "superkingdom", "Bacteria", 1.5),
1437 ("superkingdom".into(), "Bacteria".into()));
1438 }
1439
1440 #[test]
1441 fn test_group_credible_set_clusters_by_distance() {
1442 let mut nb: FxHashMap<String, FxHashMap<String, f64>> = FxHashMap::default();
1446 let mut set = |a: &str, b: &str, d: f64| {
1447 nb.entry(a.into()).or_default().insert(b.into(), d);
1448 nb.entry(b.into()).or_default().insert(a.into(), d);
1449 };
1450 set("Escherichia", "Pseudescherichia", 0.08);
1451 set("Escherichia", "Cronobacter", 0.90);
1452 set("Escherichia", "Klebsiella", 0.70);
1453 set("Pseudescherichia", "Cronobacter", 0.90);
1454 set("Pseudescherichia", "Klebsiella", 0.70);
1455 set("Cronobacter", "Klebsiella", 0.80);
1456 let members = vec![
1457 ("Escherichia".to_string(), 0.30),
1458 ("Cronobacter".to_string(), 0.25),
1459 ("Pseudescherichia".to_string(), 0.25),
1460 ("Klebsiella".to_string(), 0.20),
1461 ];
1462 let out = group_credible_set(&members, &nb, GENUS_COHERENCE_RADIUS);
1463 assert_eq!(out, "Escherichia(0.30),Pseudescherichia(0.25) | Cronobacter(0.25) | Klebsiella(0.20)");
1464 }
1465}