manim 0.18.0.post0__py3-none-any.whl → 0.19.0__py3-none-any.whl

manim/mobject/types/point_cloud_mobject.py

@@ -4,6 +4,8 @@ from __future__ import annotations
 
 __all__ = ["PMobject", "Mobject1D", "Mobject2D", "PGroup", "PointCloudDot", "Point"]
 
+from typing import TYPE_CHECKING
+
 import numpy as np
 
 from manim.mobject.opengl.opengl_compatibility import ConvertToOpenGL
@@ -17,12 +19,24 @@ from ...utils.color import (
     WHITE,
     YELLOW,
     ManimColor,
+    ParsableManimColor,
     color_gradient,
     color_to_rgba,
     rgba_to_color,
 )
 from ...utils.iterables import stretch_array_to_length
 
+__all__ = ["PMobject", "Mobject1D", "Mobject2D", "PGroup", "PointCloudDot", "Point"]
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import Any
+
+    import numpy.typing as npt
+    from typing_extensions import Self
+
+    from manim.typing import ManimFloat, Point3DLike, Vector3D
+
 
 class PMobject(Mobject, metaclass=ConvertToOpenGL):
     """A disc made of a cloud of Dots
@@ -53,19 +67,25 @@ class PMobject(Mobject, metaclass=ConvertToOpenGL):
 
     """
 
-    def __init__(self, stroke_width=DEFAULT_STROKE_WIDTH, **kwargs):
+    def __init__(self, stroke_width: int = DEFAULT_STROKE_WIDTH, **kwargs: Any) -> None:
         self.stroke_width = stroke_width
         super().__init__(**kwargs)
 
-    def reset_points(self):
+    def reset_points(self) -> Self:
         self.rgbas = np.zeros((0, 4))
         self.points = np.zeros((0, 3))
         return self
 
-    def get_array_attrs(self):
+    def get_array_attrs(self) -> list[str]:
         return super().get_array_attrs() + ["rgbas"]
 
