diff --git a/.gitignore b/.gitignore index abec5da495..03fa2f809b 100644 --- a/.gitignore +++ b/.gitignore @@ -134,3 +134,29 @@ dist/ # Ignore the built dependencies third_party/* + +# Virtual environments +.venv/ +venv/ +env/ +testenv/ + +# Python cache +__pycache__/ +*.pyc +*.pyo +*.pyd + +# mypy +.mypy_cache/ + +# pytest +.pytest_cache/ + +# IDEs +.vscode/ +.idea/ + +# OS +.DS_Store +Thumbs.db diff --git a/manim/mobject/geometry/arc.py b/manim/mobject/geometry/arc.py index 32b1133a6b..5d0aeab89b 100644 --- a/manim/mobject/geometry/arc.py +++ b/manim/mobject/geometry/arc.py @@ -118,15 +118,16 @@ def add_tip( tip_length: float | None = None, tip_width: float | None = None, at_start: bool = False, + is_loop: bool = False, ) -> Self: """Adds a tip to the TipableVMobject instance, recognising that the endpoints might need to be switched if it's a 'starting tip' or not. """ if tip is None: - tip = self.create_tip(tip_shape, tip_length, tip_width, at_start) + tip = self.create_tip(tip_shape, tip_length, tip_width, at_start, is_loop) else: - self.position_tip(tip, at_start) + self.position_tip(tip, at_start, is_loop) self.reset_endpoints_based_on_tip(tip, at_start) self.assign_tip_attr(tip, at_start) self.add(tip) @@ -138,12 +139,13 @@ def create_tip( tip_length: float | None = None, tip_width: float | None = None, at_start: bool = False, + is_loop: bool = False, ) -> tips.ArrowTip: """Stylises the tip, positions it spatially, and returns the newly instantiated tip to the caller. """ tip = self.get_unpositioned_tip(tip_shape, tip_length, tip_width) - self.position_tip(tip, at_start) + self.position_tip(tip, at_start, is_loop) return tip def get_unpositioned_tip( @@ -175,7 +177,9 @@ def get_unpositioned_tip( tip = tip_shape(length=tip_length, **style) return tip - def position_tip(self, tip: tips.ArrowTip, at_start: bool = False) -> tips.ArrowTip: + def position_tip( + self, tip: tips.ArrowTip, at_start: bool = False, is_loop: bool = False + ) -> tips.ArrowTip: # Last two control points, defining both # the end, and the tangency direction if at_start: @@ -184,25 +188,36 @@ def position_tip(self, tip: tips.ArrowTip, at_start: bool = False) -> tips.Arrow else: handle = self.get_last_handle() anchor = self.get_end() + angles = cartesian_to_spherical(handle - anchor) - tip.rotate( - angles[1] - PI - tip.tip_angle, - ) # Rotates the tip along the azimuthal - if not hasattr(self, "_init_positioning_axis"): - axis = np.array( - [ - np.sin(angles[1]), - -np.cos(angles[1]), - 0, - ] - ) # Obtains the perpendicular of the tip - tip.rotate( - -angles[2] + PI / 2, - axis=axis, - ) # Rotates the tip along the vertical wrt the axis - self._init_positioning_axis = axis - tip.shift(anchor - tip.tip_point) + if is_loop: + alpha = angles[1] - 10 * PI / 9 - tip.tip_angle + tip.rotate( + alpha, + ) + tip.move_to(handle) + else: + tip.rotate( + angles[1] - PI - tip.tip_angle, + ) # Rotates the tip along the azimuthal + + if not hasattr(self, "_init_positioning_axis"): + axis = np.array( + [ + np.sin(angles[1]), + -np.cos(angles[1]), + 0, + ] + ) # Obtains the perpendicular of the tip + + tip.rotate( + -angles[2] + PI / 2, + axis=axis, + ) # Rotates the tip along the vertical wrt the axis + self._init_positioning_axis = axis + + tip.shift(anchor - tip.tip_point) return tip def reset_endpoints_based_on_tip(self, tip: tips.ArrowTip, at_start: bool) -> Self: diff --git a/manim/mobject/geometry/line.py b/manim/mobject/geometry/line.py index 2cd7aff807..c87f4b2fbc 100644 --- a/manim/mobject/geometry/line.py +++ b/manim/mobject/geometry/line.py @@ -12,6 +12,7 @@ "DoubleArrow", "Angle", "RightAngle", + "LoopEdge", ] from typing import TYPE_CHECKING, Any, Literal, cast @@ -604,7 +605,7 @@ def __init__( self._set_stroke_width_from_length() def scale(self, factor: float, scale_tips: bool = False, **kwargs: Any) -> Self: # type: ignore[override] - r"""Scale an arrow, but keep stroke width and arrow tip size fixed. + """Scale an arrow, but keep stroke width and arrow tip size fixed. .. seealso:: @@ -1210,3 +1211,86 @@ def __init__( **kwargs: Any, ) -> None: super().__init__(line1, line2, radius=length, elbow=True, **kwargs) + + +class LoopEdge(TipableVMobject): + def __init__( + self, + vertex: Point3DLike, + graph_center: Point3DLike, + edge_type: type[Mobject] = Line, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.anchor_vertex = self._pointify(vertex) + self.graph_center = self._pointify(graph_center) + self.edge_type = edge_type + edge = edge_type() + points = self._generate_arch_points() + edge.set_points(points) + self.set_points(edge.points) + + def _generate_arch_points(self, vertex: Point3DLike | None = None) -> np.array: + radius = 0.5 + anchor_vertex = self.anchor_vertex if vertex is None else self._pointify(vertex) + + direction = anchor_vertex - self.graph_center + norm = np.linalg.norm(direction) + direction = UP if norm == 0 else direction / norm + + self.arc_center = anchor_vertex + direction * radius + + vector = anchor_vertex - self.arc_center + self.angle = np.arctan2(vector[1], vector[0]) + + arc = Arc( + arc_center=self.arc_center, + radius=radius, + start_angle=self.angle, + angle=2 * PI - 1e-3, + ) + return arc.points + + def set_points_by_vertex(self, vertex: Point3DLike | Mobject) -> None: + points = self._generate_arch_points(vertex) + + edge = self.edge_type() + edge.set_points(points) + self.set_points(edge.points) + + def get_center(self) -> Point3DLike: + return self.arc_center + + def get_start(self) -> Point3DLike: + return self.points[0] + + def get_end(self) -> Point3DLike: + return self.points[0] # technically should be [-1], but that breaks the loop + + def get_arch_point(self, t: float) -> Point3DLike: + radius = 0.5 + theta = self.angle + 2 * t * PI + return ( + np.array([radius * np.cos(theta), radius * np.sin(theta), 0]) + + self.arc_center + ) + + def get_first_handle(self) -> Point3DLike: + return self.get_arch_point(0.18) + + def get_last_handle(self) -> Point3DLike: + return self.get_arch_point(0.82) + + # copied from Line + def _pointify( + self, + mob_or_point: Mobject | Point3DLike, + direction: Vector3DLike | None = None, + ) -> Point3D: + if isinstance(mob_or_point, (Mobject, OpenGLMobject)): + mob = mob_or_point + if direction is None: + return mob.get_center() + else: + return mob.get_boundary_point(direction) + return np.array(mob_or_point) diff --git a/manim/mobject/graph.py b/manim/mobject/graph.py index 7eaee412b4..097a8ba3c5 100644 --- a/manim/mobject/graph.py +++ b/manim/mobject/graph.py @@ -26,7 +26,7 @@ from manim.animation.composition import AnimationGroup from manim.animation.creation import Create, Uncreate from manim.mobject.geometry.arc import Dot, LabeledDot -from manim.mobject.geometry.line import Line +from manim.mobject.geometry.line import Line, LoopEdge from manim.mobject.mobject import Mobject, override_animate from manim.mobject.opengl.opengl_compatibility import ConvertToOpenGL from manim.mobject.opengl.opengl_mobject import OpenGLMobject @@ -1093,7 +1093,7 @@ def add_edges( ( self._add_edge( edge, - edge_type=edge_type, + edge_type=edge_type, # TODO: change edge type for loop edges , has to be defined (in line.py) edge_config=edge_config[edge], ).submobjects for edge in edges @@ -1540,24 +1540,36 @@ def _populate_edge_dict( self, edges: list[tuple[Hashable, Hashable]], edge_type: type[Mobject] ): self.edges = { - (u, v): edge_type( - start=self[u].get_center(), - end=self[v].get_center(), - z_index=-1, - **self._edge_config[(u, v)], + (u, v): ( + edge_type( + start=self[u].get_center(), + end=self[v].get_center(), + z_index=-1, + **self._edge_config[(u, v)], + ) + if u != v + else LoopEdge( + vertex=self[u].get_center(), + graph_center=self.get_center(), + edge_type=edge_type, + 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 - 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), - ) + if u != v: + 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), + ) + else: + edge.set_points_by_vertex(graph[u].get_center()) def __repr__(self: Graph) -> str: return f"Undirected graph on {len(self.vertices)} vertices and {len(self.edges)} edges" @@ -1747,17 +1759,27 @@ def _populate_edge_dict( self, edges: list[tuple[Hashable, Hashable]], edge_type: type[Mobject] ): self.edges = { - (u, v): edge_type( - start=self[u], - end=self[v], - z_index=-1, - **self._edge_config[(u, v)], + (u, v): ( + edge_type( + start=self[u], + end=self[v], + z_index=-1, + **self._edge_config[(u, v)], + ) + if u != v + else LoopEdge( + vertex=self[u].get_center(), + graph_center=self.get_center(), + edge_type=edge_type, + 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)]) + edge.add_tip(is_loop=(u == v), **self._tip_config[(u, v)]) def update_edges(self, graph): """Updates the edges to stick at their corresponding vertices. @@ -1767,15 +1789,20 @@ def update_edges(self, graph): """ for (u, v), edge in graph.edges.items(): 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), - ) - edge.add_tip(tip) + if u != v: + # 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), + ) + edge.add_tip(tip) + else: + edge.set_points_by_vertex(graph[u]) + + edge.add_tip(tip, is_loop=True) def __repr__(self: DiGraph) -> str: return f"Directed graph on {len(self.vertices)} vertices and {len(self.edges)} edges" diff --git a/tests/module/mobject/test_graph.py b/tests/module/mobject/test_graph.py index 16c32da88e..94e1a4aa71 100644 --- a/tests/module/mobject/test_graph.py +++ b/tests/module/mobject/test_graph.py @@ -18,6 +18,18 @@ def test_graph_creation(): assert str(G_directed) == "Directed graph on 4 vertices and 4 edges" +def test_graph_creation_with_loopedge(): + vertices = [1, 2, 3, 4] + edges = [(1, 2), (2, 3), (3, 4), (4, 4)] + layout = {1: [0, 0, 0], 2: [1, 1, 0], 3: [1, -1, 0], 4: [-1, 0, 0]} + G_manual = Graph(vertices=vertices, edges=edges, layout=layout) + assert str(G_manual) == "Undirected graph on 4 vertices and 4 edges" + G_spring = Graph(vertices=vertices, edges=edges) + assert str(G_spring) == "Undirected graph on 4 vertices and 4 edges" + G_directed = DiGraph(vertices=vertices, edges=edges) + assert str(G_directed) == "Directed graph on 4 vertices and 4 edges" + + def test_graph_add_vertices(): G = Graph([1, 2, 3], [(1, 2), (2, 3)]) G.add_vertices(4) @@ -59,10 +71,16 @@ def test_graph_add_edges(): assert set(G.vertices.keys()) == {1, 2, 3, 4, 5, 42} assert set(G.edges.keys()) == {(1, 2), (2, 3), (1, 3), (1, 42)} + added_mobjects = G.add_edges((20, 20)) + assert str(added_mobjects.submobjects) == "[Dot, Line]" + assert str(G) == "Undirected graph on 7 vertices and 5 edges" + assert set(G.vertices.keys()) == {1, 2, 3, 4, 5, 42, 20} + assert set(G.edges.keys()) == {(1, 2), (2, 3), (1, 3), (1, 42), (20, 20)} + added_mobjects = G.add_edges((4, 5), (5, 6), (6, 7)) assert len(added_mobjects) == 5 - assert str(G) == "Undirected graph on 8 vertices and 7 edges" - assert set(G.vertices.keys()) == {1, 2, 3, 4, 5, 42, 6, 7} + assert str(G) == "Undirected graph on 9 vertices and 8 edges" + assert set(G.vertices.keys()) == {1, 2, 3, 4, 5, 42, 6, 7, 20} assert set(G._graph.nodes()) == set(G.vertices.keys()) assert set(G.edges.keys()) == { (1, 2), @@ -72,6 +90,7 @@ def test_graph_add_edges(): (4, 5), (5, 6), (6, 7), + (20, 20), } assert set(G._graph.edges()) == set(G.edges.keys()) @@ -90,6 +109,30 @@ def test_graph_remove_edges(): assert set(G._graph.edges()) == set() assert set(G.edges.keys()) == set() + G = Graph( + [1, 2, 3, 4, 5], + [(1, 2), (2, 3), (3, 4), (4, 5), (1, 5), (4, 4), (2, 2), (1, 1)], + ) + removed_mobjects = G.remove_edges((4, 4)) + assert str(removed_mobjects.submobjects) == "[LoopEdge]" + assert str(G) == "Undirected graph on 5 vertices and 7 edges" + assert set(G.edges.keys()) == { + (1, 2), + (2, 3), + (3, 4), + (4, 5), + (1, 5), + (2, 2), + (1, 1), + } + assert set(G._graph.edges()) == set(G.edges.keys()) + + removed_mobjects = G.remove_edges((2, 3), (2, 2), (1, 1)) + assert len(removed_mobjects) == 3 + assert str(G) == "Undirected graph on 5 vertices and 4 edges" + assert set(G.edges.keys()) == {(1, 2), (3, 4), (4, 5), (1, 5)} + assert set(G._graph.edges()) == set(G.edges.keys()) + def test_graph_accepts_labeledline_as_edge_type(): vertices = [1, 2, 3, 4] @@ -113,6 +156,28 @@ def test_graph_accepts_labeledline_as_edge_type(): assert isinstance(edge_obj, LabeledLine) assert hasattr(edge_obj, "label") + vertices = [1, 2, 3] + edges = [(1, 1), (2, 3), (1, 3), (2, 2)] + edge_config = { + (1, 1): {"label": "A"}, + (2, 3): {"label": "B"}, + (1, 3): {"label": "C"}, + (2, 2): {"label": "D"}, + } + + G_manual = Graph(vertices, edges, edge_type=LabeledLine, edge_config=edge_config) + G_directed = DiGraph( + vertices, edges, edge_type=LabeledLine, edge_config=edge_config + ) + + for edge_obj in G_manual.edges.values(): + assert isinstance(edge_obj, LabeledLine) + assert hasattr(edge_obj, "label") + + for edge_obj in G_directed.edges.values(): + assert isinstance(edge_obj, LabeledLine) + assert hasattr(edge_obj, "label") + def test_custom_animation_mobject_list(): G = Graph([1, 2, 3], [(1, 2), (2, 3)]) @@ -191,6 +256,13 @@ def test_graph_change_layout(): assert str(G) == "Undirected graph on 3 vertices and 2 edges" +def test_graph_change_layout_with_loop_edge(): + for layout in (layout for layout in _layouts if layout not in ["tree", "partite"]): + G = Graph([1, 2, 3], [(1, 2), (2, 3), (3, 3)]) + G.change_layout(layout=layout) + assert str(G) == "Undirected graph on 3 vertices and 3 edges" + + def test_tree_layout_no_root_error(): with pytest.raises(ValueError) as excinfo: G = Graph([1, 2, 3], [(1, 2), (2, 3)], layout="tree")