Vector fields¶

def pos_to_color(self, pos: Point3D) -> ManimColor:
    """The field's color at a point."""
    return rgb_to_color(self.pos_to_rgb(pos))

@staticmethod
def shift_func(func: FieldFunction, shift_vector: np.ndarray) -> FieldFunction:
    """Move a field: make the function of the field shifted by a vector.

    Args:
        func: The field's function.
        shift_vector: The vector to move it by, in scene units.

    Returns:
        A new function, which at a point p is `func(p - shift_vector)`.
    """
    return lambda p: func(p - shift_vector)

@staticmethod
def scale_func(func: FieldFunction, scalar: float) -> FieldFunction:
    """Rescale a field's pattern: make the function of the field whose vector at a
    point p is the old field's at p × `scalar`.

    A scalar below 1 spreads the pattern out, as if zoomed in; above 1, it draws it
    together. The vectors keep their lengths.

    Args:
        func: The field's function.
        scalar: The factor the points are multiplied by.

    Returns:
        A new function, which at a point p is `func(p * scalar)`.
    """
    return lambda p: func(p * scalar)

def fit_to_coordinate_system(self, coordinate_system: Axes) -> Self:
    """Place the field on axes: move each point of its parts to the axes' point
    whose coordinates are the point's own.

    A field made in the axes' coordinates (its ranges theirs) lands on them, its
    arrows stretched with the axes' units.

    Args:
        coordinate_system: The axes.

    Examples:
        ```python
        import numpy as np

        import manimgx as m


        class VectorFieldFitToCoordinateSystemExample(m.Scene):
            def construct(self) -> None:
                axes = m.Axes(
                    x_range=[-2, 2], y_range=[-2, 2], x_length=7, y_length=7
                )
                field = m.ArrowVectorField(
                    lambda p: np.array([-p[1], p[0], 0]) / 2,
                    x_range=[-2, 2, 0.5],
                    y_range=[-2, 2, 0.5],
                )
                self.add(axes, field.fit_to_coordinate_system(axes))
        ```
    """
    self.apply_function(lambda pos: coordinate_system.coords_to_point(*pos))
    return self

def nudge(
    self, mob: Mobject, dt: float = 1, substeps: int = 1, pointwise: bool = False
) -> Self:
    """Carry a mobject along the field for a while, as a particle on its flow.

    The field's vectors are velocities: the mobject moves with the field at its
    center, or, with `pointwise`, each of its points moves with the field where it
    is, which bends it. Each step is a fourth-order Runge–Kutta step.

    Args:
        mob: The mobject to move.
        dt: How long it is carried, in the field's time; negative to carry it back.
        substeps: How many steps the time is divided into: more follow the field
            more closely.
        pointwise: Whether each point of the mobject moves on its own, rather than
            the whole mobject with its center.
    """

    def runge_kutta(p: Point3D, step_size: float) -> Vector3D:
        k_1 = self.func(p)
        k_2 = self.func(p + step_size * (k_1 * 0.5))
        k_3 = self.func(p + step_size * (k_2 * 0.5))
        k_4 = self.func(p + step_size * k_3)
        return step_size / 6.0 * (k_1 + 2.0 * k_2 + 2.0 * k_3 + k_4)

    step_size = dt / substeps
    for _ in range(substeps):
        if pointwise:
            mob.apply_function(lambda p: p + runge_kutta(p, step_size))
        else:
            mob.shift(runge_kutta(mob.get_center(), step_size))
    return self

def nudge_submobjects(
    self, dt: float = 1, substeps: int = 1, pointwise: bool = False
) -> Self:
    """Carry each of the field's own parts along it for a while, as
    [nudge][manimgx.VectorField.nudge] carries a mobject.

    Args:
        dt: How long they are carried, in the field's time; negative to carry them
            back.
        substeps: How many steps the time is divided into.
        pointwise: Whether each point of a part moves on its own, rather than the
            whole part with its center.
    """
    for mob in self.submobjects:
        self.nudge(mob, dt, substeps, pointwise)
    return self

