"Got a basic working animation of a CNN in 3D.

manim_ml/gridded_rectangle.py

@@ -0,0 +1,197 @@
+from manim import *
+import numpy as np
+
+class CornersRectangle(Rectangle):
+    """Rectangle with functionality for getting the corner coordinates"""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.corners = VGroup(
+            *[Dot(corner_point) for corner_point in self.get_vertices()]
+        )
+        self.corners.set_fill(opacity=0.0)
+        self.add(self.corners)
+
+    def get_corners_dict(self):
+        """Returns a dictionary of the corners"""
+        return {
+            "top_left": self.corners[3],
+            "top_right": self.corners[0],
+            "bottom_left": self.corners[2],
+            "bottom_right": self.corners[1], 
+        }
+
+class GriddedRectangle(VGroup):
+    """Rectangle object with grid lines"""
+
+    def __init__(self, center, color=WHITE, height=2.0, width=4.0, 
+                mark_paths_closed=True, close_new_points=True,
+                grid_xstep=None, grid_ystep=None, grid_stroke_width=0.0, #DEFAULT_STROKE_WIDTH/2, 
+                grid_stroke_color=None, grid_stroke_opacity=None, **kwargs):
+        super().__init__()
+        # Fields
+        self.center = center
+        self.mark_paths_closed = mark_paths_closed
+        self.close_new_points = close_new_points
+        self.grid_xstep = grid_xstep
+        self.grid_ystep = grid_ystep
+        self.grid_stroke_width = grid_stroke_width
+        self.grid_stroke_color = grid_stroke_color
+        self.grid_stroke_opacity = grid_stroke_opacity
+        self.rotation_angles = [0, 0, 0]
+        # Make inner_rectangle
+        self.inner_rectangle = Rectangle(
+            width=width, 
+            height=height, 
+            stroke_opacity=0.0,
+            stroke_width=0.0
+        )
+        print(self.inner_rectangle.get_vertices())
+        # self.inner_rectangle = Polygon(
+        """
+        points = [
+            self.center + np.array([width / 2, height / 2, 0]),
+            self.center + np.array([width / 2, -1*(height / 2), 0]),
+            self.center + np.array([-1 * (width / 2), -1 * (height / 2), 0]),
+            self.center + np.array([-1 * (width / 2), height / 2, 0]),
+        ]
+        self.inner_rectangle = Polygram(
+            points,
+            stroke_opacity=0.0,
+            stroke_width=0.0
+        )
+        """
+        self.add(self.inner_rectangle)
+        # Make outline rectangle
+        self.outline_rectangle = SurroundingRectangle(
+            self.inner_rectangle, 
+            color=color,
+            buff=0.0,
+            **kwargs
+        )
+        self.add(self.outline_rectangle)
+        # Move to center
+        self.move_to(self.center)
+        # Setup Object
+        # TODO re-implement gridded rectangle
+        # self.grid_lines = self.make_grid_lines()
+        # self.add(self.grid_lines)
+        # Make dots for the corners
+        # Make outer corner dots
+        self.outer_corners = VGroup(
+            *[Dot(corner_point) for corner_point in self.outline_rectangle.get_vertices()]
+        )
+
+        self.outer_corners.set_fill(opacity=0.0)
+        self.add(self.outer_corners)
+        # Make inner corner dots
+        self.inner_corners = VGroup(
+            *[Dot(corner_point) for corner_point in self.inner_rectangle.get_vertices()]
+        )
+        self.inner_corners.set_fill(opacity=0.0)
+        self.add(self.inner_corners)
+    
+    def get_corners_dict(self, inner_rectangle=False):
+        """Returns a dictionary of the corners"""
+        if inner_rectangle:
+            return {
+                "top_left": self.inner_corners[3],
+                "top_right": self.inner_corners[0],
+                "bottom_left": self.inner_corners[2],
+                "bottom_right": self.inner_corners[1], 
+            }
+        else:
+            return {
+                "top_left": self.outer_corners[3],
+                "top_right": self.outer_corners[0],
+                "bottom_left": self.outer_corners[2],
+                "bottom_right": self.outer_corners[1], 
+            }
+
+    def make_grid_lines(self):
+        """Make grid lines in rectangle"""
+        grid_lines = VGroup()
+        width = self.width
+        height = self.width
+
+        v = self.inner_rectangle.get_vertices()
+        if self.grid_xstep is not None:
+            grid_xstep = abs(self.grid_xstep)
+            count = int(width / grid_xstep)
+            grid = VGroup(
+                *(
+                    Line(
+                        v[1] + i * grid_xstep * RIGHT,
+                        v[1] + i * grid_xstep * RIGHT + height * DOWN,
+                        color=self.color,
+                        stroke_width=self.grid_stroke_width
+                    )
+                    for i in range(1, count)
+                )
+            )
+            grid_lines.add(grid)
+    
+        if self.grid_ystep is not None:
+            grid_ystep = abs(self.grid_ystep)
+            count = int(height / grid_ystep)
+            grid = VGroup(
+                *(
+                    Line(
+                        v[1] + i * grid_ystep * DOWN,
+                        v[1] + i * grid_ystep * DOWN + width * RIGHT,
+                        color=self.color,
+                        stroke_width = self.grid_stroke_width
+                    )
+                    for i in range(1, count)
+                )
+            ) 
+            grid_lines.add(grid)
+
+        return grid_lines
+
+    def rotate_about_origin(self, angle, axis=OUT, axes=[]):
+        self.rotation_angles[np.nonzero(axis)[0][0]] = angle
+        return super().rotate_about_origin(angle, axis, axes)
+
+    def get_normal_vector(self):
+        """Gets the vector normal to main rectangle face"""
+        # Get three corner points
+        corner_1 = self.rectangle.get_top()
+        corner_2 = self.rectangle.get_left()
+        corner_3 = self.rectangle.get_right()
+        # Make vectors from them
+        a = corner_1 - corner_3
+        b = corner_1 - corner_2
+        # Compute cross product
+        normal_vector = np.cross(b, a)
+        normal_vector /= np.linalg.norm(normal_vector)
+
+        return normal_vector
+
+    def get_rotation_axis_and_angle(self):
+        """Gets the angle of rotation necessary to rotate something from the default z-axis to the rectangle"""
+        def unit_vector(vector):
+            """ Returns the unit vector of the vector.  """
+            return vector / np.linalg.norm(vector)
+
+        def angle_between(v1, v2):
+            """ Returns the angle in radians between vectors 'v1' and 'v2'::
+
+                    >>> angle_between((1, 0, 0), (0, 1, 0))
+                    1.5707963267948966
+                    >>> angle_between((1, 0, 0), (1, 0, 0))
+                    0.0
+                    >>> angle_between((1, 0, 0), (-1, 0, 0))
+                    3.141592653589793
+            """
+            v1_u = unit_vector(v1)
+            v2_u = unit_vector(v2)
+            return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))
+
+        normal_vector = self.get_normal_vector()
+        z_axis = Z_AXIS
+        # Get angle between normal vector and z axis
+        axis = np.cross(normal_vector, z_axis)
+        angle = angle_between(normal_vector, z_axis)
+
+        return axis, angle

