diff --git a/big_ol_pile_of_manim_imports.py b/big_ol_pile_of_manim_imports.py index bc2ad705..746936c2 100644 --- a/big_ol_pile_of_manim_imports.py +++ b/big_ol_pile_of_manim_imports.py @@ -55,6 +55,7 @@ from manimlib.mobject.types.point_cloud_mobject import * from manimlib.mobject.types.vectorized_mobject import * from manimlib.mobject.mobject_update_utils import * from manimlib.mobject.value_tracker import * +from manimlib.mobject.vector_field import * from manimlib.for_3b1b_videos.common_scenes import * from manimlib.for_3b1b_videos.pi_creature import * diff --git a/manimlib/mobject/vector_field.py b/manimlib/mobject/vector_field.py new file mode 100644 index 00000000..31ad3c37 --- /dev/null +++ b/manimlib/mobject/vector_field.py @@ -0,0 +1,307 @@ +import numpy as np +import os +import itertools as it +from PIL import Image + +from manimlib.constants import * + +from manimlib.animation.composition import AnimationGroup +from manimlib.animation.indication import ShowPassingFlash +from manimlib.mobject.geometry import Vector +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.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 func(values): + alphas = inverse_interpolate(min_value, max_value, 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) + inter_alphas = scaled_alphas % 1 + inter_alphas = inter_alphas.repeat(3).reshape((len(indices), 3)) + result = interpolate(rgbs[indices], rgbs[next_indices], inter_alphas) + return result + + 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 move_along_vector_field(mobject, func): + mobject.add_updater( + lambda m, dt: m.shift( + func(m.get_center()) * dt + ) + ) + return mobject + + +def move_submobjects_along_vector_field(mobject, func): + def apply_nudge(mob, dt): + for submob in mob: + x, y = submob.get_center()[:2] + if abs(x) < FRAME_WIDTH and abs(y) < FRAME_HEIGHT: + submob.shift(func(submob.get_center()) * dt) + + mobject.add_updater(apply_nudge) + return mobject + + +def move_points_along_vector_field(mobject, func): + def apply_nudge(self, dt): + self.mobject.apply_function( + lambda p: p + func(p) * dt + ) + mobject.add_updater(apply_nudge) + return mobject + + +# Mobjects + +class StreamLines(VGroup): + CONFIG = { + # "start_points_generator": get_flow_start_points, + "start_points_generator_config": {}, + "dt": 0.05, + "virtual_time": 3, + "n_anchors_per_line": 100, + "stroke_width": 1, + "stroke_color": WHITE, + "color_lines_by_magnitude": False, + "min_magnitude": 0.5, + "max_magnitude": 1.5, + "colors": DEFAULT_SCALAR_FIELD_COLORS, + "cutoff_norm": 15, + } + + def __init__(self, func, **kwargs): + VGroup.__init__(self, **kwargs) + self.func = func + dt = self.dt + + start_points = self.get_start_points( + **self.start_points_generator_config + ) + for point in start_points: + points = [point] + for t in np.arange(0, self.virtual_time, dt): + last_point = points[-1] + points.append(last_point + dt * func(last_point)) + if get_norm(last_point) > self.cutoff_norm: + break + line = VMobject() + step = max(1, int(len(points) / self.n_anchors_per_line)) + line.set_points_smoothly(points[::step]) + self.add(line) + + self.set_stroke(self.stroke_color, self.stroke_width) + + if self.color_lines_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) + + def get_start_points(self, + x_min=-8, + x_max=8, + y_min=-5, + y_max=5, + delta_x=0.5, + delta_y=0.5, + n_repeats=1, + noise_factor=None): + if noise_factor is None: + noise_factor = delta_y / 2 + 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) + ]) + + +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, + # 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): + VGroup.__init__(self, **kwargs) + self.func = func + self.rgb_gradient_function = get_rgb_gradient_function( + self.min_magnitude, + self.max_magnitude, + self.colors, + flip_alphas=False + ) + 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] + ) + vect.set_color(fill_color) + return vect + + +# 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 +class ShowPassingFlashWithThinningStrokeWidth(AnimationGroup): + CONFIG = { + "n_segments": 10, + "time_width": 0.1, + "remover": True + } + + def __init__(self, vmobject, **kwargs): + digest_config(self, kwargs) + max_stroke_width = vmobject.get_stroke_width() + max_time_width = kwargs.pop("time_width", self.time_width) + AnimationGroup.__init__(self, *[ + ShowPassingFlash( + vmobject.deepcopy().set_stroke(width=stroke_width), + time_width=time_width, + **kwargs + ) + for stroke_width, time_width in zip( + np.linspace(0, max_stroke_width, self.n_segments), + 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.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) diff --git a/old_projects/div_curl.py b/old_projects/div_curl.py index ebb0cb05..8e8114a7 100644 --- a/old_projects/div_curl.py +++ b/old_projects/div_curl.py @@ -1,7 +1,6 @@ from big_ol_pile_of_manim_imports import * -DEFAULT_SCALAR_FIELD_COLORS = [BLUE_E, GREEN, YELLOW, RED] # Quick note to anyone coming to this file with the # intent of recreating animations from the video. Some @@ -24,22 +23,6 @@ RABBIT_COLOR = "#C6D6EF" # Helper functions -def get_flow_start_points(x_min=-8, x_max=8, - y_min=-5, y_max=5, - delta_x=0.5, delta_y=0.5, - n_repeats=1, - noise_factor=None - ): - if noise_factor is None: - noise_factor = delta_y / 2 - 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) - ]) - - def joukowsky_map(z): if z == 0: return 0 @@ -99,77 +82,6 @@ def cylinder_flow_magnitude_field(point): return get_norm(cylinder_flow_vector_field(point)) -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 func(values): - alphas = inverse_interpolate(min_value, max_value, 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) - inter_alphas = scaled_alphas % 1 - inter_alphas = inter_alphas.repeat(3).reshape((len(indices), 3)) - result = interpolate(rgbs[indices], rgbs[next_indices], inter_alphas) - return result - - 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 vec_tex(s): return "\\vec{\\textbf{%s}}" % s @@ -244,204 +156,6 @@ def preditor_prey_vector_field(point): # Mobjects - -class StreamLines(VGroup): - CONFIG = { - "start_points_generator": get_flow_start_points, - "start_points_generator_config": {}, - "dt": 0.05, - "virtual_time": 3, - "n_anchors_per_line": 100, - "stroke_width": 1, - "stroke_color": WHITE, - "color_lines_by_magnitude": True, - "min_magnitude": 0.5, - "max_magnitude": 1.5, - "colors": DEFAULT_SCALAR_FIELD_COLORS, - "cutoff_norm": 15, - } - - def __init__(self, func, **kwargs): - VGroup.__init__(self, **kwargs) - self.func = func - dt = self.dt - - start_points = self.start_points_generator( - **self.start_points_generator_config - ) - for point in start_points: - points = [point] - for t in np.arange(0, self.virtual_time, dt): - last_point = points[-1] - points.append(last_point + dt * func(last_point)) - if get_norm(last_point) > self.cutoff_norm: - break - line = VMobject() - step = max(1, int(len(points) / self.n_anchors_per_line)) - line.set_points_smoothly(points[::step]) - self.add(line) - - self.set_stroke(self.stroke_color, self.stroke_width) - - if self.color_lines_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) - - -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, - # Takes in actual norm, spits out displayed norm - "length_func": lambda norm: 0.5 * sigmoid(norm), - "stroke_color": BLACK, - "stroke_width": 0.5, - "fill_opacity": 1.0, - "vector_config": {}, - } - - def __init__(self, func, **kwargs): - VGroup.__init__(self, **kwargs) - self.func = func - self.rgb_gradient_function = get_rgb_gradient_function( - self.min_magnitude, - self.max_magnitude, - self.colors, - flip_alphas=False - ) - for x in np.arange(self.x_min, self.x_max, self.delta_x): - for y in np.arange(self.y_min, self.y_max, self.delta_y): - point = x * RIGHT + y * UP - self.add(self.get_vector(point)) - - 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] - ) - vect.set_color(fill_color) - vect.set_fill(opacity=self.fill_opacity) - vect.set_stroke( - self.stroke_color, - self.stroke_width - ) - return vect - - -# Redefining what was once a ContinualAnimation class -# as a function -def VectorFieldFlow(mobject, func): - mobject.add_updater( - lambda m, dt: m.shift( - func(m.get_center()) * dt - ) - ) - return mobject - - -# Redefining what was once a ContinualAnimation class -# as a function -def VectorFieldSubmobjectFlow(mobject, func): - def apply_nudge(mob, dt): - for submob in mob: - x, y = submob.get_center()[:2] - if abs(x) < FRAME_WIDTH and abs(y) < FRAME_HEIGHT: - submob.shift(func(submob.get_center()) * dt) - - mobject.add_updater(apply_nudge) - return mobject - - -# Redefining what was once a ContinualAnimation class -# as a function -def VectorFieldPointFlow(mobject, func): - def apply_nudge(self, dt): - self.mobject.apply_function( - lambda p: p + func(p) * dt - ) - mobject.add_updater(apply_nudge) - return mobject - - -# 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 -class ShowPassingFlashWithThinningStrokeWidth(AnimationGroup): - CONFIG = { - "n_segments": 10, - "time_width": 0.1, - "remover": True - } - - def __init__(self, vmobject, **kwargs): - digest_config(self, kwargs) - max_stroke_width = vmobject.get_stroke_width() - max_time_width = kwargs.pop("time_width", self.time_width) - AnimationGroup.__init__(self, *[ - ShowPassingFlash( - vmobject.deepcopy().set_stroke(width=stroke_width), - time_width=time_width, - **kwargs - ) - for stroke_width, time_width in zip( - np.linspace(0, max_stroke_width, self.n_segments), - 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.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 is untested after turning it from a # ContinualAnimation into a VGroup class JigglingSubmobjects(VGroup): @@ -3066,7 +2780,7 @@ class ShowTwoPopulations(Scene): self.start_num_rabbits * RIGHT + self.start_num_foxes * UP ) - self.add(VectorFieldFlow( + self.add(move_along_vector_field( phase_point, preditor_prey_vector_field, )) @@ -3420,7 +3134,7 @@ class PhaseSpaceOfPopulationModel(ShowTwoPopulations, PiCreatureScene, MovingCam dot_vector = new_dot_vector self.play(dot.move_to, dot_vector.get_end()) - dot_movement = VectorFieldFlow( + dot_movement = move_along_vector_field( dot, lambda p: 0.3 * vector_field.func(p) ) self.add(dot_movement) diff --git a/old_projects/turbulence.py b/old_projects/turbulence.py index 52ec382f..bd554414 100644 --- a/old_projects/turbulence.py +++ b/old_projects/turbulence.py @@ -1,7 +1,7 @@ from big_ol_pile_of_manim_imports import * from old_projects.div_curl import PureAirfoilFlow -from old_projects.div_curl import VectorFieldSubmobjectFlow -from old_projects.div_curl import VectorFieldPointFlow +from old_projects.div_curl import move_submobjects_along_vector_field +from old_projects.div_curl import move_points_along_vector_field from old_projects.div_curl import four_swirls_function from old_projects.lost_lecture import ShowWord @@ -836,7 +836,7 @@ class LaminarFlowLabel(Scene): class HighCurlFieldBreakingLayers(Scene): CONFIG = { - "flow_anim": VectorFieldSubmobjectFlow, + "flow_anim": move_submobjects_along_vector_field, } def construct(self): @@ -870,7 +870,7 @@ class HighCurlFieldBreakingLayers(Scene): class HighCurlFieldBreakingLayersLines(HighCurlFieldBreakingLayers): CONFIG = { - "flow_anim": VectorFieldPointFlow + "flow_anim": move_points_along_vector_field } def get_line(self):