Improvements to camera.set_background_from_func, and incremental progress on WindingNumber

active_projects/WindingNumber.py

@@ -211,14 +211,34 @@ def rev_to_color(alpha):
 
     return interpolate_color(colors[start_index], colors[end_index], beta)
 
-def point_to_rev((x, y)):
+colorslist = map(color_to_rgba, ["#FF0000", ORANGE, YELLOW, "#00FF00", "#0000FF", "#FF00FF"])
+
+def rev_to_rgba(alpha):
+    # TODO: Merge with above
+    alpha = alpha % 1
+    colors = colorslist
+    num_colors = len(colors)
+    beta = (alpha % (1.0/num_colors)) * num_colors
+    start_index = int(np.floor(num_colors * alpha)) % num_colors
+    end_index = (start_index + 1) % num_colors
+
+    return interpolate(colors[start_index], colors[end_index], beta)
+
+def point_to_rev((x, y), allow_origin = False):
     # Warning: np.arctan2 would happily discontinuously returns the value 0 for (0, 0), due to 
     # design choices in the underlying atan2 library call, but for our purposes, this is 
     # illegitimate, and all winding number calculations must be set up to avoid this
-    if (x, y) == (0, 0):
+    if not(allow_origin) and (x, y) == (0, 0):
         print "Error! Angle of (0, 0) computed!"
-        return None
-    return np.true_divide(np.arctan2(y, x), TAU)
+        return
+    return fdiv(np.arctan2(y, x), TAU)
+
+def point_to_rgba(point):
+    rev = point_to_rev(point, allow_origin = True)
+    rgba = rev_to_rgba(rev)
+    base_size = np.sqrt(point[0]**2 + point[1]**2)
+    rescaled_size = np.sqrt(base_size/(base_size + 1))
+    return rgba * rescaled_size
 
 # Returns the value with the same fractional component as x, closest to m
 def resit_near(x, m):
@@ -232,7 +252,7 @@ def resit_near(x, m):
 def make_alpha_winder(func, start, end, num_checkpoints):
     check_points = [None for i in range(num_checkpoints)]
     check_points[0] = func(start)
