3b1b-manim/mobject/mobject.py

import numpy as np
import operator as op
import os
from PIL import Image
from copy import deepcopy
from colour import Color

from helpers import *


#TODO: Explain array_attrs

class Mobject(object):
    """
    Mathematical Object
    """
    CONFIG = {
        "color"        : WHITE,
        "stroke_width" : DEFAULT_POINT_THICKNESS,
        "name"         : None,
        "dim"          : 3,
        "target"       : None,
    }
    def __init__(self, *submobjects, **kwargs):
        digest_config(self, kwargs)
        if not all(map(lambda m : isinstance(m, Mobject), submobjects)):
            raise Exception("All submobjects must be of type Mobject")
        self.submobjects = list(submobjects)
        self.color = Color(self.color)
        if self.name is None:
            self.name = self.__class__.__name__
        self.init_points()
        self.generate_points()
        self.init_colors()

    def __str__(self):
        return str(self.name)

    def init_points(self):
        self.points = np.zeros((0, self.dim))

    def init_colors(self):
        #For subclasses
        pass

    def generate_points(self):
        #Typically implemented in subclass, unless purposefully left blank
        pass

    def add(self, *mobjects):
        self.submobjects = list_update(self.submobjects, mobjects)
        return self

    def add_to_back(self, *mobjects):
        self.remove(*mobjects)
        self.submobjects = list(mobjects) + self.submobjects
        return self

    def remove(self, *mobjects):
        for mobject in mobjects:
            if mobject in self.submobjects:
                self.submobjects.remove(mobject)
        return self

    def get_array_attrs(self):
        return ["points"]

    def digest_mobject_attrs(self):
        """
        Ensures all attributes which are mobjects are included
        in the submobjects list.
        """
        mobject_attrs = filter(
            lambda x : isinstance(x, Mobject),
            self.__dict__.values()
        )
        self.submobjects = list_update(self.submobjects, mobject_attrs)
        return self

    def apply_over_attr_arrays(self, func):
        for attr in self.get_array_attrs():
            setattr(self, attr, func(getattr(self, attr)))
        return self

    def get_image(self):
        from camera import Camera
        camera = Camera()
        camera.capture_mobject(self)
        return Image.fromarray(camera.get_image())

    def show(self):
        self.get_image().show()

    def save_image(self, name = None):
        self.get_image().save(
            os.path.join(MOVIE_DIR, (name or str(self)) + ".png")
        )

    def copy(self):
        return deepcopy(self)

    def generate_target(self):
        self.target = None #Prevent exponential explosion
        self.target = self.copy()
        return self.target

    #### Transforming operations ######

    def apply_to_family(self, func):
        for mob in self.family_members_with_points():
            func(mob)

    def shift(self, *vectors):
        total_vector = reduce(op.add, vectors)
        for mob in self.family_members_with_points():        
           mob.points += total_vector
        return self        


    def scale(self, scale_factor, about_point = None):
        if about_point is not None:
            self.shift(-about_point)
        for mob in self.family_members_with_points():
            mob.points *= scale_factor
        if about_point is not None:
            self.shift(about_point)
        return self

    def rotate_about_origin(self, angle, axis = OUT, axes = []):
        if len(axes) == 0:
            axes = [axis]
        rot_matrix = np.identity(self.dim)
        for axis in axes:
            rot_matrix = np.dot(rot_matrix, rotation_matrix(angle, axis))
        t_rot_matrix = np.transpose(rot_matrix)
        for mob in self.family_members_with_points():
            mob.points = np.dot(mob.points, t_rot_matrix)
        return self

    def rotate(self, angle, axis = OUT, axes = [], about_point = None):
        if about_point is None:
            self.rotate_about_origin(angle, axis, axes)
        else:
            self.do_about_point(about_point, self.rotate, angle, axis, axes)
        return self

    def stretch(self, factor, dim):
        for mob in self.family_members_with_points():
            mob.points[:,dim] *= factor
        return self

    def apply_function(self, function):
        for mob in self.family_members_with_points():
            mob.points = np.apply_along_axis(function, 1, mob.points)
        return self

    def wag(self, direction = RIGHT, axis = DOWN, wag_factor = 1.0):
        for mob in self.family_members_with_points():
            alphas = np.dot(mob.points, np.transpose(axis))
            alphas -= min(alphas)
            alphas /= max(alphas)
            alphas = alphas**wag_factor
            mob.points += np.dot(
                alphas.reshape((len(alphas), 1)),
                np.array(direction).reshape((1, mob.dim))
            )
        return self

    def reverse_points(self):
        for mob in self.family_members_with_points():
            mob.apply_over_attr_arrays(
                lambda arr : np.array(list(reversed(arr)))
            )
        return self

    def repeat(self, count):
        """
        This can make transition animations nicer
        """
        def repeat_array(array):
            return reduce(
                lambda a1, a2 : np.append(a1, a2, axis = 0),
                [array]*count
            )
        for mob in self.family_members_with_points():
            mob.apply_over_attr_arrays(repeat_array)
        return self

    #### In place operations ######

    def do_about_point(self, point, method, *args, **kwargs):
        self.shift(-point)
        method(*args, **kwargs)
        self.shift(point)
        return self

    def do_in_place(self, method, *args, **kwargs):
        self.do_about_point(self.get_center(), method, *args, **kwargs)
        return self

    def rotate_in_place(self, angle, axis = OUT, axes = []):
        self.do_in_place(self.rotate, angle, axis, axes)
        return self

    def flip(self, axis = UP):
        self.rotate_in_place(np.pi, axis)
        return self

    def scale_in_place(self, scale_factor):
        self.do_in_place(self.scale, scale_factor)
        return self

    def scale_about_point(self, scale_factor, point):
        self.do_about_point(point, self.scale, scale_factor)
        return self

    def pose_at_angle(self):
        self.rotate_in_place(np.pi / 7, RIGHT+UP)
        return self

    def center(self):
        self.shift(-self.get_center())
        return self

    def align_on_border(self, direction, buff = DEFAULT_MOBJECT_TO_EDGE_BUFFER):
        """
        Direction just needs to be a vector pointing towards side or
        corner in the 2d plane.
        """
        target_point = np.sign(direction) * (SPACE_WIDTH, SPACE_HEIGHT, 0)
        point_to_align = self.get_critical_point(direction)
        shift_val = target_point - point_to_align - buff * np.array(direction)
        shift_val = shift_val * abs(np.sign(direction))
        self.shift(shift_val)
        return self

    def to_corner(self, corner = LEFT+DOWN, buff = DEFAULT_MOBJECT_TO_EDGE_BUFFER):
        return self.align_on_border(corner, buff)

    def to_edge(self, edge = LEFT, buff = DEFAULT_MOBJECT_TO_EDGE_BUFFER):
        return self.align_on_border(edge, buff)

    def next_to(self, mobject_or_point, 
                direction = RIGHT, 
                buff = DEFAULT_MOBJECT_TO_MOBJECT_BUFFER,
                aligned_edge = ORIGIN,
                align_using_submobjects = False,
                ):
        if isinstance(mobject_or_point, Mobject):
            mob = mobject_or_point
            target_point = mob.get_critical_point(
                aligned_edge+direction,
                use_submobject = align_using_submobjects
            )
        else:
            target_point = mobject_or_point
        point_to_align = self.get_critical_point(
            aligned_edge-direction,
            use_submobject = align_using_submobjects
        )
        self.shift(target_point - point_to_align + buff*direction)
        return self


    def shift_onto_screen(self, **kwargs):
        space_lengths = [SPACE_WIDTH, SPACE_HEIGHT]
        for vect in UP, DOWN, LEFT, RIGHT:
            dim = np.argmax(np.abs(vect))
            buff = kwargs.get("buff", DEFAULT_MOBJECT_TO_EDGE_BUFFER)
            max_val = space_lengths[dim] - buff
            edge_center = self.get_edge_center(vect)
            if np.dot(edge_center, vect) > max_val:
                self.to_edge(vect, **kwargs)
        return self

    def is_off_screen(self):
        if self.get_left()[0] > SPACE_WIDTH:
            return True
        if self.get_right()[0] < -SPACE_WIDTH:
            return True
        if self.get_bottom()[1] > SPACE_HEIGHT:
            return True
        if self.get_top()[1] < -SPACE_HEIGHT:
            return True
        return False

    def stretch_about_point(self, factor, dim, point):
        self.do_about_point(point, self.stretch, factor, dim)
        return self

    def stretch_in_place(self, factor, dim):
        self.do_in_place(self.stretch, factor, dim)
        return self

    def stretch_to_fit(self, length, dim, stretch = True):
        old_length = self.length_over_dim(dim)
        if old_length == 0:
            return self
        if stretch:
            self.stretch_in_place(length/old_length, dim)
        else:
            self.scale_in_place(length/old_length)
        return self

    def stretch_to_fit_width(self, width):
        return self.stretch_to_fit(width, 0)

    def stretch_to_fit_height(self, height):
        return self.stretch_to_fit(height, 1)

    def scale_to_fit_width(self, width):
        return self.stretch_to_fit(width, 0, stretch = False)

    def scale_to_fit_height(self, height):
        return self.stretch_to_fit(height, 1, stretch = False)

    def space_out_submobjects(self, factor = 1.5, **kwargs):
        self.scale_in_place(factor)
        for submob in self.submobjects:
            submob.scale_in_place(1./factor)
        return self

    def move_to(self, point_or_mobject, aligned_edge = ORIGIN):
        if isinstance(point_or_mobject, Mobject):
            target = point_or_mobject.get_critical_point(aligned_edge)
        else:
            target = point_or_mobject
        point_to_align = self.get_critical_point(aligned_edge)
        self.shift(target - point_to_align)
        return self

    def replace(self, mobject, dim_to_match = 0, stretch = False):
        if not mobject.get_num_points() and not mobject.submobjects:
            raise Warning("Attempting to replace mobject with no points")
            return self
        if stretch:
            self.stretch_to_fit_width(mobject.get_width())
            self.stretch_to_fit_height(mobject.get_height())
        else:
            self.stretch_to_fit(
                mobject.length_over_dim(dim_to_match),
                dim_to_match, 
                stretch = False
            )
        self.shift(mobject.get_center() - self.get_center())
        return self

    def position_endpoints_on(self, start, end):
        curr_vect = self.points[-1] - self.points[0]
        if np.all(curr_vect == 0):
            raise Exception("Cannot position endpoints of closed loop")
        target_vect = end - start
        self.scale(np.linalg.norm(target_vect)/np.linalg.norm(curr_vect))
        self.rotate(
            angle_of_vector(target_vect) - \
            angle_of_vector(curr_vect)
        )
        self.shift(start-self.points[0])
        return self

    ## Color functions

    def highlight(self, color = YELLOW_C, family = True, condition = None):
        """
        Condition is function which takes in one arguments, (x, y, z).
        """
        raise Exception("Not implemented")

    def gradient_highlight(self, *colors):
        self.submobject_gradient_highlight(*colors)
        return self

    def submobject_gradient_highlight(self, *colors):
        if len(colors) == 0:
            raise Exception("Need at least one color")
        elif len(colors) == 1:
            return self.highlight(*colors)

        mobs = self.family_members_with_points()
        new_colors = color_gradient(colors, len(mobs))
        for mob, color in zip(mobs, new_colors):
            mob.highlight(color, family = False)
        return self

    def set_color(self, color):
        self.highlight(color)
        self.color = Color(color)
        return self

    def to_original_color(self):
        self.highlight(self.color)
        return self

    def fade_to(self, color, alpha):
        for mob in self.family_members_with_points():
            start = color_to_rgb(mob.get_color())
            end = color_to_rgb(color)
            new_rgb = interpolate(start, end, alpha)
            mob.highlight(Color(rgb = new_rgb), family = False)
        return self

    def fade(self, darkness = 0.5):
        self.fade_to(BLACK, darkness)
        return self

    def get_color(self):
        return self.color
    ##

    def save_state(self):
        if hasattr(self, "saved_state"):
            #Prevent exponential growth of data
            self.saved_state = None
        self.saved_state = self.copy()
        return self

    def restore(self):
        if not hasattr(self, "saved_state") or self.save_state is None:
            raise Exception("Trying to restore without having saved")
        self.align_data(self.saved_state)
        for sm1, sm2 in zip(self.submobject_family(), self.saved_state.submobject_family()):
            sm1.interpolate(sm1, sm2, 1)
        return self

    def apply_complex_function(self, function, **kwargs):
        return self.apply_function(
            lambda (x, y, z) : complex_to_R3(function(complex(x, y))),
            **kwargs
        )

    def reduce_across_dimension(self, points_func, reduce_func, dim):
        try:
            points = self.get_points_defining_boundary()
            values = [points_func(points[:, dim])]
        except:
            values = []
        values += [
            mob.reduce_across_dimension(points_func, reduce_func, dim)
            for mob in self.submobjects
        ]
        try:
            return reduce_func(values)
        except:
            return 0

    def get_merged_array(self, array_attr):
        result = np.zeros((0, self.dim))
        for mob in self.family_members_with_points():
            result = np.append(result, getattr(mob, array_attr), 0)
        return result

    def get_all_points(self):
        return self.get_merged_array("points")

    ### Getters ###

    def get_points_defining_boundary(self):
        return self.points

    def get_num_points(self):
        return len(self.points)

    def get_critical_point(self, direction, use_submobject = False):
        if use_submobject:
            return self.get_submobject_critical_point(direction)
        result = np.zeros(self.dim)
        for dim in range(self.dim):
            if direction[dim] <= 0:
                min_point = self.reduce_across_dimension(np.min, np.min, dim)
            if direction[dim] >= 0:
                max_point = self.reduce_across_dimension(np.max, np.max, dim)

            if direction[dim] == 0:
                result[dim] = (max_point+min_point)/2
            elif direction[dim] < 0:
                result[dim] = min_point
            else:
                result[dim] = max_point
        return result

    def get_submobject_critical_point(self, direction):
        if len(self.split()) == 1:
            return self.get_critical_point(direction)
        index = np.argmax([
            np.dot(submob.get_center(), direction)
            for submob in self
        ])
        return self[index].get_critical_point(
            direction, use_submobject = True
        )

    # Pseudonyms for more general get_critical_point method
    def get_edge_center(self, direction):
        return self.get_critical_point(direction)

    def get_corner(self, direction):
        return self.get_critical_point(direction)

    def get_center(self):
        return self.get_critical_point(np.zeros(self.dim))

    def get_center_of_mass(self):
        return np.apply_along_axis(np.mean, 0, self.get_all_points())

    def get_boundary_point(self, direction):
        all_points = self.get_all_points()
        return all_points[np.argmax(np.dot(all_points, direction))]

    def get_top(self):
        return self.get_edge_center(UP)

    def get_bottom(self):
        return self.get_edge_center(DOWN)

    def get_right(self):
        return self.get_edge_center(RIGHT)

    def get_left(self):
        return self.get_edge_center(LEFT)

    def length_over_dim(self, dim):
        return (
            self.reduce_across_dimension(np.max, np.max, dim) - 
            self.reduce_across_dimension(np.min, np.min, dim)
        )

    def get_width(self):
        return self.length_over_dim(0)

    def get_height(self):
        return self.length_over_dim(1)

    def point_from_proportion(self, alpha):
        raise Exception("Not implemented")


    ## Family matters

    def __getitem__(self, index):
        return self.split()[index]

    def __iter__(self):
        return iter(self.split())

    def __len__(self):
        return len(self.split())

    def split(self):
        result = [self] if len(self.points) > 0 else []
        return result + self.submobjects

    def submobject_family(self):
        sub_families = map(Mobject.submobject_family, self.submobjects)
        all_mobjects = [self] + list(it.chain(*sub_families))
        return remove_list_redundancies(all_mobjects)

    def family_members_with_points(self):
        return filter(
            lambda m : m.get_num_points() > 0, 
            self.submobject_family()
        )

    def arrange_submobjects(self, direction = RIGHT, center = True, **kwargs):
        for m1, m2 in zip(self.submobjects, self.submobjects[1:]):
            m2.next_to(m1, direction, **kwargs)
        if center:
            self.center()
        return self

    def sort_submobjects(self, point_to_num_func = lambda p : p[0]):
        self.submobjects.sort(
            lambda *mobs : cmp(*[
                point_to_num_func(mob.get_center())
                for mob in mobs
            ])
        )
        return self

    ## Alignment  
    def align_data(self, mobject):
        self.align_submobjects(mobject)
        self.align_points(mobject)
        #Recurse
        for m1, m2 in zip(self.submobjects, mobject.submobjects):
            m1.align_data(m2)

    def get_point_mobject(self, center = None):
        """
        The simplest mobject to be transformed to or from self.
        Should by a point of the appropriate type
        """
        raise Exception("Not implemented")


    def align_points(self, mobject):
        count1 = self.get_num_points()
        count2 = mobject.get_num_points()
        if count1 < count2:
            self.align_points_with_larger(mobject)
        elif count2 < count1:
            mobject.align_points_with_larger(self)
        return self

    def align_points_with_larger(self, larger_mobject):
        raise Exception("Not implemented")

    def align_submobjects(self, mobject):
        #If one is empty, and the other is not,
        #push it into its submobject list
        self_has_points, mob_has_points = [
            mob.get_num_points() > 0
            for mob in self, mobject
        ]
        if self_has_points and not mob_has_points:
            mobject.null_point_align(self)
        elif mob_has_points and not self_has_points:
            self.null_point_align(mobject)
        self_count = len(self.submobjects)
        mob_count = len(mobject.submobjects)
        diff = abs(self_count-mob_count)
        if self_count < mob_count:
            self.add_n_more_submobjects(diff)
        elif mob_count < self_count:
            mobject.add_n_more_submobjects(diff)
        return self

    def null_point_align(self, mobject):
        """
        If self has no points, but needs to align 
        with mobject, which has points
        """
        if self.submobjects:
            mobject.push_self_into_submobjects()
        else:
            self.points = np.array([mobject.points[0]])
        return self


    def push_self_into_submobjects(self):
        copy = self.copy()
        copy.submobjects = []
        self.init_points()
        self.add(copy)
        return self

    def add_n_more_submobjects(self, n):
        curr = len(self.submobjects)
        if n > 0 and curr == 0:
            self.add(self.copy())
            n -= 1
            curr += 1
        indices = curr*np.arange(curr+n)/(curr+n)
        new_submobjects = []
        for index in indices:
            submob = self.submobjects[index]
            if submob in new_submobjects:
                submob = self.repeat_submobject(submob)
            new_submobjects.append(submob)
        self.submobjects = new_submobjects
        return self

    def repeat_submobject(self, submob):
        return submob.copy()

    def interpolate(self, mobject1, mobject2, 
                    alpha, path_func = straight_path):
        """
        Turns self into an interpolation between mobject1 
        and mobject2.
        """
        self.points = path_func(
            mobject1.points, mobject2.points, alpha
        )
        self.interpolate_color(mobject1, mobject2, alpha)

    def interpolate_color(self, mobject1, mobject2, alpha):
        pass #To implement in subclass

    def become_partial(self, mobject, a, b):
        """
        Set points in such a way as to become only
        part of mobject.  
        Inputs 0 <= a < b <= 1 determine what portion
        of mobject to become.
        """
        pass #To implement in subclasses

        #TODO, color?

    def pointwise_become_partial(self, mobject, a, b):
        pass #To implement in subclass



class Group(Mobject):
    #Alternate name to improve readibility in cases where
    #the mobject is used primarily for its submobject housing
    #functionality.
    pass