manim 0.17.0__py3-none-any.whl → 0.19.1__py3-none-any.whl

manim/mobject/graphing/functions.py

@@ -5,7 +5,8 @@ from __future__ import annotations
 __all__ = ["ParametricFunction", "FunctionGraph", "ImplicitFunction"]
 
 
-from typing import Callable, Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
+from typing import TYPE_CHECKING
 
 import numpy as np
 from isosurfaces import plot_isoline
@@ -14,6 +15,15 @@ from manim import config
 from manim.mobject.graphing.scale import LinearBase, _ScaleBase
 from manim.mobject.opengl.opengl_compatibility import ConvertToOpenGL
 from manim.mobject.types.vectorized_mobject import VMobject
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    from typing_extensions import Self
+
+    from manim.typing import Point3D, Point3DLike
+    from manim.utils.color import ParsableManimColor
+
 from manim.utils.color import YELLOW
 
 
@@ -23,9 +33,9 @@ class ParametricFunction(VMobject, metaclass=ConvertToOpenGL):
     Parameters
     ----------
     function
-        The function to be plotted in the form of ``(lambda x: x**2)``
+        The function to be plotted in the form of ``(lambda t: (x(t), y(t), z(t)))``
     t_range
-        Determines the length that the function spans. By default ``[0, 1]``
+        Determines the length that the function spans in the form of (t_min, t_max, step=0.01). By default ``[0, 1]``
     scaling
         Scaling class applied to the points of the function. Default of :class:`~.LinearBase`.
     use_smoothing
@@ -49,10 +59,10 @@ class ParametricFunction(VMobject, metaclass=ConvertToOpenGL):
 
         class PlotParametricFunction(Scene):
             def func(self, t):
-                return np.array((np.sin(2 * t), np.sin(3 * t), 0))
+                return (np.sin(2 * t), np.sin(3 * t), 0)
 
             def construct(self):
-                func = ParametricFunction(self.func, t_range = np.array([0, TAU]), fill_opacity=0).set_color(RED)
+                func = ParametricFunction(self.func, t_range = (0, TAU), fill_opacity=0).set_color(RED)
                 self.add(func.scale(3))
 
     .. manim:: ThreeDParametricSpring
@@ -61,11 +71,11 @@ class ParametricFunction(VMobject, metaclass=ConvertToOpenGL):
         class ThreeDParametricSpring(ThreeDScene):
             def construct(self):
                 curve1 = ParametricFunction(
-                    lambda u: np.array([
+                    lambda u: (
                         1.2 * np.cos(u),
                         1.2 * np.sin(u),
                         u * 0.05
-                    ]), color=RED, t_range = np.array([-3*TAU, 5*TAU, 0.01])
+                    ), color=RED, t_range = (-3*TAU, 5*TAU, 0.01)
                 ).set_shade_in_3d(True)
                 axes = ThreeDAxes()
                 self.add(axes, curve1)
@@ -97,19 +107,22 @@ class ParametricFunction(VMobject, metaclass=ConvertToOpenGL):
 
     def __init__(
         self,
-        function: Callable[[float, float], float],
-        t_range: Sequence[float] | None = None,
+        function: Callable[[float], Point3DLike],
+        t_range: tuple[float, float] | tuple[float, float, float] = (0, 1),
         scaling: _ScaleBase = LinearBase(),
         dt: float = 1e-8,
         discontinuities: Iterable[float] | None = None,
         use_smoothing: bool = True,
         use_vectorized: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ):
-        self.function = function
-        t_range = [0, 1, 0.01] if t_range is None else t_range
+        def internal_parametric_function(t: float) -> Point3D:
+            """Wrap ``function``'s output inside a NumPy array."""
+            return np.asarray(function(t))
+
+        self.function = internal_parametric_function
         if len(t_range) == 2:
-            t_range = np.array([*t_range, 0.01])
+            t_range = (*t_range, 0.01)
 
         self.scaling = scaling
 
@@ -121,26 +134,25 @@ class ParametricFunction(VMobject, metaclass=ConvertToOpenGL):
 
         super().__init__(**kwargs)
 
-    def get_function(self):
+    def get_function(self) -> Callable[[float], Point3D]:
         return self.function
 
-    def get_point_from_function(self, t):
+    def get_point_from_function(self, t: float) -> Point3D:
         return self.function(t)
 
-    def generate_points(self):
-
+    def generate_points(self) -> Self:
         if self.discontinuities is not None:
             discontinuities = filter(
                 lambda t: self.t_min <= t <= self.t_max,
                 self.discontinuities,
             )