manim_ml/neural_network/layers/__init__.py

@@ -1,7 +1,7 @@
 from .convolutional3d_to_convolutional3d import Convolutional3DToConvolutional3D
 from .convolutional2d_to_convolutional2d import Convolutional2DToConvolutional2D
-from .convolutional_3d import Convolutional3DLayer
+from .convolutional3d import Convolutional3DLayer
-from .convolutional_2d import Convolutional2DLayer
+from .convolutional2d import Convolutional2DLayer
 from .feed_forward_to_vector import FeedForwardToVector
 from .paired_query_to_feed_forward import PairedQueryToFeedForward
 from .embedding_to_feed_forward import EmbeddingToFeedForward

manim_ml/neural_network/layers/convolutional2d.py

@@ -1,19 +1,17 @@
 from manim import *
 from matplotlib import animation
-from xarray import align
 from manim_ml.neural_network.layers.parent_layers import VGroupNeuralNetworkLayer
 
 class Convolutional2DLayer(VGroupNeuralNetworkLayer):
     
     def __init__(self, feature_map_height, feature_map_width, filter_width, filter_height,
-            stride=1, cell_width=0.5, pixel_width=0.5, feature_map_color=BLUE, filter_color=ORANGE, 
+            stride=1, cell_width=0.5, feature_map_color=BLUE, filter_color=ORANGE, 
             **kwargs):
         super(VGroupNeuralNetworkLayer, self).__init__(**kwargs)
         self.feature_map_height = feature_map_height
         self.feature_map_width = feature_map_width
         self.filter_width = filter_width
         self.filter_height = filter_height
-        self.pixel_width = pixel_width
         self.feature_map_color = feature_map_color
         self.filter_color = filter_color
         self.stride = stride

manim_ml/neural_network/layers/convolutional2d_to_convolutional2d.py

@@ -1,6 +1,6 @@
 from cv2 import line
 from manim import *
-from manim_ml.neural_network.layers.convolutional_2d import Convolutional2DLayer
+from manim_ml.neural_network.layers.convolutional2d import Convolutional2DLayer
 from manim_ml.neural_network.layers.parent_layers import ConnectiveLayer
 
 class Convolutional2DToConvolutional2D(ConnectiveLayer):

manim_ml/neural_network/layers/convolutional3d.py

