rustyhmmer 0.1.0

use rayon::prelude::*;
use std::process::exit;

const USAGE: &str = "\
Usage: rustyhmmer-hmmsearch [options] <query.hmm> <target.faa>

Search a protein FASTA database with a HMMER3/f profile HMM and write
--tblout-format tabular results.

Options:
  -h, --help        show this help and exit
  -E <x>            report sequences with E-value <= x (default 10.0)
  -T <x>            report sequences with bit score >= x
  --cut_ga          use GA (gathering) bit score cutoffs from the model
  --cut_tc          use TC (trusted) bit score cutoffs from the model
  --cut_nc          use NC (noise) bit score cutoffs from the model
  --noali           do not output alignments
  --cpu <n>         number of worker threads (0 = all cores)
  --tblout <f>      write tabular output to file f instead of stdout";

#[derive(Default)]
struct Args {
    hmm: String,
    seq: String,
    tblout: Option<String>,
    cut_ga: bool,
    cut_tc: bool,
    cut_nc: bool,
    evalue: f64,
    bitscore: Option<f64>,
    noali: bool,
    cpu: usize,
}

fn parse_args() -> Args {
    let mut a = Args {
        evalue: 10.0,
        cpu: 0,
        ..Default::default()
    };
    let mut pos: Vec<String> = Vec::new();
    let mut it = std::env::args().skip(1);
    while let Some(arg) = it.next() {
        match arg.as_str() {
            "--cut_ga" => a.cut_ga = true,
            "--cut_tc" => a.cut_tc = true,
            "--cut_nc" => a.cut_nc = true,
            "--noali" => a.noali = true,
            "-E" => a.evalue = it.next().and_then(|s| s.parse().ok()).unwrap_or(10.0),
            "-T" => a.bitscore = it.next().and_then(|s| s.parse().ok()),
            "--tblout" => a.tblout = it.next(),
            "--cpu" => a.cpu = it.next().and_then(|s| s.parse().ok()).unwrap_or(1),
            "-h" | "--help" => {
                println!("{USAGE}");
                exit(0);
            }
            // Unknown flags are tolerated (ignored) to stay compatible with
            // extra HMMER options this port does not implement.
            s if s.starts_with('-') => {}
            _ => pos.push(arg),
        }
    }
    if pos.len() != 2 {
        eprintln!("error: expected <query.hmm> <target.faa>, got {} positional argument(s)", pos.len());
        eprintln!("{USAGE}");
        exit(1);
    }
    a.hmm = pos[0].clone();
    a.seq = pos[1].clone();
    a
}

fn main() {
    let args = parse_args();

    if args.cpu > 0 {
        if let Err(e) = rayon::ThreadPoolBuilder::new()
            .num_threads(args.cpu)
            .build_global()
        {
            eprintln!("error: cannot start {} worker threads: {e}", args.cpu);
            exit(1);
        }
    }

    for p in [&args.hmm, &args.seq] {
        if !std::path::Path::new(p).exists() {
            eprintln!("error: no such file: {p}");
            exit(1);
        }
    }

    let hmms = match rustyhmmer::hmmfile::P7Hmm::read_all(&args.hmm) {
        Ok(h) if !h.is_empty() => h,
        Ok(_) => {
            eprintln!("error: no models in {}", args.hmm);
            exit(1);
        }
        Err(e) => {
            eprintln!("error: {e}");
            exit(1);
        }
    };
    let seqs = match rustyhmmer::seqio::read_fasta(&args.seq) {
        Ok(s) => s,
        Err(e) => {
            eprintln!("error: cannot read {}: {e}", args.seq);
            exit(1);
        }
    };
    let z = seqs.len() as f64;

    use rustyhmmer::tblout::{format_row, header, sort_hits, ReportThresh, Widths};
    let use_cut = args.cut_ga || args.cut_tc || args.cut_nc;

    let mut out = String::new();
    let mut header_written = false;
    for hmm in hmms {
        let model = rustyhmmer::pipeline::Model::new(hmm);
        let qname = model.hmm.name.clone();
        let qacc = model.hmm.acc.clone().unwrap_or_default();

        let (seq_bit_t, thresh) = if use_cut {
            let cut = if args.cut_ga {
                model.hmm.cutoffs.ga
            } else if args.cut_tc {
                model.hmm.cutoffs.tc
            } else {
                model.hmm.cutoffs.nc
            };
            match cut {
                Some((t1, t2)) => (Some(t1), ReportThresh::Bits { dom_t: t2 }),
                None => {
                    eprintln!("error: bit score cutoffs unavailable on model {qname}");
                    exit(1);
                }
            }
        } else {
            (None, ReportThresh::Evalue)
        };

        let mut hits: Vec<rustyhmmer::pipeline::Hit> =
            seqs.par_iter().filter_map(|s| model.search_one(s)).collect();
        sort_hits(&mut hits);

        let reported: Vec<&rustyhmmer::pipeline::Hit> = hits
            .iter()
            .filter(|h| match seq_bit_t {
                Some(t) => (h.score as f64) >= t,
                None => h.lnp.exp() * z <= args.evalue,
            })
            .collect();

        let domz = reported.len() as f64;

        let w = Widths::compute(&qname, &qacc, &reported);
        if !header_written {
            out.push_str(&header(&w));
            header_written = true;
        }
        for h in &reported {
            out.push_str(&format_row(h, &qname, &qacc, z, domz, thresh, &w));
            out.push('\n');
        }
    }
    if !header_written {
        out.push_str(&header(&Widths { tnamew: 20, taccw: 10, qnamew: 20, qaccw: 10 }));
    }
    write_tblout_trailer(&mut out, &args);

    match &args.tblout {
        Some(path) => {
            if let Err(e) = std::fs::write(path, &out) {
                eprintln!("error: cannot write {path}: {e}");
                exit(1);
            }
        }
        None => print!("{out}"),
    }
}

fn write_tblout_trailer(out: &mut String, args: &Args) {
    out.push_str("#\n");
    out.push_str("# Program:         hmmsearch\n");
    out.push_str("# Version:         3.4 (Aug 2023)\n");
    out.push_str("# Pipeline mode:   SEARCH\n");
    out.push_str(&format!("# Query file:      {}\n", args.hmm));
    out.push_str(&format!("# Target file:     {}\n", args.seq));
    out.push_str("# [ok]\n");
}