3b1b-manim/manimlib/animation/animation.py

from copy import deepcopy

import numpy as np

from manimlib.constants import *
from manimlib.mobject.mobject import Mobject
from manimlib.utils.config_ops import digest_config
from manimlib.utils.rate_functions import smooth


class Animation(object):
    CONFIG = {
        "run_time": DEFAULT_ANIMATION_RUN_TIME,
        "rate_func": smooth,
        "name": None,
        # Does this animation add or remove a mobject form the screen
        "remover": False,
        # TODO, replace this with a single lag parameter
        "submobject_mode": "all_at_once",
        "lag_factor": 2,
        # Used by EmptyAnimation to announce itself ignorable
        # in Successions and AnimationGroups
        "empty": False
    }

    def __init__(self, mobject, **kwargs):
        assert(isinstance(mobject, Mobject))
        self.mobject = mobject
        digest_config(self, kwargs)
        self.all_families_zipped = self.get_all_families_zipped()

    def begin(self):
        mobject = self.mobject
        # Make sure it's all up to date
        mobject.update()
        mobject.suspend_updating()
        # Keep track of where it started
        self.starting_mobject = mobject.copy()
        self.update(0)

    def finish(self):
        self.mobject.resume_updating()

    def clean_up_from_scene(self, scene):
        if self.is_remover():
            scene.remove(self.mobject)
        return self

    def __str__(self):
        if self.name:
            return self.name
        return self.__class__.__name__ + str(self.mobject)

    def copy(self):
        return deepcopy(self)

    def update_config(self, **kwargs):
        digest_config(self, kwargs)
        return self

    def update(self, alpha):
        alpha = np.clip(alpha, 0, 1)
        self.update_mobject(self.rate_func(alpha))

    def update_mobject(self, alpha):
        families = self.all_families_zipped
        for i, mobs in enumerate(families):
            sub_alpha = self.get_sub_alpha(alpha, i, len(families))
            self.update_submobject(*list(mobs) + [sub_alpha])
        return self

    def get_sub_alpha(self, alpha, index, num_submobjects):
        if self.submobject_mode in ["lagged_start", "smoothed_lagged_start"]:
            prop = float(index) / num_submobjects
            if self.submobject_mode is "smoothed_lagged_start":
                prop = smooth(prop)
            lf = self.lag_factor
            return np.clip(lf * alpha - (lf - 1) * prop, 0, 1)
        elif self.submobject_mode == "one_at_a_time":
            lower = float(index) / num_submobjects
            upper = float(index + 1) / num_submobjects
            return np.clip((alpha - lower) / (upper - lower), 0, 1)
        elif self.submobject_mode == "all_at_once":
            return alpha
        raise Exception("Invalid submobject mode")

    def update_submobject(self, submobject, starting_sumobject, alpha):
        # Typically ipmlemented by subclass
        pass

    def get_all_mobjects(self):
        """
        Ordering must match the ording of arguments to update_submobject
        """
        return self.mobject, self.starting_mobject

    def get_all_families_zipped(self):
        return list(zip(*map(
            Mobject.family_members_with_points,
            self.get_all_mobjects()
        )))

    def filter_out(self, *filter_functions):
        self.filter_functions += filter_functions
        return self

    def set_run_time(self, time):
        self.run_time = time
        return self

    def get_run_time(self):
        return self.run_time

    def set_rate_func(self, rate_func):
        self.rate_func = rate_func
        return self

    def get_rate_func(self):
        return self.rate_func

    def set_name(self, name):
        self.name = name
        return self

    def is_remover(self):
        return self.remover