PlanetAnimation.py

df = pd.read_parquet("plotme.parquet")

class PlanetAnimation2(Scene):
    def construct(self):
        my_template = TexTemplate()
        my_template.add_to_preamble(r"\usepackage{mathabx}")
        Tex.set_default(tex_template=my_template)

        scale = 12
        shift = np.array([-4, -1])
        times = np.array(df["date_frac"].values)
        def year_to_proportion(year):
            return (year - times[0]) / (times[-1] - times[0])

        # Create a list of positions for each planet from the dataframe

        sun_positions = np.array(df[["sun_x_1", "sun_y_1"]].values) / scale + shift
        jupiter_positions = (
            np.array(df[["jupiter_x_1", "jupiter_y_1"]].values) / scale + shift
        )
        saturn_positions = (
            np.array(df[["saturn_x_1", "saturn_y_1"]].values) / scale + shift
        )
        uranus_positions = (
            np.array(df[["uranus_x_1", "uranus_y_1"]].values) / scale + shift
        )
        neptune_positions = (
            np.array(df[["neptune_x", "neptune_y"]].values) / scale + shift
        )
        arcsec_error = np.array(df[["arcsec_error"]].values)

        # Create a labeled dot for each planet
        sun = LabeledDot(
            Tex(r"$\Sun$", color=BLACK),
            color=colors["sun"],
        )
        jupiter = LabeledDot(Tex(r"$\Jupiter$", color=BLACK), color=colors["jupiter"])
        saturn = LabeledDot(Tex(r"$\Saturn$", color=BLACK), color=colors["saturn"])
        uranus = LabeledDot(Tex(r"$\Uranus$", color=BLACK), color=colors["uranus"])
        uranus_real = LabeledDot(Tex(r"$\Uranus$", color=BLACK), color=colors["uranus"])
        uranus_pred = LabeledDot(Tex(r"$P$", color=BLACK), color=colors["uranus"])
        neptune = LabeledDot(Tex(r"$\Neptune$", color=BLACK), color=colors["neptune"])

        # Set initial positions
        sun.move_to(np.append(sun_positions[0], 0))
        jupiter.move_to(np.append(jupiter_positions[0], 0))
        saturn.move_to(np.append(saturn_positions[0], 0))
        uranus.move_to(np.append(uranus_positions[0], 0))
        neptune.move_to(np.append(neptune_positions[0], 0))

        self.add(sun, jupiter, saturn, uranus, neptune)

        # Define animations
        jupiter_path = (
            VMobject()
            .set_points_smoothly([np.append(pos, 0) for pos in jupiter_positions])
            .set_z_index(-1)
        )
        saturn_path = (
            VMobject()
            .set_points_smoothly([np.append(pos, 0) for pos in saturn_positions])
            .set_z_index(-1)
        )
        uranus_path = (
            VMobject()
            .set_points_smoothly([np.append(pos, 0) for pos in uranus_positions])
            .set_z_index(-1)
        )
        neptune_path = (
            VMobject()
            .set_points_smoothly([np.append(pos, 0) for pos in neptune_positions])
            .set_z_index(-1)
        )

        l1 = NumberLine(
            x_range=[1650, 1850, 50],
            length=4,
            numbers_with_elongated_ticks=[1700, 1800],
            include_numbers=True,
            font_size=30,
            color=WHITE,
            decimal_number_config={"group_with_commas": False, "num_decimal_places": 0},
        ).move_to(np.array([-4.0, 2.5, 0]))

        firstnumber = times[0]
        selector = (
            Triangle(fill_opacity=1, color=WHITE)
            .scale(0.2)
            .rotate(PI / 3)
            .next_to(l1.n2p(firstnumber), UP, buff=0.1)
        )
        dn = DecimalNumber(
            firstnumber, num_decimal_places=0, group_with_commas=False
        ).next_to(selector, UP, buff=0.2)

        update_grp = VGroup(selector, dn)

        def update_vgrp(vgrp):
            s, d = vgrp
            s.next_to(l1.n2p(dn.get_value()), UP, buff=0.1)
            d.next_to(selector, UP, buff=0.2)

        update_grp.add_updater(update_vgrp)

        self.add(l1, update_grp)

        self.add(
            jupiter_path.set_color(colors["jupiter"]),
            saturn_path.set_color(colors["saturn"]),
            uranus_path.set_color(colors["uranus"]),
            neptune_path.set_color(colors["neptune"]),
        )

        ########### Arcsec error
        # Set up axes
        axes = Axes(
            x_range=[1650, 1850, 50],
            y_range=(0, df.arcsec_error.max(), 50),
            x_length=5,
            y_length=6 / 2,
            x_axis_config={
                "include_numbers": True,
                "decimal_number_config": {
                    "group_with_commas": False,
                    "num_decimal_places": 0,
                },
            },
            y_axis_config={"include_numbers": True},
            tips=False,
        )
        axes.move_to(np.array([3.5, 0, 0]))
        # Set up labels
        x_label = axes.get_x_axis_label(MathTex(r"\text{Year}")).next_to(axes, DOWN)
        y_label = (
            axes.get_y_axis_label(
                MathTex(r"\text{Residual }[\text{arcsec}] ")
            )
            .rotate(np.pi / 2)
            .next_to(axes, LEFT)
        )

        # Prepare coordinates for Manim
        coords = list(
            df[["date_frac", "arcsec_error"]].itertuples(index=False, name=None)
        )
        points = [axes.coords_to_point(x, y) for x, y in coords]

        # Create line plot
        line_plot = VMobject().set_points_as_corners(points).set_color(BLUE)

        # Add everything to the scene
        self.add(axes, x_label, y_label, line_plot)

        # Add a value tracker
        tracker = ValueTracker(times[0])

        # Create a dot that moves along the line
        dot = Dot().move_to(line_plot.point_from_proportion(0))
        # Create the vertical and horizontal dashed lines
        h_line = DashedLine(axes.get_corner(DL), axes.get_corner(DR)).set_color(YELLOW)
        v_line = DashedLine(axes.get_corner(DL), axes.get_corner(UL)).set_color(YELLOW)
        # Group them together
        tracing_dot_grp = VGroup(dot, h_line, v_line)

        # Add updater function to the group        def dot_updater(mob):
            dot, h_line, v_line = mob
            x_value = dn.get_value()
            
            # Find nearest date_frac to x_value
            nearest_index = df['date_frac'].sub(x_value).abs().idxmin()
            
            y_value = df.loc[nearest_index, 'arcsec_error']
            
            new_point = axes.coords_to_point(x_value, y_value)
            dot.move_to(new_point)
            h_line.put_start_and_end_on(
                dot.get_center(),
                axes.c2p(axes.x_range[1], axes.p2c(dot.get_center())[1]),
            )
            v_line.put_start_and_end_on(
                dot.get_center(),
                axes.c2p(axes.p2c(dot.get_center())[0], axes.y_range[0]),
            )

        tracing_dot_grp.add_updater(dot_updater)
        self.add(tracing_dot_grp)

        # Planet animations and timeline update
        self.play(
            MoveAlongPath(jupiter, jupiter_path, rate_func=linear),
            MoveAlongPath(saturn, saturn_path, rate_func=linear),
            MoveAlongPath(uranus, uranus_path, rate_func=linear),
            # MoveAlongPath(uranus_force_vector, force_path, rate_func=linear),
            MoveAlongPath(neptune, neptune_path, rate_func=linear),
            ChangeDecimalToValue(dn, times[-1], rate_func=linear),
            tracker.animate.set_value(times[-1]),  # Animate value tracker
            run_time=5,
        )
        self.wait()