CountingScene improvements by CrypticSwarm

topics/counting.py

@@ -0,0 +1,243 @@
+from helpers import *
+
+from mobject.tex_mobject import TexMobject, TextMobject
+from mobject import Mobject
+from mobject.vectorized_mobject import VMobject, VGroup
+
+from animation.transform import Transform, FadeIn, MoveToTarget
+from animation.simple_animations import ShowCreation
+from topics.geometry import Arrow, Circle, Dot
+# from topics.fractals import *
+from scene import Scene
+
+
+class CountingScene(Scene):
+    CONFIG = {
+        "digit_place_colors" : [YELLOW, MAROON_B, RED, GREEN, BLUE, PURPLE_D],
+        "counting_dot_starting_position" : (SPACE_WIDTH-1)*RIGHT + (SPACE_HEIGHT-1)*UP,
+        "count_dot_starting_radius" : 0.5,
+        "dot_configuration_height" : 2,
+        "ones_configuration_location" : UP+2*RIGHT,
+        "num_scale_factor" : 2,
+        "num_start_location" : 2*DOWN,
+    }
+    def setup(self):
+        self.dots = VGroup()
+        self.number = 0        
+        self.max_place = 0
+        self.number_mob = VGroup(TexMobject(str(self.number)))
+        self.number_mob.scale(self.num_scale_factor)
+        self.number_mob.shift(self.num_start_location)
+
+        self.dot_templates = []
+        self.dot_template_iterators = []
+        self.curr_configurations = []
+
+        self.arrows = VGroup()
+
+        self.add(self.number_mob)
+
+    def get_template_configuration(self, place):
+        #This should probably be replaced for non-base-10 counting scenes
+        down_right = (0.5)*RIGHT + (np.sqrt(3)/2)*DOWN
+        result = []
+        for down_right_steps in range(5):
+            for left_steps in range(down_right_steps):
+                result.append(
+                    down_right_steps*down_right + left_steps*LEFT
+                )
+        return reversed(result[:self.get_place_max(place)])
+
+    def get_dot_template(self, place):
+        #This should be replaced for non-base-10 counting scenes
+        down_right = (0.5)*RIGHT + (np.sqrt(3)/2)*DOWN
+        dots = VGroup(*[
+            Dot(
+                point, 
+                radius = 0.25,
+                fill_opacity = 0,
+                stroke_width = 2,
+                stroke_color = WHITE,
+            )
+            for point in self.get_template_configuration(place)
+        ])
+        dots.scale_to_fit_height(self.dot_configuration_height)
+        return dots
+
+    def add_configuration(self):
+        new_template = self.get_dot_template(len(self.dot_templates))
+        new_template.move_to(self.ones_configuration_location)
+        left_vect = (new_template.get_width()+LARGE_BUFF)*LEFT
+        new_template.shift(
+            left_vect*len(self.dot_templates)
+        )
+        self.dot_templates.append(new_template)
+        self.dot_template_iterators.append(
+            it.cycle(new_template)
+        )
+        self.curr_configurations.append(VGroup())
+
+    def count(self, max_val, run_time_per_anim = 1):
+        for x in range(max_val):
+            self.increment(run_time_per_anim)
+
+    def increment(self, run_time_per_anim = 1):
+        moving_dot = Dot(
+            self.counting_dot_starting_position,
+            radius = self.count_dot_starting_radius,
+            color = self.digit_place_colors[0],
+        )
+        moving_dot.generate_target()
+        moving_dot.set_fill(opacity = 0)
+        kwargs = {
+            "run_time" : run_time_per_anim
+        }
+
+        continue_rolling_over = True
+        first_move = True
+        place = 0
+        while continue_rolling_over:
+            added_anims = []                
+            if first_move:
+                added_anims += self.get_digit_increment_animations()
+                first_move = False
+            moving_dot.target.replace(
+                self.dot_template_iterators[place].next()
+            )
+            self.play(MoveToTarget(moving_dot), *added_anims, **kwargs)
+            self.curr_configurations[place].add(moving_dot)
+
+
+            if len(self.curr_configurations[place].split()) == self.get_place_max(place):
+                full_configuration = self.curr_configurations[place]
+                self.curr_configurations[place] = VGroup()
+                place += 1
+                center = full_configuration.get_center_of_mass()
+                radius = 0.6*max(
+                    full_configuration.get_width(),
+                    full_configuration.get_height(),
+                )
+                circle = Circle(
+                    radius = radius,
+                    stroke_width = 0,
+                    fill_color = self.digit_place_colors[place],
+                    fill_opacity = 0.5,
+                )
+                circle.move_to(center)
+                moving_dot = VGroup(circle, full_configuration)
+                moving_dot.generate_target()
+                moving_dot[0].set_fill(opacity = 0)
+            else:
+                continue_rolling_over = False
+
+    def get_digit_increment_animations(self):
+        result = []
+        self.number += 1
+        is_next_digit = self.is_next_digit()
+        if is_next_digit: self.max_place += 1
+        new_number_mob = self.get_number_mob(self.number)
+        new_number_mob.move_to(self.number_mob, RIGHT)
+        if is_next_digit:
+            self.add_configuration()
+            place = len(new_number_mob.split())-1
+            result.append(FadeIn(self.dot_templates[place]))
+            arrow = Arrow(
+                new_number_mob[place].get_top(),
+                self.dot_templates[place].get_bottom(),
+                color = self.digit_place_colors[place]
+            )
+            self.arrows.add(arrow)
+            result.append(ShowCreation(arrow))
+        result.append(Transform(
+            self.number_mob, new_number_mob,
+            submobject_mode = "lagged_start"
+        ))
+        return result
+
+    def get_number_mob(self, num):
+        result = VGroup()
+        place = 0
+        max_place = self.max_place
+        while place < max_place:
+            digit = TexMobject(str(self.get_place_num(num, place)))
+            if place >= len(self.digit_place_colors):
+                self.digit_place_colors += self.digit_place_colors
+            digit.highlight(self.digit_place_colors[place])
+            digit.scale(self.num_scale_factor)
+            digit.next_to(result, LEFT, buff = SMALL_BUFF, aligned_edge = DOWN)
+            result.add(digit)
+            place += 1
+        return result
+
+    def is_next_digit(self):
+        return False
+    def get_place_num(self, num, place):
+        return 0
+    def get_place_max(self, place):
+        return 0
+
+class PowerCounter(CountingScene):
+    def is_next_digit(self):
+        number = self.number
+        while number > 1:
+            if number%self.base != 0:
+                return False
+            number /= self.base
+        return True
+    def get_place_max(self, place):
+        return self.base
+    def get_place_num(self, num, place):
+        return (num / (self.base ** place)) % self.base
+
+class CountInDecimal(PowerCounter):
+    CONFIG = {
+        "base" : 10,
+    }
+    def construct(self):
+        for x in range(11):
+            self.increment()
+        for x in range(85):
+            self.increment(0.25)
+        for x in range(20):
+            self.increment()
+
+class CountInTernary(PowerCounter):
+    CONFIG = {
+        "base" : 3,
+        "dot_configuration_height" : 1,
+        "ones_configuration_location" : UP+4*RIGHT
+    }
+    def construct(self):
+        self.count(27)
+
+    # def get_template_configuration(self):
+    #     return [ORIGIN, UP]
+
+class CountInBinaryTo256(PowerCounter):
+    CONFIG = {
+        "base" : 2,
+        "dot_configuration_height" : 1,
+        "ones_configuration_location" : UP+5*RIGHT
+    }
+    def construct(self):
+        self.count(128, 0.3)
+
+    def get_template_configuration(self):
+        return [ORIGIN, UP]
+
+class FactorialBase(CountingScene):
+    CONFIG = {
+        "dot_configuration_height" : 1,
+        "ones_configuration_location" : UP+4*RIGHT
+    }
+    def construct(self):
+        self.count(30, 0.4)
+    def is_next_digit(self):
+        return self.number == self.factorial(self.max_place + 1)
+    def get_place_max(self, place):
+        return place + 2
+    def get_place_num(self, num, place):
+        return (num / self.factorial(place + 1)) % self.get_place_max(place)
+    def factorial(self, n):
+        if (n == 1): return 1
+        else: return n * self.factorial(n - 1)