1use ahash::{HashMap, HashSet};
11use anyhow::{Context, Result};
12use bio_types::annot::contig::Contig;
13use bio_types::strand::ReqStrand;
14use log::debug;
15
16use bio_types::strand::Strand as BioStrand;
17use noodles::core::{Position, Region};
18use noodles::sam::header::record::value::{Map, map::Program, map::program::tag as pg_tag};
19use noodles::{bam, bed, sam};
20use polars::prelude::*;
21use rust_lapper::Interval;
22use rust_lapper::Lapper;
23use std::ops::Bound;
24use std::path::{Path, PathBuf};
25use std::str::FromStr;
26
27pub const CB: [u8; 2] = [b'C', b'B'];
29pub type Iv = Interval<usize, u32>;
31const BAMNADO_PROGRAM_ID: &str = "bamnado";
32const BAMNADO_PROGRAM_NAME: &str = "bamnado";
33const BAMNADO_VERSION: &str = env!("CARGO_PKG_VERSION");
34
35#[derive(Debug, Clone, PartialEq, Eq, Copy)]
37pub enum Strand {
38 Forward,
40 Reverse,
42 Both,
44}
45impl FromStr for Strand {
46 type Err = String;
47
48 fn from_str(s: &str) -> Result<Self, Self::Err> {
49 match s.to_lowercase().as_str() {
50 "forward" => Ok(Strand::Forward),
51 "reverse" => Ok(Strand::Reverse),
52 "both" => Ok(Strand::Both),
53 _ => Err(format!("Invalid strand value: {s}")),
54 }
55 }
56}
57impl std::fmt::Display for Strand {
58 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
59 match self {
60 Strand::Forward => write!(f, "forward"),
61 Strand::Reverse => write!(f, "reverse"),
62 Strand::Both => write!(f, "both"),
63 }
64 }
65}
66
67impl From<Strand> for BioStrand {
68 fn from(val: Strand) -> Self {
69 match val {
70 Strand::Forward => BioStrand::Forward,
71 Strand::Reverse => BioStrand::Reverse,
72 Strand::Both => BioStrand::Unknown,
73 }
74 }
75}
76
77pub fn progress_bar(length: u64, message: String) -> indicatif::ProgressBar {
84 let progress_style = indicatif::ProgressStyle::with_template(
86 "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({eta}) {msg}",
87 )
88 .unwrap()
89 .progress_chars("##-");
90
91 indicatif::ProgressBar::new(length)
92 .with_style(progress_style)
93 .with_message(message)
94}
95
96pub struct CellBarcodes {
98 barcodes: HashSet<String>,
99}
100
101impl Default for CellBarcodes {
102 fn default() -> Self {
103 Self::new()
104 }
105}
106
107impl CellBarcodes {
108 pub fn new() -> Self {
110 Self {
111 barcodes: HashSet::default(),
112 }
113 }
114
115 pub fn barcodes(&self) -> HashSet<String> {
117 self.barcodes.clone()
118 }
119
120 pub fn is_empty(&self) -> bool {
122 self.barcodes.is_empty()
123 }
124
125 pub fn from_csv(file_path: &PathBuf) -> Result<Self> {
129 let path = Path::new(file_path).to_path_buf();
130
131 let df = CsvReadOptions::default()
132 .with_has_header(true)
133 .try_into_reader_with_file_path(Some(path))?
134 .finish()?;
135 let mut barcodes = HashSet::default();
136
137 for barcode in df.column("barcode").unwrap().str().unwrap() {
138 let barcode = barcode.unwrap().to_string();
139 barcodes.insert(barcode);
140 }
141
142 println!("Number of barcodes: {}", barcodes.len());
143 println!(
144 "First 10 barcodes: {:?}",
145 barcodes.iter().take(10).collect::<Vec<_>>()
146 );
147
148 Ok(Self { barcodes })
149 }
150}
151
152pub struct CellBarcodesMulti {
154 barcodes: Vec<HashSet<String>>,
155}
156
157impl Default for CellBarcodesMulti {
158 fn default() -> Self {
159 Self::new()
160 }
161}
162
163impl CellBarcodesMulti {
164 pub fn new() -> Self {
166 Self {
167 barcodes: Vec::new(),
168 }
169 }
170
171 pub fn barcodes(&self) -> Vec<HashSet<String>> {
173 self.barcodes.clone()
174 }
175
176 pub fn is_empty(&self) -> bool {
178 self.barcodes.is_empty()
179 }
180
181 pub fn from_csv(file_path: &PathBuf) -> Result<Self> {
185 let path = Path::new(file_path).to_path_buf();
186
187 let df = CsvReadOptions::default()
188 .with_has_header(true)
189 .try_into_reader_with_file_path(Some(path))?
190 .finish()?;
191
192 let mut barcodes = Vec::new();
193
194 for column in df.iter() {
195 let barcode = column.str().unwrap();
196
197 let bc = barcode
198 .iter()
199 .filter_map(|x| x.map(|x| x.to_string()))
200 .collect::<HashSet<String>>();
201
202 barcodes.push(bc);
203 }
204
205 Ok(Self { barcodes })
206 }
207}
208
209#[derive(Debug)]
210#[allow(dead_code)]
211struct ChromosomeStats {
212 chrom: String,
213 length: u64,
214 mapped: u64,
215 unmapped: u64,
216}
217
218pub fn bam_header(file_path: PathBuf) -> Result<sam::Header> {
223 if !file_path.exists() {
230 return Err(anyhow::Error::from(std::io::Error::new(
231 std::io::ErrorKind::NotFound,
232 format!("File not found: {}", file_path.display()),
233 )));
234 };
235
236 let mut reader = bam::io::indexed_reader::Builder::default()
237 .build_from_path(file_path.clone())
238 .expect("Failed to open file");
239
240 let header = match reader.read_header() {
241 std::result::Result::Ok(header) => header,
242 Err(e) => {
243 debug!("Failed to read header using noodels falling back to samtools: {e}");
244
245 let header_samtools = std::process::Command::new("samtools")
246 .arg("view")
247 .arg("-H")
248 .arg(file_path.clone())
249 .output()
250 .expect("Failed to run samtools")
251 .stdout;
252
253 let header_str =
254 String::from_utf8(header_samtools).expect("Failed to convert header to string");
255
256 let header_string =
258 header_str.replace("@HD\tSO:coordinate\n", "@HD\tVN:1.6\tSO:coordinate\n");
259 let header_str = header_string.as_bytes();
260 let mut reader = sam::io::Reader::new(header_str);
261 reader
262 .read_header()
263 .expect("Failed to read header with samtools")
264 }
265 };
266 Ok(header)
267}
268
269pub struct BamStats {
271 #[allow(dead_code)]
275 file_path: PathBuf,
276
277 header: sam::Header,
279
280 #[allow(dead_code)]
282 contigs: Vec<Contig<String, ReqStrand>>,
283
284 chrom_stats: HashMap<String, ChromosomeStats>,
286
287 n_reads: u64,
289 n_mapped: u64,
291 n_unmapped: u64,
293}
294
295impl BamStats {
296 pub fn new(file_path: PathBuf) -> Result<Self> {
300 let header = bam_header(file_path.clone())?;
302
303 let contigs = header
305 .reference_sequences()
306 .iter()
307 .map(|(name, map)| {
308 let name = name.to_string();
309 let length = map.length().get();
310 Contig::new(name, 1, length, ReqStrand::Forward)
311 })
312 .collect();
313
314 let bam_reader = bam::io::indexed_reader::Builder::default()
316 .build_from_path(file_path.clone())
317 .expect("Failed to open file");
318 let index = bam_reader.index();
319
320 let mut chrom_stats = HashMap::default();
322
323 for ((reference_sequence_name_buf, reference_sequence), index_reference_sequence) in header
324 .reference_sequences()
325 .iter()
326 .zip(index.reference_sequences())
327 {
328 let reference_sequence_name = std::str::from_utf8(reference_sequence_name_buf)
329 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidInput, e))?;
330
331 let (mapped_record_count, unmapped_record_count) = index_reference_sequence
332 .metadata()
333 .map(|m| (m.mapped_record_count(), m.unmapped_record_count()))
334 .unwrap_or_default();
335
336 let stats = ChromosomeStats {
337 chrom: reference_sequence_name.to_string(),
338 length: usize::from(reference_sequence.length()) as u64,
339 mapped: mapped_record_count,
340 unmapped: unmapped_record_count,
341 };
342
343 chrom_stats.insert(reference_sequence_name.to_string(), stats);
344 }
345
346 let mut unmapped_record_count = index.unplaced_unmapped_record_count().unwrap_or_default();
347
348 let mut mapped_record_count = 0;
349
350 for (_, stats) in chrom_stats.iter() {
351 mapped_record_count += stats.mapped;
352 unmapped_record_count += stats.unmapped;
353 }
354
355 let n_reads = mapped_record_count + unmapped_record_count;
356
357 Ok(Self {
358 file_path,
359 header,
360 contigs,
361 chrom_stats,
362 n_reads,
363 n_mapped: mapped_record_count,
364 n_unmapped: unmapped_record_count,
365 })
366 }
367
368 pub fn estimate_genome_chunk_length(&self, bin_size: u64) -> Result<u64> {
373 let stats = &self.chrom_stats;
379 let genome_length = stats.values().map(|x| x.length).sum::<u64>();
380 let max_reads_per_bp = self.n_mapped as f64 / genome_length as f64;
381 let genome_chunk_length = 5e6_f64.min(2e6 / (max_reads_per_bp));
382
383 let correction = genome_chunk_length % bin_size as f64;
384 let genome_chunk_length = genome_chunk_length - correction;
385
386 Ok(genome_chunk_length as u64)
387 }
388
389 pub fn genome_chunks(&self, bin_size: u64) -> Result<Vec<Region>> {
391 let genome_chunk_length = self.estimate_genome_chunk_length(bin_size)?;
392 let chrom_chunks = self
393 .chrom_stats
394 .iter()
395 .flat_map(|(chrom, stats)| {
396 let mut chunks = Vec::new();
397 let mut chunk_start = 1;
398 let chrom_end = stats.length;
399
400 while chunk_start <= chrom_end {
401 let chunk_end = chunk_start + genome_chunk_length - 1; let chunk_end = chunk_end.min(chrom_end); let start = Position::try_from(chunk_start as usize).unwrap();
406 let end = Position::try_from(chunk_end as usize).unwrap();
407
408 let region = Region::new(&*chrom.clone(), start..=end);
409 chunks.push(region);
410 chunk_start = chunk_end + 1; }
412 chunks
413 })
414 .collect();
415
416 Ok(chrom_chunks)
417 }
418
419 pub fn chromosome_chunks(&self, chrom: &str, bin_size: u64) -> Result<Vec<Region>> {
421 let genome_chunk_length = self.estimate_genome_chunk_length(bin_size)?;
422 let stats = self
423 .chrom_stats
424 .get(chrom)
425 .ok_or_else(|| anyhow::Error::msg(format!("Chromosome {chrom} not found")))?;
426
427 let mut chunks = Vec::new();
428 let mut chunk_start = 1;
429 let chrom_end = stats.length;
430
431 while chunk_start <= chrom_end {
432 let chunk_end = chunk_start + genome_chunk_length - 1; let chunk_end = chunk_end.min(chrom_end); let start = Position::try_from(chunk_start as usize).unwrap();
437 let end = Position::try_from(chunk_end as usize).unwrap();
438
439 let region = Region::new(chrom, start..=end);
440 chunks.push(region);
441 chunk_start = chunk_end + 1; }
443
444 Ok(chunks)
445 }
446
447 pub fn chromosome_id_to_chromosome_name_mapping(&self) -> HashMap<usize, String> {
449 let mut ref_id_mapping = HashMap::default();
450 for (i, (name, _)) in self.header.reference_sequences().iter().enumerate() {
451 let name = std::str::from_utf8(name)
452 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
453 .unwrap()
454 .to_string();
455 ref_id_mapping.insert(i, name);
456 }
457 ref_id_mapping
458 }
459
460 pub fn chromosome_name_to_id_mapping(&self) -> Result<HashMap<String, usize>> {
462 let mut mapping = HashMap::default();
463 let id_to_name = self.chromosome_id_to_chromosome_name_mapping();
464
465 for (id, name) in id_to_name.iter() {
467 mapping.insert(name.clone(), *id);
468 }
469
470 Ok(mapping)
471 }
472
473 pub fn chromsizes_ref_id(&self) -> Result<HashMap<usize, u64>> {
475 let mut chromsizes = HashMap::default();
476 for (i, (_, map)) in self.header.reference_sequences().iter().enumerate() {
477 let length = map.length().get();
478 chromsizes.insert(i, length as u64);
479 }
480 Ok(chromsizes)
481 }
482
483 pub fn chromsizes_ref_name(&self) -> Result<HashMap<String, u64>> {
485 let mut chromsizes = HashMap::default();
486 for (name, map) in self.header.reference_sequences() {
487 let length = map.length().get();
488 let name = std::str::from_utf8(name)
489 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
490 .unwrap()
491 .to_string();
492 chromsizes.insert(name, length as u64);
493 }
494 Ok(chromsizes)
495 }
496
497 pub fn n_mapped(&self) -> u64 {
499 self.n_mapped
500 }
501
502 pub fn n_unmapped(&self) -> u64 {
504 self.n_unmapped
505 }
506
507 pub fn n_total_reads(&self) -> u64 {
509 self.n_reads
510 }
511
512 pub fn header(&self) -> &sam::Header {
514 &self.header
515 }
516
517 pub fn is_paired_end(&self) -> Result<bool> {
522 let mut reader = bam::io::reader::Builder
523 .build_from_path(&self.file_path)
524 .context("Failed to open BAM file to check paired-end status")?;
525 let _header = reader.read_header().context("Failed to read BAM header")?;
526 for (i, record) in reader.records().enumerate() {
527 let record = match record {
528 Ok(r) => r,
529 Err(_) => break,
530 };
531 if record.flags().is_segmented() {
532 return Ok(true);
533 }
534 if i >= 1000 {
535 break;
536 }
537 }
538 Ok(false)
539 }
540}
541
542pub enum FileType {
544 Bedgraph,
546 Bigwig,
548 TSV,
550}
551
552impl std::str::FromStr for FileType {
553 type Err = String;
554
555 fn from_str(s: &str) -> Result<FileType, String> {
556 match s.to_lowercase().as_str() {
557 "bedgraph" => Ok(FileType::Bedgraph),
558 "bdg" => Ok(FileType::Bedgraph),
559 "bigwig" => Ok(FileType::Bigwig),
560 "bw" => Ok(FileType::Bigwig),
561 "tsv" => Ok(FileType::TSV),
562 "txt" => Ok(FileType::TSV),
563 _ => Err(format!("Unknown file type: {s}")),
564 }
565 }
566}
567
568pub fn regions_to_lapper(regions: Vec<Region>) -> Result<HashMap<String, Lapper<usize, u32>>> {
570 let mut lapper: HashMap<String, Lapper<usize, u32>> = HashMap::default();
571 let mut intervals: HashMap<String, Vec<Iv>> = HashMap::default();
572
573 for reg in regions {
574 let chrom = reg.name().to_string();
575 let start = reg.start().map(|x| x.get());
576 let end = reg.end().map(|x| x.get());
577
578 let start = match start {
579 Bound::Included(start) => start,
580 Bound::Excluded(start) => start + 1,
581 _ => 0,
582 };
583
584 let end = match end {
585 Bound::Included(end) => end,
586 Bound::Excluded(end) => end - 1,
587 _ => 0,
588 };
589
590 let iv = Iv {
591 start,
592 stop: end,
593 val: 0,
594 };
595 intervals.entry(chrom.clone()).or_default().push(iv);
596 }
597
598 for (chrom, ivs) in intervals.iter() {
599 let lap = Lapper::new(ivs.to_vec());
600 lapper.insert(chrom.clone(), lap);
601 }
602
603 Ok(lapper)
604}
605
606pub fn bed_to_lapper(bed: PathBuf) -> Result<HashMap<String, Lapper<usize, u32>>> {
608 let reader = std::fs::File::open(bed)?;
610 let buf_reader = std::io::BufReader::new(reader);
611 let mut bed_reader = bed::io::Reader::<4, _>::new(buf_reader);
612 let mut record = bed::Record::default();
613 let mut intervals: HashMap<String, Vec<Iv>> = HashMap::default();
614 let mut lapper: HashMap<String, Lapper<usize, u32>> = HashMap::default();
615
616 while bed_reader.read_record(&mut record)? != 0 {
617 let chrom = record.reference_sequence_name().to_string();
618 let start = record.feature_start()?;
619 let end = record
620 .feature_end()
621 .context("Failed to get feature end")??;
622
623 let iv = Iv {
624 start: start.get(),
625 stop: end.get(),
626 val: 0,
627 };
628 intervals.entry(chrom.clone()).or_default().push(iv);
629 }
630
631 for (chrom, ivs) in intervals.iter() {
632 let lap = Lapper::new(ivs.to_vec());
633 lapper.insert(chrom.clone(), lap);
634 }
635 if lapper.is_empty() {
637 return Err(anyhow::Error::msg("Lapper is empty"));
638 }
639
640 Ok(lapper)
641}
642
643pub fn convert_lapper_chrom_names_to_ids(
647 lapper: HashMap<String, Lapper<usize, u32>>,
648 bam_stats: &BamStats,
649) -> Result<HashMap<usize, Lapper<usize, u32>>> {
650 let chrom_id_mapping = bam_stats.chromosome_name_to_id_mapping()?;
652 let mut lapper_chrom_id: HashMap<usize, Lapper<usize, u32>> = HashMap::default();
653 for (chrom, lap) in lapper.iter() {
654 let chrom_id = chrom_id_mapping.get(chrom).ok_or_else(|| {
655 anyhow::Error::msg(format!("Chromosome {chrom} not found in BAM file"))
656 })?;
657 lapper_chrom_id.insert(*chrom_id, lap.clone());
658 }
659 Ok(lapper_chrom_id)
660}
661
662pub fn get_bam_header<P>(file_path: P) -> Result<sam::Header>
667where
668 P: AsRef<Path>,
669{
670 let file_path = file_path.as_ref();
671
672 if !file_path.exists() {
674 return Err(anyhow::Error::from(std::io::Error::new(
675 std::io::ErrorKind::NotFound,
676 format!("File not found: {}", file_path.display()),
677 )));
678 };
679
680 let mut reader = bam::io::indexed_reader::Builder::default()
681 .build_from_path(file_path)
682 .expect("Failed to open file");
683
684 let header = match reader.read_header() {
685 std::result::Result::Ok(header) => header,
686 Err(e) => {
687 debug!("Failed to read header using noodels falling back to samtools: {e}");
688
689 let header_samtools = std::process::Command::new("samtools")
690 .arg("view")
691 .arg("-H")
692 .arg(file_path)
693 .output()
694 .expect("Failed to run samtools")
695 .stdout;
696
697 let header_str =
698 String::from_utf8(header_samtools).expect("Failed to convert header to string");
699
700 let header_string =
702 header_str.replace("@HD\tSO:coordinate\n", "@HD\tVN:1.6\tSO:coordinate\n");
703 let header_str = header_string.as_bytes();
704 let mut reader = sam::io::Reader::new(header_str);
705 reader
706 .read_header()
707 .expect("Failed to read header with samtools")
708 }
709 };
710 Ok(header)
711}
712
713pub fn add_bamnado_program_group(header: &sam::Header) -> Result<sam::Header> {
715 let mut output_header = header.clone();
716 let mut program = Map::<Program>::default();
717
718 program
719 .other_fields_mut()
720 .insert(pg_tag::NAME, BAMNADO_PROGRAM_NAME.into());
721 program
722 .other_fields_mut()
723 .insert(pg_tag::VERSION, BAMNADO_VERSION.into());
724
725 if let Some(command_line) = bamnado_command_line() {
726 program
727 .other_fields_mut()
728 .insert(pg_tag::COMMAND_LINE, command_line.into());
729 }
730
731 output_header
732 .programs_mut()
733 .add(BAMNADO_PROGRAM_ID, program)
734 .context("Failed to append bamnado @PG header record")?;
735
736 Ok(output_header)
737}
738
739fn bamnado_command_line() -> Option<String> {
740 let args = std::env::args_os()
741 .map(|arg| arg.to_string_lossy().into_owned())
742 .collect::<Vec<_>>();
743
744 (!args.is_empty()).then(|| args.join(" "))
745}
746
747#[cfg(test)]
748mod tests {
749 use super::*;
750 use noodles::sam::header::record::value::map::program::tag as pg_tag;
751
752 fn test_bam() -> PathBuf {
753 PathBuf::from(env!("CARGO_MANIFEST_DIR"))
754 .parent()
755 .expect("Failed to get workspace root")
756 .join("test/data/test.bam")
757 }
758
759 #[test]
760 fn test_add_bamnado_program_group_appends_pg_record() {
761 let header = sam::Header::default();
762 let output_header =
763 add_bamnado_program_group(&header).expect("Failed to add bamnado program group");
764
765 let program = output_header
766 .programs()
767 .as_ref()
768 .get(&b"bamnado"[..])
769 .expect("Missing bamnado @PG record");
770
771 assert_eq!(
772 program
773 .other_fields()
774 .get(&pg_tag::NAME)
775 .map(|v| v.as_ref()),
776 Some(&b"bamnado"[..])
777 );
778 assert_eq!(
779 program
780 .other_fields()
781 .get(&pg_tag::VERSION)
782 .map(|v| v.as_ref()),
783 Some(env!("CARGO_PKG_VERSION").as_bytes())
784 );
785 assert!(program.other_fields().contains_key(&pg_tag::COMMAND_LINE));
786 }
787
788 #[test]
789 fn test_is_paired_end_returns_true_for_paired_bam() {
790 let stats = BamStats::new(test_bam()).expect("Failed to create BamStats");
791 assert!(
792 stats.is_paired_end().expect("is_paired_end failed"),
793 "test.bam should be detected as paired-end"
794 );
795 }
796
797 #[test]
798 fn test_bam_stats_reads_are_nonzero() {
799 let stats = BamStats::new(test_bam()).expect("Failed to create BamStats");
800 assert!(stats.n_mapped() > 0, "Expected mapped reads in test.bam");
801 }
802}