| 4 | from .algo_node import * |
| 5 | |
| 6 | class AlgoGraph(AlgoVGroup): |
| 7 | def __init__(self, scene, nodes=[], edges=[], **kwargs): |
| 8 | self.nodes = nodes |
| 9 | self.arrows = {} |
| 10 | self.node_objs = {} |
| 11 | self.scene = scene |
| 12 | super().__init__(**kwargs) |
| 13 | |
| 14 | self.init_networkx(nodes, edges) |
| 15 | |
| 16 | for k in nodes: |
| 17 | n = AlgoNode(str(k)) |
| 18 | p = self.get_node_pos(k) |
| 19 | n.shift(p) |
| 20 | self.node_objs[k] = n |
| 21 | self.add(n) |
| 22 | |
| 23 | for k in edges: |
| 24 | self.add_edge_internal(k[0], k[1]) |
| 25 | |
| 26 | self.center() |
| 27 | |
| 28 | def init_networkx(self, nodes, edges): |
| 29 | self.g = nx.Graph() |
| 30 | for k in nodes: |
| 31 | self.g.add_node(k) |
| 32 | for k in edges: |
| 33 | self.g.add_edge(*k) |
| 34 | self.pos_infos = nx.nx_agraph.graphviz_layout(self.g, prog='dot') |
| 35 | # self.pos_infos = nx.random_layout(self.g, center=(0, 0), dim=2) |
| 36 | |
| 37 | def get_node_pos(self, k): |
| 38 | p = self.pos_infos[k] |
| 39 | ratio = 60 |
| 40 | return [p[0]/ratio, p[1]/ratio, 0] |
| 41 | |
| 42 | def clear_edges(self): |
| 43 | self.g.clear_edges() |
| 44 | |
| 45 | for k in self.arrows: |
| 46 | self.scene.play(FadeOut(k, run_time=0.3)) |
| 47 | self.arrows = [] |
| 48 | |
| 49 | def add_edge_internal(self, i, j): |
| 50 | if i == j: |
| 51 | a = Arrow(self.get_node_pos(i), self.get_node_pos(j)+RIGHT*0.1, path_arc=np.pi*1.5, thickness=0.03).scale(0.5) |
| 52 | self.arrows[(i, j)] = a |
| 53 | self.add(a) |
| 54 | else: |
| 55 | a = Arrow(self.get_node_pos(i), self.get_node_pos(j), thickness=0.03) |
| 56 | self.add(a) |
| 57 | self.arrows[(i, j)] = a |
| 58 | |
| 59 | def add_edge(self, i, j): |
| 60 | ni = self.node_objs[i] |
| 61 | nj = self.node_objs[j] |
| 62 | if i == j: |
| 63 | a = Arrow(ni.get_center(), nj.get_center()+RIGHT*0.1, path_arc=np.pi*1.5, thickness=0.03).scale(0.5) |