Finished oracle guidance video. Integrated various changes necessary to complete this.

examples/paper_visualizations/oracle_guidance/oracle_guidance.py

@@ -0,0 +1,505 @@
+"""
+    Here is a animated explanatory figure for the "Oracle Guided Image Synthesis with Relative Queries" paper. 
+"""
+from manim import *
+from manim_ml.neural_network.layers import triplet
+from manim_ml.neural_network.layers.image import ImageLayer
+from manim_ml.neural_network.layers.paired_query import PairedQueryLayer
+from manim_ml.neural_network.layers.triplet import TripletLayer
+from manim_ml.neural_network.neural_network import NeuralNetwork
+from manim_ml.neural_network.layers import FeedForwardLayer, EmbeddingLayer
+from manim_ml.neural_network.layers.util import get_connective_layer
+import os
+
+from manim_ml.probability import GaussianDistribution
+
+# Make the specific scene
+config.pixel_height = 1200
+config.pixel_width = 1900
+config.frame_height = 6.0
+config.frame_width = 6.0
+
+class Localizer():
+    """
+        Holds the localizer object, which contains the queries, images, etc.
+        needed to represent a localization run. 
+    """
+    
+    def __init__(self, axes):
+        # Set dummy values for these 
+        self.index = -1
+        self.axes = axes
+        self.num_queries = 3
+        self.assets_path = "../../../assets/oracle_guidance"
+        self.ground_truth_image_path = os.path.join(self.assets_path, "ground_truth.jpg")
+        self.ground_truth_location = np.array([2, 3])
+        # Prior distribution
+        print("initial gaussian")
+        self.prior_distribution = GaussianDistribution(
+            self.axes,
+            mean=np.array([0.0, 0.0]),
+            cov=np.array([[3, 0], [0, 3]]),
+            dist_theme="ellipse",
+            color=GREEN,
+        ) 
+        # Define the query images and embedded locations
+        # Contains image paths [(positive_path, negative_path), ...]
+        self.query_image_paths = [
+            (os.path.join(self.assets_path, "positive_1.jpg"), os.path.join(self.assets_path, "negative_1.jpg")),
+            (os.path.join(self.assets_path, "positive_2.jpg"), os.path.join(self.assets_path, "negative_2.jpg")),
+            (os.path.join(self.assets_path, "positive_3.jpg"), os.path.join(self.assets_path, "negative_3.jpg")),
+        ] 
+        # Contains 2D locations for each image [([2, 3], [2, 4]), ...]
+        self.query_locations = [
+            (np.array([-1, -1]), np.array([1, 1])),
+            (np.array([1, -1]), np.array([-1, 1])), 
+            (np.array([0.3, -0.6]), np.array([-0.5, 0.7])),
+        ] 
+        # Make the covariances for each query
+        self.query_covariances = [
+            (np.array([[0.3, 0], [0.0, 0.2]]), np.array([[0.2, 0], [0.0, 0.2]])),
+            (np.array([[0.2, 0], [0.0, 0.2]]), np.array([[0.2, 0], [0.0, 0.2]])), 
+            (np.array([[0.2, 0], [0.0, 0.2]]), np.array([[0.2, 0], [0.0, 0.2]])),
+        ] 
+        # Posterior distributions over time GaussianDistribution objects
+        self.posterior_distributions = [
+            GaussianDistribution(
+                self.axes, 
+                dist_theme="ellipse", 
+                color=GREEN,
+                mean=np.array([-0.3, -0.3]),
+                cov=np.array([[5, -4], [-4, 6]])
+            ).scale(0.6),
+            GaussianDistribution(
+                self.axes, 
+                dist_theme="ellipse", 
+                color=GREEN,
+                mean=np.array([0.25, -0.25]),
+                cov=np.array([[3, -2], [-2, 4]])
+            ).scale(0.35),
+            GaussianDistribution(
+                self.axes, 
+                dist_theme="ellipse", 
+                color=GREEN,
+                mean=np.array([0.4, -0.35]),
+                cov=np.array([[1, 0], [0, 1]])
+            ).scale(0.3),
+        ]
+        # Some assumptions
+        assert len(self.query_locations) == len(self.query_image_paths) 
+        assert len(self.query_locations) == len(self.posterior_distributions)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        """Steps through each localization time instance"""
+        if self.index < len(self.query_image_paths):
+            self.index += 1
+        else:
+            raise StopIteration
+
+        # Return query_paths, query_locations, posterior
+        out_tuple = (
+            self.query_image_paths[self.index], 
+            self.query_locations[self.index], 
+            self.posterior_distributions[self.index],
+            self.query_covariances[self.index]
+        )
+
+        return out_tuple
+
+class OracleGuidanceVisualization(Scene):
+
+    def __init__(self):
+        super().__init__()
+        self.neural_network, self.embedding_layer = self.make_vae()
+        self.localizer = iter(Localizer(self.embedding_layer.axes))
+        self.subtitle = None
+        self.title = None
+        # Set image paths
+        # VAE embedding animation image paths
+        self.assets_path = "../../../assets/oracle_guidance"
+        self.input_embed_image_path = os.path.join(self.assets_path, "input_image.jpg")
+        self.output_embed_image_path = os.path.join(self.assets_path, "output_image.jpg")
+
+    def make_vae(self):
+        """Makes a simple VAE architecture"""
+        embedding_layer = EmbeddingLayer(dist_theme="ellipse")
+        self.encoder = NeuralNetwork([
+            FeedForwardLayer(5),
+            FeedForwardLayer(3),
+            embedding_layer,
+        ])
+
+        self.decoder = NeuralNetwork([
+            FeedForwardLayer(3), 
+            FeedForwardLayer(5),
+        ])
+
+        neural_network = NeuralNetwork([
+            self.encoder, 
+            self.decoder
+        ])
+
+        neural_network.shift(DOWN*0.4)
+        return neural_network, embedding_layer
+
+    @override_animation(Create)
+    def _create_animation(self):
+        animation_group = AnimationGroup(
+            Create(self.neural_network)
+        )
+        
+        return animation_group
+
+    def insert_at_start(self, layer, create=True):
+        """Inserts a layer at the beggining of the network"""
+        # Note: Does not move the rest of the network
+        current_first_layer = self.encoder.all_layers[0]
+        # Get connective layer
+        connective_layer = get_connective_layer(layer, current_first_layer)
+        # Insert both layers
+        self.encoder.all_layers.insert(0, layer)
+        self.encoder.all_layers.insert(1, connective_layer)
+        # Move layers to the correct location
+        # TODO: Fix this cause its hacky
+        layer.shift(DOWN*0.4)
+        layer.shift(LEFT*2.35)
+        # Make insert animation
+        if not create:
+            animation_group = AnimationGroup(
+                Create(connective_layer)
+            )
+        else:
+            animation_group = AnimationGroup(
+                Create(layer),
+                Create(connective_layer)
+            )
+        self.play(animation_group)
+        
+
+    def remove_start_layer(self):
+        """Removes the first layer of the network"""
+        first_layer = self.encoder.all_layers.remove_at_index(0)
+        first_connective = self.encoder.all_layers.remove_at_index(0)
+        # Make remove animations
+        animation_group = AnimationGroup(
+            FadeOut(first_layer),
+            FadeOut(first_connective)
+        )
+
+        self.play(animation_group)
+
+    def insert_at_end(self, layer):
+        """Inserts the given layer at the end of the network"""
+        current_last_layer = self.decoder.all_layers[-1]
+        # Get connective layer
+        connective_layer = get_connective_layer(current_last_layer, layer)
+        # Insert both layers
+        self.decoder.all_layers.add(connective_layer)
+        self.decoder.all_layers.add(layer)
+        # Move layers to the correct location
+        # TODO: Fix this cause its hacky
+        layer.shift(DOWN*0.4)
+        layer.shift(RIGHT*2.35)
+        # Make insert animation
+        animation_group = AnimationGroup(
+            Create(layer),
+            Create(connective_layer)
+        )
+        self.play(animation_group)
+
+    def remove_end_layer(self):
+        """Removes the lsat layer of the network"""
+        first_layer = self.decoder.all_layers.remove_at_index(-1)
+        first_connective = self.decoder.all_layers.remove_at_index(-1)
+        # Make remove animations
+        animation_group = AnimationGroup(
+            FadeOut(first_layer),
+            FadeOut(first_connective)
+        )
+
+        self.play(animation_group)
+
+    def change_title(self, text, title_location=np.array([0, 1.25, 0]), font_size=24):
+        """Changes title to the given text"""
+        if self.title is None:
+            self.title = Text(text, font_size=font_size)
+            self.title.move_to(title_location)
+            self.add(self.title)
+            self.play(Write(self.title), run_time=1)
+            self.wait(1)
+            return
+
+        self.play(Unwrite(self.title))
+        new_title = Text(text, font_size=font_size)
+        new_title.move_to(self.title)
+        self.title = new_title
+        self.wait(0.1)
+        self.play(Write(self.title))
+
+    def change_subtitle(self, text, title_location=np.array([0, 0.9, 0]), font_size=20):
+        """Changes subtitle to the next algorithm step"""
+        if self.subtitle is None:
+            self.subtitle = Text(text, font_size=font_size)
+            self.subtitle.move_to(title_location)
+            self.play(Write(self.subtitle))
+            return
+
+        self.play(Unwrite(self.subtitle))
+        new_title = Text(text, font_size=font_size)
+        new_title.move_to(title_location)
+        self.subtitle = new_title
+        self.wait(0.1)
+        self.play(Write(self.subtitle))
+
+    def make_embed_input_image_animation(self, input_image_path, output_image_path):
+        """Makes embed input image animation"""
+        # insert the input image at the begginging
+        input_image_layer = ImageLayer.from_path(input_image_path)
+        input_image_layer.scale(0.6)
+        current_first_layer = self.encoder.all_layers[0]
+        # Get connective layer
+        connective_layer = get_connective_layer(input_image_layer, current_first_layer)
+        # Insert both layers
+        self.encoder.all_layers.insert(0, input_image_layer)
+        self.encoder.all_layers.insert(1, connective_layer)
+        # Move layers to the correct location
+        # TODO: Fix this cause its hacky
+        input_image_layer.shift(DOWN*0.4)
+        input_image_layer.shift(LEFT*2.35)
+        # Play full forward pass
+        forward_pass = self.neural_network.make_forward_pass_animation(
+            layer_args=
+            {
+                self.encoder: {
+                    self.embedding_layer: {
+                        "dist_args": {
+                            "cov": np.array([[1.5, 0], [0, 1.5]]),
+                            "mean": np.array([0.5, 0.5]),
+                            "dist_theme": "ellipse",
+                            "color": ORANGE
+                        }
+                    }
+                }
+            }
+        )
+        self.play(forward_pass)
+        # insert the output image at the end
+        output_image_layer = ImageLayer.from_path(output_image_path)
+        output_image_layer.scale(0.6)
+        self.insert_at_end(output_image_layer)
+        # Remove the input and output layers
+        self.remove_start_layer()
+        self.remove_end_layer()
+        # Remove the latent distribution
+        self.play(FadeOut(self.embedding_layer.latent_distribution))
+
+    def make_localization_time_step(self, old_posterior):
+        """
+            Performs one query update for the localization procedure
+            
+            Procedure:
+            a. Embed query input images
+            b. Oracle is asked a query
+            c. Query is embedded
+            d. Show posterior update
+            e. Show current recomendation
+        """
+        # Helper functions
+        def embed_query_to_latent_space(query_locations, query_covariance):
+            """Makes animation for a paired query"""
+            # Assumes first layer of neural network is a PairedQueryLayer
+            # Make the embedding animation
+            # Wait
+            self.play(Wait(1))
+            # Embed the query to latent space
+            self.change_subtitle("3. Embed the Query in Latent Space")
+            # Make forward pass animation
+            self.embedding_layer.paired_query_mode = True
+            # Make embedding embed query animation
+            embed_query_animation = self.encoder.make_forward_pass_animation(
+                run_time=5,
+                layer_args={
+                    self.embedding_layer: {
+                        "positive_dist_args": {
+                            "cov": query_covariance[0],
+                            "mean": query_locations[0],
+                            "dist_theme": "ellipse",
+                            "color": BLUE
+                        },
+                        "negative_dist_args": {
+                            "cov": query_covariance[1],
+                            "mean": query_locations[1],
+                            "dist_theme": "ellipse",
+                            "color": RED
+                        }
+                    }
+                } 
+            )
+            self.play(embed_query_animation)
+
+        # Access localizer information
+        query_paths, query_locations, posterior_distribution, query_covariances = next(self.localizer)
+        positive_path, negative_path = query_paths
+        # Make subtitle for present user with query
+        self.change_subtitle("2. Present User with Query")
+        # Insert the layer into the encoder
+        query_layer = PairedQueryLayer.from_paths(positive_path, negative_path, grayscale=False)
+        query_layer.scale(0.5)
+        self.insert_at_start(query_layer)
+        # Embed query to latent space
+        query_to_latent_space_animation = embed_query_to_latent_space(
+            query_locations,
+            query_covariances
+        )
+        # Wait
+        self.play(Wait(1))
+        # Update the posterior
+        self.change_subtitle("4. Update the Posterior")
+        # Remove the old posterior
+        self.play(
+            ReplacementTransform(old_posterior, posterior_distribution)
+        )
+        """
+        self.play(
+            self.embedding_layer.remove_gaussian_distribution(self.localizer.posterior_distribution)
+        )
+        """
+        # self.embedding_layer.add_gaussian_distribution(posterior_distribution)
+        # self.localizer.posterior_distribution = posterior_distribution
+        # Remove query layer
+        self.remove_start_layer()
+        # Remove query ellipses
+        
+        fade_outs = []
+        for dist in self.embedding_layer.gaussian_distributions:
+            self.embedding_layer.gaussian_distributions.remove(dist)
+            fade_outs.append(FadeOut(dist))
+        
+        if not len(fade_outs) == 0:
+            fade_outs = AnimationGroup(*fade_outs)
+            self.play(fade_outs)
+
+        return posterior_distribution
+
+    def make_generate_estimate_animation(self, estimate_image_path):
+        """Makes the generate estimate animation"""
+        # Change embedding layer mode
+        self.embedding_layer.paired_query_mode = False
+        # Sample from posterior distribution
+        # self.embedding_layer.latent_distribution = self.localizer.posterior_distribution
+        emb_to_ff_ind = self.neural_network.all_layers.index_of(self.encoder)
+        embedding_to_ff = self.neural_network.all_layers[emb_to_ff_ind + 1]
+        self.play(embedding_to_ff.make_forward_pass_animation())
+        # Pass through decoder
+        self.play(self.decoder.make_forward_pass_animation(), run_time=1)
+        # Create Image layer after the decoder
+        output_image_layer = ImageLayer.from_path(estimate_image_path)
+        output_image_layer.scale(0.5)
+        self.insert_at_end(output_image_layer)
+        # Wait
+        self.wait(1)
+        # Remove the image at the end
+        print(self.neural_network)
+        self.remove_end_layer()
+        
+    def make_triplet_forward_animation(self):
+        """Make triplet forward animation"""
+        # Make triplet layer
+        anchor_path = os.path.join(self.assets_path, "anchor.jpg")
+        positive_path = os.path.join(self.assets_path, "positive.jpg")
+        negative_path = os.path.join(self.assets_path, "negative.jpg")
+        triplet_layer = TripletLayer.from_paths(anchor_path, positive_path, negative_path, grayscale=False, font_size=100, buff=1.05)
+        triplet_layer.scale(0.10)
+        self.insert_at_start(triplet_layer)
+        # Make latent triplet animation
+        self.play(
+            self.encoder.make_forward_pass_animation(
+                layer_args={
+                    self.embedding_layer: {
+                        "triplet_args": {
+                            "anchor_dist": {
+                                "cov": np.array([[0.3, 0], [0, 0.3]]),
+                                "mean": np.array([0.7, 1.4]),
+                                "dist_theme": "ellipse",
+                                "color": BLUE
+                            },
+                            "positive_dist": {
+                                "cov": np.array([[0.2, 0], [0, 0.2]]),
+                                "mean": np.array([0.8, -0.4]),
+                                "dist_theme": "ellipse",
+                                "color": GREEN
+                            },
+                            "negative_dist": {
+                                "cov": np.array([[0.4, 0], [0, 0.25]]),
+                                "mean": np.array([-1, -1.2]),
+                                "dist_theme": "ellipse",
+                                "color": RED
+                            }
+                        }
+                    }
+                },
+                run_time=3
+            )
+        )
+
+    def construct(self):
+        """
+            Makes the whole visualization.
+
+            1. Create the Architecture
+                a. Create the traditional VAE architecture with images
+            2. The Localization Procedure
+            3. The Training Procedure
+        """
+        # 1. Create the Architecture
+        self.neural_network.scale(1.2)
+        create_vae = Create(self.neural_network)
+        self.play(create_vae, run_time=3)
+        # Make changing title
+        self.change_title("Oracle Guided Image Synthesis\n      with Relative Queries")
+        # 2. The Localization Procedure
+        self.change_title("The Localization Procedure")
+        # Make algorithm subtitle
+        self.change_subtitle("Algorithm Steps")
+        # Wait
+        self.play(Wait(1))
+        # Make prior distribution subtitle
+        self.change_subtitle("1. Calculate Prior Distribution")
+        # Draw the prior distribution
+        self.play(Create(self.localizer.prior_distribution))
+        old_posterior = self.localizer.prior_distribution
+        # For N queries update the posterior
+        for query_index in range(self.localizer.num_queries):
+            # Make localization iteration
+            old_posterior = self.make_localization_time_step(old_posterior)
+            self.play(Wait(1))
+            if not query_index == self.localizer.num_queries - 1:
+                # Repeat
+                self.change_subtitle("5. Repeat") 
+                # Wait a second
+                self.play(Wait(1))
+        # Generate final estimate
+        self.change_subtitle("5. Generate Estimate Image")
+        # Generate an estimate image
+        estimate_image_path = os.path.join(self.assets_path, "estimate_image.jpg")
+        self.make_generate_estimate_animation(estimate_image_path)
+        self.wait(1)
+        # Remove old posterior
+        self.play(FadeOut(old_posterior))
+        # 3. The Training Procedure
+        self.change_title("The Training Procedure")
+        # Make training animation
+        # Do an Image forward pass
+        self.change_subtitle("1. Unsupervised Image Reconstruction")
+        self.make_embed_input_image_animation(
+            self.input_embed_image_path, 
+            self.output_embed_image_path
+        )
+        self.wait(1)
+        # Do triplet forward pass
+        self.change_subtitle("2. Triplet Loss in Latent Space")
+        self.make_triplet_forward_animation()
+        self.wait(1)

