]> git.donarmstrong.com Git - bamtools.git/blob - src/toolkit/bamtools_sort.cpp
Reviewed BAM to YAML conversion contribution (thank you Chase Miller). Reorganized...
[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: 21 June 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>]");
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     vector<BamAlignment> buffer;
211     buffer.reserve(m_settings->MaxBufferCount);
212     
213     // while data available
214     BamAlignment al;
215     while ( inputReader.GetNextAlignmentCore(al)) {
216         
217         // store alignments in buffer
218         buffer.push_back(al);
219         
220         // if buffer is full, handle contents (sort & write to temp file)
221         if ( buffer.size() == m_settings->MaxBufferCount )
222             HandleBufferContents(buffer);
223     }
224     
225     // handle any remaining buffer contents
226     if ( buffer.size() > 0 )
227         HandleBufferContents(buffer);
228     
229     // close reader & return success
230     inputReader.Close();
231     return true;
232 }
233
234 bool SortTool::SortToolPrivate::HandleBufferContents(vector<BamAlignment>& buffer ) {
235  
236     // do sorting
237     SortBuffer(buffer);
238   
239     // write sorted contents to temp file, store success/fail
240     stringstream tempStr;
241     tempStr << m_tempFilenameStub << m_numberOfRuns;
242     bool success = WriteTempFile( buffer, tempStr.str() );
243     
244     // save temp filename for merging later
245     m_tempFilenames.push_back(tempStr.str());
246     
247     // clear buffer contents & update run counter
248     buffer.clear();
249     ++m_numberOfRuns;
250     
251     // return success/fail of writing to temp file
252     return success;
253 }
254
255 // merges sorted temp BAM files into single sorted output BAM file
256 bool SortTool::SortToolPrivate::MergeSortedRuns(void) {
257   
258     // open up multi reader for all of our temp files
259     // this might get broken up if we do a multi-pass system later ??
260     BamMultiReader multiReader;
261     multiReader.Open(m_tempFilenames, false, true);
262     
263     // open writer for our completely sorted output BAM file
264     BamWriter mergedWriter;
265     mergedWriter.Open(m_settings->OutputBamFilename, m_headerText, m_references);
266     
267     // while data available in temp files
268     BamAlignment al;
269     while ( multiReader.GetNextAlignmentCore(al) ) {
270         mergedWriter.SaveAlignment(al);
271     }
272   
273     // close readers
274     multiReader.Close();
275     mergedWriter.Close();
276     
277     // delete all temp files
278     vector<string>::const_iterator tempIter = m_tempFilenames.begin();
279     vector<string>::const_iterator tempEnd  = m_tempFilenames.end();
280     for ( ; tempIter != tempEnd; ++tempIter ) {
281         const string& tempFilename = (*tempIter);
282         remove(tempFilename.c_str());
283     }
284   
285     return true;
286 }
287
288 bool SortTool::SortToolPrivate::Run(void) {
289  
290     // this does a single pass, chunking up the input file into smaller sorted temp files, 
291     // then write out using BamMultiReader to handle merging
292     
293     if ( GenerateSortedRuns() )
294         return MergeSortedRuns();
295     else 
296         return false;
297
298     
299 void SortTool::SortToolPrivate::SortBuffer(vector<BamAlignment>& buffer) {
300  
301     // ** add further custom sort options later ?? **
302     
303     // sort buffer by desired method
304     if ( m_settings->IsSortingByName )
305         sort ( buffer.begin(), buffer.end(), SortLessThanName() );
306     else 
307         sort ( buffer.begin(), buffer.end(), SortLessThanPosition() );
308 }
309     
310     
311 bool SortTool::SortToolPrivate::WriteTempFile(const vector<BamAlignment>& buffer, const string& tempFilename) {
312
313     // open temp file for writing
314     BamWriter tempWriter;
315     tempWriter.Open(tempFilename, m_headerText, m_references);
316   
317     // write data
318     vector<BamAlignment>::const_iterator buffIter = buffer.begin();
319     vector<BamAlignment>::const_iterator buffEnd  = buffer.end();
320     for ( ; buffIter != buffEnd; ++buffIter )  {
321         const BamAlignment& al = (*buffIter);
322         tempWriter.SaveAlignment(al);
323     }
324   
325     // close temp file & return success
326     tempWriter.Close();
327     return true;
328 }