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