manim_ml/image.py

@@ -34,15 +34,17 @@ class GrayscaleImageMobject(ImageMobject):
 class LabeledColorImage(Group):
     """Labeled Color Image"""
 
-    def __init__(self, image, color=RED, label="Positive", stroke_width=5):
+    def __init__(self, image, color=RED, label="Positive", stroke_width=5,
+                font_size=24, buff=0.2):
         super().__init__()
         self.image = image
         self.color = color
         self.label = label
         self.stroke_width = stroke_width
+        self.font_size = font_size
 
-        text = Text(label).scale(2)
-        text.next_to(self.image, UP, buff=1.0)
+        text = Text(label, font_size=self.font_size)
+        text.next_to(self.image, UP, buff=buff)
         rectangle = SurroundingRectangle(
             self.image, 
             color=color,

manim_ml/list_group.py

@@ -18,7 +18,7 @@ class ListGroup(Mobject):
 
     def remove_at_index(self, index):
         """Removes item at index"""
-        if index < 0 or index > len(self.items):
+        if index > len(self.items):
             raise Exception(f"ListGroup index out of range: {index}")
         item = self.items[index]
         del self.items[index]

manim_ml/neural_network/layers/embedding.py

@@ -9,6 +9,8 @@ class EmbeddingLayer(VGroupNeuralNetworkLayer):
                 covariance=np.array([[1.0, 0], [0, 1.0]]), dist_theme="gaussian", 
                 paired_query_mode=False, **kwargs):
         super(VGroupNeuralNetworkLayer, self).__init__(**kwargs)