-            discontinuities = np.array(list(discontinuities))
+            discontinuities_array = np.array(list(discontinuities))
             boundary_times = np.array(
                 [
                     self.t_min,
                     self.t_max,
-                    *(discontinuities - self.dt),
-                    *(discontinuities + self.dt),
+                    *(discontinuities_array - self.dt),
+                    *(discontinuities_array + self.dt),
                 ],
             )
             boundary_times.sort()
@@ -159,7 +171,7 @@ class ParametricFunction(VMobject, metaclass=ConvertToOpenGL):
                 x, y, z = self.function(t_range)
                 if not isinstance(z, np.ndarray):
                     z = np.zeros_like(x)
-                points = np.stack(zip(x, y, z), axis=0)
+                points = np.stack([x, y, z], axis=1)
             else:
                 points = np.array([self.function(t) for t in t_range])
 
@@ -170,7 +182,8 @@ class ParametricFunction(VMobject, metaclass=ConvertToOpenGL):
             self.make_smooth()
         return self
 
-    init_points = generate_points
+    def init_points(self) -> None:
+        self.generate_points()
 
 
 class FunctionGraph(ParametricFunction):
@@ -202,20 +215,27 @@ class FunctionGraph(ParametricFunction):
                 self.add(cos_func, sin_func_1, sin_func_2)
     """
 
-    def __init__(self, function, x_range=None, color=YELLOW, **kwargs):
-
+    def __init__(
+        self,
+        function: Callable[[float], Any],
+        x_range: tuple[float, float] | tuple[float, float, float] | None = None,
+        color: ParsableManimColor = YELLOW,
+        **kwargs: Any,
+    ) -> None:
         if x_range is None:
-            x_range = np.array([-config["frame_x_radius"], config["frame_x_radius"]])
+            x_range = (-config["frame_x_radius"], config["frame_x_radius"])
 
         self.x_range = x_range
-        self.parametric_function = lambda t: np.array([t, function(t), 0])
-        self.function = function
+        self.parametric_function: Callable[[float], Point3D] = lambda t: np.array(
+            [t, function(t), 0]
+        )
+        self.function = function  # type: ignore[assignment]
         super().__init__(self.parametric_function, self.x_range, color=color, **kwargs)
 
-    def get_function(self):
+    def get_function(self) -> Callable[[float], Any]:
         return self.function
 
-    def get_point_from_function(self, x):
+    def get_point_from_function(self, x: float) -> Point3D:
         return self.parametric_function(x)
 
 
@@ -228,7 +248,7 @@ class ImplicitFunction(VMobject, metaclass=ConvertToOpenGL):
         min_depth: int = 5,
         max_quads: int = 1500,
         use_smoothing: bool = True,
-        **kwargs,
+        **kwargs: Any,
     ):
         """An implicit function.
 
@@ -287,7 +307,7 @@ class ImplicitFunction(VMobject, metaclass=ConvertToOpenGL):
 
         super().__init__(**kwargs)
 
-    def generate_points(self):
+    def generate_points(self) -> Self:
         p_min, p_max = (
             np.array([self.x_range[0], self.y_range[0]]),
             np.array([self.x_range[1], self.y_range[1]]),
@@ -309,4 +329,5 @@ class ImplicitFunction(VMobject, metaclass=ConvertToOpenGL):
             self.make_smooth()
         return self
 
-    init_points = generate_points
+    def init_points(self) -> None:
+        self.generate_points()

manim/mobject/graphing/number_line.py

@@ -2,9 +2,22 @@
 
 from __future__ import annotations
 
+from manim.mobject.mobject import Mobject
+from manim.mobject.opengl.opengl_vectorized_mobject import OpenGLVMobject
+
 __all__ = ["NumberLine", "UnitInterval"]
 
-from typing import Iterable, Sequence
+
+from collections.abc import Callable, Iterable, Sequence
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    from typing_extensions import Self
+
+    from manim.mobject.geometry.tips import ArrowTip
+    from manim.typing import Point3D, Point3DLike, Vector3D
 
 import numpy as np
 
@@ -12,8 +25,9 @@ from manim import config
 from manim.constants import *
 from manim.mobject.geometry.line import Line
 from manim.mobject.graphing.scale import LinearBase, _ScaleBase
-from manim.mobject.text.numbers import DecimalNumber
+from manim.mobject.text.numbers import DecimalNumber, Integer
 from manim.mobject.text.tex_mobject import MathTex, Tex
+from manim.mobject.text.text_mobject import Text
 from manim.mobject.types.vectorized_mobject import VGroup, VMobject
 from manim.utils.bezier import interpolate
 from manim.utils.config_ops import merge_dicts_recursively
@@ -49,6 +63,10 @@ class NumberLine(Line):
         The width of the tip.
     tip_height
         The height of the tip.
+    tip_shape
+        The mobject class used to construct the tip, or ``None`` (the
+        default) for the default arrow tip. Passed classes have to inherit
+        from :class:`.ArrowTip`.
     include_numbers
         Whether to add numbers to the tick marks. The number of decimal places is determined
         by the step size, this default can be overridden by ``decimal_number_config``.
@@ -140,17 +158,18 @@ class NumberLine(Line):
         include_tip: bool = False,
         tip_width: float = DEFAULT_ARROW_TIP_LENGTH,
         tip_height: float = DEFAULT_ARROW_TIP_LENGTH,
+        tip_shape: type[ArrowTip] | None = None,
         # numbers/labels
         include_numbers: bool = False,
         font_size: float = 36,
-        label_direction: Sequence[float] = DOWN,
-        label_constructor: VMobject = MathTex,
+        label_direction: Point3DLike = DOWN,
+        label_constructor: type[MathTex] = MathTex,
         scaling: _ScaleBase = LinearBase(),
         line_to_number_buff: float = MED_SMALL_BUFF,
         decimal_number_config: dict | None = None,
         numbers_to_exclude: Iterable[float] | None = None,
         numbers_to_include: Iterable[float] | None = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         # avoid mutable arguments in defaults
         if numbers_to_exclude is None:
@@ -173,8 +192,11 @@ class NumberLine(Line):
                 "num_decimal_places": self._decimal_places_from_step(x_range[2]),
             }
 