def get_nudge_updater(
    self, speed: float = 1, pointwise: bool = False
) -> Callable[[Mobject, float], Mobject]:
    """Make an updater that carries a mobject along the field: added to a mobject,
    it moves it with the field every frame, by the time since the last.

    Args:
        speed: How fast the mobject is carried: the field's time per second.
        pointwise: Whether each point of the mobject moves on its own, rather than
            the whole mobject with its center.

    Returns:
        The updater, a function of the mobject and the time step, to give
        [add_updater][manimgx.Mobject.add_updater].

    Examples:
        ```python
        import numpy as np

        import manimgx as m


        class VectorFieldGetNudgeUpdaterExample(m.Scene):
            def construct(self) -> None:
                def swirl(p: np.ndarray) -> np.ndarray:
                    return np.sin(p[1] / 2) * m.RIGHT + np.cos(p[0] / 2) * m.UP

                field = m.ArrowVectorField(
                    swirl, x_range=[-7, 7, 1], y_range=[-4, 4, 1]
                )
                circle = m.Circle(radius=1, color=m.YELLOW).shift(2 * m.LEFT)
                dot = m.Dot(2 * m.RIGHT, radius=0.15, color=m.RED)
                circle.add_updater(field.get_nudge_updater(pointwise=True))
                dot.add_updater(field.get_nudge_updater())
                self.add(field, circle, dot)
                self.wait(2.5)
        ```
    """
    return lambda mob, dt: self.nudge(mob, dt * speed, pointwise=pointwise)

def start_submobject_movement(
    self, speed: float = 1, pointwise: bool = False
) -> Self:
    """Set the field's own parts moving along it: an updater carries them every
    frame, as [nudge_submobjects][manimgx.VectorField.nudge_submobjects] does.

    A movement already started is replaced.

    Args:
        speed: How fast they are carried: the field's time per second.
        pointwise: Whether each point of a part moves on its own, rather than the
            whole part with its center.
    """
    self.stop_submobject_movement()
    self.submob_movement_updater = lambda mob, dt: mob.nudge_submobjects(
        dt * speed, pointwise=pointwise
    )
    self.add_updater(self.submob_movement_updater)
    return self

def stop_submobject_movement(self) -> Self:
    """Stop the movement
    [start_submobject_movement][manimgx.VectorField.start_submobject_movement]
    started, if any.
    """
    if self.submob_movement_updater is not None:
        self.remove_updater(self.submob_movement_updater)
    self.submob_movement_updater = None
    return self

def get_vector(self, point: Point3D) -> Vector:
    """Make the field's arrow at a point: from the point, in the field's direction
    there, as long as `length_func` makes it, in the field's color there.

    Args:
        point: The point, in scene coordinates.

    Returns:
        A new [Vector][manimgx.Vector], not added to the field.
    """
    output = np.array(self.func(point), dtype=float)
    norm = np.linalg.norm(output)
    if norm != 0:
        output *= self.length_func(float(norm)) / norm
    vect = Vector(output, **self.vector_config)
    vect.shift(point)
    vect.set_color(self.color if self.single_color else self.pos_to_color(point))
    return vect

def create(self, **kwargs: Unpack[AnimationOptions]) -> AnimationGroup:
    """Make an animation that draws the lines in, one after another in a random
    order.

    Each line is drawn with [Create][manimgx.Create] over `run_time`, by default
    the field's `virtual_time`, each beginning `lag_ratio` of that after the one
    before: by default `run_time / 2` divided by the number of lines. The order is
    shuffled with Python's `random` (seed it for the same order every time).

    Args:
        **kwargs: [Animation options][manimgx.animation.timeline.AnimationOptions]
            for each line's animation, and `lag_ratio` for the whole.

    Returns:
        A new animation group.

    Examples:
        ```python
        import manimgx as m


        class StreamLinesCreateExample(m.Scene):
            def construct(self) -> None:
                stream_lines = m.StreamLines(
                    lambda p: (p[0] * m.UR + p[1] * m.LEFT) - p,
                    color=m.YELLOW,
                    x_range=[-7, 7, 1],
                    y_range=[-4, 4, 1],
                    stroke_width=3,
                    virtual_time=1,
                    max_anchors_per_line=6,
                )
                self.play(stream_lines.create())
        ```
    """
    run_time = kwargs.pop("run_time", self.virtual_time)
    lag_ratio = kwargs.pop("lag_ratio", run_time / 2 / len(self.submobjects))
    kwargs["run_time"] = run_time
    animations = [Create(line, **kwargs) for line in self.stream_lines]
    random.shuffle(animations)
    return AnimationGroup(*animations, lag_ratio=lag_ratio)