+        self.gaussian_distributions = VGroup()
+        self.add(self.gaussian_distributions)
         self.point_radius = point_radius
         self.dist_theme = dist_theme
         self.paired_query_mode = paired_query_mode
@@ -16,8 +18,8 @@ class EmbeddingLayer(VGroupNeuralNetworkLayer):
             tips=False,
             x_length=0.8,
             y_length=0.8,
-            x_range=(-2.0, 2.0),
-            y_range=(-2.0, 2.0),
+            x_range=(-1.4, 1.4),
+            y_range=(-1.8, 1.8),
             x_axis_config={
                 "include_ticks": False,
                 "stroke_width": 0.0
@@ -33,8 +35,20 @@ class EmbeddingLayer(VGroupNeuralNetworkLayer):
         self.point_cloud = self.construct_gaussian_point_cloud(mean, covariance)
         self.add(self.point_cloud)
         # Make latent distribution
-        self.latent_distribution = GaussianDistribution(self.axes, mean=mean, cov=covariance,
-                                                        dist_theme=self.dist_theme) # Use defaults
+        self.latent_distribution = GaussianDistribution(self.axes, mean=mean, cov=covariance) # Use defaults
+
+    def add_gaussian_distribution(self, gaussian_distribution):
+        """Adds given GaussianDistribution to the list"""
+        self.gaussian_distributions.add(gaussian_distribution)
+
+        return Create(gaussian_distribution)
+
+    def remove_gaussian_distribution(self, gaussian_distribution):
+        """Removes the given gaussian distribution from the embedding"""
+        for gaussian in self.gaussian_distributions:
+            if gaussian == gaussian_distribution:
+                self.gaussian_distributions.remove(gaussian_distribution)
+                return FadeOut(gaussian)
 
     def sample_point_location_from_distribution(self):
         """Samples from the current latent distribution"""
@@ -50,57 +64,112 @@ class EmbeddingLayer(VGroupNeuralNetworkLayer):
         """Returns mean of latent distribution in axes frame"""
         return self.axes.coords_to_point(self.latent_distribution.mean)
 
-    def construct_gaussian_point_cloud(self, mean, covariance, point_color=BLUE,
-                                    num_points=200):
+    def construct_gaussian_point_cloud(self, mean, covariance, point_color=WHITE,
+                                    num_points=400):
         """Plots points sampled from a Gaussian with the given mean and covariance"""
         # Sample points from a Gaussian
+        np.random.seed(5)
         points = np.random.multivariate_normal(mean, covariance, num_points)
         # Add each point to the axes
         point_dots = VGroup()
         for point in points:
             point_location = self.axes.coords_to_point(*point)
             dot = Dot(point_location, color=point_color, radius=self.point_radius/2) 
+            dot.set_z_index(-1)
             point_dots.add(dot)
 
         return point_dots
 
-    def make_paired_query_embedding_animation(self):
-        """Embed paired query"""
-        animations = []
-        # Make the animation
-        
-        # Animation group
-        animation_group = AnimationGroup(
-            *animations, 
-            lag_ratio=1.0
-        )
-        
-        return animation_group
-
     def make_forward_pass_animation(self, layer_args={}, **kwargs):
         """Forward pass animation"""
         animations = [] 
-        if not self.paired_query_mode:
-            # Normal embedding mode
-            # Make ellipse object corresponding to the latent distribution
-            self.latent_distribution = GaussianDistribution(
+        if "triplet_args" in layer_args:
+            triplet_args = layer_args["triplet_args"]
+            positive_dist_args = triplet_args["positive_dist"]
+            negative_dist_args = triplet_args["negative_dist"]
+            anchor_dist_args = triplet_args["anchor_dist"]
+            # Create each dist
+            anchor_dist = GaussianDistribution(
                 self.axes,
-                dist_theme=self.dist_theme, 
-                cov=np.array([[0.8, 0], [0.0, 0.8]])
-            ) # Use defaults
+                **anchor_dist_args
+            )
+            animations.append(Create(anchor_dist))
+
+            positive_dist = GaussianDistribution(
+                self.axes,
+                **positive_dist_args
+            )
+            animations.append(Create(positive_dist))
+
+            negative_dist = GaussianDistribution(
+                self.axes,
+                **negative_dist_args
+            )
+            animations.append(Create(negative_dist))
+            # Draw edges in between anchor and positive, anchor and negative
+            anchor_positive = Line(
+                anchor_dist.get_center(),
+                positive_dist.get_center(),
+                color=GOLD,
+                stroke_width=DEFAULT_STROKE_WIDTH/2
+            )
+            anchor_positive.set_z_index(3)
+            animations.append(Create(anchor_positive))
+
+            anchor_negative = Line(
+                anchor_dist.get_center(),
+                negative_dist.get_center(),
+                color=GOLD,
+                stroke_width=DEFAULT_STROKE_WIDTH/2
+            )
+            anchor_negative.set_z_index(3)
+
+            animations.append(Create(anchor_negative))
+        elif not self.paired_query_mode:
+            # Normal embedding mode
+            if "dist_args" in layer_args:
+                scale_factor = 1.0
+                if "scale_factor" in layer_args:
+                    scale_factor = layer_args["scale_factor"]
+                self.latent_distribution = GaussianDistribution(
+                    self.axes,
+                    **layer_args["dist_args"]
+                ).scale(scale_factor)
+            else:
+                # Make ellipse object corresponding to the latent distribution
+                # self.latent_distribution = GaussianDistribution(
+                #     self.axes, 
+                #     dist_theme=self.dist_theme, 
+                #     cov=np.array([[0.8, 0], [0.0, 0.8]])
+                # )
+                pass
             # Create animation
-            #create_distribution = Create(self.latent_distribution.construct_gaussian_distribution(self.latent_distribution.mean, self.latent_distribution.cov)) #Create(self.latent_distribution)
-            create_distribution = Create(self.latent_distribution.ellipses) 
+            create_distribution = Create(self.latent_distribution)
             animations.append(create_distribution)
-            
-            animation_group = AnimationGroup(*animations)
-
-            return animation_group
         else:
             # Paired Query Mode
+            assert "positive_dist_args" in layer_args
+            assert "negative_dist_args" in layer_args
+            positive_dist_args = layer_args["positive_dist_args"]
+            negative_dist_args = layer_args["negative_dist_args"]
             # Handle logic for embedding a paired query into the embedding layer
-            paired_query_embedding_animation = self.make_paired_query_embedding_animation()
-            return paired_query_embedding_animation
+            positive_dist = GaussianDistribution(
+                self.axes,
+                **positive_dist_args
+            )
+            self.gaussian_distributions.add(positive_dist)
+            negative_dist = GaussianDistribution(
+                self.axes,
+                **negative_dist_args
+            )
+            self.gaussian_distributions.add(negative_dist)
+
+            animations.append(Create(positive_dist))
+            animations.append(Create(negative_dist))
+
+        animation_group = AnimationGroup(*animations)
+
+        return animation_group
 
     @override_animation(Create)
     def _create_override(self, **kwargs):

manim_ml/neural_network/layers/feed_forward.py

@@ -41,6 +41,7 @@ class FeedForwardLayer(VGroupNeuralNetworkLayer):
         self.surrounding_rectangle = SurroundingRectangle(self.node_group, color=self.rectangle_color, 
                                                         fill_color=self.rectangle_fill_color, fill_opacity=1.0, 
                                                         buff=self.layer_buffer, stroke_width=self.rectangle_stroke_width)
+        self.surrounding_rectangle.set_z_index(1)                                                
         # Add the objects to the class
         self.add(self.surrounding_rectangle, self.node_group)
 

manim_ml/neural_network/layers/image.py

@@ -2,6 +2,8 @@ from manim import *
 from manim_ml.image import GrayscaleImageMobject
 from manim_ml.neural_network.layers.parent_layers import NeuralNetworkLayer
 
+from PIL import Image
+
 class ImageLayer(NeuralNetworkLayer):
     """Single Image Layer for Neural Network"""
 
@@ -17,6 +19,17 @@ class ImageLayer(NeuralNetworkLayer):
             self.image_mobject = ImageMobject(self.numpy_image).scale_to_fit_height(height)
         self.add(self.image_mobject)
 
+    @classmethod
+    def from_path(cls, image_path, grayscale=True, **kwargs):
+        """Creates a query using the paths"""
+        # Load images from path
+        image = Image.open(image_path)
+        numpy_image = np.asarray(image)
+        # Make the layer
+        image_layer = cls(numpy_image, **kwargs)
+
+        return image_layer
+
     @override_animation(Create)
     def _create_override(self, **kwargs):
         debug_mode = False

manim_ml/neural_network/layers/paired_query.py

@@ -6,18 +6,22 @@ import numpy as np
 class PairedQueryLayer(NeuralNetworkLayer):
     """Paired Query Layer"""
 
-    def __init__(self, positive, negative, stroke_width=5, **kwargs):
+    def __init__(self, positive, negative, stroke_width=5, font_size=18, 
+                spacing=0.5, **kwargs):
         super().__init__(**kwargs)
         self.positive = positive
         self.negative = negative
+        self.font_size = font_size
+        self.spacing = spacing
 
         self.stroke_width = stroke_width
         # Make the assets
         self.assets = self.make_assets()
         self.add(self.assets)
+        self.add(self.title)
     
     @classmethod
-    def from_paths(cls, positive_path, negative_path, grayscale=True):
+    def from_paths(cls, positive_path, negative_path, grayscale=True, **kwargs):
         """Creates a query using the paths"""
         # Load images from path
         if grayscale:
@@ -27,7 +31,7 @@ class PairedQueryLayer(NeuralNetworkLayer):
             positive = ImageMobject(positive_path)
             negative = ImageMobject(negative_path)
         # Make the layer
-        query_layer = cls(positive, negative)
+        query_layer = cls(positive, negative, **kwargs)
 
         return query_layer
 
@@ -38,8 +42,9 @@ class PairedQueryLayer(NeuralNetworkLayer):
         # Handle positive
         positive_group = LabeledColorImage(
             self.positive, 
-            color=GREEN,
+            color=BLUE,
             label="Positive",
+            font_size=self.font_size,
             stroke_width=self.stroke_width
         )
         # Handle negative
@@ -47,11 +52,12 @@ class PairedQueryLayer(NeuralNetworkLayer):
             self.negative, 
             color=RED,
             label="Negative", 
+            font_size=self.font_size,
             stroke_width=self.stroke_width
         )
         # Distribute the groups uniformly vertically
         assets = Group(positive_group, negative_group)
