Graph theory: Weighted edges, loop edges

.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

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:

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)

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"