@@ -0,0 +1,96 @@
+from manim import *
+from manim_ml.neural_network.layers.parent_layers import ThreeDLayer, VGroupNeuralNetworkLayer
+from manim_ml.gridded_rectangle import GriddedRectangle
+import numpy as np
+
+class Convolutional3DLayer(VGroupNeuralNetworkLayer, ThreeDLayer):
+    """Handles rendering a convolutional layer for a nn"""
+
+    def __init__(self, num_feature_maps, feature_map_width, feature_map_height, 
+                filter_width, filter_height, cell_width=0.2, filter_spacing=0.1, color=BLUE, 
+                pulse_color=ORANGE, filter_color=ORANGE, stride=1, stroke_width=2.0, **kwargs):
+        super(VGroupNeuralNetworkLayer, self).__init__(**kwargs)
+        self.num_feature_maps = num_feature_maps
+        self.feature_map_height = feature_map_height
+        self.filter_color = filter_color
+        self.feature_map_width = feature_map_width
+        self.filter_width = filter_width
+        self.filter_height = filter_height
+        self.cell_width = cell_width
+        self.filter_spacing = filter_spacing
+        self.color = color
+        self.pulse_color = pulse_color
+        self.stride = stride
+        self.stroke_width = stroke_width
+        # Make the feature maps
+        self.feature_maps = self.construct_feature_maps()
+        self.add(self.feature_maps)
+
+    def construct_feature_maps(self):
+        """Creates the neural network layer"""
+        # Draw rectangles that are filled in with opacity
+        feature_maps = VGroup()
+        for filter_index in range(self.num_feature_maps):
+            rectangle = GriddedRectangle(
+                center=[0, 0, filter_index * self.filter_spacing], # Center coordinate
+                color=self.color, 
+                height=self.feature_map_height * self.cell_width,
+                width=self.feature_map_width * self.cell_width,
+                fill_color=self.color,
+                fill_opacity=0.2, 
+                stroke_color=self.color,
+                stroke_width=self.stroke_width,
+                # grid_xstep=self.cell_width,
+                # grid_ystep=self.cell_width,
+                # grid_stroke_width=DEFAULT_STROKE_WIDTH/2
+            )
+            # Rotate about z axis
+            rectangle.rotate_about_origin(
+                90 * DEGREES, 
+                np.array([0, 1, 0])
+            ) 
+            feature_maps.add(rectangle)
+        
+        return feature_maps
+
+    def make_forward_pass_animation(
+            self, 
+            run_time=5, 
+            corner_pulses=False,
+            layer_args={}, 
+            **kwargs
+        ):
+        """Convolution forward pass animation"""
+        # Note: most of this animation is done in the Convolution3DToConvolution3D layer
+        print(f"Corner pulses: {corner_pulses}")
+        if corner_pulses:
+            passing_flashes = []
+            for line in self.corner_lines:
+                pulse = ShowPassingFlash(
+                    line.copy()
+                        .set_color(self.pulse_color)
+                        .set_stroke(opacity=1.0), 
+                    time_width=0.5,
+                    run_time=run_time,
+                    rate_func=rate_functions.linear
+                )
+                passing_flashes.append(pulse)
+
+            # per_filter_run_time = run_time / len(self.feature_maps)
+            # Make animation group
+            animation_group = AnimationGroup(
+                *passing_flashes,
+            #    filter_flashes
+            )
+        else:
+            animation_group = AnimationGroup()
+
+        return animation_group
+
+    def scale(self, scale_factor, **kwargs):
+        self.cell_width *= scale_factor
+        super().scale(scale_factor, **kwargs)
+
+    @override_animation(Create)
+    def _create_override(self, **kwargs):
+        return FadeIn(self.feature_maps)

manim_ml/neural_network/layers/convolutional3d_to_convolutional3d.py

@@ -1,66 +1,287 @@
-from cProfile import run
 from manim import *
-from manim_ml.neural_network.layers.convolutional_3d import Convolutional3DLayer
+from manim_ml.neural_network.layers.convolutional3d import Convolutional3DLayer
-from manim_ml.neural_network.layers.parent_layers import ConnectiveLayer
+from manim_ml.neural_network.layers.parent_layers import ConnectiveLayer, ThreeDLayer
+from manim_ml.gridded_rectangle import GriddedRectangle, CornersRectangle
 
