]> git.donarmstrong.com Git - bamtools.git/blob - src/toolkit/bamtools_sort.cpp
ad0e4f0a69b9fa9a38ce92831ae5728853a37259
[bamtools.git] / src / toolkit / bamtools_sort.cpp
1 // ***************************************************************************
2 // bamtools_sort.cpp (c) 2010 Derek Barnett, Erik Garrison
3 // Marth Lab, Department of Biology, Boston College
4 // ---------------------------------------------------------------------------
5 // Last modified: 28 September 2011 (DB)
6 // ---------------------------------------------------------------------------
7 // Sorts an input BAM file
8 // ***************************************************************************
9
10 #include "bamtools_sort.h"
11
12 #include <api/SamConstants.h>
13 #include <api/BamMultiReader.h>
14 #include <api/BamWriter.h>
15 #include <utils/bamtools_options.h>
16 using namespace BamTools;
17
18 #include <cstdio>
19 #include <algorithm>
20 #include <iostream>
21 #include <sstream>
22 #include <string>
23 #include <vector>
24 using namespace std;
25
26 namespace BamTools {
27   
28 // defaults
29 //
30 // ** These defaults should be tweaked & 'optimized' per testing ** //
31 //
32 //    I say 'optimized' because each system will naturally perform
33 //    differently.  We will attempt to determine a sensible
34 //    compromise that should perform well on average.
35 const unsigned int SORT_DEFAULT_MAX_BUFFER_COUNT  = 500000;  // max numberOfAlignments for buffer
36 const unsigned int SORT_DEFAULT_MAX_BUFFER_MEMORY = 1024;    // Mb
37
38 // -----------------------------------
39 // comparison objects (for sorting)
40
41 struct SortLessThanPosition {
42     bool operator() (const BamAlignment& lhs, const BamAlignment& rhs) {
43
44         // force unmapped alignments to end
45         if ( lhs.RefID == -1 ) return false;
46         if ( rhs.RefID == -1 ) return true;
47
48         // sort first on RefID, then by Position
49         if ( lhs.RefID != rhs.RefID )
50             return lhs.RefID < rhs.RefID;
51         else
52             return lhs.Position < rhs.Position;
53     }
54 };
55
56 struct SortLessThanName {
57     bool operator() (const BamAlignment& lhs, const BamAlignment& rhs) {
58         return lhs.Name < rhs.Name;
59     }
60 };
61     
62 } // namespace BamTools
63
64 // ---------------------------------------------
65 // SortSettings implementation
66
67 struct SortTool::SortSettings {
68
69     // flags
70     bool HasInputBamFilename;
71     bool HasMaxBufferCount;
72     bool HasMaxBufferMemory;
73     bool HasOutputBamFilename;
74     bool IsSortingByName;
75
76     // filenames
77     string InputBamFilename;
78     string OutputBamFilename;
79
80     // parameters
81     unsigned int MaxBufferCount;
82     unsigned int MaxBufferMemory;
83
84     // constructor
85     SortSettings(void)
86         : HasInputBamFilename(false)
87         , HasMaxBufferCount(false)
88         , HasMaxBufferMemory(false)
89         , HasOutputBamFilename(false)
90         , IsSortingByName(false)
91         , InputBamFilename(Options::StandardIn())
92         , OutputBamFilename(Options::StandardOut())
93         , MaxBufferCount(SORT_DEFAULT_MAX_BUFFER_COUNT)
94         , MaxBufferMemory(SORT_DEFAULT_MAX_BUFFER_MEMORY)
95     { }
96 };
97
98 // ---------------------------------------------
99 // SortToolPrivate implementation
100
101 class SortTool::SortToolPrivate {
102       
103     // ctor & dtor
104     public:
105         SortToolPrivate(SortTool::SortSettings* settings);
106         ~SortToolPrivate(void) { }
107         
108     // 'public' interface
109     public:
110         bool Run(void);
111         
112     // internal methods
113     private:
114         bool CreateSortedTempFile(vector<BamAlignment>& buffer);
115         bool GenerateSortedRuns(void);
116         bool MergeSortedRuns(void);
117         bool WriteTempFile(const vector<BamAlignment>& buffer, const string& tempFilename);
118         void SortBuffer(vector<BamAlignment>& buffer);
119         
120     // data members
121     private:
122         SortTool::SortSettings* m_settings;
123         string m_tempFilenameStub;
124         int m_numberOfRuns;
125         string m_headerText;
126         RefVector m_references;
127         vector<string> m_tempFilenames;
128 };
129
130 // constructor
131 SortTool::SortToolPrivate::SortToolPrivate(SortTool::SortSettings* settings) 
132     : m_settings(settings)
133     , m_numberOfRuns(0) 
134
135     // set filename stub depending on inputfile path
136     // that way multiple sort runs don't trip on each other's temp files
137     if ( m_settings) {
138         size_t extensionFound = m_settings->InputBamFilename.find(".bam");
139         if (extensionFound != string::npos )
140             m_tempFilenameStub = m_settings->InputBamFilename.substr(0,extensionFound);
141         m_tempFilenameStub.append(".sort.temp.");
142     }
143 }
144
145 // generates mutiple sorted temp BAM files from single unsorted BAM file
146 bool SortTool::SortToolPrivate::GenerateSortedRuns(void) {
147     
148     // open input BAM file
149     BamReader reader;
150     if ( !reader.Open(m_settings->InputBamFilename) ) {
151         cerr << "bamtools sort ERROR: could not open " << m_settings->InputBamFilename
152              << " for reading... Aborting." << endl;
153         return false;
154     }
155     
156     // get basic data that will be shared by all temp/output files 
157     SamHeader header = reader.GetHeader();
158     header.SortOrder = ( m_settings->IsSortingByName
159                        ? Constants::SAM_HD_SORTORDER_QUERYNAME
160                        : Constants::SAM_HD_SORTORDER_COORDINATE );
161     m_headerText = header.ToString();
162     m_references = reader.GetReferenceData();
163     
164     // set up alignments buffer
165     BamAlignment al;
166     vector<BamAlignment> buffer;
167     buffer.reserve( (size_t)(m_settings->MaxBufferCount*1.1) );
168     bool bufferFull = false;
169     
170     // if sorting by name, we need to generate full char data
171     // so can't use GetNextAlignmentCore()
172     if ( m_settings->IsSortingByName ) {
173
174         // iterate through file
175         while ( reader.GetNextAlignment(al)) {
176
177             // check buffer's usage
178             bufferFull = ( buffer.size() >= m_settings->MaxBufferCount );
179
180             // store alignments until buffer is "full"
181             if ( !bufferFull )
182                 buffer.push_back(al);
183
184             // if buffer is "full"
185             else {
186
187                 // push any unmapped reads into buffer,
188                 // don't want to split these into a separate temp file
189                 if ( !al.IsMapped() )
190                     buffer.push_back(al);
191
192                 // "al" is mapped, so create a sorted temp file with current buffer contents
193                 // then push "al" into fresh buffer
194                 else {
195                     CreateSortedTempFile(buffer);
196                     buffer.push_back(al);
197                 }
198             }
199         }
200     }
201
202     // sorting by position, can take advantage of GNACore() speedup
203     else {
204
205         // iterate through file
206         while ( reader.GetNextAlignmentCore(al) ) {
207
208             // check buffer's usage
209             bufferFull = ( buffer.size() >= m_settings->MaxBufferCount );
210
211             // store alignments until buffer is "full"
212             if ( !bufferFull )
213                 buffer.push_back(al);
214
215             // if buffer is "full"
216             else {
217
218                 // push any unmapped reads into buffer,
219                 // don't want to split these into a separate temp file
220                 if ( !al.IsMapped() )
221                     buffer.push_back(al);
222
223                 // "al" is mapped, so create a sorted temp file with current buffer contents
224                 // then push "al" into fresh buffer
225                 else {
226                     CreateSortedTempFile(buffer);
227                     buffer.push_back(al);
228                 }
229             }
230         }
231     }
232
233     // handle any leftover buffer contents
234     if ( !buffer.empty() )
235         CreateSortedTempFile(buffer);
236     
237     // close reader & return success
238     reader.Close();
239     return true;
240 }
241
242 bool SortTool::SortToolPrivate::CreateSortedTempFile(vector<BamAlignment>& buffer) {
243  
244     // do sorting
245     SortBuffer(buffer);
246   
247     // write sorted contents to temp file, store success/fail
248     stringstream tempStr;
249     tempStr << m_tempFilenameStub << m_numberOfRuns;
250     bool success = WriteTempFile( buffer, tempStr.str() );
251     
252     // save temp filename for merging later
253     m_tempFilenames.push_back(tempStr.str());
254     
255     // clear buffer contents & update run counter
256     buffer.clear();
257     ++m_numberOfRuns;
258     
259     // return success/fail of writing to temp file
260     // TODO: a failure returned here is not actually caught and handled anywhere
261     return success;
262 }
263
264 // merges sorted temp BAM files into single sorted output BAM file
265 bool SortTool::SortToolPrivate::MergeSortedRuns(void) {
266   
267     // open up multi reader for all of our temp files
268     // this might get broken up if we do a multi-pass system later ??
269     BamMultiReader multiReader;
270     if ( !multiReader.Open(m_tempFilenames) ) {
271         cerr << "bamtools sort ERROR: could not open BamMultiReader for merging temp files... Aborting."
272              << endl;
273         return false;
274     }
275
276     // set sort order for merge
277     if ( m_settings->IsSortingByName )
278         multiReader.SetSortOrder(BamMultiReader::SortedByReadName);
279     else
280         multiReader.SetSortOrder(BamMultiReader::SortedByPosition);
281     
282     // open writer for our completely sorted output BAM file
283     BamWriter mergedWriter;
284     if ( !mergedWriter.Open(m_settings->OutputBamFilename, m_headerText, m_references) ) {
285         cerr << "bamtools sort ERROR: could not open " << m_settings->OutputBamFilename
286              << " for writing... Aborting." << endl;
287         multiReader.Close();
288         return false;
289     }
290     
291     // while data available in temp files
292     BamAlignment al;
293     while ( multiReader.GetNextAlignmentCore(al) )
294         mergedWriter.SaveAlignment(al);
295   
296     // close readers
297     multiReader.Close();
298     mergedWriter.Close();
299     
300     // delete all temp files
301     vector<string>::const_iterator tempIter = m_tempFilenames.begin();
302     vector<string>::const_iterator tempEnd  = m_tempFilenames.end();
303     for ( ; tempIter != tempEnd; ++tempIter ) {
304         const string& tempFilename = (*tempIter);
305         remove(tempFilename.c_str());
306     }
307   
308     return true;
309 }
310
311 bool SortTool::SortToolPrivate::Run(void) {
312  
313     // this does a single pass, chunking up the input file into smaller sorted temp files, 
314     // then write out using BamMultiReader to handle merging
315     
316     if ( GenerateSortedRuns() )
317         return MergeSortedRuns();
318     else 
319         return false;
320
321     
322 void SortTool::SortToolPrivate::SortBuffer(vector<BamAlignment>& buffer) {
323  
324     // ** add further custom sort options later ?? **
325     
326     // sort buffer by desired method
327     if ( m_settings->IsSortingByName )
328         std::stable_sort ( buffer.begin(), buffer.end(), SortLessThanName() );
329     else 
330         std::stable_sort ( buffer.begin(), buffer.end(), SortLessThanPosition() );
331 }
332     
333 bool SortTool::SortToolPrivate::WriteTempFile(const vector<BamAlignment>& buffer,
334                                               const string& tempFilename)
335 {
336     // open temp file for writing
337     BamWriter tempWriter;
338     if ( !tempWriter.Open(tempFilename, m_headerText, m_references) ) {
339         cerr << "bamtools sort ERROR: could not open " << tempFilename
340              << " for writing." << endl;
341         return false;
342     }
343   
344     // write data
345     vector<BamAlignment>::const_iterator buffIter = buffer.begin();
346     vector<BamAlignment>::const_iterator buffEnd  = buffer.end();
347     for ( ; buffIter != buffEnd; ++buffIter )  {
348         const BamAlignment& al = (*buffIter);
349         tempWriter.SaveAlignment(al);
350     }
351   
352     // close temp file & return success
353     tempWriter.Close();
354     return true;
355 }
356
357 // ---------------------------------------------
358 // SortTool implementation
359
360 SortTool::SortTool(void)
361     : AbstractTool()
362     , m_settings(new SortSettings)
363     , m_impl(0)
364 {
365     // set program details
366     Options::SetProgramInfo("bamtools sort", "sorts a BAM file", "[-in <filename>] [-out <filename>] [sortOptions]");
367
368     // set up options
369     OptionGroup* IO_Opts = Options::CreateOptionGroup("Input & Output");
370     Options::AddValueOption("-in",  "BAM filename", "the input BAM file",  "",
371                             m_settings->HasInputBamFilename,  m_settings->InputBamFilename,
372                             IO_Opts, Options::StandardIn());
373     Options::AddValueOption("-out", "BAM filename", "the output BAM file", "",
374                             m_settings->HasOutputBamFilename, m_settings->OutputBamFilename,
375                             IO_Opts, Options::StandardOut());
376
377     OptionGroup* SortOpts = Options::CreateOptionGroup("Sorting Methods");
378     Options::AddOption("-byname", "sort by alignment name", m_settings->IsSortingByName, SortOpts);
379
380     OptionGroup* MemOpts = Options::CreateOptionGroup("Memory Settings");
381     Options::AddValueOption("-n",   "count", "max number of alignments per tempfile", "",
382                             m_settings->HasMaxBufferCount,  m_settings->MaxBufferCount,
383                             MemOpts, SORT_DEFAULT_MAX_BUFFER_COUNT);
384     Options::AddValueOption("-mem", "Mb", "max memory to use", "",
385                             m_settings->HasMaxBufferMemory, m_settings->MaxBufferMemory,
386                             MemOpts, SORT_DEFAULT_MAX_BUFFER_MEMORY);
387 }
388
389 SortTool::~SortTool(void) {
390
391     delete m_settings;
392     m_settings = 0;
393
394     delete m_impl;
395     m_impl = 0;
396 }
397
398 int SortTool::Help(void) {
399     Options::DisplayHelp();
400     return 0;
401 }
402
403 int SortTool::Run(int argc, char* argv[]) {
404
405     // parse command line arguments
406     Options::Parse(argc, argv, 1);
407
408     // initialize SortTool with settings
409     m_impl = new SortToolPrivate(m_settings);
410
411     // run SortTool, return success/fail
412     if ( m_impl->Run() )
413         return 0;
414     else
415         return 1;
416 }