fractals.py

from functools import reduce
import random

from manimlib.constants import *
# from manimlib.for_3b1b_videos.pi_creature import PiCreature
# from manimlib.for_3b1b_videos.pi_creature import Randolph
# from manimlib.for_3b1b_videos.pi_creature import get_all_pi_creature_modes
from manimlib.mobject.geometry import Circle
from manimlib.mobject.geometry import Polygon
from manimlib.mobject.geometry import RegularPolygon
from manimlib.mobject.types.vectorized_mobject import VGroup
from manimlib.mobject.types.vectorized_mobject import VMobject
from manimlib.utils.bezier import interpolate
from manimlib.utils.color import color_gradient
from manimlib.utils.dict_ops import digest_config
from manimlib.utils.space_ops import center_of_mass
from manimlib.utils.space_ops import compass_directions
from manimlib.utils.space_ops import rotate_vector
from manimlib.utils.space_ops import rotation_matrix


def rotate(points, angle=np.pi, axis=OUT):
    if axis is None:
        return points
    matrix = rotation_matrix(angle, axis)
    points = np.dot(points, np.transpose(matrix))
    return points


def fractalify(vmobject, order=3, *args, **kwargs):
    for x in range(order):
        fractalification_iteration(vmobject)
    return vmobject


def fractalification_iteration(vmobject, dimension=1.05, num_inserted_anchors_range=list(range(1, 4))):
    num_points = vmobject.get_num_points()
    if num_points > 0:
        # original_anchors = vmobject.get_anchors()
        original_anchors = [
            vmobject.point_from_proportion(x)
            for x in np.linspace(0, 1 - 1. / num_points, num_points)
        ]
        new_anchors = []
        for p1, p2, in zip(original_anchors, original_anchors[1:]):
            num_inserts = random.choice(num_inserted_anchors_range)
            inserted_points = [
                interpolate(p1, p2, alpha)
                for alpha in np.linspace(0, 1, num_inserts + 2)[1:-1]
            ]
            mass_scaling_factor = 1. / (num_inserts + 1)
            length_scaling_factor = mass_scaling_factor**(1. / dimension)
            target_length = get_norm(p1 - p2) * length_scaling_factor
            curr_length = get_norm(p1 - p2) * mass_scaling_factor
            # offset^2 + curr_length^2 = target_length^2
            offset_len = np.sqrt(target_length**2 - curr_length**2)
            unit_vect = (p1 - p2) / get_norm(p1 - p2)
            offset_unit_vect = rotate_vector(unit_vect, np.pi / 2)
            inserted_points = [
                point + u * offset_len * offset_unit_vect
                for u, point in zip(it.cycle([-1, 1]), inserted_points)
            ]
            new_anchors += [p1] + inserted_points
        new_anchors.append(original_anchors[-1])
        vmobject.set_points_as_corners(new_anchors)
    vmobject.set_submobjects([
        fractalification_iteration(
            submob, dimension, num_inserted_anchors_range)
        for submob in vmobject.submobjects
    ])
    return vmobject


class SelfSimilarFractal(VMobject):
    order = 5
    num_subparts = 3
    height = 4
    colors = [RED, WHITE]
    stroke_width = 1  # Not the right way to assign stroke width
    fill_opacity = 1  # Not the right way to assign fill opacity

    def init_colors(self):
        VMobject.init_colors(self)
        self.set_color_by_gradient(*self.colors)

    def init_points(self):
        order_n_self = self.get_order_n_self(self.order)
        if self.order == 0:
            self.set_submobjects([order_n_self])
        else:
            self.set_submobjects(order_n_self.submobjects)
        return self

    def get_order_n_self(self, order):
        if order == 0:
            result = self.get_seed_shape()
        else:
            lower_order = self.get_order_n_self(order - 1)
            subparts = [
                lower_order.copy()
                for x in range(self.num_subparts)
            ]
            self.arrange_subparts(*subparts)
            result = VGroup(*subparts)

        result.set_height(self.height)
        result.center()
        return result

    def get_seed_shape(self):
        raise Exception("Not implemented")

    def arrange_subparts(self, *subparts):
        raise Exception("Not implemented")