-    step_size = np.true_divide(end - start, num_checkpoints)
+    step_size = fdiv(end - start, num_checkpoints)
     for i in range(num_checkpoints - 1):
         check_points[i + 1] = \
         resit_near(
@@ -334,7 +354,9 @@ class WalkerAnimation(Animation):
         self.rev_func = rev_func
         self.coords_to_point = coords_to_point
         self.compound_walker = VGroup()
-        self.compound_walker.walker = PiCreature(color = RED)
+        dot = Dot()
+        dot.scale(5)
+        self.compound_walker.walker = dot #PiCreature()
         self.compound_walker.walker.scale(scale_factor)
         self.compound_walker.arrow = Arrow(ORIGIN, RIGHT) #, buff = 0)
         self.compound_walker.digest_mobject_attrs()
@@ -403,21 +425,35 @@ def LinearWalker(
         number_update_func = number_update_func,
         **kwargs)
 
-class PiWalker(Scene):
+class ColorMappedByFuncScene(Scene):
+    CONFIG = {
+        "func" : lambda p : p
+    }
+
+    def construct(self):
+        self.num_plane = NumberPlane()
+        self.num_plane.fade()
+        self.add(self.num_plane)
+        self.camera.set_background_from_func(
+            lambda (x, y): point_to_rgba(
+                self.func(
+                    self.num_plane.point_to_coords(np.array([x, y, 0]))
+                )
+            )
+        )
+
+class PiWalker(ColorMappedByFuncScene):
     CONFIG = {
-        "func" : plane_func_from_complex_func(lambda c : c**2),
         "walk_coords" : [],
         "step_run_time" : 1
     }
 
     def construct(self):
+        ColorMappedByFuncScene.construct(self)
+        num_plane = self.num_plane
 
         rev_func = lambda p : point_to_rev(self.func(p))
 
-        num_plane = NumberPlane()
-        num_plane.fade()
-        self.add(num_plane)
-
         walk_coords = self.walk_coords
         for i in range(len(walk_coords)):
             start_x, start_y = start_coords = walk_coords[i]
@@ -425,6 +461,7 @@ class PiWalker(Scene):
             end_x, end_y = end_coords = walk_coords[(i + 1) % len(walk_coords)]
             end_point = num_plane.coords_to_point(end_x, end_y)
             self.play(
+                ShowCreation(Line(start_point, end_point), rate_func = None),
                 LinearWalker(
                     start_coords = start_coords, 
                     end_coords = end_coords,
@@ -432,7 +469,6 @@ class PiWalker(Scene):
                     rev_func = rev_func,
                     remover = (i < len(walk_coords) - 1)
                 ),
-                ShowCreation(Line(start_point, end_point), rate_func = None),
                 run_time = self.step_run_time)
 
         # TODO: Allow smooth paths instead of breaking them up into lines, and 
@@ -473,9 +509,8 @@ class PiWalkerCircle(PiWalker):
 # TODO: Perhaps restructure this to avoid using AnimationGroup, and instead 
 # use lists of animations or lists or other such data, to be merged and processed into parallel 
 # animations later
-class EquationSolver2d(Scene):
+class EquationSolver2d(ColorMappedByFuncScene):
     CONFIG = {
-        "func" : plane_poly_with_roots((1, 2), (-1, 3)),
         "initial_lower_x" : -5.1,
         "initial_upper_x" : 5.1,
         "initial_lower_y" : -3.1,
@@ -487,9 +522,8 @@ class EquationSolver2d(Scene):
     }
 
     def construct(self):
-        num_plane = NumberPlane()
-        num_plane.fade()
-        self.add(num_plane)
+        ColorMappedByFuncScene.construct(self)
+        num_plane = self.num_plane
 
         rev_func = lambda p : point_to_rev(self.func(p))
         clockwise_rev_func = lambda p : -rev_func(p)
@@ -613,7 +647,7 @@ class LinePulser(ContinualAnimation):
         end = self.line.get_end()
         for i in range(self.num_bullets):
             position = interpolate(start, end, 
-                np.true_divide((i + alpha),(self.num_bullets)))
+                fdiv((i + alpha),(self.num_bullets)))
             self.bullets[i].move_to(position)
             if self.output_func:
                 position_2d = (position[0], position[1])
@@ -635,7 +669,7 @@ class ArrowCircleTest(Scene):
             return x
 
         num_arrows = 8 * 3
-        arrows = [rev_rotate(base_arrow.copy(), (np.true_divide(i, num_arrows))) for i in range(num_arrows)]
+        arrows = [rev_rotate(base_arrow.copy(), (fdiv(i, num_arrows))) for i in range(num_arrows)]
         arrows_vgroup = VGroup(*arrows)
 
         self.play(ShowCreation(arrows_vgroup), run_time = 2.5, rate_func = None)
@@ -1098,4 +1132,15 @@ class DiffOdometer(OdometerScene):
 # TODO: Add to camera an option for low-quality background than other rendering, helpful
 # for previews
 
+####################
+
+class PureTest(Scene):
+    def construct(self):
+        point_list = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)]
+        output_list = map(lambda p : (p, point_to_rgba(p)), point_list)
+
+        print output_list
+
+        self.wait()
+
 # FIN

animation/simple_animations.py

@@ -446,9 +446,9 @@ class Succession(Animation):
     def jump_to_start_of_anim(self, index):
         if index != self.current_anim_index:
             self.mobject.remove(*self.mobject.submobjects) # Should probably have a cleaner "remove_all" method...
-            self.mobject.add(self.animations[index].mobject)
             for m in self.scene_mobjects_at_time[index].submobjects:
                 self.mobject.add(m)
+            self.mobject.add(self.animations[index].mobject)
 
         self.current_anim_index = index
         self.current_alpha = self.critical_alphas[index]

camera/camera.py

@@ -22,12 +22,13 @@ class Camera(object):
         "max_allowable_norm" : 2*SPACE_WIDTH,
         "image_mode" : "RGBA",
         "n_rgb_coords" : 4,
-        "background_alpha" : 0, #Out of 255
+        "background_alpha" : 0, #Out of color_max_val
         "pixel_array_dtype" : 'uint8'
     }
 
     def __init__(self, background = None, **kwargs):
         digest_config(self, kwargs, locals())
