bam2mr.cpp

/* to-mr: a program for converting SAM and BAM format to MappedRead
 * format.
 * Currently supported mappers: bsmap, bismark.
 *
 * Copyright (C) 2009-2012 University of Southern California and
 * Andrew D. Smith
 *
 * Authors: Meng Zhou, Qiang Song, Andrew Smith
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 */

#include <cmath>

#include <string>
#include <vector>
#include <iostream>
#include <fstream>
#include <queue>
#include <tr1/unordered_map>

#include "OptionParser.hpp"
#include "smithlab_utils.hpp"
#include "smithlab_os.hpp"
#include "SAM.hpp"

using std::string;
using std::vector;
using std::cout;
using std::cerr;
using std::endl;
using std::ifstream;
using std::max;
using std::min;
using std::priority_queue;
using std::tr1::unordered_map;





/********Below are functions for merging pair-end reads********/
static void
fill_overlap(const bool pos_str, const MappedRead &mr, const size_t start, 
 const size_t end, const size_t offset, string &seq, string &scr) {
 const size_t a = pos_str ? (start - mr.r.get_start()) : (mr.r.get_end() - end);
 const size_t b = pos_str ? (end - mr.r.get_start()) : (mr.r.get_end() - start);
 copy(mr.seq.begin() + a, mr.seq.begin() + b, seq.begin() + offset);
 copy(mr.scr.begin() + a, mr.scr.begin() + b, scr.begin() + offset);
}

static bool
merge_mates(const size_t suffix_len, const size_t range,
 const MappedRead &one, const MappedRead &two,
 MappedRead &merged, int &len) {
 
 const bool pos_str = one.r.pos_strand();
 const size_t overlap_start = max(one.r.get_start(), two.r.get_start());
 const size_t overlap_end = min(one.r.get_end(), two.r.get_end());

 const size_t one_left = pos_str ? 
 one.r.get_start() : max(overlap_end, one.r.get_start());
 const size_t one_right = 
 pos_str ? min(overlap_start, one.r.get_end()) : one.r.get_end();
 
 const size_t two_left = pos_str ? 
 max(overlap_end, two.r.get_start()) : two.r.get_start();
 const size_t two_right = pos_str ? 
 two.r.get_end() : min(overlap_start, two.r.get_end());

 len = pos_str ? (two_right - one_left) : (one_right - two_left);
 
 if(len < 0){
 cerr << one << endl;
 cerr << two << endl;
 return false;
 }
 // assert(len > 0);
 // assert(one_left <= one_right && two_left <= two_right);
 // assert(overlap_start >= overlap_end || static_cast<size_t>(len) == 
 // ((one_right - one_left) + (two_right - two_left) + (overlap_end - overlap_start)));
 
 string seq(len, 'N');
 string scr(len, 'B');
 if (len >= 0 && len <= static_cast<int>(range)) {
 // lim_one: offset in merged sequence where overlap starts
 const size_t lim_one = one_right - one_left;
 copy(one.seq.begin(), one.seq.begin() + lim_one, seq.begin());
 copy(one.scr.begin(), one.scr.begin() + lim_one, scr.begin());
 
 const size_t lim_two = two_right - two_left;
 copy(two.seq.end() - lim_two, two.seq.end(), seq.end() - lim_two);
 copy(two.scr.end() - lim_two, two.scr.end(), scr.end() - lim_two);
 
 // deal with overlapping part
 if (overlap_start < overlap_end) {
 const size_t one_bads = count(one.seq.begin(), one.seq.end(), 'N');
 const int info_one = one.seq.length() - (one_bads + one.r.get_score());
 
 const size_t two_bads = count(two.seq.begin(), two.seq.end(), 'N');
 const int info_two = two.seq.length() - (two_bads + two.r.get_score());
 
 // use the mate with the most info to fill in the overlap
 if (info_one >= info_two)
	fill_overlap(pos_str, one, overlap_start, overlap_end, lim_one, seq, scr);
 else
	fill_overlap(pos_str, two, overlap_start, overlap_end, lim_one, seq, scr);
 }
 }
 
 merged = one;
 merged.r.set_start(pos_str ? one.r.get_start() : two.r.get_start());
 merged.r.set_end(merged.r.get_start() + len);
 merged.r.set_score(one.r.get_score() + two.r.get_score());
 merged.seq = seq;
 merged.scr = scr; 
 const string name(one.r.get_name());
 merged.r.set_name("FRAG:" + name.substr(0, name.size() - suffix_len));

 return true;
}