class Sierpinski(SelfSimilarFractal):
    def get_seed_shape(self):
        return Polygon(
            RIGHT, np.sqrt(3) * UP, LEFT,
        )

    def arrange_subparts(self, *subparts):
        tri1, tri2, tri3 = subparts
        tri1.move_to(tri2.get_corner(DOWN + LEFT), UP)
        tri3.move_to(tri2.get_corner(DOWN + RIGHT), UP)


class DiamondFractal(SelfSimilarFractal):
    num_subparts = 4
    height = 4
    colors = [GREEN_E, YELLOW]

    def get_seed_shape(self):
        return RegularPolygon(n=4)

    def arrange_subparts(self, *subparts):
        # VGroup(*subparts).rotate(np.pi/4)
        for part, vect in zip(subparts, compass_directions(start_vect=UP + RIGHT)):
            part.next_to(ORIGIN, vect, buff=0)
        VGroup(*subparts).rotate(np.pi / 4, about_point=ORIGIN)


class PentagonalFractal(SelfSimilarFractal):
    num_subparts = 5
    colors = [MAROON_B, YELLOW, RED]
    height = 6

    def get_seed_shape(self):
        return RegularPolygon(n=5, start_angle=np.pi / 2)

    def arrange_subparts(self, *subparts):
        for x, part in enumerate(subparts):
            part.shift(0.95 * part.get_height() * UP)
            part.rotate(2 * np.pi * x / 5, about_point=ORIGIN)


class PentagonalPiCreatureFractal(PentagonalFractal):
    def init_colors(self):
        SelfSimilarFractal.init_colors(self)
        internal_pis = [
            pi
            for pi in self.get_family()
            if isinstance(pi, PiCreature)
        ]
        colors = color_gradient(self.colors, len(internal_pis))
        for pi, color in zip(internal_pis, colors):
            pi.init_colors()
            pi.body.set_stroke(color, width=0.5)
            pi.set_color(color)

    def get_seed_shape(self):
        return Randolph(mode="shruggie")

    def arrange_subparts(self, *subparts):
        for part in subparts:
            part.rotate(2 * np.pi / 5, about_point=ORIGIN)
        PentagonalFractal.arrange_subparts(self, *subparts)


class PiCreatureFractal(VMobject):
    order = 7
    scale_val = 2.5
    start_mode = "hooray"
    height = 6
    colors = [
        BLUE_D, BLUE_B, MAROON_B, MAROON_D, GREY,
        YELLOW, RED, GREY_BROWN, RED, RED_E,
    ]
    random_seed = 0
    stroke_width = 0

    def init_colors(self):
        VMobject.init_colors(self)
        internal_pis = [
            pi
            for pi in self.get_family()
            if isinstance(pi, PiCreature)
        ]
        random.seed(self.random_seed)
        for pi in reversed(internal_pis):
            color = random.choice(self.colors)
            pi.set_color(color)
            pi.set_stroke(color, width=0)

    def init_points(self):
        random.seed(self.random_seed)
        modes = get_all_pi_creature_modes()
        seed = PiCreature(mode=self.start_mode)
        seed.set_height(self.height)
        seed.to_edge(DOWN)
        creatures = [seed]
        self.add(VGroup(seed))
        for x in range(self.order):
            new_creatures = []
            for creature in creatures:
                for eye, vect in zip(creature.eyes, [LEFT, RIGHT]):
                    new_creature = PiCreature(
                        mode=random.choice(modes)
                    )
                    new_creature.set_height(
                        self.scale_val * eye.get_height()
                    )
                    new_creature.next_to(
                        eye, vect,
                        buff=0,
                        aligned_edge=DOWN
                    )
                    new_creatures.append(new_creature)
                creature.look_at(random.choice(new_creatures))
            self.add_to_back(VGroup(*new_creatures))
            creatures = new_creatures

    # def init_colors(self):
    #     VMobject.init_colors(self)
    #     self.set_color_by_gradient(*self.colors)


