manim/topics/complex_numbers.py

from constants import *


from animation.animation import Animation
from animation.creation import ShowCreation
from animation.movement import Homotopy
from animation.movement import SmoothedVectorizedHomotopy
from animation.transform import ApplyPointwiseFunction
from animation.transform import MoveToTarget
from mobject.coordinate_systems import NumberPlane
from mobject.coordinate_systems import ComplexPlane
from mobject.svg.tex_mobject import TexMobject
from mobject.svg.tex_mobject import TextMobject
from mobject.types.vectorized_mobject import VGroup
from scene.scene import Scene
from utils.config_ops import digest_config
from utils.config_ops import instantiate
from utils.paths import path_along_arc
from utils.space_ops import R3_to_complex
from utils.space_ops import complex_to_R3

class ComplexTransformationScene(Scene):
    CONFIG = {
        "plane_config" : {},
        "background_fade_factor" : 0.5,
        "use_multicolored_plane" : False,
        "vert_start_color" : BLUE, ##TODO
        "vert_end_color" : BLUE,
        "horiz_start_color" : BLUE,
        "horiz_end_color" : BLUE,
        "num_anchors_to_add_per_line" : 50,
        "post_transformation_stroke_width" : None,
        "default_apply_complex_function_kwargs" : {
            "run_time" : 5,
        },
        "background_label_scale_val" : 0.5,
        "include_coordinate_labels" : True,
    }
    def setup(self):
        self.foreground_mobjects = []
        self.transformable_mobjects = []
        self.add_background_plane()
        if self.include_coordinate_labels:
            self.add_coordinate_labels()

    def add_foreground_mobject(self, mobject):
        self.add_foreground_mobjects(mobject)

    def add_transformable_mobjects(self, *mobjects):
        self.transformable_mobjects += list(mobjects)
        self.add(*mobjects)

    def add_foreground_mobjects(self, *mobjects):
        self.foreground_mobjects += list(mobjects)
        Scene.add(self, *mobjects)

    def add(self, *mobjects):
        Scene.add(self, *list(mobjects)+self.foreground_mobjects)

    def play(self, *animations, **kwargs):
        Scene.play(
            self,
            *list(animations)+map(Animation, self.foreground_mobjects),
            **kwargs
        )

    def add_background_plane(self):
        background = ComplexPlane(**self.plane_config)
        background.fade(self.background_fade_factor)
        self.add(background)
        self.background = background

    def add_coordinate_labels(self):
        self.background.add_coordinates()
        self.add(self.background)

    def add_transformable_plane(self, **kwargs):
        self.plane = self.get_transformable_plane()
        self.add(self.plane)

    def get_transformable_plane(self, x_range = None, y_range = None):
        """
        x_range and y_range would be tuples (min, max)
        """
        plane_config = dict(self.plane_config)
        shift_val = ORIGIN
        if x_range is not None:
            x_min, x_max = x_range
            plane_config["x_radius"] = x_max - x_min
            shift_val += (x_max+x_min)*RIGHT/2.
        if y_range is not None:
            y_min, y_max = y_range
            plane_config["y_radius"] = y_max - y_min
            shift_val += (y_max+y_min)*UP/2.
        plane = ComplexPlane(**plane_config)
        plane.shift(shift_val)
        if self.use_multicolored_plane:
            self.paint_plane(plane)
        return plane

    def prepare_for_transformation(self, mob):
        if hasattr(mob, "prepare_for_nonlinear_transform"):
            mob.prepare_for_nonlinear_transform(
                self.num_anchors_to_add_per_line
            )
        #TODO...

    def paint_plane(self, plane):
        for lines in plane.main_lines, plane.secondary_lines:
            lines.set_color_by_gradient(
                self.vert_start_color,
                self.vert_end_color,
                self.horiz_start_color,
                self.horiz_end_color,
            )
        # plane.axes.set_color_by_gradient(
        #     self.horiz_start_color,
        #     self.vert_start_color
        # )

    def z_to_point(self, z):
        return self.background.number_to_point(z)
        
    def get_transformer(self, **kwargs):
        transform_kwargs = dict(self.default_apply_complex_function_kwargs)
        transform_kwargs.update(kwargs)
        plane = self.plane
        self.prepare_for_transformation(plane)
        transformer = VGroup(
            plane, *self.transformable_mobjects
        )        
        return transformer, transform_kwargs


    def apply_complex_function(self, func, added_anims = [], **kwargs):
        transformer, transform_kwargs = self.get_transformer(**kwargs)
        transformer.generate_target()
        #Rescale, apply function, scale back
        transformer.target.shift(-self.background.get_center_point())
        transformer.target.scale(1./self.background.unit_size)
        transformer.target.apply_complex_function(func)
        transformer.target.scale(self.background.unit_size)
        transformer.target.shift(self.background.get_center_point())
        #

        for mob in transformer.target[0].family_members_with_points():
            mob.make_smooth()
        if self.post_transformation_stroke_width is not None:
            transformer.target.set_stroke(width = self.post_transformation_stroke_width)
        self.play(
            MoveToTarget(transformer, **transform_kwargs),
            *added_anims
        )

    def apply_complex_homotopy(self, complex_homotopy, added_anims = [], **kwargs):
        transformer, transform_kwargs = self.get_transformer(**kwargs)
        def homotopy(x, y, z, t):
            output = complex_homotopy(complex(x, y), t)
            rescaled_output = self.z_to_point(output)
            return (rescaled_output.real, rescaled_output.imag, z)

        self.play(
            SmoothedVectorizedHomotopy(
                homotopy, transformer,
                **transform_kwargs
            ),
            *added_anims
        )

##### Unsure about what comes under here...

def complex_string(complex_num):
    return filter(lambda c : c not in "()", str(complex_num))

class ComplexFunction(ApplyPointwiseFunction):
    def __init__(self, function, mobject = ComplexPlane, **kwargs):
        if "path_func" not in kwargs:
            self.path_func = path_along_arc(
                np.log(function(complex(1))).imag
            )
        ApplyPointwiseFunction.__init__(
            self,
            lambda (x, y, z) : complex_to_R3(function(complex(x, y))),
            instantiate(mobject),
            **kwargs
        )

class ComplexHomotopy(Homotopy):
    def __init__(self, complex_homotopy, mobject = ComplexPlane, **kwargs):
        """
        Complex Hootopy a function Cx[0, 1] to C
        """
        def homotopy(event):
            x, y, z, t = event
            c = complex_homotopy((complex(x, y), t))
            return (c.real, c.imag, z)
        Homotopy.__init__(self, homotopy, mobject, *args, **kwargs)