add code boxes to all elementwise scenes in chapter 6.

Author: michelp

Chapter6/Scene1.py

@@ -7,7 +7,7 @@
 load_dotenv()
 
 import numpy as np
-from scene_utils import setup_scene, create_sparse_matrix, create_small_graph_from_matrix
+from scene_utils import setup_scene, create_sparse_matrix, create_square_digraph
 
 
 # Shared data from notebook
@@ -17,8 +17,9 @@
 # Graph B: edges (1→2), (2→3), (3→0)
 B_DATA = [[0, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0]]
 
-# Union result: edges (0→1), (1→2), (2→3), (3→0)
-UNION_DATA = [[0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0]]
+# Union result with binary.plus: edges sum where both exist
+# (0→1)=1, (1→2)=2, (2→3)=2, (3→0)=1
+UNION_DATA = [[0, 1, 0, 0], [0, 0, 2, 0], [0, 0, 0, 2], [1, 0, 0, 0]]
 
 # Node positions (square layout)
 POS = {0: np.array([0, 1, 0]), 1: np.array([1, 1, 0]), 2: np.array([1, 0, 0]), 3: np.array([0, 0, 0])}
@@ -41,7 +42,7 @@ def construct(self):
             self.play(Write(title))
 
             # Create Graph A (left, blue)
-            graph_a = self.create_graph(A_DATA, BLUE)
+            graph_a = create_square_digraph(A_DATA, BLUE)
             graph_a.scale(1.2).shift(LEFT * 4 + DOWN * 0.5)
             label_a = Text("Graph A", font_size=24, color=BLUE).next_to(graph_a, UP)
 
@@ -50,7 +51,7 @@ def construct(self):
             mat_a.next_to(graph_a, DOWN, buff=0.5)
 
             # Create Graph B (right, green)
-            graph_b = self.create_graph(B_DATA, GREEN)
+            graph_b = create_square_digraph(B_DATA, GREEN)
             graph_b.scale(1.2).shift(RIGHT * 4 + DOWN * 0.5)
             label_b = Text("Graph B", font_size=24, color=GREEN).next_to(graph_b, UP)
 
@@ -93,31 +94,35 @@ def construct(self):
             self.wait(1)
 
         with self.voiceover(
-            """When we compute eWiseAdd, we get all edges from both: 0 to 1 from A,
-            1 to 2 and 2 to 3 from both, and 3 to 0 from B. The result has four
-            distinct edge positions representing the union."""
+            """When we compute eWiseAdd with binary plus, we get all edges from both
+            graphs. Edge 0 to 1 comes from A, edge 3 to 0 comes from B. The shared
+            edges 1 to 2 and 2 to 3 sum to 2 in the result."""
         ):
-            # Show formula
-            formula = MathTex(r"A \cup B", font_size=48, color=YELLOW)
-            formula.move_to(ORIGIN + UP * 0.5)
-            self.play(Write(formula))
+            # Show Python code below title
+            code = Code(
+                code_string="A.ewise_add(B, binary.plus)",
+                language="python",
+                background="window",
+            ).scale(0.7)
+            code.next_to(title, DOWN, buff=0.3)
+            self.play(FadeIn(code))
             self.wait(0.5)
 
             # Fade out highlights and prepare for result
             self.play(FadeOut(common_highlight_a), FadeOut(common_highlight_b))
 
-            # Create union result graph (center, yellow)
-            graph_union = self.create_graph(UNION_DATA, YELLOW)
+            # Create union result graph (center, yellow) with edge weights
+            graph_union = create_square_digraph(UNION_DATA, YELLOW, show_weights=True)
             graph_union.scale(1.2).move_to(ORIGIN + DOWN * 0.5)
             label_union = Text("A ∪ B", font_size=24, color=YELLOW).next_to(graph_union, UP)
 
             # Create union matrix
             mat_union = create_sparse_matrix(UNION_DATA, scale=0.5)
             mat_union.next_to(graph_union, DOWN, buff=0.5)
 
-            # Animate combining: fade formula, bring in result
+            # Animate bringing in result
             self.play(
-                ReplacementTransform(formula, label_union),
+                Write(label_union),
                 Create(graph_union),
                 FadeIn(mat_union),
             )
@@ -131,48 +136,14 @@ def construct(self):
 
         # Cleanup
         self.play(
-            FadeOut(title), FadeOut(graph_a), FadeOut(label_a), FadeOut(mat_a),
+            FadeOut(title), FadeOut(code),
+            FadeOut(graph_a), FadeOut(label_a), FadeOut(mat_a),
             FadeOut(graph_b), FadeOut(label_b), FadeOut(mat_b),
             FadeOut(graph_union), FadeOut(label_union), FadeOut(mat_union),
             FadeOut(edge_count),
         )
         self.wait(0.5)
 
