rustyhmmer 0.1.2

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

#[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" => {
                print!("\
rustyhmmer-hmmsearch — pure-Rust HMMER3 hmmsearch (byte-identical --tblout)

USAGE:
    rustyhmmer-hmmsearch [options] <db.hmm> <seq.faa>

OPTIONS:
    -E <x>        report hits with E-value <= x (default: 10)
    -T <x>        report hits with bit score >= x
    --cut_ga      use the model's GA (gathering) bit-score cutoffs
    --cut_tc      use the model's TC (trusted) cutoffs
    --cut_nc      use the model's NC (noise) cutoffs
    --tblout <f>  write per-hit tabular output to file <f>
    --noali       accepted for compatibility (alignments are not emitted)
    --cpu <n>     number of worker threads (default: all cores)
    -h, --help    print this help and exit
");
                exit(0);
            }
            s if s.starts_with('-') => {  }
            _ => pos.push(arg),
        }
    }
    if pos.len() != 2 {
        eprintln!("error: expected <db.hmm> <seq.faa>, got {} positional args", pos.len());
        exit(1);
    }
    a.hmm = pos[0].clone();
    a.seq = pos[1].clone();
    a
}

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

    
    
    
    
    
    if args.cpu > 0 {
        rayon::ThreadPoolBuilder::new()
            .num_threads(args.cpu)
            .build_global()
            .unwrap_or_else(|e| {
                eprintln!("error: cannot start {} worker threads: {}", args.cpu, e);
                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");
}