Begin incorporating div_curl vector field functionality into the main library

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 *

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)

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)

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 move_submobjects_along_vector_field
-from old_projects.div_curl import VectorFieldPointFlow
+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):