-        assets.arrange(DOWN, buff=1.5)
+        assets.arrange(DOWN, buff=self.spacing)
 
         return assets
 

manim_ml/neural_network/layers/parent_layers.py

@@ -8,7 +8,7 @@ class NeuralNetworkLayer(ABC, Group):
         super(Group, self).__init__()
         self.title_text = kwargs["title"] if "title" in kwargs else " "
         self.title = Text(self.title_text, font_size=DEFAULT_FONT_SIZE/3).scale(0.6)
-        # self.title.next_to(self, UP, 1.2)
+        self.title.next_to(self, UP, 1.2)
         # self.add(self.title)
 
     @abstractmethod

manim_ml/neural_network/layers/triplet.py

@@ -7,19 +7,22 @@ class TripletLayer(NeuralNetworkLayer):
     """Shows triplet images"""
 
     def __init__(self, anchor, positive, negative, stroke_width=5,
-                **kwargs):
+                font_size=22, buff=0.2, **kwargs):
         super().__init__(**kwargs)
         self.anchor = anchor
         self.positive = positive
         self.negative = negative
+        self.buff = buff
 
         self.stroke_width = stroke_width
+        self.font_size = font_size
         # Make the assets
         self.assets = self.make_assets()
         self.add(self.assets)
     
     @classmethod
