Animated mAP in information retrieval example with manim

gistfile1.txt

from manim import *
from manim.utils.color import rgb_to_color

# Set custom resolution
config.pixel_width = 1600
config.pixel_height = 1200

# Tailwind colors
teal_correct = rgb_to_color((0.129, 0.588, 0.588))  # Teal for relevant items
grey_incorrect = GREY  # Grey for irrelevant items
pale_teal = rgb_to_color((0.7, 0.9, 0.9))  # Pale teal for faded items
pale_grey = rgb_to_color((0.9, 0.9, 0.9))  # Pale grey for faded items
pale_grey_f = rgb_to_color((0.3, 0.3, 0.3))  # Pale grey for faded items

class IRResultTableWithAPAnimation(Scene):
    def construct(self):
        # Fonts
        font_name = "Inter"  # Replace with your desired font name

        # Data for the table
        rankings = list(range(1, 11))  # Ranks 1 to 10
        items = ["Item {}".format(i) for i in range(1, 11)]
        relevant_indices = [2, 4, 7]  # Relevant items (0-based indices)

        # Table configuration
        cell_width = 3  # Width of rectangles
        cell_height = 0.6  # Height of rectangles
        vertical_spacing = 0.2  # Spacing between rows

        # Calculate table height to center it dynamically with padding
        total_height = len(rankings) * (cell_height + vertical_spacing) - vertical_spacing
        start_y = (total_height / 2) + 0.5  # Add padding at the top

        # Store all elements for alignment
        table_elements = VGroup()

        # Draw the table
        rects = VGroup()  # Use VGroup for animations
        labels = VGroup()  # Use VGroup for animations

        for i, (rank, item) in enumerate(zip(rankings, items)):
            # Position of the current row
            y = start_y - i * (cell_height + vertical_spacing)

            # Draw the rank
            rank_text = Text(str(rank), font_size=24, font=font_name, color=WHITE)  # Rank column in white
            rank_text.next_to([-2.5, y, 0], RIGHT, buff=0.5)
            table_elements.add(rank_text)

            # Draw the item as a rectangle
            item_color = teal_correct if i in relevant_indices else grey_incorrect
            rect = Rectangle(
                width=cell_width,
                height=cell_height,
                color=item_color,
                fill_opacity=0.8,
                stroke_width=0
            )
            rect.move_to([0, y, 0])
            table_elements.add(rect)
            rects.add(rect)

            # Add label inside the rectangle
            item_label = Text(item, font_size=20, font=font_name, color=WHITE)
            item_label.move_to(rect.get_center())
            table_elements.add(item_label)
            labels.add(item_label)

        # Add a `}` shape to connect the table and formula
        connecting_brace = Brace(
            table_elements, direction=RIGHT, buff=0.3
        )

        # LaTeX formula
        formula3 = MathTex(
            r"AP = \frac{\frac{1}{3} + \frac{2}{5} + \frac{3}{8}}{3}",
            font_size=54,
        )
        formula3.move_to(RIGHT * 3)  # Consistent position
        formula_final = MathTex(r"AP = 0.295", font_size=54).next_to(formula3, DOWN, aligned_edge=LEFT)

        # Group and shift the table and brace to the left
        table_and_brace = VGroup(table_elements, connecting_brace)
        table_and_brace.shift(LEFT * 2.5)  # Shift the table and brace to the left

        # Add initial elements to the scene
        self.add(table_elements, connecting_brace, formula3)

        # Adjust the position of the formula group to align with the brace
        formula_group = VGroup(formula3, formula_final)
        formula_group.shift(UP * (connecting_brace.get_center()[1] - formula_group.get_center()[1]))


        # Group all elements (table, brace, and formulas)
        content_group = VGroup(table_and_brace, formula_group)

        # Center the content vertically
        content_group.move_to(ORIGIN)

        print('debug')
        print(formula3[0][7], formula3[0][8], formula3[0][9], formula3[0][10], formula3[0][11])

        # Animation step 1: Highlight the first term
        self.play(
            rects[3:].animate.set_fill(pale_grey),  # Fade irrelevant items after rank 3
            rects[2].animate.set_fill(teal_correct),  # Highlight rank 3
            labels[3:].animate.set_color(pale_grey),  # Fade label colors after rank 3
            labels[2].animate.set_color(WHITE),  # Ensure rank 3 label remains white
            formula3[0][:6].animate.set_color(WHITE),  # Highlight \frac{1}{3}
            formula3[0][6:].animate.set_color(pale_grey_f),  # Fade other terms
        )
        self.wait(1)

        # Animation step 2: Highlight the first two terms
        self.play(
            rects[5:].animate.set_fill(pale_grey),  # Fade irrelevant items after rank 5
            rects[4].animate.set_fill(teal_correct),  # Highlight rank 5
            rects[2].animate.set_fill(teal_correct),  # Keep rank 3 highlighted
            rects[3].animate.set_fill(grey_incorrect),  # De-highlight rank 4
            labels[5:].animate.set_color(pale_grey),  # Fade labels after rank 5
            labels[4].animate.set_color(WHITE),  # Ensure rank 5 label remains white
            formula3[0][:10].animate.set_color(WHITE),  # Highlight \frac{1}{3} + \frac{2}{5}
            formula3[0][10:].animate.set_color(pale_grey_f),  # Fade other terms
        )
        self.wait(1)

        # Animation step 3: Highlight all numerator terms
        self.play(
            rects[8:].animate.set_fill(pale_grey),  # Fade irrelevant items after rank 8
            rects[7].animate.set_fill(teal_correct),  # Highlight rank 8
            rects[4].animate.set_fill(teal_correct),  # Keep rank 5 highlighted
            rects[2].animate.set_fill(teal_correct),  # Keep rank 3 highlighted
            rects[3].animate.set_fill(grey_incorrect),
            rects[5].animate.set_fill(grey_incorrect),
            rects[6].animate.set_fill(grey_incorrect),
            labels[8:].animate.set_color(pale_grey),  # Fade labels after rank 8
            labels[7].animate.set_color(WHITE),  # Ensure rank 8 label remains white
            formula3[0][:].animate.set_color(WHITE),  # Highlight full numerator
        )
        self.wait(1)

        # Animation step 4: Add final AP value
        self.play(
            Write(formula_final)  # Add AP = 0.295 below the formula
        )
        self.wait(2)