Using a thread pool instead of setting number of threads. Also changed how good reads and bases are tabulated for in and out to be reported in the stats output.

Author: andrewdavidsmith

src/utils/uniq.cpp

@@ -23,10 +23,14 @@
 #include <iostream>
 #include <stdexcept>
 
+// generated by autotools
 #include <config.h>
 
+// from HTSlib
 #include <htslib/sam.h>
+#include <htslib/thread_pool.h>
 
+// from smithlab
 #include "OptionParser.hpp"
 #include "smithlab_utils.hpp"
 #include "smithlab_os.hpp"
@@ -43,8 +47,16 @@ using std::runtime_error;
 using std::to_string;
 
 
+inline bool
+format_is_bam_or_sam(htsFile *hts) {
+  const htsFormat *fmt = hts_get_format(hts);
+  return fmt->category == sequence_data &&
+    (fmt->format == bam || fmt->format == sam);
+}
+
+
 inline string
-read_name(const bam1_t *b) {return string(bam_get_qname(b));}
+qname(const bam1_t *b) {return string(bam_get_qname(b));}
 
 
 inline int32_t
@@ -80,24 +92,31 @@ equivalent_end_and_strand(const bam1_t *a, const bam1_t *b) {
 }
 
 
+struct rd_stats { // keep track of good bases/reads in and out
+  size_t bases;
+  size_t reads;
+  rd_stats() : bases(0), reads(0) {}
+  void update(bam1_t *b) { bases += qlen(b); ++reads; }
+};
+
+
 static void