-    def create_graph(self, matrix_data, color):
-        """Create a directed graph from adjacency matrix with square layout."""
-        n = len(matrix_data)
-        positions = {
-            0: np.array([-0.7, 0.7, 0]),
-            1: np.array([0.7, 0.7, 0]),
-            2: np.array([0.7, -0.7, 0]),
-            3: np.array([-0.7, -0.7, 0]),
-        }
-
-        # Create vertices
-        vertices = {}
-        for i in range(n):
-            label = MathTex(str(i), color=BLACK).scale(0.5)
-            dot = LabeledDot(label, radius=0.2, fill_color=WHITE, fill_opacity=1)
-            dot.move_to(positions[i])
-            vertices[i] = dot
-
-        # Create edges
-        edges = VGroup()
-        for i in range(n):
-            for j in range(n):
-                if matrix_data[i][j] != 0:
-                    arrow = Arrow(
-                        positions[i], positions[j],
-                        color=color, buff=0.25, stroke_width=3,
-                        tip_length=0.15, max_tip_length_to_length_ratio=0.25
-                    )
-                    edges.add(arrow)
-
-        graph = VGroup(edges, *vertices.values())
-        graph.vertices = vertices
-        graph.edges = edges
-        return graph
-
     def get_nearest_node(self, pos):
         """Get the nearest node index for a position."""
         positions = {

Chapter6/Scene2.py

@@ -7,7 +7,7 @@
 load_dotenv()
 
 import numpy as np
-from scene_utils import setup_scene, create_sparse_matrix
+from scene_utils import setup_scene, create_sparse_matrix, create_square_digraph
 
 
 # Shared data from notebook
@@ -37,7 +37,7 @@ def construct(self):
             self.play(Write(title))
 
             # Create Graph A (left, blue)
-            graph_a = self.create_graph(A_DATA, BLUE)
+            graph_a = create_square_digraph(A_DATA, BLUE)
             graph_a.scale(1.2).shift(LEFT * 4 + DOWN * 0.5)
             label_a = Text("Graph A", font_size=24, color=BLUE).next_to(graph_a, UP)
 
@@ -46,7 +46,7 @@ def construct(self):
             mat_a.next_to(graph_a, DOWN, buff=0.5)
 
             # Create Graph B (right, green)
-            graph_b = self.create_graph(B_DATA, GREEN)
+            graph_b = create_square_digraph(B_DATA, GREEN)
             graph_b.scale(1.2).shift(RIGHT * 4 + DOWN * 0.5)
             label_b = Text("Graph B", font_size=24, color=GREEN).next_to(graph_b, UP)
 
@@ -61,14 +61,18 @@ def construct(self):
             self.wait(1)
 
         with self.voiceover(
-            """The result contains just two edges: 1 to 2 and 2 to 3. Contrast this
-            with union which gave us four edges. Intersection keeps only what
-            both graphs share."""
+            """The result contains just two edges: 1 to 2 and 2 to 3, each with
+            value 1 times 1 equals 1. Contrast this with union which gave us four
+            edges. Intersection keeps only what both graphs share."""
         ):
-            # Show formula
-            formula = MathTex(r"A \cap B", font_size=48, color=RED_C)
-            formula.move_to(ORIGIN + UP * 0.5)
-            self.play(Write(formula))
+            # Show Python code below title
+            code = Code(
+                code_string="A.ewise_mult(B, binary.times)",
+                language="python",
+                background="window",
+            ).scale(0.7)
+            code.next_to(title, DOWN, buff=0.3)
+            self.play(FadeIn(code))
             self.wait(0.5)
 
             # Highlight common edges and fade unique ones
@@ -110,8 +114,8 @@ def construct(self):
             )
             self.wait(0.5)
 
-            # Create intersection result graph (center, red)
-            graph_int = self.create_graph(INTERSECTION_DATA, RED_C)
+            # Create intersection result graph (center, red) with edge weights
+            graph_int = create_square_digraph(INTERSECTION_DATA, RED_C, show_weights=True)
             graph_int.scale(1.2).move_to(ORIGIN + DOWN * 0.5)
             label_int = Text("A ∩ B", font_size=24, color=RED_C).next_to(graph_int, UP)
 
@@ -120,7 +124,7 @@ def construct(self):
             mat_int.next_to(graph_int, DOWN, buff=0.5)
 
             self.play(
-                ReplacementTransform(formula, label_int),
+                Write(label_int),
                 Create(graph_int),
                 FadeIn(mat_int),
             )
@@ -138,48 +142,14 @@ def construct(self):
 
         # Cleanup
         self.play(
-            FadeOut(title), FadeOut(graph_a), FadeOut(label_a), FadeOut(mat_a),
+            FadeOut(title), FadeOut(code),
+            FadeOut(graph_a), FadeOut(label_a), FadeOut(mat_a),
             FadeOut(graph_b), FadeOut(label_b), FadeOut(mat_b),
             FadeOut(graph_int), FadeOut(label_int), FadeOut(mat_int),
             FadeOut(comparison),
         )
         self.wait(0.5)
 
