mirror of
https://github.com/helblazer811/ManimML.git
synced 2025-05-25 17:24:59 +08:00
121 lines
4.0 KiB
Python
121 lines
4.0 KiB
Python
from manim import *
|
|
import numpy as np
|
|
import math
|
|
|
|
|
|
class GaussianDistribution(VGroup):
|
|
"""Object for drawing a Gaussian distribution"""
|
|
|
|
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, color=self.color
|
|
)
|
|
self.add(self.ellipses)
|
|
elif self.dist_theme is "ellipse":
|
|
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)
|
|
"""
|
|
|
|
def compute_covariance_rotation_and_scale(self, covariance):
|
|
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, num_ellipses=4
|
|
):
|
|
"""Returns a 2d Gaussian distribution object with given mean and covariance"""
|
|
# map mean and covariance to frame coordinates
|
|
mean = self.axes.coords_to_point(*mean)
|
|
# Figure out the scale and angle of rotation
|
|
rotation, width, height = self.compute_covariance_rotation_and_scale(covariance)
|
|
# Make covariance ellipses
|
|
opacity = 0.0
|
|
ellipses = VGroup()
|
|
for ellipse_number in range(num_ellipses):
|
|
opacity += 1.0 / num_ellipses
|
|
ellipse_width = width * (1 - opacity)
|
|
ellipse_height = height * (1 - opacity)
|
|
ellipse = Ellipse(
|
|
width=ellipse_width,
|
|
height=ellipse_height,
|
|
color=color,
|
|
fill_opacity=opacity,
|
|
stroke_width=2.0,
|
|
)
|
|
ellipse.move_to(mean)
|
|
ellipse.rotate(rotation)
|
|
ellipses.add(ellipse)
|
|
|
|
return ellipses
|
|
|
|
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
|
|
mean = self.axes.coords_to_point(*mean)
|
|
angle, width, height = self.compute_covariance_rotation_and_scale(covariance)
|
|
# Make covariance ellipses
|
|
ellipses = VGroup()
|
|
opacity = 0.4
|
|
ellipse = Ellipse(
|
|
width=width,
|
|
height=height,
|
|
color=color,
|
|
fill_opacity=opacity,
|
|
stroke_width=1.0,
|
|
)
|
|
ellipse.move_to(mean)
|
|
ellipse.rotate(angle)
|
|
ellipses.add(ellipse)
|
|
ellipses.set_z_index(3)
|
|
|
|
return ellipses
|