marimo-team/colliding-blocks-and-pi

1# /// script2# requires-python = ">=3.13"3# dependencies = [4#     "marimo",5#     "matplotlib==3.10.1",6#     "numpy==2.2.3",7# ]8# ///9 10import marimo11 12__generated_with = "0.11.20"13app = marimo.App()14 15 16@app.cell17def _():18    import marimo as mo19    return (mo,)20 21 22@app.cell(hide_code=True)23def _(mo):24    mo.md(25        r"""26        # Finding $\pi$ in colliding blocks27 28        One of the remarkable things about mathematical constants like $\pi$ is how frequently they arise in nature, in the most surprising of places.29 30        Inspired by 3Blue1Brown, this [marimo notebook](https://github.com/marimo-team/marimo) shows how the number of collisions incurred in a particular system involving two blocks converges to the digits in $\pi$.31 32        **Tip!**: Use the menu in the top right to reveal the notebook's code.33        """34    )35    return36 37 38@app.cell(hide_code=True)39def _(mo):40    slider = mo.ui.slider(start=0, stop=3, value=3, show_value=True)41    return (slider,)42 43@app.cell(hide_code=True)44def _(mo, slider):45    mo.md("## Simulate!")46    return47    48@app.cell(hide_code=True)49def _(mo, slider):50    mo.md(f"Use this slider to control the weight of the heavier block: {slider}")51    return52 53 54@app.cell(hide_code=True)55def _(mo, slider):56    mo.md(rf"The heavier block weighs **$100^{{ {slider.value} }}$** kg.")57    return58 59 60@app.cell(hide_code=True)61def _(mo):62    run_button = mo.ui.run_button(label="Run simulation!")63    run_button.right()64    return (run_button,)65 66 67@app.cell68def _(run_button, simulate_collisions, slider):69    if run_button.value:70        mass_ratio = 100**slider.value71        _, ani, collisions = simulate_collisions(72            mass_ratio, total_time=15, dt=0.00173        )74    return ani, collisions, mass_ratio75 76 77@app.cell78def _(ani, mo, run_button):79    video = None80    if run_button.value:81        with mo.status.spinner(title="Rendering collision video ..."):82            video = mo.Html(ani.to_html5_video())83    video84    return (video,)85 86@app.cell(hide_code=True)87def _(mo):88    mo.md(89        r"""90        ## The 3Blue1Brown video91 92        If you haven't seen it, definitely check out the video that inspired this notebook:93        """94    )95    return96 97 98@app.cell(hide_code=True)99def _(mo):100    mo.accordion(101    {102        "🎥 Watch the video": mo.Html(103            '<iframe width="700" height="400" src="https://www.youtube.com/embed/6dTyOl1fmDo?si=xl9v6Y8x2e3r3A9I" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" referrerpolicy="strict-origin-when-cross-origin" allowfullscreen></iframe>'104        )105    })106    return107 108    109@app.cell110def _():111    import numpy as np112    import matplotlib.pyplot as plt113    import matplotlib.animation as animation114    from matplotlib.patches import Rectangle115    return Rectangle, animation, np, plt116 117 118@app.cell119def _():120    class Block:121        def __init__(self, mass, velocity, position, size=1.0):122            self.mass = mass123            self.velocity = velocity124            self.position = position125            self.size = size126 127        def update(self, dt):128            self.position += self.velocity * dt129 130        def collide(self, other):131            # Calculate velocities after elastic collision132            m1, m2 = self.mass, other.mass133            v1, v2 = self.velocity, other.velocity134 135            new_v1 = (m1 - m2) / (m1 + m2) * v1 + (2 * m2) / (m1 + m2) * v2136            new_v2 = (2 * m1) / (m1 + m2) * v1 + (m2 - m1) / (m1 + m2) * v2137 138            self.velocity = new_v1139            other.velocity = new_v2140 141            return 1  # Return 1 collision142    return (Block,)143 144 145@app.cell146def check_collisions():147    def check_collisions(small_block, big_block, wall_pos=0):148        collisions = 0149 150        # Check for collision between blocks151        if small_block.position + small_block.size > big_block.position:152            small_block.position = big_block.position - small_block.size153            collisions += small_block.collide(big_block)154 155        # Check for collision with the wall156        if small_block.position < wall_pos:157            small_block.position = wall_pos158            small_block.velocity *= -1159            collisions += 1160 161        return collisions162    return (check_collisions,)163 164 165@app.cell166def _(Block, check_collisions, create_animation):167    def simulate_collisions(mass_ratio, total_time=15, dt=0.001, animate=True):168        # Initialize blocks169        small_block = Block(mass=1, velocity=0, position=2)170        big_block = Block(mass=mass_ratio, velocity=-0.5, position=4)171 172        # Simulation variables173        time = 0174        collision_count = 0175 176        # For animation177        times = []178        small_positions = []179        big_positions = []180        collision_counts = []181 182        # Run simulation183        while time < total_time:184            # Update positions185            small_block.update(dt)186            big_block.update(dt)187 188            # Check for and handle collisions189            new_collisions = check_collisions(small_block, big_block)190            collision_count += new_collisions191 192            # Store data for animation193            times.append(time)194            small_positions.append(small_block.position)195            big_positions.append(big_block.position)196            collision_counts.append(collision_count)197 198            time += dt199 200 201        print(f"Mass ratio: {mass_ratio}, Total collisions: {collision_count}")202 203        if animate:204            axis, ani = create_animation(205                times, small_positions, big_positions, collision_counts, mass_ratio206            )207        else:208            axis, ani = None209 210        return axis, ani, collision_count211    return (simulate_collisions,)212 213 214@app.cell215def _(Rectangle, animation, plt):216    def create_animation(217        times, small_positions, big_positions, collision_counts, mass_ratio218    ):219        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))220 221        # Setup for blocks visualization222        ax1.set_xlim(-1, 10)223        ax1.set_ylim(-1, 2)224        ax1.set_xlabel("Position")225        ax1.set_title(f"Block Collisions (Mass Ratio = {mass_ratio})")226        wall = plt.Line2D([0, 0], [-1, 2], color="black", linewidth=3)227        ax1.add_line(wall)228 229        small_block = Rectangle((small_positions[0], 0), 1, 1, color="blue")230        big_block = Rectangle((big_positions[0], 0), 1, 1, color="red")231        ax1.add_patch(small_block)232        ax1.add_patch(big_block)233 234        # Add weight labels for each block235        small_label = ax1.text(236            small_positions[0] + 0.5,237            1.2,238            f"{1}kg",239            ha="center",240            va="center",241            color="blue",242            fontweight="bold",243        )244        big_label = ax1.text(245            big_positions[0] + 0.5,246            1.2,247            f"{mass_ratio}kg",248            ha="center",249            va="center",250            color="red",251            fontweight="bold",252        )253 254        # Setup for collision count255        ax2.set_xlim(0, times[-1])256        # ax2.set_ylim(0, collision_counts[-1] * 1.1)257        ax2.set_ylim(0, collision_counts[-1] * 1.1)258        ax2.set_xlabel("Time")259        ax2.set_ylabel("# Collisions:")260        ax2.set_yscale("symlog")261        (collision_line,) = ax2.plot([], [], "g-")262 263        # Add text for collision count264        collision_text = ax2.text(265            0.02, 0.9, "", transform=ax2.transAxes, fontsize="x-large"266        )267 268        def init():269            small_block.set_xy((small_positions[0], 0))270            big_block.set_xy((big_positions[0], 0))271            small_label.set_position((small_positions[0] + 0.5, 1.2))272            big_label.set_position((big_positions[0] + 0.5, 1.2))273            collision_line.set_data([], [])274            collision_text.set_text("")275            return small_block, big_block, collision_line, collision_text276 277        frame_step = 300278 279        def animate(i):280            # Speed up animation but ensure we reach the final frame281            frame_index = min(i * frame_step, len(times) - 1)282 283            small_block.set_xy((small_positions[frame_index], 0))284            big_block.set_xy((big_positions[frame_index], 0))285 286            # Update the weight labels to follow the blocks287            small_label.set_position((small_positions[frame_index] + 0.5, 1.2))288            big_label.set_position((big_positions[frame_index] + 0.5, 1.2))289 290            # Show data up to the current frame291            collision_line.set_data(292                times[: frame_index + 1], collision_counts[: frame_index + 1]293            )294 295            # For the last frame, show the final collision count296            if frame_index >= len(times) - 1:297                collision_text.set_text(298                    f"# Collisions: {collision_counts[-1]}"299                )300            else:301                collision_text.set_text(302                    f"# Collisions: {collision_counts[frame_index]}"303                )304 305            return (306                small_block,307                big_block,308                small_label,309                big_label,310                collision_line,311                collision_text,312            )313 314        plt.tight_layout()315 316        frames = max(1, len(times) // frame_step)  # Ensure at least 1 frame317        ani = animation.FuncAnimation(318            fig,319            animate,320            frames=frames + 1,  # +1 to ensure we reach the end321            init_func=init,322            blit=True,323            interval=30,324        )325 326        plt.tight_layout()327        return plt.gca(), ani328 329        # Uncomment to save animation330        # ani.save('pi_collisions.mp4', writer='ffmpeg', fps=30)331    return (create_animation,)332 333 334if __name__ == "__main__":335    app.run()