From 068b0746f01762a63ec478add9c41a3a316f17ec Mon Sep 17 00:00:00 2001 From: Alec Helbling Date: Mon, 28 Mar 2022 21:24:49 -0400 Subject: [PATCH] Made new neural network forward pass animation --- manim_ml/logo.py | 41 -------------------------- manim_ml/neural_network.py | 59 +++++++++++++++++++++++++++++++------- 2 files changed, 49 insertions(+), 51 deletions(-) delete mode 100644 manim_ml/logo.py diff --git a/manim_ml/logo.py b/manim_ml/logo.py deleted file mode 100644 index 98475bb..0000000 --- a/manim_ml/logo.py +++ /dev/null @@ -1,41 +0,0 @@ -""" - Logo for Manim Machine Learning -""" -from manim import * -from neural_network import NeuralNetwork - -config.pixel_height = 500 -config.pixel_width = 1920 -config.frame_height = 10.0 -config.frame_width = 10.0 - -class ManimMLLogo(Scene): - - def construct(self): - self.text = Text("ManimML") - self.text.scale(1.3) - self.neural_network = NeuralNetwork([3, 5, 3, 6, 3], layer_spacing=0.6, node_color=BLUE) - self.neural_network.scale(0.8) - self.neural_network.move_to(self.text.get_right()) - self.neural_network.shift(RIGHT * 1.3) - self.logo_group = VGroup(self.text, self.neural_network) - self.logo_group.scale(1.5) - self.logo_group.move_to(ORIGIN) - self.play(Write(self.text)) - self.play(Create(self.neural_network)) - # self.surrounding_rectangle = SurroundingRectangle(self.logo_group, buff=0.3, color=BLUE) - underline = Underline(self.text, color=BLUE) - animation_group = AnimationGroup( - self.neural_network.make_forward_propagation_animation(run_time=5), - Create(underline), - # Create(self.surrounding_rectangle) - ) - # self.surrounding_rectangle = SurroundingRectangle(self.logo_group, buff=0.3, color=BLUE) - underline = Underline(self.text, color=BLUE) - animation_group = AnimationGroup( - self.neural_network.make_forward_propagation_animation(run_time=5), - Create(underline), - # Create(self.surrounding_rectangle) - ) - self.play(animation_group) - self.wait(5) diff --git a/manim_ml/neural_network.py b/manim_ml/neural_network.py index df393e8..4ed11fa 100644 --- a/manim_ml/neural_network.py +++ b/manim_ml/neural_network.py @@ -10,6 +10,18 @@ Example: NeuralNetwork(layer_node_count) """ from manim import * +from matplotlib import animation + +class ChangeColor(Animation): + CONFIG={ + "rate_func":linear + } + def interpolate_submobject(self, submobject, starting_sumobject, alpha): + m = int(alpha * 10) % 2 + if m == 0: + submobject.set_color(RED) + else: + submobject.set_color(YELLOW) class NeuralNetworkLayer(VGroup): """Handles rendering a layer for a neural network""" @@ -18,7 +30,7 @@ class NeuralNetworkLayer(VGroup): self, num_nodes, layer_buffer=SMALL_BUFF/2, node_radius=0.08, node_color=BLUE, node_outline_color=WHITE, rectangle_color=WHITE, node_spacing=0.3, rectangle_fill_color=BLACK, node_stroke_width=2.0, - rectangle_stroke_width=2.0): + rectangle_stroke_width=2.0, animation_dot_color=RED): super(VGroup, self).__init__() self.num_nodes = num_nodes self.layer_buffer = layer_buffer @@ -30,6 +42,7 @@ class NeuralNetworkLayer(VGroup): self.rectangle_color = rectangle_color self.node_spacing = node_spacing self.rectangle_fill_color = rectangle_fill_color + self.animation_dot_color = animation_dot_color self.node_group = VGroup() @@ -54,6 +67,16 @@ class NeuralNetworkLayer(VGroup): # Add the objects to the class self.add(surrounding_rectangle, self.node_group) + def _make_highlight_nodes_animation(self): + # make highlight animation + succession = Succession( + ApplyMethod(self.node_group.set_color, self.animation_dot_color, run_time=0.25), + Wait(1.0), + ApplyMethod(self.node_group.set_color, self.node_color, run_time=0.25), + ) + + return succession + class NeuralNetwork(VGroup): def __init__( @@ -84,7 +107,11 @@ class NeuralNetwork(VGroup): layers = VGroup() # Create each layer for layer_index, node_count in enumerate(self.layer_node_count): - layer = NeuralNetworkLayer(node_count, node_color=self.node_color) + layer = NeuralNetworkLayer( + node_count, + node_color=self.node_color, + animation_dot_color=self.animation_dot_color + ) # Manage spacing layer.move_to([self.layer_spacing * layer_index, 0, 0]) # Add layer to VGroup @@ -110,12 +137,12 @@ class NeuralNetwork(VGroup): edge_layers.set_z_index(0) return edge_layers - def make_forward_propagation_animation(self, run_time=2): + def make_forward_propagation_animation(self, run_time=2, passing_flash = True): """Generates an animation for feed forward propogation""" all_animations = [] per_layer_run_time = run_time / len(self.edge_layers) self.dots = VGroup() - for edge_layer in self.edge_layers: + for i, edge_layer in enumerate(self.edge_layers): path_animations = [] for edge in edge_layer: dot = Dot(color=self.animation_dot_color, fill_opacity=1.0, radius=self.dot_radius) @@ -124,12 +151,24 @@ class NeuralNetwork(VGroup): # Add to dots group self.dots.add(dot) # Make the animation - anim = MoveAlongPath(dot, edge, run_time=per_layer_run_time, rate_function=sigmoid) + if passing_flash: + anim = ShowPassingFlash(edge.copy().set_color(self.animation_dot_color), time_width=0.2, run_time=3) + else: + anim = MoveAlongPath(dot, edge, run_time=per_layer_run_time, rate_function=sigmoid) path_animations.append(anim) + # Highlight each node + layer = self.layers[i] + highlight_animation = layer._make_highlight_nodes_animation() + all_animations.append(highlight_animation) + path_animation_group = AnimationGroup(*path_animations) all_animations.append(path_animation_group) - animation_group = AnimationGroup(*all_animations, run_time=run_time, lag_ratio=1) + layer = self.layers[-1] + highlight_animation = layer._make_highlight_nodes_animation() + all_animations.append(highlight_animation) + + animation_group = AnimationGroup(*all_animations, run_time=run_time, lag_ratio=1.0) return animation_group @@ -148,10 +187,10 @@ class TestNeuralNetworkScene(Scene): nn.move_to(ORIGIN) # Make Animation self.add(nn) - forward_propagation_animation = nn.make_forward_propagation_animation() + forward_propagation_animation = nn.make_forward_propagation_animation(5) - second_nn = NeuralNetwork([3, 4]) - self.add(second_nn) + # second_nn = NeuralNetwork([3, 4]) + # self.add(second_nn) self.play(forward_propagation_animation) - self.play(second_nn.make_forward_propagation_animation()) + # self.play(second_nn.make_forward_propagation_animation())