Plotting¶

def plot(
    self,
    function: Callable[[float], float],
    x_range: Sequence[float] | None = None,
    use_vectorized: bool = False,
    colorscale: Colorscale | None = None,
    colorscale_axis: int = 1,
    **kwargs: Unpack[CurveOptions],
) -> ParametricFunction:
    """Plot the graph of a function y = f(x) on the axes, in their coordinates.

    The function is sampled over `x_range` and the samples are joined smoothly (see
    [ParametricFunction][manimgx.ParametricFunction]); on a logarithmic x-axis, x
    runs over powers. The graph keeps the function as its `underlying_function`,
    for the methods that work on graphs. With a `colorscale`, the graph's stroke is
    a gradient, from left to right, of the colors of its values (y, or x), taken at
    every step of `x_range` (every 0.01 without one).

    Args:
        function: The function, from x to y, in the axes' coordinates.
        x_range: The range of x the graph spans, `[x_min, x_max]` or
            `[x_min, x_max, x_step]`, the function sampled every `x_step`; None
            for the x-axis's range. Without a step, it is a tenth of the x-axis's.
        use_vectorized: Whether `function` takes all the values of x at once, as an
            array; it is tried on an array anyway, and called once per value if it
            does not take one.
        colorscale: Colors to paint the graph with by its value: colors spread
            evenly over the axis's range, or `(color, value)` pairs, blended
            between; None for the graph's own color.
        colorscale_axis: The value the colorscale goes by: 1 for y, 0 for x.
        **kwargs: [Curve keywords][manimgx.mobjects.plotting.CurveOptions]:
            `discontinuities`, `use_smoothing`, and the style keywords.

    Returns:
        A new [ParametricFunction][manimgx.ParametricFunction], not added to the
        axes.

    Examples:
        ```python
        import numpy as np

        import manimgx as m


        class AxesPlotExample(m.Scene):
            def construct(self) -> None:
                axes = m.Axes(x_range=[-4, 4], y_range=[-2, 2], y_length=6)
                sine = axes.plot(np.sin, color=m.BLUE)
                parabola = axes.plot(
                    lambda x: x**2 - 1.5,
                    x_range=[-1.8, 1.8],
                    colorscale=[m.GREEN, m.YELLOW, m.RED],
                )
                self.add(axes)
                self.play(m.Create(sine), m.Create(parabola))
        ```
    """
    t_range = np.array(self.x_range, dtype=float)
    if x_range is not None:
        t_range[: len(x_range)] = x_range
    if x_range is None or len(x_range) < 3:
        t_range[2] /= self.num_sampled_graph_points_per_tick

    def point(t: float | np.ndarray) -> np.ndarray:
        """The graph's point at t — or at every t of an array, mapped in one call."""
        if isinstance(t, np.ndarray):
            return self.coords_to_point(t, sample(function, t))
        return self.coords_to_point(t, function(t))

    graph = ParametricFunction(
        point,
        t_range=t_range,
        scaling=self.x_axis.scaling,
        use_vectorized=use_vectorized,
        **kwargs,
    )
    graph.underlying_function = function
    if colorscale:
        resolution = (
            x_range[2] if x_range is not None and len(x_range) == 3 else 0.01
        )
        graph.set_stroke(
            self._colorscale(
                function, colorscale, colorscale_axis, t_range, resolution
            )
        )
        graph.set_sheen_direction(RIGHT)
    return graph

