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