moving digest.rb

[biopieces.git] / code_ruby / lib / maasha / align.rb

# Copyright (C) 2007-2011 Martin A. Hansen.

# This program 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 2
# of the License, or (at your option) any later version.

# This program 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, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.

# http://www.gnu.org/copyleft/gpl.html

# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# This software is part of the Biopieces framework (www.biopieces.org).

# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

require 'pp'
require 'open3'
require 'narray'
require 'maasha/fasta'

class AlignError < StandardError; end;

ALPH_DNA  = %w{A T C G}
ALPH_AMBI = %w{A T C G M R W S Y K V H D B N}

BIT_INDEL = 0
# BIT_A = 1 << 0
# BIT_T = 1 << 1
# BIT_C = 1 << 2
# BIT_G = 1 << 3

BIT_M = BIT_A | BIT_C
BIT_R = BIT_A | BIT_G
BIT_W = BIT_A | BIT_T
BIT_S = BIT_C | BIT_G
BIT_Y = BIT_C | BIT_T
BIT_K = BIT_G | BIT_T
BIT_V = BIT_A | BIT_C | BIT_G
BIT_H = BIT_A | BIT_C | BIT_T
BIT_D = BIT_A | BIT_G | BIT_T
BIT_B = BIT_C | BIT_G | BIT_T
BIT_N = BIT_G | BIT_A | BIT_T | BIT_C

BITMAP = [
  BIT_A,
  BIT_T,
  BIT_C,
  BIT_G,
  BIT_M,
  BIT_R,
  BIT_W,
  BIT_S,
  BIT_Y,
  BIT_K,
  BIT_V,
  BIT_H,
  BIT_D,
  BIT_B,
  BIT_N
]

TR_NUC = "-"                   + ALPH_AMBI.join("").downcase
TR_HEX = [BIT_INDEL].pack("C") + BITMAP.pack("C*")

ROW_A = 0
ROW_T = 1
ROW_C = 2
ROW_G = 3

class Fasta
  def initialize(io)
    @io = io
  end

  def close
    @io.close
  end
end

