From 8b5a789640e83c2952854dd02dfe090f1efca973 Mon Sep 17 00:00:00 2001 From: Tigar Cyr Date: Sun, 30 Jun 2024 22:27:49 -0700 Subject: [PATCH] Apply separate updaters to each edge in a Graph --- manim/mobject/graph.py | 202 +++++++----------- .../control_data/graph/digraph_add_edge.npz | Bin 0 -> 4359 bytes .../graph/graph_concurrent_animations.npz | Bin 0 -> 5699 bytes tests/test_graphical_units/test_graph.py | 34 +++ 4 files changed, 108 insertions(+), 128 deletions(-) create mode 100644 tests/test_graphical_units/control_data/graph/digraph_add_edge.npz create mode 100644 tests/test_graphical_units/control_data/graph/graph_concurrent_animations.npz create mode 100644 tests/test_graphical_units/test_graph.py diff --git a/manim/mobject/graph.py b/manim/mobject/graph.py index 72a26d27db..22ae33f84e 100644 --- a/manim/mobject/graph.py +++ b/manim/mobject/graph.py @@ -590,7 +590,6 @@ def __init__( nx_graph = self._empty_networkx_graph() nx_graph.add_nodes_from(vertices) - nx_graph.add_edges_from(edges) self._graph = nx_graph if isinstance(labels, dict): @@ -627,6 +626,20 @@ def __init__( self.vertices = {v: vertex_type(**self._vertex_config[v]) for v in vertices} self.vertices.update(vertex_mobjects) + if edge_config is None: + edge_config = {} + + default_configs, per_edge_tip_configs = GenericGraph._separate_child_configs( + edge_config, lambda k: isinstance(k, tuple) + ) + self.default_edge_config, self.default_tip_config = ( + GenericGraph._separate_tip_configs(default_configs) + ) + + self.edges = {} + for edge in edges: + self._add_edge(edge, edge_type, per_edge_tip_configs.get(edge)) + self.change_layout( layout=layout, layout_scale=layout_scale, @@ -634,51 +647,25 @@ def __init__( partitions=partitions, root_vertex=root_vertex, ) - - # build edge_config - if edge_config is None: - edge_config = {} - default_tip_config = {} - default_edge_config = {} - if edge_config: - default_tip_config = edge_config.pop("tip_config", {}) - default_edge_config = { - k: v - for k, v in edge_config.items() - if not isinstance( - k, tuple - ) # everything that is not an edge is an option - } - self._edge_config = {} - self._tip_config = {} - for e in edges: - if e in edge_config: - self._tip_config[e] = edge_config[e].pop( - "tip_config", copy(default_tip_config) - ) - self._edge_config[e] = edge_config[e] - else: - self._tip_config[e] = copy(default_tip_config) - self._edge_config[e] = copy(default_edge_config) - - self.default_edge_config = default_edge_config - self._populate_edge_dict(edges, edge_type) - self.add(*self.vertices.values()) - self.add(*self.edges.values()) - - self.add_updater(self.update_edges) @staticmethod def _empty_networkx_graph() -> nx.classes.graph.Graph: """Return an empty networkx graph for the given graph type.""" raise NotImplementedError("To be implemented in concrete subclasses") - def _populate_edge_dict( - self, edges: list[tuple[Hashable, Hashable]], edge_type: type[Mobject] - ): - """Helper method for populating the edges of the graph.""" - raise NotImplementedError("To be implemented in concrete subclasses") + @staticmethod + def _separate_tip_configs(config: dict) -> (dict, dict): + edge_config, tip_config_holder = GenericGraph._separate_child_configs( + config, lambda k: k == "tip_config" + ) + return edge_config, tip_config_holder.get("tip_config", {}) + + @staticmethod + def _separate_child_configs(config: dict, is_child_key) -> (dict, dict): + default_config = {k: v for k, v in config.items() if not is_child_key(k)} + per_child_configs = {k: v for k, v in config.items() if is_child_key(k)} + return default_config, per_child_configs def __getitem__(self: Graph, v: Hashable) -> Mobject: return self.vertices[v] @@ -955,8 +942,6 @@ def _remove_vertex(self, vertex): self._vertex_config.pop(vertex) edge_tuples = [e for e in self.edges if vertex in e] - for e in edge_tuples: - self._edge_config.pop(e) to_remove = [self.edges.pop(e) for e in edge_tuples] to_remove.append(self.vertices.pop(vertex)) @@ -1028,28 +1013,26 @@ def _add_edge( """ if edge_config is None: - edge_config = self.default_edge_config.copy() - added_mobjects = [] - for v in edge: - if v not in self.vertices: - added_mobjects.append(self._add_vertex(v)) + edge_config = {} + added_vertices = [self._add_vertex(v) for v in edge if v not in self.vertices] u, v = edge self._graph.add_edge(u, v) - base_edge_config = self.default_edge_config.copy() - base_edge_config.update(edge_config) - edge_config = base_edge_config - self._edge_config[(u, v)] = edge_config + edge_mobject = self._create_edge(self[u], self[v], edge_type, edge_config) + self.edges[(u, v)] = edge_mobject + self.add(edge_mobject) + return self.get_group_class()(*added_vertices, edge_mobject) + def _create_edge(self, u, v, edge_type, config): edge_mobject = edge_type( - self[u].get_center(), self[v].get_center(), z_index=-1, **edge_config + u.get_center(), + v.get_center(), + z_index=-1, + **{**self.default_edge_config, **config}, ) - self.edges[(u, v)] = edge_mobject - - self.add(edge_mobject) - added_mobjects.append(edge_mobject) - return self.get_group_class()(*added_mobjects) + edge_mobject.add_updater(self._generate_edge_updater(u, v, config)) + return edge_mobject def add_edges( self, @@ -1087,29 +1070,14 @@ def add_edges( """ if edge_config is None: edge_config = {} - non_edge_settings = {k: v for (k, v) in edge_config.items() if k not in edges} - base_edge_config = self.default_edge_config.copy() - base_edge_config.update(non_edge_settings) - base_edge_config = {e: base_edge_config.copy() for e in edges} - for e in edges: - base_edge_config[e].update(edge_config.get(e, {})) - edge_config = base_edge_config - - edge_vertices = set(it.chain(*edges)) - new_vertices = [v for v in edge_vertices if v not in self.vertices] - added_vertices = self.add_vertices(*new_vertices, **kwargs) - - added_mobjects = sum( - ( - self._add_edge( - edge, - edge_type=edge_type, - edge_config=edge_config[edge], - ).submobjects - for edge in edges - ), - added_vertices, - ) + + added_mobjects = [ + mobject + for edge in edges + for mobject in self._add_edge( + edge, edge_type, edge_config.get(edge) + ).submobjects + ] return self.get_group_class()(*added_mobjects) @override_animate(add_edges) @@ -1145,7 +1113,6 @@ def _remove_edge(self, edge: tuple[Hashable]): edge_mobject = self.edges.pop(edge) self._graph.remove_edge(*edge) - self._edge_config.pop(edge, None) self.remove(edge_mobject) return edge_mobject @@ -1544,29 +1511,17 @@ def construct(self): def _empty_networkx_graph() -> nx.Graph: return nx.Graph() - def _populate_edge_dict( - self, edges: list[tuple[Hashable, Hashable]], edge_type: type[Mobject] - ): - self.edges = { - (u, v): edge_type( - self[u].get_center(), - self[v].get_center(), - z_index=-1, - **self._edge_config[(u, v)], - ) - for (u, v) in edges - } - - def update_edges(self, graph): - for (u, v), edge in graph.edges.items(): - # Undirected graph has a Line edge + def _generate_edge_updater(self, u, v, config): + def edge_updater(edge): edge.set_points_by_ends( - graph[u].get_center(), - graph[v].get_center(), - buff=self._edge_config.get("buff", 0), - path_arc=self._edge_config.get("path_arc", 0), + u.get_center(), + v.get_center(), + buff=config.get("buff", 0), + path_arc=config.get("path_arc", 0), ) + return edge_updater + def __repr__(self: Graph) -> str: return f"Undirected graph on {len(self.vertices)} vertices and {len(self.edges)} edges" @@ -1751,39 +1706,30 @@ def construct(self): def _empty_networkx_graph() -> nx.DiGraph: return nx.DiGraph() - def _populate_edge_dict( - self, edges: list[tuple[Hashable, Hashable]], edge_type: type[Mobject] - ): - self.edges = { - (u, v): edge_type( - self[u], - self[v], - z_index=-1, - **self._edge_config[(u, v)], - ) - for (u, v) in edges - } - - for (u, v), edge in self.edges.items(): - edge.add_tip(**self._tip_config[(u, v)]) - - def update_edges(self, graph): - """Updates the edges to stick at their corresponding vertices. + def _create_edge(self, u, v, edge_type, config): + edge_config, tip_config = GenericGraph._separate_tip_configs(config) + edge_mobject = edge_type( + u.get_center(), + v.get_center(), + z_index=-1, + **{**self.default_edge_config, **edge_config}, + ) + edge_mobject.add_updater(self._generate_edge_updater(u, v, edge_config)) + edge_mobject.add_tip(**{**self.default_tip_config, **tip_config}) + return edge_mobject - Arrow tips need to be repositioned since otherwise they can be - deformed. - """ - for (u, v), edge in graph.edges.items(): + def _generate_edge_updater(self, u, v, config): + def edge_updater(edge): tip = edge.pop_tips()[0] - # Passing the Mobject instead of the vertex makes the tip - # stop on the bounding box of the vertex. edge.set_points_by_ends( - graph[u], - graph[v], - buff=self._edge_config.get("buff", 0), - path_arc=self._edge_config.get("path_arc", 0), + u.get_center(), + v.get_center(), + buff=config.get("buff", 0), + path_arc=config.get("path_arc", 0), ) edge.add_tip(tip) + return edge_updater + def __repr__(self: DiGraph) -> str: return f"Directed graph on {len(self.vertices)} vertices and {len(self.edges)} edges" diff --git a/tests/test_graphical_units/control_data/graph/digraph_add_edge.npz b/tests/test_graphical_units/control_data/graph/digraph_add_edge.npz new file mode 100644 index 0000000000000000000000000000000000000000..b65c18fc042daeb10c3c461afd0284bcca62fd3c GIT binary patch literal 4359 zcmWIWW@gc4U|`??Vnv4Qrj?uiLjfOy2t!&?Vs2`DN@7W(US2^ZBZB}#1D6p{B?CjL z0GMKUduwBs%1u$Wgso2P%wDeDT3d5wB<$6>@qCa4k6<+eT4~o(vHnLQfAgT zsZ%EO%FKj&4qRVYzA{X5@DIoja4fbdPh7bD3*Lv|33Zv^!ojpuQfl@Cw9TlWm*AQ&!gpR#`4>leH9R;@Z^*A0xDt+~>dOZTeI!nlD#1tSgVeLferzU_}J?gOJhSqr`Aq8{_V}p%@J$EtS5^vV9C9G z%*1!1*4&lnZ7(PP`tZ;>$87f8liUGKi4sY>j}`Xpd*5AHP*9K$ERO%NG2A&H`*a;s z#j9ue4ZfFC=6l95JiF9)Q17kOo2a;f_Ts4h4gZf`ycjrV?%bpQ{?&bpWcX2oXgM`);P(rH2m_@7%fLa@k~KgwCT! zkCe=2`);gc^fKXb&be;Uckz3z;JNeXAHRF2=jZQlKhyl{&i#xG-;2X{u4b6G>!)-= zmG)`d%g(O5b7Sr9Q`#yV-8<7JYwzzrLM5L%f>x>(RIx;1GCH zR#GC8o}RvO|Ni;4Vhnc3#lwaey~p3and#~4C-03rp0!mA7>k*=w)JMtQVO*So?pGJ zy`-W-!ra{aqKU7LSoe+{J37v$1?K=8-@JzfZkzTcf2oac`TY6wj@`SvFJHdw|ID}c zT|UEszjs34X*Dd*`^Prp-l?MVUnX3Ve(?3{*Y7(o#T~tR^=kh;2Jh3{0jsto2Ygz0 zux9(CA3t^glW5J%H*1&&IIu>%^%HIne7WM4-*El<_506(B};gCxP29~jxM}#`FywQ z@#DvfJ~Mqt-L#o+h)0*h$$!iN-i%DT%&5Z_paBWc-~_V&8o+HU)Ip2@Z&o%?5HJFv L4Upan9>@RyG_H@_ literal 0 HcmV?d00001 diff --git a/tests/test_graphical_units/control_data/graph/graph_concurrent_animations.npz b/tests/test_graphical_units/control_data/graph/graph_concurrent_animations.npz new file mode 100644 index 0000000000000000000000000000000000000000..f083cd18cd1d8c3da2f281d5e2644bc4b0687e4c GIT binary patch literal 5699 zcmeHLc~p~E7Jq<}%392dIK9MWBj+fI!)T5I{&m65A;%vXm-C zL4v~~AYufR$PywIh{}>8AVv()pt3Jv2?B=9`_RsG`tKZj#!i!y?>pz^zV+VU@BZHP zaFN-x2LP}I{{9R!Kkl2By!|Kw6%ZC56cZYFG$)$$+I;XNzSg?CI$1I=-tX{A*VjtRQ&b6=k6ZY_JIoxV|1Ec!}Er|=6=gP*6y}NXGwOn zpoPM2aTv`S-$OC8xGE{YWa=){N7iG6FVzCpoAOjK4M&75NoOEBp#F&PhZwYY!KGkb z7|USZSowBwwQ;{>ZeD@KTA$~~u9mBj1JNlY?$EHugS0pB!8w*dMC)>5sfQb+9=Q5X=j@#}aIEobvl@j_;mTqt+%hX}`#M*kGE% z0_T>7_%mc83rKVq9kxfA?c^e@i~ES%RwL-2v6zAH^GWbnR%%-gF=c^!9- z-TIAOUAXj9(c@{+c^1z(Tc?{g?q-uT_2~R^0W_03Mif8v_G93*7s!DmVf)ktaEmz{ zPR(>e5ii@u@EQ|q@S{D!J%zGV%uK`4^4*qI3#;gauTub`hDb+2^q+1hYGtv#tj(6@+wqzpqWbGx2fTNS+} zAG%c>v;KOeCT^r7d9s=zSRC(b4-lG_ClXS7E?RR&S^O5j~K%XLKtWsbNW^4hWVDqRk$3vl zN1rzEx_p!)i=`Aeq!N1fva_^}L{6pG2!hjkwWHEh0#m0cX*#_#56y1Cj|PwS6q*o4 zQ&SCKEOz^D*Hbs6233m+al_3TCl$2~MQ*P7W(GtHusYOr!Su{1Y25C#k?*y(VJ0zz zTN}#7OWpo{mrF2*W*-Oh?km9=bnu#J+9c}X)i(p`xn(R(4zlHRNy@DC)hLtlrB}j{ zx^S6GtTH~PjI*ClUW<9X*yr9pTKTPlS{O$bz|lWxp)``k1TE|M(pLeI6aAHNnxFXh z!*vszc)@{zkNJ7FY5*>${FU|ro3$}P`0+YHbVEFkbW_@HQuidgtb@ z*JXo>jNZEJ&2n-h3_ckT|H{qjad=!G<9E>--@$bT2htInuNRqKRRifzCSEN}-Gt3G zd#e#wt7nE{S=x}f6OpdHI7W%xHUB2n0*jdW0)1du!l-(aVT>>`uK1GPyEJ(~D6`-3 z;txgePB84f+-1sX>nXGifg+4s2e9;NKZa9EY2pmb>o9ZIyv^&_zeYSLO`d5O6+CIc zhcr$+ZT*H@2g7KlWo_nBa|)icG~g=~FZ0<2u`L=9ScU|{?reL}m!b@?Sl;icX^H#9 z#xU`V$1W!Lmt6w><^2V!5HjeOPpqZJFqM@0m?2dt1N;v+C!5sA@^zM<;N2k{KTTa- zEG{uVd`mfyB2E-aye_-mqbY!}uJfLxq(;(8{6&XUoBC+(7gHs73t>Bl=1-Da8p+0V zj22M>NS5Ytf;Nbpr|WC8+5SyQ?GTLuC>D%tB7gk5Jc716N*FB&*VB`OboR45Cqu-m zOD5#kuV9}e&wl*Y(V2HE4yCAVb0FaWJdf|}=I^3vV)-rV8X5*R6P4w(`!G7cKXKxO zpSHP`)!1QqV2CP2)leksw^A0udOacdx)2NVsD~=Nb8!Cy5Z$$qyFO0rKdDPr=iwO1 zFBxC;7Q^iK(ZQyq7}ESbz25PJ;MrlON$hYl(k%9|oHmbT3?YONth5S2hOiDiQH$tP zNIuLp@=ebo%(cqZ!RDr!#wx3|vHnUgA>5hDsHi*N+=yrl)S3LbPBiN^dRbE%>vmSp z-FQZz&yTf&Dh}D|BKgk3D2HvyCykK*1`J1riL=Z^*0T$y;ptiXimKz`$n)YfcD63M zr667o4ZOdbv2@j5THs*Vy(P}PRt0QmR}n_?P@LwRrDZ@idI3j?ulQL9?wFKSP*u@h zf(y9CSEf$ocZbDU1nwNy;x6+kNGErUI6*`L-}Q_cW#5Pu^g8QyLnk{eC3znj?b!rE zUxZU5aLCP33QNCHUfj%(ekV}){~mN=iE!gkb;UsB1_^}KX%xEYJ@_2|#skwyw)asR z>djL|M#jX>?rIgNxLD?r*U^%?7JT$*>DXl)gE*;i1R-oC>bmz#MQdy8mo>qTLc$wS zXDx2XlBPaMPyya$C7{AkTNdj*t`Dy4bC&Y(w%JO(^1u@u$jy?zbmhM`as2uITK(lD zoS}5%vT@q1k%#C5IHzG559ge+O-YG;`!Q7+nNE;07U}WMA1Eqz zr0+0FAMB%TB^rUd*P?BOE28EKV*Upx2~Zq$JO3Aa*MQ~&?~ literal 0 HcmV?d00001 diff --git a/tests/test_graphical_units/test_graph.py b/tests/test_graphical_units/test_graph.py new file mode 100644 index 0000000000..9644e44bc5 --- /dev/null +++ b/tests/test_graphical_units/test_graph.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from manim import * +from manim.utils.testing.frames_comparison import frames_comparison + +__module_test__ = "graph" + + +@frames_comparison(last_frame=False) +def test_graph_concurrent_animations(scene): + vertices = [0, 1] + positions = {0: [-1, 0, 0], 1: [1, 0, 0]} + g = Graph(vertices, [], layout=positions) + scene.play(g[1].animate.move_to([1, 1, 0]), g.animate.add_edges((0, 1))) + scene.wait(0.1) + + +@frames_comparison(last_frame=False) +def test_digraph_add_edge(scene): + vertices = [0, 1] + positions = {0: [-1, 0, 0], 1: [1, 0, 0]} + g = DiGraph( + vertices, + [], + layout=positions, + edge_config={ + "tip_config": { + "tip_shape": ArrowSquareTip, + "tip_length": 0.15, + } + }, + ) + scene.play(g.animate.add_edges((0, 1))) + scene.wait(0.1)