mirror of
https://github.com/helblazer811/ManimML.git
synced 2025-05-29 04:02:40 +08:00
66 lines
2.4 KiB
Python
66 lines
2.4 KiB
Python
from manim import *
|
|
|
|
from manim_ml.neural_network import NeuralNetwork, NeuralNetworkLayer
|
|
import numpy as np
|
|
import math
|
|
|
|
class NeuralNetworkEmbedding(NeuralNetworkLayer, Axes):
|
|
"""NeuralNetwork embedding object that can show probability distributions"""
|
|
|
|
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]
|
|
print(eigenvectors[0])
|
|
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])
|
|
return angle, width, height
|
|
|
|
def construct_gaussian_point_cloud(self, mean, covariance, color=BLUE):
|
|
"""Plots points sampled from a Gaussian with the given mean and covariance"""
|
|
pass
|
|
|
|
def construct_gaussian_distribution(self, mean, covariance, color=ORANGE,
|
|
dot_radius=0.05, num_ellipses=4):
|
|
"""Returns a 2d Gaussian distribution object with given mean and covariance"""
|
|
# map mean and covariance to frame coordinates
|
|
mean = self.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=0.0
|
|
)
|
|
ellipse.move_to(mean)
|
|
ellipse.rotate(rotation)
|
|
ellipses.add(ellipse)
|
|
|
|
return ellipses
|
|
|
|
def make_forward_pass_animation(self):
|
|
|
|
pass
|
|
|
|
class NeuralNetworkEmbeddingTestScene(Scene):
|
|
|
|
def construct(self):
|
|
nne = NeuralNetworkEmbedding()
|
|
mean = np.array([0, 0])
|
|
cov = np.array([[0.1, 0.8], [0.0, 0.8]])
|
|
|
|
gaussian = nne.construct_gaussian_distribution(mean, cov)
|
|
gaussian.scale(3)
|
|
|
|
self.add(gaussian) |