Adds DiagramOfProgression

from_3b1b/active/chess.py

@@ -12,6 +12,11 @@ def string_to_bools(message):
     return [bool(int(b)) for b in bits]
 
 
+def layer_mobject(mobject, nudge=1e-6):
+    for i, sm in enumerate(mobject.family_members_with_points()):
+        sm.shift(i * nudge * OUT)
+
+
 class Chessboard(SGroup):
     CONFIG = {
         "shape": (8, 8),
@@ -69,11 +74,14 @@ class Coin(Group):
                 TextMobject("T"),
             )
             for label, vect in zip(labels, [OUT, IN]):
-                label.shift(1.01 * vect)
+                label.shift(1.02 * vect)
                 label.set_height(0.8)
             labels[1].flip(RIGHT)
-            self.add(*labels)
+            labels.apply_depth_test()
-            self.labels = labels
+        else:
+            labels = Group(Mobject(), Mobject())
+        self.add(*labels)
+        self.labels = labels
 
         self.set_height(self.height)
         self.set_depth(self.depth, stretch=True)
@@ -81,6 +89,9 @@ class Coin(Group):
     def is_heads(self):
         return self.top.get_center()[2] > self.bottom.get_center()[2]
 
+    def flip(self, axis=RIGHT):
+        super().flip(axis)
+
 
 class CoinsOnBoard(Group):
     CONFIG = {
@@ -100,9 +111,31 @@ class CoinsOnBoard(Group):
     def flip_at_random(self, p=0.5):
         for coin in self:
             if random.random() < p:
-                coin.flip(RIGHT)
+                coin.flip()
         return self
 
+    def flip_by_message(self, message):
+        heads = string_to_bools(message)
+        for coin, head in zip(self, heads):
+            if not head:
+                coin.flip()
+        return self
+
+
+class Key(SVGMobject):
+    CONFIG = {
+        "file_name": "key",
+        "fill_color": GOLD,
+        "fill_opacity": 1,
+        "stroke_color": GOLD,
+        "stroke_width": 0,
+        "gloss": 0.5,
+    }
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.rotate(PI / 2, OUT)
+
 
 class FlipCoin(Animation):
     CONFIG = {
@@ -150,7 +183,7 @@ class IntroducePuzzle(Scene):
         grid.set_gloss(0.5)
         grid.prepare_for_nonlinear_transform(0)
 
-        coins = CoinsOnBoard(chessboard, include_labels=False)
+        coins = CoinsOnBoard(chessboard)
         coins.set_gloss(0.2)
         coins_to_flip = Group()
         head_bools = string_to_bools('3b1b  :)')
@@ -192,11 +225,7 @@ class IntroducePuzzle(Scene):
         self.wait()
 
         # Show key
-        key = SVGMobject("key")
+        key = Key()
-        key.set_stroke(GOLD, 0)
-        key.set_fill(GOLD, 1)
-        key.set_gloss(0.5)
-        key.rotate(PI / 2, OUT)
         key.rotate(PI / 4, RIGHT)
         key.move_to(3 * OUT)
         key.scale(0.8)
@@ -373,19 +402,153 @@ class FromCoinToSquareMapsSingleFlips(FromCoinToSquareMaps):
     }
 
 
-class DiagramOfProgression(Scene):
+class DiagramOfProgression(ThreeDScene):
     def construct(self):
-        stages = VGroup(*[Rectangle(4, 2) for x in range(4)])
+        # Setup panels
-        stages.arrange_in_grid(buff=1)
+        P1_COLOR = BLUE_C
-        stages.set_height(6)
+        P2_COLOR = RED
-        stages.to_edge(DOWN)
 
+        rect = Rectangle(4, 2)
+        rect.set_fill(GREY_E, 1)
+        panels = Group()
+        for x in range(4):
+            panels.add(Group(rect.copy()))
+        panels.arrange_in_grid(buff=1)
+        panels[::2].shift(0.5 * LEFT)
+        panels.set_width(FRAME_WIDTH - 2)
+        panels.center().to_edge(DOWN)
+
+        chessboard = Chessboard()
+        chessboard.set_height(0.9 * panels[0].get_height())
+        coins = CoinsOnBoard(
+            chessboard,
+            coin_config={
+                "disk_resolution": (2, 25),
+                "include_labels": False,
+            }
+        )
+        coins.flip_by_message("Tau > Pi")
+
+        for panel in panels[1:]:
+            cb = chessboard.copy()
+            co = coins.copy()
+            cb.next_to(panel.get_right(), LEFT)
+            co.next_to(cb, OUT, 0)
+            panel.chessboard = cb
+            panel.coins = co
+            panel.add(cb, co)
+
+        kw = {
+            "tex_to_color_map": {
+                "Prisoner 1": P1_COLOR,
+                "Prisoner 2": P2_COLOR,
+            }
+        }
         titles = VGroup(
-            TextMobject("Prisoners conspire"),
+            TextMobject("Prisoners conspire", **kw),
-            TextMobject("Prisoner 1 sees key"),
+            TextMobject("Prisoner 1 sees key", **kw),
-            TextMobject("Prisoner 1 flips coin"),
+            TextMobject("Prisoner 1 flips coin", **kw),
+            TextMobject("Prisoner 2 guesses key square", **kw),
         )
 
-        self.add(stages)
+        for panel, title in zip(panels, titles):
+            title.next_to(panel, UP)
+            panel.title = title
+            panel.add(title)
 
-        self.embed()
+        # Darken first chessboard
+        for coin in panels[1].coins:
+            coin.remove(coin.edge)
+            if coin.is_heads():
+                coin.remove(coin.bottom)
+                coin.remove(coin.labels[1])
+            else:
+                coin.remove(coin.top)
+                coin.remove(coin.labels[0])
+            coin.set_opacity(0.25)
+
+        # Add characters
+        prisoner1 = PiCreature(color=P1_COLOR)
+        prisoner2 = PiCreature(color=P2_COLOR)
+        pis = VGroup(prisoner1, prisoner2)
+        pis.arrange(RIGHT, buff=1)
+        pis.set_height(1.5)
+
+        p0_pis = pis.copy()
+        p0_pis.set_height(2.0, about_edge=DOWN)
+        p0_pis[1].flip()
+        p0_pis.next_to(panels[0].get_bottom(), UP, SMALL_BUFF)
+        p0_pis[0].change("pondering", p0_pis[1].eyes)
+        p0_pis[1].change("speaking", p0_pis[0].eyes)
+        panels[0].add(p0_pis)
+
+        p1_pi = pis[0].copy()
+        p1_pi.next_to(panels[1].get_corner(DL), UR, SMALL_BUFF)
+        p1_pi.change("happy")
+        key = Key()
+        key.set_height(0.5)
+        key.next_to(p1_pi, UP)
+        key.set_color(YELLOW)
+        key_cube = panels[1].chessboard[18]
+        key_square = Square()
+        key_square.replace(key_cube)
+        key_square.set_stroke(width=3)
+        key_square.match_color(key)
+        p1_pi.look_at(key_square)
+        key_arrow = Arrow(
+            key.get_right() + SMALL_BUFF * UP,
+            key_square.get_corner(UL),
+            path_arc=-45 * DEGREES,
+            buff=SMALL_BUFF
+        )
+        key_arrow.tip.set_stroke(width=0)
+        key_arrow.match_color(key)
+        panels[1].add(p1_pi, key)
+
+        p2_pi = pis[0].copy()
+        p2_pi.next_to(panels[2].get_corner(DL), UR, SMALL_BUFF)
+        p2_pi.change("tease")
+        flip_coin = panels[2].coins[38]
+        panels[3].coins[38].flip()
+        flip_square = Square()
+        flip_square.replace(panels[2].chessboard[38])
+        flip_square.set_stroke(BLUE, 5)
+        for coin in panels[2].coins:
+            if coin is not flip_coin:
+                coin.remove(coin.edge)
+                if coin.is_heads():
+                    coin.remove(coin.bottom)
+                    coin.remove(coin.labels[1])
+                else:
+                    coin.remove(coin.top)
+                    coin.remove(coin.labels[0])
+                coin.set_opacity(0.25)
+        panels[2].add(p2_pi)
+
+        p3_pi = pis[1].copy()
+        p3_pi.next_to(panels[3].get_corner(DL), UR, SMALL_BUFF)
+        p3_pi.shift(MED_LARGE_BUFF * RIGHT)
+        p3_pi.change("confused")
+        panels[3].add(p3_pi)
+
+        # Animate each panel in
+        self.play(FadeIn(panels[1], DOWN))
+        self.play(
+            ShowCreation(key_arrow),
+            FadeInFromLarge(key_square),
+        )
+        self.wait()
+
+        self.play(FadeIn(panels[2], UP))
+        self.play(
+            ShowCreation(flip_square),
+            FlipCoin(flip_coin),
+            p2_pi.look_at, flip_coin,
+        )
+        self.wait()
+
+        self.play(FadeIn(panels[3], LEFT))
+        self.wait()
+
+        self.play(FadeIn(panels[0]))
+        self.wait()