-class Convolutional3DToConvolutional3D(ConnectiveLayer):
+class Filters(VGroup):
+    """Group for showing a collection of filters connecting two layers"""
+
+    def __init__(self, input_layer, output_layer, line_color=ORANGE, stroke_width=2.0):
+        super().__init__()
+        self.input_layer = input_layer
+        self.output_layer = output_layer
+        self.line_color = line_color
+        self.stroke_width = stroke_width
+        # Make the filter
+        self.input_rectangles = self.make_input_feature_map_rectangles()
+        self.add(self.input_rectangles)
+        self.output_rectangles = self.make_output_feature_map_rectangles()
+        self.add(self.output_rectangles)
+        self.connective_lines = self.make_connective_lines()
+        self.add(self.connective_lines)
+
+    def make_input_feature_map_rectangles(self):
+        rectangles = []
+
+        rectangle_width = self.input_layer.filter_width * self.input_layer.cell_width
+        rectangle_height = self.input_layer.filter_height * self.input_layer.cell_width
+        filter_color = self.input_layer.filter_color
+
+        for index, feature_map in enumerate(self.input_layer.feature_maps):
+            rectangle = GriddedRectangle(
+                center=feature_map.get_center(),
+                width=rectangle_width, 
+                height=rectangle_height,
+                fill_color=filter_color,
+                stroke_color=filter_color,
+                fill_opacity=0.2,
+                z_index=2,
+                stroke_width=self.stroke_width,
+            )
+            # Center on feature map
+            # rectangle.move_to(feature_map.get_center())
+            # Rotate so it is in the yz plane
+            rectangle.rotate(
+                90 * DEGREES,
+                axis=[0, 1, 0]
+            )
+            # Get the feature map top left corner
+            feature_map_top_left = feature_map.get_corners_dict(inner_rectangle=True)["top_left"]
+            rectangle_top_left = rectangle.get_corners_dict()["top_left"]
+            # Move the rectangle to the corner location
+            rectangle.next_to(
+                feature_map_top_left,
+                submobject_to_align=rectangle_top_left,
+                buff=0.0
+            )
+
+            rectangles.append(rectangle)
+            
+        feature_map_rectangles = VGroup(*rectangles)
+
+        return feature_map_rectangles
+
+    def make_output_feature_map_rectangles(self):
+        rectangles = []
+
+        rectangle_width = self.output_layer.cell_width
+        rectangle_height = self.output_layer.cell_width
+        filter_color = self.output_layer.filter_color
+
+        for index, feature_map in enumerate(self.output_layer.feature_maps):
+            rectangle = GriddedRectangle(
+                center=feature_map.get_center(),
+                width=rectangle_width, 
+                height=rectangle_height,
+                fill_color=filter_color,
+                stroke_color=filter_color,
+                fill_opacity=0.2,
+                stroke_width=self.stroke_width,
+                z_index=2,
+            )
+            # Center on feature map
+            # rectangle.move_to(feature_map.get_center())
+            # Rotate so it is in the yz plane
+            rectangle.rotate(
+                90 * DEGREES,
+                axis=[0, 1, 0]
+            )
+            # Get the feature map top left corner
+            feature_map_top_left = feature_map.get_corners_dict(inner_rectangle=True)["top_left"]
+            rectangle_top_left = rectangle.get_corners_dict()["top_left"]
+            # Move the rectangle to the corner location
+            rectangle.next_to(
+                feature_map_top_left,
+                submobject_to_align=rectangle_top_left,
+                buff=0.0
+            )
+
+            rectangles.append(rectangle)
+            
+        feature_map_rectangles = VGroup(*rectangles)
+
+        return feature_map_rectangles
+
+    def make_connective_lines(self):
+        """Lines connecting input filter with output node"""
+
+        corner_names = ["top_left", "bottom_left", "top_right", "bottom_right"]
+
+        def make_input_connective_lines():
+            """Makes connective lines between the corners of the input filters"""
+            first_input_rectangle = self.input_rectangles[0]
+            last_input_rectangle = self.input_rectangles[-1]
+            # Get the corner dots for each rectangle
+            first_input_corners = first_input_rectangle.get_corners_dict()
+            last_input_corners = last_input_rectangle.get_corners_dict()
+            # Iterate through each corner and make the lines
+            lines = []
+            for corner_name in corner_names:
+                line = Line(
+                    first_input_corners[corner_name].get_center(),
+                    last_input_corners[corner_name].get_center(),
+                    color=self.line_color,
+                    stroke_width=self.stroke_width
+                )
+                lines.append(line)
+
+            return VGroup(*lines)
+
+        def make_output_connective_lines():
+            """Makes connective lines between the corners of the output filters"""
+            first_output_rectangle = self.output_rectangles[0]
+            last_output_rectangle = self.output_rectangles[-1]
+            # Get the corner dots for each rectangle
+            first_output_corners = first_output_rectangle.get_corners_dict()
+            last_output_corners = last_output_rectangle.get_corners_dict()
+            # Iterate through each corner and make the lines
+            lines = []
+            for corner_name in corner_names:
+                line = Line(
+                    first_output_corners[corner_name].get_center(),
+                    last_output_corners[corner_name].get_center(),
+                    color=self.line_color,
+                    stroke_width=self.stroke_width
+                )
+                lines.append(line)
+
+            return VGroup(*lines)
+
+        def make_input_to_output_connective_lines():
+            """Make connective lines between last input filter and first output filter"""
+            last_input_rectangle = self.input_rectangles[-1]
+            first_output_rectangle = self.output_rectangles[0]
+            # Get the corner dots for each rectangle
+            last_input_corners = last_input_rectangle.get_corners_dict()
+            first_output_corners = first_output_rectangle.get_corners_dict()
+            # Iterate through each corner and make the lines
+            lines = []
+            for corner_name in corner_names:
+                line = Line(
+                    last_input_corners[corner_name].get_center(),
+                    first_output_corners[corner_name].get_center(),
+                    color=self.line_color,
+                    stroke_width=self.stroke_width
+                )
+                lines.append(line)
+
+            return VGroup(*lines)   
+            
+        input_lines = make_input_connective_lines()
+        output_lines = make_output_connective_lines()
+        input_output_lines = make_input_to_output_connective_lines()
+
+        connective_lines = VGroup(
+            *input_lines, 
+            *output_lines, 
+            *input_output_lines
+        )
+ 
+        return connective_lines
+
+class Convolutional3DToConvolutional3D(ConnectiveLayer, ThreeDLayer):
     """Feed Forward to Embedding Layer"""
     input_class = Convolutional3DLayer
     output_class = Convolutional3DLayer
 
-    def __init__(self, input_layer, output_layer, color=WHITE, pulse_color=RED,
+    def __init__(self, input_layer: Convolutional3DLayer, output_layer: Convolutional3DLayer, 
-                **kwargs):
+                color=WHITE, filter_opacity=0.3, line_color=WHITE, 
-        super().__init__(input_layer, output_layer, input_class=Convolutional3DLayer, output_class=Convolutional3DLayer,
+                pulse_color=ORANGE, **kwargs):
-                        **kwargs)
+        super().__init__(input_layer, output_layer, input_class=Convolutional3DLayer, 
+                output_class=Convolutional3DLayer, **kwargs)
         self.color = color
+        self.filter_color = self.input_layer.filter_color
+        self.filter_width = self.input_layer.filter_width
+        self.filter_height = self.input_layer.filter_height
+        self.feature_map_width = self.input_layer.feature_map_width
+        self.feature_map_height = self.input_layer.feature_map_height
+        self.num_input_feature_maps = self.input_layer.num_feature_maps
+        self.num_output_feature_maps = self.output_layer.num_feature_maps
+        self.cell_width = self.input_layer.cell_width
+        self.stride = self.input_layer.stride
+        self.filter_opacity = filter_opacity
+        self.line_color = line_color
         self.pulse_color = pulse_color
 