-        # turn into into an np array to scale by just applying the function
+        # turn into a NumPy array to scale by just applying the function
         self.x_range = np.array(x_range, dtype=float)
+        self.x_min: float
+        self.x_max: float
+        self.x_step: float
         self.x_min, self.x_max, self.x_step = scaling.function(self.x_range)
         self.length = length
         self.unit_size = unit_size
@@ -217,7 +239,11 @@ class NumberLine(Line):
         self.center()
 
         if self.include_tip:
-            self.add_tip(tip_length=self.tip_height, tip_width=self.tip_width)
+            self.add_tip(
+                tip_length=self.tip_height,
+                tip_width=self.tip_width,
+                tip_shape=tip_shape,
+            )
             self.tip.set_stroke(self.stroke_color, self.stroke_width)
 
         if self.include_ticks:
@@ -228,16 +254,16 @@ class NumberLine(Line):
             if self.scaling.custom_labels:
                 tick_range = self.get_tick_range()
 
+                custom_labels = self.scaling.get_custom_labels(
+                    tick_range,
+                    unit_decimal_places=decimal_number_config["num_decimal_places"],
+                )
+
                 self.add_labels(
                     dict(
                         zip(
                             tick_range,
-                            self.scaling.get_custom_labels(
-                                tick_range,
-                                unit_decimal_places=decimal_number_config[
-                                    "num_decimal_places"
-                                ],
-                            ),
+                            custom_labels,
                         )
                     ),
                 )
@@ -249,22 +275,28 @@ class NumberLine(Line):
                     font_size=self.font_size,
                 )
 
-    def rotate_about_zero(self, angle: float, axis: Sequence[float] = OUT, **kwargs):
+    def rotate_about_zero(
+        self, angle: float, axis: Vector3D = OUT, **kwargs: Any
+    ) -> Self:
         return self.rotate_about_number(0, angle, axis, **kwargs)
 
     def rotate_about_number(
-        self, number: float, angle: float, axis: Sequence[float] = OUT, **kwargs
-    ):
+        self, number: float, angle: float, axis: Vector3D = OUT, **kwargs: Any
+    ) -> Self:
         return self.rotate(angle, axis, about_point=self.n2p(number), **kwargs)
 
-    def add_ticks(self):
+    def add_ticks(self) -> None:
         """Adds ticks to the number line. Ticks can be accessed after creation
-        via ``self.ticks``."""
+        via ``self.ticks``.
+        """
         ticks = VGroup()
         elongated_tick_size = self.tick_size * self.longer_tick_multiple
+        elongated_tick_offsets = (
+            np.array(self.numbers_with_elongated_ticks) - self.x_min
+        )
         for x in self.get_tick_range():
             size = self.tick_size
-            if x in self.numbers_with_elongated_ticks:
+            if np.any(np.isclose(x - self.x_min, elongated_tick_offsets)):
                 size = elongated_tick_size
             ticks.add(self.get_tick(x, size))
         self.add(ticks)
@@ -327,6 +359,7 @@ class NumberLine(Line):
     def number_to_point(self, number: float | np.ndarray) -> np.ndarray:
         """Accepts a value along the number line and returns a point with
         respect to the scene.
