from typing import TYPE_CHECKING
from pydistsim.gui.drawing import NODE_COLORS
from pydistsim.utils.tree import get_root_node
if TYPE_CHECKING:
from pydistsim.network import NetworkType
[docs]
def show_mst(net: "NetworkType", tree_key="mst"):
"""
Show tree representation of network.
:param net: network to show
:param tree_key: key in nodes memory (dictionary) where parent and
children data is stored in format:
{'parent': parent_node,
'children': [child_node1, child_node2 ...]}
"""
nodesToCheck = [(get_root_node(net, tree_key), 0)]
edgelist = []
levels = {} # level of node in tree, root is 0
while nodesToCheck:
(node, level) = nodesToCheck.pop()
edgelist += [(node, child.unbox()) for child in node.memory[tree_key]["children"]]
levels[node] = NODE_COLORS[level]
nodesToCheck += [(child.unbox(), level + 1) for child in node.memory[tree_key]["children"]]
net.show(edge_filter=edgelist, node_colors=levels)
from matplotlib.pyplot import gca
gca().set_title("Minimum spanning tree in memory['%s']" % tree_key)