+        self.color_max_val = np.iinfo(self.pixel_array_dtype).max
         self.init_background()
         self.resize_space_shape()
         self.reset()
@@ -75,11 +76,39 @@ class Camera(object):
     def get_pixel_array(self):
         return self.pixel_array
 
-    def set_pixel_array(self, pixel_array):
-        self.pixel_array = np.array(pixel_array)
+    def convert_pixel_array(self, pixel_array, convert_from_floats = False):
+        retval = np.array(pixel_array)
+        if convert_from_floats:
+            retval = np.apply_along_axis(
+                lambda f : (f * self.color_max_val).astype(self.pixel_array_dtype),
+                2,
+                retval)
+        return retval
 
-    def set_background(self, pixel_array):
-        self.background = np.array(pixel_array)
+    def set_pixel_array(self, pixel_array, convert_from_floats = False):
+        self.pixel_array = self.convert_pixel_array(pixel_array, convert_from_floats)
+
+    def set_background(self, pixel_array, convert_from_floats = False):
+        self.background = self.convert_pixel_array(pixel_array, convert_from_floats)
+
+    def set_background_from_func(self, coords_to_colors_func):
+        """
+        Sets background by using coords_to_colors_func to determine each pixel's color. Each input 
+        to coords_to_colors_func is an (x, y) pair in space (in ordinary space coordinates; not 
+        pixel coordinates), and each output is expected to be an RGBA array of 4 floats.
+        """
+
+        print "Starting set_background_from_func"
+
+        coords = self.get_coords_of_all_pixels()
+        new_background = np.apply_along_axis(
+            coords_to_colors_func,
+            2,
+            coords
+        )
+        self.set_background(new_background, convert_from_floats = True)
+
+        print "Ending set_background_from_func"
 
     def reset(self):
         self.set_pixel_array(self.background)
@@ -173,7 +202,7 @@ class Camera(object):
         )
         fill = aggdraw.Brush(
             self.color_to_hex_l(self.get_fill_color(vmobject)),
-            opacity = int(255*vmobject.get_fill_opacity())
+            opacity = int(self.color_max_val*vmobject.get_fill_opacity())
         )
         return (pen, fill)
 
@@ -222,7 +251,7 @@ class Camera(object):
         )
         rgba_len = self.pixel_array.shape[2]
 
-        rgbas = (255*rgbas).astype('uint8')
+        rgbas = (self.color_max_val*rgbas).astype(self.pixel_array_dtype)
         target_len = len(pixel_coords)
         factor = target_len/len(rgbas)
         rgbas = np.array([rgbas]*factor).reshape((target_len, rgba_len))
@@ -311,7 +340,7 @@ class Camera(object):
 
     def overlay_rgba_array(self, arr):
         # """ Overlays arr onto self.pixel_array with relevant alphas"""
-        bg, fg = self.pixel_array/255.0, arr/255.0
+        bg, fg = fdiv(self.pixel_array, self.color_max_val), fdiv(arr, self.color_max_val)
         bga, fga = [arr[:,:,3:] for arr in bg, fg]
         alpha_sum = fga + (1-fga)*bga
         with np.errstate(divide = 'ignore', invalid='ignore'):
@@ -320,7 +349,7 @@ class Camera(object):
                 np.divide(bg[:,:,:3]*(1-fga)*bga, alpha_sum),
             ])
         bg[:,:,3:] = 1 - (1 - bga)*(1 - fga)
-        self.pixel_array = (255*bg).astype(self.pixel_array_dtype)
+        self.pixel_array = (self.color_max_val*bg).astype(self.pixel_array_dtype)
 
     def align_points_to_camera(self, points):
         ## This is where projection should live
@@ -382,36 +411,25 @@ class Camera(object):
         size = pixel_coords.size
         return pixel_coords.reshape((size/2, 2))
 