def plot_line_graph(
    self,
    x_values: Iterable[float],
    y_values: Iterable[float],
    z_values: Iterable[float] | None = None,
    line_color: ParsableManimColor = PURE_YELLOW,
    add_vertex_dots: bool = True,
    vertex_dot_radius: float = DEFAULT_DOT_RADIUS,
    vertex_dot_style: Style | None = None,
    **kwargs: Unpack[Style],
) -> VDict:
    """Plot a line graph: straight segments through points given by their
    coordinates, with a dot at each point; bright yellow (`PURE_YELLOW`), with white
    dots, unless styled. Empty inputs produce an empty graph.

    Args:
        x_values: The points' x coordinates.
        y_values: Their y coordinates, as many.
        z_values: Their z coordinates, on three-dimensional axes; None for 0.
        line_color: The line's color, unless the keywords give a `color`.
        add_vertex_dots: Whether a dot marks each point.
        vertex_dot_radius: The dots' radius, in scene units.
        vertex_dot_style: [Style keywords][manimgx.drawing.paint.Style] for the dots;
            None for none.
        **kwargs: [Style keywords][manimgx.drawing.paint.Style] for the line.

    Returns:
        A new [VDict][manimgx.VDict], not added to the axes, with the line
        under `"line_graph"` and the dots, in a group, under `"vertex_dots"`.

    Examples:
        ```python
        import manimgx as m


        class AxesPlotLineGraphExample(m.Scene):
            def construct(self) -> None:
                axes = m.Axes(x_range=[0, 7], y_range=[0, 5], y_length=6)
                graph = axes.plot_line_graph(
                    x_values=[0, 1.5, 2, 2.8, 4, 6.25],
                    y_values=[1, 3, 2.25, 4, 2.5, 1.75],
                    line_color=m.ORANGE,
                    vertex_dot_radius=0.12,
                    vertex_dot_style={"color": m.PURPLE},
                    stroke_width=6,
                )
                self.add(axes.add_coordinates())
                self.play(m.Create(graph))
        ```
    """
    xs, ys = (
        np.array(list(x_values), dtype=float),
        np.array(list(y_values), dtype=float),
    )
    zs = (
        np.zeros(xs.shape)
        if z_values is None
        else np.array(list(z_values), dtype=float)
    )
    line_graph = VDict()
    if kwargs.get("color") is None:  # not given, None too
        kwargs["color"] = line_color
    graph = VMobject(**kwargs)
    vertices = self.coords_to_point(*np.array((xs, ys, zs))).T
    graph.set_points_as_corners(vertices)
    line_graph["line_graph"] = graph
    if add_vertex_dots:
        line_graph["vertex_dots"] = VGroup(
            *(
                Dot(
                    point=vertex.copy(),  # the dot retains its center
                    radius=vertex_dot_radius,
                    **(vertex_dot_style or Style()),
                )
                for vertex in vertices
            )
        )
    return line_graph

def plot_polar_graph(
    self,
    r_func: Callable[[float], float],
    theta_range: Sequence[float] | None = None,
    **kwargs: Unpack[CurveOptions],
) -> ParametricFunction:
    """Plot a curve in polar coordinates, r = r_func(θ), on the axes: for each
    angle θ of its range, the point in the direction θ from the origin, at the
    distance `r_func(θ)`.

    Args:
        r_func: The function, from the angle θ, in radians, to the distance r, in
            the axes' units.
        theta_range: The range of θ, `[θ_min, θ_max]` or `[θ_min, θ_max, θ_step]`,
            in radians, sampled every 0.01 without a step; None for a full turn,
            [0, 2π].
        **kwargs: [Curve keywords][manimgx.mobjects.plotting.CurveOptions]:
            `use_smoothing`, and the style keywords.

    Returns:
        A new [ParametricFunction][manimgx.ParametricFunction], not added to the
        axes.

    Examples:
        ```python
        import numpy as np

        import manimgx as m


        class AxesPlotPolarGraphExample(m.Scene):
            def construct(self) -> None:
                plane = m.PolarPlane(radius_max=3, size=7.5)
                rose = plane.plot_polar_graph(
                    lambda theta: 3 * np.sin(5 * theta), color=m.ORANGE
                )
                self.add(plane)
                self.play(m.Create(rose), run_time=3)
        ```
    """
    theta_range = theta_range if theta_range is not None else [0, 2 * PI]
    return ParametricFunction(
        function=lambda th: self.pr2pt(r_func(th), th),
        t_range=theta_range,
        **kwargs,
    )

