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