-    def add_points(self, points, rgbas=None, color=None, alpha=1):
+    def add_points(
+        self,
+        points: npt.NDArray,
+        rgbas: npt.NDArray | None = None,
+        color: ParsableManimColor | None = None,
+        alpha: float = 1,
+    ) -> Self:
         """Add points.
 
         Points must be a Nx3 numpy array.
@@ -83,24 +103,26 @@ class PMobject(Mobject, metaclass=ConvertToOpenGL):
         self.rgbas = np.append(self.rgbas, rgbas, axis=0)
         return self
 
-    def set_color(self, color=YELLOW, family=True):
+    def set_color(
+        self, color: ParsableManimColor = YELLOW, family: bool = True
+    ) -> Self:
         rgba = color_to_rgba(color)
         mobs = self.family_members_with_points() if family else [self]
         for mob in mobs:
             mob.rgbas[:, :] = rgba
-        self.color = color
+        self.color = ManimColor.parse(color)
         return self
 
-    def get_stroke_width(self):
+    def get_stroke_width(self) -> int:
         return self.stroke_width
 
-    def set_stroke_width(self, width, family=True):
+    def set_stroke_width(self, width: int, family: bool = True) -> Self:
         mobs = self.family_members_with_points() if family else [self]
         for mob in mobs:
             mob.stroke_width = width
         return self
 
-    def set_color_by_gradient(self, *colors):
+    def set_color_by_gradient(self, *colors: ParsableManimColor) -> Self:
         self.rgbas = np.array(
             list(map(color_to_rgba, color_gradient(*colors, len(self.points)))),
         )
@@ -108,11 +130,11 @@ class PMobject(Mobject, metaclass=ConvertToOpenGL):
 
     def set_colors_by_radial_gradient(
         self,
-        center=None,
-        radius=1,
-        inner_color=WHITE,
-        outer_color=BLACK,
-    ):
+        center: Point3DLike | None = None,
+        radius: float = 1,
+        inner_color: ParsableManimColor = WHITE,
+        outer_color: ParsableManimColor = BLACK,
+    ) -> Self:
         start_rgba, end_rgba = list(map(color_to_rgba, [inner_color, outer_color]))
         if center is None:
             center = self.get_center()
@@ -127,48 +149,48 @@ class PMobject(Mobject, metaclass=ConvertToOpenGL):
             )
         return self
 
-    def match_colors(self, mobject):
+    def match_colors(self, mobject: Mobject) -> Self:
         Mobject.align_data(self, mobject)
         self.rgbas = np.array(mobject.rgbas)
         return self
 
-    def filter_out(self, condition):
+    def filter_out(self, condition: npt.NDArray) -> Self:
         for mob in self.family_members_with_points():
             to_eliminate = ~np.apply_along_axis(condition, 1, mob.points)
             mob.points = mob.points[to_eliminate]
             mob.rgbas = mob.rgbas[to_eliminate]
         return self
 
-    def thin_out(self, factor=5):
-        """
-        Removes all but every nth point for n = factor
-        """
+    def thin_out(self, factor: int = 5) -> Self:
+        """Removes all but every nth point for n = factor"""
         for mob in self.family_members_with_points():
             num_points = self.get_num_points()
             mob.apply_over_attr_arrays(
-                lambda arr: arr[np.arange(0, num_points, factor)],
+                lambda arr, n=num_points: arr[np.arange(0, n, factor)],
             )
         return self
 
-    def sort_points(self, function=lambda p: p[0]):
-        """
-        Function is any map from R^3 to R
-        """
+    def sort_points(
+        self, function: Callable[[npt.NDArray[ManimFloat]], float] = lambda p: p[0]
+    ) -> Self:
+        """Function is any map from R^3 to R"""
         for mob in self.family_members_with_points():
             indices = np.argsort(np.apply_along_axis(function, 1, mob.points))
-            mob.apply_over_attr_arrays(lambda arr: arr[indices])
+            mob.apply_over_attr_arrays(lambda arr, idx=indices: arr[idx])
         return self
 
-    def fade_to(self, color, alpha, family=True):
+    def fade_to(
+        self, color: ParsableManimColor, alpha: float, family: bool = True
+    ) -> Self:
         self.rgbas = interpolate(self.rgbas, color_to_rgba(color), alpha)
         for mob in self.submobjects:
             mob.fade_to(color, alpha, family)
         return self
 
-    def get_all_rgbas(self):
+    def get_all_rgbas(self) -> npt.NDArray:
         return self.get_merged_array("rgbas")
 
-    def ingest_submobjects(self):
+    def ingest_submobjects(self) -> Self:
         attrs = self.get_array_attrs()
         arrays = list(map(self.get_merged_array, attrs))
         for attr, array in zip(attrs, arrays):
@@ -176,30 +198,32 @@ class PMobject(Mobject, metaclass=ConvertToOpenGL):
         self.submobjects = []
         return self
 
-    def get_color(self):
+    def get_color(self) -> ManimColor:
         return rgba_to_color(self.rgbas[0, :])
 
-    def point_from_proportion(self, alpha):
+    def point_from_proportion(self, alpha: float) -> Any:
         index = alpha * (self.get_num_points() - 1)
-        return self.points[index]
+        return self.points[np.floor(index)]
 
     @staticmethod
-    def get_mobject_type_class():
+    def get_mobject_type_class() -> type[PMobject]:
         return PMobject
 
     # Alignment
-    def align_points_with_larger(self, larger_mobject):
+    def align_points_with_larger(self, larger_mobject: Mobject) -> None:
         assert isinstance(larger_mobject, PMobject)
         self.apply_over_attr_arrays(
             lambda a: stretch_array_to_length(a, larger_mobject.get_num_points()),
         )
 
-    def get_point_mobject(self, center=None):
+    def get_point_mobject(self, center: Point3DLike | None = None) -> Point:
         if center is None:
             center = self.get_center()
         return Point(center)
 
-    def interpolate_color(self, mobject1, mobject2, alpha):
+    def interpolate_color(
+        self, mobject1: Mobject, mobject2: Mobject, alpha: float
+    ) -> Self:
         self.rgbas = interpolate(mobject1.rgbas, mobject2.rgbas, alpha)
         self.set_stroke_width(
             interpolate(
@@ -210,7 +234,7 @@ class PMobject(Mobject, metaclass=ConvertToOpenGL):
         )
         return self
 
-    def pointwise_become_partial(self, mobject, a, b):
+    def pointwise_become_partial(self, mobject: Mobject, a: float, b: float) -> None:
         lower_index, upper_index = (int(x * mobject.get_num_points()) for x in (a, b))
         for attr in self.get_array_attrs():
             full_array = getattr(mobject, attr)
@@ -220,24 +244,31 @@ class PMobject(Mobject, metaclass=ConvertToOpenGL):
 
 # TODO, Make the two implementations below non-redundant
 class Mobject1D(PMobject, metaclass=ConvertToOpenGL):
-    def __init__(self, density=DEFAULT_POINT_DENSITY_1D, **kwargs):
+    def __init__(self, density: int = DEFAULT_POINT_DENSITY_1D, **kwargs: Any) -> None:
         self.density = density
         self.epsilon = 1.0 / self.density
         super().__init__(**kwargs)
 
-    def add_line(self, start, end, color=None):
+    def add_line(
+        self,
+        start: npt.NDArray,
+        end: npt.NDArray,
+        color: ParsableManimColor | None = None,
+    ) -> None:
         start, end = list(map(np.array, [start, end]))
         length = np.linalg.norm(end - start)
         if length == 0:
-            points = [start]
+            points = np.array([start])
         else:
             epsilon = self.epsilon / length
-            points = [interpolate(start, end, t) for t in np.arange(0, 1, epsilon)]
+            points = np.array(
+                [interpolate(start, end, t) for t in np.arange(0, 1, epsilon)]
+            )
         self.add_points(points, color=color)
 
 
 class Mobject2D(PMobject, metaclass=ConvertToOpenGL):
-    def __init__(self, density=DEFAULT_POINT_DENSITY_2D, **kwargs):
+    def __init__(self, density: int = DEFAULT_POINT_DENSITY_2D, **kwargs: Any) -> None:
         self.density = density
         self.epsilon = 1.0 / self.density
         super().__init__(**kwargs)
@@ -266,7 +297,7 @@ class PGroup(PMobject):
 
     """
 