class WonkyHexagonFractal(SelfSimilarFractal):
    num_subparts = 7

    def get_seed_shape(self):
        return RegularPolygon(n=6)

    def arrange_subparts(self, *subparts):
        for i, piece in enumerate(subparts):
            piece.rotate(i * np.pi / 12, about_point=ORIGIN)
        p1, p2, p3, p4, p5, p6, p7 = subparts
        center_row = VGroup(p1, p4, p7)
        center_row.arrange(RIGHT, buff=0)
        for p in p2, p3, p5, p6:
            p.set_width(p1.get_width())
        p2.move_to(p1.get_top(), DOWN + LEFT)
        p3.move_to(p1.get_bottom(), UP + LEFT)
        p5.move_to(p4.get_top(), DOWN + LEFT)
        p6.move_to(p4.get_bottom(), UP + LEFT)


class CircularFractal(SelfSimilarFractal):
    num_subparts = 3
    colors = [GREEN, BLUE, GREY]

    def get_seed_shape(self):
        return Circle()

    def arrange_subparts(self, *subparts):
        if not hasattr(self, "been_here"):
            self.num_subparts = 3 + self.order
            self.been_here = True
        for i, part in enumerate(subparts):
            theta = np.pi / self.num_subparts
            part.next_to(
                ORIGIN, UP,
                buff=self.height / (2 * np.tan(theta))
            )
            part.rotate(i * 2 * np.pi / self.num_subparts, about_point=ORIGIN)
        self.num_subparts -= 1

######## Space filling curves ############


class JaggedCurvePiece(VMobject):
    def insert_n_curves(self, n):
        if self.get_num_curves() == 0:
            self.set_points(np.zeros((1, 3)))
        anchors = self.get_anchors()
        indices = np.linspace(
            0, len(anchors) - 1, n + len(anchors)
        ).astype('int')
        self.set_points_as_corners(anchors[indices])


class FractalCurve(VMobject):
    radius = 3
    order = 5
    colors = [RED, GREEN]
    num_submobjects = 20
    monochromatic = False
    order_to_stroke_width_map = {
        3: 3,
        4: 2,
        5: 1,
    }

    def init_points(self):
        points = self.get_anchor_points()
        self.set_points_as_corners(points)
        if not self.monochromatic:
            alphas = np.linspace(0, 1, self.num_submobjects)
            for alpha_pair in zip(alphas, alphas[1:]):
                submobject = JaggedCurvePiece()
                submobject.pointwise_become_partial(
                    self, *alpha_pair
                )
                self.add(submobject)
            self.set_points(np.zeros((0, 3)))

    def init_colors(self):
        VMobject.init_colors(self)
        self.set_color_by_gradient(*self.colors)
        for order in sorted(self.order_to_stroke_width_map.keys()):
            if self.order >= order:
                self.set_stroke(width=self.order_to_stroke_width_map[order])

    def get_anchor_points(self):
        raise Exception("Not implemented")


class LindenmayerCurve(FractalCurve):
    axiom = "A"
    rule = {}
    scale_factor = 2
    radius = 3
    start_step = RIGHT
    angle = np.pi / 2

    def expand_command_string(self, command):
        result = ""
        for letter in command:
            if letter in self.rule:
                result += self.rule[letter]
            else:
                result += letter
        return result

    def get_command_string(self):
        result = self.axiom
        for x in range(self.order):
            result = self.expand_command_string(result)
        return result

    def get_anchor_points(self):
        step = float(self.radius) * self.start_step
        step /= (self.scale_factor**self.order)
        curr = np.zeros(3)
        result = [curr]
        for letter in self.get_command_string():
            if letter == "+":
                step = rotate(step, self.angle)
            elif letter == "-":
                step = rotate(step, -self.angle)
            else:
                curr = curr + step
                result.append(curr)
        return np.array(result) - center_of_mass(result)


