newton-raphson.py

import math
import scipy.misc
from manimlib.imports import *


class NewtonRaphson(GraphScene):
    CONFIG = {
        "func": lambda x: ((x - 2)**2) * math.sin(2 * x) / 3 + 10,
        "initial_guess": 5,
        "num_iters": 7,

        "y_max": 50,
        "y_min": -50,
        "x_max": 15,
        "x_min": -15,
        "y_tick_frequency": 10,
        "x_tick_frequency": 3,
        "x_axis_label": None,
        "y_axis_label": None,
        "x_axis_width": 7.5,
        "y_axis_height": 7.25,
        "axes_color": BLUE,
        "graph_origin": 3 * LEFT,
    }

    def construct(self):
        self.setup_axes(animate=True)

        graph = self.get_graph(self.func, color=RED)
        graph_label = self.get_graph_label(graph, label="f(x)")
        self.play(ShowCreation(graph), ShowCreation(graph_label), run_time=2)

        step_1 = TextMobject("1. Make a guess $g$")
        step_2 = TextMobject("2. Find tangent of $f$ at $g$")
        step_3 = TextMobject("3. Set $g$ to $x$-intersect of the tangent")
        step_4 = TextMobject("4. Return to step 2 if $f(g)$ is not $0$")

        step_1.scale(.6).shift(3.2 * RIGHT + 2.9 * UP)
        step_2.scale(.6).next_to(step_1, DOWN).align_to(step_1, LEFT)
        step_3.scale(.6).next_to(step_2, DOWN).align_to(step_2, LEFT)
        step_4.scale(.6).next_to(step_3, DOWN).align_to(step_3, LEFT)

        g = TexMobject("g = ").align_to(step_1, LEFT).shift(1 * DOWN + 1 * RIGHT)
        f_g = TexMobject("f(g) = ").next_to(g, DOWN).align_to(g, RIGHT)

        self.play(FadeIn(step_1), FadeIn(step_2), FadeIn(step_3), FadeIn(step_4), FadeIn(g),
                  FadeIn(f_g))

        cur_guess = self.initial_guess

        tangent_graph = VectorizedPoint(self.input_to_graph_point(cur_guess, graph))
        x_intersect = VectorizedPoint()
        guess = VectorizedPoint().next_to(g, RIGHT)
        value = VectorizedPoint().next_to(f_g, RIGHT)
        arrow = Arrow().scale(.3).next_to(step_1, .5 * LEFT)

        self.add(tangent_graph)
        self.play(FadeIn(arrow))

        x_intersect = Dot(self.coords_to_point(cur_guess, 0))

        new_guess = TextMobject(f"{cur_guess:.3f}")
        new_guess.next_to(g, RIGHT)
        new_value = TextMobject(f"{self.func(cur_guess):.3f}")
        new_value.next_to(f_g, RIGHT)

        self.play(ShowCreation(x_intersect), Transform(guess, new_guess),
                  Transform(value, new_value))

        for _ in range(self.num_iters):
            slope = scipy.misc.derivative(self.func, cur_guess, dx=1e-6)
            val = self.func(cur_guess)

            vert_line = self.get_vertical_line_to_graph(cur_guess, graph, color=WHITE)
            new_tangent_graph = self.get_graph(lambda x: slope * (x - cur_guess) + val, color=GREEN)

            cur_guess -= val / slope

            new_guess = TextMobject(f"{cur_guess:.3f}")
            new_guess.next_to(g, RIGHT)
            new_value = TextMobject(f"{self.func(cur_guess):.3f}")
            new_value.next_to(f_g, RIGHT)

            new_x_intersect = Dot(self.coords_to_point(cur_guess, 0))

            self.play(ApplyMethod(arrow.next_to, step_2, .5 * LEFT))
            self.play(ShowCreation(vert_line))
            self.play(Transform(tangent_graph, new_tangent_graph))
            self.play(FadeOut(vert_line))

            self.play(ApplyMethod(arrow.next_to, step_3, .5 * LEFT))
            self.play(
                Transform(x_intersect, new_x_intersect),
                Transform(guess, new_guess),
                Transform(value, new_value),
            )

            self.play(ApplyMethod(arrow.next_to, step_4, .5 * LEFT))
            self.wait()