# Class for aligning sequences.
class Align
  attr_accessor :options
  attr_reader   :entries

  # Class method to align sequences in a list of Seq objects and
  # return these as a new list of Seq objects.
  def self.muscle(entries, verbose = false)
    result   = []
    index    = {}

    Open3.popen3("muscle", "-quiet") do |stdin, stdout, stderr|
      entries.each do |entry|
        raise AlignError, "Duplicate sequence name: #{entry.seq_name}" if index.has_key? entry.seq_name

        index[entry.seq_name] = entry.dup

        stdin.puts entry.to_fasta
      end

      stdin.close

      stderr.each_line { |line| $stderr.puts line } if verbose

      aligned_entries = Fasta.new(stdout)

      aligned_entries.each do |fa_entry|
        fq_entry = index[fa_entry.seq_name]

        fa_entry.seq.scan(/-+/) do |m|
          fq_entry.seq  = fq_entry.seq[0 ... $`.length]  + ('-' * m.length) + fq_entry.seq[$`.length .. -1]
          fq_entry.qual = fq_entry.qual[0 ... $`.length] + ('@' * m.length) + fq_entry.qual[$`.length .. -1] unless fq_entry.qual.nil?
        end

        result << fq_entry
      end
    end

    self.new(result)
  end
  
  # Class method to create a pairwise alignment of two given Seq objects. The
  # alignment is created by casting a search space the size of the sequences
  # and save the longest perfect match between the sequences and recurse into
  # the left and right search spaced around this match. When all search spaces
  # are exhausted the saved matches are used to insert gaps in the sequences.
  def self.pair(q_entry, s_entry)
    q_space_beg = 0
    q_space_end = q_entry.length
    s_space_beg = 0
    s_space_end = s_entry.length

    q_entry.seq.downcase!
    s_entry.seq.downcase!

    ap = AlignPair.new(q_entry.seq, s_entry.seq)

    ap.align(q_space_beg, q_space_end, s_space_beg, s_space_end, 256)

    ap.upcase_match
    ap.insert_gaps

    #ap.dump_bp and exit

    self.new([q_entry, s_entry])
  end

  # Method to initialize an Align object with a list of aligned Seq objects.
  def initialize(entries, options = {})
    @entries = entries
    @options = options
  end

  # Method that returns the length of the alignment.
  def length
    @entries.first.length
  end

  # Method that returns the number of members or sequences in the alignment.
  def members
    @entries.size
  end

  # Method that returns the identity of an alignment with two members.
  def identity
    if self.members != 2
      raise AlignError "Bad number of members for similarity calculation: #{self.members}"
    end

    na1 = NArray.to_na(@entries[0].seq.upcase, "byte")
    na2 = NArray.to_na(@entries[1].seq.upcase, "byte")

    shared   = (na1 - na2).count_false
    total    = (@entries[0].length < @entries[1].length) ? @entries[0].length : @entries[1].length
    identity = shared.to_f / total

    identity
  end

  # Method to create a consensus sequence from an Align object and
  # return a new Seq object with the consensus.
  def consensus
    cons = Consensus.new(@entries, @options)
    cons.consensus
  end

  # Method to pretty print an alignment from an Align object.
  def to_s
    cons = Consensus.new(@entries, @options)
    cons.mask_entries!

    max_name = @entries.group_by { |entry| entry.seq_name.length }.max.first

    output = ""

    @entries.each do |entry|
      output << entry.seq_name + (" " * (max_name + 3 - entry.seq_name.length )) + entry.seq + $/
    end

    cons_entry = cons.consensus 

    output << " " * (max_name + 3) + cons_entry.seq
    output << $/ + " " * (max_name + 3) + cons_entry.qual.tr("[@-h]", "           ..........ooooooooooOOOOOOOOOO") unless cons_entry.qual.nil?
    output
  end

  # Method for iterating each of the aligned sequences.
  def each
    if block_given?
      @entries.each { |entry| yield entry }
    else
      return @entries
    end
  end

  private

  # Class for creating a pairwise alignment of two specified sequences. The
  # alignment is created by casting a search space the size of the sequences
  # and save the longest perfect match between the sequences and recurse into
  # the left and right search spaced around this match. When all search spaces
  # are exhausted the saved matches are used to insert gaps in the sequences.
  class AlignPair
    attr_reader :matches

    # Method to initialize an AlignPair object given two sequence strings.
    def initialize(q_seq, s_seq)
      @q_seq   = q_seq
      @s_seq   = s_seq
      @matches = []
    end

    # Method that locates the longest perfect match in a specified search space
    # by comparing two sequence indeces. The first index is created by hashing
    # for the first sequence the position of all non-overlapping unique oligos
    # with the oligo as key an the position as value. The second index is
    # created for the second sequence by hashing the position of all
    # overlapping unique oligos in a similar way. The longest match is located 
    # by comparing these indeces and the longest match is expanded maximally
    # within the specified search space. Finally, the longest match is used
    # to define a left and a right search space which are recursed into
    # until no more matches are found.
    def align(q_space_beg, q_space_end, s_space_beg, s_space_end, length)
#      puts "search space: #{q_space_beg}-#{q_space_end} x #{s_space_beg}-#{s_space_end}"
      new_matches = []

      while new_matches.empty? and length > 0
        if @q_seq.length >= length and @s_seq.length >= length
          q_index = index_seq(q_space_beg, q_space_end, @q_seq, length, length)
          s_index = index_seq(s_space_beg, s_space_end, @s_seq, length, 1)
          new_matches = find_matches(q_index, s_index, length)
        end

        unless new_matches.empty?
          longest_match = new_matches.sort_by { |match| match.length }.last
          longest_match.expand(@q_seq, @s_seq, q_space_beg, q_space_end, s_space_beg, s_space_end)