-        self.lines = self.make_lines()
+    def make_filter_propagation_animation(self):
-        self.add(self.lines)
+        """Make filter propagation animation"""
-
+        # TODO implement this
-    def make_lines(self):
+        raise NotImplementedError()
-        """Make lines connecting the input and output layers"""
+        # Deprecated code
-        lines = VGroup()
+        old_z_index = self.filter_lines.z_index
-        # Get the first and last rectangle
+        lines_copy = self.filter_lines.copy().set_color(ORANGE).set_z_index(old_z_index + 1)
-        input_rectangle = self.input_layer.rectangles[-1]
-        output_rectangle = self.output_layer.rectangles[0]
-        input_vertices = input_rectangle.get_vertices()
-        output_vertices = output_rectangle.get_vertices()
-        # Go through each vertex
-        for vertex_index in range(len(input_vertices)):
-            # Make a line 
-            line = Line(
-                start=input_vertices[vertex_index],
-                end=output_vertices[vertex_index],
-                color=self.color,
-                stroke_opacity=0.0
-            )
-            lines.add(line)
-
-        return lines
-
-    def make_forward_pass_animation(self, layer_args={}, run_time=1.5, **kwargs):
-        """Forward pass animation from conv to conv"""
-        animations = []
-        # Go thorugh the lines
-        for line in self.lines:
-            pulse = ShowPassingFlash(
-                line.copy()
-                    .set_color(self.pulse_color)
-                    .set_stroke(opacity=1.0), 
-                    time_width=0.5,
-                    run_time=run_time
-                )
-            animations.append(pulse)
-        # Make animation group
         animation_group = AnimationGroup(
-            *animations,
+            Create(lines_copy, lag_ratio=0.0),
-            run_time=run_time
+            # FadeOut(self.filter_lines),
+            FadeOut(lines_copy),
+            lag_ratio=1.0
         )
 
         return animation_group
 
+    def make_forward_pass_animation(self, layer_args={}, run_time=10.5, **kwargs):
+        """Forward pass animation from conv2d to conv2d"""
+        animations = []
+        # Create the filters, output nodes (feature map square), and lines
+        filters = Filters(self.input_layer, self.output_layer)
+        self.add(filters)
+        # Make animations for creating the filters, output_nodes, and filter_lines
+        # TODO decide if I want to create the filters at the start of a conv 
+        # animation or have them there by default
+        # animations.append(
+        #     Create(filters)
+        # )
+        # Make shift amounts
+        right_shift = np.array([0, self.input_layer.cell_width, 0])# * 1.55
+        left_shift = np.array([0, -1*self.input_layer.cell_width, 0])# * 1.55
+        up_shift = np.array([0, 0, -1*self.input_layer.cell_width])# * 1.55
+        down_shift = np.array([0, 0, self.input_layer.cell_width])# * 1.55
+        # Make filter shifting animations
+        num_y_moves = int((self.feature_map_height - self.filter_height) / self.stride)
+        num_x_moves = int((self.feature_map_width - self.filter_width) / self.stride)
+        for y_move in range(num_y_moves):
+            # Go right num_x_moves
+            for x_move in range(num_x_moves):
+                # Shift right
+                shift_animation = ApplyMethod(
+                    filters.shift,
+                    self.stride * right_shift
+                )
+                # shift_animation = self.animate.shift(right_shift)
+                animations.append(shift_animation)
+            
+            # Go back left num_x_moves and down one
+            shift_amount = self.stride * num_x_moves * left_shift + self.stride * down_shift
+            # Make the animation
+            shift_animation = ApplyMethod(
+                filters.shift,
+                shift_amount  
+            )
+            animations.append(shift_animation)
+        # Do last row move right
+        for x_move in range(num_x_moves):
+            # Shift right
+            shift_animation = ApplyMethod(
+                filters.shift,
+                self.stride * right_shift
+            )
+            # shift_animation = self.animate.shift(right_shift)
+            animations.append(shift_animation)
+            
+        # Remove filters
+        return Succession(
+            *animations,
+            lag_ratio=1.0
+        )
+
+    def set_z_index(self, z_index, family=False):
+        """Override set_z_index"""
+        super().set_z_index(4)
+
+    def scale(self, scale_factor, **kwargs):
+        self.cell_width *= scale_factor
+        super().scale(scale_factor, **kwargs)
+
     @override_animation(Create)
     def _create_override(self, **kwargs):
         return AnimationGroup()
-

manim_ml/neural_network/layers/convolutional_3d.py

