]> git.donarmstrong.com Git - bamtools.git/blob - BamMultiReader.cpp
58d82a2d016b49de532d1ec85be39650a982d915
[bamtools.git] / BamMultiReader.cpp
1 // ***************************************************************************
2 // BamMultiReader.cpp (c) 2010 Erik Garrison
3 // Marth Lab, Department of Biology, Boston College
4 // All rights reserved.
5 // ---------------------------------------------------------------------------
6 // Last modified: 23 Februrary 2010 (EG)
7 // ---------------------------------------------------------------------------
8 // Uses BGZF routines were adapted from the bgzf.c code developed at the Broad
9 // Institute.
10 // ---------------------------------------------------------------------------
11 // Functionality for simultaneously reading multiple BAM files.
12 //
13 // This functionality allows applications to work on very large sets of files
14 // without requiring intermediate merge, sort, and index steps for each file
15 // subset.  It also improves the performance of our merge system as it
16 // precludes the need to sort merged files.
17 // ***************************************************************************
18
19 // C++ includes
20 #include <algorithm>
21 #include <iterator>
22 #include <string>
23 #include <vector>
24 #include <iostream>
25 #include <sstream>
26
27 // BamTools includes
28 #include "BGZF.h"
29 #include "BamMultiReader.h"
30 using namespace BamTools;
31 using namespace std;
32
33 // -----------------------------------------------------
34 // BamMultiReader implementation
35 // -----------------------------------------------------
36
37 // constructor
38 BamMultiReader::BamMultiReader(void)
39     : CurrentRefID(0)
40     , CurrentLeft(0)
41 { }
42
43 // destructor
44 BamMultiReader::~BamMultiReader(void) {
45     Close(); // close the bam files
46     // clean up reader objects
47     for (vector<pair<BamReader*, BamAlignment*> >::iterator it = readers.begin(); it != readers.end(); ++it) {
48         delete it->first;
49         delete it->second;
50     }
51 }
52
53 // close the BAM files
54 void BamMultiReader::Close(void) {
55     for (vector<pair<BamReader*, BamAlignment*> >::iterator it = readers.begin(); it != readers.end(); ++it) {
56         BamReader* reader = it->first;
57         reader->Close();  // close the reader
58     }
59 }
60
61 // updates the reference id stored in the BamMultiReader
62 // to reflect the current state of the readers
63 void BamMultiReader::UpdateReferenceID(void) {
64     // the alignments are sorted by position, so the first alignment will always have the lowest reference ID
65     if (alignments.begin()->second.second->RefID != CurrentRefID) {
66         // get the next reference id
67         // while there aren't any readers at the next ref id
68         // increment the ref id
69         int nextRefID = CurrentRefID;
70         while (alignments.begin()->second.second->RefID != nextRefID) {
71             ++nextRefID;
72         }
73         //cerr << "updating reference id from " << CurrentRefID << " to " << nextRefID << endl;
74         CurrentRefID = nextRefID;
75     }
76 }
77
78 // checks if any readers still have alignments
79 bool BamMultiReader::HasOpenReaders() {
80     return alignments.size() > 0;
81 }
82
83 // get next alignment among all files (from specified region, if given)
84 bool BamMultiReader::GetNextAlignment(BamAlignment& nextAlignment) {
85
86     // bail out if we are at EOF in all files, means no more alignments to process
87     if (!HasOpenReaders())
88         return false;
89
90     // when all alignments have stepped into a new target sequence, update our
91     // current reference sequence id
92     UpdateReferenceID();
93
94     // our lowest alignment and reader will be at the front of our alignment index
95     BamAlignment* lowestAlignment = alignments.begin()->second.second;
96     BamReader* lowestReader = alignments.begin()->second.first;
97
98     // now that we have the lowest alignment in the set, save it by copy to our argument
99     nextAlignment = BamAlignment(*lowestAlignment);
100
101     // remove this alignment index entry from our alignment index
102     alignments.erase(alignments.begin());
103
104     // and add another entry if we can get another alignment from the reader
105     if (lowestReader->GetNextAlignment(*lowestAlignment)) {
106         alignments.insert(make_pair(make_pair(lowestAlignment->RefID, lowestAlignment->Position), 
107                                     make_pair(lowestReader, lowestAlignment)));
108     } else { // do nothing
109         //cerr << "reached end of file " << lowestReader->GetFilename() << endl;
110     }
111
112     return true;
113 }
114
115 // jumps to specified region(refID, leftBound) in BAM files, returns success/fail
116 bool BamMultiReader::Jump(int refID, int position) {
117
118     //if ( References.at(refID).RefHasAlignments && (position <= References.at(refID).RefLength) ) {
119     CurrentRefID = refID;
120     CurrentLeft  = position;
121
122     bool result = true;
123     for (vector<pair<BamReader*, BamAlignment*> >::iterator it = readers.begin(); it != readers.end(); ++it) {
124         BamReader* reader = it->first;
125         result &= reader->Jump(refID, position);
126         if (!result) {
127             cerr << "ERROR: could not jump " << reader->GetFilename() << " to " << refID << ":" << position << endl;
128             exit(1);
129         }
130     }
131     if (result) {
132         // Update Alignments
133         alignments.clear();
134         for (vector<pair<BamReader*, BamAlignment*> >::iterator it = readers.begin(); it != readers.end(); ++it) {
135             BamReader* br = it->first;
136             BamAlignment* ba = it->second;
137             if (br->GetNextAlignment(*ba)) {
138                 alignments.insert(make_pair(make_pair(ba->RefID, ba->Position), 
139                                             make_pair(br, ba)));
140             } else {
141                 // assume BamReader EOF
142             }
143         }
144     }
145     return result;
146 }
147
148 // opens BAM files
149 void BamMultiReader::Open(const vector<string> filenames, bool openIndexes) {
150     // for filename in filenames
151     fileNames = filenames; // save filenames in our multireader
152     for (vector<string>::const_iterator it = filenames.begin(); it != filenames.end(); ++it) {
153         string filename = *it;
154         BamReader* reader = new BamReader;
155         if (openIndexes) {
156             reader->Open(filename, filename + ".bai");
157         } else {
158             reader->Open(filename); // for merging, jumping is disallowed
159         }
160         BamAlignment* alignment = new BamAlignment;
161         reader->GetNextAlignment(*alignment);
162         readers.push_back(make_pair(reader, alignment)); // store pointers to our readers for cleanup
163         alignments.insert(make_pair(make_pair(alignment->RefID, alignment->Position),
164                                     make_pair(reader, alignment)));
165     }
166     ValidateReaders();
167 }
168
169 void BamMultiReader::PrintFilenames(void) {
170     for (vector<pair<BamReader*, BamAlignment*> >::iterator it = readers.begin(); it != readers.end(); ++it) {
171         BamReader* reader = it->first;
172         cout << reader->GetFilename() << endl;
173     }
174 }
175
176 // for debugging
177 void BamMultiReader::DumpAlignmentIndex(void) {
178     for (AlignmentIndex::const_iterator it = alignments.begin(); it != alignments.end(); ++it) {
179         cerr << it->first.first << ":" << it->first.second << " " << it->second.first->GetFilename() << endl;
180     }
181 }
182
183 // returns BAM file pointers to beginning of alignment data
184 bool BamMultiReader::Rewind(void) { 
185     bool result = true;
186     for (vector<pair<BamReader*, BamAlignment*> >::iterator it = readers.begin(); it != readers.end(); ++it) {
187         BamReader* reader = it->first;
188         result &= reader->Rewind();
189     }
190     return result;
191 }
192
193 // saves index data to BAM index files (".bai") where necessary, returns success/fail
194 bool BamMultiReader::CreateIndexes(void) {
195     bool result = true;
196     for (vector<pair<BamReader*, BamAlignment*> >::iterator it = readers.begin(); it != readers.end(); ++it) {
197         BamReader* reader = it->first;
198         result &= reader->CreateIndex();
199     }
200     return result;
201 }
202
203 // makes a virtual, unified header for all the bam files in the multireader
204 const string BamMultiReader::GetHeaderText(void) const {
205
206     string mergedHeader = "";
207     map<string, bool> readGroups;
208
209     // foreach extraction entry (each BAM file)
210     bool isFirstTime = true;
211     for (vector<pair<BamReader*, BamAlignment*> >::const_iterator it = readers.begin(); it != readers.end(); ++it) {
212
213         BamReader* reader = it->first;
214
215         stringstream header(reader->GetHeaderText());
216         vector<string> lines;
217         string item;
218         while (getline(header, item))
219             lines.push_back(item);
220
221         for (vector<string>::const_iterator it = lines.begin(); it != lines.end(); ++it) {
222
223             // get next line from header, skip if empty
224             string headerLine = *it;
225             if ( headerLine.empty() ) { continue; }
226
227             // if first file, save HD & SQ entries
228             if ( isFirstTime ) {
229                 if ( headerLine.find("@HD") == 0 || headerLine.find("@SQ") == 0) {
230                     mergedHeader.append(headerLine.c_str());
231                     mergedHeader.append(1, '\n');
232                 }
233             }
234
235             // (for all files) append RG entries if they are unique
236             if ( headerLine.find("@RG") == 0 ) {
237                 stringstream headerLineSs(headerLine);
238                 string part, readGroupPart, readGroup;
239                 while(std::getline(headerLineSs, part, '\t')) {
240                     if (part == "@RG") {
241                        std::getline(headerLineSs, readGroupPart, '\t');
242                        stringstream readGroupPartSs(readGroupPart);
243                        std::getline(readGroupPartSs, readGroup, ':');
244                        std::getline(readGroupPartSs, readGroup, ':');
245                        break;
246                     }
247                 }
248                 if (readGroups.find(readGroup) == readGroups.end()) { // prevents duplicate @RG entries
249                     mergedHeader.append(headerLine.c_str() );
250                     mergedHeader.append(1, '\n');
251                     readGroups[readGroup] = true;
252                 }
253             }
254
255         }
256
257         // set iteration flag
258         isFirstTime = false;
259     }
260
261     // return merged header text
262     return mergedHeader;
263 }
264
265 // ValidateReaders checks that all the readers point to BAM files representing
266 // alignments against the same set of reference sequences, and that the
267 // sequences are identically ordered.  If these checks fail the operation of
268 // the multireader is undefined, so we force program exit.
269 void BamMultiReader::ValidateReaders(void) const {
270     int firstRefCount = readers.front().first->GetReferenceCount();
271     BamTools::RefVector firstRefData = readers.front().first->GetReferenceData();
272     for (vector<pair<BamReader*, BamAlignment*> >::const_iterator it = readers.begin(); it != readers.end(); ++it) {
273         BamReader* reader = it->first;
274         BamTools::RefVector currentRefData = reader->GetReferenceData();
275         BamTools::RefVector::const_iterator f = firstRefData.begin();
276         BamTools::RefVector::const_iterator c = currentRefData.begin();
277         if (reader->GetReferenceCount() != firstRefCount || firstRefData.size() != currentRefData.size()) {
278             cerr << "ERROR: mismatched number of references in " << reader->GetFilename()
279                       << " expected " << firstRefCount 
280                       << " reference sequences but only found " << reader->GetReferenceCount() << endl;
281             exit(1);
282         }
283         // this will be ok; we just checked above that we have identically-sized sets of references
284         // here we simply check if they are all, in fact, equal in content
285         while (f != firstRefData.end()) {
286             if (f->RefName != c->RefName || f->RefLength != c->RefLength) {
287                 cerr << "ERROR: mismatched references found in " << reader->GetFilename()
288                           << " expected: " << endl;
289                 for (BamTools::RefVector::const_iterator a = firstRefData.begin(); a != firstRefData.end(); ++a)
290                     cerr << a->RefName << " " << a->RefLength << endl;
291                 cerr << "but found: " << endl;
292                 for (BamTools::RefVector::const_iterator a = currentRefData.begin(); a != currentRefData.end(); ++a)
293                     cerr << a->RefName << " " << a->RefLength << endl;
294                 exit(1);
295             }
296             ++f; ++c;
297         }
298     }
299 }
300
301 // NB: The following functions assume that we have identical references for all
302 // BAM files.  We enforce this by invoking the above validation function
303 // (ValidateReaders) to verify that our reference data is the same across all
304 // files on Open, so we will not encounter a situation in which there is a
305 // mismatch and we are still live.
306
307 // returns the number of reference sequences
308 const int BamMultiReader::GetReferenceCount(void) const {
309     return readers.front().first->GetReferenceCount();
310 }
311
312 // returns vector of reference objects
313 const BamTools::RefVector BamMultiReader::GetReferenceData(void) const {
314     return readers.front().first->GetReferenceData();
315 }
316
317 const int BamMultiReader::GetReferenceID(const string& refName) const { 
318     return readers.front().first->GetReferenceID(refName);
319 }