Embedding Neural Network Layer.

manim_ml/neural_network/feed_forward.py

@@ -0,0 +1,128 @@
+from manim import *
+from manim_ml.neural_network.layers import VGroupNeuralNetworkLayer, ConnectiveLayer
+
+class FeedForwardLayer(VGroupNeuralNetworkLayer):
+    """Handles rendering a layer for a neural network"""
+
+    def __init__(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, animation_dot_color=RED):
+        super(VGroupNeuralNetworkLayer, self).__init__()
+        self.num_nodes = num_nodes
+        self.layer_buffer = layer_buffer
+        self.node_radius = node_radius
+        self.node_color = node_color
+        self.node_stroke_width = node_stroke_width
+        self.node_outline_color = node_outline_color
+        self.rectangle_stroke_width = rectangle_stroke_width
+        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()
+
+        self._construct_neural_network_layer()
+
+    def _construct_neural_network_layer(self):
+        """Creates the neural network layer"""
+        # Add Nodes
+        for node_number in range(self.num_nodes):
+            node_object = Circle(radius=self.node_radius, color=self.node_color, 
+                                stroke_width=self.node_stroke_width)
+            self.node_group.add(node_object)
+        # Space the nodes
+        # Assumes Vertical orientation
+        for node_index, node_object in enumerate(self.node_group):
+            location = node_index * self.node_spacing
+            node_object.move_to([0, location, 0])
+        # Create Surrounding Rectangle
+        self.surrounding_rectangle = SurroundingRectangle(self.node_group, color=self.rectangle_color, 
+                                                        fill_color=self.rectangle_fill_color, fill_opacity=1.0, 
+                                                        buff=self.layer_buffer, stroke_width=self.rectangle_stroke_width)
+        # Add the objects to the class
+        self.add(self.surrounding_rectangle, self.node_group)
+
+    def make_forward_pass_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
+
+    @override_animation(Create)
+    def _create_animation(self, **kwargs):
+        animations = []
+
+        animations.append(Create(self.surrounding_rectangle))
+
+        for node in self.node_group:
+            animations.append(Create(node))
+
+        animation_group = AnimationGroup(*animations, lag_ratio=0.0)
+        return animation_group
+
+class FeedForwardToFeedForward(ConnectiveLayer):
+    """Layer for connecting FeedForward layer to FeedForwardLayer"""
+
+    def __init__(self, input_layer, output_layer, passing_flash=True,
+                dot_radius=0.05, animation_dot_color=RED, edge_color=WHITE,
+                edge_width=0.5):
+        super().__init__(input_layer, output_layer)
+        self.passing_flash = passing_flash
+        self.edge_color = edge_color
+        self.dot_radius = dot_radius
+        self.animation_dot_color = animation_dot_color
+        self.edge_width = edge_width
+
+        self.edges = self.construct_edges()
+        self.add(self.edges)
+
+    def construct_edges(self):
+        # Go through each node in the two layers and make a connecting line
+        edges = []
+        for node_i in self.input_layer.node_group:
+            for node_j in self.output_layer.node_group:
+                line = Line(node_i.get_center(), node_j.get_center(), 
+                            color=self.edge_color, stroke_width=self.edge_width)
+                edges.append(line)
+
+        edges = VGroup(*edges)
+        return edges
+
+    def make_forward_pass_animation(self, run_time=1):
+        """Animation for passing information from one FeedForwardLayer to the next"""
+        path_animations = []
+        dots = []
+        for edge in self.edges:
+            dot = Dot(color=self.animation_dot_color, fill_opacity=1.0, radius=self.dot_radius)   
+            # Add to dots group
+            dots.append(dot)
+            # Make the animation
+            if self.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=run_time, rate_function=sigmoid)
+            path_animations.append(anim)
+
+        if not self.passing_flash:
+            dots = VGroup(*dots)
+            self.add(dots)
+
+        path_animations = AnimationGroup(*path_animations)
+
+        return path_animations
+
+    @override_animation(Create)
+    def _create_animation(self, **kwargs):
+        animations = []
+
+        for edge in self.edges:
+            animations.append(Create(edge))
+
+        animation_group = AnimationGroup(*animations, lag_ratio=0.0)
+        return animation_group