Update app.py

app.py

@@ -1,180 +1,324 @@
 # /// script
-# requires-python = ">=3.10"
+# requires-python = ">=3.13"
 # dependencies = [
 #     "marimo",
-#     "numba==0.60.0",
-#     "numpy==2.0.2",
-#     "scikit-image==0.24.0",
+#     "matplotlib==3.10.1",
+#     "numpy==2.2.3",
 # ]
 # ///
 
 import marimo
 
-__generated_with = "0.9.6"
-app = marimo.App(width="medium")
+__generated_with = "0.11.20"
+app = marimo.App()
+
+
+@app.cell
+def _():
+    import marimo as mo
+    return (mo,)
 
 
 @app.cell(hide_code=True)
-def __(mo):
+def _(mo):
     mo.md(
-        """
-        # Seam Carving 
-
-        _Example adapted from work by [Vincent Warmerdam](https://x.com/fishnets88)_.
+        r"""
+        # Finding $\pi$ in colliding blocks
 
-        ## The seam carving algorithm
-        This marimo demonstration is partially an homage to [a great video by Grant
-        Sanderson](https://www.youtube.com/watch?v=rpB6zQNsbQU) of 3Blue1Brown, which demonstrates
-        the seam carving algorithm in [Pluto.jl](https://plutojl.org/):
+        One of the remarkable things about mathematical constants like $\pi$ is how frequently they arise in nature, in the most surprising of places.
 
-        <iframe width="560" height="315" src="https://www.youtube.com/embed/rpB6zQNsbQU?si=oiZclGIj2atJR47m" 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>
+        Inspired by 3Blue1Brown, this marimo notebook shows how the number of collisions incurred in a particular system involving two blocks converges to the digits in $\pi$.
+        """
+    )
+    return
 
-        As Grant explains, the seam carving algorithm preserves the shapes of the main content in the image, while killing the "dead space": the image is resized, but the clocks and other content are not resized or deformed.
 
-        This notebook is a Python version of the seam carving algorithm, but it is also a
-        demonstration of marimo's [caching
-        feature](https://docs.marimo.io/guides/best_practices/performance.html#cache-computations-with-mo-cache),
-        which is helpful because the algorithm is compute intensive even when you
-        use [Numba](https://numba.pydata.org/).
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## The 3Blue1Brown video
 
-        Try it out by playing with the slider!
+        If you haven't seen it, definitely check out the video that inspired this notebook!
         """
     )
     return
 
 
 @app.cell(hide_code=True)
-def __():
-    input_image = "https://upload.wikimedia.org/wikipedia/en/d/dd/The_Persistence_of_Memory.jpg"
+def _(mo):
+    mo.Html('<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>')
+    return
+
 
-    return input_image,
+@app.cell(hide_code=True)
+def _(mo):
+    slider = mo.ui.slider(start=0, stop=4, value=0, show_value=True)
+    return (slider,)
 
 
 @app.cell(hide_code=True)
-def __(mo):
-    mo.md("""## Try it!""")
+def _(mo, slider):
+    mo.md(f"Use this slider to control the weight of the heavier block: {slider}")
     return
 
 
-@app.cell
-def __():
-    import marimo as mo
+@app.cell(hide_code=True)
+def _(mo, slider):
+    mo.md(rf"The heavier block weighs **$100^{{ {slider.value} }}$** kg.")
+    return
 
-    slider = mo.ui.slider(
-        0.7,
-        1.0,
-        step=0.05,
-        value=1.0,
-        label="Amount of resizing to perform:",
-        show_value=True,
-    )
-    slider
-    return mo, slider
+
+@app.cell(hide_code=True)
+def _(mo):
+    run_button = mo.ui.run_button(label="Run simulation!")
+    run_button.right()
+    return (run_button,)
 
 
 @app.cell
-def __(efficient_seam_carve, input_image, mo, slider):
-    scale_factor = slider.value
-    result = efficient_seam_carve(input_image, scale_factor)
+def _(run_button, simulate_collisions, slider):
+    if run_button.value:
+        mass_ratio = 100**slider.value
+        _, ani, collisions = simulate_collisions(
+            mass_ratio, total_time=15, dt=0.001
+        )
+    return ani, collisions, mass_ratio
 
-    mo.hstack([mo.image(input_image), mo.image(result)], justify="start")
-    return result, scale_factor
+
+@app.cell
+def _(ani, mo, run_button):
+    video = None
+    if run_button.value:
+        with mo.status.spinner(title="Rendering collision video ..."):
+            video = mo.Html(ani.to_html5_video())
+    video
+    return (video,)
 
 
 @app.cell
-def __(mo):
+def _():
     import numpy as np
-    from numba import jit
-    from skimage import io, filters, transform
-    import time
-
-
-    def rgb2gray(rgb):
-        return np.dot(rgb[..., :3], [0.2989, 0.5870, 0.1140])
-
-
-    def compute_energy_map(gray):
-        return np.abs(filters.sobel_h(gray)) + np.abs(filters.sobel_v(gray))
-
-
-    @jit(nopython=True)
-    def find_seam(energy_map):
-        height, width = energy_map.shape
-        dp = energy_map.copy()
-        backtrack = np.zeros((height, width), dtype=np.int32)
-
-        for i in range(1, height):
-            for j in range(width):
-                if j == 0:
-                    idx = np.argmin(dp[i - 1, j : j + 2])
-                    backtrack[i, j] = idx + j
-                    min_energy = dp[i - 1, idx + j]
-                elif j == width - 1:
-                    idx = np.argmin(dp[i - 1, j - 1 : j + 1])
-                    backtrack[i, j] = idx + j - 1
-                    min_energy = dp[i - 1, idx + j - 1]
-                else:
-                    idx = np.argmin(dp[i - 1, j - 1 : j + 2])
-                    backtrack[i, j] = idx + j - 1
-                    min_energy = dp[i - 1, idx + j - 1]
-
-                dp[i, j] += min_energy
-
-        return backtrack
-
-
-    @jit(nopython=True)
-    def remove_seam(image, backtrack):
-        height, width, _ = image.shape
-        output = np.zeros((height, width - 1, 3), dtype=np.uint8)
-        j = np.argmin(backtrack[-1])
-
-        for i in range(height - 1, -1, -1):
-            for k in range(3):
-                output[i, :, k] = np.delete(image[i, :, k], j)
-            j = backtrack[i, j]
-
-        return output
-
-
-    def seam_carving(image, new_width):
-        height, width, _ = image.shape
-
-        while width > new_width:
-            gray = rgb2gray(image)
-            energy_map = compute_energy_map(gray)
-            backtrack = find_seam(energy_map)
-            image = remove_seam(image, backtrack)
-            width -= 1
-
-        return image
-
-    @mo.cache
-    def efficient_seam_carve(image_path, scale_factor):
-        img = io.imread(image_path)
-        new_width = int(img.shape[1] * scale_factor)
-
-        start_time = time.time()
-        carved_img = seam_carving(img, new_width)
-        end_time = time.time()
-
-        print(f"Seam carving completed in {end_time - start_time:.2f} seconds")
-
-        return carved_img
-    return (
-        compute_energy_map,
-        efficient_seam_carve,
-        filters,
-        find_seam,
-        io,
-        jit,
-        np,
-        remove_seam,
-        rgb2gray,
-        seam_carving,
-        time,
-        transform,
-    )
+    import matplotlib.pyplot as plt
+    import matplotlib.animation as animation
+    from matplotlib.patches import Rectangle
+    return Rectangle, animation, np, plt
+
+
+@app.cell
+def _():
+    class Block:
+        def __init__(self, mass, velocity, position, size=1.0):
+            self.mass = mass
+            self.velocity = velocity
+            self.position = position
+            self.size = size
+
+        def update(self, dt):
+            self.position += self.velocity * dt
+
+        def collide(self, other):
+            # Calculate velocities after elastic collision
+            m1, m2 = self.mass, other.mass
+            v1, v2 = self.velocity, other.velocity
+
+            new_v1 = (m1 - m2) / (m1 + m2) * v1 + (2 * m2) / (m1 + m2) * v2
+            new_v2 = (2 * m1) / (m1 + m2) * v1 + (m2 - m1) / (m1 + m2) * v2
+
+            self.velocity = new_v1
+            other.velocity = new_v2
+
+            return 1  # Return 1 collision
+    return (Block,)
+
+
+@app.cell
+def check_collisions():
+    def check_collisions(small_block, big_block, wall_pos=0):
+        collisions = 0
+
+        # Check for collision between blocks
+        if small_block.position + small_block.size > big_block.position:
+            small_block.position = big_block.position - small_block.size
+            collisions += small_block.collide(big_block)
+
+        # Check for collision with the wall
+        if small_block.position < wall_pos:
+            small_block.position = wall_pos
+            small_block.velocity *= -1
+            collisions += 1
+
+        return collisions
+    return (check_collisions,)
+
+
+@app.cell
+def _(Block, check_collisions, create_animation):
+    def simulate_collisions(mass_ratio, total_time=15, dt=0.001, animate=True):
+        # Initialize blocks
+        small_block = Block(mass=1, velocity=0, position=2)
+        big_block = Block(mass=mass_ratio, velocity=-0.5, position=4)
+
+        # Simulation variables
+        time = 0
+        collision_count = 0
+
+        # For animation
+        times = []
+        small_positions = []
+        big_positions = []
+        collision_counts = []
+
+        # Run simulation
+        while time < total_time:
+            # Update positions
+            small_block.update(dt)
+            big_block.update(dt)
+
+            # Check for and handle collisions
+            new_collisions = check_collisions(small_block, big_block)
+            collision_count += new_collisions
+
+            # Store data for animation
+            times.append(time)
+            small_positions.append(small_block.position)
+            big_positions.append(big_block.position)
+            collision_counts.append(collision_count)
+
+            time += dt
+
+
+        print(f"Mass ratio: {mass_ratio}, Total collisions: {collision_count}")
+
+        if animate:
+            axis, ani = create_animation(
+                times, small_positions, big_positions, collision_counts, mass_ratio
+            )
+        else:
+            axis, ani = None
+
+        return axis, ani, collision_count
+    return (simulate_collisions,)
+
+
+@app.cell
+def _(Rectangle, animation, plt):
+    def create_animation(
+        times, small_positions, big_positions, collision_counts, mass_ratio
+    ):
+        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))
+
+        # Setup for blocks visualization
+        ax1.set_xlim(-1, 10)
+        ax1.set_ylim(-1, 2)
+        ax1.set_xlabel("Position")
+        ax1.set_title(f"Block Collisions (Mass Ratio = {mass_ratio})")
+        wall = plt.Line2D([0, 0], [-1, 2], color="black", linewidth=3)
+        ax1.add_line(wall)
+
+        small_block = Rectangle((small_positions[0], 0), 1, 1, color="blue")
+        big_block = Rectangle((big_positions[0], 0), 1, 1, color="red")
+        ax1.add_patch(small_block)
+        ax1.add_patch(big_block)
+
+        # Add weight labels for each block
+        small_label = ax1.text(
+            small_positions[0] + 0.5,
+            1.2,
+            f"{1}kg",
+            ha="center",
+            va="center",
+            color="blue",
+            fontweight="bold",
+        )
+        big_label = ax1.text(
+            big_positions[0] + 0.5,
+            1.2,
+            f"{mass_ratio}kg",
+            ha="center",
+            va="center",
+            color="red",
+            fontweight="bold",
+        )
+
+        # Setup for collision count
+        ax2.set_xlim(0, times[-1])
+        # ax2.set_ylim(0, collision_counts[-1] * 1.1)
+        ax2.set_ylim(0, collision_counts[-1] * 1.1)
+        ax2.set_xlabel("Time")
+        ax2.set_ylabel("# Collisions:")
+        ax2.set_yscale("symlog")
+        (collision_line,) = ax2.plot([], [], "g-")
+
+        # Add text for collision count
+        collision_text = ax2.text(
+            0.02, 0.9, "", transform=ax2.transAxes, fontsize="x-large"
+        )
+
+        def init():
+            small_block.set_xy((small_positions[0], 0))
+            big_block.set_xy((big_positions[0], 0))
+            small_label.set_position((small_positions[0] + 0.5, 1.2))
+            big_label.set_position((big_positions[0] + 0.5, 1.2))
+            collision_line.set_data([], [])
+            collision_text.set_text("")
+            return small_block, big_block, collision_line, collision_text
+
+        frame_step = 300
+
+        def animate(i):
+            # Speed up animation but ensure we reach the final frame
+            frame_index = min(i * frame_step, len(times) - 1)
+
+            small_block.set_xy((small_positions[frame_index], 0))
+            big_block.set_xy((big_positions[frame_index], 0))
+
+            # Update the weight labels to follow the blocks
+            small_label.set_position((small_positions[frame_index] + 0.5, 1.2))
+            big_label.set_position((big_positions[frame_index] + 0.5, 1.2))
+
+            # Show data up to the current frame
+            collision_line.set_data(
+                times[: frame_index + 1], collision_counts[: frame_index + 1]
+            )
+
+            # For the last frame, show the final collision count
+            if frame_index >= len(times) - 1:
+                collision_text.set_text(
+                    f"# Collisions: {collision_counts[-1]}"
+                )
+            else:
+                collision_text.set_text(
+                    f"# Collisions: {collision_counts[frame_index]}"
+                )
+
+            return (
+                small_block,
+                big_block,
+                small_label,
+                big_label,
+                collision_line,
+                collision_text,
+            )
+
+        plt.tight_layout()
+
+        frames = max(1, len(times) // frame_step)  # Ensure at least 1 frame
+        ani = animation.FuncAnimation(
+            fig,
+            animate,
+            frames=frames + 1,  # +1 to ensure we reach the end
+            init_func=init,
+            blit=True,
+            interval=30,
+        )
+
+        plt.tight_layout()
+        return plt.gca(), ani
+
+        # Uncomment to save animation
+        # ani.save('pi_collisions.mp4', writer='ffmpeg', fps=30)
+    return (create_animation,)
 
 
 if __name__ == "__main__":