def plot_implicit_curve(
    self,
    func: Callable[[float, float], float],
    min_depth: int = 5,
    max_quads: int = 1500,
    **kwargs: Unpack[ImplicitOptions],
) -> ImplicitFunction:
    """Plot the curve where a function of x and y is zero, `func(x, y) = 0`, on the
    axes, in their coordinates: found within their ranges.

    The rectangle of the ranges is divided in four, and each part in four again,
    finer where the curve passes (see [ImplicitFunction][manimgx.ImplicitFunction]).
    On a logarithmic axis, the function is given powers.

    Args:
        func: The function of x and y, in the axes' coordinates.
        min_depth: How many times the rectangle is at least divided in four: more
            finds smaller pieces of the curve.
        max_quads: The most cells the rectangle may be divided into: more follow
            the curve more closely, and take longer.
        **kwargs: [Implicit curve keywords][manimgx.mobjects.plotting.ImplicitOptions]:
            `use_smoothing`, and the style keywords.

    Returns:
        A new [ImplicitFunction][manimgx.ImplicitFunction], not added to the axes.

    Examples:
        ```python
        import manimgx as m


        class AxesPlotImplicitCurveExample(m.Scene):
            def construct(self) -> None:
                axes = m.Axes(x_range=[-3, 3], y_range=[-3, 3], y_length=7)
                curve = axes.plot_implicit_curve(
                    lambda x, y: y**2 - x**3 + 2 * x - 1, color=m.YELLOW
                )
                self.add(axes)
                self.play(m.Create(curve))
        ```
    """
    x_scale = self.get_x_axis().scaling
    y_scale = self.get_y_axis().scaling
    graph = ImplicitFunction(
        func=lambda x, y: func(x_scale.function(x), y_scale.function(y)),
        x_range=self.x_range[:2],
        y_range=self.y_range[:2],
        min_depth=min_depth,
        max_quads=max_quads,
        **kwargs,
    )
    graph.stretch(self.get_x_unit_size(), 0, about_point=ORIGIN).stretch(
        self.get_y_unit_size(), 1, about_point=ORIGIN
    ).shift(self.get_origin())
    return graph

def plot_derivative_graph(
    self, graph: ParametricFunction, **kwargs: Unpack[PlotOptions]
) -> ParametricFunction:
    """Plot the derivative of a graph: at each x, its
    [slope][manimgx.Axes.slope_of_tangent]; green unless styled.

    Args:
        **kwargs: [Plot keywords][manimgx.mobjects.plotting.PlotOptions]:
            `x_range`, `color`, ….

    Returns:
        A new [ParametricFunction][manimgx.ParametricFunction], not added to the
        axes.

    Examples:
        ```python
        import manimgx as m


        class AxesPlotDerivativeGraphExample(m.Scene):
            def construct(self) -> None:
                axes = m.Axes(x_range=[-3, 3], y_range=[-3, 4], y_length=6.5)
                cubic = axes.plot(lambda x: x**3 / 6 - x, color=m.BLUE)
                derivative = axes.plot_derivative_graph(cubic)
                labels = m.VGroup(
                    axes.get_graph_label(cubic, "f(x)", x_val=-1.4, direction=m.UP),
                    axes.get_graph_label(derivative, "f'(x)", x_val=2.5),
                )
                self.add(axes, cubic, labels[0])
                self.play(m.Create(derivative), m.FadeIn(labels[1]))
        ```
    """
    if kwargs.get("color") is None:  # not given, None too
        kwargs["color"] = GREEN
    return self.plot(lambda x: self.slope_of_tangent(x, graph), **kwargs)