-    def __init__(self, *pmobs, **kwargs):
+    def __init__(self, *pmobs: Any, **kwargs: Any) -> None:
         if not all(isinstance(m, (PMobject, OpenGLPMobject)) for m in pmobs):
             raise ValueError(
                 "All submobjects must be of type PMobject or OpenGLPMObject"
@@ -275,10 +306,13 @@ class PGroup(PMobject):
         super().__init__(**kwargs)
         self.add(*pmobs)
 
-    def fade_to(self, color, alpha, family=True):
+    def fade_to(
+        self, color: ParsableManimColor, alpha: float, family: bool = True
+    ) -> Self:
         if family:
             for mob in self.submobjects:
                 mob.fade_to(color, alpha, family)
+        return self
 
 
 class PointCloudDot(Mobject1D):
@@ -315,13 +349,13 @@ class PointCloudDot(Mobject1D):
 
     def __init__(
         self,
-        center=ORIGIN,
-        radius=2.0,
-        stroke_width=2,
-        density=DEFAULT_POINT_DENSITY_1D,
-        color=YELLOW,
-        **kwargs,
-    ):
+        center: Vector3D = ORIGIN,
+        radius: float = 2.0,
+        stroke_width: int = 2,
+        density: int = DEFAULT_POINT_DENSITY_1D,
+        color: ManimColor = YELLOW,
+        **kwargs: Any,
+    ) -> None:
         self.radius = radius
         self.epsilon = 1.0 / density
         super().__init__(
@@ -329,22 +363,24 @@ class PointCloudDot(Mobject1D):
         )
         self.shift(center)
 
-    def init_points(self):
+    def init_points(self) -> None:
         self.reset_points()
         self.generate_points()
 
-    def generate_points(self):
+    def generate_points(self) -> None:
         self.add_points(
-            [
-                r * (np.cos(theta) * RIGHT + np.sin(theta) * UP)
-                for r in np.arange(self.epsilon, self.radius, self.epsilon)
-                # Num is equal to int(stop - start)/ (step + 1) reformulated.
-                for theta in np.linspace(
-                    0,
-                    2 * np.pi,
-                    num=int(2 * np.pi * (r + self.epsilon) / self.epsilon),
-                )
-            ],
+            np.array(
+                [
+                    r * (np.cos(theta) * RIGHT + np.sin(theta) * UP)
+                    for r in np.arange(self.epsilon, self.radius, self.epsilon)
+                    # Num is equal to int(stop - start)/ (step + 1) reformulated.
+                    for theta in np.linspace(
+                        0,
+                        2 * np.pi,
+                        num=int(2 * np.pi * (r + self.epsilon) / self.epsilon),
+                    )
+                ]
+            ),
         )
 
 
@@ -369,14 +405,16 @@ class Point(PMobject):
                 self.add(point)
     """
 
-    def __init__(self, location=ORIGIN, color=BLACK, **kwargs):
+    def __init__(
+        self, location: Vector3D = ORIGIN, color: ManimColor = BLACK, **kwargs: Any
+    ) -> None:
         self.location = location
         super().__init__(color=color, **kwargs)
 
-    def init_points(self):
+    def init_points(self) -> None:
         self.reset_points()
         self.generate_points()
         self.set_points([self.location])
 
-    def generate_points(self):
-        self.add_points([self.location])
+    def generate_points(self) -> None:
+        self.add_points(np.array([self.location]))