Using Manim to animate the Gauss length of a Bezier curve

BezierGaussLength.py

from manim import *
from manim.mobject.geometry import ArrowTriangleTip
import numpy as np

def quad_bezier(x_0, x_1, x_2, t):
    r = 1.0 - t
    return r * r * x_0 + 2 * r * t * x_1 + t * t * x_2

def quad_bezier_grad(x_0, x_1, x_2, t):
    r = 1.0 - t
    return 2.0 * (r * (x_1 - x_0) + t * (x_2 - x_1))

def quad_bezier_normal(x_0, x_1, x_2, t):
    tangent = quad_bezier_grad(x_0, x_1, x_2, t)
    return np.array([-tangent[1], tangent[0], 0.0])


## This is mostly copied from the Arrow class
## I wanted a dashed arrow, which I couldn't find supported
## in manim
class DashedArrow(DashedLine):
    def __init__(
        self,
        *args,
        stroke_width=3,
        buff=MED_SMALL_BUFF,
        max_tip_length_to_length_ratio=0.25,
        max_stroke_width_to_length_ratio=5,
        tip_length=0.25,
        **kwargs,
    ):
        self.max_tip_length_to_length_ratio = max_tip_length_to_length_ratio
        self.max_stroke_width_to_length_ratio = max_stroke_width_to_length_ratio
        tip_shape = kwargs.pop("tip_shape", ArrowTriangleTip)
        super().__init__(*args, buff=buff, stroke_width=stroke_width, **kwargs)
        # TODO, should this be affected when
        # Arrow.set_stroke is called?
        self.initial_stroke_width = self.stroke_width
        self.add_tip(tip_shape=tip_shape, tip_length=tip_length)
        self.set_stroke_width_from_length()
    
    def scale(self, factor, scale_tips=False, **kwargs):
        r"""Scale an arrow, but keep stroke width and arrow tip size fixed.

        See Also
        --------
        :meth:`~.Mobject.scale`

        Examples
        --------
        ::

            >>> arrow = Arrow(np.array([-1, -1, 0]), np.array([1, 1, 0]), buff=0)
            >>> scaled_arrow = arrow.scale(2)
            >>> np.round(scaled_arrow.get_start_and_end(), 8) + 0
            array([[-2., -2.,  0.],
                   [ 2.,  2.,  0.]])
            >>> arrow.tip.length == scaled_arrow.tip.length
            True

        Manually scaling the object using the default method
        :meth:`~.Mobject.scale` does not have the same properties::

            >>> new_arrow = Arrow(np.array([-1, -1, 0]), np.array([1, 1, 0]), buff=0)
            >>> another_scaled_arrow = VMobject.scale(new_arrow, 2)
            >>> another_scaled_arrow.tip.length == arrow.tip.length
            False

        """
        if self.get_length() == 0:
            return self

        if scale_tips:
            super().scale(factor, **kwargs)
            self.set_stroke_width_from_length()
            return self

        has_tip = self.has_tip()
        has_start_tip = self.has_start_tip()
        if has_tip or has_start_tip:
            old_tips = self.pop_tips()

        super().scale(factor, **kwargs)
        self.set_stroke_width_from_length()

        if has_tip:
            self.add_tip(tip=old_tips[0])
        if has_start_tip:
            self.add_tip(tip=old_tips[1], at_start=True)
        return self

    def get_normal_vector(self) -> np.ndarray:
        """Returns the normal of a vector.

        Examples
        --------
        ::

            >>> Arrow().get_normal_vector() + 0. # add 0. to avoid negative 0 in output
            array([ 0.,  0., -1.])
        """

        p0, p1, p2 = self.tip.get_start_anchors()[:3]
        return normalize(np.cross(p2 - p1, p1 - p0))

    def reset_normal_vector(self):
        """Resets the normal of a vector"""
        self.normal_vector = self.get_normal_vector()
        return self

    def get_default_tip_length(self) -> float:
        """Returns the default tip_length of the arrow.

        Examples
        --------

        ::

            >>> Arrow().get_default_tip_length()
            0.35
        """

        max_ratio = self.max_tip_length_to_length_ratio
        return min(self.tip_length, max_ratio * self.get_length())

    def set_stroke_width_from_length(self):
        """Used internally. Sets stroke width based on length."""
        max_ratio = self.max_stroke_width_to_length_ratio
        if config.renderer == "opengl":
            self.set_stroke(
                width=min(self.initial_stroke_width, max_ratio * self.get_length()),
                recurse=False,
            )
        else:
            self.set_stroke(
                width=min(self.initial_stroke_width, max_ratio * self.get_length()),
                family=False,
            )
        return self