-    def get_points_of_all_pixels(self):
-        """
-        Returns an array a such that a[i, j] gives the spatial
-        coordinates associated with the pixel self.pixel_array[i, j]
-        """
-        shape = self.pixel_array.shape
-        indices = np.indices(shape[:2], dtype = 'float64')
-        all_point_coords = np.zeros((shape[0], shape[1], 3))
-        for i, space_dim in enumerate([SPACE_HEIGHT, SPACE_WIDTH]):
-            all_point_coords[:,:,i] = \
-                indices[i,:,:]*2*space_dim/shape[i] - space_dim
-        return all_point_coords
+    def get_coords_of_all_pixels(self):
+        uncentered_pixel_indices = np.indices(self.pixel_shape).transpose(1, 2, 0)
+        uncentered_space_indices = np.true_divide(
+            uncentered_pixel_indices * self.space_shape, 
+            self.pixel_shape)
+        # Could structure above line's computation slightly differently, but figured (without much 
+        # thought) multiplying by space_shape first, THEN dividing by pixel_shape, is probably 
+        # better than the other order, for avoiding underflow quantization in the division (whereas 
+        # overflow is unlikely to be a problem)
+        centered_space_indices = uncentered_space_indices - np.true_divide(self.space_shape, 2)
 
-    def set_background_by_color_function(self, point_to_rgba_func):
-        """
-        point_to_rgba_func should take in a point in R^2, an array
-        of two floats, and output a four element array representing 
-        rgba values, all between 0 and 1.
-        """
+        # Have to account for increasing y now going up instead of down, and also for swapping the 
+        # order of x and y
+        coords = np.apply_along_axis(
+            lambda (y, x) : (x, -y),
+            2,
+            centered_space_indices)
 
-        # point_to_rgba = lambda p : [1, 1, 0, 0]
-        def float_rgba_to_int_rgba(rgba):
-            return (255*np.array(rgba)).astype(self.pixel_array_dtype)
-
-        points_of_all_pixels = self.get_points_of_all_pixels()
-        self.set_background(np.apply_along_axis(
-            lambda p : float_rgba_to_int_rgba(point_to_rgba_func(p)),
-            2, points_of_all_pixels
-        ))
-        self.reset() # Perhaps this really belongs in set_background?
+        return coords
 
 
 class MovingCamera(Camera):
@@ -497,20 +515,24 @@ class MultiCamera(Camera):
                 shifted_camera.start_x:shifted_camera.end_x] \
             = shifted_camera.camera.pixel_array
 
-    def set_background(self, pixel_array):
+    def set_background(self, pixel_array, **kwargs):
         for shifted_camera in self.shifted_cameras:
             shifted_camera.camera.set_background(
                 pixel_array[
                     shifted_camera.start_y:shifted_camera.end_y, 
-                    shifted_camera.start_x:shifted_camera.end_x])
+                    shifted_camera.start_x:shifted_camera.end_x],
+                **kwargs
+            )
 
-    def set_pixel_array(self, pixel_array):
-        Camera.set_pixel_array(self, pixel_array)
+    def set_pixel_array(self, pixel_array, **kwargs):
+        Camera.set_pixel_array(self, pixel_array, **kwargs)
         for shifted_camera in self.shifted_cameras:
             shifted_camera.camera.set_pixel_array(
                 pixel_array[
                     shifted_camera.start_y:shifted_camera.end_y, 
-                    shifted_camera.start_x:shifted_camera.end_x])
+                    shifted_camera.start_x:shifted_camera.end_x],
+                **kwargs
+            )
 
     def init_background(self):
         Camera.init_background(self)

helpers.py

@@ -638,3 +638,6 @@ class DictAsObject(object):
     def __init__(self, dict):
          self.__dict__ = dict
 
+# Just to have a less heavyweight name for this extremely common operation
+def fdiv(a, b):
+    return np.true_divide(a,b)

topics/number_line.py

@@ -321,6 +321,12 @@ class NumberPlane(VMobject):
         y = new_point[1]/self.get_y_unit_size()
         return x, y
 
+    def point_to_coords_cheap(self, point):
+        new_point = point - self.center_point
+        x = new_point[0]/self.x_unit_size
+        y = new_point[1]/self.y_unit_size
+        return x, y
+
     def get_x_unit_size(self):
         return self.axes.get_width() / (2.0*self.x_radius)