bp_bin/classify_taxonomy

   1 #!/usr/bin/env ruby
   2
   3 # Copyright (C) 2007-2012 Martin A. Hansen.
   4
   5 # This program is free software; you can redistribute it and/or
   6 # modify it under the terms of the GNU General Public License
   7 # as published by the Free Software Foundation; either version 2
   8 # of the License, or (at your option) any later version.
   9
  10 # This program is distributed in the hope that it will be useful,
  11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13 # GNU General Public License for more details.
  14
  15 # You should have received a copy of the GNU General Public License
  16 # along with this program; if not, write to the Free Software
  17 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  18
  19 # http://www.gnu.org/copyleft/gpl.html
  20
  21 # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
  22
  23 # This program is part of the Biopieces framework (www.biopieces.org).
  24
  25 # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> DESCRIPTION <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
  26
  27 # Classify records with taxonomy information.
  28
  29 # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
  30
  31 require 'pp'
  32 require 'maasha/biopieces'
  33
  34 class TaxTree
  35   def initialize
  36     @tree = TaxNode.new("root", "Root", 0, 0, 0.0)
  37   end
  38
  39   # Method to add to the TaxTree a GreenGenes entry.
  40   def add_gg(s_id, score, size)
  41     node = @tree
  42
  43     s_id.scan(/ ([\w])__([^;]+)/) do
  44       level = expand_level($1)
  45       name  = $2
  46
  47       if node.children[name].nil?
  48         node.children[name] = TaxNode.new(level, name, 1, size, score)
  49       else
  50         node.children[name].count += 1
  51         node.children[name].score += score
  52       end
  53
  54       node = node.children[name]
  55     end
  56   end
  57
  58   def merge(tree)
  59     node = @tree
  60
  61     tree.flatten.each do |new_node|
  62       next if new_node.level == 'root'
  63
  64       new_node.score = new_node.score / new_node.count
  65       new_node.count = 1
  66
  67       if node.children[new_node.name].nil?
  68         node.children[new_node.name] = new_node
  69       else
  70         node.children[new_node.name].count += new_node.count
  71         node.children[new_node.name].score += new_node.score
  72         node.children[new_node.name].size  += new_node.size
  73       end
  74
  75       node = node.children[new_node.name]
  76     end
  77   end
  78
  79   def flatten(node = @tree, list = [])
  80     list << TaxNode.new(node.level, node.name, node.count, node.size, node.score)
  81
  82     node.children.each do |name, child|
  83       list = flatten(child, list.dup)
  84     end
  85
  86     list
  87   end
  88
  89   def lowest_common_ancestor(node = @tree)
  90     node.children = {} if node.children.size > 1
  91
  92     node.children.each do |name, child|
  93       lowest_common_ancestor(child)
  94     end
  95   end
  96
  97   def each
  98     self.flatten.each do |node|
  99       yield node
 100     end
 101
 102     self
 103   end
 104
 105   private
 106
 107   def expand_level(level)
 108     hash = {
 109       'd' => "domain",
 110       'k' => "kingdom",
 111       'p' => "phylum",
 112       'c' => "class",
 113       'o' => "order",
 114       'f' => "family",
 115       'g' => "genus",
 116       's' => "species"
 117     }
 118
 119     raise "unknown level: #{level}" unless hash[level]
 120
 121     hash[level]
 122   end
 123
 124   class TaxNode
 125     attr_accessor :level, :name, :count, :size, :score, :children
 126
 127     def initialize(level, name, count, size, score)
 128       @level    = level
 129       @name     = name
 130       @count    = count
 131       @size     = size
 132       @score    = score
 133       @children = {}
 134     end
 135
 136     def to_bp
 137       record = {}
 138       record[:REC_TYPE] = "Classification"
 139       record[:LEVEL]    = @level
 140       record[:NAME]     = @name
 141       record[:COUNT]    = @size
 142       record[:SCORE]    = @count == 0 ? 0.0 : (@score / @count).round(2)
 143
 144       record
 145     end
 146   end
 147 end
 148
 149 casts = []
 150 casts << {:long=>'LCA',  :short=>'l', :type=>'flag', :mandatory=>false, :default=>nil, :allowed=>nil, :disallowed=>nil}
 151 casts << {:long=>'size', :short=>'s', :type=>'flag', :mandatory=>false, :default=>nil, :allowed=>nil, :disallowed=>nil}
 152
 153 options = Biopieces.options_parse(ARGV, casts)
 154
 155 tax_tree = TaxTree.new
 156 tax_hash = {}
 157
 158 Biopieces.open(options[:stream_in], options[:stream_out]) do |input, output|
 159   input.each_record do |record|
 160     if record[:Q_ID] and record[:S_ID] and record[:SCORE]
 161       size = 1
 162
 163       if options[:size]
 164         if record[:Q_ID].match(/_(\d+)$/)
 165           size = $1.to_i
 166         else
 167           raise BiopiecesError, "Could not extract size from Q_ID: #{record[:Q_ID]}"
 168         end
 169       end
 170
 171       tax_hash[record[:Q_ID]] = TaxTree.new unless tax_hash[record[:Q_ID]]
 172       tax_hash[record[:Q_ID]].add_gg(record[:S_ID], record[:SCORE].to_f, size)
 173     else
 174       output.puts record
 175     end
 176   end
 177
 178   if options[:LCA]
 179     tax_hash.each do |q_id, tree|
 180       tree.lowest_common_ancestor
 181     end
 182   end
 183
 184   tax_hash.each do |q_id, tree|
 185     tax_tree.merge(tree)
 186   end
 187
 188   tax_tree.each do |node|
 189     output.puts node.to_bp unless node.level == 'root'
 190   end
 191 end
 192
 193
 194 # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
 195
 196
 197 __END__