3b1b-manim/topics/fractals.py

677 lines
20 KiB
Python

# from mobject.mobject import Mobject, Point, Mobject1D
from animation.creation import ShowCreation
from animation.transform import Transform
from characters import PiCreature
from characters import Randolph
from characters import get_all_pi_creature_modes
from mobject.vectorized_mobject import VGroup
from mobject.vectorized_mobject import VMobject
from mobject.vectorized_mobject import VectorizedPoint
from scene.scene import Scene
from topics.geometry import Circle
from topics.geometry import Line
from topics.geometry import Polygon
from topics.geometry import RegularPolygon
from topics.geometry import Square
from utils.bezier import interpolate
from utils.color import color_gradient
from utils.config_ops import digest_config
from utils.space_ops import center_of_mass
from utils.space_ops import compass_directions
from utils.space_ops import rotate_vector
from utils.space_ops import rotation_matrix
from constants import *
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 = 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 = np.linalg.norm(p1-p2)*length_scaling_factor
curr_length = np.linalg.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)/np.linalg.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.submobjects = [
fractalification_iteration(submob, dimension, num_inserted_anchors_range)
for submob in vmobject.submobjects
]
return vmobject
class SelfSimilarFractal(VMobject):
CONFIG = {
"order" : 5,
"num_subparts" : 3,
"height" : 4,
"colors" : [RED, WHITE],
"stroke_width" : 1,
"fill_opacity" : 1,
"propagate_style_to_family" : True,
}
def init_colors(self):
VMobject.init_colors(self)
self.set_color_by_gradient(*self.colors)
def generate_points(self):
order_n_self = self.get_order_n_self(self.order)
if self.order == 0:
self.submobjects = [order_n_self]
else:
self.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.scale_to_fit_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):
CONFIG = {
"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):
CONFIG = {
"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.submobject_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):
CONFIG = {
"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.submobject_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 generate_points(self):
random.seed(self.random_seed)
modes = get_all_pi_creature_modes()
seed = PiCreature(mode = self.start_mode)
seed.scale_to_fit_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.scale_to_fit_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):
CONFIG = {
"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_submobjects(RIGHT, buff = 0)
for p in p2, p3, p5, p6:
p.scale_to_fit_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):
CONFIG = {
"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_anchor_points(self, n):
if self.get_num_anchor_points() == 0:
self.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):
CONFIG = {
"radius" : 3,
"order" : 5,
"colors" : [RED, GREEN],
"num_submobjects" : 20,
"monochromatic" : False,
"order_to_stroke_width_map" : {
3 : 3,
4 : 2,
5 : 1,
},
"propagate_style_to_family" : True,
}
def generate_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.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):
CONFIG = {
"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 is "+":
step = rotate(step, self.angle)
elif letter is "-":
step = rotate(step, -self.angle)
else:
curr = curr + step
result.append(curr)
return np.array(result) - center_of_mass(result)
class SelfSimilarSpaceFillingCurve(FractalCurve):
CONFIG = {
"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):
CONFIG = {
"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):
CONFIG = {
"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):
CONFIG = {
"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):
CONFIG = {
"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):
# CONFIG = {
# "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):
CONFIG = {
"colors" : [WHITE, BLUE_D],
"axis_offset_pairs" : [
],
"scale_factor" : 3,
"radius_scale_factor" : 2/(3*np.sqrt(3)),
}
class FlowSnake(LindenmayerCurve):
CONFIG = {
"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):
CONFIG = {
"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):
CONFIG = {
"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):
CONFIG = {
"axiom" : "A--"
}
class QuadraticKoch(LindenmayerCurve):
CONFIG = {
"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):
CONFIG = {
"axiom" : "A+A+A+A"
}
class StellarCurve(LindenmayerCurve):
CONFIG = {
"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):
CONFIG = {
"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 = 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