+        Equivalent to `NumberLine @ number`
 
         Parameters
         ----------
@@ -347,7 +380,9 @@ class NumberLine(Line):
             array([0., 0., 0.])
             >>> number_line.number_to_point(1)
             array([1., 0., 0.])
-            >>> number_line.number_to_point([1,2,3])
+            >>> number_line @ 1
+            array([1., 0., 0.])
+            >>> number_line.number_to_point([1, 2, 3])
             array([[1., 0., 0.],
                    [2., 0., 0.],
                    [3., 0., 0.]])
@@ -379,42 +414,45 @@ class NumberLine(Line):
 
             >>> from manim import NumberLine
             >>> number_line = NumberLine()
-            >>> number_line.point_to_number((0,0,0))
-            0.0
-            >>> number_line.point_to_number((1,0,0))
-            1.0
-            >>> number_line.point_to_number([[0.5,0,0],[1,0,0],[1.5,0,0]])
+            >>> number_line.point_to_number((0, 0, 0))
+            np.float64(0.0)
+            >>> number_line.point_to_number((1, 0, 0))
+            np.float64(1.0)
+            >>> number_line.point_to_number([[0.5, 0, 0], [1, 0, 0], [1.5, 0, 0]])
             array([0.5, 1. , 1.5])
 
         """
         point = np.asarray(point)
         start, end = self.get_start_and_end()
         unit_vect = normalize(end - start)
-        proportion = np.dot(point - start, unit_vect) / np.dot(end - start, unit_vect)
+        proportion: float = np.dot(point - start, unit_vect) / np.dot(
+            end - start, unit_vect
+        )
         return interpolate(self.x_min, self.x_max, proportion)
 
-    def n2p(self, number: float | np.ndarray) -> np.ndarray:
+    def n2p(self, number: float | np.ndarray) -> Point3D:
         """Abbreviation for :meth:`~.NumberLine.number_to_point`."""
         return self.number_to_point(number)
 
-    def p2n(self, point: Sequence[float]) -> float:
+    def p2n(self, point: Point3DLike) -> float:
         """Abbreviation for :meth:`~.NumberLine.point_to_number`."""
         return self.point_to_number(point)
 
     def get_unit_size(self) -> float:
-        return self.get_length() / (self.x_range[1] - self.x_range[0])
+        val: float = self.get_length() / (self.x_range[1] - self.x_range[0])
+        return val
 
-    def get_unit_vector(self) -> np.ndarray:
+    def get_unit_vector(self) -> Vector3D:
         return super().get_unit_vector() * self.unit_size
 
     def get_number_mobject(
         self,
         x: float,
-        direction: Sequence[float] | None = None,
+        direction: Vector3D | None = None,
         buff: float | None = None,
         font_size: float | None = None,
-        label_constructor: VMobject | None = None,
-        **number_config,
+        label_constructor: type[MathTex] | None = None,
+        **number_config: dict[str, Any],
     ) -> VMobject:
         """Generates a positioned :class:`~.DecimalNumber` mobject
         generated according to ``label_constructor``.
@@ -439,7 +477,7 @@ class NumberLine(Line):
         :class:`~.DecimalNumber`
             The positioned mobject.
         """
-        number_config = merge_dicts_recursively(
+        number_config_merged = merge_dicts_recursively(
             self.decimal_number_config,
             number_config,
         )
@@ -453,16 +491,19 @@ class NumberLine(Line):
             label_constructor = self.label_constructor
 
         num_mob = DecimalNumber(
-            x, font_size=font_size, mob_class=label_constructor, **number_config
+            x,
+            font_size=font_size,
+            mob_class=label_constructor,
+            **number_config_merged,
         )
 
         num_mob.next_to(self.number_to_point(x), direction=direction, buff=buff)
         if x < 0 and self.label_direction[0] == 0:
             # Align without the minus sign
-            num_mob.shift(num_mob[0].get_width() * LEFT / 2)
+            num_mob.shift(num_mob[0].width * LEFT / 2)
         return num_mob
 
-    def get_number_mobjects(self, *numbers, **kwargs) -> VGroup:
+    def get_number_mobjects(self, *numbers: float, **kwargs: Any) -> VGroup:
         if len(numbers) == 0:
             numbers = self.default_numbers_to_display()
         return VGroup([self.get_number_mobject(number, **kwargs) for number in numbers])
@@ -475,9 +516,9 @@ class NumberLine(Line):
         x_values: Iterable[float] | None = None,
         excluding: Iterable[float] | None = None,
         font_size: float | None = None,
-        label_constructor: VMobject | None = None,
-        **kwargs,
-    ):
+        label_constructor: type[MathTex] | None = None,
+        **kwargs: Any,
+    ) -> Self:
         """Adds :class:`~.DecimalNumber` mobjects representing their position
         at each tick of the number line. The numbers can be accessed after creation
         via ``self.numbers``.
