zodiac.streams.plot_stream

 1# SPDX-License-Identifier: MPL-2.0 AND LicenseRef-Commons-Clause-License-Condition-1.0
 2# <!-- // /*  d a r k s h a p e s */ -->
 3
 4
 5import asyncio
 6
 7import matplotlib as mpl
 8import matplotlib.pyplot as plt
 9import networkx as nx
10
11from zodiac.streams import ModelStream
12
13
14async def main():
15    model_source = ModelStream()
16    await model_source.model_graph()
17    await model_source.trace_models("image", "speech")
18    graph_copy = model_source._graph.intent_graph.to_undirected()
19    nx_graph = nx.Graph(graph_copy)
20
21    pos = nx.spring_layout(nx_graph, seed=7)  # positions for all nodes - seed for reproducibility
22
23    # nodes
24    nx.draw_networkx_nodes(nx_graph, pos, node_size=700)
25
26    # edges
27    # print(nx_graph.edges())
28    # print(nx_graph.adjacency())
29    path = nx.bidirectional_shortest_path(nx_graph, "image", "speech")
30    path_list = [nx_graph[path[x]][path[x + 1]] for x in range(len(path) - 1)]
31    nx.draw_networkx_edges(nx_graph, pos, edgelist=nx_graph.edges(), width=6)
32    nx.draw_networkx_edges(nx_graph, pos, edgelist=path_list, width=6, alpha=0.5, edge_color="b", style="dashed")
33
34    # node labels
35    nx.draw_networkx_labels(nx_graph, pos, font_size=20, font_family="sans-serif")
36    # edge weight labels
37    edge_labels = nx.get_edge_attributes(nx_graph, "weight")
38    nx.draw_networkx_edge_labels(nx_graph, pos, edge_labels)
39
40    ax = plt.gca()
41    ax.margins(0.08)
42    plt.axis("off")
43    plt.tight_layout()
44    plt.show()
45
46
47if __name__ == "__main__":
48    asyncio.run(main())