#! /usr/bin/env python # -*- coding: utf-8 -*- ############################################################################## ## DendroPy Phylogenetic Computing Library. ## ## Copyright 2010-2015 Jeet Sukumaran and Mark T. Holder. ## All rights reserved. ## ## See "LICENSE.rst" for terms and conditions of usage. ## ## If you use this work or any portion thereof in published work, ## please cite it as: ## ## Sukumaran, J. and M. T. Holder. 2010. DendroPy: a Python library ## for phylogenetic computing. Bioinformatics 26: 1569-1571. ## ############################################################################## """ Writing data in NEXUS format to an output stream. """ import re import warnings import collections from dendropy.utility import textprocessing from dendropy.dataio import ioservice from dendropy.dataio import newick from dendropy.dataio import nexusprocessing from dendropy.dataio import newickwriter ############################################################################### ## NexusWriter class NexusWriter(ioservice.DataWriter): """ Formatter for NEXUS data. """ def __init__(self, **kwargs): """ Keyword Arguments ----------------- simple : boolean, default: |False| If |True|, write in simple NEXUS format, i.e. in a single "DATA" block, instead of separate "TAXA" and "CHARACTER" blocks. suppress_taxa_blocks: boolean, default: |False| If |True|, do not write a "TAXA" block. Note that this may make the file impossible to parse if there are multiple taxon namespaces in the data. suppress_unreferenced_taxon_namespaces: boolean, default: |False| If |True|, then when writing |DataSet| objects, any |TaxonNamespace| object in the DataSet's ``taxon_namespaces`` collection will *not* be written as a "TAXA" block if it is not referenced by any character matrix (``char_matrices``) or tree list (``tree_lists``). suppress_block_titles : bool or |None| If |True| then 'TITLE' element to blocks will not be written. Note that this may make the file impossible to parse if there are multiple taxon namespaces in the data. If |False|, then the 'TITLE' element will always be written. Default is |None|: the 'TITLE' element will only be written if needed because there is more than on taxon namespace in the data. file_comments: iterable [``str``] List of lines of text to be added as comments to the file. preamble_blocks: iterable [``str``] List of strings to be written before data (e.g., PAUP blocks suppressing warnings etc.). supplemental_blocks: iterable [``str``] List of strings to be written after data (e.g., PAUP blocks, MrBayes blocks etc.). allow_multiline_comments : bool If |False| then comments will be merged into a single string before being written. Default is |True|: each comment element will be written on its own line. continuous_character_state_value_format_fn : function object When writing |ContinuousCharacterMatrix| data: a function that takes a continuous character value and returns the string representation of it. discrete_character_state_value_format_fn : function object When writing discrete character data (e.g., a |StandardCharacterMatrix|): a function that takes a standard character state value (i.e., a |StateIdentity| instance) and returns the string representation of it. suppress_leaf_taxon_labels : boolean, default: |False| If |True|, then taxon labels will not be rendered for leaves. Default is |False|: render leaf taxon labels. See notes below for details. suppress_leaf_node_labels : boolean, default: |True| If |False|, then node labels (if available) will be printed for leaves. Defaults to |True|: do not render leaf node labels. See notes below for details. suppress_internal_taxon_labels : boolean, default: |False| If |True|, then taxon labels will not be printed for internal nodes. Default is |False|: print taxon labels for internal nodes. See notes below for details. suppress_internal_node_labels : boolean, default: |False| If |True|, then node labels will not be printed for internal nodes. Default is |False|: print node labels for internal nodes. See notes below for details. suppress_rooting : boolean, default: |False| If |True|, will not write rooting token ('[&R]' or '[&U]'). Default is |False|: rooting token will be written. suppress_edge_lengths : boolean, default: |False| If |True|, will not write edge lengths. Default is |False|: edge lengths will be written. unquoted_underscores : boolean, default: |False| If |True|, labels with underscores will not be quoted, which will mean that they will be interpreted as spaces if read again ("soft" underscores). If |False|, then labels with underscores will be quoted, resulting in "hard" underscores. Default is |False|. preserve_spaces : boolean, default: |False| If |True|, spaces will not be replaced with underscores in labels (which means any labels containing spaces will have to be quoted). Default is |False|: spaces will be converted to underscores. False. store_tree_weights : boolean, default: |False| If |True|, tree weights are written. Default is |False|: tree weights will not be written. translate_tree_taxa : boolean or dict or |None|, default: |None|. If |False| or |None|, then a "TRANSLATE" statement will not be used, and tree statements will contain the full taxon labels. If not |False| or |None|, a "TRANSLATE" statement will be written and referenced in tree statements (instead of using the taxon labels). If |True|, then a default translate statement will be used, with tokens given by the taxon indexes. If a dictionary is given, then the keys should be |Taxon| objects and the values should be the token (strings). suppress_annotations : boolean, default: |False| If |True|, metadata annotations will be ignored. Defaults to |False|: metadata annotations will be written. annotations_as_nhx : boolean, default: |False| If |True|, and if ``suppress_annotations`` is |False|, will write annotations as NHX statements. Default is |False|: annotations will not be written as NHX statements. suppress_item_comments : boolean, default: |False| If |True|: comments will be ignored. Default is |False|: any additional comments associated with trees, nodes, edges, etc. will be written. node_label_element_separator : string, default: ' ' If both ``suppress_leaf_taxon_labels`` and ``suppress_leaf_node_labels`` are |False|, then this will be the string used to join them. Defaults to ' ' (space). node_label_compose_fn : function object or |None|, default: |None| If not |None|, should be a function that takes a |Node| object as an argument and returns the string to be used to represent the node in the tree statement. The return value from this function is used unconditionally to print a node representation in a tree statement, by-passing the default labelling function, ignoring ``suppress_leaf_taxon_labels``, ``suppress_leaf_node_labels=True``, ``suppress_internal_taxon_labels``, ``suppress_internal_node_labels``, etc. Defaults to |None|. edge_label_compose_fn : function object or |None|, default: |None| If not |None|, should be a function that takes an Edge object as an argument, and returns the string to be used to represent the edge length in the tree statement. real_value_format_specifier : string, default: '' Format specification for real/float values. Will be applied to edge lengths (if ``edge_label_compose_fn`` is not given) as well as annotations. The format specifier should be given in Python's string format specification mini-language. E.g. ".8f", ".4E", "8.4f". exclude_from_taxa_blocks : set (or other iterable) of |Taxon| instances If specified, any |Taxon| object in this set will be excluded from all TAXA blocks, *even* if they are referenced in the data. ignore_unrecognized_keyword_arguments : boolean, default: |False| If |True|, then unsupported or unrecognized keyword arguments will not result in an error. Default is |False|: unsupported keyword arguments will result in an error. """ # base ioservice.DataWriter.__init__(self) # Following are NEXUS specific (i.e., not used by NEWICK formatters), # and need to be removed so as not to cause problems with our keyword # validation scheme self.simple = kwargs.pop("simple", False) self.suppress_taxa_blocks = kwargs.pop("suppress_taxa_blocks", None) self.suppress_block_titles = kwargs.pop("suppress_block_titles", None) self.file_comments = kwargs.pop("file_comments", []) if self.file_comments is None: self.file_comments = [] self.preamble_blocks = kwargs.pop("preamble_blocks", []) if self.preamble_blocks is None: self.preamble_blocks = [] self.supplemental_blocks = kwargs.pop("supplemental_blocks", []) if self.supplemental_blocks is None: self.supplemental_blocks = [] self.allow_multiline_comments = kwargs.pop("allow_multiline_comments", True) self.suppress_unreferenced_taxon_namespaces = kwargs.pop("suppress_unreferenced_taxon_namespaces", False) self.continuous_character_state_value_format_fn = kwargs.pop("continuous_character_state_value_format_fn", self._format_continuous_character_value) if self.continuous_character_state_value_format_fn is None: self.continuous_character_state_value_format_fn = self._format_continuous_character_value self.discrete_character_state_value_format_fn = kwargs.pop("discrete_character_state_value_format_fn", self._format_discrete_character_value) if self.discrete_character_state_value_format_fn is None: self.discrete_character_state_value_format_fn = self._format_discrete_character_value self.translate_tree_taxa = kwargs.pop("translate_tree_taxa", None) self.exclude_from_taxa_blocks = kwargs.pop("exclude_from_taxa_blocks", None) # The following are used by NewickWriter in addition to NexusWriter, so # they are extracted/set here and then forwarded on ... self.unquoted_underscores = kwargs.get('unquoted_underscores', False) self.preserve_spaces = kwargs.get("preserve_spaces", False) self.annotations_as_nhx = kwargs.get("annotations_as_nhx", False) self.real_value_format_specifier = kwargs.pop("real_value_format_specifier", "") # As above, but the NEXUS format default is different from the NEWICK # default, so this rather convoluted approach self.suppress_annotations = kwargs.pop("suppress_annotations", False) kwargs["suppress_annotations"] = self.suppress_annotations self.suppress_item_comments = kwargs.pop("suppress_item_comments", False) kwargs["suppress_item_comments"] = self.suppress_item_comments # The newick writer to which tree-writing will be delegated self._newick_writer = newickwriter.NewickWriter(**kwargs) # Book-keeping self.taxon_namespaces_to_write = [] self._block_title_map = {} self._title_block_map = {} def _write(self, stream, taxon_namespaces=None, tree_lists=None, char_matrices=None, global_annotations_target=None): # Header stream.write('#NEXUS\n\n') # File/Document-level annotations and comments if self.file_comments: self._write_comments(stream, self.file_comments) if global_annotations_target is not None: self._write_item_annotations(stream, global_annotations_target) self._write_item_comments(stream, global_annotations_target) # Other blocks if self.preamble_blocks: for block in self.preamble_blocks: stream.write(block) stream.write("\n") stream.write("\n") # Taxon namespace discovery candidate_taxon_namespaces = collections.OrderedDict() if self.attached_taxon_namespace is not None: # should this be False? candidate_taxon_namespaces[self.attached_taxon_namespace] = True elif taxon_namespaces is not None: if self.suppress_unreferenced_taxon_namespaces: # preload to preserve order for tns in taxon_namespaces: candidate_taxon_namespaces[tns] = False else: for tns in taxon_namespaces: candidate_taxon_namespaces[tns] = True for data_collection in (tree_lists, char_matrices): if data_collection is not None: for i in data_collection: if self.attached_taxon_namespace is None or i.taxon_namespace is self.attached_taxon_namespace: candidate_taxon_namespaces[i.taxon_namespace] = True self.taxon_namespaces_to_write = [tns for tns in candidate_taxon_namespaces if candidate_taxon_namespaces[tns]] # Write out taxon namespaces if not self.simple and not self.suppress_taxa_blocks: if self.suppress_block_titles and len(taxon_namespace_to_write) > 1: warnings.warn("Multiple taxon namespaces will be written, but block titles are suppressed: data file may not be interpretable") for tns in self.taxon_namespaces_to_write: self._write_taxa_block(stream, tns) # Write out character matrices if char_matrices is not None: for char_matrix in char_matrices: if (self.attached_taxon_namespace is None or char_matrix.taxon_namespace is self.attached_taxon_namespace): self._write_char_block(stream=stream, char_matrix=char_matrix) # Write out tree lists if tree_lists is not None: for tree_list in tree_lists: if (self.attached_taxon_namespace is None or tree_list.taxon_namespace is self.attached_taxon_namespace): self._write_trees_block(stream=stream, tree_list=tree_list) # Write out remaining if self.supplemental_blocks: for block in self.supplemental_blocks: stream.write(block) stream.write("\n") def _get_taxa_to_include(self, taxon_namespace): if not self.exclude_from_taxa_blocks: return list(taxon_namespace) else: return [t for t in taxon_namespace if t not in self.exclude_from_taxa_blocks] def _write_taxa_block(self, stream, taxon_namespace): stream.write("BEGIN TAXA;\n") self._write_block_title(stream, taxon_namespace) self._write_item_annotations(stream, taxon_namespace) self._write_item_comments(stream, taxon_namespace) taxon_to_include = self._get_taxa_to_include(taxon_namespace) stream.write(" DIMENSIONS NTAX={};\n".format(len(taxon_to_include))) stream.write(" TAXLABELS\n") for taxon in taxon_to_include: stream.write(" {}\n".format( nexusprocessing.escape_nexus_token(taxon.label, preserve_spaces=self.preserve_spaces, quote_underscores=not self.unquoted_underscores), )) self._write_item_annotations(stream, taxon) self._write_item_comments(stream, taxon) stream.write(" ;\n") stream.write("END;\n\n") def _set_and_write_translate_block(self, stream, taxon_namespace): if not self.translate_tree_taxa: self._newick_writer.taxon_token_map = None return if self.translate_tree_taxa is True: m = {} for taxon_idx, t in enumerate(taxon_namespace): m[t] = str(taxon_namespace.accession_index(t)+1) # m[t] = str(taxon_idx) self._newick_writer.taxon_token_map = m else: self._newick_writer.taxon_token_map = dict(self.translate_tree_taxa) stream.write(" Translate\n") statement = [] for taxon in taxon_namespace: label = nexusprocessing.escape_nexus_token(str(taxon.label), preserve_spaces=self.preserve_spaces, quote_underscores=not self.unquoted_underscores) statement.append(" {} {}".format(self._newick_writer.taxon_token_map[taxon], label)) statement = ",\n".join(statement) stream.write("{}\n ;\n".format(statement)) def _write_trees_block(self, stream, tree_list): stream.write("BEGIN TREES;\n") self._write_block_title(stream, tree_list) self._write_item_annotations(stream, tree_list) self._write_item_comments(stream, tree_list) self._write_link_to_taxa_block(stream, tree_list.taxon_namespace) self._set_and_write_translate_block(stream, tree_list.taxon_namespace) for tree_idx, tree in enumerate(tree_list): if tree.label: tree_name = tree.label else: tree_name = str(tree_idx+1) tree_name = nexusprocessing.escape_nexus_token( tree_name, preserve_spaces=self.preserve_spaces, quote_underscores=not self.unquoted_underscores) stream.write(" TREE {} = ".format(tree_name)) self._newick_writer._write_tree(stream, tree) stream.write("\n") stream.write("END;\n\n") def _write_char_block(self, stream, char_matrix): taxon_label_map = collections.OrderedDict() for taxon in char_matrix: taxon_label_map[taxon] = nexusprocessing.escape_nexus_token(taxon.label, preserve_spaces=self.preserve_spaces, quote_underscores=not self.unquoted_underscores) nchar = max([len(seq) for seq in char_matrix.values()]) if self.simple: stream.write("BEGIN DATA;\n") # note that this will only list the number of taxa for # which sequences are available ntaxstr = " NTAX={}".format(len(taxon_label_map)) else: stream.write("BEGIN CHARACTERS;\n") ntaxstr = "" self._write_block_title(stream, char_matrix) self._write_item_annotations(stream, char_matrix) self._write_item_comments(stream, char_matrix) self._write_link_to_taxa_block(stream, char_matrix.taxon_namespace) stream.write(" DIMENSIONS{} NCHAR={};\n".format(ntaxstr, nchar)) stream.write(" FORMAT {};\n".format(self._compose_format_terms(char_matrix))) stream.write(" MATRIX\n") if char_matrix.data_type == "continuous": state_value_writer = lambda x : stream.write("{} ".format(self.continuous_character_state_value_format_fn(x))) else: state_value_writer = lambda x : stream.write("{}".format(self.discrete_character_state_value_format_fn(x))) max_label_len = max(len(v) for v in taxon_label_map.values()) for taxon in char_matrix: stream.write(" {taxon_label:{field_len}} ".format(taxon_label=taxon_label_map[taxon], field_len=max_label_len)) for state in char_matrix[taxon]: state_value_writer(state) stream.write("\n") stream.write(" ;\n") stream.write("END;\n\n\n") self._write_character_subsets(stream, char_matrix) def _compose_format_terms(self, char_matrix): format = [] if char_matrix.data_type == "dna": format.append("DATATYPE=DNA") format.append("GAP=- MISSING=? MATCHCHAR=.") elif char_matrix.data_type == "rna": format.append("DATATYPE=RNA") format.append("GAP=- MISSING=? MATCHCHAR=.") elif char_matrix.data_type == "nucleotide": format.append("DATATYPE=NUCLEOTIDE") format.append("GAP=- MISSING=? MATCHCHAR=.") elif char_matrix.data_type == "protein": format.append("DATATYPE=PROTEIN") format.append("GAP=- MISSING=? MATCHCHAR=.") elif char_matrix.data_type == "continuous": format.append("DATATYPE=CONTINUOUS ITEMS=(STATES)") else: format.append("DATATYPE=STANDARD") fundamental_symbols = set() for state_alphabet in char_matrix.state_alphabets: for s in state_alphabet.fundamental_state_iter(): if s.symbol is not None: fundamental_symbols.add(s.symbol) else: raise Exception("Could not match character state to symbol: '%s'." % s) format.append('SYMBOLS="%s"' % "".join(fundamental_symbols)) equates = set() for state_alphabet in char_matrix.state_alphabets: for a in state_alphabet.ambiguous_state_iter(): if a.symbol == "?": format.append("MISSING=?") elif a.symbol == "-": format.append("GAP=-") else: if a.symbol is not None: equates.add("%s={%s}" % (a.symbol, "".join(a.fundamental_symbols))) for state_alphabet in char_matrix.state_alphabets: for p in state_alphabet.polymorphic_state_iter(): if p.symbol is not None: equates.add("%s=(%s)" % (p.symbol, "".join(p.fundamental_symbols))) if equates: format.append('EQUATE="%s"' % equates) return ' '.join(format) def _write_comments(self, stream, comments): if self.allow_multiline_comments: if textprocessing.is_str_type(comments): stream.write("[{}]\n".format(comments)) else: comments = "\n".join([str(c) for c in comments]) stream.write("[\n{}\n]\n".format(comments)) else: if textprocessing.is_str_type(comments): comments = comments.replace("\r\n", "\n").replace("\n\r","\n").replace("\r","\n") # comments = re.split(r'[\r\n]+', comments) comments = [c for c in re.split(r'[\r\n]+', comments) if c] for c in comments: stream.write("[{}]\n".format(c)) def _write_item_comments(self, stream, item): if not self.suppress_item_comments and item.comments: self._write_comments(stream, item.comments) def _write_item_annotations(self, stream, item): if not self.suppress_annotations and item.annotations: a = nexusprocessing.format_item_annotations_as_comments(item, nhx=self.annotations_as_nhx, real_value_format_specifier=self.real_value_format_specifier) stream.write("{}\n".format(a)) def _write_block_title(self, stream, block): if not self._link_blocks(): return title = self._get_block_title(block) if not title: return stream.write(" TITLE {};\n".format(title)) def _write_link_to_taxa_block(self, stream, taxon_namespace): if not self._link_blocks(): return link_title = self._get_block_title(taxon_namespace) if not link_title: return stream.write(" LINK TAXA = {};\n".format(link_title)) def _get_block_title(self, block): # if self.is_write_block_titles is False then no block titles; # if only one taxon set, or attached taxon set mode, unless self.is_write_block_titles # is explicitly True, then again, we do not write block titles if not self._link_blocks(): return None if block in self._block_title_map: return self._block_title_map[block] if not block.label: title = str(id(block)) else: title = block.label idx = 1 original_title = title title = nexusprocessing.escape_nexus_token( original_title, preserve_spaces=self.preserve_spaces, quote_underscores=not self.unquoted_underscores) while title in self._title_block_map: raw_title = "{}.{}".format(original_title, idx) title = nexusprocessing.escape_nexus_token( raw_title, preserve_spaces=self.preserve_spaces, quote_underscores=not self.unquoted_underscores) idx += 1 self._title_block_map[title] = block self._block_title_map[block] = title return title def _link_blocks(self): """ If only one taxon set in dataset, or in attached taxon set mode, then unless the 'block_titles' directive has been explicitly set to True by the user, block titles and links will not be written. """ if self.suppress_block_titles is None: if len(self.taxon_namespaces_to_write) > 1: return True else: return False else: return self.suppress_block_titles def _format_continuous_character_value(self, v): return "{}".format(v) def _format_discrete_character_value(self, v): return str(v) def _write_character_subsets(self, stream, char_matrix): if not hasattr(char_matrix, "character_subsets") or not char_matrix.character_subsets: return stream.write("BEGIN SETS;\n") for label, char_set in char_matrix.character_subsets.items(): label = nexusprocessing.escape_nexus_token(char_set.label, preserve_spaces=self.preserve_spaces, quote_underscores=not self.unquoted_underscores) ranges = nexusprocessing.group_ranges(char_set.character_indices) pos = " ".join("-".join(str(c+1) for c in r) for r in ranges) stream.write(" charset {} = {};\n".format(label, pos)) stream.write("END;\n\n\n")