def plot_antiderivative_graph(
    self,
    graph: ParametricFunction,
    y_intercept: float = 0,
    samples: int = 50,
    **kwargs: Unpack[PlotOptions],
) -> ParametricFunction:
    """Plot an antiderivative of a graph: at each x, the integral of the graph from
    0 to x, plus `y_intercept`.

    Each value is a trapezoidal sum of the graph's values at `samples` points from
    0 to x.

    Args:
        y_intercept: The antiderivative's value at 0.
        samples: How many of the graph's values each integral sums.
        **kwargs: [Plot keywords][manimgx.mobjects.plotting.PlotOptions]:
            `x_range`, `color`, ….

    Returns:
        A new [ParametricFunction][manimgx.ParametricFunction], not added to the
        axes.

    Examples:
        ```python
        import manimgx as m


        class AxesPlotAntiderivativeGraphExample(m.Scene):
            def construct(self) -> None:
                axes = m.Axes(x_range=[-3, 3], y_range=[-3, 3], y_length=6.5)
                graph = axes.plot(lambda x: (x**2 - 2) / 3, color=m.RED)
                antiderivative = axes.plot_antiderivative_graph(
                    graph, y_intercept=1, color=m.BLUE
                )
                self.add(axes, graph)
                self.play(m.Create(antiderivative))
        ```
    """
    axis = 1 if kwargs.get("use_vectorized", False) else 0
    f_vec = np.vectorize(
        graph.underlying_function
        or (lambda x: self.input_to_graph_coords(x, graph)[1])
    )

    def antideriv(x: float) -> float:
        x_vals = np.linspace(0, x, samples, axis=axis)
        return float(np.trapezoid(f_vec(x_vals), x_vals) + y_intercept)

    return self.plot(antideriv, **kwargs)

def input_to_graph_point(self, x: float, graph: ParametricFunction) -> Point3D:
    """Find the point of a graph at an input: at x, for a graph of
    [plot][manimgx.Axes.plot]; at a value of its parameter, for another curve.

    `i2gp` is its short name. The input may be past the graph's ends: the point is
    then its function's, beyond what is drawn.

    Args:
        x: The input: x, in the axes' coordinates, for a graph.
        graph: The graph, or curve.

    Returns:
        The point, in scene coordinates.

    Examples:
        ```python
        import numpy as np

        import manimgx as m


        class AxesInputToGraphPointExample(m.Scene):
            def construct(self) -> None:
                axes = m.Axes(x_range=[0, 7], y_range=[-1.5, 1.5, 0.5])
                curve = axes.plot(np.cos, color=m.BLUE)
                dots = m.VGroup(
                    *(m.Dot(axes.i2gp(x, curve), radius=0.12) for x in range(7))
                )
                self.add(axes, curve, dots.set_color(m.YELLOW))
        ```
    """
    return graph.function(x)

def input_to_graph_coords(
    self, x: float, graph: ParametricFunction
) -> tuple[float, float]:
    """Find the coordinates of a graph's point at an input: (x, f(x)), for a graph
    of [plot][manimgx.Axes.plot]; for another curve, those of its point at a value
    of its parameter.

    Args:
        x: The input: x, for a graph.
        graph: The graph, or curve.

    Returns:
        The point's coordinates, (x, y).
    """
    if (
        graph.underlying_function is None
    ):  # a curve, not a graph of y(x): its point's coordinates
        coords = self.point_to_coords(graph.function(x))
        return float(coords[0]), float(coords[1])
    return (x, graph.underlying_function(x))

def get_graph_label(
    self,
    graph: ParametricFunction,
    label: Label = "f(x)",
    x_val: float | None = None,
    direction: Vector3DLike = RIGHT,
    buff: float = MED_SMALL_BUFF,
    color: ParsableManimColor | None = None,
    dot: bool = False,
    dot_config: Style | None = None,
) -> Mobject:
    r"""Make a label for a graph, by one of its points: in the graph's color, unless
    given another.

    A string or a number is typeset with the axes' `label_constructor` (as math,
    with [MathTex][manimgx.MathTex], by default); a mobject is used as it is, and
    recolored too. The label is put next to the graph's point at `x_val`, toward
    `direction`, then moved onto the screen if it is off it. Without `x_val`, the
    point is the rightmost, of a hundred over the x-axis's range, that is below the
    top of the frame.

    Args:
        label: The label: a string or a number, typeset as math, or a mobject.
        x_val: The x of the point the label is put by; None for the rightmost on
            screen.
        direction: The direction the label is put toward, from the point.
        buff: The gap between the point and the label, in scene units.
        color: The label's color; None for the graph's.
        dot: Whether a dot marks the point, as part of the label.
        dot_config: [Style keywords][manimgx.drawing.paint.Style] for the dot; None
            for none.

    Returns:
        The label, not added to the axes.

    Examples:
        ```python
        import numpy as np

        import manimgx as m


        class AxesGetGraphLabelExample(m.Scene):
            def construct(self) -> None:
                axes = m.Axes(x_range=[-4, 4], y_range=[-2, 4], x_length=11)
                parabola = axes.plot(lambda x: x**2 / 4, color=m.BLUE)
                sine = axes.plot(np.sin, color=m.YELLOW)
                parabola_label = axes.get_graph_label(parabola, "x^2 / 4")
                sine_label = axes.get_graph_label(
                    sine, r"\sin x", x_val=m.PI / 2, direction=m.UP, dot=True
                )
                self.add(axes, parabola, sine, parabola_label, sine_label)
        ```
    """
    label_object: Mobject = self.x_axis._create_label_tex(label).set_color(
        graph.get_color() if color is None else color
    )
    if x_val is None:  # the rightmost point of the graph that is on screen
        xs = np.linspace(self.x_range[1], self.x_range[0], 100)
        points = (self.input_to_graph_point(x, graph) for x in xs)
        point = next(
            (p for p in points if p[1] < config["frame_y_radius"]),
            self.input_to_graph_point(xs[-1], graph),
        )
    else:
        point = self.input_to_graph_point(x_val, graph)
    label_object.next_to(point, direction, buff=buff)
    label_object.shift_onto_screen()
    if dot:
        label_object.add(Dot(point=point, **(dot_config or Style())))
    return label_object

