3b1b-manim/active_projects/clacks.py

from big_ol_pile_of_manim_imports import *
import subprocess
from pydub import AudioSegment


class SlidingBlocks(VGroup):
    CONFIG = {
        "block1_config": {
            "mass": 1,
            "velocity": -2,
            "distance": 7,
            "width": 1,
        },
        "block2_config": {
            "mass": 1,
            "velocity": 0,
            "distance": 3,
            "width": 1,
        },
        "block_style": {
            "fill_opacity": 1,
            "fill_color": (GREY, WHITE),
            "stroke_width": 3,
            "stroke_color": WHITE,
            "sheen_direction": UL,
        }
    }

    def __init__(self, surrounding_scene, **kwargs):
        VGroup.__init__(self, **kwargs)
        self.surrounding_scene = surrounding_scene
        self.floor = surrounding_scene.floor
        self.wall = surrounding_scene.wall

        self.block1 = self.get_block(**self.block1_config)
        self.block2 = self.get_block(**self.block2_config)
        self.mass_ratio = self.block2.mass / self.block1.mass
        self.phase_space_point_tracker = self.get_phase_space_point_tracker()
        self.add(
            self.block1, self.block2,
            self.phase_space_point_tracker,
        )
        self.add_updater(self.__class__.update_positions)

        self.clack_data = self.get_clack_data()

    def get_block(self, mass, width, distance, velocity):
        block = Square(side_length=width)
        block.mass = mass
        block.velocity = velocity

        block.set_style(**self.block_style)
        block.set_fill(color=interpolate_color(
            block.get_fill_color(),
            BLUE_E,
            (1 - 1.0 / (max(np.log10(mass), 1)))
        ))
        block.move_to(
            self.floor.get_top()[1] * UP +
            (self.wall.get_right()[0] + distance) * RIGHT,
            DL,
        )
        label = block.label = TextMobject(
            "{:,}\\,kg".format(int(mass))
        )
        label.scale(0.8)
        label.next_to(block, UP, SMALL_BUFF)
        block.add(label)
        return block

    def get_phase_space_point_tracker(self):
        block1, block2 = self.block1, self.block2
        w2 = block2.get_width()
        s1 = block1.get_left()[0] - self.wall.get_right()[0] - w2
        s2 = block2.get_right()[0] - self.wall.get_right()[0] - w2
        result = VectorizedPoint([
            s1 * np.sqrt(block1.mass),
            s2 * np.sqrt(block2.mass),
            0
        ])

        result.velocity = np.array([
            np.sqrt(block1.mass) * block1.velocity,
            np.sqrt(block2.mass) * block2.velocity,
            0
        ])
        return result

    def update_positions(self, dt):
        self.phase_space_point_tracker.shift(
            self.phase_space_point_tracker.velocity * dt
        )
        self.update_blocks_from_phase_space_point_tracker()

    def old_update_positions(self, dt):
        # Based on velocity diagram bouncing...didn't work for
        # large masses, due to frame rate mismatch
        blocks = self.submobjects
        for block in blocks:
            block.shift(block.velocity * dt * RIGHT)
        if blocks[0].get_left()[0] < blocks[1].get_right()[0]:
            # Two blocks collide
            m1 = blocks[0].mass
            m2 = blocks[1].mass
            v1 = blocks[0].velocity
            v2 = blocks[1].velocity
            v_phase_space_point = np.array([
                np.sqrt(m1) * v1, -np.sqrt(m2) * v2
            ])
            angle = 2 * np.arctan(np.sqrt(m2 / m1))
            new_vps_point = rotate_vector(v_phase_space_point, angle)
            for block, value in zip(blocks, new_vps_point):
                block.velocity = value / np.sqrt(block.mass)
            blocks[1].move_to(blocks[0].get_corner(DL), DR)
            self.surrounding_scene.clack(blocks[0].get_left())
        if blocks[1].get_left()[0] < self.wall.get_right()[0]:
            # Second block hits wall
            blocks[1].velocity *= -1
            blocks[1].move_to(self.wall.get_corner(DR), DL)
            if blocks[0].get_left()[0] < blocks[1].get_right()[0]:
                blocks[0].move_to(blocks[1].get_corner(DR), DL)
            self.surrounding_scene.clack(blocks[1].get_left())
        return self

    def update_blocks_from_phase_space_point_tracker(self):
        block1, block2 = self.block1, self.block2

        ps_point = self.phase_space_point_tracker.get_location()
        theta = np.arctan(np.sqrt(self.mass_ratio))
        ps_point_angle = angle_of_vector(ps_point)
        n_clacks = int(ps_point_angle / theta)
        reflected_point = rotate_vector(
            ps_point,
            -2 * np.ceil(n_clacks / 2) * theta
        )
        reflected_point = np.abs(reflected_point)

        shadow_wall_x = self.wall.get_right()[0] + block2.get_width()
        floor_y = self.floor.get_top()[1]
        s1 = reflected_point[0] / np.sqrt(block1.mass)
        s2 = reflected_point[1] / np.sqrt(block2.mass)
        block1.move_to(
            (shadow_wall_x + s1) * RIGHT +
            floor_y * UP,
            DL,
        )
        block2.move_to(
            (shadow_wall_x + s2) * RIGHT +
            floor_y * UP,
            DR,
        )

        self.surrounding_scene.update_num_clacks(n_clacks)

    def get_clack_data(self):
        ps_point = self.phase_space_point_tracker.get_location()
        ps_velocity = self.phase_space_point_tracker.velocity
        if ps_velocity[1] != 0:
            raise Exception(
                "Haven't implemented anything to gather clack "
                "data from a start state with block2 moving"
            )
        y = ps_point[1]
        theta = np.arctan(np.sqrt(self.mass_ratio))

        clack_data = []
        for k in range(1, int(PI / theta) + 1):
            clack_ps_point = np.array([
                y / np.tan(k * theta),
                y,
                0
            ])
            time = get_norm(ps_point - clack_ps_point) / get_norm(ps_velocity)
            reflected_point = rotate_vector(
                clack_ps_point,
                -2 * np.ceil((k - 1) / 2) * theta
            )
            block2 = self.block2
            s2 = reflected_point[1] / np.sqrt(block2.mass)
            location = np.array([
                self.wall.get_right()[0] + s2,
                block2.get_center()[1],
                0
            ])
            if k % 2 == 1:
                location += block2.get_width() * RIGHT
            clack_data.append((location, time))
        return clack_data


