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

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
-
+        return ellipses