def get_T_label(
    self,
    x_val: float,
    graph: ParametricFunction,
    label: float | str | Mobject | None = None,
    label_color: ParsableManimColor | None = None,
    triangle_size: float = MED_SMALL_BUFF,
    triangle_color: ParsableManimColor | None = WHITE,
    line_func: type[Line] = Line,
    line_color: ParsableManimColor = PURE_YELLOW,
) -> VGroup:
    """Make a T-label for a value of x: a small triangle under the x-axis pointing
    up at the value, a label below the triangle, and a line from the x-axis up to
    the graph.

    Args:
        x_val: The value of x it marks.
        graph: The graph the line reaches.
        label: The label below the triangle: a string or a number, typeset as math,
            or a mobject; None for none.
        label_color: The label's color; None to leave it as it is.
        triangle_size: The triangle's height, in scene units.
        triangle_color: The triangle's color.
        line_func: The class of the line: [Line][manimgx.Line] (solid),
            [DashedLine][manimgx.DashedLine], ….
        line_color: The line's color.

    Returns:
        A new group of the label (if any), the triangle and the line; not added to
        the axes.

    Examples:
        ```python
        import numpy as np

        import manimgx as m


        class AxesGetTLabelExample(m.Scene):
            def construct(self) -> None:
                axes = m.Axes(x_range=[0, 10], y_range=[0, 10], y_length=6)
                graph = axes.plot(lambda x: 3 * np.sqrt(x), color=m.BLUE)
                t_label = axes.get_T_label(4, graph, label=m.MathTex("x = 4"))
                self.add(axes, graph)
                self.play(m.FadeIn(t_label))
        ```
    """
    T_label_group = VGroup()
    triangle = RegularPolygon(n=3, start_angle=np.pi / 2, stroke_width=0).set_fill(
        color=triangle_color, opacity=1
    )
    triangle.height = triangle_size
    triangle.move_to(self.coords_to_point(x_val, 0), UP)
    if label is not None:
        t_label = self.x_axis._create_label_tex(label)
        if label_color is not None:
            t_label.set_color(label_color)
        t_label.next_to(triangle, DOWN)
        T_label_group.add(t_label)
    v_line = self.get_vertical_line(
        self.i2gp(x_val, graph), color=line_color, line_func=line_func
    )
    T_label_group.add(triangle, v_line)
    return T_label_group