-write_stats_output(const size_t reads_in, const size_t reads_out,
-                   const size_t good_bases_in, const size_t good_bases_out,
-                   const size_t reads_with_dups, const string &statfile) {
+write_stats_output(const rd_stats &rs_in, const rd_stats &rs_out,
+                   const size_t reads_duped, const string &statfile) {
   if (!statfile.empty()) {
-    const size_t reads_removed = reads_in - reads_out;
+    const size_t reads_removed = rs_in.reads - rs_out.reads;
     const double non_dup_frac =
-      (reads_out - reads_with_dups)/static_cast<double>(reads_in);
+      (rs_out.reads - reads_duped)/static_cast<double>(rs_in.reads);
     const double dup_rate =
-      (reads_removed + reads_with_dups)/static_cast<double>(reads_with_dups);
+      (reads_removed + reads_duped)/static_cast<double>(reads_duped);
     ofstream out_stat(statfile);
     if (!out_stat) throw runtime_error("bad stats output file");
-    out_stat << "total_reads: " << reads_in << endl
-             << "total_bases: " << good_bases_in << endl
-             << "unique_reads: " << reads_out << endl
-             << "unique_read_bases: " << good_bases_out << endl
+    out_stat << "total_reads: " << rs_in.reads << endl
+             << "total_bases: " << rs_in.bases << endl
+             << "unique_reads: " << rs_out.reads << endl
+             << "unique_read_bases: " << rs_out.bases << endl
              << "non_duplicate_fraction: " << non_dup_frac << endl
-             << "duplicate_reads: " << reads_with_dups << endl
+             << "duplicate_reads: " << reads_duped << endl
              << "reads_removed: " << reads_removed << endl
              << "duplication_rate: " << dup_rate << endl;
   }
@@ -124,9 +143,8 @@ static void
 process_inner_buffer(const vector<bam1_t*>::const_iterator it,
                      const vector<bam1_t*>::const_iterator jt,
                      sam_hdr_t *hdr, samFile *out,
-                     size_t &reads_out,
-                     size_t &good_bases_out,
-                     size_t &reads_with_dups,
+                     rd_stats &rs_out,
+                     size_t &reads_duped,
                      vector<size_t> &hist) {
   const size_t n_reads = std::distance(it, jt);
   const size_t selected = rand() % n_reads;
@@ -135,29 +153,25 @@ process_inner_buffer(const vector<bam1_t*>::const_iterator it,
   if (hist.size() <= n_reads)
     hist.resize(n_reads + 1);
   hist[n_reads]++;
-  good_bases_out += qlen(*(it + selected));
-  ++reads_out;
-  reads_with_dups += (n_reads > 1);
+  rs_out.update(*(it + selected));
+  reads_duped += (n_reads > 1);
 }
 
 
 /* The buffer corresponds to reads sharing the same mapping chromosome
    and start position. These are gathered and then processed together. */
 static void
-process_buffer(size_t &reads_out, size_t &good_bases_out,
-               size_t &reads_with_dups, vector<size_t> &hist,
+process_buffer(rd_stats &rs_out, size_t &reads_duped, vector<size_t> &hist,
                vector<bam1_t*> &buffer, sam_hdr_t *hdr, samFile *out) {
   sort(begin(buffer), end(buffer), precedes_by_end_and_strand);
   auto it(begin(buffer));
   auto jt = it + 1;
   for (; jt != end(buffer); ++jt)
     if (!equivalent_end_and_strand(*it, *jt)) {
-      process_inner_buffer(it, jt, hdr, out, reads_out, good_bases_out,
-                           reads_with_dups, hist);
+      process_inner_buffer(it, jt, hdr, out, rs_out, reads_duped, hist);
       it = jt;
     }
-  process_inner_buffer(it, jt, hdr, out, reads_out, good_bases_out,
-                       reads_with_dups, hist);
+  process_inner_buffer(it, jt, hdr, out, rs_out, reads_duped, hist);
 
   // free the bam1_t pointers before clearing the buffer
   for (size_t i = 0; i < buffer.size(); ++i)
@@ -193,21 +207,21 @@ uniq(const bool VERBOSE, const size_t n_threads,
   if (!hts || errno)
     throw runtime_error("bad htslib file: " + infile);
 
-  if (n_threads > 1 && hts_set_threads(hts, n_threads/2) < 0)
+  htsThreadPool the_thread_pool{hts_tpool_init(n_threads), 0};
+  if (hts_set_thread_pool(hts, &the_thread_pool) < 0)
     throw runtime_error("error setting threads");
 
-  if (hts_get_format(hts)->category != sequence_data)
+  if (!format_is_bam_or_sam(hts))
     throw runtime_error("bad file format: " + infile);
 
   sam_hdr_t *hdr = sam_hdr_read(hts);
   if (!hdr)
     throw runtime_error("failed to read header: " + infile);
 
   // open the output file
-  samFile *out = bam_format ? hts_open(outfile.c_str(), "wb") :
-    hts_open(outfile.c_str(), "w");
+  samFile *out = hts_open(outfile.c_str(), bam_format ? "wb" : "w");
 
-  if (n_threads > 1 && hts_set_threads(out, (n_threads + 1)/2) < 0)
+  if (hts_set_thread_pool(out, &the_thread_pool) < 0)
     throw runtime_error("error setting threads");
 
   // take care of the output file's header
@@ -224,13 +238,12 @@ uniq(const bool VERBOSE, const size_t n_threads,
   if (!aln)
     throw runtime_error("failed parsing read from input file");
 
-  // values to tabulate stats; these cost almost nothing
+  // values to tabulate stats; no real cost
+  rd_stats rs_in, rs_out;
+  size_t reads_duped = 0;
   vector<size_t> hist;
-  size_t reads_in = 1; // count the one we just got
-  size_t good_bases_in = qlen(aln); // count its good bases
-  size_t reads_out = 0;
-  size_t good_bases_out = 0;
-  size_t reads_with_dups = 0;
+
+  rs_in.update(aln); // update for the input we just got
 
   vector<bam1_t*> buffer(1, aln); // select output from this buffer
 
@@ -239,36 +252,33 @@ uniq(const bool VERBOSE, const size_t n_threads,
   int32_t cur_chrom = get_tid(aln);
 
   while (aln = get_read(hts, hdr)) {
-    ++reads_in;
-    good_bases_in += qlen(aln);
+    rs_in.update(aln);
+
+    // below works because buffer reset at every new chrom
+    if (precedes_by_start(aln, buffer[0]))
+      throw runtime_error("not sorted: " + qname(buffer[0]) + " " + qname(aln));
 
-    // below works because buffer is reset every chrom
-    if (precedes_by_start(aln, buffer.front()))
-      throw runtime_error("input not properly sorted:\n" +
-                          read_name(buffer[0]) + "\n" + read_name(aln));
     const int32_t chrom = get_tid(aln);
     if (chrom != cur_chrom) {
       if (chroms_seen[chrom]) throw runtime_error("input not sorted");
       chroms_seen[chrom] = true;
       cur_chrom = chrom;
     }
 
-    if (!equivalent_chrom_and_start(buffer.front(), aln))
-      process_buffer(reads_out, good_bases_out, reads_with_dups,
-                     hist, buffer, hdr, out);
+    if (!equivalent_chrom_and_start(buffer[0], aln))
+      process_buffer(rs_out, reads_duped, hist, buffer, hdr, out);
     buffer.push_back(aln);
   }
-  process_buffer(reads_out, good_bases_out, reads_with_dups,
-                 hist, buffer, hdr, out);
+  process_buffer(rs_out, reads_duped, hist, buffer, hdr, out);
 
-  // remember to close these
+  // remember to turn off the lights
   bam_hdr_destroy(hdr);
   hts_close(out);
   hts_close(hts);
+  hts_tpool_destroy(the_thread_pool.pool);
 
-  write_stats_output(reads_in, reads_out, good_bases_in, good_bases_out,
-                     reads_with_dups, statfile);
-
+  // write any additional output requested
+  write_stats_output(rs_in, rs_out, reads_duped, statfile);
   write_hist_output(hist, histfile);
 }