def start_animation(
    self,
    warm_up: bool = True,
    flow_speed: float = 1,
    time_width: float = 0.3,
    line_animation_class: type[ShowPassingFlash] = ShowPassingFlash,
    **kwargs: Unpack[Untimed],
) -> Self:
    """Set the lines flowing: each flashes along itself, over and over, while the
    scene plays or waits.

    An updater plays each line's animation — by default a
    [ShowPassingFlash][manimgx.ShowPassingFlash] — along the line again and again,
    a cycle every `virtual_time / flow_speed` seconds, each line at a phase of its
    own, drawn with Python's `random`.
    [end_animation][manimgx.StreamLines.end_animation] ends the flow.

    Args:
        warm_up: Whether each line waits, empty, until its own time to begin,
            rather than all flowing from the start.
        flow_speed: How fast the lines flow: the field's time per second.
        time_width: The length of each flash, as a fraction of its line.
        line_animation_class: The animation played along each line: a
            ShowPassingFlash, or a class of its kind.
        **kwargs: [Animation options][manimgx.animation.timeline.Untimed] for each
            line's animation, but `run_time`; `rate_func` is `linear` unless given.

    Examples:
        ```python
        import numpy as np

        import manimgx as m


        class StreamLinesStartAnimationExample(m.Scene):
            def construct(self) -> None:
                def func(p: np.ndarray) -> np.ndarray:
                    return np.sin(p[0] / 2) * m.UR + np.cos(p[1] / 2) * m.LEFT

                stream_lines = m.StreamLines(
                    func, stroke_width=3, max_anchors_per_line=7
                )
                self.add(stream_lines)
                stream_lines.start_animation(warm_up=False, flow_speed=1.5)
                self.wait(stream_lines.virtual_time / stream_lines.flow_speed)
        ```
    """
    kwargs.setdefault("rate_func", linear)
    for line in self.stream_lines:
        line.anim = line_animation_class(
            line,
            run_time=line.duration / flow_speed,
            time_width=time_width,
            **kwargs,
        )
        line.anim.begin()
        line.time = random.random() * self.virtual_time * (-1 if warm_up else 1)
        self.add(line.anim.mobject)

    def updater(mob: StreamLines, dt: float) -> None:
        for line in mob.stream_lines:
            line.time += dt * flow_speed
            if line.time >= mob.virtual_time:
                line.time -= mob.virtual_time
            line.anim.interpolate(
                float(np.clip(line.time / line.anim.run_time, 0, 1))
            )

    self.add_updater(flow(updater))  # each line at the time it has run
    self.flow_animation = updater
    self.flow_speed = flow_speed
    self.time_width = time_width
    return self

def end_animation(self) -> AnimationGroup:
    """Make an animation that ends the flow: each line finishes its flash, then is
    drawn in with [Create][manimgx.Create], easing out.

    The flow's updater is removed at once, and a line still waiting to begin stays
    hidden until its time. Called before
    [start_animation][manimgx.StreamLines.start_animation], it raises a
    ValueError.

    Warning:
        The lines end hidden, not whole: a finished flash leaves its line trimmed
        past its end, and the Create that follows draws the line only as it finds
        it.

    Returns:
        A new animation group.
    """
    if self.flow_animation is None:
        raise ValueError("You have to start the animation before fading it out.")

    def hide_and_wait(mob: Mobject, alpha: float) -> None:  # unseen until its turn
        mob.set_stroke(opacity=float(alpha >= 1))

    def finish_cycle(start: float) -> Callable[[StreamLine, float], None]:
        """The flow, on to the end of the line's cycle: its time at alpha (linear)."""

        def flow(line: StreamLine, alpha: float) -> None:
            line.time = start + alpha * (self.virtual_time - start)
            line.anim.interpolate(min(line.time / line.anim.run_time, 1))
            if alpha == 1:
                self.remove(line.anim.mobject)
                line.anim.finish()

        return flow

    max_run_time = self.virtual_time / self.flow_speed
    creation_rate_func = ease_out_sine
    creation_staring_speed = creation_rate_func(0.001) * 1000
    creation_run_time = (
        max_run_time / (1 + self.time_width) * creation_staring_speed
    )
    animations: list[Animation] = []
    self.remove_updater(self.flow_animation)
    self.flow_animation = None
    for line in self.stream_lines:
        create = Create(
            line, run_time=creation_run_time, rate_func=creation_rate_func
        )
        if line.time <= 0:
            animations.append(
                Succession(
                    UpdateFromAlphaFunc(
                        line, hide_and_wait, run_time=-line.time / self.flow_speed
                    ),
                    create,
                )
            )
            self.remove(line.anim.mobject)
            line.anim.finish()
        else:
            remaining_time = max_run_time - line.time / self.flow_speed
            animations.append(
                Succession(
                    UpdateFromAlphaFunc(
                        line,
                        finish_cycle(line.time),
                        run_time=remaining_time,
                        rate_func=linear,
                    ),
                    create,
                )
            )
    return AnimationGroup(*animations)