use needletail::parse_fastx_file;
use std;
use std::collections::HashSet;
use std::io::Read;
use std::path::Path;
use std::process;
use bam_generator::is_long_read_only_preset;
use bam_generator::rammap_preset_name;
use bam_generator::MappingProgram;
use CONCATENATED_FASTA_FILE_SEPARATOR;
use tempdir::TempDir;
use tempfile::{Builder, NamedTempFile};
pub trait MappingIndex {
fn index_path(&self) -> &String;
fn command_prefix(&self) -> &str {
""
}
}
pub struct VanillaIndexStruct {
index_path_internal: String,
}
impl VanillaIndexStruct {
pub fn new(reference_path: &str) -> VanillaIndexStruct {
VanillaIndexStruct {
index_path_internal: reference_path.to_string(),
}
}
}
impl MappingIndex for VanillaIndexStruct {
fn index_path(&self) -> &String {
&self.index_path_internal
}
}
pub struct TemporaryIndexStruct {
#[allow(dead_code)] tempdir: TempDir,
index_path_internal: String,
}
impl TemporaryIndexStruct {
pub fn new(
mapping_program: MappingProgram,
reference_path: &str,
num_threads: Option<u16>,
index_creation_options: Option<&str>,
) -> TemporaryIndexStruct {
let td =
TempDir::new("coverm-mapping-index").expect("Unable to create temporary directory");
let index_path = std::path::Path::new(td.path()).join(
std::path::Path::new(reference_path)
.file_name()
.expect("Failed to glean file stem from reference DB. Strange."),
);
run_index_command(
mapping_program,
reference_path,
&index_path,
num_threads,
index_creation_options,
);
TemporaryIndexStruct {
index_path_internal: index_path.to_string_lossy().to_string(),
tempdir: td,
}
}
}
fn build_index_command(
mapping_program: MappingProgram,
reference_path: &str,
index_path: &Path,
num_threads: Option<u16>,
index_creation_options: Option<&str>,
) -> Option<std::process::Command> {
let mut cmd = match mapping_program {
MappingProgram::BWA_MEM => std::process::Command::new("bwa"),
MappingProgram::BWA_MEM2 => std::process::Command::new("bwa-mem2"),
MappingProgram::MINIMAP2_SR
| MappingProgram::MINIMAP2_ONT
| MappingProgram::MINIMAP2_PB
| MappingProgram::MINIMAP2_HIFI
| MappingProgram::MINIMAP2_LR_HQ
| MappingProgram::MINIMAP2_NO_PRESET => std::process::Command::new("minimap2"),
MappingProgram::MINIBWA => std::process::Command::new("minibwa"),
MappingProgram::RAMMAP_SR
| MappingProgram::RAMMAP_ONT
| MappingProgram::RAMMAP_PB
| MappingProgram::RAMMAP_HIFI
| MappingProgram::RAMMAP_LR_HQ
| MappingProgram::RAMMAP_NO_PRESET => std::process::Command::new("rammap"),
MappingProgram::STROBEALIGN => {
warn!("STROBEALIGN pre-indexing is not supported currently, so skipping index generation.");
return None;
}
};
match &mapping_program {
MappingProgram::BWA_MEM | MappingProgram::BWA_MEM2 => {
cmd.arg("index")
.arg("-p")
.arg(index_path)
.arg(reference_path);
}
MappingProgram::MINIBWA => {
cmd.arg("index");
if let Some(t) = num_threads {
cmd.arg("-t").arg(format!("{t}"));
}
cmd.arg(reference_path).arg(index_path);
}
MappingProgram::RAMMAP_SR
| MappingProgram::RAMMAP_ONT
| MappingProgram::RAMMAP_PB
| MappingProgram::RAMMAP_HIFI
| MappingProgram::RAMMAP_LR_HQ
| MappingProgram::RAMMAP_NO_PRESET => {
if let Some(preset) = rammap_preset_name(mapping_program) {
cmd.arg("-x").arg(preset);
}
cmd.arg("--dump-index").arg(index_path).arg(reference_path);
}
MappingProgram::MINIMAP2_SR
| MappingProgram::MINIMAP2_ONT
| MappingProgram::MINIMAP2_HIFI
| MappingProgram::MINIMAP2_PB
| MappingProgram::MINIMAP2_LR_HQ
| MappingProgram::MINIMAP2_NO_PRESET => {
match &mapping_program {
MappingProgram::MINIMAP2_SR => {
cmd.arg("-x").arg("sr");
}
MappingProgram::MINIMAP2_ONT => {
cmd.arg("-x").arg("map-ont");
}
MappingProgram::MINIMAP2_HIFI => {
cmd.arg("-x").arg("map-hifi");
}
MappingProgram::MINIMAP2_PB => {
cmd.arg("-x").arg("map-pb");
}
MappingProgram::MINIMAP2_LR_HQ => {
cmd.arg("-x").arg("lr:hq");
}
MappingProgram::MINIMAP2_NO_PRESET
| MappingProgram::BWA_MEM
| MappingProgram::BWA_MEM2
| MappingProgram::STROBEALIGN
| MappingProgram::MINIBWA
| MappingProgram::RAMMAP_SR
| MappingProgram::RAMMAP_ONT
| MappingProgram::RAMMAP_PB
| MappingProgram::RAMMAP_HIFI
| MappingProgram::RAMMAP_LR_HQ
| MappingProgram::RAMMAP_NO_PRESET => {}
};
if let Some(t) = num_threads {
cmd.arg("-t").arg(format!("{t}"));
}
cmd.arg("-d").arg(index_path).arg(reference_path);
}
MappingProgram::STROBEALIGN => unreachable!(),
};
if let Some(params) = index_creation_options {
for s in params.split_whitespace() {
cmd.arg(s);
}
};
Some(cmd)
}
fn run_index_command(
mapping_program: MappingProgram,
reference_path: &str,
index_path: &Path,
num_threads: Option<u16>,
index_creation_options: Option<&str>,
) {
info!("Generating {mapping_program:?} index for {reference_path} ..");
let cmd = match build_index_command(
mapping_program,
reference_path,
index_path,
num_threads,
index_creation_options,
) {
Some(cmd) => cmd,
None => return,
};
execute_index_command(cmd, mapping_program);
}
fn execute_index_command(mut cmd: std::process::Command, mapping_program: MappingProgram) {
cmd.stdout(std::process::Stdio::piped());
cmd.stderr(std::process::Stdio::piped());
debug!("Running DB indexing command: {cmd:?}");
let mut process = cmd
.spawn()
.unwrap_or_else(|_| panic!("Failed to start {:?} index process", mapping_program));
let es = process.wait().unwrap_or_else(|_| {
panic!(
"Failed to glean exitstatus from failing {:?} index process",
mapping_program
)
});
if !es.success() {
error!("Error when running {mapping_program:?} index process.");
let mut err = String::new();
process
.stderr
.unwrap_or_else(|| {
panic!(
"Failed to grab stderr from failed {:?} index process",
mapping_program
)
})
.read_to_string(&mut err)
.expect("Failed to read stderr into string");
error!("The STDERR was: {err:?}");
error!("Cannot continue after {mapping_program:?} index failed.");
process::exit(1);
}
info!("Finished generating {mapping_program:?} index.");
}
fn generate_strobealign_persistent_index(
reference_path: &str,
output_directory: &str,
num_threads: Option<u16>,
index_creation_options: Option<&str>,
) -> String {
let reference_file_name = std::path::Path::new(reference_path)
.file_name()
.expect("Failed to glean file name from reference path");
let copied_reference = std::path::Path::new(output_directory).join(reference_file_name);
if copied_reference == std::path::Path::new(reference_path) {
error!(
"The strobealign database for {reference_path} would overwrite the reference \
itself. Please choose an output directory other than the reference's directory."
);
process::exit(1);
}
info!(
"Copying reference {} into {} so the strobealign database is self-contained ..",
reference_path, output_directory
);
std::fs::copy(reference_path, &copied_reference).unwrap_or_else(|e| {
panic!(
"Failed to copy reference {} into output directory: {}",
reference_path, e
)
});
info!("Generating STROBEALIGN index for {reference_path} ..");
let mut cmd = std::process::Command::new("strobealign");
if let Some(t) = num_threads {
cmd.arg("-t").arg(format!("{t}"));
}
cmd.arg("--create-index").arg(&copied_reference);
if let Some(params) = index_creation_options {
for s in params.split_whitespace() {
cmd.arg(s);
}
};
execute_index_command(cmd, MappingProgram::STROBEALIGN);
copied_reference.to_string_lossy().to_string()
}
impl MappingIndex for TemporaryIndexStruct {
fn index_path(&self) -> &String {
&self.index_path_internal
}
}
impl Drop for TemporaryIndexStruct {
fn drop(&mut self) {
debug!(
"Dropping index tempdir {}",
self.tempdir.path().to_string_lossy()
)
}
}
fn check_for_bwa_index_existence(reference_path: &str, mapping_program: &MappingProgram) -> bool {
let bwa_extensions = match mapping_program {
MappingProgram::BWA_MEM => vec!["amb", "ann", "bwt", "pac", "sa"],
MappingProgram::BWA_MEM2 => vec!["0123", "amb", "ann", "bwt.2bit.64", "pac"],
MappingProgram::MINIBWA => vec!["l2b", "mbw"],
_ => unreachable!(),
};
let num_extensions = bwa_extensions.len();
let mut num_existing: usize = 0;
for extension in bwa_extensions {
if std::path::Path::new(&format!("{reference_path}.{extension}")).exists() {
num_existing += 1;
}
}
if num_existing == 0 {
false
} else if num_existing == num_extensions {
true
} else {
error!("BWA index appears to be incomplete, cannot continue.");
process::exit(1);
}
}
pub fn check_reference_existence(reference_path: &str, mapping_program: &MappingProgram) {
let ref_path = std::path::Path::new(reference_path);
match mapping_program {
MappingProgram::BWA_MEM | MappingProgram::BWA_MEM2 | MappingProgram::MINIBWA => {
if check_for_bwa_index_existence(reference_path, mapping_program) {
return;
}
}
MappingProgram::MINIMAP2_SR
| MappingProgram::MINIMAP2_ONT
| MappingProgram::MINIMAP2_HIFI
| MappingProgram::MINIMAP2_PB
| MappingProgram::MINIMAP2_LR_HQ
| MappingProgram::MINIMAP2_NO_PRESET
| MappingProgram::STROBEALIGN
| MappingProgram::RAMMAP_SR
| MappingProgram::RAMMAP_ONT
| MappingProgram::RAMMAP_PB
| MappingProgram::RAMMAP_HIFI
| MappingProgram::RAMMAP_LR_HQ
| MappingProgram::RAMMAP_NO_PRESET => {}
};
if !ref_path.exists() {
panic!(
"The reference specified '{}' does not appear to exist",
&reference_path
);
} else if !ref_path.is_file() {
panic!(
"The reference specified '{}' should be a file, not e.g. a directory",
&reference_path
);
}
}
pub fn generate_bwa_index(
reference_path: &str,
index_creation_parameters: Option<&str>,
mapping_program: MappingProgram,
) -> Box<dyn MappingIndex> {
if check_for_bwa_index_existence(reference_path, &mapping_program) {
info!("BWA index appears to be complete, so going ahead and using it.");
Box::new(VanillaIndexStruct::new(reference_path))
} else {
Box::new(TemporaryIndexStruct::new(
mapping_program,
reference_path,
None,
index_creation_parameters,
))
}
}
pub fn generate_minibwa_index(
reference_path: &str,
num_threads: Option<u16>,
index_creation_parameters: Option<&str>,
) -> Box<dyn MappingIndex> {
if check_for_bwa_index_existence(reference_path, &MappingProgram::MINIBWA) {
info!("minibwa index appears to be complete, so going ahead and using it.");
Box::new(VanillaIndexStruct::new(reference_path))
} else {
Box::new(TemporaryIndexStruct::new(
MappingProgram::MINIBWA,
reference_path,
num_threads,
index_creation_parameters,
))
}
}
pub fn generate_minimap2_index(
reference_path: &str,
num_threads: Option<u16>,
index_creation_parameters: Option<&str>,
mapping_program: MappingProgram,
) -> Box<dyn MappingIndex> {
Box::new(TemporaryIndexStruct::new(
mapping_program,
reference_path,
num_threads,
index_creation_parameters,
))
}
#[derive(Clone, Copy, Debug)]
pub enum MakedbUsageMode {
Contig,
Genome,
}
pub fn makedb_usage_command(
mapping_program: MappingProgram,
db_path: &str,
mode: MakedbUsageMode,
) -> String {
let mapper = mapping_program_db_name(mapping_program);
let index_flag = match mapping_program {
MappingProgram::MINIMAP2_SR
| MappingProgram::MINIMAP2_ONT
| MappingProgram::MINIMAP2_PB
| MappingProgram::MINIMAP2_HIFI
| MappingProgram::MINIMAP2_LR_HQ
| MappingProgram::MINIMAP2_NO_PRESET => " --minimap2-reference-is-index",
MappingProgram::STROBEALIGN => " --strobealign-use-index",
MappingProgram::BWA_MEM
| MappingProgram::BWA_MEM2
| MappingProgram::MINIBWA
| MappingProgram::RAMMAP_SR
| MappingProgram::RAMMAP_ONT
| MappingProgram::RAMMAP_PB
| MappingProgram::RAMMAP_HIFI
| MappingProgram::RAMMAP_LR_HQ
| MappingProgram::RAMMAP_NO_PRESET => "",
};
let (subcommand, extra) = match mode {
MakedbUsageMode::Contig => ("contig", ""),
MakedbUsageMode::Genome => ("genome", " -s '~'"),
};
let reads = if is_long_read_only_preset(mapping_program) {
"--single reads.fq"
} else {
"-1 read1.fq -2 read2.fq"
};
format!("coverm {subcommand} -p {mapper} --reference {db_path}{index_flag}{extra} {reads}")
}
pub fn mapping_program_db_name(mapping_program: MappingProgram) -> &'static str {
match mapping_program {
MappingProgram::BWA_MEM => "bwa-mem",
MappingProgram::BWA_MEM2 => "bwa-mem2",
MappingProgram::MINIMAP2_SR => "minimap2-sr",
MappingProgram::MINIMAP2_ONT => "minimap2-ont",
MappingProgram::MINIMAP2_PB => "minimap2-pb",
MappingProgram::MINIMAP2_HIFI => "minimap2-hifi",
MappingProgram::MINIMAP2_LR_HQ => "minimap2-lr-hq",
MappingProgram::MINIMAP2_NO_PRESET => "minimap2-no-preset",
MappingProgram::STROBEALIGN => "strobealign",
MappingProgram::MINIBWA => "minibwa",
MappingProgram::RAMMAP_SR => "rammap-sr",
MappingProgram::RAMMAP_ONT => "rammap-ont",
MappingProgram::RAMMAP_PB => "rammap-pb",
MappingProgram::RAMMAP_HIFI => "rammap-hifi",
MappingProgram::RAMMAP_LR_HQ => "rammap-lr-hq",
MappingProgram::RAMMAP_NO_PRESET => "rammap-no-preset",
}
}
pub fn generate_persistent_index(
mapping_program: MappingProgram,
reference_path: &str,
output_directory: &str,
num_threads: Option<u16>,
index_creation_options: Option<&str>,
) -> String {
if let MappingProgram::STROBEALIGN = mapping_program {
return generate_strobealign_persistent_index(
reference_path,
output_directory,
num_threads,
index_creation_options,
);
}
let reference_stem = std::path::Path::new(reference_path)
.file_name()
.expect("Failed to glean file name from reference path")
.to_string_lossy();
let db_program_name = mapping_program_db_name(mapping_program);
let index_path = match mapping_program {
MappingProgram::MINIMAP2_SR
| MappingProgram::MINIMAP2_ONT
| MappingProgram::MINIMAP2_PB
| MappingProgram::MINIMAP2_HIFI
| MappingProgram::MINIMAP2_LR_HQ
| MappingProgram::MINIMAP2_NO_PRESET => std::path::Path::new(output_directory)
.join(format!("{reference_stem}.{db_program_name}.mmi")),
MappingProgram::BWA_MEM | MappingProgram::BWA_MEM2 => {
std::path::Path::new(output_directory)
.join(format!("{reference_stem}.{db_program_name}"))
}
MappingProgram::MINIBWA => std::path::Path::new(output_directory)
.join(format!("{reference_stem}.{db_program_name}")),
MappingProgram::RAMMAP_SR
| MappingProgram::RAMMAP_ONT
| MappingProgram::RAMMAP_PB
| MappingProgram::RAMMAP_HIFI
| MappingProgram::RAMMAP_LR_HQ
| MappingProgram::RAMMAP_NO_PRESET => std::path::Path::new(output_directory)
.join(format!("{reference_stem}.{db_program_name}.idx")),
MappingProgram::STROBEALIGN => unreachable!(),
};
run_index_command(
mapping_program,
reference_path,
&index_path,
num_threads,
index_creation_options,
);
index_path.to_string_lossy().to_string()
}
pub fn generate_concatenated_fasta_file(fasta_file_paths: &Vec<String>) -> NamedTempFile {
let tmpfile: NamedTempFile = Builder::new()
.prefix("coverm-concatenated-fasta")
.tempfile()
.unwrap();
write_concatenated_fasta(&tmpfile, fasta_file_paths);
tmpfile
}
pub fn generate_concatenated_fasta_file_at(fasta_file_paths: &Vec<String>, output_path: &Path) {
let file = std::fs::File::create(output_path).unwrap_or_else(|e| {
error!(
"Failed to create concatenated FASTA file at {}: {}",
output_path.display(),
e
);
process::exit(1);
});
write_concatenated_fasta(std::io::BufWriter::new(file), fasta_file_paths);
}
fn write_concatenated_fasta<W: std::io::Write>(out: W, fasta_file_paths: &Vec<String>) {
let mut something_written_at_all = false;
{
let mut writer = bio::io::fasta::Writer::new(out);
let mut genome_names: HashSet<String> = HashSet::new();
for file in fasta_file_paths {
let mut something_written = false;
let path = std::path::Path::new(file);
let mut reader = parse_fastx_file(path)
.unwrap_or_else(|_| panic!("Unable to read fasta file {}", file));
let mut genome_name1 =
String::from(path.to_str().expect("File name string conversion problem"));
if let Some(i) = genome_name1.rfind(".gz") {
genome_name1.truncate(i);
} else if let Some(i) = genome_name1.rfind(".bz") {
genome_name1.truncate(i);
} else if let Some(i) = genome_name1.rfind(".xz") {
genome_name1.truncate(i);
}
let path1 = Path::new(&genome_name1);
let genome_name = String::from(
path1
.file_stem()
.expect("Problem while determining file stem")
.to_str()
.expect("File name string conversion problem"),
);
if genome_names.contains(&genome_name) {
error!("The genome name {genome_name} was derived from >1 file");
process::exit(1);
}
while let Some(record) = reader.next() {
let record_expected = record
.unwrap_or_else(|_| panic!("Failed to parse record in fasta file {:?}", path));
if record_expected.format() != needletail::parser::Format::Fasta {
panic!(
"File {:?} is not a fasta file, but a {:?}",
path,
record_expected.format()
);
}
something_written = true;
something_written_at_all = true;
let contig_name = String::from(
std::str::from_utf8(record_expected.id())
.expect("UTF-8 conversion problem in contig name"),
);
writer
.write(
&format!(
"{}{}{}",
genome_name,
CONCATENATED_FASTA_FILE_SEPARATOR,
match contig_name.split_once(' ') {
Some((contig, _)) => contig.to_string(),
None => contig_name,
}
),
None,
&record_expected.seq(),
)
.unwrap()
}
genome_names.insert(genome_name);
if !something_written {
error!(
"FASTA file {file} appears to be empty as no sequences were contained in it"
);
process::exit(1);
}
}
}
if !something_written_at_all {
error!("Concatenated FASTA file to use as a reference is empty");
process::exit(1);
}
}
pub struct PregeneratedStrobealignIndexStruct {
index_path_internal: String,
}
impl PregeneratedStrobealignIndexStruct {
pub fn new(reference_path: &str) -> PregeneratedStrobealignIndexStruct {
PregeneratedStrobealignIndexStruct {
index_path_internal: reference_path.to_string(),
}
}
}
impl MappingIndex for PregeneratedStrobealignIndexStruct {
fn index_path(&self) -> &String {
&self.index_path_internal
}
fn command_prefix(&self) -> &str {
"--use-index"
}
}