Source code for geppy.support.visualization

# coding=utf-8
"""
.. moduleauthor:: Shuhua Gao

This module :mod:`visualization` provides utility functions to visualization the expression tree from a given
K-expression, a gene or a chromosome in GEP.
"""

from ..core.entity import KExpression, Chromosome, Gene
from ..core.symbol import Function, Terminal


def _graph_kexpression(expr, starting_index):
    """
    Create a graph for a K-expression *expr* with the node's number starting from *starting_index*.

    :param expr: k-expression
    :param starting_index: the first number of nodes in the expression tree
    :return: A node list, an edge list, and a dictionary of labels.
    """
    assert len(expr) > 0
    nodes = [starting_index + i for i in range(len(expr))]
    edges = []
    labels = {}
    for i, p in enumerate(expr):
        if isinstance(p, Function):
            labels[starting_index + i] = p.name
        elif isinstance(p, Terminal):
            labels[starting_index + i] = p.format()
        else:
            raise RuntimeError('Unrecognized symbol. Normally, a symbol in the K-expression is either a function '
                               'or a terminal')
    i = 0
    j = 0
    while i < len(expr):
        for _ in range(expr[i].arity):
            j += 1
            edges.append((i + starting_index, j + starting_index))
        i += 1
    return nodes, edges, labels


[docs]def graph(genome, label_renaming_map=None): """ Construct the graph of a genome. It returns in order a node list, an edge list, and a dictionary of the per node labels. The node are represented by numbers, the edges are tuples connecting two nodes (number), and the labels are values of a dictionary for which keys are the node numbers. :param genome: :class:`~geppy.core.entity.KExpression`, :class:`~geppy.core.entity.Gene`, or :class:`~geppy.core.entity.Chromosome`, the genotype of an individual :param label_renaming_map: dict, which maps the old name of a primitive (or a linking function) to a new one for better visualization. The default label for each node is just the name of the primitive placed on this node. For example, you may provide ``renamed_labels={'and_': 'and'}``. :return: A node list, an edge list, and a dictionary of labels. You can visualize a genome and export the tree visualization to an image file directly using the :func:`export_expression_tree` function. """ nodes = [] edges = [] labels = {} if isinstance(genome, KExpression): nodes, edges, labels = _graph_kexpression(genome, 0) elif isinstance(genome, Gene): nodes, edges, labels = _graph_kexpression(genome.kexpression, 0) elif isinstance(genome, Chromosome): if len(genome) == 1: nodes, edges, labels = _graph_kexpression(genome[0].kexpression, 0) else: # multigenic chromosome, we need to concatenate multiple trees starting_index = 1 sub_roots = [] for gene in genome: expr = gene.kexpression sub_roots.append(starting_index) sub_nodes, sub_edges, sub_labels = _graph_kexpression(expr, starting_index) nodes.extend(sub_nodes) edges.extend(sub_edges) labels.update(sub_labels) starting_index += len(expr) # connect subtrees by inserting the linker node as 0 nodes.append(0) for root in sub_roots: edges.append((0, root)) labels[0] = genome.linker.__name__ else: raise TypeError('Only an argument of type KExpression, Gene, and Chromosome is acceptable. The provided ' 'genome type is {}.'.format(type(genome))) # rename_labels labels if label_renaming_map is not None: for k, v in labels.items(): if v in label_renaming_map: labels[k] = label_renaming_map[v] return nodes, edges, labels
[docs]def export_expression_tree(genome, label_renaming_map=None, file='tree.png'): """ Construct the graph of a *genome* and then export it to a *file*. :param genome: :class:`~geppy.core.entity.KExpression`, :class:`~geppy.core.entity.Gene`, or :class:`~geppy.core.entity.Chromosome`, the genotype of an individual :param label_renaming_map: dict, which maps the old name of a primitive (or a linking function) to a new one for better visualization. The default label for each node is just the name of the primitive placed on this node. For example, you may provide ``renamed_labels={'and_': 'and'}``. :param file: str, the file path to draw the expression tree, which may be a relative or absolute one. If no extension is included in *file*, then the default extension 'png' is used. .. note:: This function currently depends on the :mod:`graphviz` module to render the tree. Please first install the `graphviz <https://pypi.org/project/graphviz/>`_ module before using this function. Alternatively, you can always obtain the raw graph data with the :func:`graph` function, then postprocess the data and render them with other tools as you want. """ import graphviz as gv import os.path nodes, edges, labels = graph(genome, label_renaming_map) file_name, ext = os.path.splitext(file) ext = ext.lstrip('.') g = gv.Graph(format=ext) for name, label in labels.items(): g.node(str(name), str(label)) # add node for name1, name2 in edges: g.edge(str(name1), str(name2)) # add edge g.render(file_name)
__all__ = ['graph', 'export_expression_tree']