]> git.donarmstrong.com Git - bamtools.git/blob - src/toolkit/bamtools_sort.cpp
Fixed: Improper sorting -byname
[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 // All rights reserved.
5 // ---------------------------------------------------------------------------
6 // Last modified: 16 December 2010 (DB)
7 // ---------------------------------------------------------------------------
8 // Sorts an input BAM file (default by position) and stores in a new BAM file.
9 // ***************************************************************************
10
11 #include <cstdio>
12 #include <algorithm>
13 #include <iostream>
14 #include <sstream>
15 #include <string>
16 #include <vector>
17
18 #include "bamtools_sort.h"
19 #include "bamtools_options.h"
20 #include "BamReader.h"
21 #include "BamMultiReader.h"
22 #include "BamWriter.h"
23
24 using namespace std;
25 using namespace BamTools;
26
27 namespace BamTools {
28   
29     // defaults
30     //
31     // ** These defaults should be tweaked & 'optimized' per testing ** //
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  = 10000; // 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 // SortToolPrivate declaration
60 class SortTool::SortToolPrivate {
61       
62     // ctor & dtor
63     public:
64         SortToolPrivate(SortTool::SortSettings* settings);
65         ~SortToolPrivate(void);
66         
67     // 'public' interface
68     public:
69         bool Run(void);
70         
71     // internal methods
72     private:
73         void ClearBuffer(vector<BamAlignment>& buffer);
74         bool GenerateSortedRuns(void);
75         bool HandleBufferContents(vector<BamAlignment>& buffer);
76         bool MergeSortedRuns(void);
77         bool WriteTempFile(const vector<BamAlignment>& buffer, const string& tempFilename);
78         void SortBuffer(vector<BamAlignment>& buffer);
79         
80     // data members
81     private:
82         SortTool::SortSettings* m_settings;
83         string m_tempFilenameStub;
84         int m_numberOfRuns;
85         string m_headerText;
86         RefVector m_references;
87         vector<string> m_tempFilenames;
88 };
89
90 // ---------------------------------------------
91 // SortSettings implementation
92
93 struct SortTool::SortSettings {
94
95     // flags
96     bool HasInputBamFilename;
97     bool HasMaxBufferCount;
98     bool HasMaxBufferMemory;
99     bool HasOutputBamFilename;
100     bool IsSortingByName;
101
102     // filenames
103     string InputBamFilename;
104     string OutputBamFilename;
105     
106     // parameters
107     unsigned int MaxBufferCount;
108     unsigned int MaxBufferMemory;
109     
110     // constructor
111     SortSettings(void)
112         : HasInputBamFilename(false)
113         , HasMaxBufferCount(false)
114         , HasMaxBufferMemory(false)
115         , HasOutputBamFilename(false)
116         , IsSortingByName(false)
117         , InputBamFilename(Options::StandardIn())
118         , OutputBamFilename(Options::StandardOut())
119         , MaxBufferCount(SORT_DEFAULT_MAX_BUFFER_COUNT)
120         , MaxBufferMemory(SORT_DEFAULT_MAX_BUFFER_MEMORY)
121     { } 
122 };  
123
124 // ---------------------------------------------
125 // SortTool implementation
126
127 SortTool::SortTool(void)
128     : AbstractTool()
129     , m_settings(new SortSettings)
130     , m_impl(0)
131 {
132     // set program details
133     Options::SetProgramInfo("bamtools sort", "sorts a BAM file", "[-in <filename>] [-out <filename>] [sortOptions]");
134     
135     // set up options 
136     OptionGroup* IO_Opts = Options::CreateOptionGroup("Input & Output");
137     Options::AddValueOption("-in",  "BAM filename", "the input BAM file",  "", m_settings->HasInputBamFilename,  m_settings->InputBamFilename,  IO_Opts, Options::StandardIn());
138     Options::AddValueOption("-out", "BAM filename", "the output BAM file", "", m_settings->HasOutputBamFilename, m_settings->OutputBamFilename, IO_Opts, Options::StandardOut());
139     
140     OptionGroup* SortOpts = Options::CreateOptionGroup("Sorting Methods");
141     Options::AddOption("-byname", "sort by alignment name", m_settings->IsSortingByName, SortOpts);
142     
143     OptionGroup* MemOpts = Options::CreateOptionGroup("Memory Settings");
144     Options::AddValueOption("-n",   "count", "max number of alignments per tempfile", "", m_settings->HasMaxBufferCount,  m_settings->MaxBufferCount,  MemOpts, SORT_DEFAULT_MAX_BUFFER_COUNT);
145     Options::AddValueOption("-mem", "Mb",    "max memory to use",                     "", m_settings->HasMaxBufferMemory, m_settings->MaxBufferMemory, MemOpts, SORT_DEFAULT_MAX_BUFFER_MEMORY);
146 }
147
148 SortTool::~SortTool(void) {
149     
150     delete m_settings;
151     m_settings = 0;
152     
153     delete m_impl;
154     m_impl = 0;
155 }
156
157 int SortTool::Help(void) {
158     Options::DisplayHelp();
159     return 0;
160 }
161
162 int SortTool::Run(int argc, char* argv[]) {
163   
164     // parse command line arguments
165     Options::Parse(argc, argv, 1);
166     
167     // run internal SortTool implementation, return success/fail
168     m_impl = new SortToolPrivate(m_settings);
169     
170     if ( m_impl->Run() ) return 0;
171     else return 1;
172 }
173
174 // ---------------------------------------------
175 // SortToolPrivate implementation
176
177 // constructor
178 SortTool::SortToolPrivate::SortToolPrivate(SortTool::SortSettings* settings) 
179     : m_settings(settings)
180     , m_numberOfRuns(0) 
181
182     // set filename stub depending on inputfile path
183     // that way multiple sort runs don't trip on each other's temp files
184     if ( m_settings) {
185         size_t extensionFound = m_settings->InputBamFilename.find(".bam");
186         if (extensionFound != string::npos )
187             m_tempFilenameStub = m_settings->InputBamFilename.substr(0,extensionFound);
188         m_tempFilenameStub.append(".sort.temp.");
189     }
190 }
191
192 // destructor
193 SortTool::SortToolPrivate::~SortToolPrivate(void) { }
194
195 // generates mutiple sorted temp BAM files from single unsorted BAM file
196 bool SortTool::SortToolPrivate::GenerateSortedRuns(void) {
197     
198     // open input BAM file
199     BamReader inputReader;
200     if (!inputReader.Open(m_settings->InputBamFilename)) {
201         cerr << "Could not open " << m_settings->InputBamFilename << " for reading." << endl;
202         return false;
203     }
204     
205     // get basic data that will be shared by all temp/output files 
206     m_headerText = inputReader.GetHeaderText();
207     m_references = inputReader.GetReferenceData();
208     
209     // set up alignments buffer
210     BamAlignment al;
211     vector<BamAlignment> buffer;
212     buffer.reserve(m_settings->MaxBufferCount);
213     
214     // if sorting by name, we need to generate full char data
215     // so can't use GetNextAlignmentCore()
216     if ( m_settings->IsSortingByName ) {
217
218         // iterate through file
219         while ( inputReader.GetNextAlignment(al)) {
220
221             // store alignments in buffer
222             buffer.push_back(al);
223
224             // if buffer is full, handle contents (sort & write to temp file)
225             if ( buffer.size() == m_settings->MaxBufferCount )
226                 HandleBufferContents(buffer);
227         }
228
229     }
230
231     // sorting by position, can take advantage of GNACore() speedup
232     else {
233
234         // iterate through file
235         while ( inputReader.GetNextAlignmentCore(al)) {
236
237             // store alignments in buffer
238             buffer.push_back(al);
239
240             // if buffer is full, handle contents (sort & write to temp file)
241             if ( buffer.size() == m_settings->MaxBufferCount )
242                 HandleBufferContents(buffer);
243         }
244     }
245
246     // handle any remaining buffer contents
247     if ( buffer.size() > 0 )
248         HandleBufferContents(buffer);
249     
250     // close reader & return success
251     inputReader.Close();
252     return true;
253 }
254
255 bool SortTool::SortToolPrivate::HandleBufferContents(vector<BamAlignment>& buffer ) {
256  
257     // do sorting
258     SortBuffer(buffer);
259   
260     // write sorted contents to temp file, store success/fail
261     stringstream tempStr;
262     tempStr << m_tempFilenameStub << m_numberOfRuns;
263     bool success = WriteTempFile( buffer, tempStr.str() );
264     
265     // save temp filename for merging later
266     m_tempFilenames.push_back(tempStr.str());
267     
268     // clear buffer contents & update run counter
269     buffer.clear();
270     ++m_numberOfRuns;
271     
272     // return success/fail of writing to temp file
273     return success;
274 }
275
276 // merges sorted temp BAM files into single sorted output BAM file
277 bool SortTool::SortToolPrivate::MergeSortedRuns(void) {
278   
279     // open up multi reader for all of our temp files
280     // this might get broken up if we do a multi-pass system later ??
281     BamMultiReader multiReader;
282     multiReader.Open(m_tempFilenames, false, true);
283     
284     // open writer for our completely sorted output BAM file
285     BamWriter mergedWriter;
286     mergedWriter.Open(m_settings->OutputBamFilename, m_headerText, m_references);
287     
288     // while data available in temp files
289     BamAlignment al;
290     while ( multiReader.GetNextAlignmentCore(al) )
291         mergedWriter.SaveAlignment(al);
292   
293     // close readers
294     multiReader.Close();
295     mergedWriter.Close();
296     
297     // delete all temp files
298     vector<string>::const_iterator tempIter = m_tempFilenames.begin();
299     vector<string>::const_iterator tempEnd  = m_tempFilenames.end();
300     for ( ; tempIter != tempEnd; ++tempIter ) {
301         const string& tempFilename = (*tempIter);
302         remove(tempFilename.c_str());
303     }
304   
305     return true;
306 }
307
308 bool SortTool::SortToolPrivate::Run(void) {
309  
310     // this does a single pass, chunking up the input file into smaller sorted temp files, 
311     // then write out using BamMultiReader to handle merging
312     
313     if ( GenerateSortedRuns() )
314         return MergeSortedRuns();
315     else 
316         return false;
317
318     
319 void SortTool::SortToolPrivate::SortBuffer(vector<BamAlignment>& buffer) {
320  
321     // ** add further custom sort options later ?? **
322     
323     // sort buffer by desired method
324     if ( m_settings->IsSortingByName )
325         sort ( buffer.begin(), buffer.end(), SortLessThanName() );
326     else 
327         sort ( buffer.begin(), buffer.end(), SortLessThanPosition() );
328 }
329     
330     
331 bool SortTool::SortToolPrivate::WriteTempFile(const vector<BamAlignment>& buffer, const string& tempFilename) {
332
333     // open temp file for writing
334     BamWriter tempWriter;
335     tempWriter.Open(tempFilename, m_headerText, m_references);
336   
337     // write data
338     vector<BamAlignment>::const_iterator buffIter = buffer.begin();
339     vector<BamAlignment>::const_iterator buffEnd  = buffer.end();
340     for ( ; buffIter != buffEnd; ++buffIter )  {
341         const BamAlignment& al = (*buffIter);
342         tempWriter.SaveAlignment(al);
343     }
344   
345     // close temp file & return success
346     tempWriter.Close();
347     return true;
348 }