class BezierCurveGaussLength(Scene):
    """
    Animate the Gauss length of a Bezier curve.
    """
    def construct(self):
        self.init_curve()
        self.move_dot_and_draw_curve()
        self.dot_back_to_start()
        self.add_circle()
        tangent, normed_tangent, sq_point = self.add_tangents()
        self.animate_gl_arc(tangent, normed_tangent, sq_point)
    
    def init_curve(self, curve_segments=10):
        self.x_0 = np.array([-3.0, -3.0, 0.0])
        self.x_1 = np.array([-2.0,  1.0, 0.0])
        self.x_2 = np.array( [0.0,  0.0, 0.0])

    def make_square_marker(self):
        return Square(side_length=0.1, color=BLUE, fill_color=BLUE, fill_opacity=1.0)
        
    
    def move_dot_and_draw_curve(self):
        dot = Dot(radius=0.08, color=YELLOW)
        self.dot = dot
        dot.move_to(self.x_0)
        self.t_offset = 0
        self._x = self.x_0
        rate = 0.75

        def advance_bezier(mob, dt):
            self.t_offset += (dt * rate)
            self.t_offset = np.clip(self.t_offset, 0.0, 1.0)
            self._x = quad_bezier(self.x_0, self.x_1, self.x_2, self.t_offset * self.t_offset)
            mob.move_to(self._x)

        def draw_function_label():
            normal = quad_bezier_normal(self.x_0, self.x_1, self.x_2, self.t_offset * self.t_offset)
            normal /= np.linalg.norm(normal)
            tex = MathTex(r"f({:.2f})".format(self.t_offset)).move_to(self._x + 1.0 * normal)
            return tex
        
        def update_f_label(mob, dt):
            tex = draw_function_label()
            mob.become(tex)
        
        f_label = draw_function_label()


        self.curve = VGroup()
        self.curve.add(Line(self.x_0, self.x_0))
        def update_curve(mob, dt):
            last_line = self.curve[-1]
            x = quad_bezier(self.x_0, self.x_1, self.x_2, self.t_offset * self.t_offset)
            new_line = Line(last_line.get_end(), self._x, color=BLUE_A)
            self.curve.add(new_line)
            mob.become(self.curve)

        dot.add_updater(advance_bezier)
        self.curve.add_updater(update_curve)
        f_label.add_updater(update_f_label)

        self.add(dot)
        self.add(self.curve, f_label)
        self.wait(1.8)
        dot.remove_updater(advance_bezier)
        f_label.remove_updater(update_f_label)
        self.curve.remove_updater(update_curve)

        self.play(ScaleInPlace(f_label, 1.1), run_time=0.2)
        self.play(ScaleInPlace(f_label, 0.0), run_time=0.3)
        self.wait(0.1)
    
    def dot_back_to_start(self):
        rate = 1.0
        def retreat_bezier(mob, dt):
            self.t_offset -= (dt * rate)
            self.t_offset = np.clip(self.t_offset, 0.0, 1.0)
            self._x = quad_bezier(self.x_0, self.x_1, self.x_2, self.t_offset * self.t_offset)
            mob.move_to(self._x)
        self.dot.add_updater(retreat_bezier)
        self.wait(1.0)
        self.dot.remove_updater(retreat_bezier)
    
    def add_circle(self):
        self.circle = Circle().move_to(self.dot.get_center())
        self.play(Create(self.circle))
    
    def compute_tangent_endpoints(self):
        tangent = quad_bezier_grad(self.x_0, self.x_1, self.x_2, self.t_offset * self.t_offset)
        return self._x, self._x + 0.5 * tangent
    
    def compute_normalized_tangent_endpoints(self):
        tangent = quad_bezier_grad(self.x_0, self.x_1, self.x_2, self.t_offset * self.t_offset)
        normed = tangent / np.linalg.norm(tangent)
        return self._x, self._x + normed
    
    def add_tangents(self):
        def draw_tangent():
            x, x_t = self.compute_tangent_endpoints()
            return DashedArrow(x, x_t, buff=0.2, color=RED_B, fill_opacity=0.4)

        def draw_normalized_tangent_line():
            x, p = self.compute_normalized_tangent_endpoints()
            sq = self.make_square_marker().move_to(p)
            return Line(x, p, buff=0.0, color=BLUE), sq
            
        normed_tangent, sq = draw_normalized_tangent_line()
        tangent = draw_tangent()
        x, xt = self.compute_tangent_endpoints()
        p = xt + 1.0 * (xt - x) / np.linalg.norm(xt - x)
        self.grad_label = MathTex(r"\nabla f").move_to(p)

        self.play(Create(tangent), run_time=0.8)
        self.play(Create(self.grad_label), run_time=0.5)
        self.play(Create(normed_tangent), Create(sq), run_time=0.3)
        self.wait(0.5)
        return tangent, normed_tangent, sq
    
    def animate_gl_arc(self, tangent, normed_tangent, sq_point):
        rate = 0.5
        st_sq = sq_point.get_center() - self.circle.get_center()

        def advance_bezier(mob, dt):
            self.t_offset += (dt * rate)
            self.t_offset = np.clip(self.t_offset, 0.0, 1.0)
            self._x = quad_bezier(self.x_0, self.x_1, self.x_2, self.t_offset * self.t_offset)
            mob.move_to(self._x)
        
        def advance_circle(mob, dt):
            c = self.dot.get_center()
            self.circle.move_to(c)
        
        def get_arc():
            z = sq_point.get_center() - self.circle.get_center()

            # We flip these here. Maybe a more general solution exists...
            start_angle = np.arctan2(st_sq[1], st_sq[0])
            end_angle = np.arctan2(z[1], z[0])
            return Arc(
                start_angle=start_angle,
                angle=end_angle - start_angle,
                radius=1.0,
                arc_center=self.circle.get_center(),
                color=BLUE
            )


        def draw_arc():
            arc_curve = get_arc()
            start_sq = self.make_square_marker().move_to(st_sq + self.circle.get_center())
            return VGroup(arc_curve, start_sq)
        
        def advance_tangent(mob, dt):
            x, xt = self.compute_tangent_endpoints()
            mob.put_start_and_end_on(x, xt)
        
        def place_grad_label(mob, dt):
            x, xt = self.compute_tangent_endpoints()
            p = xt + 1.0 * (xt - x) / np.linalg.norm(xt - x)
            mob.move_to(p)
        
        def advance_normed_tangent(mob, dt):
            x, xnt = self.compute_normalized_tangent_endpoints()
            mob.put_start_and_end_on(x, xnt)

        def advance_square(mob, dt):
            x, xnt = self.compute_normalized_tangent_endpoints()
            # Rotating would be nice
            mob.move_to(xnt)

        self.dot.add_updater(advance_bezier)
        self.circle.add_updater(advance_circle)
        tangent.add_updater(advance_tangent)
        normed_tangent.add_updater(advance_normed_tangent)
        sq_point.add_updater(advance_square)

        
        self.grad_label.add_updater(place_grad_label)

        arc = always_redraw(draw_arc)
        self.add(arc)
        self.wait(2.0)
        self.remove(arc)

        arc_curve = get_arc()
        self.add(arc_curve)
        self.add(
            self.make_square_marker().move_to(st_sq + self.circle.get_center())
        )

        self.play(
            Uncreate(self.curve),
            Uncreate(tangent),
            Uncreate(normed_tangent),
            Uncreate(self.grad_label),
            FadeOut(self.dot)
        )
        brace = ArcBrace(arc_curve)
        self.play(Write(brace))

        btext = brace.get_text("Gauss Length")
        self.play(Write(btext))