@@ -1,122 +0,0 @@
-from manim import *
-from manim_ml.neural_network.layers.parent_layers import VGroupNeuralNetworkLayer
-
-class Convolutional3DLayer(VGroupNeuralNetworkLayer):
-    """Handles rendering a convolutional layer for a nn"""
-
-    def __init__(self, num_filters, filter_width, filter_height, filter_spacing=0.1, color=BLUE, 
-                pulse_color=ORANGE, **kwargs):
-        super(VGroupNeuralNetworkLayer, self).__init__(**kwargs)
-        self.num_filters = num_filters
-        self.filter_width = filter_width
-        self.filter_height = filter_height
-        self.filter_spacing = filter_spacing
-        self.color = color
-        self.pulse_color = pulse_color
-
-        self._construct_layer(num_filters=self.num_filters, filter_width=self.filter_width, filter_height=self.filter_height)
-
-    def _construct_layer(self, num_filters=5, filter_width=4, filter_height=4):
-        """Creates the neural network layer"""
-        # Make axes, but hide the lines
-        axes = ThreeDAxes(
-            tips=False,
-            x_length=1,
-            y_length=1,
-            x_axis_config={
-                "include_ticks": False,
-                "stroke_width": 0.0
-            },
-            y_axis_config={
-                "include_ticks": False,
-                "stroke_width": 0.0
-            },
-            z_axis_config={
-                "include_ticks": False,
-                "stroke_width": 0.0 
-            }
-        )
-        self.add(axes)
-        # Set the camera angle so that the 
-        # self.set_camera_orientation(phi=75 * DEGREES, theta=30 * DEGREES)
-        # Draw rectangles that are filled in with opacity
-        self.rectangles = VGroup()
-        for filter_index in range(num_filters):
-            rectangle = Rectangle(
-                color=self.color, 
-                height=filter_height,
-                width=filter_width,
-                fill_color=self.color,
-                fill_opacity=0.2, 
-                stroke_color=WHITE,
-            )
-            rectangle.rotate_about_origin((80 - filter_index*0.5) * DEGREES, np.array([0, 1, 0])) # Rotate about z axis
-            rectangle.rotate_about_origin(15 * DEGREES, np.array([1, 0, 0])) # Rotate about x axis
-            rectangle.shift(np.array([filter_index*self.filter_spacing, filter_height*0.5, -3]))
-
-            self.rectangles.add(rectangle)
-        
-        self.add(self.rectangles)
-
-        self.corner_lines = self.make_filter_corner_lines()
-        self.add(self.corner_lines)
-
-    def make_filter_corner_lines(self):
-        """Make filter corner lines"""
-        corner_lines = VGroup()
-
-        first_rectangle = self.rectangles[0]
-        last_rectangle = self.rectangles[-1]
-        first_vertices = first_rectangle.get_vertices()
-        last_vertices = last_rectangle.get_vertices()
-        for vertex_index in range(len(first_vertices)):
-            # Make a line 
-            line = Line(
-                start=first_vertices[vertex_index],
-                end=last_vertices[vertex_index],
-                color=WHITE,
-                stroke_opacity=0.0
-            )
-            corner_lines.add(line)
-
-        return corner_lines
-
-    def make_forward_pass_animation(self, run_time=5, layer_args={}, **kwargs):
-        """Convolution forward pass animation"""
-        passing_flashes = []
-        for line in self.corner_lines:
-            pulse = ShowPassingFlash(
-                line.copy()
-                    .set_color(self.pulse_color)
-                    .set_stroke(opacity=1.0), 
-                time_width=0.5,
-                run_time=run_time,
-                rate_func=rate_functions.linear
-            )
-            passing_flashes.append(pulse)
-
-        per_filter_run_time = run_time / len(self.rectangles)
-        filter_flashes = []
-        for filter in self.rectangles:
-            single_flash = Succession(
-                ApplyMethod(filter.set_color, self.pulse_color, run_time=per_filter_run_time/4),
-                Wait(per_filter_run_time/2),
-                ApplyMethod(filter.set_color, self.color, run_time=per_filter_run_time/4),
-                ApplyMethod(filter.set_stroke_color, WHITE, run_time=0.0)
-            )
-            filter_flashes.append(single_flash)
-
-        filter_flashes = Succession(
-            *filter_flashes,
-        )
-        # Make animation group
-        animation_group = AnimationGroup(
-            *passing_flashes,
-            filter_flashes
-        )
-
-        return animation_group
-
-    @override_animation(Create)
-    def _create_override(self, **kwargs):
-        return FadeIn(self.rectangles)

manim_ml/neural_network/layers/parent_layers.py

@@ -35,6 +35,12 @@ class VGroupNeuralNetworkLayer(NeuralNetworkLayer):
     def _create_override(self):
         return super()._create_override()
 
+class ThreeDLayer(ABC):
+    """Abstract class for 3D layers"""
+
+    def __init__(self):
+        pass
+
 class ConnectiveLayer(VGroupNeuralNetworkLayer):
     """Forward pass animation for a given pair of layers"""
 

manim_ml/neural_network/neural_network.py

@@ -14,16 +14,18 @@ import textwrap
 from manim_ml.neural_network.layers.embedding import EmbeddingLayer
 
 from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer
-from manim_ml.neural_network.layers.parent_layers import ConnectiveLayer
+from manim_ml.neural_network.layers.parent_layers import ConnectiveLayer, ThreeDLayer
 from manim_ml.neural_network.layers.util import get_connective_layer
 from manim_ml.list_group import ListGroup
 from manim_ml.neural_network.neural_network_transformations import InsertLayer, RemoveLayer
 
 class NeuralNetwork(Group):
+    """Neural Network Visualization Container Class"""
 
     def __init__(self, input_layers, edge_color=WHITE, layer_spacing=0.2,
                     animation_dot_color=RED, edge_width=2.5, dot_radius=0.03,
-                    title=" "):
+                    title=" ", camera=None, camera_phi=-70 * DEGREES, 
+                    camera_theta=-80 * DEGREES):
         super(Group, self).__init__()
         self.input_layers = ListGroup(*input_layers)
         self.edge_width = edge_width
@@ -33,6 +35,11 @@ class NeuralNetwork(Group):
         self.dot_radius = dot_radius
         self.title_text = title
         self.created = False
+        self.camera = camera
+        # Set the camera orientation for 3D Layers
+        if not self.camera is None:
+            self.camera.set_phi(camera_phi)
+            self.camera.set_theta(camera_theta)
         # TODO take layer_node_count [0, (1, 2), 0] 
         # and make it have explicit distinct subspaces
         self._place_layers()
@@ -59,7 +66,8 @@ class NeuralNetwork(Group):
             current_layer = self.input_layers[layer_index]
             current_layer.move_to(previous_layer)
             # TODO Temp fix
-            if isinstance(current_layer, EmbeddingLayer) or isinstance(previous_layer, EmbeddingLayer):
+            if isinstance(current_layer, EmbeddingLayer) \
+                or isinstance(previous_layer, EmbeddingLayer):
                 shift_vector = np.array([(previous_layer.get_width()/2 + current_layer.get_width()/2 - 0.2), 0, 0]) 
             else:
                 shift_vector = np.array([(previous_layer.get_width()/2 + current_layer.get_width()/2) + self.layer_spacing, 0, 0])
