use std::collections::HashSet;
use std::io::{ErrorKind, Result, Write};
use std::path::{Path, PathBuf};
use indexmap::IndexMap;
#[cfg(feature = "python-extension")]
use pyo3::prelude::*;
use rust_htslib::bam::{index, Header, IndexedReader, Read};
use rust_htslib::htslib;
use crate::cli::DepthOptions;
use crate::io::get_writer;
use crate::utils::styled_progress_bar;
fn add_extension(path: &mut PathBuf, extension: impl AsRef<Path>) {
match path.extension() {
Some(ext) => {
let mut ext = ext.to_os_string();
ext.push(".");
ext.push(extension.as_ref());
path.set_extension(ext)
}
None => path.set_extension(extension.as_ref()),
};
}
pub fn create_index(bam_path: &PathBuf) {
let mut csi = PathBuf::from(bam_path);
add_extension(&mut csi, "csi");
if Path::new(&csi).exists() {
return;
}
let mut bai = PathBuf::from(bam_path);
add_extension(&mut bai, "bai");
if Path::new(&bai).exists() {
return;
}
match index::build(bam_path, None, index::Type::Csi(14), 1) {
Err(e) => eprintln!("Error writing BAM index: {e:?}"),
Ok(_) => eprintln!("Successfully created BAM index"),
}
}
pub fn open_bam(
bam_path: &Option<PathBuf>,
cram_path: &Option<PathBuf>,
fasta_path: &Option<PathBuf>,
make_index: bool,
) -> IndexedReader {
let bam_cram_path = match bam_path {
None => cram_path.as_ref().unwrap(),
&Some(_) => bam_path.as_ref().unwrap(),
};
if make_index {
create_index(bam_cram_path);
}
let mut reader = IndexedReader::from_path(bam_cram_path).unwrap();
if cram_path.is_some() && fasta_path.is_some() {
reader.set_reference(fasta_path.as_ref().unwrap()).unwrap();
}
reader
}
pub fn reads_from_bam<F: Fn()>(
seq_names: &HashSet<Vec<u8>>,
mut bam: IndexedReader,
callback: &Option<F>,
) -> HashSet<Vec<u8>> {
let mut wanted_reads = HashSet::new();
let total = seq_names.len();
let progress_bar = styled_progress_bar(total, "Locating alignments");
for seq_name in seq_names {
if bam.fetch(seq_name).is_err() {
eprintln!("Sequence {:?} not found in BAM file", seq_name)
}
for read in bam
.rc_records()
.map(|x| x.expect("Failure parsing Bam file"))
.filter(|read| {
read.flags()
& (htslib::BAM_FUNMAP
| htslib::BAM_FSECONDARY
| htslib::BAM_FQCFAIL
| htslib::BAM_FDUP) as u16
== 0
})
{
wanted_reads.insert(read.qname().to_vec());
}
if let Some(cb) = callback {
cb()
}
progress_bar.inc(1);
}
progress_bar.finish();
wanted_reads
}
fn seq_lengths_from_header(
bam: &IndexedReader,
seq_names: &HashSet<Vec<u8>>,
) -> IndexMap<String, usize> {
let header = Header::from_template(bam.header());
let mut seq_lengths: IndexMap<String, usize> = IndexMap::new();
for (_, records) in header.to_hashmap() {
for record in records {
if record.contains_key("SN") {
if !seq_names.is_empty() && !seq_names.contains(&record["SN"].as_bytes().to_vec()) {
continue;
}
seq_lengths
.entry(record["SN"].to_string())
.or_insert(record["LN"].parse::<usize>().unwrap());
}
}
}
seq_lengths
}
#[derive(Clone, Debug)]
#[cfg_attr(feature = "python-extension", pyclass)]
pub struct BinnedCov {
seq_name: String,
bins: Vec<f64>,
bin_count: usize,
last_bin: usize,
seq_length: usize,
step: usize,
}
#[cfg(feature = "python-extension")]
#[pymethods]
impl BinnedCov {
#[getter]
fn seq_name(&self) -> &str {
&self.seq_name
}
#[getter]
fn bins(&self) -> &Vec<f64> {
&self.bins
}
#[getter]
fn bin_count(&self) -> usize {
self.bin_count
}
#[getter]
fn last_bin(&self) -> usize {
self.last_bin
}
#[getter]
fn seq_length(&self) -> usize {
self.seq_length
}
#[getter]
fn step(&self) -> usize {
self.step
}
}
fn depth_to_bed(
raw_cov: Vec<usize>,
length: &usize,
step: usize,
seq_name: &String,
writer: &mut Box<dyn Write>,
) -> Result<()> {
let mut bins: Vec<f64> = vec![];
let mut divisor = step;
let mut end: usize = 0;
let seq_length = length.to_owned();
for cov in raw_cov {
end += step;
if end > seq_length {
divisor -= end - seq_length;
}
bins.push(cov as f64 / divisor as f64);
}
let mut start = 0;
let mut end;
let bin_count = bins.len();
for bin in bins.iter().take(bin_count) {
end = start + step;
if end > seq_length {
end = seq_length;
}
let line = format!("{}\t{}\t{}\t{:.2}", seq_name, start, end, bin);
writeln!(writer, "{}", line)?;
start = end;
}
Ok(())
}
pub fn bed_from_bam<F: Fn()>(
seq_lengths: &IndexMap<String, usize>,
mut bam: IndexedReader,
options: &DepthOptions,
callback: &Option<F>,
) {
let total = seq_lengths.len();
let progress_bar = styled_progress_bar(total, "Locating alignments");
let bin_size = options.bin_size;
let step = bin_size;
let mut writer = get_writer(&options.bed);
for (seq_name, length) in seq_lengths.clone() {
let mut raw_cov: Vec<usize> = vec![];
for _ in (0..length).step_by(step) {
raw_cov.push(0)
}
if bam.fetch(&seq_name).is_err() {
eprintln!("Sequence {:?} not found in BAM file", seq_name)
}
for p in bam.pileup() {
let pileup = p.unwrap();
let bin = pileup.pos() as usize / step;
raw_cov[bin] += pileup.depth() as usize;
}
if let Some(cb) = callback {
cb()
}
match depth_to_bed(raw_cov, &length, step, &seq_name, &mut writer) {
Err(err) if err.kind() == ErrorKind::BrokenPipe => return,
Err(err) => panic!("unable to write {} to bed file: {}", &seq_name, err),
Ok(_) => (),
};
progress_bar.inc(1);
}
progress_bar.finish();
}
fn depth_to_cov(raw_cov: Vec<usize>, length: &usize, step: usize, seq_name: &String) -> BinnedCov {
let mut bins: Vec<f64> = vec![];
let mut divisor = step;
let mut end: usize = 0;
let seq_length = length.to_owned();
for cov in raw_cov {
end += step;
if end > seq_length {
divisor -= end - seq_length;
}
bins.push(cov as f64 / divisor as f64);
}
BinnedCov {
seq_name: seq_name.to_owned(),
step,
bin_count: bins.len(),
bins,
seq_length,
last_bin: divisor,
}
}
pub fn depth_from_bam<F: Fn()>(
seq_lengths: &IndexMap<String, usize>,
mut bam: IndexedReader,
options: &DepthOptions,
callback: &Option<F>,
) -> Vec<BinnedCov> {
let total = seq_lengths.len();
let progress_bar = styled_progress_bar(total, "Locating alignments");
let bin_size = options.bin_size;
let step = bin_size;
let mut binned_covs = vec![];
for (seq_name, length) in seq_lengths.clone() {
let mut raw_cov: Vec<usize> = vec![];
for _ in (0..length).step_by(step) {
raw_cov.push(0)
}
if bam.fetch(&seq_name).is_err() {
eprintln!("Sequence {:?} not found in BAM file", seq_name)
}
for p in bam.pileup() {
let pileup = p.unwrap();
let bin = pileup.pos() as usize / step;
raw_cov[bin] += pileup.depth() as usize;
}
if let Some(cb) = callback {
cb()
}
binned_covs.push(depth_to_cov(raw_cov, &length, step, &seq_name));
progress_bar.inc(1);
}
progress_bar.finish();
binned_covs
}
pub fn get_bed_file<F: Fn()>(
bam: IndexedReader,
seq_names: &HashSet<Vec<u8>>,
options: &DepthOptions,
callback: &Option<F>,
) {
let seq_lengths = seq_lengths_from_header(&bam, seq_names);
bed_from_bam(&seq_lengths, bam, options, callback);
}
pub fn get_depth<F: Fn()>(
bam: IndexedReader,
seq_names: &HashSet<Vec<u8>>,
options: &DepthOptions,
callback: &Option<F>,
) -> Vec<BinnedCov> {
let seq_lengths = seq_lengths_from_header(&bam, seq_names);
depth_from_bam(&seq_lengths, bam, options, callback)
}