-    def from_paths(cls, anchor_path, positive_path, negative_path, grayscale=True):
+    def from_paths(cls, anchor_path, positive_path, negative_path, grayscale=True,
+                    font_size=22, buff=0.2):
         """Creates a triplet using the anchor paths"""
         # Load images from path
         if grayscale:
@@ -31,7 +34,7 @@ class TripletLayer(NeuralNetworkLayer):
             positive = ImageMobject(positive_path)
             negative = ImageMobject(negative_path)
         # Make the layer
-        triplet_layer = cls(anchor, positive, negative)
+        triplet_layer = cls(anchor, positive, negative, font_size=font_size, buff=buff)
 
         return triplet_layer
 
@@ -44,21 +47,27 @@ class TripletLayer(NeuralNetworkLayer):
             self.anchor, 
             color=WHITE,
             label="Anchor", 
-            stroke_width=self.stroke_width
+            stroke_width=self.stroke_width,
+            font_size=self.font_size,
+            buff=self.buff
         )
         # Handle positive
         positive_group = LabeledColorImage(
             self.positive, 
             color=GREEN,
             label="Positive", 
-            stroke_width=self.stroke_width
+            stroke_width=self.stroke_width,
+            font_size=self.font_size,
+            buff=self.buff
         )
         # Handle negative
         negative_group = LabeledColorImage(
             self.negative, 
             color=RED,
             label="Negative", 
-            stroke_width=self.stroke_width
+            stroke_width=self.stroke_width,
+            font_size=self.font_size,
+            buff=self.buff
         )
         # Distribute the groups uniformly vertically
         assets = Group(anchor_group, positive_group, negative_group)

manim_ml/neural_network/neural_network.py

@@ -13,6 +13,7 @@ from cv2 import AGAST_FEATURE_DETECTOR_NONMAX_SUPPRESSION
 from manim import *
 import warnings
 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
@@ -59,7 +60,11 @@ class NeuralNetwork(Group):
             previous_layer = self.input_layers[layer_index - 1]
             current_layer = self.input_layers[layer_index]
             current_layer.move_to(previous_layer)
-            shift_vector = np.array([(previous_layer.get_width()/2 + current_layer.get_width()/2) + self.layer_spacing, 0, 0])
+            # TODO Temp fix
+            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])
             current_layer.shift(shift_vector)
 
     def _construct_connective_layers(self):

manim_ml/neural_network/neural_network_transformations.py