@@ -71,6 +79,11 @@ class NeuralNetwork(Group):
         all_layers = ListGroup()
         for layer_index in range(len(self.input_layers) - 1):
             current_layer = self.input_layers[layer_index]
+            # Make the layer fixed in frame if its not 3D
+            if not isinstance(current_layer, ThreeDLayer):
+                self.camera.add_fixed_orientation_mobjects(current_layer)
+                self.camera.add_fixed_in_frame_mobjects(current_layer)
+            # Add the layer to the list of layers
             all_layers.add(current_layer)
             next_layer = self.input_layers[layer_index + 1]
             # Check if layer is actually a nested NeuralNetwork
@@ -83,7 +96,16 @@ class NeuralNetwork(Group):
             # Find connective layer with correct layer pair
             connective_layer = get_connective_layer(current_layer, next_layer)
             connective_layers.add(connective_layer)
+            # Make the layer fixed in frame if its not 3D
+            if not isinstance(current_layer, ThreeDLayer):
+                self.camera.add_fixed_orientation_mobjects(connective_layer)
+                self.camera.add_fixed_in_frame_mobjects(connective_layer)
+            # Add the layer to the list of layers
             all_layers.add(connective_layer)
+        # Check if final layer is a 3D layer
+        if not isinstance(self.input_layers[-1], ThreeDLayer):
+            self.camera.add_fixed_orientation_mobjects(self.input_layers[-1])
+            self.camera.add_fixed_in_frame_mobjects(self.input_layers[-1])
         # Add final layer
         all_layers.add(self.input_layers[-1])
         # Handle layering
@@ -114,11 +136,11 @@ class NeuralNetwork(Group):
 
         return animation_group
 
-    def make_forward_pass_animation(self, run_time=10, passing_flash=True, layer_args={},
+    def make_forward_pass_animation(self, run_time=None, passing_flash=True, layer_args={},
                                     **kwargs):
         """Generates an animation for feed forward propagation"""
         all_animations = []
-        per_layer_runtime = run_time/len(self.all_layers)
+        per_layer_runtime = run_time / len(self.all_layers) if not run_time is None else None
         for layer_index, layer in enumerate(self.all_layers):
             # Get the layer args
             if isinstance(layer, ConnectiveLayer):
@@ -148,7 +170,6 @@ class NeuralNetwork(Group):
         # Make the animation group
         animation_group = Succession(
             *all_animations, 
-            run_time=run_time, 
             lag_ratio=1.0
         )
 

tests/test_3d_camera_move.py

@@ -0,0 +1,41 @@
+from manim import *
+from PIL import Image
+
+from manim_ml.neural_network.layers.convolutional3d import Convolutional3DLayer
+from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer
+from manim_ml.neural_network.layers.image import ImageLayer
+from manim_ml.neural_network.neural_network import NeuralNetwork
+
+# Make the specific scene
+config.pixel_height = 1200
+config.pixel_width = 1900
+config.frame_height = 6.0
+config.frame_width = 6.0
+
+class CombinedScene(ThreeDScene):
+    def construct(self):
+        image = Image.open('../assets/mnist/digit.jpeg')
+        numpy_image = np.asarray(image)
+        # Make nn
+        nn = NeuralNetwork(
+            [
+                ImageLayer(numpy_image, height=1.4),
+                Convolutional3DLayer(1, 5, 5, 3, 3, filter_spacing=0.2),
+                Convolutional3DLayer(2, 3, 3, 1, 1, filter_spacing=0.2),
+                FeedForwardLayer(3, rectangle_stroke_width=4, node_stroke_width=4),
+            ], 
+            layer_spacing=0.5,
+            camera=self.camera
+        )
+
+        nn.scale(1.3)
+        # Center the nn
+        nn.move_to(ORIGIN)
+        self.add(nn)
+        # Play animation
+        forward_pass = nn.make_forward_pass_animation(
+            corner_pulses=False
+        )
+        self.play(
+            forward_pass
+        ) 

tests/test_convolutional_3d_layer.py

@@ -1,7 +1,7 @@
 from manim import *
 from PIL import Image
 
-from manim_ml.neural_network.layers.convolutional_3d import Convolutional3DLayer
+from manim_ml.neural_network.layers.convolutional3d import Convolutional3DLayer
 from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer
 from manim_ml.neural_network.layers.image import ImageLayer
 from manim_ml.neural_network.neural_network import NeuralNetwork
@@ -22,34 +22,64 @@ class SingleConvolutionalLayerScene(ThreeDScene):
         self.set_camera_orientation(phi=280*DEGREES, theta=-10*DEGREES, gamma=90*DEGREES)
         # self.play(nn.make_forward_pass_animation(run_time=5))
 
+class Simple3DConvScene(ThreeDScene):
+
+    def construct(self):
+        """
+            TODO
+            - [X] Make grid lines for the CNN filters
+            - [ ] Make Scanning filter effect
+            - [ ] Have filter box go accross each input feature map
+            - [ ] Make filter lines effect
+            - [ ] Make flowing animation down filter lines
+        """
+        # Make nn
+        layers = [
+            Convolutional3DLayer(1, 5, 5, 5, 5, feature_map_height=3, filter_width=3, filter_height=3),
+            Convolutional3DLayer(1, 3, 3, 1, 1, feature_map_width=3, filter_width=3, filter_height=3),
+        ]
+        nn = NeuralNetwork(layers)
+        # Center the nn
+        nn.move_to(ORIGIN)
+        self.add(nn)
+        # Play animation
+        # self.set_camera_orientation(phi=280*DEGREES, theta=-10*DEGREES, gamma=90*DEGREES)
+        self.play(nn.make_forward_pass_animation(run_time=30))
+
 # Make the specific scene
 config.pixel_height = 1200
 config.pixel_width = 1900
-config.frame_height = 12.0
+config.frame_height = 6.0
-config.frame_width = 12.0
+config.frame_width = 6.0
 
-class CombinedScene(ThreeDScene, Scene):
+class CombinedScene(ThreeDScene):
     def construct(self):
         image = Image.open('../assets/mnist/digit.jpeg')
         numpy_image = np.asarray(image)
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 ImageLayer(numpy_image, height=1.4),
-            Convolutional3DLayer(3, 3, 3, filter_spacing=0.2),
+                Convolutional3DLayer(1, 5, 5, 3, 3, filter_spacing=0.2),
-            Convolutional3DLayer(5, 2, 2, filter_spacing=0.2),
+                Convolutional3DLayer(2, 3, 3, 1, 1, filter_spacing=0.2),
-            Convolutional3DLayer(10, 2, 1, filter_spacing=0.2),
+                FeedForwardLayer(3, rectangle_stroke_width=4, node_stroke_width=4),
-            FeedForwardLayer(3, rectangle_stroke_width=4, node_stroke_width=4).scale(2),
+                FeedForwardLayer(1, rectangle_stroke_width=4, node_stroke_width=4)
-            FeedForwardLayer(1, rectangle_stroke_width=4, node_stroke_width=4).scale(2)
+            ], 
-        ], layer_spacing=0.2)
+            layer_spacing=0.5,
+            camera=self.camera
+        )
 
         nn.scale(1.3)
         # Center the nn
         nn.move_to(ORIGIN)
