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