@@ -0,0 +1,268 @@
+"""
+    Transformations for manipulating a neural network object. 
+"""
+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.
+
+        Note: I needed to do something strange for creating the new connective layer.
+        The issue with creating it intially is that the positions of the sides of the 
+        connective layer depend upon the location of the moved layers **after** the
+        move animations are performed. However, all of these animations are performed
+        after the animations have been created. This means that the animation depends upon
+        the state of the neural network layers after previous animations have been run. 
+        To fix this issue I needed to use an UpdateFromFunc.
+    """
+
+    def __init__(self, layer, neural_network, layer_spacing=0.2):
+        self.layer = layer
+        self.neural_network = neural_network
+        self.layer_spacing = layer_spacing
+        # Get the before and after layers
+        layers_tuple = self.get_connective_layers()
+        self.before_layer = layers_tuple[0]
+        self.after_layer = layers_tuple[1]
+        self.before_connective = layers_tuple[2]
+        self.after_connective = layers_tuple[3]
+        # Make the animations
+        remove_animations = self.make_remove_animation()
+        move_animations = self.make_move_animation()
+        new_connective_animation = self.make_new_connective_animation()
+        # Add all of the animations to the group
+        animations_list = [
+            remove_animations,
+            move_animations,
+            new_connective_animation
+        ]
+
+        super().__init__(*animations_list, lag_ratio=1.0)
+
+    def get_connective_layers(self):
+        """Gets the connective layers before and after self.layer"""
+        # Get layer index
+        layer_index = self.neural_network.all_layers.index_of(self.layer)
+        if layer_index == -1:
+            raise Exception("Layer object not found")
+        # Get the layers before and after
+        before_layer = None
+        after_layer = None
+        before_connective = None
+        after_connective = None
+        if layer_index - 2 >= 0:
+            before_layer = self.neural_network.all_layers[layer_index - 2]
+            before_connective = self.neural_network.all_layers[layer_index - 1]
+        if layer_index + 2 < len(self.neural_network.all_layers): 
+            after_layer = self.neural_network.all_layers[layer_index + 2]
+            after_connective = self.neural_network.all_layers[layer_index + 1]
+        
+        return before_layer, after_layer, before_connective, after_connective
+
+    def make_remove_animation(self):
+        """Removes layer and the surrounding connective layers"""
+        remove_layer_animation = self.make_remove_layer_animation()
+        remove_connective_animation = self.make_remove_connective_layers_animation()
+        # Remove animations
+        remove_animations = AnimationGroup(
+            remove_layer_animation, 
+            remove_connective_animation
+        )
+
+        return remove_animations
+
+    def make_remove_layer_animation(self):
+        """Removes the layer"""
+        # Remove the layer
+        self.neural_network.all_layers.remove(self.layer)
+        # Fade out the removed layer
+        fade_out_removed = FadeOut(self.layer)
+        return fade_out_removed
+
+    def make_remove_connective_layers_animation(self):
+        """Removes the connective layers before and after layer if they exist"""
+        # Fade out the removed connective layers
+        fade_out_before_connective = AnimationGroup()
+        if not self.before_connective is None:
+            self.neural_network.all_layers.remove(self.before_connective)
+            fade_out_before_connective = FadeOut(self.before_connective)
+        fade_out_after_connective = AnimationGroup()
+        if not self.after_connective is None:
+            self.neural_network.all_layers.remove(self.after_connective)
+            fade_out_after_connective = FadeOut(self.after_connective)
+        # Group items
+        remove_connective_group = AnimationGroup(
+            fade_out_after_connective, 
+            fade_out_before_connective
+        )
+
+        return remove_connective_group
+
+    def make_move_animation(self):
+        """Collapses layers"""
+        # Animate the movements
+        move_before_layers = AnimationGroup()
+        shift_right_amount = None
+        if not self.before_layer is None:
+            # Compute shift amount
+            layer_dist = np.abs(self.layer.get_center() - self.before_layer.get_right())[0]
+            shift_right_amount = np.array([layer_dist - self.layer_spacing/2, 0, 0])
+            # Shift all layers before forward
+            before_layer_index = self.neural_network.all_layers.index_of(self.before_layer)
+            layers_before = Group(*self.neural_network.all_layers[:before_layer_index + 1])
+            move_before_layers = layers_before.animate.shift(shift_right_amount)
+        move_after_layers = AnimationGroup()
+        shift_left_amount = None
+        if not self.after_layer is None:
+            layer_dist = np.abs(self.after_layer.get_left() - self.layer.get_center())[0]
+            shift_left_amount = np.array([-layer_dist + self.layer_spacing / 2, 0, 0])
+            # Shift all layers after backward
+            after_layer_index = self.neural_network.all_layers.index_of(self.after_layer)
+            layers_after = Group(*self.neural_network.all_layers[after_layer_index:])
+            move_after_layers = layers_after.animate.shift(shift_left_amount)
+        # Group the move animations
+        move_group = AnimationGroup(
+            move_before_layers, 
+            move_after_layers
+        )
+
+        return move_group
+
+    def make_new_connective_animation(self):
+        """Makes new connective layer"""
+        self.anim_count = 0
+        def create_new_connective(neural_network):
+            """
+                Creates new connective layer
+            
+                This is a closure that creates a new connective layer and animates it.
+            """
+            self.anim_count += 1
+            if self.anim_count == 1:
+                if not self.before_layer is None and not self.after_layer is None:
+                    print(neural_network)
+                    new_connective = get_connective_layer(self.before_layer, self.after_layer)
+                    before_layer_index = neural_network.all_layers.index_of(self.before_layer) + 1
+                    neural_network.all_layers.insert(before_layer_index, new_connective)
+                    print(neural_network)
+
+        update_func_anim = UpdateFromFunc(self.neural_network, create_new_connective)
+
+        return update_func_anim
+
+class InsertLayer(AnimationGroup):
+    """Animation for inserting layer at given index"""
+
+    def __init__(self, layer, index, neural_network):
+        self.layer = layer
+        self.index = index
+        self.neural_network = neural_network
+        # Check valid index
+        assert index < len(self.neural_network.all_layers)
+        # Layers before and after
+        self.layers_before = self.neural_network.all_layers[:self.index]
+        self.layers_after = self.neural_network.all_layers[self.index:]
+        # Get the non-connective layer before and after
+        if len(self.layers_before) > 0:
+            self.layer_before = self.layers_before[-2]
+        if len(self.layers_after) > 0:
+            self.layer_after = self.layers_after[0]
+        # Move layer
+        if not self.layer_after is None:
+            self.layer.move_to(self.layer_after)
+        # Make animations
+        self.old_connective_layer, remove_connective_layer = self.remove_connective_layer_animation()
+        move_layers = self.make_move_layers_animation()
+        create_layer = self.make_create_layer_animation()
+        # create_connective_layers = self.make_create_connective_layers()
+        animations = [
+            remove_connective_layer,
+            move_layers,
+            create_layer, 
+        #    create_connective_layers
+        ]
+
+        super().__init__(*animations, lag_ratio=1.0)
+
+    def get_connective_layer_widths(self):
+        """Gets the widths of the connective layers"""
+        # Make the layers
+        before_connective = None
+        after_connective = None
+        # Get the connective layer objects
+        if len(self.layers_before) > 0:
+            before_connective = get_connective_layer(self.layer_before, self.layer)
+        if len(self.layers_after) > 0:
+            after_connective = get_connective_layer(self.layer, self.layer_after)
+        # Compute the widths
+        before_connective_width = 0
+        if not before_connective is None:
+            before_connective_width = before_connective.width
+        after_connective_width = 0
+        if not after_connective is None:
+            after_connective_width = after_connective.width
+        return before_connective_width, after_connective_width
+
+    def remove_connective_layer_animation(self):
+        """Removes the connective layer before the insertion index"""
+        # Check if connective layer before exists
+        if len(self.layers_before) > 0:
+            removed_connective = self.layers_before[-1]
+            self.layers_before.remove(removed_connective)
+            self.neural_network.all_layers.remove(removed_connective)
+            # Make remove animation
+            remove_animation = FadeOut(removed_connective)
+            return removed_connective, remove_animation
+
+        return None, AnimationGroup()
+
+    def make_move_layers_animation(self):
+        """Shifts layers before and after"""
+        before_connective_width, after_connective_width = self.get_connective_layer_widths()
+        old_connective_width = 0
+        if not self.old_connective_layer is None:
+            old_connective_width = self.old_connective_layer.width
+        # Before layer shift
+        before_shift_animation = AnimationGroup()
+        if len(self.layers_before) > 0:
+            before_shift = np.array([-self.layer.width/2 - before_connective_width + old_connective_width, 0, 0])
+            # Shift layers before
+            before_shift_animation = Group(*self.layers_before).animate.shift(before_shift)
+        # After layer shift
+        after_shift_animation = AnimationGroup()
+        if len(self.layers_after) > 0:
+            after_shift = np.array([self.layer.width/2 + after_connective_width, 0, 0])
+            # Shift layers after
+            after_shift_animation = Group(*self.layers_after).animate.shift(after_shift)
+        # Make animation group
+        shift_animations = AnimationGroup(
+            before_shift_animation, 
+            after_shift_animation
+        )
+
+        return shift_animations
+
+    def make_create_layer_animation(self):
+        """Animates the creation of the layer"""
+        return Create(self.layer)
+
+    def make_create_connective_layers_animation(self, before_connective, after_connective):
+        """Create connective layers"""
+        # Make the layers
+        before_connective = None
+        after_connective = None
+        # Get the connective layer objects
+        if len(self.layers_before) > 0:
+            before_connective = get_connective_layer(self.layers_before[-1], self.layer)
+        if len(self.layers_after) > 0:
+            after_connective = get_connective_layer(self.layers_after[0], self.layer)
+        # Insert the layers
+        # Make the animation
+        animation_group = AnimationGroup(
+            Create(before_connective), 
+            Create(after_connective)
+        )
+
+        return animation_group
+    

manim_ml/neural_network/variational_autoencoder.py

@@ -7,11 +7,8 @@ and Traditional Autoencoders.
 from manim import *
 import numpy as np
 from PIL import Image
-import os
-from manim_ml.neural_network.feed_forward import FeedForwardLayer
-from manim_ml.neural_network.image import ImageLayer
+from manim_ml.neural_network.layers import FeedForwardLayer, EmbeddingLayer, ImageLayer
 from manim_ml.neural_network.neural_network import NeuralNetwork
-from manim_ml.neural_network.embedding import EmbeddingLayer
 
 class VariationalAutoencoder(VGroup):
     """Variational Autoencoder Manim Visualization"""
@@ -66,47 +63,3 @@ class VariationalAutoencoder(VGroup):
         )
 
         return animation_group
-
-        """
-        # Make encoder forward pass
-        encoder_forward_pass = self.encoder.make_forward_propagation_animation(run_time=per_unit_runtime)
-        # Make red dot in embedding
-        mean = [1.0, 1.5]
-        mean_point = self.embedding.axes.coords_to_point(*mean)
-        std = [0.8, 1.2]
-        # Make the dot convergence animation
-        dot_convergence_animation = self.make_dot_convergence_animation(mean, run_time=per_unit_runtime)
-        encoding_succesion = Succession(
-            encoder_forward_pass, 
-            dot_convergence_animation
-        )
-        # Make an ellipse centered at mean_point witAnimationGraph std outline
-        center_dot = Dot(mean_point, radius=self.dot_radius, color=RED)
-        ellipse = Ellipse(width=std[0], height=std[1], color=RED, fill_opacity=0.3, stroke_width=self.ellipse_stroke_width)
-        ellipse.move_to(mean_point)
-        self.distribution_objects = VGroup(
-            center_dot, 
-            ellipse
-        )
-        # Make ellipse animation
-        ellipse_animation = AnimationGroup(
-            GrowFromCenter(center_dot), 
-            GrowFromCenter(ellipse),
-        )
-        # Make the dot divergence animation
-        sampled_point = [0.51, 1.0]
-        divergence_point = self.embedding.axes.coords_to_point(*sampled_point)
-        dot_divergence_animation = self.make_dot_divergence_animation(divergence_point, run_time=per_unit_runtime)
-        # Make decoder foward pass
-        decoder_forward_pass = self.decoder.make_forward_propagation_animation(run_time=per_unit_runtime)
-        # Add the animations to the group
-        animation_group = AnimationGroup(
-            FadeIn(self.input_image),
-            encoding_succesion,
-            ellipse_animation,
-            dot_divergence_animation,
-            decoder_forward_pass,
-            FadeIn(self.output_image),
-            lag_ratio=1,
-        )
-        """

manim_ml/probability.py

@@ -5,43 +5,60 @@ import math
 class GaussianDistribution(VGroup):
     """Object for drawing a Gaussian distribution"""
 
-    def __init__(self, axes, mean=None, cov=None, dist_theme="gaussian", **kwargs):
+    def __init__(self, axes, mean=None, cov=None, dist_theme="gaussian", color=ORANGE, **kwargs):
         super(VGroup, self).__init__(**kwargs)
         self.axes = axes
         self.mean = mean
         self.cov = cov
         self.dist_theme = dist_theme
+        self.color = color
         if mean is None:
             self.mean = np.array([0.0, 0.0])
         if cov is None:
             self.cov = np.array([[1, 0], [0, 1]])
         # Make the Gaussian
         if self.dist_theme is "gaussian":
-            self.ellipses = self.construct_gaussian_distribution(self.mean, self.cov)
+            self.ellipses = self.construct_gaussian_distribution(self.mean, self.cov, color=self.color)
+            self.add(self.ellipses)
         elif self.dist_theme is "ellipse":
-            self.ellipses = self.construct_simple_gaussian_ellipse(self.mean, self.cov)
+            self.ellipses = self.construct_simple_gaussian_ellipse(self.mean, self.cov, color=self.color)
+            self.add(self.ellipses)
         else:
             raise Exception(f"Uncrecognized distribution theme: {self.dist_theme}")
 
+    """  
     @override_animation(Create)
     def _create_gaussian_distribution(self):
