Coloring Nodes on a Phylogeny

colornodes.rb

#An automation of a tedious task I have to do often: coloring the nodes on a phylogeny. This script takes a dendroscope tree file and a "color description" file, a simple csv file with taxa labels and a corresponding color. The color may either be an RGB triplet or a scalar value which will be mapped to a pallete. Usage is:

#color-dendro.rb [options] CLRFILE TREEFILE1 [...]

#where the options are:

#-h, --help Display this screen 
#-m, --default-color STR The default color nodes will be given 
#--map-to-colors The coloring instructions give a float value which will be mapped to a color 
#--save STR 
#Name output files according this template 
#-o, --overwrite Overwrite pre-existing files


#!/usr/bin/env ruby
# Color the node labels in a Dendroscope tree.

### IMPORTS

require 'test/unit/assertions'
require 'optparse'
require 'pp'
require 'csv'
require 'ostruct'
require 'date'
require 'time'

include Test::Unit::Assertions


### CONSTANTS & DEFINES


### IMPLEMENTATION

def interpolate(str, sub_hash)
   return str.gsub(/\{([^}]+)\}/) { |m|
      sub_hash[$1]
   }
end


### MAIN


class ColorMapper
   attr_accessor(:min, :max, :range)
   
   def initialize(vals)
      @max, @min = vals.max.to_f(), vals.min.to_f()
      @range = @max - @min
   end
   
   def map(val)
      norm_val = normalise_val(val)
      val255 = (norm_val * 255).round()
      clrs = [val255, 255 - val255, val255].each { |v|
         v.to_i.to_s
      }
      return clrs.join(' ')
   end
   
   def normalise_val(val)
      # get a val from 0 to 1
      return (val - @min) / @range
   end
end


# Return a mapping between label names and colors
#
def parse_clr_file(clr_file, map_to_clr)
   clr_map = {}
   
   # read file
   CSV::Reader.parse(File.open(clr_file, 'rb')) { |row|
      #assert_equal(2, row.length, "row length should be 2")
      label, color = row[0].strip(), row[1].strip()
      if map_to_clr == false
         assert_equal(3, color.split(' ').length(),
            "color '#{color}' looks malformed")
      end
      clr_map[label] = color
   }
   
   # map to colors if need be
   if (map_to_clr)
      clr_map.each_pair { |k,v|
         clr_map[k] = v.to_f()
      }
      mapper = ColorMapper.new(clr_map.values)
      clr_map.each_pair { |k,v|
         clr_map[k] = mapper.map(v)
      }
   end
   
   return clr_map
end


# Parse commandline arguments.
#
def parse_clargs(arg_arr)
   clopts = {
      :save => "{root}-colored.{ext}",
      :overwrite => false,
      :def_clr => "0 0 0",
      :map_to_clr => false,
   }
   
   OptionParser.new { |opts|
      opts.program_name = __FILE__
      opts.banner = "Color the nodes in a dendroscope file"
      opts.separator("")
      opts.separator("The input is a CSV file of coloring instructions and one")
      opts.separator("or more Dendroscope tree files. The coloring instructions")
      opts.separator("consist of one row for each node to be labelled,")
      opts.separator("consisting of:")
      opts.separator("")
      opts.separator("   <label>, <fgcolor>")
      opts.separator("")
      opts.separator("where the color is a RGB 256 triplet (e.g. '0 128 255').")
      opts.separator("")
      
      opts.separator("Usage: #{opts.program_name} [options] CLRFILE TREEFILE1 ...]")
      
      opts.on('-h', '--help', 'Display this screen') {
         puts opts
         exit
      }
      opts.on('-m', '--default-color STR',
         "The default color nodes will be given") { |v|
         clopts[:def_clr] = v
      }
      opts.on('-m', '--map-to-colors',
         "The coloring instructions give a float value which will be mapped to a color") {
         clopts[:map_to_clr] = true
      }
      opts.on('', '--save STR',
         "Name output files according this template") { |v|
         clopts[:save] = v
      }
      opts.on('-o', '--overwrite',
         "Overwrite pre-existing files") {
         clopts[:overwrite] = true
      }
      
      begin
         opts.parse!(arg_arr)
      rescue OptionParser::InvalidOption => e
         puts e
         puts opts
         exit 1
      end
   }
   
   pargs = arg_arr
   assert(2 <= pargs.length, "need files to work on")
 
   return clopts, pargs[0], pargs[1, pargs.length]
end


# Main script functionality.
#
def main()
   clopts, clr_file, tree_files = parse_clargs(ARGV)
   
   clr_table = parse_clr_file(clr_file, clopts[:map_to_clr])
   
   tree_files.each { |f|
      # slurp ...
      puts "== Reading '#{f}' ..."
      converted_lines = []
      tree_data = File.open(f, 'rb').read()
      
      # give each node a color, whether it wants it or not
      tree_data.gsub!(Regexp.new('^(\d+: .*) lb=(.+)$')) { |m|
         match_str = m.to_s()
         if (match_str.include?(" lc="))
            # color already defined
            match_str
         else
            "#{$1} lc=#{clopts[:def_clr]} lb=#{$2}"
         end
      }
      
      # add color labels
      clr_table.each_pair { |label, color|
         pp color
         tree_data.gsub!(
            Regexp.new("lc=(.*) (lb='#{Regexp.escape(label)}')"),
            "lc=#{color} " + '\2'
         )
      }
      
      # write output
      
      # make filename
      ext = File.extname(f)
      subs = {
         "ext" => ext[1, ext.length],
         "base" => File.basename(f),
         "root" => File.basename(f, ext),
         "date" => Date.today.to_s(),
         "time" => Time.now.strftime(fmt='%T'),
         "datetime" => DateTime.now.strftime(fmt='%F T%T'),
      }
      out_name = interpolate(clopts[:save], subs)
      puts "Saving drawing to '#{out_name}' ..."
      
      # do the writing
      if File.exists?(out_name)
         assert(clopts[:overwrite], "Can't overwrite existing file '#{out_name}'")
      end
      File.open(out_name, 'w').write(tree_data)
      
      #drwr.save(out_name)
      puts "Saved."
      
   }
   
   puts "== Finished."
end

if $0 == __FILE__
   main()
end


### END