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         bool GenerateSortedRuns(void);
  74         bool HandleBufferContents(vector<BamAlignment>& buffer);
  75         bool MergeSortedRuns(void);
  76         bool WriteTempFile(const vector<BamAlignment>& buffer, const string& tempFilename);
  77         void SortBuffer(vector<BamAlignment>& buffer);
  78
  79     // data members
  80     private:
  81         SortTool::SortSettings* m_settings;
  82         string m_tempFilenameStub;
  83         int m_numberOfRuns;
  84         string m_headerText;
  85         RefVector m_references;
  86         vector<string> m_tempFilenames;
  87 };
  88
  89 // ---------------------------------------------
  90 // SortSettings implementation
  91
  92 struct SortTool::SortSettings {
  93
  94     // flags
  95     bool HasInputBamFilename;
  96     bool HasMaxBufferCount;
  97     bool HasMaxBufferMemory;
  98     bool HasOutputBamFilename;
  99     bool IsSortingByName;
 100
 101     // filenames
 102     string InputBamFilename;
 103     string OutputBamFilename;
 104
 105     // parameters
 106     unsigned int MaxBufferCount;
 107     unsigned int MaxBufferMemory;
 108
 109     // constructor
 110     SortSettings(void)
 111         : HasInputBamFilename(false)
 112         , HasMaxBufferCount(false)
 113         , HasMaxBufferMemory(false)
 114         , HasOutputBamFilename(false)
 115         , IsSortingByName(false)
 116         , InputBamFilename(Options::StandardIn())
 117         , OutputBamFilename(Options::StandardOut())
 118         , MaxBufferCount(SORT_DEFAULT_MAX_BUFFER_COUNT)
 119         , MaxBufferMemory(SORT_DEFAULT_MAX_BUFFER_MEMORY)
 120     { }
 121 };
 122
 123 // ---------------------------------------------
 124 // SortTool implementation
 125
 126 SortTool::SortTool(void)
 127     : AbstractTool()
 128     , m_settings(new SortSettings)
 129     , m_impl(0)
 130 {
 131     // set program details
 132     Options::SetProgramInfo("bamtools sort", "sorts a BAM file", "[-in <filename>] [-out <filename>]");
 133
 134     // set up options
 135     OptionGroup* IO_Opts = Options::CreateOptionGroup("Input & Output");
 136     Options::AddValueOption("-in",  "BAM filename", "the input BAM file",  "", m_settings->HasInputBamFilename,  m_settings->InputBamFilename,  IO_Opts, Options::StandardIn());
 137     Options::AddValueOption("-out", "BAM filename", "the output BAM file", "", m_settings->HasOutputBamFilename, m_settings->OutputBamFilename, IO_Opts, Options::StandardOut());
 138
 139     OptionGroup* SortOpts = Options::CreateOptionGroup("Sorting Methods");
 140     Options::AddOption("-byname", "sort by alignment name", m_settings->IsSortingByName, SortOpts);
 141
 142     OptionGroup* MemOpts = Options::CreateOptionGroup("Memory Settings");
 143     Options::AddValueOption("-n",   "count", "max number of alignments per tempfile", "", m_settings->HasMaxBufferCount,  m_settings->MaxBufferCount,  MemOpts, SORT_DEFAULT_MAX_BUFFER_COUNT);
 144     Options::AddValueOption("-mem", "Mb",    "max memory to use",                     "", m_settings->HasMaxBufferMemory, m_settings->MaxBufferMemory, MemOpts, SORT_DEFAULT_MAX_BUFFER_MEMORY);
 145 }
 146
 147 SortTool::~SortTool(void) {
 148
 149     delete m_settings;
 150     m_settings = 0;
 151
 152     delete m_impl;
 153     m_impl = 0;
 154 }
 155
 156 int SortTool::Help(void) {
 157     Options::DisplayHelp();
 158     return 0;
 159 }
 160
 161 int SortTool::Run(int argc, char* argv[]) {
 162
 163     // parse command line arguments
 164     Options::Parse(argc, argv, 1);
 165
 166     // run internal SortTool implementation, return success/fail
 167     m_impl = new SortToolPrivate(m_settings);
 168
 169     if ( m_impl->Run() ) return 0;
 170     else return 1;
 171 }
 172
 173 // ---------------------------------------------
 174 // SortToolPrivate implementation
 175
 176 // constructor
 177 SortTool::SortToolPrivate::SortToolPrivate(SortTool::SortSettings* settings)
 178     : m_settings(settings)
 179     , m_tempFilenameStub("bamtools.sort.temp.")
 180     , m_numberOfRuns(0)
 181 { }
 182
 183 // destructor
 184 SortTool::SortToolPrivate::~SortToolPrivate(void) { }
 185
 186 // generates mutiple sorted temp BAM files from single unsorted BAM file
 187 bool SortTool::SortToolPrivate::GenerateSortedRuns(void) {
 188
 189     // open input BAM file
 190     BamReader inputReader;
 191     inputReader.Open(m_settings->InputBamFilename);
 192
 193     // get basic data that will be shared by all temp/output files
 194     m_headerText = inputReader.GetHeaderText();
 195     m_references = inputReader.GetReferenceData();
 196
 197     // set up alignments buffer
 198     vector<BamAlignment> buffer;
 199     buffer.reserve(m_settings->MaxBufferCount);
 200
 201     // while data available
 202     BamAlignment al;
 203     while ( inputReader.GetNextAlignmentCore(al) ) {
 204
 205         // store alignments in buffer
 206         buffer.push_back(al);
 207
 208         // if buffer is full, handle contents (sort & write to temp file)
 209         if ( buffer.size() == m_settings->MaxBufferCount )
 210             HandleBufferContents(buffer);
 211     }
 212
 213     // handle any remaining buffer contents
 214     if ( buffer.size() > 0 )
 215         HandleBufferContents(buffer);
 216
 217     // close reader & return success
 218     inputReader.Close();
 219     return true;
 220 }
 221
 222 bool SortTool::SortToolPrivate::HandleBufferContents(vector<BamAlignment>& buffer ) {
 223
 224     // do sorting
 225     SortBuffer(buffer);
 226
 227     // write sorted contents to temp file, store success/fail
 228     stringstream tempStr;
 229     tempStr << m_tempFilenameStub << m_numberOfRuns;
 230     bool success = WriteTempFile( buffer, tempStr.str() );
 231
 232     // save temp filename for merging later
 233     m_tempFilenames.push_back(tempStr.str());
 234
 235     // clear buffer contents & update run counter
 236     buffer.clear();
 237     ++m_numberOfRuns;
 238
 239     // return success/fail of writing to temp file
 240     return success;
 241 }
 242
 243 // merges sorted temp BAM files into single sorted output BAM file
 244 bool SortTool::SortToolPrivate::MergeSortedRuns(void) {
 245
 246     // open up multi reader for all of our temp files
 247     // this might get broken up if we do a multi-pass system later ??
 248     BamMultiReader multiReader;
 249     multiReader.Open(m_tempFilenames, false, true);
 250
 251     // open writer for our completely sorted output BAM file
 252     BamWriter mergedWriter;
 253     mergedWriter.Open(m_settings->OutputBamFilename, m_headerText, m_references);
 254
 255     // while data available in temp files
 256     BamAlignment al;
 257     while ( multiReader.GetNextAlignment(al) ) {
 258         mergedWriter.SaveAlignment(al);
 259     }
 260
 261     // close readers
 262     multiReader.Close();
 263     mergedWriter.Close();
 264
 265     // delete all temp files
 266     vector<string>::const_iterator tempIter = m_tempFilenames.begin();
 267     vector<string>::const_iterator tempEnd  = m_tempFilenames.end();
 268     for ( ; tempIter != tempEnd; ++tempIter ) {
 269         const string& tempFilename = (*tempIter);
 270         remove(tempFilename.c_str());
 271     }
 272
 273     return true;
 274 }
 275
 276 bool SortTool::SortToolPrivate::Run(void) {
 277
 278     // this does a single pass, chunking up the input file into smaller sorted temp files,
 279     // then write out using BamMultiReader to handle merging
 280
 281     if ( GenerateSortedRuns() )
 282         return MergeSortedRuns();
 283     else
 284         return false;
 285 }
 286
 287 void SortTool::SortToolPrivate::SortBuffer(vector<BamAlignment>& buffer) {
 288
 289     // ** add further custom sort options later ?? **
 290
 291     // sort buffer by desired method
 292     if ( m_settings->IsSortingByName )
 293         sort ( buffer.begin(), buffer.end(), SortLessThanName() );
 294     else
 295         sort ( buffer.begin(), buffer.end(), SortLessThanPosition() );
 296 }
 297
 298
 299 bool SortTool::SortToolPrivate::WriteTempFile(const vector<BamAlignment>& buffer, const string& tempFilename) {
 300
 301     // open temp file for writing
 302     BamWriter tempWriter;
 303     tempWriter.Open(tempFilename, m_headerText, m_references);
 304
 305     // write data
 306     vector<BamAlignment>::const_iterator buffIter = buffer.begin();
 307     vector<BamAlignment>::const_iterator buffEnd  = buffer.end();
 308     for ( ; buffIter != buffEnd; ++buffIter )  {
 309         const BamAlignment& al = (*buffIter);
 310         tempWriter.SaveAlignment(al);
 311     }
 312
 313     // close temp file & return success
 314     tempWriter.Close();
 315     return true;
 316 }