class SelfSimilarSpaceFillingCurve(FractalCurve):
    offsets = []
    # keys must awkwardly be in string form...
    offset_to_rotation_axis = {}
    scale_factor = 2
    radius_scale_factor = 0.5

    def transform(self, points, offset):
        """
        How to transform the copy of points shifted by
        offset.  Generally meant to be extended in subclasses
        """
        copy = np.array(points)
        if str(offset) in self.offset_to_rotation_axis:
            copy = rotate(
                copy,
                axis=self.offset_to_rotation_axis[str(offset)]
            )
        copy /= self.scale_factor,
        copy += offset * self.radius * self.radius_scale_factor
        return copy

    def refine_into_subparts(self, points):
        transformed_copies = [
            self.transform(points, offset)
            for offset in self.offsets
        ]
        return reduce(
            lambda a, b: np.append(a, b, axis=0),
            transformed_copies
        )

    def get_anchor_points(self):
        points = np.zeros((1, 3))
        for count in range(self.order):
            points = self.refine_into_subparts(points)
        return points

    def generate_grid(self):
        raise Exception("Not implemented")


class HilbertCurve(SelfSimilarSpaceFillingCurve):
    offsets = [
        LEFT + DOWN,
        LEFT + UP,
        RIGHT + UP,
        RIGHT + DOWN,
    ]
    offset_to_rotation_axis = {
        str(LEFT + DOWN): RIGHT + UP,
        str(RIGHT + DOWN): RIGHT + DOWN,
    }


class HilbertCurve3D(SelfSimilarSpaceFillingCurve):
    offsets = [
        RIGHT + DOWN + IN,
        LEFT + DOWN + IN,
        LEFT + DOWN + OUT,
        RIGHT + DOWN + OUT,
        RIGHT + UP + OUT,
        LEFT + UP + OUT,
        LEFT + UP + IN,
        RIGHT + UP + IN,
    ]
    offset_to_rotation_axis_and_angle = {
        str(RIGHT + DOWN + IN): (LEFT + UP + OUT, 2 * np.pi / 3),
        str(LEFT + DOWN + IN): (RIGHT + DOWN + IN, 2 * np.pi / 3),
        str(LEFT + DOWN + OUT): (RIGHT + DOWN + IN, 2 * np.pi / 3),
        str(RIGHT + DOWN + OUT): (UP, np.pi),
        str(RIGHT + UP + OUT): (UP, np.pi),
        str(LEFT + UP + OUT): (LEFT + DOWN + OUT, 2 * np.pi / 3),
        str(LEFT + UP + IN): (LEFT + DOWN + OUT, 2 * np.pi / 3),
        str(RIGHT + UP + IN): (RIGHT + UP + IN, 2 * np.pi / 3),
    }
    # Rewrote transform method to include the rotation angle

    def transform(self, points, offset):
        copy = np.array(points)
        copy = rotate(
            copy,
            axis=self.offset_to_rotation_axis_and_angle[str(offset)][0],
            angle=self.offset_to_rotation_axis_and_angle[str(offset)][1],
        )
        copy /= self.scale_factor,
        copy += offset * self.radius * self.radius_scale_factor
        return copy


class PeanoCurve(SelfSimilarSpaceFillingCurve):
    colors = [PURPLE, TEAL]
    offsets = [
        LEFT + DOWN,
        LEFT,
        LEFT + UP,
        UP,
        ORIGIN,
        DOWN,
        RIGHT + DOWN,
        RIGHT,
        RIGHT + UP,
    ]
    offset_to_rotation_axis = {
        str(LEFT): UP,
        str(UP): RIGHT,
        str(ORIGIN): LEFT + UP,
        str(DOWN): RIGHT,
        str(RIGHT): UP,
    }
    scale_factor = 3
    radius_scale_factor = 2.0 / 3


class TriangleFillingCurve(SelfSimilarSpaceFillingCurve):
    colors = [MAROON, YELLOW]
    offsets = [
        LEFT / 4. + DOWN / 6.,
        ORIGIN,
        RIGHT / 4. + DOWN / 6.,
        UP / 3.,
    ]
    offset_to_rotation_axis = {
        str(ORIGIN): RIGHT,
        str(UP / 3.): UP,
    }
    scale_factor = 2
    radius_scale_factor = 1.5


