public-mirrors/3b1b-manim
1
0
Fork 0
mirror of https://github.com/3b1b/manim.git synced 2025-09-01 00:48:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #1415 from 3b1b/matrix-exp-development

Browse source 
Matrix exp development
This commit is contained in:
Grant Sanderson 2021-02-25 08:52:29 -08:00 • committed by GitHub
commit 531a031b50
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 297 additions and 230 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
manimlib/__main__.py
View file

@ -6,12 +6,12 @@ import manimlib.utils.init_config
def main():
args = manimlib.config.parse_cli()
if args.config:
manimlib.utils.init_config.init_customization()
else:
config = manimlib.config.get_configuration(args)
scenes = manimlib.extract_scene.main(config)
for scene in scenes:
scene.run()

51
manimlib/animation/indication.py
View file

@ -1,4 +1,5 @@
import numpy as np
import math
from manimlib.constants import *
from manimlib.animation.animation import Animation
@ -152,6 +153,56 @@ class ShowPassingFlash(ShowPartial):
submob.pointwise_become_partial(start, 0, 1)
class VShowPassingFlash(Animation):
CONFIG = {
"time_width": 0.3,
"taper_width": 0.1,
"remover": True,
}
def begin(self):
self.mobject.align_stroke_width_data_to_points()
super().begin()
def interpolate_submobject(self, submobject, starting_sumobject, alpha):
original_widths = starting_sumobject.get_stroke_widths()
# anchor_widths = np.array([*original_widths[0::3, 0], original_widths[-1, 0]])
anchor_widths = np.array([0, *original_widths[3::3, 0], 0])
n_anchors = len(anchor_widths)
time_width = self.time_width
# taper_width = self.taper_width
# Create a gaussian such that 3 sigmas out on either side
# will equals time_width * (number of points)
sigma = time_width / 6
mu = interpolate(-time_width / 2, 1 + time_width / 2, alpha)
offset = math.exp(-4.5) # 3 sigmas out
def kernel_func(x):
result = math.exp(-0.5 * ((x - mu) / sigma)**2) - offset
result = max(result, 0)
# if x < taper_width:
# result *= x / taper_width
# elif x > 1 - taper_width:
# result *= (1 - x) / taper_width
return result
kernel_array = np.array([
kernel_func(n / (n_anchors - 1))
for n in range(n_anchors)
])
scaled_widths = anchor_widths * kernel_array
new_widths = np.zeros(submobject.get_num_points())
new_widths[0::3] = scaled_widths[:-1]
new_widths[1::3] = (scaled_widths[:-1] + scaled_widths[1:]) / 2
new_widths[2::3] = scaled_widths[1:]
submobject.set_stroke(width=new_widths)
def finish(self):
super().finish()
for submob, start in self.get_all_families_zipped():
submob.match_style(start)
class ShowCreationThenDestruction(ShowPassingFlash):
CONFIG = {
"time_width": 2.0,

2
manimlib/constants.py
View file

@ -136,3 +136,5 @@ RED = RED_C
MAROON = MAROON_C
PURPLE = PURPLE_C
GREY = GREY_C
COLORMAP_3B1B = [BLUE_E, GREEN, YELLOW, RED]

16
manimlib/mobject/coordinate_systems.py
View file

@ -46,9 +46,15 @@ class CoordinateSystem():
"""Abbreviation for point_to_coords"""
return self.point_to_coords(point)
def get_origin(self):
return self.c2p(*[0] * self.dimension)
def get_axes(self):
raise Exception("Not implemented")
def get_all_ranges(self):
raise Exception("Not implemented")
def get_axis(self, index):
return self.get_axes()[index]
@ -319,6 +325,9 @@ class Axes(VGroup, CoordinateSystem):
def get_axes(self):
return self.axes
def get_all_ranges(self):
return [self.x_range, self.y_range]
def add_coordinate_labels(self,
x_values=None,
y_values=None,
@ -346,9 +355,11 @@ class ThreeDAxes(Axes):
def __init__(self, x_range=None, y_range=None, z_range=None, **kwargs):
Axes.__init__(self, x_range, y_range, **kwargs)
if z_range is not None:
self.z_range[:len(z_range)] = z_range
z_axis = self.create_axis(
z_range or self.z_range,
self.z_range,
self.z_axis_config,
self.depth,
)
@ -365,6 +376,9 @@ class ThreeDAxes(Axes):
for axis in self.axes:
axis.insert_n_curves(self.num_axis_pieces - 1)
def get_all_ranges(self):
return [self.x_range, self.y_range, self.z_range]
class NumberPlane(Axes):
CONFIG = {

33
manimlib/mobject/mobject.py
View file

@ -504,7 +504,7 @@ class Mobject(object):
self.refresh_has_updater_status()
if call_updater:
self.update()
self.update(dt=0)
return self
def remove_updater(self, update_function):
@ -841,7 +841,30 @@ class Mobject(object):
# Color functions
def set_rgba_array(self, color=None, opacity=None, name="rgbas", recurse=True):
def set_rgba_array(self, rgba_array, name="rgbas", recurse=False):
for mob in self.get_family(recurse):
mob.data[name] = np.array(rgba_array)
return self
def set_color_by_rgba_func(self, func, recurse=True):
"""
Func should take in a point in R3 and output an rgba value
"""
for mob in self.get_family(recurse):
rgba_array = [func(point) for point in mob.get_points()]
mob.set_rgba_array(rgba_array)
return self
def set_color_by_rgb_func(self, func, opacity=1, recurse=True):
"""
Func should take in a point in R3 and output an rgb value
"""
for mob in self.get_family(recurse):
rgba_array = [[*func(point), opacity] for point in mob.get_points()]
mob.set_rgba_array(rgba_array)
return self
def set_rgba_array_by_color(self, color=None, opacity=None, name="rgbas", recurse=True):
if color is not None:
rgbs = np.array([color_to_rgb(c) for c in listify(color)])
if opacity is not None:
@ -870,8 +893,8 @@ class Mobject(object):
return self
def set_color(self, color, opacity=None, recurse=True):
self.set_rgba_array(color, opacity, recurse=False)
# Recurse to submobjects differently from how set_rgba_array
self.set_rgba_array_by_color(color, opacity, recurse=False)
# Recurse to submobjects differently from how set_rgba_array_by_color
# in case they implement set_color differently
if recurse:
for submob in self.submobjects:
@ -879,7 +902,7 @@ class Mobject(object):
return self
def set_opacity(self, opacity, recurse=True):
self.set_rgba_array(color=None, opacity=opacity, recurse=False)
self.set_rgba_array_by_color(color=None, opacity=opacity, recurse=False)
if recurse:
for submob in self.submobjects:
submob.set_opacity(opacity, recurse=True)

30
manimlib/mobject/types/vectorized_mobject.py
View file

@ -106,24 +106,42 @@ class VMobject(Mobject):
self.set_flat_stroke(self.flat_stroke)
return self
def set_rgba_array(self, rgba_array, name=None, recurse=False):
if name is None:
names = ["fill_rgba", "stroke_rgba"]
else:
names = [name]
for name in names:
super().set_rgba_array(rgba_array, name, recurse)
return self
def set_fill(self, color=None, opacity=None, recurse=True):
self.set_rgba_array(color, opacity, 'fill_rgba', recurse)
self.set_rgba_array_by_color(color, opacity, 'fill_rgba', recurse)
return self
def set_stroke(self, color=None, width=None, opacity=None, background=None, recurse=True):
self.set_rgba_array(color, opacity, 'stroke_rgba', recurse)
self.set_rgba_array_by_color(color, opacity, 'stroke_rgba', recurse)
if width is not None:
for mob in self.get_family(recurse):
mob.data['stroke_width'] = np.array([
[width] for width in listify(width)
])
if isinstance(width, np.ndarray):
arr = width.reshape((len(width), 1))
else:
arr = np.array([[w] for w in listify(width)])
mob.data['stroke_width'] = arr
if background is not None:
for mob in self.get_family(recurse):
mob.draw_stroke_behind_fill = background
return self
def align_stroke_width_data_to_points(self, recurse=True):
for mob in self.get_family(recurse):
mob.data["stroke_width"] = resize_with_interpolation(
mob.data["stroke_width"], len(mob.get_points())
)
def set_style(self,
fill_color=None,
fill_opacity=None,
@ -259,7 +277,7 @@ class VMobject(Mobject):
return self.get_fill_color()
def has_stroke(self):
return any(self.get_stroke_widths()) and any(self.get_stroke_opacities())
return self.get_stroke_widths().any() and self.get_stroke_opacities().any()
def has_fill(self):
return any(self.get_fill_opacities())

363
manimlib/mobject/vector_field.py
View file

@ -1,66 +1,32 @@
import numpy as np
import os
import itertools as it
from PIL import Image
import random
from manimlib.constants import *
from manimlib.animation.composition import AnimationGroup
from manimlib.animation.indication import ShowPassingFlash
from manimlib.mobject.geometry import Vector
from manimlib.animation.indication import VShowPassingFlash
from manimlib.mobject.geometry import Arrow
from manimlib.mobject.types.vectorized_mobject import VGroup
from manimlib.mobject.types.vectorized_mobject import VMobject
from manimlib.utils.bezier import inverse_interpolate
from manimlib.utils.bezier import interpolate
from manimlib.utils.color import color_to_rgb
from manimlib.utils.color import rgb_to_color
from manimlib.utils.color import get_colormap_list
from manimlib.utils.config_ops import merge_dicts_recursively
from manimlib.utils.config_ops import digest_config
from manimlib.utils.rate_functions import linear
from manimlib.utils.simple_functions import sigmoid
from manimlib.utils.space_ops import get_norm
# from manimlib.utils.space_ops import normalize
DEFAULT_SCALAR_FIELD_COLORS = [BLUE_E, GREEN, YELLOW, RED]
def get_colored_background_image(scalar_field_func,
number_to_rgb_func,
pixel_height=DEFAULT_PIXEL_HEIGHT,
pixel_width=DEFAULT_PIXEL_WIDTH):
ph = pixel_height
pw = pixel_width
fw = FRAME_WIDTH
fh = FRAME_HEIGHT
points_array = np.zeros((ph, pw, 3))
x_array = np.linspace(-fw / 2, fw / 2, pw)
x_array = x_array.reshape((1, len(x_array)))
x_array = x_array.repeat(ph, axis=0)
y_array = np.linspace(fh / 2, -fh / 2, ph)
y_array = y_array.reshape((len(y_array), 1))
y_array.repeat(pw, axis=1)
points_array[:, :, 0] = x_array
points_array[:, :, 1] = y_array
scalars = np.apply_along_axis(scalar_field_func, 2, points_array)
rgb_array = number_to_rgb_func(scalars.flatten()).reshape((ph, pw, 3))
return Image.fromarray((rgb_array * 255).astype('uint8'))
def get_rgb_gradient_function(min_value=0, max_value=1,
colors=[BLUE, RED],
flip_alphas=True, # Why?
):
rgbs = np.array(list(map(color_to_rgb, colors)))
def get_vectorized_rgb_gradient_function(min_value, max_value, color_map):
rgbs = np.array(get_colormap_list(color_map))
def func(values):
alphas = inverse_interpolate(
min_value, max_value, np.array(values)
)
alphas = np.clip(alphas, 0, 1)
# if flip_alphas:
# alphas = 1 - alphas
scaled_alphas = alphas * (len(rgbs) - 1)
indices = scaled_alphas.astype(int)
next_indices = np.clip(indices + 1, 0, len(rgbs) - 1)
@ -71,29 +37,9 @@ def get_rgb_gradient_function(min_value=0, max_value=1,
return func
def get_color_field_image_file(scalar_func,
min_value=0, max_value=2,
colors=DEFAULT_SCALAR_FIELD_COLORS
):
# try_hash
np.random.seed(0)
sample_inputs = 5 * np.random.random(size=(10, 3)) - 10
sample_outputs = np.apply_along_axis(scalar_func, 1, sample_inputs)
func_hash = hash(
str(min_value) + str(max_value) + str(colors) + str(sample_outputs)
)
file_name = "%d.png" % func_hash
full_path = os.path.join(RASTER_IMAGE_DIR, file_name)
if not os.path.exists(full_path):
print("Rendering color field image " + str(func_hash))
rgb_gradient_func = get_rgb_gradient_function(
min_value=min_value,
max_value=max_value,
colors=colors
)
image = get_colored_background_image(scalar_func, rgb_gradient_func)
image.save(full_path)
return full_path
def get_rgb_gradient_function(min_value, max_value, color_map):
vectorized_func = get_vectorized_rgb_gradient_function(min_value, max_value, color_map)
return lambda value: vectorized_func([value])[0]
def move_along_vector_field(mobject, func):
@ -125,166 +71,195 @@ def move_points_along_vector_field(mobject, func):
return mobject
def get_sample_points_from_coordinate_system(coordinate_system, step_multiple):
ranges = []
for range_args in coordinate_system.get_all_ranges():
_min, _max, step = range_args
step *= step_multiple
ranges.append(np.arange(_min, _max + step, step))
return it.product(*ranges)
# Mobjects
class VectorField(VGroup):
CONFIG = {
"delta_x": 0.5,
"delta_y": 0.5,
"x_min": int(np.floor(-FRAME_WIDTH / 2)),
"x_max": int(np.ceil(FRAME_WIDTH / 2)),
"y_min": int(np.floor(-FRAME_HEIGHT / 2)),
"y_max": int(np.ceil(FRAME_HEIGHT / 2)),
"min_magnitude": 0,
"max_magnitude": 2,
"colors": DEFAULT_SCALAR_FIELD_COLORS,
"step_multiple": 0.5,
"magnitude_range": (0, 2),
"color_map": "3b1b_colormap",
# Takes in actual norm, spits out displayed norm
"length_func": lambda norm: 0.45 * sigmoid(norm),
"opacity": 1.0,
"vector_config": {},
}
def __init__(self, func, **kwargs):
def __init__(self, func, coordinate_system, **kwargs):
super().__init__(**kwargs)
self.func = func
self.rgb_gradient_function = get_rgb_gradient_function(
self.min_magnitude,
self.max_magnitude,
self.colors,
flip_alphas=False
self.coordinate_system = coordinate_system
self.value_to_rgb = get_rgb_gradient_function(
*self.magnitude_range, self.color_map,
)
x_range = np.arange(
self.x_min,
self.x_max + self.delta_x,
self.delta_x
)
y_range = np.arange(
self.y_min,
self.y_max + self.delta_y,
self.delta_y
)
for x, y in it.product(x_range, y_range):
point = x * RIGHT + y * UP
self.add(self.get_vector(point))
self.set_opacity(self.opacity)
def get_vector(self, point, **kwargs):
output = np.array(self.func(point))
norm = get_norm(output)
if norm == 0:
output *= 0
else:
output *= self.length_func(norm) / norm
vector_config = dict(self.vector_config)
vector_config.update(kwargs)
vect = Vector(output, **vector_config)
vect.shift(point)
fill_color = rgb_to_color(
self.rgb_gradient_function(np.array([norm]))[0]
samples = get_sample_points_from_coordinate_system(
coordinate_system, self.step_multiple
)
vect.set_color(fill_color)
self.add(*(
self.get_vector(coords)
for coords in samples
))
def get_vector(self, coords, **kwargs):
vector_config = merge_dicts_recursively(
self.vector_config,
kwargs
)
output = np.array(self.func(*coords))
norm = get_norm(output)
if norm > 0:
output *= self.length_func(norm) / norm
origin = self.coordinate_system.get_origin()
_input = self.coordinate_system.c2p(*coords)
_output = self.coordinate_system.c2p(*output)
vect = Arrow(
origin, _output, buff=0,
**vector_config
)
vect.shift(_input)
vect.set_rgba_array([[*self.value_to_rgb(norm), self.opacity]])
return vect
class StreamLines(VGroup):
CONFIG = {
# TODO, this is an awkward way to inherit
# defaults to a method.
"start_points_generator_config": {},
# Config for choosing start points
"x_min": -8,
"x_max": 8,
"y_min": -5,
"y_max": 5,
"delta_x": 0.5,
"delta_y": 0.5,
"step_multiple": 0.5,
"n_repeats": 1,
"noise_factor": None,
# Config for drawing lines
"dt": 0.05,
"virtual_time": 3,
"n_anchors_per_line": 100,
"arc_len": 3,
"max_time_steps": 200,
"n_samples_per_line": 10,
"cutoff_norm": 15,
# Style info
"stroke_width": 1,
"stroke_color": WHITE,
"color_by_arc_length": True,
# Min and max arc lengths meant to define
# the color range, should color_by_arc_length be True
"min_arc_length": 0,
"max_arc_length": 12,
"color_by_magnitude": False,
# Min and max magnitudes meant to define
# the color range, should color_by_magnitude be True
"min_magnitude": 0.5,
"max_magnitude": 1.5,
"colors": DEFAULT_SCALAR_FIELD_COLORS,
"cutoff_norm": 15,
"stroke_opacity": 1,
"color_by_magnitude": True,
"magnitude_range": (0, 2.0),
"taper_stroke_width": False,
"color_map": "3b1b_colormap",
}
def __init__(self, func, **kwargs):
VGroup.__init__(self, **kwargs)
def __init__(self, func, coordinate_system, **kwargs):
super().__init__(**kwargs)
self.func = func
dt = self.dt
self.coordinate_system = coordinate_system
self.draw_lines()
self.init_style()
start_points = self.get_start_points(
**self.start_points_generator_config
def point_func(self, point):
return self.coordinate_system.c2p(
*self.func(*self.coordinate_system.p2c(point))
)
for point in start_points:
def draw_lines(self):
lines = []
for point in self.get_start_points():
points = [point]
for t in np.arange(0, self.virtual_time, dt):
total_arc_len = 0
# for t in np.arange(0, self.virtual_time, self.dt):
for x in range(self.max_time_steps):
last_point = points[-1]
points.append(last_point + dt * func(last_point))
new_point = last_point + self.dt * self.point_func(last_point)
points.append(new_point)
total_arc_len += get_norm(new_point - last_point)
if get_norm(last_point) > self.cutoff_norm:
break
if total_arc_len > self.arc_len:
break
line = VMobject()
step = max(1, int(len(points) / self.n_anchors_per_line))
step = max(1, int(len(points) / self.n_samples_per_line))
line.set_points_smoothly(points[::step])
self.add(line)
self.set_stroke(self.stroke_color, self.stroke_width)
if self.color_by_arc_length:
len_to_rgb = get_rgb_gradient_function(
self.min_arc_length,
self.max_arc_length,
colors=self.colors,
)
for line in self:
arc_length = line.get_arc_length()
rgb = len_to_rgb([arc_length])[0]
color = rgb_to_color(rgb)
line.set_color(color)
elif self.color_by_magnitude:
image_file = get_color_field_image_file(
lambda p: get_norm(func(p)),
min_value=self.min_magnitude,
max_value=self.max_magnitude,
colors=self.colors,
)
self.color_using_background_image(image_file)
lines.append(line)
self.set_submobjects(lines)
def get_start_points(self):
x_min = self.x_min
x_max = self.x_max
y_min = self.y_min
y_max = self.y_max
delta_x = self.delta_x
delta_y = self.delta_y
n_repeats = self.n_repeats
noise_factor = self.noise_factor
cs = self.coordinate_system
sample_coords = get_sample_points_from_coordinate_system(
cs, self.step_multiple,
)
noise_factor = self.noise_factor
if noise_factor is None:
noise_factor = delta_y / 2
noise_factor = cs.x_range[2] * self.step_multiple * 0.5
return np.array([
x * RIGHT + y * UP + noise_factor * np.random.random(3)
for n in range(n_repeats)
for x in np.arange(x_min, x_max + delta_x, delta_x)
for y in np.arange(y_min, y_max + delta_y, delta_y)
cs.c2p(*coords) + noise_factor * np.random.random(3)
for n in range(self.n_repeats)
for coords in sample_coords
])
def init_style(self):
if self.color_by_magnitude:
values_to_rgbs = get_vectorized_rgb_gradient_function(
*self.magnitude_range, self.color_map,
)
cs = self.coordinate_system
for line in self.submobjects:
norms = [
get_norm(self.func(*cs.p2c(point)))
for point in line.get_points()
]
rgbs = values_to_rgbs(norms)
rgbas = np.zeros((len(rgbs), 4))
rgbas[:, :3] = rgbs
rgbas[:, 3] = self.stroke_opacity
line.set_rgba_array(rgbas, "stroke_rgba")
else:
self.set_stroke(self.stroke_color)
# TODO: Make it so that you can have a group of stream_lines
# varying in response to a changing vector field, and still
# animate the resulting flow
if self.taper_stroke_width:
width = [0, self.stroke_width, 0]
else:
width = self.stroke_width
self.set_stroke(width=width)
class AnimatedStreamLines(VGroup):
CONFIG = {
"lag_range": 4,
"line_anim_class": VShowPassingFlash,
"line_anim_config": {
"run_time": 4,
"rate_func": linear,
"time_width": 0.5,
},
}
def __init__(self, stream_lines, **kwargs):
super().__init__(**kwargs)
self.stream_lines = stream_lines
for line in stream_lines:
line.anim = self.line_anim_class(line, **self.line_anim_config)
line.anim.begin()
line.time = -self.lag_range * random.random()
self.add(line.anim.mobject)
self.add_updater(lambda m, dt: m.update(dt))
def update(self, dt):
stream_lines = self.stream_lines
for line in stream_lines:
line.time += dt
adjusted_time = max(line.time, 0) % line.anim.run_time
line.anim.update(adjusted_time / line.anim.run_time)
# TODO: This class should be deleted
class ShowPassingFlashWithThinningStrokeWidth(AnimationGroup):
CONFIG = {
"n_segments": 10,
@ -307,35 +282,3 @@ class ShowPassingFlashWithThinningStrokeWidth(AnimationGroup):
np.linspace(max_time_width, 0, self.n_segments)
)
])
# TODO, this is untested after turning it from a
# ContinualAnimation into a VGroup
class AnimatedStreamLines(VGroup):
CONFIG = {
"lag_range": 4,
"line_anim_class": ShowPassingFlash,
"line_anim_config": {
"run_time": 4,
"rate_func": linear,
"time_width": 0.3,
},
}
def __init__(self, stream_lines, **kwargs):
VGroup.__init__(self, **kwargs)
self.stream_lines = stream_lines
for line in stream_lines:
line.anim = self.line_anim_class(line, **self.line_anim_config)
line.anim.begin()
line.time = -self.lag_range * random.random()
self.add(line.anim.mobject)
self.add_updater(lambda m, dt: m.update(dt))
def update(self, dt):
stream_lines = self.stream_lines
for line in stream_lines:
line.time += dt
adjusted_time = max(line.time, 0) % line.anim.run_time
line.anim.update(adjusted_time / line.anim.run_time)

2
manimlib/scene/scene.py
View file

@ -117,7 +117,7 @@ class Scene(object):
self.stop_skipping()
self.linger_after_completion = False
self.update_frame()
from IPython.terminal.embed import InteractiveShellEmbed
shell = InteractiveShellEmbed()
# Have the frame update after each command

24
manimlib/utils/color.py
View file

@ -4,7 +4,9 @@ from colour import Color
import numpy as np
from manimlib.constants import WHITE
from manimlib.constants import COLORMAP_3B1B
from manimlib.utils.bezier import interpolate
from manimlib.utils.iterables import resize_with_interpolation
from manimlib.utils.simple_functions import clip_in_place
from manimlib.utils.space_ops import normalize
@ -114,10 +116,22 @@ def get_shaded_rgb(rgb, point, unit_normal_vect, light_source):
def get_colormap_list(map_name="viridis", n_colors=9):
"""
Options for map_name:
3b1b_colormap
magma
inferno
plasma
viridis
cividis
twilight
twilight_shifted
turbo
"""
from matplotlib.cm import get_cmap
rgbs = get_cmap(map_name).colors # Make more general?
return [
rgbs[int(n)]
for n in np.linspace(0, len(rgbs) - 1, n_colors)
]
if map_name == "3b1b_colormap":
rgbs = [color_to_rgb(color) for color in COLORMAP_3B1B]
else:
rgbs = get_cmap(map_name).colors # Make more general?
return resize_with_interpolation(np.array(rgbs), n_colors)

2
manimlib/utils/iterables.py
View file

@ -98,6 +98,8 @@ def resize_preserving_order(nparray, length):
def resize_with_interpolation(nparray, length):
if len(nparray) == length:
return nparray
if length == 0:
return np.zeros((0, *nparray.shape[1:]))
cont_indices = np.linspace(0, len(nparray) - 1, length)
return np.array([
(1 - a) * nparray[lh] + a * nparray[rh]