def parse_table_to_biom(table_lines, table_format="tab-delimited",\
    biom_format = 'otu table'):
    """Read the lines of an open trait table file, and output a .biom table object

    The trait table must be either a biom file, or a picrust tab-delimited file
    table_format -- must be either 'tab-delimited' or 'biom'

    """
    return parse_table(table_lines) def make_output_dir(dirpath, strict=False):
    """Make an output directory if it doesn't exist

    Returns the path to the directory
    dirpath -- a string describing the path to the directory
    strict -- if True, raise an exception if dir already
    exists
    """
    dirpath = abspath(dirpath)

    #Check if directory already exists
    if isdir(dirpath):
        if strict == True:
            err_str = "Directory '%s' already exists" % dirpath
            raise IOError(err_str)

        return dirpath
    try:
        makedirs(dirpath)
    except IOError,e:
        err_str = "Could not create directory '%s'. Are permissions set correctly? Got error: '%s'" %e
        raise IOError(err_str)

    return dirpath class PicrustNode(PhyloNode):
    def multifurcating(self, num, eps=None, constructor=None):
        """Return a new tree with every node having num or few children

        num : the number of children a node can have max
        eps : default branch length to set if self or constructor is of
            PhyloNode type
        constructor : a TreeNode or subclass constructor. If None, uses self
        """
        if num < 2:
            raise TreeError, "Minimum number of children must be >= 2"

        if eps is None:
            eps = 0.0

        if constructor is None:
            constructor = self.__class__

        if hasattr(constructor, 'Length'):
            set_branchlength = True
        else:
            set_branchlength = False

        new_tree = self.copy()

        for n in new_tree.preorder(include_self=True):
            while len(n.Children) > num:
                new_node = constructor(Children=n.Children[-num:])

                if set_branchlength:
                    new_node.Length = eps

                n.append(new_node)

        return new_tree

    def bifurcating(self, eps=None, constructor=None):
        """Wrap multifurcating with a num of 2"""
        return self.multifurcating(2, eps, constructor)

    def nameUnnamedNodes(self):
        """sets the Data property of unnamed nodes to an arbitrary value

        Internal nodes are often unnamed and so this function assigns a
        value for referencing.
        Note*: This method is faster then pycogent nameUnnamedNodes()
        because it uses a dict instead of an array. Also, we traverse
        only over internal nodes (and not including tips)
        """

        #make a list of the names that are already in the tree
        names_in_use = {}
        for node in self.iterNontips(include_self=True):
            if node.Name:
                names_in_use[node.Name]=1

        #assign unique names to the Data property of nodes where Data = None
        name_index = 1
        for node in self.iterNontips(include_self=True):
            #if (not node.Name) or re.match('edge',node.Name):
            if not node.Name:
                new_name = 'node' + str(name_index)
                #choose a new name if name is already in tree
                while new_name in names_in_use:
                    name_index += 1
                    new_name = 'node' + str(name_index)
                node.Name = new_name
                names_in_use[node.Name]=1
                name_index += 1

    def getSubTree(self,names):
        """return a new subtree with just the tips in names

        assumes names is a set
        assumes all names in names are present as tips in tree
        """
        tcopy = self.deepcopy()

        while len(tcopy.tips()) != len(names):
            # for each tip, remove it if we do not want to keep it
            for n in tcopy.tips():
                if n.Name not in names:
                    n.Parent.removeNode(n)

            # reduce single-child nodes
            tcopy.prune()

        return tcopy def write_biom_table(biom_table, biom_table_fp, compress=True,
                     write_hdf5=HAVE_H5PY, format_fs=None):
    """Writes a BIOM table to the specified filepath

    Parameters
    ----------
    biom_table : biom.Table
        The table object to write out
    biom_table_fp : str
        The path to the output file
    compress : bool, optional
        Defaults to ``True``. If True, built-in compression on the output HDF5
        file will be enabled. This option is only relevant if ``write_hdf5`` is
        ``True``.
    write_hdf5 : bool, optional
        Defaults to ``True`` if H5PY is installed and to ``False`` if H5PY is
        not installed. If ``True`` the output biom table will be written as an
        HDF5 binary file, otherwise it will be a JSON string.
    format_fs : dict, optional
        Formatting functions to be passed to `Table.to_hdf5`

    Notes
    -----
    This code was adapted from QIIME 1.9
    """
    generated_by = "PICRUSt " + __version__

    if write_hdf5:
        with biom_open(biom_table_fp, 'w') as biom_file:
            biom_table.to_hdf5(biom_file, generated_by, compress,
                               format_fs=format_fs)
    else:
        with open(biom_table_fp, 'w') as biom_file:
            biom_table.to_json(generated_by, biom_file) def transpose_trait_table_fields(data_fields,header,id_row_idx=0,\
    input_header_delimiter="\t",output_delimiter="\t"):
    """Transpose the fields of a trait table, returning new data_fields,header

    data_fields:  list of lists for data fields
    header:  a string describing the header_line
    id_row_idx:  index of row labels.  Almost always 0 but included for
    but included for completeness

    input_header_delimiter: delimiter for fields in the header string
    output_delimiter: use this delimiter to join header fields

    NOTE:  typically the header and data fields are generated
    by parse_trait_table in picrust.parse
    """
    header_fields = header.split(input_header_delimiter)

    # ensure no trailing newlines
    old_header_fields = [h.strip() for h in header_fields]
    new_header_fields = [old_header_fields[0]] + \
                        [df[id_row_idx].strip() for df in data_fields]

    non_label_data_fields = []
    for row in data_fields:
        non_label_fields = [e for i, e in enumerate(row) if i != id_row_idx]
        non_label_data_fields.append(non_label_fields)

    data_array = array(non_label_data_fields)
    new_data_array = data_array.T

    new_rows = []
    for i,row in enumerate(new_data_array):
        label = old_header_fields[i+1]
        # this is i+1 not i because i is the blank/meaningless
        # upper left corner entry.
        new_row = [label] + list(row)
        new_rows.append(new_row)
    new_header = output_delimiter.join(new_header_fields)

    return new_header + "\n", new_rows def parse_trait_table(trait_table_lines,delimiter="\t",has_header=True):
    """Return a header line, and a generator that will yield data fields

    trait_table_lines -- tab-seperated lines that may have newline characters

    if has_header is True, then the first non-blank, non-comment line
    must be a header line of equal length to the number of columns, with
    labels for the contents of each column.

    Comment lines (starting with '#') and blank lines will be ignored,
    and won't show up in the output.
    """
    header_line = None
    if not has_header:
        header_line = ''
    else:
        for i,line in enumerate(trait_table_lines):
            if not line or line.startswith("#"):
                continue
            if i == 0:
                header_line = line
                break
    if header_line is None:
        raise RuntimeError("Could not find header line in input trait table lines.  Was it skipped due to starting with a comment ('#') sign?")
    #Now that we have the header (if present) yield_trait_table_fields
    #can assume no header exists, and just parse data fields
    return header_line, yield_trait_table_fields(trait_table_lines,\
      delimiter=delimiter,has_header=False) def filter_table_by_presence_in_tree(tree,trait_table_fields,name_field_index = 0,delimiter="\t"):
    """yield lines of a trait table lacking organisms missing from the tree"""

    tree_tips = [str(node.Name.strip()) for node in tree.preorder()]
    #print tree_tips
    result_fields = []
    for fields in trait_table_fields:
        curr_name = fields[name_field_index].strip()
        if curr_name not in tree_tips:
            #print curr_name,"could not be found in tree nodes"
            #print curr_name in tree_tips
            #try:
            #    print int(curr_name) in tree_tips
            #except:
            #    pass
            #print curr_name.strip() in tree_tips
            continue
        result_fields.append(fields)
    return result_fields def set_label_conversion_fns(remove_whitespace_from_labels=True,\
  replace_problematic_label_characters=True,verbose=False):
    """Return a list of functions for formatting tree node or trait table labels"""
     #Set the functions that will be applied to trait table labels
    label_conversion_fns = []
    if remove_whitespace_from_labels:
        if verbose:
            print "Removing whitespace from trait table organism labels..."
        label_conversion_fns.append(remove_spaces)


    if replace_problematic_label_characters:
        #  Replace ambiguous characters with
        replacement_dict ={":":"_",";":"_"}
        if verbose:
            print "Replacing problematic labels in organism labels:"
            for k,v in replacement_dict.items():
                print k,'-->',v

        chars_to_delete = """'"'"""
        replace_problematic_chars_fn =\
           make_char_translation_fn(replacement_dict,chars_to_delete)
        label_conversion_fns.append(replace_problematic_chars_fn)
    return label_conversion_fns def fix_tree_labels(tree,label_conversion_fns,verbose=False):
    """Fix tree labels by removing problematic characters"""
    if verbose:
        print "reformatting tree node names..."
    tree = format_tree_node_names(tree,label_conversion_fns)
    #print "Number of tree tips with single quotes:",len([t.Name for t in tree if "'" in t.Name])
    return tree def convert_trait_table_entries(trait_table_fields,\
        label_conversion_fns=[str],value_conversion_fns = [float]):
    """Convert trait values by running conversion_fns on labels and values

    trait_table_fields -- list of strings (from a trait table line)
      the first field is assumed to be an organism name, and so isn't
      formatted.

    label_conversion_fns -- a list of functions to be run on each
    organism name label (in the order they should be run).  Each
    function should need only a single entry as input, and output
    the resulting label

    value_conversion_fns -- another list of functions, but for
    trait values.  Again these will be run in order on each table
    value.


    """
    name_field_index = 0
    #print "Value conversion fns:",[f.__name__ for f in value_conversion_fns]
    #print "label_conversion_fns:",[f.__name__ for f in label_conversion_fns]
    for fields in trait_table_fields:
        new_fields = []
        for i,field in enumerate(fields):
            if i != name_field_index:
                converters_to_use = value_conversion_fns
            else:
                converters_to_use = label_conversion_fns

            #Run appropriate converters on this field
            new_val = field
            for curr_conv_fn in converters_to_use:
                new_val = str(curr_conv_fn(new_val))
            new_fields.append(new_val)

        yield new_fields def yield_test_trees(tree,tree_modifier_fn = exclude_tip):
    """yield successive test trees from the full tree

    tree -- a PyCogent PhyloNode object for the full tree

    tree_modifier_fn -- a function that takes a tip and a tree and returns
    a new, test tree (for example a 'leave-one-out' tree)

    """


    for tip in tree.iterTips():
        # Need to pass name only to avoid
        # modifying the original tree
        tip_name = tip.Name
        tree_copy = tree.deepcopy()
        test_tree = \
          tree_modifier_fn(tree_copy.getNodeMatchingName(tip_name),tree_copy)
        yield tip, test_tree def make_distance_based_tip_label_randomizer(randomization_distance):
    """Returns a function that randomizes the labels of all nodes within d of a given node.
    randomization_distance -- the units of branchlength from which neighbors should be excluded
    """

    def randomize_tip_and_neighbors(tip,tree):

        tips_to_randomize = tip.tipsWithinDistance(randomization_distance)
        #print "tips_to_randomize (from tipsWithinDistance):",[t.Name for t in tips_to_randomize]
        if tip not in tips_to_randomize:
            tips_to_randomize.append(tip)
        #print "tips_to_randomize:",[t.Name for t in tips_to_randomize]
        # Resample tips without replacement to get random labels
        new_tip_labels = [tip.Name for tip in tips_to_randomize]
        shuffle(new_tip_labels)

        for i,tip in enumerate(tips_to_randomize):
            tip.Name = new_tip_labels[i]

        return tree

    new_fn = randomize_tip_and_neighbors

    #Document the new function
    new_fn_doc = "Randomize tip labels within %f of tip of interest" % randomization_distance
    new_fn.__doc__ = new_fn_doc

    return new_fn def make_distance_based_exclusion_fn(exclusion_distance):
    """Returns a function that the neighbors of a tip within d.
    exclusion_distance -- the units of branchlength from which neighbors should be excluded

    """


    def exclude_tip_neighbors(tip,tree, verbose=False):
        # check that the exclusion distance won't exclude the entire tree
        max_dist = tree.maxTipTipDistance()[0]
        tips_to_delete = tip.tipsWithinDistance(exclusion_distance)
        tips_to_delete = [t.Name for t in tips_to_delete]
        if verbose:
            print "Tips to delete:",len(tips_to_delete)
        tips_to_keep = [t.Name for t in tree.tips() if t.Name not in tips_to_delete]

        if float(exclusion_distance) > float(max_dist) or \
          len(tips_to_delete) == len(tree.tips()):

            raise ValueError("specified tree would be entirely excluded because branch lengths are too short")


        #holdout_tip = tip
        holdout_tip = tip.Name
        #We're working with the *pruned* tree so we want the holdout tip
        #removed
        if holdout_tip in tips_to_keep:
            tips_to_keep.remove(holdout_tip)
            if verbose:
                print "Will remove our tip of interest from the tree: {0}".format(holdout_tip)
        if verbose:
            print "Tips to keep:", len(tips_to_keep)

        subtree = get_sub_tree(tree,tips_to_keep)
        del tree
        # Then it's neighbors
        #for tip in tips_to_delete:
        #    if tip == holdout_tip:
        #        #Don't exclude the tip you're testing
        #        continue
        #    tree = exclude_tip(tip,tree)
        return subtree

    new_fn = exclude_tip_neighbors

    #Document the new function
    new_fn_doc =\
      "Exclude neighbors of tip within %f branch length units" \
      % exclusion_distance
    new_fn.__doc__ = new_fn_doc

    return new_fn def exclude_tip(tip, tree,verbose=False):
    """Return pruned tree with tip excluded

    tip --  a PyCogent PhyloNode object from the tree
    tree -- a PyCogent PhyloNode object (the tree)

    NOTE: This function exists to support distance-based exclusion of
    neightbors.  For holdout analysis we really want to
    exclude the organism in the trait table rather than the tree.
    """

    # Make sure the tip is really a tip
    if not tip.isTip():
        raise ValueError("The 'holdout' fn got a non-tip node for it's 'tip' parameter.")

    #Old exclusion based method...but for large trees this is quite slow

    # Assuming that's OK, remove the tip
    #if tip.Parent is not None:
    #    removal_ok = tip.Parent.remove(tip)
    #else:
    #    removal_ok = False
    #if not removal_ok:
    #    raise ValueError("holdout could not remove node '%s' from tree." %(tip.Name))
    #tree.prune()

    #Newer subtree based method
    tips_to_keep = [t.Name for t in tree.tips() if t.Name != tip.Name]
    if verbose:
        print "Generating subtree with %i tips" %len(tips_to_keep)
    subtree = get_sub_tree(tree,tips_to_keep)

    return subtree def write_tree(output_dir,base_name,test_tree,tip_to_predict,delimiter='--' ):
    """Write a test trees to files with descriptive filenames"""
    file_base = delimiter.join([base_name,tip_to_predict])
    filename=os.path.join(output_dir,file_base)
    f=open(filename,"w")
    f.write(test_tree.getNewick(with_distances=True))
    f.close()
    return filename def yield_genome_test_data_by_distance(tree,trait_table_fields,\
  test_fn_factory=make_distance_based_exclusion_fn, min_dist=0.0,max_dist=0.70,\
  increment=0.03,modify_tree=True,limit_to_tips=False,verbose=True):
    """Yield successive test datasets as analysis_name,test_tree,test_trait_table, test_tip tuples

    tree  -- a PyCogent PhyloNode object

    trait_table_fields -- a list of strings from a trait table
      (the first is organism name, the rest are trait values)

    test_fn_factory -- a factory function that returns a new test tree generating function,
      given a distance parameter.  The test tree generating function, in turn, should
      return a new test tree, given a tip of interest and the full tree.

    min_dist -- the minimum distance at which to start generating test
      trees.

    max_dist -- the maximum distance for test trees

    increment -- the increment by which the distance parameter should be
      increased.

    modify_tree -- if set to False, just return the original tree without modification.

    limit_to_tips -- takes a list of organisms(tree tips) to include in the analysis.  If this is provided, instead of making test datasets
    for all tips, only generate test datasets for tree tips with names exactly matching an entry in the supplied list of tips.

    For example, to limit the datasets to just E.coli K-12 MG1655 (when using the IMG / greengenes dataset)
    pass limit_to_included_tips=['NC_000913|646311926']
    """

    if verbose:
    #Include only organisms present in the tree
        print "Prefiltering orgs in trait table:",len([f[0] for f in trait_table_fields])
        print "Orgs with single quotes:",len([f[0] for f in trait_table_fields if "'" in f[0]])
        print "Tree tips with single quotes:",len([t.Name for t in tree.iterTips() if "'" in t.Name])

    trait_table_fields = filter_table_by_presence_in_tree(tree,\
                  trait_table_fields)

    orgs_in_table = [f[0] for f in trait_table_fields]
    if verbose:
        print "number of organisms matching between tree and table:",len(orgs_in_table)
        #print "Filtered trait table organisms:", orgs_in_table
        #print "Tree tips:",[t.Name for t in tree.iterTips()]

    if not orgs_in_table:
        raise RuntimeError("Could not match trait table organisms to tree")

    table_by_org = trait_dict_from_fields(trait_table_fields)

    #Generate a test data generation function
    #using the specified method and distance parameter

    dists = [min_dist+i*increment for i in range(int((max_dist-min_dist) // increment))]

    if verbose:
        print "Building test trees for each of the following distances: %s" \
        %(",".join(map(str,dists)))


    for curr_dist in dists:
        if verbose:
            print "Generating test datset for distance: %f" % curr_dist

        test_fn = test_fn_factory(curr_dist)


        #Select which tips will be processed

        if limit_to_tips is False:
            tips_to_predict = orgs_in_table
            if verbose:
                print "Generating test data for all tips"
        else:
            tips_to_predict = limit_to_tips
            if verbose:
                print "Generating test data for tips:",tips_to_predict
            tips_to_predict = limit_to_tips

        #For each tip of interest, generate a test dataset
        test_data = []
        total_organisms = len(tips_to_predict)
        for i,tip_to_predict in enumerate(tips_to_predict):

            if verbose:
                print "Generating test dataset (%i/%i) for: %s" %(i+1,\
                  total_organisms,tip_to_predict)

            if modify_tree:

                if verbose:
                    print "Copying tree..."

                tree_copy = tree.deepcopy()

                if verbose:
                    print "Locating tip of interest in tree..."

                tip_of_interest = tree_copy.getNodeMatchingName(tip_to_predict)

                if verbose:
                    print "Modifying tree using function:",test_fn.__name__,\
                      test_fn.__doc__

                test_tree = \
                  test_fn(tip_of_interest,tree_copy)
            else:
                test_tree = tree

            if verbose:
                print "Setting expected traits"
            expected_traits = table_by_org[tip_to_predict]

            if verbose:
                print "Generating test trait table"
            test_trait_table_fields = filter_table_by_presence_in_tree(test_tree,\
                  trait_table_fields)

            yield curr_dist,test_tree,tip_to_predict, expected_traits, test_trait_table_fields
import os
from copy import deepcopy
from cogent.parse.tree import DndParser
from cogent.util.option_parsing import parse_command_line_parameters,\
  make_option
from biom.table import Table
from picrust.util import parse_table_to_biom, make_output_dir,\
     PicrustNode, write_biom_table,transpose_trait_table_fields
from picrust.parse import parse_trait_table
from picrust.format_tree_and_trait_table import filter_table_by_presence_in_tree,\
  set_label_conversion_fns, fix_tree_labels, convert_trait_table_entries
from numpy import array
from picrust.make_test_datasets import yield_test_trees,\
  make_distance_based_tip_label_randomizer,make_distance_based_exclusion_fn,\
  exclude_tip,write_tree, yield_genome_test_data_by_distance
#!/usr/bin/env python
# File created on 22 Feb 2012
from __future__ import division

__author__ = "Jesse Zaneveld"
__copyright__ = "Copyright 2011-2013, The PICRUSt Project"
__credits__ = ["Jesse Zaneveld"]
__license__ = "GPL"
__version__ = "1.1.4"
__maintainer__ = "Jesse Zaneveld"
__email__ = "zaneveld@gmail.com"
__status__ = "Development"


script_info = {}
script_info['brief_description'] = "Generates test datasets for cross-validation studies of PICRUSt's accuracy"
script_info['script_description'] = ""
script_info['script_usage'] = [("","Generate holdout test trees from genome_tree.newick, and save results in the directory ./test_holdout_trees/.","%prog -t genome_tree.newick -o ./test_holdout_trees")]
script_info['output_description']= ""
method_choices = ['exclude_tips_by_distance','randomize_tip_labels_by_distance','collapse_tree_by_distance']
script_info['required_options'] = [

 make_option('-i','--input_trait_table',type='existing_filepath',\
   help='the input trait table.'),\
 make_option('-t','--input_tree',type='existing_filepath',\
   help='the input tree in Newick format'),\
]
script_info['optional_options'] = [\
  make_option('-o','--output_dir',default='./test_datasets/',type='new_dirpath',\
  help='the output directory.  Duplicate trees, trait tables, expected values and prediction files will be saved here.[default:%default]'),\
  make_option('--min_dist',default=0.0,type='float',\
  help='the minimum phylogenetic distance to use with the holdout method, if applicable.  Usually 0.0.[default:%default]'),\
   make_option('--suppress_tree_modification',default=False,action="store_true",help='If passed, modify only the trait table, not the tree . [default: %default]'),\
  make_option('--dist_increment',default=0.03,type='float',\
  help='the phylogenetic distance increment to use with the holdout method, if applicable.[default:%default]'),\
  make_option('--max_dist',default=0.45,type='float',\
  help='the maximum phylogenetic distance to use with the holdout method, if applicable.[default:%default]'),\
  make_option('--limit_to_tips',default='',type='string',\
  help='if specified, limit test dataset generation to specified tips (comma-separated).[default:%default]'),\
  make_option('-m','--method',type='choice',\
    choices=method_choices,default=method_choices[0],\
    help='The test method to use in generating test data.  Valid choices are:'\
      +','.join(method_choices)+' [default: %default]'),\

]

script_info['version'] = __version__



def main():
    """Generate test trees given parameters"""
    option_parser, opts, args =\
       parse_command_line_parameters(**script_info)

    if opts.verbose:
        print "Loading trait table..."
    input_trait_table = open(opts.input_trait_table,"U")

    if opts.verbose:
        print "Loading tree..."
    #PicrustNode seems to run into very slow/memory intentsive perfromance...
    #tree = DndParser(open(opts.input_tree),constructor=PicrustNode)
    tree = DndParser(open(opts.input_tree))

    if opts.verbose:
        print "Parsing trait table..."
    #Find which taxa are to be used in tests
    #(by default trait table taxa)
    trait_table_header,trait_table_fields = \
            parse_trait_table(input_trait_table)

    if opts.verbose:
       print "Ensuring tree and trait table labels are formatted consistently..."

    label_conversion_fns = set_label_conversion_fns(verbose=opts.verbose)

    fix_tree_labels(tree,label_conversion_fns)

    trait_table_fields = convert_trait_table_entries(trait_table_fields,\
      value_conversion_fns = [],\
      label_conversion_fns = label_conversion_fns)

    trait_table_fields = [t for t in trait_table_fields]
    print "Number of trait table fields with single quotes:",\
     len([t for t in trait_table_fields if "'" in t[0]])

    if opts.verbose:
        print "Making output directory..."
    make_output_dir(opts.output_dir)


    if opts.limit_to_tips:

        included_tips = opts.limit_to_tips.split(",")
        if opts.verbose:
            print "Limiting test datasets to %i tips: %s" %(len(included_tips),included_tips)
    else:
        included_tips = False

    method_fns =\
      {"exclude_tips_by_distance":\
         make_distance_based_exclusion_fn,\
       "randomize_tip_labels_by_distance":\

         make_distance_based_tip_label_randomizer