plot_vars does not show value edges if var has dist

[wiep]

Currently, the plotting function plot_vars does not show edges that determine the value of a node if that node has a distribution.

This potential fix draws edges in red if they influence value as well as distribution.

def _draw_edges(graph, axis, pos, is_var):
    """Adds edges to the figure."""

    edges = list(graph.edges)

    if is_var:
        dist_edges = []
        value_edges = []

        for edge in edges:

            # find distribution edges
            if edge[1].has_dist:
                edge_0_output_nodes = set(edge[0].all_output_nodes())
                edge_0_nodes = edge[0].nodes
                edge_1_input_nodes = set(edge[1].dist_node.all_input_nodes())

                if bool(edge_0_output_nodes.union(edge_0_nodes) & edge_1_input_nodes):
                    dist_edges.append(edge)

            # find value edges
            edge_0_output_nodes = set(edge[0].all_output_nodes())
            edge_0_nodes = edge[0].nodes
            edge_1_input_nodes = set(edge[1].value_node.all_input_nodes())

            if bool(edge_0_output_nodes.union(edge_0_nodes) & edge_1_input_nodes):
                value_edges.append(edge)

        edges_in_both = set(dist_edges) & set(value_edges)
        dist_edges = set(dist_edges) - edges_in_both
        value_edges = set(value_edges) - edges_in_both

        # assigns value_edges to edges to make it comparible with is_var=False
        edges = value_edges

        nx.draw_networkx_edges(
            graph,
            pos,
            edgelist=edges_in_both,
            edge_color="#FF0000",
            arrows=True,
            ax=axis,
            node_size=500,
        )

        nx.draw_networkx_edges(
            graph,
            pos,
            edgelist=dist_edges,
            edge_color="#aaaaaa",
            arrows=True,
            ax=axis,
            node_size=500,
        )


    nx.draw_networkx_edges(
        graph,
        pos,
        edgelist=edges,
        edge_color="#111111",
        arrows=True,
        ax=axis,
        node_size=500,
    )

[jobrachem]

I'm not sure if I understand the scenario. When would a node's value be a function of another node, if the first node has a distribution?

[wiep]

for example when using batching. Imagine three nodes:

• node a hold the batch index
• node b holds the variable
• node c is a calculator selecting from a the elements indicated by b: a[b]
• node c is also associated with a distribution for a[b]. this distribution may depend on a, b, or other nodes.

Regardless of the viability of the scenario, liesel let's you build such a model graph and the plotting should reflect all edges in the model.

[wiep]

See PR #231 for an example.