-        self.play(Create(nn))
+        self.add(nn)
         # Play animation
         # self.set_camera_orientation(phi=280* DEGREES, theta=-20*DEGREES, gamma=90 * DEGREES)
         # self.begin_ambient_camera_rotation()
-        forward_pass = nn.make_forward_pass_animation(run_time=10)
+        forward_pass = nn.make_forward_pass_animation(
+            run_time=10,
+            corner_pulses=False
+        )
         print(forward_pass)
         self.play(
             forward_pass

tests/test_flow.py

@@ -0,0 +1,6 @@
+from manim_ml.flow.flow import *
+
+class TestScene(Scene):
+
+    def construct(self):
+        self.add(Rectangle())

tests/test_opengl_shader.py

@@ -0,0 +1,87 @@
+import manim.utils.opengl as opengl
+from manim import *
+from manim.opengl import *  # type: ignore
+
+class InlineShaderExample(Scene):
+    def construct(self):
+        config["background_color"] = "#333333"
+
+        c = Circle(fill_opacity=0.7).shift(UL)
+        self.add(c)
+
+        shader = Shader(
+            self.renderer.context,
+            source={
+                "vertex_shader": """
+                #version 330
+                in vec4 in_vert;
+                in vec4 in_color;
+                out vec4 v_color;
+                uniform mat4 u_model_view_matrix;
+                uniform mat4 u_projection_matrix;
+                void main() {
+                    v_color = in_color;
+                    vec4 camera_space_vertex = u_model_view_matrix * in_vert;
+                    vec4 clip_space_vertex = u_projection_matrix * camera_space_vertex;
+                    gl_Position = clip_space_vertex;
+                }
+            """,
+                "fragment_shader": """
+            #version 330
+            in vec4 v_color;
+            out vec4 frag_color;
+            void main() {
+              frag_color = v_color;
+            }
+
+            void main () {
+                // Previously, you'd have rendered your complete scene into a texture
+                // bound to "fullScreenTexture."
+                vec4 rValue = texture2D(fullscreenTexture, gl_TexCoords[0] - rOffset);  
+                vec4 gValue = texture2D(fullscreenTexture, gl_TexCoords[0] - gOffset);
+                vec4 bValue = texture2D(fullscreenTexture, gl_TexCoords[0] - bOffset);  
+
+                // Combine the offset colors.
+                gl_FragColor = vec4(rValue.r, gValue.g, bValue.b, 1.0);
+            }
+            """,
+            },
+        )
+        shader.set_uniform("u_model_view_matrix", opengl.view_matrix())
+        shader.set_uniform(
+            "u_projection_matrix",
+            opengl.orthographic_projection_matrix(),
+        )
+
+        attributes = np.zeros(
+            6,
+            dtype=[
+                ("in_vert", np.float32, (4,)),
+                ("in_color", np.float32, (4,)),
+            ],
+        )
+        attributes["in_vert"] = np.array(
+            [
+                [-1, -1, 0, 1],
+                [-1, 1, 0, 1],
+                [1, 1, 0, 1],
+                [-1, -1, 0, 1],
+                [1, -1, 0, 1],
+                [1, 1, 0, 1],
+            ],
+        )
+        attributes["in_color"] = np.array(
+            [
+                [0, 0, 1, 1],
+                [0, 0, 1, 1],
+                [0, 0, 1, 1],
+                [0, 0, 1, 1],
+                [0, 0, 1, 1],
+                [0, 0, 1, 1],
+            ],
+        )
+        mesh = Mesh(shader, attributes)
+        self.add(mesh)
+
+        self.wait(5)
+        # self.embed_2()

tests/test_surrounding_rectangle.py

@@ -0,0 +1,10 @@
+from manim import *
+
+class SurroundingRectangleTest(Scene):
+
+    def construct(self):
+        rectangle = Rectangle(width=1, height=1, color=WHITE, fill_opacity=1.0)
+        self.add(rectangle)
+
+        surrounding_rectangle = SurroundingRectangle(rectangle, color=GREEN, buff=0.0)
+        self.add(surrounding_rectangle)