@@ -528,11 +569,11 @@ class NumberLine(Line):
     def add_labels(
         self,
         dict_values: dict[float, str | float | VMobject],
-        direction: Sequence[float] = None,
+        direction: Point3DLike | None = None,
         buff: float | None = None,
         font_size: float | None = None,
-        label_constructor: VMobject | None = None,
-    ):
+        label_constructor: type[MathTex] | None = None,
+    ) -> Self:
         """Adds specifically positioned labels to the :class:`~.NumberLine` using a ``dict``.
         The labels can be accessed after creation via ``self.labels``.
 
@@ -558,26 +599,24 @@ class NumberLine(Line):
         AttributeError
             If the label does not have a ``font_size`` attribute, an ``AttributeError`` is raised.
         """
-
         direction = self.label_direction if direction is None else direction
         buff = self.line_to_number_buff if buff is None else buff
         font_size = self.font_size if font_size is None else font_size
-        label_constructor = (
-            self.label_constructor if label_constructor is None else label_constructor
-        )
+        if label_constructor is None:
+            label_constructor = self.label_constructor
 
         labels = VGroup()
         for x, label in dict_values.items():
-
             # TODO: remove this check and ability to call
             # this method via CoordinateSystem.add_coordinates()
             # must be explicitly called
-            if isinstance(label, str) and self.label_constructor is MathTex:
+            if isinstance(label, str) and label_constructor is MathTex:
                 label = Tex(label)
             else:
-                label = self._create_label_tex(label)
+                label = self._create_label_tex(label, label_constructor)
 
             if hasattr(label, "font_size"):
+                assert isinstance(label, (MathTex, Tex, Text, Integer)), label
                 label.font_size = font_size
             else:
                 raise AttributeError(f"{label} is not compatible with add_labels.")
@@ -589,42 +628,61 @@ class NumberLine(Line):
         return self
 
     def _create_label_tex(
-        self, label_tex: str | float | VMobject, **kwargs
+        self,
+        label_tex: str | float | VMobject,
+        label_constructor: Callable | None = None,
+        **kwargs: Any,
     ) -> VMobject:
         """Checks if the label is a :class:`~.VMobject`, otherwise, creates a
-        label according to :attr:`label_constructor`.
+        label by passing ``label_tex`` to ``label_constructor``.
 
         Parameters
         ----------
         label_tex
-            The label to be compared against the above types.
+            The label for which a mobject should be created. If the label already
+            is a mobject, no new mobject is created.
+        label_constructor
+            Optional. A class or function returning a mobject when
+            passing ``label_tex`` as an argument. If ``None`` is passed
+            (the default), the label constructor from the :attr:`.label_constructor`
+            attribute is used.
 
         Returns
         -------
         :class:`~.VMobject`
             The label.
         """
-
-        if isinstance(label_tex, VMobject):
+        if isinstance(label_tex, (VMobject, OpenGLVMobject)):
             return label_tex
-        else:
-            return self.label_constructor(label_tex, **kwargs)
+        if label_constructor is None:
+            label_constructor = self.label_constructor
+        if isinstance(label_tex, str):
+            return label_constructor(label_tex, **kwargs)
+        return label_constructor(str(label_tex), **kwargs)
 
     @staticmethod
-    def _decimal_places_from_step(step) -> int:
-        step = str(step)
-        if "." not in step:
+    def _decimal_places_from_step(step: float) -> int:
+        step_str = str(step)
+        if "." not in step_str:
             return 0
-        return len(step.split(".")[-1])
+        return len(step_str.split(".")[-1])
+
+    def __matmul__(self, other: float) -> Point3D:
+        return self.n2p(other)
+
+    def __rmatmul__(self, other: Point3DLike | Mobject) -> float:
+        if isinstance(other, Mobject):
+            other = other.get_center()
+        return self.p2n(other)
 
 
 class UnitInterval(NumberLine):
     def __init__(
         self,
-        unit_size=10,
-        numbers_with_elongated_ticks=None,
-        decimal_number_config=None,
-        **kwargs,
+        unit_size: float = 10,
+        numbers_with_elongated_ticks: list[float] | None = None,
+        decimal_number_config: dict[str, Any] | None = None,
+        **kwargs: Any,
     ):
         numbers_with_elongated_ticks = (
             [0, 1]