def get_vertical_lines_to_graph(
    self,
    graph: ParametricFunction,
    x_range: Sequence[float] | None = None,
    num_lines: int = 20,
    **kwargs: Unpack[AxisLine],
) -> VGroup:
    """Make lines from the x-axis to a graph, evenly spaced over a range of x:
    dashed, white and thin unless styled.

    Args:
        x_range: The range of x, `[x_min, x_max]`: the first line at x_min, the
            last at x_max; None for the x-axis's range.
        num_lines: How many lines.
        **kwargs: [Axis line keywords][manimgx.mobjects.plotting.AxisLine]:
            the lines' class, keywords, `color` and `stroke_width`.

    Returns:
        A new group of the lines, from left to right; not added to the axes.

    Examples:
        ```python
        import numpy as np

        import manimgx as m


        class AxesGetVerticalLinesToGraphExample(m.Scene):
            def construct(self) -> None:
                axes = m.Axes(x_range=[0, 8], y_range=[-1, 1, 0.5], y_length=5)
                curve = axes.plot(
                    lambda x: np.sin(2 * x) * np.exp(-x / 4), color=m.YELLOW
                )
                lines = axes.get_vertical_lines_to_graph(
                    curve, x_range=[0.5, 7.5], num_lines=15, color=m.BLUE
                )
                self.add(axes, curve)
                self.play(m.Create(lines))
        ```
    """
    x_range = x_range if x_range is not None else self.x_range
    return VGroup(
        *(
            self.get_vertical_line(self.i2gp(x, graph), **kwargs)
            for x in np.linspace(x_range[0], x_range[1], num_lines)
        )
    )

def get_area(
    self,
    graph: ParametricFunction,
    x_range: Sequence[float] | None = None,
    color: ParsableManimColor | Iterable[ParsableManimColor] = (BLUE, GREEN),
    opacity: float = 0.3,
    bounded_graph: ParametricFunction | None = None,
    **kwargs: Unpack[StyleBase],
) -> Polygon:
    """Make the area between a graph and the x-axis, or between two graphs, over a
    range of x: a polygon, blue to green and translucent unless styled.

    The area's edge follows the graph's own points, and closes along the x-axis, or
    back along `bounded_graph`; its fill and outline both take `color` and
    `opacity`. Two graphs whose ranges do not meet raise a ValueError.

    Args:
        x_range: The range of x, `(x_min, x_max)`; None for the graph's.
        color: The area's color, or colors for a gradient.
        opacity: The area's opacity, from 0 to 1.
        bounded_graph: A second graph the area reaches to, rather than the x-axis;
            None for the x-axis. Only the range both graphs span is covered.
        **kwargs: [Style keywords][manimgx.drawing.paint.Style] for the polygon, but
            its colors and opacities: `stroke_width`, ….

    Returns:
        A new [Polygon][manimgx.Polygon], not added to the axes.

    Examples:
        ```python
        import manimgx as m


        class AxesGetAreaExample(m.Scene):
            def construct(self) -> None:
                axes = m.Axes(x_range=[0, 5], y_range=[0, 6], y_length=6.5)
                curve = axes.plot(lambda x: 5 - (x - 2.5) ** 2 / 2, color=m.BLUE)
                line = axes.plot(lambda x: x / 2, color=m.YELLOW)
                under = axes.get_area(curve, x_range=(0.5, 2))
                between = axes.get_area(
                    curve, x_range=(3, 4.5), bounded_graph=line, color=m.RED
                )
                self.add(axes, under, between, curve, line)
        ```
    """
    if x_range is None:
        a = graph.t_min
        b = graph.t_max
    else:
        a, b = x_range
    if bounded_graph is not None:
        if bounded_graph.t_min > b:
            raise ValueError(f"Ranges not matching: {bounded_graph.t_min} < {b}")
        if bounded_graph.t_max < a:
            raise ValueError(f"Ranges not matching: {bounded_graph.t_max} > {a}")
        a = max(a, bounded_graph.t_min)
        b = min(b, bounded_graph.t_max)
    if bounded_graph is None:
        points = (
            [self.c2p(a), graph.function(a)]
            + [p for p in graph.points if a <= self.p2c(p)[0] <= b]
            + [graph.function(b), self.c2p(b)]
        )
    else:
        graph_points, bounded_graph_points = (
            [g.function(a)]
            + [p for p in g.points if a <= self.p2c(p)[0] <= b]
            + [g.function(b)]
            for g in (graph, bounded_graph)
        )
        points = graph_points + bounded_graph_points[::-1]
    return Polygon(*points, **kwargs).set_opacity(opacity).set_color(color)