-    def create_graph(self, matrix_data, color):
-        """Create a directed graph from adjacency matrix with square layout."""
-        n = len(matrix_data)
-        positions = {
-            0: np.array([-0.7, 0.7, 0]),
-            1: np.array([0.7, 0.7, 0]),
-            2: np.array([0.7, -0.7, 0]),
-            3: np.array([-0.7, -0.7, 0]),
-        }
-
-        # Create vertices
-        vertices = {}
-        for i in range(n):
-            label = MathTex(str(i), color=BLACK).scale(0.5)
-            dot = LabeledDot(label, radius=0.2, fill_color=WHITE, fill_opacity=1)
-            dot.move_to(positions[i])
-            vertices[i] = dot
-
-        # Create edges
-        edges = VGroup()
-        for i in range(n):
-            for j in range(n):
-                if matrix_data[i][j] != 0:
-                    arrow = Arrow(
-                        positions[i], positions[j],
-                        color=color, buff=0.25, stroke_width=3,
-                        tip_length=0.15, max_tip_length_to_length_ratio=0.25
-                    )
-                    edges.add(arrow)
-
-        graph = VGroup(edges, *vertices.values())
-        graph.vertices = vertices
-        graph.edges = edges
-        return graph
-
     def get_nearest_node(self, pos):
         """Get the nearest node index for a position."""
         positions = {

Chapter6/Scene3.py

@@ -34,21 +34,25 @@ def construct(self):
 
             # Create weighted graph
             graph = self.create_weighted_graph(W_DATA)
-            graph.scale(1.5).shift(LEFT * 3 + DOWN * 0.3)
+            graph.scale(1.5).shift(LEFT * 3 + UP * 0.3)
             label = Text("Weighted Graph W", font_size=24).next_to(graph, UP)
 
             # Create matrix
             mat = create_sparse_matrix(W_DATA, scale=0.55)
-            mat.shift(RIGHT * 3 + UP * 0.5)
+            mat.shift(RIGHT * 3 + UP * 1.0)
             mat_label = Text("W", font_size=28).next_to(mat, UP)
 
             self.play(Create(graph), Write(label), FadeIn(mat), Write(mat_label))
             self.wait(0.5)
 
-            # Show threshold condition
-            threshold = MathTex(r"\text{weight} > 3", font_size=36, color=YELLOW)
-            threshold.next_to(mat, DOWN, buff=0.5)
-            self.play(Write(threshold))
+            # Show Python code for select
+            code = Code(
+                code_string='W.select(">", 3)',
+                language="python",
+                background="window",
+            ).scale(0.7)
+            code.next_to(mat, DOWN, buff=0.4)
+            self.play(FadeIn(code))
             self.wait(1)
 
         with self.voiceover(
@@ -110,7 +114,7 @@ def construct(self):
 
             # Show result matrix
             result_mat = create_sparse_matrix(FILTERED_DATA, scale=0.55)
-            result_mat.next_to(threshold, DOWN, buff=0.5)
+            result_mat.next_to(code, DOWN, buff=0.4)
             result_label = Text("Result: 2 edges", font_size=20, color=GREEN)
             result_label.next_to(result_mat, DOWN, buff=0.3)
 
@@ -120,7 +124,7 @@ def construct(self):
         # Cleanup
         self.play(
             FadeOut(title), FadeOut(graph), FadeOut(label),
-            FadeOut(mat), FadeOut(mat_label), FadeOut(threshold),
+            FadeOut(mat), FadeOut(mat_label), FadeOut(code),
             FadeOut(result_mat), FadeOut(result_label),
         )
         self.wait(0.5)

Chapter6/Scene4.py

@@ -34,16 +34,20 @@ def construct(self):
 
             # Create original matrix
             mat_m = create_sparse_matrix(M_DATA, scale=0.6)
-            mat_m.shift(LEFT * 3)
+            mat_m.shift(LEFT * 4.5)
             mat_label = MathTex("M", font_size=36).next_to(mat_m, UP, buff=0.3)
 
             self.play(FadeIn(mat_m), Write(mat_label))
             self.wait(0.5)
 
-            # Show formula
-            formula = MathTex(r"\sqrt{M}", font_size=48, color=GREEN)
-            formula.move_to(ORIGIN)
-            self.play(Write(formula))
+            # Show Python code
+            code = Code(
+                code_string="M.apply(unary.sqrt)",
+                language="python",
+                background="window",
+            ).scale(0.7)
+            code.move_to(ORIGIN + UP * 0.5)
+            self.play(FadeIn(code))
             self.wait(0.5)
 
         with self.voiceover(
@@ -53,8 +57,8 @@ def construct(self):
         ):
             # Create result matrix
             mat_result = create_sparse_matrix(SQRT_DATA, scale=0.6)
-            mat_result.shift(RIGHT * 3)
-            result_label = MathTex(r"\sqrt{M}", font_size=36, color=GREEN)
+            mat_result.shift(RIGHT * 4.5)
+            result_label = Text("Result", font_size=24, color=GREEN)
             result_label.next_to(mat_result, UP, buff=0.3)
 
             # Animate sqrt transformation for each non-zero entry
@@ -107,7 +111,7 @@ def construct(self):
         # Cleanup
         self.play(
             FadeOut(title), FadeOut(mat_m), FadeOut(mat_label),
-            FadeOut(formula), FadeOut(mat_result), FadeOut(result_label),
+            FadeOut(code), FadeOut(mat_result), FadeOut(result_label),
             FadeOut(note),
         )
         self.wait(0.5)