#          pp longest_match

          @matches << longest_match

          q_space_beg_left  = (q_space_beg == 0) ? 0 : q_space_beg + 1
          s_space_beg_left  = (s_space_beg == 0) ? 0 : s_space_beg + 1
          q_space_end_left  = longest_match.q_beg - 1
          s_space_end_left  = longest_match.s_beg - 1

          q_space_beg_right = longest_match.q_end + 1
          s_space_beg_right = longest_match.s_end + 1
          q_space_end_right = (q_space_end == @q_seq.length) ? @q_seq.length : q_space_end - 1
          s_space_end_right = (s_space_end == @s_seq.length) ? @s_seq.length : s_space_end - 1

          if q_space_end_left - q_space_beg_left > 0 and s_space_end_left - s_space_beg_left > 0
            align(q_space_beg_left,  q_space_end_left,  s_space_beg_left,  s_space_end_left,  length)
          end

          if q_space_end_right - q_space_beg_right > 0 and s_space_end_right - s_space_beg_right > 0
            align(q_space_beg_right, q_space_end_right, s_space_beg_right, s_space_end_right, length)
          end
        end

        length /= 2
      end
    end

    # Method for debugging purposes that upcase matching sequence while non-matches
    # sequence is kept in lower case.
    def upcase_match
      @matches.each do |match|
        @q_seq[match.q_beg .. match.q_end] = @q_seq[match.q_beg .. match.q_end].upcase
        @s_seq[match.s_beg .. match.s_end] = @s_seq[match.s_beg .. match.s_end].upcase
      end
    end

    # Method that insert gaps in sequences based on a list of matches and thus
    # creating an alignment.
    # TODO check boundaries!
    def insert_gaps
      @matches.sort_by! { |m| m.q_beg }

      q_gaps = 0
      s_gaps = 0

      @matches.each do |match|
        diff = (q_gaps + match.q_beg) - (s_gaps + match.s_beg)

        #pp "q_gaps #{q_gaps}   s_gaps #{s_gaps}   diff #{diff}"

        if diff < 0
          @q_seq.insert(match.q_beg + q_gaps, "-" * diff.abs)
          q_gaps += diff.abs
        elsif diff > 0
          @s_seq.insert(match.s_beg + s_gaps, "-" * diff.abs)
          s_gaps += diff.abs
        end
      end

      diff = @q_seq.length - @s_seq.length

      if diff < 0
        @q_seq << ("-" * diff.abs)
      else
        @s_seq << ("-" * diff.abs)
      end
    end
    
    # Method that dumps all matches as Biopiece records for use with
    # plot_matches.
    def dump_bp
      @matches.sort_by! { |m| m.q_beg }

      @matches.each { |m| 
        puts "Q_BEG: #{m.q_beg}"
        puts "Q_END: #{m.q_end}"
        puts "S_BEG: #{m.s_beg}"
        puts "S_END: #{m.s_end}"
        puts "STRAND: +"
        puts "---"
      }
    end

    private

    # Method for indexing a sequence given a search space, oligo size
    # and step size. All oligo positions are saved in a lookup hash
    # where the oligo is the key and the position is the value. Non-
    # unique oligos are removed and the index returned. 
    def index_seq(space_beg, space_end, seq, oligo_size, step_size)
      index       = Hash.new
      index_count = Hash.new(0)

      (space_beg ... space_end).step(step_size).each do |pos|
        oligo = seq[pos ... pos + oligo_size]

        if oligo.length == oligo_size
          index[oligo]        = pos
          index_count[oligo] += 1
        end
      end

      index_count.each do |oligo, count|
        index.delete(oligo) if count > 1
      end

      index
    end

    # Method for locating matches between two indexes where one is created
    # using non-overlapping unique oligos and their positions and one is
    # using overlapping unique oligos and positions. Thus iterating over the
    # non-overlapping oligos and checking last match, this can be extended.
    def find_matches(q_index, s_index, length)
      matches = []

      q_index.each do |oligo, q_pos|
        if s_index[oligo]
          s_pos = s_index[oligo]

          last_match = matches.last
          new_match  = Match.new(q_pos, s_pos, length)

          if last_match and
             last_match.q_beg + last_match.length == new_match.q_beg and
             last_match.s_beg + last_match.length == new_match.s_beg

             last_match.length += length
          else
            matches << new_match
          end
        end
      end

      matches
    end

    # Class for Match objects.
    class Match
      attr_accessor :length
      attr_reader :q_beg, :s_beg

      def initialize(q_beg, s_beg, length)
        @q_beg  = q_beg
        @s_beg  = s_beg
        @length = length
      end

      def q_end
        @q_beg + @length - 1
      end

      def s_end
        @s_beg + @length - 1
      end

      def to_s
        "q_beg=#{q_beg} q_end=#{q_end} s_beg=#{s_beg} s_end=#{s_end} length=#{length}"
      end

      # Method to expand a match as far as possible to the left and right
      # within a given search space.
      def expand(q_seq, s_seq, q_space_beg, q_space_end, s_space_beg, s_space_end)
        expand_left(q_seq, s_seq, q_space_beg, s_space_beg)
        expand_right(q_seq, s_seq, q_space_end, s_space_end)
      end

      private

      # Method to expand a match as far as possible to the left within a given
      # search space.
      def expand_left(q_seq, s_seq, q_space_beg, s_space_beg)
        while @q_beg > q_space_beg and @s_beg > s_space_beg and q_seq[@q_beg - 1] == s_seq[@s_beg - 1]
          @q_beg  -= 1
          @s_beg  -= 1
          @length += 1
        end
      end

      # Method to expand a match as far as possible to the right within a given
      # search space.
      def expand_right(q_seq, s_seq, q_space_end, s_space_end)
        while q_end + 1 <= q_space_end and s_end + 1 <= s_space_end and q_seq[q_end + 1] == s_seq[s_end + 1]
          @length += 1
        end
      end
    end
  end

  class Consensus
    # Method to initialize a Consensus object given a list of aligned Seq object.
    def initialize(entries, options)
      @entries = entries
      @options = options

      @cols = entries.first.seq.length
      @rows = entries.size

      @has_qual = entries.first.qual.nil? ? false : true

      @na_seq  = NArray.byte(@cols, @rows)
      @na_qual = NArray.byte(@cols, @rows) if @has_qual

      na_add_entries
      consensus_calc
    end

    # Method that lowercase residues that have been removed during
    # the determination of the consensus sequence.
    def mask_entries!
      na_seq = NArray.byte(@cols, @rows)

      @entries.each_with_index do |entry, i|
        na_seq[true, i]  = NArray.to_na(entry.seq.upcase, "byte")
      end

      na_seq += ((na_seq.ne('-'.ord) - @na_seq.ne(0)) * ' '.ord)

      @entries.each_with_index do |entry, i|
        entry.seq = na_seq[true, i].to_s
      end
    end

    # Method that returns a Sequence object with a consensus sequence
    # for the entries in an Align object.
    def consensus
      new_seq      = Seq.new
      new_seq.seq  = consensus_seq
      new_seq.qual = consensus_qual if @has_qual
      new_seq.type = "dna"

      new_seq
    end

    private

    # Method to add a Seq entry object to the two NArrays; @na_seq and @na_qual
    def na_add_entries
      @entries.each_with_index do |entry, i|
        @na_seq[true, i]  = NArray.to_na(entry.seq.downcase.tr(TR_NUC, TR_HEX), "byte")
        @na_qual[true, i] = NArray.to_na(entry.qual, "byte") - SCORE_BASE if @has_qual
      end
    end

    # Method to calculate a consensus sequence from a list of sequenced stored in two
    # NArrays.
    def consensus_calc
      if @has_qual
        if @options.has_key? :quality_min
          mask = mask_quality_min

          @na_seq  *= mask
          @na_qual *= mask
        end

        if @options.has_key? :quality_mean
          mask = mask_quality_mean

          @na_seq  *= mask
          @na_qual *= mask
        end
      end

      if @options.has_key? :sequence_min
        mask = mask_sequence_min

        @na_seq  *= mask
        @na_qual *= mask if @has_qual
      end

      if @options.has_key? :gap_max
        mask = mask_gap_max

        @na_seq  *= mask
        @na_qual *= mask if @has_qual
      end

      if @options.has_key? :residue_min
        mask = mask_residue_min

        @na_seq  *= mask
        @na_qual *= mask if @has_qual
      end
    end

    # Mask that indicates which columns have more than sequence_min sequences.
    # Columns with less than sequence_min are 0'ed, otherwise set to 1.
    def mask_sequence_min
      mask  = NArray.byte(@cols, @rows) + 1
      mask *= ((@na_seq > 0).to_type("int").sum(1) >= @options[:sequence_min])
      mask
    end

    # Mask that indicates which residue frequencies that are above the residue_min.
    # The residue frequencies are calculated for each column and residue type as the
    # number of each residue type over the sum of all non-gap residues in that column.
    # Positions with residues above the residue_min are indicated with 1.
    def mask_residue_min
      cons_min = @options[:residue_min]
      factor   = 1 / @na_seq.ne(0).to_type("float").sum(1)

      mask_A = (@na_seq & BIT_A > 0).to_type("int")
      mask_T = (@na_seq & BIT_T > 0).to_type("int")
      mask_C = (@na_seq & BIT_C > 0).to_type("int")
      mask_G = (@na_seq & BIT_G > 0).to_type("int")

      mask_A = (mask_A * mask_A.sum(1)) * factor >= cons_min
      mask_T = (mask_T * mask_T.sum(1)) * factor >= cons_min
      mask_C = (mask_C * mask_C.sum(1)) * factor >= cons_min
      mask_G = (mask_G * mask_G.sum(1)) * factor >= cons_min

      mask_A | mask_T | mask_C | mask_G
    end

    # Mask that indicates which columns contain fewer gaps than max_gap.
    # Columns with more gaps are 0'ed, others are set to 1.
    def mask_gap_max
      mask  = NArray.byte(@cols, @rows) + 1
      mask *= @na_seq.ne(0).to_type("float").sum(1) / @rows > @options[:gap_max]

      mask
    end

    # Mask that indicates which residues in an alignment are above quality_min.
    # Positions with subquality are 0'ed - all others are set to 1.
    def mask_quality_min
      @na_qual > @options[:quality_min]
    end

    # Mask that indicates which columns have a quality mean above quality_mean which
    # is the mean of all non-gap quality residues in that column. Columns with less then
    # quality_mean are 0'ed, otherwise set to 1.
    def mask_quality_mean
      mask     = NArray.byte(@cols, @rows) + 1
      residues = @na_seq.ne(0).to_type("int").sum(1)
      quality  = @na_qual.to_type("float").sum(1)

      mask * (quality / residues).round > @options[:quality_mean]
    end

    # Method to calculate a consensus sequence from a Consensus object.
    def consensus_seq
      cons  = NArray.byte(@cols)
      cons |= (@na_seq & BIT_A).max(1)
      cons |= (@na_seq & BIT_T).max(1)
      cons |= (@na_seq & BIT_C).max(1)
      cons |= (@na_seq & BIT_G).max(1)

      cons.to_s.tr!(TR_HEX, TR_NUC).upcase
    end

    # Method to calculate a consensus quality from a Consensus object.
    def consensus_qual
      (@na_qual.mean(1).round.to_type("byte") + SCORE_BASE).to_s
    end
  end
end

__END__