class HexagonFillingCurve(SelfSimilarSpaceFillingCurve):
    start_color = WHITE
    end_color = BLUE_D
    axis_offset_pairs = [
        (None,                1.5*DOWN + 0.5*np.sqrt(3)*LEFT),
        (UP+np.sqrt(3)*RIGHT, 1.5*DOWN + 0.5*np.sqrt(3)*RIGHT),
        (np.sqrt(3)*UP+RIGHT, ORIGIN),
        ((UP, RIGHT),         np.sqrt(3)*LEFT),
        (None,                1.5*UP + 0.5*np.sqrt(3)*LEFT),
        (None,                1.5*UP + 0.5*np.sqrt(3)*RIGHT),
        (RIGHT,               np.sqrt(3)*RIGHT),
    ]
    scale_factor = 3
    radius_scale_factor = 2/(3*np.sqrt(3))

    def refine_into_subparts(self, points):
        return SelfSimilarSpaceFillingCurve.refine_into_subparts(
            self,
            rotate(points, np.pi/6, IN)
        )


class UtahFillingCurve(SelfSimilarSpaceFillingCurve):
    colors = [WHITE, BLUE_D]
    axis_offset_pairs = []
    scale_factor = 3
    radius_scale_factor = 2 / (3 * np.sqrt(3))


class FlowSnake(LindenmayerCurve):
    colors = [YELLOW, GREEN]
    axiom = "A"
    rule = {
        "A": "A-B--B+A++AA+B-",
        "B": "+A-BB--B-A++A+B",
    }
    radius = 6  # TODO, this is innaccurate
    scale_factor = np.sqrt(7)
    start_step = RIGHT
    angle = -np.pi / 3

    def __init__(self, **kwargs):
        LindenmayerCurve.__init__(self, **kwargs)
        self.rotate(-self.order * np.pi / 9, about_point=ORIGIN)


class SierpinskiCurve(LindenmayerCurve):
    colors = [RED, WHITE]
    axiom = "B"
    rule = {
        "A": "+B-A-B+",
        "B": "-A+B+A-",
    }
    radius = 6  # TODO, this is innaccurate
    scale_factor = 2
    start_step = RIGHT
    angle = -np.pi / 3


class KochSnowFlake(LindenmayerCurve):
    colors = [BLUE_D, WHITE, BLUE_D]
    axiom = "A--A--A--"
    rule = {
        "A": "A+A--A+A"
    }
    radius = 4
    scale_factor = 3
    start_step = RIGHT
    angle = np.pi / 3
    order_to_stroke_width_map = {
        3: 3,
        5: 2,
        6: 1,
    }

    def __init__(self, **kwargs):
        digest_config(self, kwargs)
        self.scale_factor = 2 * (1 + np.cos(self.angle))
        LindenmayerCurve.__init__(self, **kwargs)


class KochCurve(KochSnowFlake):
    axiom = "A--"


class QuadraticKoch(LindenmayerCurve):
    colors = [YELLOW, WHITE, MAROON_B]
    axiom = "A"
    rule = {"A": "A+A-A-AA+A+A-A"}
    radius = 4
    scale_factor = 4
    start_step = RIGHT
    angle = np.pi / 2


class QuadraticKochIsland(QuadraticKoch):
    axiom = "A+A+A+A"


class StellarCurve(LindenmayerCurve):
    start_color = RED
    end_color = BLUE_E
    rule = {
        "A": "+B-A-B+A-B+",
        "B": "-A+B+A-B+A-",
    }
    scale_factor = 3
    angle = 2 * np.pi / 5


class SnakeCurve(FractalCurve):
    start_color = BLUE
    end_color = YELLOW

    def get_anchor_points(self):
        result = []
        resolution = 2**self.order
        step = 2.0 * self.radius / resolution
        lower_left = ORIGIN + \
            LEFT * (self.radius - step / 2) + \
            DOWN * (self.radius - step / 2)

        for y in range(resolution):
            x_range = list(range(resolution))
            if y % 2 == 0:
                x_range.reverse()
            for x in x_range:
                result.append(
                    lower_left + x * step * RIGHT + y * step * UP
                )
        return result