Visualising biggus expressions in with networkx and dot (blue are results)

biggus_expression.png

dot_vis.py

import os
import subprocess
import sys

import networkx as nx
import numpy as np

import biggus


def which(program):
    import os
    def is_exe(fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            path = path.strip('"')
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file
    return None


def write_networkx_to_dot(dg, filename):
    import os
    try:
        p = nx.to_pydot(dg)
    except AttributeError:
        raise ImportError("Can not find pydot module. Please install.\n"
                          "    pip install pydot")
    p.set_rankdir('TB')
    dot_name = filename + '.dot'
    with open(dot_name, 'w') as f:
        f.write(p.to_string())

    dot = which('dot')
    if dot is None:
        raise ImportError('Unable to find the dot executable.')
    subprocess.check_output([dot, '-Tpdf', dot_name, '-o', '{}.pdf'.format(filename)])
    print("Writing graph to %s.pdf" % filename)


class Group(biggus.AllThreadedEngine.Group):
    """
    A collection of Array instances which are to be evaluated in
    parallel.

    """
    def dot(self, graph):
        g = graph
        for array in self.arrays:
            iteration_order = range(array.ndim)
            node = self._make_node(array, iteration_order, masked=False)

            result_node_id = id(array)

            # Result nodes are blue.
            g.add_node(result_node_id, label='Array\n{}'.format(array.shape),
                       shape='box', color="blue")

            if isinstance(node, biggus.StreamsHandlerNode):
                sh = node.streams_handler
                node_id = id(node)
                if isinstance(sh, biggus._AggregationStreamsHandler):
                    node_id = id(sh)
                    # Single aggregation.
                    sh_name = sh.__class__.__name__.strip('_').replace('StreamsHandler', '')
                    g.add_node(node_id, label='Aggregation\n({}, ax={})'.format(sh_name, sh.axis), shape='circle')
                    g.add_edge(node_id, result_node_id)

                    input_id = id(sh.array)
                    g.add_node(input_id, label='Array\n{}'.format(sh.array.shape), shape='box')
                    g.add_edge(input_id, node_id)
                elif isinstance(sh, biggus._ElementwiseStreamsHandler):
                    op = sh.operator

                    if isinstance(op, np.ufunc):
                        name = op.__name__
                    else:
                        name = 'Elementwise\n({})'.format(op)
                    g.add_node(node_id, label=name, shape='circle')
                    g.add_edge(node_id, result_node_id)

                    for arr in sh.sources:
                        input_id = id(arr)
                        if isinstance(arr, biggus.BroadcastArray):
                            shp = arr.array.shape
                            if not shp:
                                # Trivial scalars have no shape.
                                label = str(arr.array.ndarray())
                                input_id = label
                            else:
                                label = 'Array\n{}'.format(shp)
                        else:
                            label = 'Array\n{}'.format(arr.shape)
                        g.add_node(input_id, label=label, shape='box')
                        g.add_edge(input_id, node_id)
                else:
                    raise NotImplementedError('Not yet dealt with this type of StreamsHandler.')
            elif isinstance(node, biggus.ProducerNode):
                pass
            else:
                raise NotImplementedError('Not yet dealt with this type {}.'.format(type(node)))


class BiggusExpressionVis(object):
    def _groups(self, arrays):
        # XXX Placeholder implementation which assumes everything
        # is compatible and can be evaluated in parallel.
        return [Group(arrays, range(len(arrays)))]

    def dot(self, *arrays, **kwargs):
        fname = kwargs.pop('fname', 'biggus_expression')
        if kwargs:
            raise TypeError('Unexpected kwargs to dot. {}'.format(kwargs.keys()))

        graph = nx.DiGraph()
        for group in self._groups(arrays):
            group.dot(graph)
        write_networkx_to_dot(graph, filename=fname)
        graph


if __name__ == '__main__':
    # Put the sys prefix directory on the path.
    os.environ['PATH'] += ':{}/bin'.format(sys.prefix)

    e = BiggusExpressionVis()

    a = biggus.zeros([200, 20, 20])
    b = biggus.mean(a - a, axis=0)
    s = biggus.std(a, axis=0)
    e.dot(a, b, b + 1, b - 1, s, s - b)