-        return Create(self.ellipses)
+        return Create(self)
+    """
 
     def compute_covariance_rotation_and_scale(self, covariance):
-        # Get the eigenvectors and eigenvalues
-        eigenvalues, eigenvectors = np.linalg.eig(covariance)
-        y, x = eigenvectors[0, 1], eigenvectors[0, 0]
-        center_location = np.array([y, x, 0])
-        center_location = self.axes.coords_to_point(*center_location)
-        angle = math.atan(x / y) # x over y to denote the angle between y axis and vector
-        # Calculate the width and height
-        height = np.abs(eigenvalues[0])
-        width = np.abs(eigenvalues[1])
-        shape_coord = np.array([width, height, 0])
-        shape_coord = self.axes.coords_to_point(*shape_coord)
-        width = shape_coord[0]
-        height = shape_coord[1]
 
+        def eigsorted(cov):
+            '''
+            Eigenvalues and eigenvectors of the covariance matrix.
+            '''
+            vals, vecs = np.linalg.eigh(cov)
+            order = vals.argsort()[::-1]
+            return vals[order], vecs[:, order]
+
+        def cov_ellipse(cov, nstd):
+            """
+            Source: http://stackoverflow.com/a/12321306/1391441
+            """
+
+            vals, vecs = eigsorted(cov)
+            theta = np.degrees(np.arctan2(*vecs[:, 0][::-1]))
+
+            # Width and height are "full" widths, not radius
+            width, height = 2 * nstd * np.sqrt(vals)
+
+            return width, height, theta
+
+        width, height, angle = cov_ellipse(covariance, 1)
+        scale_factor = np.abs(self.axes.x_range[0] - self.axes.x_range[1]) / self.axes.x_length
+        width /= scale_factor
+        height /= scale_factor
         return angle, width, height
 
     def construct_gaussian_distribution(self, mean, covariance, color=ORANGE, 
@@ -73,28 +90,22 @@ class GaussianDistribution(VGroup):
 
     def construct_simple_gaussian_ellipse(self, mean, covariance, color=ORANGE):
         """Returns a 2d Gaussian distribution object with given mean and covariance"""
-        # map mean and covariance to frame coordinates
+        # Map mean and covariance to frame coordinates
         mean = self.axes.coords_to_point(*mean)
-        # Figure out the scale and angle of rotation
-        # TODO fix this
-        # rotation, width, height = self.compute_covariance_rotation_and_scale(covariance)
-        mean = np.array([0, 0, 0])
-        mean = self.axes.coords_to_point(*mean)
-        rotation = 0.0
+        angle, width, height = self.compute_covariance_rotation_and_scale(covariance)
         # Make covariance ellipses
-        opacity = 0.0
         ellipses = VGroup()
-        opacity = 0.2
+        opacity = 0.4
         ellipse = Ellipse(
-            width=0.6, 
-            height=0.6, 
+            width=width, 
+            height=height, 
             color=color, 
             fill_opacity=opacity, 
-            stroke_width=2.0
+            stroke_width=1.0
         )
         ellipse.move_to(mean)
-        ellipse.rotate(rotation)
+        ellipse.rotate(angle)
         ellipses.add(ellipse)
+        ellipses.set_z_index(3)
 
         return ellipses
-

tests/test_embedding_layer.py

@@ -26,6 +26,52 @@ class EmbeddingNNScene(Scene):
 
         self.play(neural_network.make_forward_pass_animation(run_time=5))
 
+
+class TripletEmbeddingNNScene(Scene):
+
+    def construct(self):
+        embedding_layer = EmbeddingLayer()
+
+        neural_network = NeuralNetwork([
+            FeedForwardLayer(5),
+            FeedForwardLayer(3),
+            embedding_layer,
+            FeedForwardLayer(3),
+            FeedForwardLayer(5)
+        ])
+
+        self.play(Create(neural_network))
+
+        self.play(
+            neural_network.make_forward_pass_animation(
+                layer_args={
+                    embedding_layer: {
+                        "triplet_args": {
+                            "anchor_dist": {
+                                "cov": np.array([[0.3, 0], [0, 0.3]]),
+                                "mean": np.array([0.7, 1.4]),
+                                "dist_theme": "ellipse",
+                                "color": BLUE
+                            },
+                            "positive_dist": {
+                                "cov": np.array([[0.2, 0], [0, 0.2]]),
+                                "mean": np.array([0.8, -0.4]),
+                                "dist_theme": "ellipse",
+                                "color": GREEN
+                            },
+                            "negative_dist": {
+                                "cov": np.array([[0.4, 0], [0, 0.25]]),
+                                "mean": np.array([-1, -1.2]),
+                                "dist_theme": "ellipse",
+                                "color": RED
+                            }
+                        }
+                    }
+                },
+                run_time=5
+            )
+        )
+
 class QueryEmbeddingNNScene(Scene):
 
     def construct(self):
@@ -47,7 +93,18 @@ class QueryEmbeddingNNScene(Scene):
                 run_time=5,
                 layer_args={
                     embedding_layer: {
-                        "query_locations": (np.array([2, 2]), np.array([1, 1]))
+                        "positive_dist_args": {
+                            "cov": np.array([[1, 0], [0, 1]]),
+                            "mean": np.array([1, 1]),
+                            "dist_theme": "ellipse",
+                            "color": GREEN
+                        },
+                        "negative_dist_args": {
+                            "cov": np.array([[1, 0], [0, 1]]),
+                            "mean": np.array([-1, -1]),
+                            "dist_theme": "ellipse",
+                            "color": RED
+                        }
                     }
                 } 
             )

tests/test_neural_network.py

@@ -94,7 +94,7 @@ class NeuralNetworkScene(Scene):
     def construct(self):
         # Make the Layer object
         layers = [
-            FeedForwardLayer(3, title="Title Test"), 
+            FeedForwardLayer(3), 
             FeedForwardLayer(5), 
             FeedForwardLayer(3)
         ]
@@ -151,23 +151,6 @@ class ImageNeuralNetworkScene(Scene):
         self.play(nn.make_forward_pass_animation(run_time=5))
         self.play(nn.make_forward_pass_animation(run_time=5))
 
-class EmbeddingNNScene(Scene):
-
-    def construct(self):
-        embedding_layer = EmbeddingLayer()
-
-        neural_network = NeuralNetwork([
-            FeedForwardLayer(5),
-            FeedForwardLayer(3),
-            embedding_layer,
-            FeedForwardLayer(3),
-            FeedForwardLayer(5)
-        ])
-
-        self.play(Create(neural_network))
-
-        self.play(neural_network.make_forward_pass_animation(run_time=5))
-
 class RecursiveNNScene(Scene):
 
     def construct(self):