From ba63116b37853a92c03aaf6e21c7e68f9f42596a Mon Sep 17 00:00:00 2001 From: Alec Helbling Date: Mon, 9 Jan 2023 15:52:37 +0900 Subject: [PATCH] Reformatted the code using black, allowd for different orientation NNs, made an option for highlighting the active filter in a CNN forward pass. --- .../translation_equivariance.py | 45 +++++ .../decision_tree/classification_areas.py | 34 +++- manim_ml/decision_tree/decision_tree.py | 191 ++++++++++-------- manim_ml/decision_tree/helpers.py | 23 ++- manim_ml/gridded_rectangle.py | 9 +- manim_ml/image.py | 28 ++- manim_ml/neural_network/animations/dropout.py | 111 ++++------ .../neural_network/layers/convolutional3d.py | 7 - .../convolutional3d_to_convolutional3d.py | 58 +++++- .../neural_network/layers/feed_forward.py | 14 +- .../layers/feed_forward_to_feed_forward.py | 15 +- manim_ml/neural_network/layers/image.py | 4 + .../layers/image_to_convolutional3d.py | 9 +- .../neural_network/layers/parent_layers.py | 11 +- manim_ml/neural_network/neural_network.py | 108 +++++++--- .../neural_network_transformations.py | 1 + tests/test_decision_tree.py | 38 ++-- tests/test_nn_dropout.py | 18 +- tests/test_nn_scale.py | 44 ++++ 19 files changed, 485 insertions(+), 283 deletions(-) create mode 100644 examples/translation_equivariance/translation_equivariance.py create mode 100644 tests/test_nn_scale.py diff --git a/examples/translation_equivariance/translation_equivariance.py b/examples/translation_equivariance/translation_equivariance.py new file mode 100644 index 0000000..8684995 --- /dev/null +++ b/examples/translation_equivariance/translation_equivariance.py @@ -0,0 +1,45 @@ +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.layers.parent_layers import ThreeDLayer +from manim_ml.neural_network.neural_network import NeuralNetwork + +# Make the specific scene +config.pixel_height = 1200 +config.pixel_width = 800 +config.frame_height = 6.0 +config.frame_width = 6.0 + + +class CombinedScene(ThreeDScene): + def construct(self): + image = Image.open("../../assets/doggo.jpeg") + numpy_image = np.asarray(image) + # Rotate the Three D layer position + ThreeDLayer.rotation_angle = 15 * DEGREES + ThreeDLayer.rotation_axis = [1, -1.0, 0] + # Make nn + nn = NeuralNetwork( + [ + ImageLayer(numpy_image, height=1.5), + Convolutional3DLayer(1, 7, 7, 3, 3, filter_spacing=0.32), + Convolutional3DLayer(3, 5, 5, 1, 1, filter_spacing=0.18), + ], + layer_spacing=0.25, + layout_direction="top_to_bottom", + ) + # Center the nn + nn.move_to(ORIGIN) + nn.scale(1.5) + self.add(nn) + # Play animation + forward_pass = nn.make_forward_pass_animation( + corner_pulses=False, + all_filters_at_once=False, + highlight_active_feature_map=True, + ) + self.wait(1) + self.play(forward_pass) diff --git a/manim_ml/decision_tree/classification_areas.py b/manim_ml/decision_tree/classification_areas.py index b03b7d1..5f51371 100644 --- a/manim_ml/decision_tree/classification_areas.py +++ b/manim_ml/decision_tree/classification_areas.py @@ -3,8 +3,10 @@ import numpy as np from collections import deque from sklearn.tree import _tree as ctree + class AABB: """Axis-aligned bounding box""" + def __init__(self, n_features): self.limits = np.array([[-np.inf, np.inf]] * n_features) @@ -18,6 +20,7 @@ class AABB: return left, right + def tree_bounds(tree, n_features=None): """Compute final decision rule for each node in tree""" if n_features is None: @@ -33,8 +36,9 @@ def tree_bounds(tree, n_features=None): queue.extend([l, r]) return aabbs + def compute_decision_areas(tree_classifier, maxrange, x=0, y=1, n_features=None): - """ Extract decision areas. + """Extract decision areas. tree_classifier: Instance of a sklearn.tree.DecisionTreeClassifier maxrange: values to insert for [left, right, top, bottom] if the interval is open (+/-inf) @@ -69,21 +73,27 @@ def compute_decision_areas(tree_classifier, maxrange, x=0, y=1, n_features=None) rectangles[:, [1, 3]] = np.minimum(rectangles[:, [1, 3]], maxrange[1::2]) return rectangles + def plot_areas(rectangles): for rect in rectangles: - color = ['b', 'r'][int(rect[4])] + color = ["b", "r"][int(rect[4])] print(rect[0], rect[1], rect[2] - rect[0], rect[3] - rect[1]) - rp = Rectangle([rect[0], rect[2]], - rect[1] - rect[0], - rect[3] - rect[2], color=color, alpha=0.3) + rp = Rectangle( + [rect[0], rect[2]], + rect[1] - rect[0], + rect[3] - rect[2], + color=color, + alpha=0.3, + ) plt.gca().add_artist(rp) + def merge_overlapping_polygons(all_polygons, colors=[BLUE, GREEN, ORANGE]): # get all polygons of each color polygon_dict = { - str(BLUE).lower():[], - str(GREEN).lower():[], - str(ORANGE).lower():[] + str(BLUE).lower(): [], + str(GREEN).lower(): [], + str(ORANGE).lower(): [], } for polygon in all_polygons: print(polygon_dict) @@ -98,7 +108,7 @@ def merge_overlapping_polygons(all_polygons, colors=[BLUE, GREEN, ORANGE]): vertices = polygon.get_vertices().tolist() vertices = [tuple(vert) for vert in vertices] for pt in vertices: - if pt in points: # Shared vertice, remove it. + if pt in points: # Shared vertice, remove it. points.remove(pt) else: points.add(pt) @@ -143,8 +153,10 @@ def merge_overlapping_polygons(all_polygons, colors=[BLUE, GREEN, ORANGE]): # Remove implementation-markers from the polygon. poly = [point for point, _ in polygon] for vertex in poly: - if vertex in edges_h: edges_h.pop(vertex) - if vertex in edges_v: edges_v.pop(vertex) + if vertex in edges_h: + edges_h.pop(vertex) + if vertex in edges_v: + edges_v.pop(vertex) polygon = Polygon(*poly, color=color, fill_opacity=0.3, stroke_opacity=1.0) return_polygons.append(polygon) return return_polygons diff --git a/manim_ml/decision_tree/decision_tree.py b/manim_ml/decision_tree/decision_tree.py index a8702cd..85f8dd1 100644 --- a/manim_ml/decision_tree/decision_tree.py +++ b/manim_ml/decision_tree/decision_tree.py @@ -6,18 +6,23 @@ TODO reimplement the decision 2D decision tree surface drawing. """ from manim import * -from manim_ml.decision_tree.classification_areas import compute_decision_areas, merge_overlapping_polygons +from manim_ml.decision_tree.classification_areas import ( + compute_decision_areas, + merge_overlapping_polygons, +) import manim_ml.decision_tree.helpers as helpers from manim_ml.one_to_one_sync import OneToOneSync import numpy as np from PIL import Image + class LeafNode(Group): """Leaf node in tree""" - def __init__(self, class_index, display_type="image", class_image_paths=[], - class_colors=[]): + def __init__( + self, class_index, display_type="image", class_image_paths=[], class_colors=[] + ): super().__init__() self.display_type = display_type self.class_image_paths = class_image_paths @@ -36,16 +41,17 @@ class LeafNode(Group): node = ImageMobject(pil_image) node.scale(1.5) rectangle = Rectangle( - width=node.width + 0.05, + width=node.width + 0.05, height=node.height + 0.05, - color=self.class_colors[class_index], - stroke_width=6 + color=self.class_colors[class_index], + stroke_width=6, ) rectangle.move_to(node.get_center()) rectangle.shift([-0.02, 0.02, 0]) self.add(rectangle) self.add(node) + class SplitNode(VGroup): """Node for splitting decision in tree""" @@ -53,25 +59,24 @@ class SplitNode(VGroup): super().__init__() node_text = f"{feature}\n<= {threshold:.2f} cm" # Draw decision text - decision_text = Text( - node_text, - color=WHITE - ) + decision_text = Text(node_text, color=WHITE) # Draw the surrounding box - bounding_box = SurroundingRectangle( - decision_text, - buff=0.3, - color=WHITE - ) + bounding_box = SurroundingRectangle(decision_text, buff=0.3, color=WHITE) self.add(bounding_box) self.add(decision_text) + class DecisionTreeDiagram(Group): """Decision Tree Diagram Class for Manim""" - def __init__(self, sklearn_tree, feature_names=None, - class_names=None, class_images_paths=None, - class_colors=[RED, GREEN, BLUE]): + def __init__( + self, + sklearn_tree, + feature_names=None, + class_names=None, + class_images_paths=None, + class_colors=[RED, GREEN, BLUE], + ): super().__init__() self.tree = sklearn_tree self.feature_names = feature_names @@ -87,14 +92,13 @@ class DecisionTreeDiagram(Group): node_index, ): """Make node""" - is_split_node = self.tree.children_left[node_index] != self.tree.children_right[node_index] + is_split_node = ( + self.tree.children_left[node_index] != self.tree.children_right[node_index] + ) if is_split_node: node_feature = self.tree.feature[node_index] node_threshold = self.tree.threshold[node_index] - node = SplitNode( - self.feature_names[node_feature], - node_threshold - ) + node = SplitNode(self.feature_names[node_feature], node_threshold) else: # Get the most abundant class for the given leaf node # Make the leaf node object @@ -102,7 +106,7 @@ class DecisionTreeDiagram(Group): node = LeafNode( class_index=tree_class_index, class_colors=self.class_colors, - class_image_paths=self.class_image_paths + class_image_paths=self.class_image_paths, ) return node @@ -113,11 +117,7 @@ class DecisionTreeDiagram(Group): bottom_node_top_location = bottom.get_center() bottom_node_top_location[1] += bottom.height / 2 - line = Line( - top_node_bottom_location, - bottom_node_top_location, - color=WHITE - ) + line = Line(top_node_bottom_location, bottom_node_top_location, color=WHITE) return line @@ -139,35 +139,53 @@ class DecisionTreeDiagram(Group): edge_map = {} # tree height tree_height = scale_factor * node_height * max_depth - tree_width = scale_factor * 2 ** max_depth * node_width + tree_width = scale_factor * 2**max_depth * node_width # traverse tree def recurse(node_index, depth, direction, parent_object, parent_node): # make the node object - is_leaf = self.tree.children_left[node_index] == self.tree.children_right[node_index] + is_leaf = ( + self.tree.children_left[node_index] + == self.tree.children_right[node_index] + ) node_object = self._make_node(node_index=node_index) nodes_map[node_index] = node_object node_height = node_object.height # set the node position direction_factor = -1 if direction == "left" else 1 - shift_right_amount = 0.9 * direction_factor * scale_factor * tree_width / (2 ** depth) / 2 + shift_right_amount = ( + 0.9 * direction_factor * scale_factor * tree_width / (2**depth) / 2 + ) if is_leaf: shift_down_amount = -1.0 * scale_factor * node_height else: shift_down_amount = -1.8 * scale_factor * node_height - node_object \ - .match_x(parent_object) \ - .match_y(parent_object) \ - .shift([shift_right_amount, shift_down_amount, 0]) + node_object.match_x(parent_object).match_y(parent_object).shift( + [shift_right_amount, shift_down_amount, 0] + ) tree_group.add(node_object) # make a connection - connection = self._make_connection(parent_object, node_object, is_leaf=is_leaf) - edge_name = str(parent_node)+","+str(node_index) + connection = self._make_connection( + parent_object, node_object, is_leaf=is_leaf + ) + edge_name = str(parent_node) + "," + str(node_index) edge_map[edge_name] = connection tree_group.add(connection) # recurse if not is_leaf: - recurse(self.tree.children_left[node_index], depth + 1, "left", node_object, node_index) - recurse(self.tree.children_right[node_index], depth + 1, "right", node_object, node_index) + recurse( + self.tree.children_left[node_index], + depth + 1, + "left", + node_object, + node_index, + ) + recurse( + self.tree.children_right[node_index], + depth + 1, + "right", + node_object, + node_index, + ) recurse(self.tree.children_left[0], 1, "left", root_node, 0) recurse(self.tree.children_right[0], 1, "right", root_node, 0) @@ -185,9 +203,7 @@ class DecisionTreeDiagram(Group): # Compute parent mapping parent_mapping = helpers.compute_node_to_parent_mapping(self.tree) # Create the root node - animations.append( - Create(self.nodes_map[0]) - ) + animations.append(Create(self.nodes_map[0])) # Iterate through the nodes queue = [0] while len(queue) > 0: @@ -211,21 +227,16 @@ class DecisionTreeDiagram(Group): FadeIn(right_node), Create(left_parent_edge), Create(right_parent_edge), - lag_ratio=0.0 - ) - animations.append( - split_animation + lag_ratio=0.0, ) + animations.append(split_animation) # Add the children to the queue if left_child != -1: queue.append(left_child) if right_child != -1: queue.append(right_child) - return AnimationGroup( - *animations, - lag_ratio=1.0 - ) + return AnimationGroup(*animations, lag_ratio=1.0) @override_animation(Create) def create_decision_tree(self, traversal_order="bfs"): @@ -237,8 +248,8 @@ class DecisionTreeDiagram(Group): else: raise Exception(f"Uncrecognized traversal: {traversal_order}") -class IrisDatasetPlot(VGroup): +class IrisDatasetPlot(VGroup): def __init__(self, iris): points = iris.data[:, 0:2] labels = iris.feature_names @@ -249,19 +260,13 @@ class IrisDatasetPlot(VGroup): self.axes_group = self._make_axes_group(points, labels) # Make legend self.legend_group = self._make_legend( - [BLUE, ORANGE, GREEN], - iris.target_names, - self.axes_group + [BLUE, ORANGE, GREEN], iris.target_names, self.axes_group ) # Make title - #title_text = "Iris Dataset Plot" - #self.title = Text(title_text).match_y(self.axes_group).shift([0.5, self.axes_group.height / 2 + 0.5, 0]) + # title_text = "Iris Dataset Plot" + # self.title = Text(title_text).match_y(self.axes_group).shift([0.5, self.axes_group.height / 2 + 0.5, 0]) # Make all group - self.all_group = Group( - self.point_group, - self.axes_group, - self.legend_group - ) + self.all_group = Group(self.point_group, self.axes_group, self.legend_group) # scale the groups self.point_group.scale(1.6) self.point_group.match_x(self.axes_group) @@ -278,8 +283,8 @@ class IrisDatasetPlot(VGroup): Wait(0.5), Create(self.axes_group, run_time=2), # add title - #Create(self.title), - Create(self.legend_group) + # Create(self.title), + Create(self.legend_group), ) return animation_group @@ -300,7 +305,9 @@ class IrisDatasetPlot(VGroup): setosa = Text("Setosa", color=BLUE) verisicolor = Text("Verisicolor", color=ORANGE) virginica = Text("Virginica", color=GREEN) - labels = VGroup(setosa, verisicolor, virginica).arrange(direction=RIGHT, aligned_edge=LEFT, buff=2.0) + labels = VGroup(setosa, verisicolor, virginica).arrange( + direction=RIGHT, aligned_edge=LEFT, buff=2.0 + ) labels.scale(0.5) legend_group.add(labels) # surrounding rectangle @@ -314,10 +321,14 @@ class IrisDatasetPlot(VGroup): return legend_group - def _make_axes_group(self, points, labels, font='Source Han Sans', font_scale=0.75): + def _make_axes_group(self, points, labels, font="Source Han Sans", font_scale=0.75): axes_group = VGroup() # make the axes - x_range = [np.amin(points, axis=0)[0] - 0.2, np.amax(points, axis=0)[0] - 0.2, 0.5] + x_range = [ + np.amin(points, axis=0)[0] - 0.2, + np.amax(points, axis=0)[0] - 0.2, + 0.5, + ] y_range = [np.amin(points, axis=0)[1] - 0.2, np.amax(points, axis=0)[1], 0.5] axes = Axes( x_range=x_range, @@ -330,23 +341,27 @@ class IrisDatasetPlot(VGroup): axes_group.add(axes) # make axis labels # x_label - x_label = Text(labels[0], font=font) \ - .match_y(axes.get_axes()[0]) \ - .shift([0.5, -0.75, 0]) \ - .scale(font_scale) + x_label = ( + Text(labels[0], font=font) + .match_y(axes.get_axes()[0]) + .shift([0.5, -0.75, 0]) + .scale(font_scale) + ) axes_group.add(x_label) # y_label - y_label = Text(labels[1], font=font) \ - .match_x(axes.get_axes()[1]) \ - .shift([-0.75, 0, 0]) \ - .rotate(np.pi / 2) \ - .scale(font_scale) + y_label = ( + Text(labels[1], font=font) + .match_x(axes.get_axes()[1]) + .shift([-0.75, 0, 0]) + .rotate(np.pi / 2) + .scale(font_scale) + ) axes_group.add(y_label) return axes_group -class DecisionTreeSurface(VGroup): +class DecisionTreeSurface(VGroup): def __init__(self, tree_clf, data, axes, class_colors=[BLUE, ORANGE, GREEN]): # take the tree and construct the surface from it self.tree_clf = tree_clf @@ -363,11 +378,7 @@ class DecisionTreeSurface(VGroup): bottom = np.amin(self.data[:, 1]) - 0.2 maxrange = [left, right, bottom, top] rectangles = compute_decision_areas( - self.tree_clf, - maxrange, - x=0, - y=1, - n_features=2 + self.tree_clf, maxrange, x=0, y=1, n_features=2 ) # turn the rectangle objects into manim rectangles def convert_rectangle_to_polygon(rect): @@ -381,9 +392,16 @@ class DecisionTreeSurface(VGroup): bottom_right_coord = self.axes.coords_to_point(*bottom_right) top_right_coord = self.axes.coords_to_point(*top_right) top_left_coord = self.axes.coords_to_point(*top_left) - points = [bottom_left_coord, bottom_right_coord, top_right_coord, top_left_coord] + points = [ + bottom_left_coord, + bottom_right_coord, + top_right_coord, + top_left_coord, + ] # construct a polygon object from those manim coordinates - rectangle = Polygon(*points, color=color, fill_opacity=0.3, stroke_opacity=0.0) + rectangle = Polygon( + *points, color=color, fill_opacity=0.3, stroke_opacity=0.0 + ) return rectangle manim_rectangles = [] @@ -392,7 +410,9 @@ class DecisionTreeSurface(VGroup): rectangle = convert_rectangle_to_polygon(rect) manim_rectangles.append(rectangle) - manim_rectangles = merge_overlapping_polygons(manim_rectangles, colors=[BLUE, GREEN, ORANGE]) + manim_rectangles = merge_overlapping_polygons( + manim_rectangles, colors=[BLUE, GREEN, ORANGE] + ) return manim_rectangles @@ -416,6 +436,7 @@ class DecisionTreeSurface(VGroup): return animation_group + class DecisionTreeContainer(OneToOneSync): """Connects the DecisionTreeDiagram to the DecisionTreeEmbedding""" diff --git a/manim_ml/decision_tree/helpers.py b/manim_ml/decision_tree/helpers.py index da2980f..11e21bb 100644 --- a/manim_ml/decision_tree/helpers.py +++ b/manim_ml/decision_tree/helpers.py @@ -1,11 +1,14 @@ - def compute_node_depths(tree): """Computes the depths of nodes for level order traversal""" + def depth(node_index, current_node_index=0): """Compute the height of a node""" if current_node_index == node_index: return 0 - elif tree.children_left[current_node_index] == tree.children_right[current_node_index]: + elif ( + tree.children_left[current_node_index] + == tree.children_right[current_node_index] + ): return -1 else: # Compute the height of each subtree @@ -23,13 +26,18 @@ def compute_node_depths(tree): return node_depths + def compute_level_order_traversal(tree): """Computes level order traversal of a sklearn tree""" + def depth(node_index, current_node_index=0): """Compute the height of a node""" if current_node_index == node_index: return 0 - elif tree.children_left[current_node_index] == tree.children_right[current_node_index]: + elif ( + tree.children_left[current_node_index] + == tree.children_right[current_node_index] + ): return -1 else: # Compute the height of each subtree @@ -47,14 +55,15 @@ def compute_level_order_traversal(tree): node_depths = sorted(node_depths, key=lambda x: x[1]) sorted_inds = [node_depth[0] for node_depth in node_depths] - return sorted_inds + return sorted_inds + def compute_node_to_parent_mapping(tree): """Returns a hashmap mapping node indices to their parent indices""" - node_to_parent = {0: -1} # Root has no parent + node_to_parent = {0: -1} # Root has no parent num_nodes = tree.node_count for node_index in range(num_nodes): - # Explore left children + # Explore left children left_child_node_index = tree.children_left[node_index] if left_child_node_index != -1: node_to_parent[left_child_node_index] = node_index @@ -62,5 +71,5 @@ def compute_node_to_parent_mapping(tree): right_child_node_index = tree.children_right[node_index] if right_child_node_index != -1: node_to_parent[right_child_node_index] = node_index - + return node_to_parent diff --git a/manim_ml/gridded_rectangle.py b/manim_ml/gridded_rectangle.py index f3e908c..86b0511 100644 --- a/manim_ml/gridded_rectangle.py +++ b/manim_ml/gridded_rectangle.py @@ -24,6 +24,7 @@ class GriddedRectangle(VGroup): ): super().__init__() # Fields + self.color = color self.mark_paths_closed = mark_paths_closed self.close_new_points = close_new_points self.grid_xstep = grid_xstep @@ -33,8 +34,6 @@ class GriddedRectangle(VGroup): self.grid_stroke_opacity = grid_stroke_opacity if show_grid_lines else 0.0 self.stroke_width = stroke_width self.rotation_angles = [0, 0, 0] - self.rectangle_width = width - self.rectangle_height = height self.show_grid_lines = show_grid_lines # Make rectangle self.rectangle = Rectangle( @@ -129,3 +128,9 @@ class GriddedRectangle(VGroup): normal_vector = np.cross((vertex_1 - vertex_2), (vertex_1 - vertex_3)) return normal_vector + + def set_color(self, color): + """Sets the color of the gridded rectangle""" + self.color = color + self.rectangle.set_color(color) + self.rectangle.set_stroke_color(color) diff --git a/manim_ml/image.py b/manim_ml/image.py index bb1ec64..e9c942f 100644 --- a/manim_ml/image.py +++ b/manim_ml/image.py @@ -3,22 +3,22 @@ import numpy as np from PIL import Image -class GrayscaleImageMobject(ImageMobject): +class GrayscaleImageMobject(Group): """Mobject for creating images in Manim from numpy arrays""" def __init__(self, numpy_image, height=2.3): + super().__init__() self.numpy_image = numpy_image - assert len(np.shape(self.numpy_image)) == 2 + input_image = self.numpy_image[None, :, :] # Convert grayscale to rgb version of grayscale input_image = np.repeat(input_image, 3, axis=0) input_image = np.rollaxis(input_image, 0, start=3) - - super().__init__(input_image, image_mode="RGB") - - self.set_resampling_algorithm(RESAMPLING_ALGORITHMS["nearest"]) - self.scale_to_fit_height(height) + self.image_mobject = ImageMobject(input_image, image_mode="RBG") + self.add(self.image_mobject) + self.image_mobject.set_resampling_algorithm(RESAMPLING_ALGORITHMS["nearest"]) + self.image_mobject.scale_to_fit_height(height) @classmethod def from_path(cls, path, height=2.3): @@ -32,6 +32,20 @@ class GrayscaleImageMobject(ImageMobject): def create(self, run_time=2): return FadeIn(self) + def scale(self, scale_factor, **kwargs): + """Scales the image mobject""" + # super().scale(scale_factor) + # height = self.height + self.image_mobject.scale(scale_factor) + # self.scale_to_fit_height(2) + # self.apply_points_function_about_point( + # lambda points: scale_factor * points, **kwargs + # ) + + def set_opacity(self, opacity): + """Set the opacity""" + self.image_mobject.set_opacity(opacity) + class LabeledColorImage(Group): """Labeled Color Image""" diff --git a/manim_ml/neural_network/animations/dropout.py b/manim_ml/neural_network/animations/dropout.py index b94174f..31ec935 100644 --- a/manim_ml/neural_network/animations/dropout.py +++ b/manim_ml/neural_network/animations/dropout.py @@ -6,10 +6,12 @@ from manim import * import random from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer -from manim_ml.neural_network.layers.feed_forward_to_feed_forward import FeedForwardToFeedForward +from manim_ml.neural_network.layers.feed_forward_to_feed_forward import ( + FeedForwardToFeedForward, +) + class XMark(VGroup): - def __init__(self, stroke_width=1.0, color=GRAY): super().__init__() line_one = Line( @@ -17,7 +19,7 @@ class XMark(VGroup): [0.1, -0.1, 0], stroke_width=1.0, stroke_color=color, - z_index=4 + z_index=4, ) self.add(line_one) line_two = Line( @@ -25,10 +27,11 @@ class XMark(VGroup): [0.1, 0.1, 0], stroke_width=1.0, stroke_color=color, - z_index=4 + z_index=4, ) self.add(line_two) + def get_edges_to_drop_out(layer: FeedForwardToFeedForward, layers_to_nodes_to_drop_out): """Returns edges to drop out for a given FeedForwardToFeedForward layer""" prev_layer = layer.input_layer @@ -43,18 +46,21 @@ def get_edges_to_drop_out(layer: FeedForwardToFeedForward, layers_to_nodes_to_dr prev_node_index = int(edge_index / next_layer.num_nodes) next_node_index = edge_index % next_layer.num_nodes # Check if the edges should be dropped out - if prev_node_index in prev_layer_nodes_to_dropout \ - or next_node_index in next_layer_nodes_to_dropout: + if ( + prev_node_index in prev_layer_nodes_to_dropout + or next_node_index in next_layer_nodes_to_dropout + ): edges_to_dropout.append(edge) edge_indices_to_dropout.append(edge_index) return edges_to_dropout, edge_indices_to_dropout + def make_pre_dropout_animation( neural_network, layers_to_nodes_to_drop_out, dropped_out_color=GRAY, - dropped_out_opacity=0.2 + dropped_out_opacity=0.2, ): """Makes an animation that sets up the NN layer for dropout""" animations = [] @@ -70,8 +76,7 @@ def make_pre_dropout_animation( layers_to_edges_to_dropout = {} for layer in feed_forward_to_feed_forward_layers: layers_to_edges_to_dropout[layer], _ = get_edges_to_drop_out( - layer, - layers_to_nodes_to_drop_out + layer, layers_to_nodes_to_drop_out ) # Dim the colors of the edges dim_edge_colors_animations = [] @@ -92,12 +97,9 @@ def make_pre_dropout_animation( ) """ - dim_edge_colors_animations.append( - FadeOut(edge) - ) + dim_edge_colors_animations.append(FadeOut(edge)) dim_edge_colors_animation = AnimationGroup( - *dim_edge_colors_animations, - lag_ratio=0.0 + *dim_edge_colors_animations, lag_ratio=0.0 ) # Dim the colors of the nodes dim_nodes_animations = [] @@ -113,10 +115,7 @@ def make_pre_dropout_animation( create_x = Create(x_mark) dim_nodes_animations.append(create_x) - dim_nodes_animation = AnimationGroup( - *dim_nodes_animations, - lag_ratio=0.0 - ) + dim_nodes_animation = AnimationGroup(*dim_nodes_animations, lag_ratio=0.0) animation_group = AnimationGroup( dim_edge_colors_animation, @@ -125,6 +124,7 @@ def make_pre_dropout_animation( return animation_group, x_marks + def make_post_dropout_animation( neural_network, layers_to_nodes_to_drop_out, @@ -143,19 +143,15 @@ def make_post_dropout_animation( layers_to_edges_to_dropout = {} for layer in feed_forward_to_feed_forward_layers: layers_to_edges_to_dropout[layer], _ = get_edges_to_drop_out( - layer, - layers_to_nodes_to_drop_out + layer, layers_to_nodes_to_drop_out ) - # Remove the x marks + # Remove the x marks uncreate_animations = [] for x_mark in x_marks: uncreate_x_mark = Uncreate(x_mark) uncreate_animations.append(uncreate_x_mark) - uncreate_x_marks = AnimationGroup( - *uncreate_animations, - lag_ratio=0.0 - ) + uncreate_x_marks = AnimationGroup(*uncreate_animations, lag_ratio=0.0) # Add the edges back create_edge_animations = [] for layer in layers_to_edges_to_dropout.keys(): @@ -164,20 +160,12 @@ def make_post_dropout_animation( for edge_index, edge in enumerate(edges_to_drop_out): edge_copy = edge.copy() edges_to_drop_out[edge_index] = edge_copy - create_edge_animations.append( - FadeIn(edge_copy) - ) + create_edge_animations.append(FadeIn(edge_copy)) - create_edge_animation = AnimationGroup( - *create_edge_animations, - lag_ratio=0.0 - ) + create_edge_animation = AnimationGroup(*create_edge_animations, lag_ratio=0.0) + + return AnimationGroup(uncreate_x_marks, create_edge_animation, lag_ratio=0.0) - return AnimationGroup( - uncreate_x_marks, - create_edge_animation, - lag_ratio=0.0 - ) def make_forward_pass_with_dropout_animation( neural_network, @@ -195,35 +183,25 @@ def make_forward_pass_with_dropout_animation( ) # Iterate through network and get feed forward layers for layer in feed_forward_layers: - layer_args[layer] = { - "dropout_node_indices": layers_to_nodes_to_drop_out[layer] - } + layer_args[layer] = {"dropout_node_indices": layers_to_nodes_to_drop_out[layer]} for layer in feed_forward_to_feed_forward_layers: - _, edge_indices = get_edges_to_drop_out( - layer, - layers_to_nodes_to_drop_out - ) - layer_args[layer] = { - "edge_indices_to_dropout": edge_indices - } + _, edge_indices = get_edges_to_drop_out(layer, layers_to_nodes_to_drop_out) + layer_args[layer] = {"edge_indices_to_dropout": edge_indices} + + return neural_network.make_forward_pass_animation(layer_args=layer_args) - return neural_network.make_forward_pass_animation( - layer_args=layer_args - ) def make_neural_network_dropout_animation( - neural_network, - dropout_rate=0.5, - do_forward_pass=True + neural_network, dropout_rate=0.5, do_forward_pass=True ): """ - Makes a dropout animation for a given neural network. + Makes a dropout animation for a given neural network. - NOTE Does dropout only on feed forward layers. - - 1. Does dropout - 2. If `do_forward_pass` then do forward pass animation - 3. Revert network to pre-dropout appearance + NOTE Does dropout only on feed forward layers. + + 1. Does dropout + 2. If `do_forward_pass` then do forward pass animation + 3. Revert network to pre-dropout appearance """ # Go through the network and get the FeedForwardLayer instances feed_forward_layers = neural_network.filter_layers( @@ -247,26 +225,19 @@ def make_neural_network_dropout_animation( layers_to_nodes_to_drop_out[feed_forward_layer] = nodes_to_drop_out # Make the animation pre_dropout_animation, x_marks = make_pre_dropout_animation( - neural_network, - layers_to_nodes_to_drop_out + neural_network, layers_to_nodes_to_drop_out ) if do_forward_pass: forward_pass_animation = make_forward_pass_with_dropout_animation( - neural_network, - layers_to_nodes_to_drop_out + neural_network, layers_to_nodes_to_drop_out ) else: forward_pass_animation = AnimationGroup() post_dropout_animation = make_post_dropout_animation( - neural_network, - layers_to_nodes_to_drop_out, - x_marks + neural_network, layers_to_nodes_to_drop_out, x_marks ) # Combine the animations into one return Succession( - pre_dropout_animation, - forward_pass_animation, - post_dropout_animation + pre_dropout_animation, forward_pass_animation, post_dropout_animation ) - diff --git a/manim_ml/neural_network/layers/convolutional3d.py b/manim_ml/neural_network/layers/convolutional3d.py index 235fba5..a3cefc0 100644 --- a/manim_ml/neural_network/layers/convolutional3d.py +++ b/manim_ml/neural_network/layers/convolutional3d.py @@ -51,13 +51,6 @@ class Convolutional3DLayer(VGroupNeuralNetworkLayer, ThreeDLayer): about_point=self.get_center(), axis=ThreeDLayer.rotation_axis, ) - """ - self.rotate( - ThreeDLayer.three_d_y_rotation, - about_point=self.get_center(), - axis=[0, 1, 0] - ) - """ def construct_feature_maps(self): """Creates the neural network layer""" diff --git a/manim_ml/neural_network/layers/convolutional3d_to_convolutional3d.py b/manim_ml/neural_network/layers/convolutional3d_to_convolutional3d.py index c3082e2..e4c6b82 100644 --- a/manim_ml/neural_network/layers/convolutional3d_to_convolutional3d.py +++ b/manim_ml/neural_network/layers/convolutional3d_to_convolutional3d.py @@ -260,6 +260,7 @@ class Convolutional3DToConvolutional3D(ConnectiveLayer, ThreeDLayer): pulse_color=ORANGE, cell_width=0.2, show_grid_lines=True, + highlight_color=ORANGE, **kwargs, ): super().__init__( @@ -284,6 +285,7 @@ class Convolutional3DToConvolutional3D(ConnectiveLayer, ThreeDLayer): self.line_color = line_color self.pulse_color = pulse_color self.show_grid_lines = show_grid_lines + self.highlight_color = highlight_color def get_rotated_shift_vectors(self): """ @@ -344,11 +346,11 @@ class Convolutional3DToConvolutional3D(ConnectiveLayer, ThreeDLayer): animations.append(FadeOut(filters)) return Succession(*animations, lag_ratio=1.0) - def animate_filters_one_at_a_time(self): + def animate_filters_one_at_a_time(self, highlight_active_feature_map=False): """Animates each of the filters one at a time""" animations = [] output_feature_maps = self.output_layer.feature_maps - for filter_index in range(len(output_feature_maps)): + for feature_map_index in range(len(output_feature_maps)): # Make filters filters = Filters( self.input_layer, @@ -356,9 +358,28 @@ class Convolutional3DToConvolutional3D(ConnectiveLayer, ThreeDLayer): line_color=self.color, cell_width=self.cell_width, show_grid_lines=self.show_grid_lines, - output_feature_map_to_connect=filter_index, # None means all at once + output_feature_map_to_connect=feature_map_index, # None means all at once ) animations.append(Create(filters)) + # Highlight the feature map + if highlight_active_feature_map: + feature_map = output_feature_maps[feature_map_index] + original_feature_map_color = feature_map.color + # Change the output feature map colors + change_color_animations = [] + change_color_animations.append( + ApplyMethod(feature_map.set_color, self.highlight_color) + ) + # Change the input feature map colors + input_feature_maps = self.input_layer.feature_maps + for input_feature_map in input_feature_maps: + change_color_animations.append( + ApplyMethod(input_feature_map.set_color, self.highlight_color) + ) + # Combine the animations + animations.append( + AnimationGroup(*change_color_animations, lag_ratio=0.0) + ) # Get the rotated shift vectors right_shift, down_shift = self.get_rotated_shift_vectors() left_shift = -1 * right_shift @@ -403,11 +424,36 @@ class Convolutional3DToConvolutional3D(ConnectiveLayer, ThreeDLayer): animations.append(shift_animation) # Remove the filters animations.append(FadeOut(filters)) + # Un-highlight the feature map + if highlight_active_feature_map: + feature_map = output_feature_maps[feature_map_index] + # Change the output feature map colors + change_color_animations = [] + change_color_animations.append( + ApplyMethod(feature_map.set_color, original_feature_map_color) + ) + # Change the input feature map colors + input_feature_maps = self.input_layer.feature_maps + for input_feature_map in input_feature_maps: + change_color_animations.append( + ApplyMethod( + input_feature_map.set_color, original_feature_map_color + ) + ) + # Combine the animations + animations.append( + AnimationGroup(*change_color_animations, lag_ratio=0.0) + ) return Succession(*animations, lag_ratio=1.0) def make_forward_pass_animation( - self, layer_args={}, all_filters_at_once=False, run_time=10.5, **kwargs + self, + layer_args={}, + all_filters_at_once=False, + highlight_active_feature_map=False, + run_time=10.5, + **kwargs, ): """Forward pass animation from conv2d to conv2d""" print(f"All filters at once: {all_filters_at_once}") @@ -415,7 +461,9 @@ class Convolutional3DToConvolutional3D(ConnectiveLayer, ThreeDLayer): if all_filters_at_once: return self.animate_filters_all_at_once() else: - return self.animate_filters_one_at_a_time() + return self.animate_filters_one_at_a_time( + highlight_active_feature_map=highlight_active_feature_map + ) def scale(self, scale_factor, **kwargs): self.cell_width *= scale_factor diff --git a/manim_ml/neural_network/layers/feed_forward.py b/manim_ml/neural_network/layers/feed_forward.py index db2b050..0b1fc8d 100644 --- a/manim_ml/neural_network/layers/feed_forward.py +++ b/manim_ml/neural_network/layers/feed_forward.py @@ -1,6 +1,7 @@ from manim import * from manim_ml.neural_network.layers.parent_layers import VGroupNeuralNetworkLayer + class FeedForwardLayer(VGroupNeuralNetworkLayer): """Handles rendering a layer for a neural network""" @@ -78,16 +79,10 @@ class FeedForwardLayer(VGroupNeuralNetworkLayer): # Make highlight animation succession = Succession( ApplyMethod( - nodes_to_highlight.set_color, - self.animation_dot_color, - run_time=0.25 + nodes_to_highlight.set_color, self.animation_dot_color, run_time=0.25 ), Wait(1.0), - ApplyMethod( - nodes_to_highlight.set_color, - self.node_color, - run_time=0.25 - ), + ApplyMethod(nodes_to_highlight.set_color, self.node_color, run_time=0.25), ) return succession @@ -97,8 +92,7 @@ class FeedForwardLayer(VGroupNeuralNetworkLayer): if "dropout_node_indices" in layer_args: # Drop out certain nodes return self.make_dropout_forward_pass_animation( - layer_args=layer_args, - **kwargs + layer_args=layer_args, **kwargs ) else: # Make highlight animation diff --git a/manim_ml/neural_network/layers/feed_forward_to_feed_forward.py b/manim_ml/neural_network/layers/feed_forward_to_feed_forward.py index ccd816d..d2b4c03 100644 --- a/manim_ml/neural_network/layers/feed_forward_to_feed_forward.py +++ b/manim_ml/neural_network/layers/feed_forward_to_feed_forward.py @@ -68,19 +68,20 @@ class FeedForwardToFeedForward(ConnectiveLayer): return animation_group def make_forward_pass_animation( - self, - layer_args={}, - run_time=1, - feed_forward_dropout=0.0, - **kwargs + self, layer_args={}, run_time=1, feed_forward_dropout=0.0, **kwargs ): """Animation for passing information from one FeedForwardLayer to the next""" path_animations = [] dots = [] for edge_index, edge in enumerate(self.edges): - if not edge_index in layer_args["edge_indices_to_dropout"]: + if ( + not "edge_indices_to_dropout" in layer_args + or not edge_index in layer_args["edge_indices_to_dropout"] + ): dot = Dot( - color=self.animation_dot_color, fill_opacity=1.0, radius=self.dot_radius + color=self.animation_dot_color, + fill_opacity=1.0, + radius=self.dot_radius, ) # Add to dots group dots.append(dot) diff --git a/manim_ml/neural_network/layers/image.py b/manim_ml/neural_network/layers/image.py index 21ebfd9..0bfbf4e 100644 --- a/manim_ml/neural_network/layers/image.py +++ b/manim_ml/neural_network/layers/image.py @@ -56,6 +56,10 @@ class ImageLayer(NeuralNetworkLayer): """Override get right""" return self.image_mobject.get_right() + def scale(self, scale_factor, **kwargs): + """Scales the image mobject""" + self.image_mobject.scale(scale_factor) + @property def width(self): return self.image_mobject.width diff --git a/manim_ml/neural_network/layers/image_to_convolutional3d.py b/manim_ml/neural_network/layers/image_to_convolutional3d.py index b70bd24..59d3403 100644 --- a/manim_ml/neural_network/layers/image_to_convolutional3d.py +++ b/manim_ml/neural_network/layers/image_to_convolutional3d.py @@ -27,10 +27,10 @@ class ImageToConvolutional3DLayer(VGroupNeuralNetworkLayer, ThreeDLayer): """Maps image to convolutional layer""" # Transform the image from the input layer to the num_image_channels = self.input_layer.num_channels - if num_image_channels == 3: - return self.rbg_image_animation() - elif num_image_channels == 1: + if num_image_channels == 1 or num_image_channels == 3: # TODO fix this later return self.grayscale_image_animation() + elif num_image_channels == 3: + return self.rbg_image_animation() else: raise Exception( f"Unrecognized number of image channels: {num_image_channels}" @@ -43,7 +43,6 @@ class ImageToConvolutional3DLayer(VGroupNeuralNetworkLayer, ThreeDLayer): # TODO create image mobjects for each channel and transform # it to the feature maps of the output_layer raise NotImplementedError() - pass def grayscale_image_animation(self): """Handles animation for 1 channel image""" @@ -80,7 +79,7 @@ class ImageToConvolutional3DLayer(VGroupNeuralNetworkLayer, ThreeDLayer): # Scale the max of width or height to the # width of the feature_map max_width_height = max(image_mobject.width, image_mobject.height) - scale_factor = target_feature_map.rectangle_width / max_width_height + scale_factor = target_feature_map.width / max_width_height scale_image = ApplyMethod(image_mobject.scale, scale_factor, run_time=0.5) # Move the image move_image = ApplyMethod(image_mobject.move_to, target_feature_map) diff --git a/manim_ml/neural_network/layers/parent_layers.py b/manim_ml/neural_network/layers/parent_layers.py index bd1a96f..fae7035 100644 --- a/manim_ml/neural_network/layers/parent_layers.py +++ b/manim_ml/neural_network/layers/parent_layers.py @@ -69,10 +69,13 @@ class ConnectiveLayer(VGroupNeuralNetworkLayer): return super()._create_override() def __repr__(self): - return f"{self.__class__.__name__}(" + \ - f"input_layer={self.input_layer.__class__.__name__}," + \ - f"output_layer={self.output_layer.__class__.__name__}," + \ - ")" + return ( + f"{self.__class__.__name__}(" + + f"input_layer={self.input_layer.__class__.__name__}," + + f"output_layer={self.output_layer.__class__.__name__}," + + ")" + ) + class BlankConnective(ConnectiveLayer): """Connective layer to be used when the given pair of layers is undefined""" diff --git a/manim_ml/neural_network/neural_network.py b/manim_ml/neural_network/neural_network.py index 04f13dd..959cd62 100644 --- a/manim_ml/neural_network/neural_network.py +++ b/manim_ml/neural_network/neural_network.py @@ -22,6 +22,7 @@ from manim_ml.neural_network.neural_network_transformations import ( RemoveLayer, ) + class NeuralNetwork(Group): """Neural Network Visualization Container Class""" @@ -34,8 +35,8 @@ class NeuralNetwork(Group): edge_width=2.5, dot_radius=0.03, title=" ", - three_d_phi=-70 * DEGREES, - three_d_theta=-80 * DEGREES, + layout="linear", + layout_direction="left_to_right", ): super(Group, self).__init__() self.input_layers = ListGroup(*input_layers) @@ -46,13 +47,12 @@ class NeuralNetwork(Group): self.dot_radius = dot_radius self.title_text = title self.created = False - # Make the layer fixed in frame if its not 3D - ThreeDLayer.three_d_theta = three_d_theta - ThreeDLayer.three_d_phi = three_d_phi + self.layout = layout + self.layout_direction = layout_direction # TODO take layer_node_count [0, (1, 2), 0] # and make it have explicit distinct subspaces # Place the layers - self._place_layers() + self._place_layers(layout=layout, layout_direction=layout_direction) self.connective_layers, self.all_layers = self._construct_connective_layers() # Make overhead title self.title = Text(self.title_text, font_size=DEFAULT_FONT_SIZE / 2) @@ -67,7 +67,7 @@ class NeuralNetwork(Group): # Print neural network print(repr(self)) - def _place_layers(self): + def _place_layers(self, layout="linear", layout_direction="top_to_bottom"): """Creates the neural network""" # TODO implement more sophisticated custom layouts # Default: Linear layout @@ -79,26 +79,65 @@ class NeuralNetwork(Group): 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, - ] - ) + if layout_direction == "left_to_right": + shift_vector = np.array( + [ + ( + previous_layer.get_width() / 2 + + current_layer.get_width() / 2 + - 0.2 + ), + 0, + 0, + ] + ) + elif layout_direction == "top_to_bottom": + shift_vector = np.array( + [ + 0, + -( + previous_layer.get_width() / 2 + + current_layer.get_width() / 2 + - 0.2 + ), + 0, + ] + ) + else: + raise Exception( + f"Unrecognized layout direction: {layout_direction}" + ) else: - shift_vector = np.array( - [ - (previous_layer.get_width() / 2 + current_layer.get_width() / 2) - + self.layer_spacing, - 0, - 0, - ] - ) + if layout_direction == "left_to_right": + shift_vector = np.array( + [ + ( + previous_layer.get_width() / 2 + + current_layer.get_width() / 2 + ) + + self.layer_spacing, + 0, + 0, + ] + ) + elif layout_direction == "top_to_bottom": + shift_vector = np.array( + [ + 0, + -( + ( + previous_layer.get_width() / 2 + + current_layer.get_width() / 2 + ) + + self.layer_spacing + ), + 0, + ] + ) + else: + raise Exception( + f"Unrecognized layout direction: {layout_direction}" + ) current_layer.shift(shift_vector) def _construct_connective_layers(self): @@ -163,7 +202,7 @@ class NeuralNetwork(Group): if isinstance(layer, ConnectiveLayer): """ NOTE: By default a connective layer will get the combined - layer_args of the layers it is connecting and itself. + layer_args of the layers it is connecting and itself. """ before_layer_args = {} current_layer_args = {} @@ -176,9 +215,9 @@ class NeuralNetwork(Group): after_layer_args = layer_args[layer.output_layer] # Merge the two dicts current_layer_args = { - **before_layer_args, - **current_layer_args, - **after_layer_args + **before_layer_args, + **current_layer_args, + **after_layer_args, } else: current_layer_args = {} @@ -229,14 +268,18 @@ class NeuralNetwork(Group): """Overriden scale""" for layer in self.all_layers: layer.scale(scale_factor, **kwargs) - # super().scale(scale_factor) + # Place layers with scaled spacing + self.layer_spacing *= scale_factor + self._place_layers(layout=self.layout, layout_direction=self.layout_direction) def filter_layers(self, function): """Filters layers of the network given function""" layers_to_return = [] for layer in self.all_layers: func_out = function(layer) - assert isinstance(func_out, bool), "Filter layers function returned a non-boolean type." + assert isinstance( + func_out, bool + ), "Filter layers function returned a non-boolean type." if func_out: layers_to_return.append(layer) @@ -257,6 +300,7 @@ class NeuralNetwork(Group): string_repr = "NeuralNetwork([\n" + inner_string + "])" return string_repr + class FeedForwardNeuralNetwork(NeuralNetwork): """NeuralNetwork with just feed forward layers""" diff --git a/manim_ml/neural_network/neural_network_transformations.py b/manim_ml/neural_network/neural_network_transformations.py index c1a6e89..28423c4 100644 --- a/manim_ml/neural_network/neural_network_transformations.py +++ b/manim_ml/neural_network/neural_network_transformations.py @@ -4,6 +4,7 @@ from manim import * from manim_ml.neural_network.layers.util import get_connective_layer + class RemoveLayer(AnimationGroup): """ Animation for removing a layer from a neural network. diff --git a/tests/test_decision_tree.py b/tests/test_decision_tree.py index cc92d8e..fec75cd 100644 --- a/tests/test_decision_tree.py +++ b/tests/test_decision_tree.py @@ -1,27 +1,31 @@ from manim import * -from manim_ml.decision_tree.decision_tree import DecisionTreeDiagram, DecisionTreeSurface, IrisDatasetPlot +from manim_ml.decision_tree.decision_tree import ( + DecisionTreeDiagram, + DecisionTreeSurface, + IrisDatasetPlot, +) from sklearn.tree import DecisionTreeClassifier from sklearn import datasets import sklearn import matplotlib.pyplot as plt + def learn_iris_decision_tree(iris): decision_tree = DecisionTreeClassifier( - random_state=1, - max_depth=3, - max_leaf_nodes=6 + random_state=1, max_depth=3, max_leaf_nodes=6 ) decision_tree = decision_tree.fit(iris.data[:, :2], iris.target) # output the decisioin tree in some format return decision_tree + def make_sklearn_tree(dataset, max_tree_depth=3): tree = learn_iris_decision_tree(dataset) feature_names = dataset.feature_names[0:2] return tree, tree.tree_ -class DecisionTreeScene(Scene): +class DecisionTreeScene(Scene): def construct(self): """Makes a decision tree object""" iris_dataset = datasets.load_iris() @@ -36,15 +40,8 @@ class DecisionTreeScene(Scene): "images/iris_dataset/VeriscolorFlower.jpeg", "images/iris_dataset/VirginicaFlower.jpeg", ], - class_names=[ - "Setosa", - "Veriscolor", - "Virginica" - ], - feature_names=[ - "Sepal Length", - "Sepal Width" - ] + class_names=["Setosa", "Veriscolor", "Virginica"], + feature_names=["Sepal Length", "Sepal Width"], ) decision_tree.move_to(ORIGIN) create_decision_tree = Create(decision_tree, traversal_order="bfs") @@ -53,7 +50,6 @@ class DecisionTreeScene(Scene): class SurfacePlot(Scene): - def construct(self): iris_dataset = datasets.load_iris() iris_dataset_plot = IrisDatasetPlot(iris_dataset) @@ -63,13 +59,7 @@ class SurfacePlot(Scene): # make the decision tree classifier decision_tree_classifier, sklearn_tree = make_sklearn_tree(iris_dataset) decision_tree_surface = DecisionTreeSurface( - decision_tree_classifier, - iris_dataset.data, - iris_dataset_plot.axes_group[0] + decision_tree_classifier, iris_dataset.data, iris_dataset_plot.axes_group[0] ) - self.play( - Create( - decision_tree_surface - ) - ) - self.wait(1) \ No newline at end of file + self.play(Create(decision_tree_surface)) + self.wait(1) diff --git a/tests/test_nn_dropout.py b/tests/test_nn_dropout.py index 9df4ceb..c61c75d 100644 --- a/tests/test_nn_dropout.py +++ b/tests/test_nn_dropout.py @@ -1,5 +1,7 @@ from manim import * -from manim_ml.neural_network.animations.dropout import make_neural_network_dropout_animation +from manim_ml.neural_network.animations.dropout import ( + make_neural_network_dropout_animation, +) from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer from manim_ml.neural_network.layers.image import ImageLayer from PIL import Image @@ -11,6 +13,7 @@ config.pixel_width = 1900 config.frame_height = 5.0 config.frame_width = 5.0 + def make_code_snippet(): code_str = """ nn = NeuralNetwork([ @@ -40,17 +43,19 @@ def make_code_snippet(): return code + class DropoutNeuralNetworkScene(Scene): def construct(self): # Make nn - nn = NeuralNetwork([ + nn = NeuralNetwork( + [ FeedForwardLayer(3, rectangle_color=BLUE), FeedForwardLayer(5, rectangle_color=BLUE), FeedForwardLayer(3, rectangle_color=BLUE), FeedForwardLayer(5, rectangle_color=BLUE), FeedForwardLayer(4, rectangle_color=BLUE), ], - layer_spacing=0.4 + layer_spacing=0.4, ) # Center the nn nn.move_to(ORIGIN) @@ -59,17 +64,16 @@ class DropoutNeuralNetworkScene(Scene): code_snippet = make_code_snippet() self.add(code_snippet) code_snippet.next_to(nn, DOWN * 0.7) - Group(code_snippet, nn).move_to(ORIGIN) + Group(code_snippet, nn).move_to(ORIGIN) # Play animation self.play( make_neural_network_dropout_animation( - nn, - dropout_rate=0.25, - do_forward_pass=True + nn, dropout_rate=0.25, do_forward_pass=True ) ) self.wait(1) + if __name__ == "__main__": """Render all scenes""" dropout_nn_scene = DropoutNeuralNetworkScene() diff --git a/tests/test_nn_scale.py b/tests/test_nn_scale.py new file mode 100644 index 0000000..692402a --- /dev/null +++ b/tests/test_nn_scale.py @@ -0,0 +1,44 @@ +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.5), + Convolutional3DLayer(1, 7, 7, 3, 3, filter_spacing=0.32), + Convolutional3DLayer(3, 5, 5, 3, 3, filter_spacing=0.32), + FeedForwardLayer(3), + ], + layer_spacing=0.25, + ) + # Center the nn + nn.move_to(ORIGIN) + nn.scale(1.3) + self.add(nn) + """ + self.play( + FadeIn(nn) + ) + """ + # Play animation + forward_pass = nn.make_forward_pass_animation( + corner_pulses=False, all_filters_at_once=False, highlight_filters=True + ) + self.wait(1) + self.play(forward_pass)