ml-kuleuven / deepstochlog Goto Github PK

def to_dot(self, term: Term):
    if term not in self._and_or_tree:
        raise RuntimeError("term {} is not in the tree".format(term))
    from deepstochlog.logic import Or, And, NNLeaf, StaticProbability, TermLeaf

    nodes, edges = [], []
    num_and_nodes = 0
    num_or_nodes = 0

    def recurse_term_node(term: Term, parent_hash):
        node = self._and_or_tree[term]
        node_hash = abs(hash(node) + parent_hash)
        if isinstance(node, Or):
            nonlocal num_or_nodes
            num_or_nodes += 1
            nodes.append(
                "{} [label={}]".format(
                    node_hash,
                    '<OR<BR/> <FONT POINT-SIZE="7"> {} </FONT>>'.format(str(term)),
                )
            )
        elif isinstance(node, And):
            nonlocal num_and_nodes
            num_and_nodes += 1
            nodes.append(
                "{} [label={}]".format(
                    node_hash,
                    '<AND<BR/> <FONT POINT-SIZE="7"> {} </FONT>>'.format(str(term)),
                )
            )

        elif isinstance(node, NNLeaf):
            nodes.append('{} [label="{}"]'.format(node_hash, str(node)))
            return

        for child in node.children:
            child_hash = abs(hash(child) + node_hash)
            if not isinstance(child, TermLeaf):
                edges.append((node_hash, child_hash))
            resolve_other_node(child, node_hash)

    def resolve_other_node(node, parent_hash):
        node_hash = abs(hash(node) + parent_hash)
        if isinstance(node, TermLeaf):
            edges.append(
                (
                    parent_hash,
                    abs(hash(self._and_or_tree[node.term]) + node_hash),
                )
            )
            recurse_term_node(node.term, node_hash)
            return
        elif isinstance(node, Or):
            nonlocal num_or_nodes
            num_or_nodes += 1
            nodes.append("{} [label={}]".format(node_hash, '"OR"'))
        elif isinstance(node, And):
            nonlocal num_and_nodes
            num_and_nodes += 1
            nodes.append("{} [label={}]".format(node_hash, '"AND"'))
        elif isinstance(node, StaticProbability):
            nodes.append("{} [label={}]".format(node_hash, node.probability))
            return
        elif isinstance(node, NNLeaf):
            nodes.append('{} [label="{}"]'.format(node_hash, str(node)))
            return
        else:
            raise RuntimeError("unexpected node type")

        for child in node.children:
            child_hash = abs(hash(child) + node_hash)
            if not isinstance(child, TermLeaf):
                edges.append((node_hash, child_hash))
            resolve_other_node(child, node_hash)

    recurse_term_node(term, 0)
    dot_string = (
        "Digraph {\n"
        + "\n".join(nodes)
        + "\n"
        + "\n".join(
            [
                "{} -> {}".format(source, destination)
                for source, destination in edges
            ]
        )
        + "\n}"
    )
    print(
        "run circuit with {} nodes and {} edges ({} and nodes, {} or nodes)".format(
            len(nodes), len(edges), num_and_nodes, num_or_nodes
        )
    )
    return dot_string, {
        "num_nodes": len(nodes),
        "num_edges": len(edges),
        "num_or_nodes": num_or_nodes,
        "num_and_nodes": num_and_nodes,
    }