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