source: trunk/yat/omic/Pileup.h @ 3336

#ifndef theplu_yat_omic_pileup
#define theplu_yat_omic_pileup

// $Id: Pileup.h 3336 2014-10-24 23:27:45Z peter $

/*
  Copyright (C) 2014 Peter Johansson

  This file is part of the yat library, http://dev.thep.lu.se/yat

  The yat library is free software; you can redistribute it and/or
  modify it under the terms of the GNU General Public License as
  published by the Free Software Foundation; either version 3 of the
  License, or (at your option) any later version.

  The yat library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

#include "BamRead.h"

#include "yat/utility/CigarIterator.h"
#include "yat/utility/yat_assert.h"

#include <boost/iterator/iterator_traits.hpp>
#include <boost/concept/assert.hpp>
#include <boost/shared_ptr.hpp>

#include <algorithm>
#include <list>

namespace theplu {
namespace yat {
namespace omic {

  /**
     Pileup takes a range of BamReads and pile them up on the current
     position. User can investigate the pile of reads at current
     (locus) position via begin() and end(), which provide access to
     reads at current position. The position can be incremented and in
     that way the user can investigate each position of the genome.

     \since new in yat 0.13

     Type Requirments:
     - \c Iterator is a \input_iterator
     - \c value_type is convertible to BamRead
   */
  template<typename Iterator>
  class Pileup
  {
  public:
    /**
       Essentially a BamRead but with additional information regarding
       base we currently pointing to
     */
    class Entry
    {
    public:
      /**
         \brief default constructor
       */
      Entry(void) {}

      /**
         Create an Entry with read \a b. Entry is initialized to point
         to first aligned base (b.pos()).
       */
      Entry(const BamRead& b);

      /**
         \return const reference to BamRead
       */
      const BamRead& bam(void) const
      {
        YAT_ASSERT(bam_.get());
        return *bam_;
      }

      /**
         \return cigar operation at current position
       */
      uint8_t cigar_op(void) const { return *cigar_; }


      /**
         \return length of current cigar element
       */
      uint32_t cigar_oplen(void) const
      {
        YAT_ASSERT(bam_.get());
        YAT_ASSERT(cigar_.base() >= bam_->cigar());
        return bam_cigar_oplen(*cigar_.base());
      }


      /**
         \return base quality of current position
       */
      unsigned short qual(void) const
      {
        YAT_ASSERT(bam_.get());
        YAT_ASSERT(qpos_ < bam_->sequence_length());
        return bam_->qual(qpos_);
      }

      /**
         \return base of current position
       */
      uint8_t sequence(void) const;

      /**
         Increment CIGAR. If current CIGAR element is consuming query
         (not deletion), increment to next base on query.
       */
      void increment(void);

      /**
         \return true if this position was deleted in this query
       */
      bool is_del(void) const
      {
        return cigar_type()==2;
      }

      /**
         \return index of base at this position
       */
      size_t qpos(void) const { return qpos_; }
    private:
      boost::shared_ptr<BamRead> bam_;
      // index of base pointed to
      size_t qpos_;
      yat::utility::CigarIterator<const uint32_t*> cigar_;
      size_t skip_;
      // return 0, 1, 2, or 3 see bam_cigar_type in yat/utility/Cigar.h
      uint8_t cigar_type(void) const { return bam_cigar_type(cigar_op()); }
      friend class Pileup;
    };

  public:
    /**
       Const iterator that can be used to iterate over reads that
       overlap with current position (as deined by tid() and pos()).
     */
    typedef typename std::list<Entry>::const_iterator const_iterator;

    /**
       Create an empty Pileup
     */
    Pileup(void) {}

    /**
       Create a Pileup that will use reads in range [first, last). The
       range must be sorted as defined by BamLessPos.

       Iterator is a \readable_iterator.
     */
    Pileup(Iterator first, Iterator last);

    /**
       \return an iterator that points to the first Entry in Pileup.

       Iteration is done in same order as range provided in
       constructor. Reads that do not overlap with current position
       (tid:pos) are ignored.
     */
    const_iterator begin(void) const;

    /**
       \return an iterator that points one past the last element in
       the Pileup.
     */
    const_iterator end(void) const;

    /**
       \return true if there are more positions to increment into
     */
    bool good(void) const { return !(bam_ == bam_end_ && data_.empty()); }

    /**
       \brief step to next position.

       If any Entry has a insertion at current position, increment
       each Entry that has an insertion. Otherwise increment to next
       position and update each Entry accordingly.

       If the position is not covered by any read, the Pileup
       fastforward to next covered locus.
     */
    void increment(void);

    /**
       \return template id as defined in bam file header
     */
    int32_t tid(void) const { return tid_; }

    /**
       \return position
     */
    int32_t pos(void) const { return pos_; }

    /**
       \return true if position does not exist in reference, i.e.,
       inserted region.
     */
    bool skip_ref(void) const;
  private:
    typedef typename std::list<Entry>::iterator iterator;

    Iterator bam_;
    Iterator bam_end_;
    int32_t pos_;
    int32_t tid_;
    std::list<Entry> data_;
    uint32_t skip_ref_;
  };


  /**
     \relates Pileup

     \since new in yat 0.13
   */
  template<typename Iterator>
  Pileup<Iterator> make_pileup(Iterator first, Iterator last)
  {
    return Pileup<Iterator>(first, last);
  }

  // implementations

  template<typename Iterator>
  Pileup<Iterator>::Pileup(Iterator first, Iterator last)
    : bam_(first), bam_end_(last), pos_(0), tid_(0), skip_ref_(0)
  {
    BOOST_CONCEPT_ASSERT((boost::InputIterator<Iterator>));
    using boost::Convertible;
    typedef typename boost::iterator_value<Iterator>::type value_type;
    BOOST_CONCEPT_ASSERT((Convertible<value_type, BamRead>));
    increment();
  }


  template<typename Iterator>
  typename Pileup<Iterator>::const_iterator Pileup<Iterator>::begin(void) const
  {
    return data_.begin();
  }


  template<typename Iterator>
  typename Pileup<Iterator>::const_iterator Pileup<Iterator>::end(void) const
  {
    return data_.end();
  }


  template<typename Iterator>
  void Pileup<Iterator>::increment(void)
  {
    if (!data_.empty()) {
      // in insertion
      if (skip_ref()) {
        for (iterator it = data_.begin(); it!=data_.end(); ++it) {
          if (it->cigar_type() == 1) {
            it->increment();
          }
          --it->skip_;
        }
        --skip_ref_;
      }
      // walk one base forward on reference
      else {
        ++pos_;
        typename std::list<Entry>::iterator read = data_.begin();
        while (read != data_.end()) {
          read->increment();
          // remove entry if read is entirely left of pos
          if (read->bam().end() <= pos_) { // FIXME speedup possible
            data_.erase(read++);
            continue;
          }
          // check if we're stepping into insertion
          if (read->cigar_op() == BAM_CINS) {
            skip_ref_ = std::max(skip_ref_, read->cigar_oplen());
          }
          ++read;
        }

        // stepped into an insertion
        if (skip_ref_)
          for (iterator iter=data_.begin(); iter!=data_.end(); ++iter)
            iter->skip_ = skip_ref_;
      }
    }


    // fast forward to next covered locus
    if (data_.empty()) {
      if (bam_==bam_end_)
        return;
      tid_ = bam_->core().tid;
      pos_ = bam_->core().pos;
    }

    // read data
    while (bam_!=bam_end_ && bam_->core().tid==tid_ && bam_->pos()<=pos_) {
      // skip unmapped reads
      if (!(bam_->core().flag & BAM_FUNMAP)) {
        data_.push_back(Entry(*bam_));
      }
      ++bam_;
    }
  }


  template<typename Iterator>
  bool Pileup<Iterator>::skip_ref(void) const
  {
    return skip_ref_;
  }


  template<typename Iterator>
  Pileup<Iterator>::Entry::Entry(const BamRead& b)
    : bam_(new BamRead(b)), qpos_(0), cigar_(bam_->cigar()), skip_(0)
  {
    YAT_ASSERT(cigar_.base() == bam_->cigar());
    if (bam_->core().n_cigar==0)
      return;
    // iterate cigar to first match
    while (qpos_ < bam_->sequence_length() && cigar_type() != 3) {
      // type is either 0 (hardclipped) or 1 (softclipped)
      YAT_ASSERT(cigar_type()!=2);
      if (cigar_type() == 1)
        ++qpos_;
      ++cigar_;
    }
  }


  template<typename Iterator>
  void Pileup<Iterator>::Entry::increment(void)
  {
    YAT_ASSERT(cigar_.base() >= bam_->cigar());
    if (cigar_type() & 1)
      ++qpos_;
    ++cigar_;
    YAT_ASSERT(cigar_.base() >= bam_->cigar());
  }


  template<typename Iterator>
  uint8_t Pileup<Iterator>::Entry::sequence(void) const
  {
    if (skip_==0) {
      if (cigar_type() & 1)
        return bam_->sequence(qpos_);
      return 0;
    }
    if (cigar_op()!=BAM_CINS)
      return 0;
    return bam_->sequence(qpos_);
  }

}}}
#endif