inline static bool
same_read(const size_t suffix_len, 
	 const MappedRead &a, const MappedRead &b) {
 const string sa(a.r.get_name());
 const string sb(b.r.get_name());
 bool SAME_NAME = false;
 if(sa == sb)
 SAME_NAME = true;
 return (SAME_NAME && a.r.same_chrom(b.r));
}



static void
revcomp(MappedRead &mr) {
 // set the strand to the opposite of the current value
 mr.r.set_strand(mr.r.pos_strand() ? '-' : '+');
 // reverse complement the sequence, and reverse the quality scores
 revcomp_inplace(mr.seq);
 std::reverse(mr.scr.begin(), mr.scr.end());
}
/********Above are functions for merging pair-end reads********/


/////comparison function for priority queue/////////////////

/**************** FOR CLARITY BELOW WHEN COMPARING READS *************/
static inline bool
chrom_greater(const MappedRead &a, const MappedRead &b) {
 return a.r.get_chrom() > b.r.get_chrom();
}
static inline bool
same_start(const MappedRead &a, const MappedRead &b) {
 return a.r.get_start() == b.r.get_start();
}
static inline bool
start_greater(const MappedRead &a, const MappedRead &b) {
 return a.r.get_start() > b.r.get_start();
}
static inline bool
end_greater(const MappedRead &a, const MappedRead &b) {
 return a.r.get_end() > b.r.get_end();
}
/******************************************************************************/


struct MappedReadOrderChecker {
 bool operator()(const MappedRead &prev, const MappedRead &gr) const {
 return start_check(prev, gr);
 }
 static bool
 is_ready(const priority_queue<MappedRead, vector<MappedRead>, MappedReadOrderChecker> &pq,
 const MappedRead &mr, const size_t max_width) {
 return !pq.top().r.same_chrom(mr.r) || pq.top().r.get_end() + max_width < mr.r.get_start();
 }
 static bool
 start_check(const MappedRead &prev, const MappedRead &mr) {
 return (chrom_greater(prev, mr)
 || (prev.r.same_chrom(mr.r) && start_greater(prev, mr))
 || (prev.r.same_chrom(mr.r) && same_start(prev, mr) && end_greater(prev, mr)));
 }
};



static void empty_pq(MappedRead &prev_mr,
 priority_queue<MappedRead, vector<MappedRead>,
 MappedReadOrderChecker> &read_pq,
 const string &input_file_name,
		 std::ostream& out){
 
 MappedRead curr_mr = read_pq.top();
 //	 cerr << "outputting from queue : " << read_pq.top() << endl;
 read_pq.pop();

 /*
 // check if reads are sorted
 if (curr_mr.r.same_chrom(prev_mr.r) &&
 curr_mr.r.get_start() < prev_mr.r.get_start()
 && curr_mr.r.get_end() < prev_mr.r.get_end()){
 cerr << "prev = \t" << prev_mr << endl;
 cerr << "curr = \t" << curr_mr << endl;
 cerr << "Increase seg_len if in paired end mode" << endl;
 throw SMITHLABException("reads unsorted in " + input_file_name);
 }
 */

 out << curr_mr << endl;

 prev_mr = curr_mr;
}