class ClackFlashes(ContinualAnimation):
    CONFIG = {
        "flash_config": {
            "run_time": 0.5,
            "line_length": 0.1,
            "flash_radius": 0.2,
        },
        "start_up_time": 0,
    }

    def __init__(self, clack_data, **kwargs):
        digest_config(self, kwargs)
        self.flashes = []
        group = Group()
        for location, time in clack_data:
            flash = Flash(location, **self.flash_config)
            flash.start_time = time
            flash.end_time = time + flash.run_time
            self.flashes.append(flash)
        ContinualAnimation.__init__(self, group, **kwargs)

    def update_mobject(self, dt):
        total_time = self.external_time
        for flash in self.flashes:
            if flash.start_time < total_time < flash.end_time:
                if flash.mobject not in self.mobject:
                    self.mobject.add(flash.mobject)
                flash.update(
                    (total_time - flash.start_time) / flash.run_time
                )
            else:
                if flash.mobject in self.mobject:
                    self.mobject.remove(flash.mobject)


class BlocksAndWallScene(Scene):
    CONFIG = {
        "include_sound_file": True,
        "count_clacks": True,
        "sliding_blocks_config": {},
        "floor_y": -2,
        "wall_x": -5,
        "clack_counter_label": "\\# Collisions: "
    }

    def setup(self):
        self.floor = self.get_floor()
        self.wall = self.get_wall()
        self.blocks = SlidingBlocks(self, **self.sliding_blocks_config)
        self.clack_data = self.blocks.clack_data
        self.clack_flashes = ClackFlashes(self.clack_data)
        self.add(self.floor, self.wall, self.blocks, self.clack_flashes)

        if self.count_clacks:
            self.add_clack_counter()
        self.track_time()

    def track_time(self):
        time_tracker = ValueTracker()
        time_tracker.add_updater(lambda m, dt: m.increment_value(dt))
        self.add(time_tracker)
        self.get_time = time_tracker.get_value

    def add_clack_counter(self):
        self.n_clacks = 0
        clack_counter_label = TextMobject(self.clack_counter_label)
        clack_counter = Integer(self.n_clacks)
        clack_counter.next_to(
            clack_counter_label[-1], RIGHT,
            aligned_edge=DOWN,
        )
        clack_group = VGroup(
            clack_counter_label,
            clack_counter,
        )
        clack_group.to_corner(UR)
        clack_group.shift(LEFT)
        self.add(clack_group)

        self.clack_counter = clack_counter

    def get_wall(self):
        wall = Line(self.floor_y * UP, FRAME_HEIGHT * UP / 2)
        wall.shift(self.wall_x * RIGHT)
        lines = VGroup(*[
            Line(ORIGIN, 0.25 * UR)
            for x in range(15)
        ])
        lines.set_stroke(width=1)
        lines.arrange_submobjects(UP, buff=MED_SMALL_BUFF)
        lines.move_to(wall, DR)
        wall.add(lines)
        return wall

    def get_floor(self):
        floor = Line(self.wall_x * RIGHT, FRAME_WIDTH * RIGHT / 2)
        floor.shift(self.floor_y * UP)
        return floor

    def update_num_clacks(self, n_clacks):
        if hasattr(self, "n_clacks"):
            if n_clacks == self.n_clacks:
                return
            self.clack_counter.set_value(n_clacks)

    def create_sound_file(self, clack_data):
        directory = get_scene_output_directory(self.__class__)
        clack_file = os.path.join(directory, 'sounds', 'bell.wav')
        output_file = self.get_movie_file_path(extension='.wav')
        times = [
            time
            for location, time in clack_data
            if time < 30
        ]

        clack = AudioSegment.from_wav(clack_file)
        total_time = max(times) + 1
        clacks = AudioSegment.silent(int(1000 * total_time))
        last_position = 0
        min_diff = 4
        for time in times:
            position = int(1000 * time)
            d_position = position - last_position
            last_position = position
            if d_position < min_diff:
                continue
            clacks = clacks.fade(-50, start=position, end=position + 10)
            clacks = clacks.overlay(
                clack,
                position=position
            )
        clacks.export(output_file, format="wav")
        return output_file

    def close_movie_pipe(self):
        Scene.close_movie_pipe(self)
        if self.include_sound_file:
            sound_file_path = self.create_sound_file(self.clack_data)
            movie_path = self.get_movie_file_path()
            temp_path = self.get_movie_file_path(str(self) + "TempSound")
            commands = [
                "ffmpeg",
                "-i", movie_path,
                "-i", sound_file_path,
                "-c:v", "copy", "-c:a", "aac",
                '-loglevel', 'error',
                "-strict", "experimental",
                temp_path,
            ]
            subprocess.call(commands)
            subprocess.call(["rm", sound_file_path])
            subprocess.call(["mv", temp_path, movie_path])

# Animated scenes


class MathAndPhysicsConspiring(Scene):
    def construct(self):
        pass


class BlocksAndWallExample(BlocksAndWallScene):
    CONFIG = {
        "sliding_blocks_config": {
            "block1_config": {
                "mass": 1e4,
                "width": 1.5,
                "velocity": -2,
            }
        }
    }

    def construct(self):
        self.wait(20)