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