int 
main(int argc, const char **argv) {
 try {
 string outfile;
 string mapper = "general";
 size_t MAX_SEGMENT_LENGTH = 5000;
 size_t suffix_len = 0;
 bool VERBOSE = false;
 size_t MAX_READS_TO_HOLD = 1000000;
 
 /****************** COMMAND LINE OPTIONS ********************/
 OptionParser opt_parse(strip_path(argv[0]),
 "Convert the SAM/BAM output from "
 "mapped read format to sam/bam_file");
 opt_parse.add_opt("output", 'o', "Name of output file", 
 false, outfile);
 opt_parse.add_opt("suff", 's', "read name suffix length (default: 0)",
 false, suffix_len); 
 opt_parse.add_opt("max-frag", 'L', "maximum allowed insert size", 
 false, MAX_SEGMENT_LENGTH); 
 opt_parse.add_opt("max_reads", 'R', "maximum number of reads to hold for merging",
		 false, MAX_READS_TO_HOLD);
 opt_parse.add_opt("verbose", 'v', "print more information",
 false, VERBOSE);

 vector<string> leftover_args;
 opt_parse.parse(argc, argv, leftover_args);
 if (argc < 3 || opt_parse.help_requested()) {
 cerr << opt_parse.help_message() << endl
 << opt_parse.about_message() << endl;
 return EXIT_SUCCESS;
 }
 if (opt_parse.about_requested()) {
 cerr << opt_parse.about_message() << endl;
 return EXIT_SUCCESS;
 }
 if (opt_parse.option_missing()) {
 cerr << opt_parse.option_missing_message() << endl;
 return EXIT_SUCCESS;
 }
 if (leftover_args.empty()) {
 cerr << opt_parse.help_message() << endl;
 return EXIT_SUCCESS;
 }
 const string mapped_reads_file = leftover_args.front();
 /****************** END COMMAND LINE OPTIONS *****************/

 std::ofstream of;
 if (!outfile.empty()) of.open(outfile.c_str());
 std::ostream out(outfile.empty() ? cout.rdbuf() : of.rdbuf());
 if (VERBOSE)
 {
 cerr << "Input file: " << mapped_reads_file << endl
 << "Output file: " << (outfile.empty() ? "stdout" : outfile) << endl;
 }

 SAMReader sam_reader(mapped_reads_file, mapper);
 std::tr1::unordered_map<string, SAMRecord> dangling_mates;
 
 const size_t progress_step = 1000000;
 SAMRecord samr;
 MappedRead prev_mr;
 size_t n_mates = 0;

 std::priority_queue<MappedRead, vector<MappedRead>, MappedReadOrderChecker> read_pq;

 while ((sam_reader >> samr, sam_reader.is_good()))
 {
 if(samr.is_primary && samr.is_mapped){
	// only convert mapped and primary reads
	++n_mates;
	if (samr.is_mapping_paired){
	 const string read_name
	 = samr.mr.r.get_name().substr(0, samr.mr.r.get_name().size() - suffix_len);

	 if (dangling_mates.find(read_name) != dangling_mates.end()){
	 // other end is in dangling mates, merge the two mates
	 if(same_read(suffix_len, samr.mr, dangling_mates[read_name].mr)){
	 if (samr.is_Trich) std::swap(samr, dangling_mates[read_name]);

	 revcomp(samr.mr);

	 MappedRead merged;
	 int len = 0;
	 bool MERGE_SUCCESS =
		merge_mates(suffix_len, MAX_SEGMENT_LENGTH,
			 dangling_mates[read_name].mr, samr.mr, merged, len);

	 if (MERGE_SUCCESS && 
		 len >= 0 && 
		 len <= static_cast<int>(MAX_SEGMENT_LENGTH)){
		//if(read_pq.empty())
		//prev_mr = merged;
		//else
		 read_pq.push(merged);
	 }
	 else{
		// informative error message!
		if(VERBOSE){
		 cerr << "problem merging read " << read_name << ", splitting read" << endl;
		 cerr << samr.mr << endl;
		 cerr << dangling_mates[read_name].mr << endl;
		 cerr << "To merge, set max segement length (seg_len) higher." << endl;
		}

		// don't throw error for problems merging
		//if(read_pq.empty()){
		// prev_mr = dangling_mates[read_name].mr;
		// read_pq.push(samr.mr);
		//}
		//else{
		 read_pq.push(samr.mr);
		 read_pq.push(dangling_mates[read_name].mr);
		 //}
	 }

	 dangling_mates.erase(read_name);
	 }
	 else{
	 //if(read_pq.empty()){
	 //prev_mr = dangling_mates[read_name].mr;
	 //read_pq.push(samr.mr);
	 //}
	 //else{
		read_pq.push(samr.mr);
		read_pq.push(dangling_mates[read_name].mr);
		//}
	 dangling_mates.erase(read_name);
	 }

	 if(!(read_pq.empty()) &&
	 MappedReadOrderChecker::is_ready(read_pq, samr.mr, MAX_SEGMENT_LENGTH)) {
	 //begin emptying priority queue
	 while(!(read_pq.empty()) &&
		 MappedReadOrderChecker::is_ready(read_pq, samr.mr, MAX_SEGMENT_LENGTH) ){
		empty_pq(prev_mr, read_pq, mapped_reads_file, out);
	 }//end while loop
	 }//end statement for emptying priority queue

	 }
	 else
	 dangling_mates[read_name] = samr;

	}
	else{ 
	 // unmatched, output read
	 if (!samr.is_Trich) revcomp(samr.mr);

	 //if(read_pq.empty())
	 //prev_mr = samr.mr;
	 //else 
	 read_pq.push(samr.mr);

	 if(!(read_pq.empty()) &&
	 MappedReadOrderChecker::is_ready(read_pq, samr.mr, MAX_SEGMENT_LENGTH)) {
	 //begin emptying priority queue
	 while(!(read_pq.empty()) &&
		 MappedReadOrderChecker::is_ready(read_pq, samr.mr, MAX_SEGMENT_LENGTH) ){
	 empty_pq(prev_mr, read_pq, mapped_reads_file, out);
	 }//end while loop
	 }//end statement for emptying priority queue

	}

 // dangling mates is too large, flush dangling_mates of reads
 // on different chroms and too far away 
	if (dangling_mates.size() > MAX_READS_TO_HOLD){
	 
	// if(VERBOSE)
	// cerr << "dangling mates too large, emptying" << endl;

	 using std::tr1::unordered_map;
	 unordered_map<string, SAMRecord> tmp;
	 for (unordered_map<string, SAMRecord>::iterator
		 itr = dangling_mates.begin();
	 itr != dangling_mates.end(); ++itr){
	 if (itr->second.mr.r.get_chrom() != samr.mr.r.get_chrom()
		|| (itr->second.mr.r.get_chrom() == samr.mr.r.get_chrom()
		 && itr->second.mr.r.get_end() + MAX_SEGMENT_LENGTH <
		 samr.mr.r.get_start())) {
	 if (!itr->second.is_Trich) revcomp(itr->second.mr);
	 if(itr->second.seg_len >= 0)
		read_pq.push(itr->second.mr);
	 }
	 else
	 tmp[itr->first] = itr->second;
	 }
	 std::swap(tmp, dangling_mates);
	 tmp.clear();

	 if(!(read_pq.empty()) &&
	 MappedReadOrderChecker::is_ready(read_pq, samr.mr, MAX_SEGMENT_LENGTH)) {
	 //begin emptying priority queue
	 while(!(read_pq.empty()) &&
		 MappedReadOrderChecker::is_ready(read_pq, samr.mr, MAX_SEGMENT_LENGTH) ){
	 empty_pq(prev_mr, read_pq, mapped_reads_file, out);
	 }//end while loop
	 }//end statement for emptying priority queue

	}
 
	if (VERBOSE && n_mates % progress_step == 0)
	 cerr << "Processed " << n_mates << " records" << endl;
 }
 }
 // flushing dangling_mates of all remaining ends
 if(!(dangling_mates.empty()) && VERBOSE){
 cerr << "dangling mates not empty" << endl;
 cerr << "dangling_mates.size = " << dangling_mates.size() << endl;
 }
 while (!dangling_mates.empty()){
 if (!dangling_mates.begin()->second.is_Trich)
 revcomp(dangling_mates.begin()->second.mr);
 read_pq.push(dangling_mates.begin()->second.mr);
 dangling_mates.erase(dangling_mates.begin());
 }

 while(!read_pq.empty()){
 empty_pq(prev_mr, read_pq, mapped_reads_file, out);
 }
 
 if (VERBOSE){
 cerr << "Done." << endl;
 cerr << "total mates processed = " << n_mates << endl;
 }
 }
 catch (const SMITHLABException &e) {
 cerr << e.what() << endl;
 return EXIT_FAILURE;
 }
 catch (std::bad_alloc &ba) {
 cerr << "ERROR: could not allocate memory" << endl;
 return EXIT_FAILURE;
 }
 return EXIT_SUCCESS;
}