opensource/ManimML
1
0
Fork 0
mirror of https://github.com/helblazer811/ManimML.git synced 2025-05-19 19:46:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

Used Black to reformat the code in the repository.

Browse Source
This commit is contained in:
Alec Helbling
2023-01-01 23:24:59 -05:00
parent 334662e8c8
commit 3d6e8072e1
71 changed files with 1701 additions and 1135 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
manim_ml/neural_network/layers/embedding.py
View File

@ -2,12 +2,19 @@ from manim import *
from manim_ml.probability import GaussianDistribution
from manim_ml.neural_network.layers.parent_layers import VGroupNeuralNetworkLayer
class EmbeddingLayer(VGroupNeuralNetworkLayer):
"""NeuralNetwork embedding object that can show probability distributions"""
def __init__(self, point_radius=0.02, mean = np.array([0, 0]),
covariance=np.array([[1.0, 0], [0, 1.0]]), dist_theme="gaussian",
paired_query_mode=False, **kwargs):
def __init__(
self,
point_radius=0.02,
mean=np.array([0, 0]),
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)
@ -20,14 +27,8 @@ class EmbeddingLayer(VGroupNeuralNetworkLayer):
y_length=0.8,
x_range=(-1.4, 1.4),
y_range=(-1.8, 1.8),
x_axis_config={
"include_ticks": False,
"stroke_width": 0.0
},
y_axis_config={
"include_ticks": False,
"stroke_width": 0.0
}
x_axis_config={"include_ticks": False, "stroke_width": 0.0},
y_axis_config={"include_ticks": False, "stroke_width": 0.0},
)
self.add(self.axes)
self.axes.move_to(self.get_center())
@ -35,7 +36,9 @@ 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) # 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"""
@ -57,15 +60,16 @@ class EmbeddingLayer(VGroupNeuralNetworkLayer):
point = np.random.multivariate_normal(mean, cov)
# Make dot at correct location
location = self.axes.coords_to_point(point[0], point[1])
return location
def get_distribution_location(self):
"""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=WHITE,
num_points=400):
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)
@ -74,7 +78,7 @@ class EmbeddingLayer(VGroupNeuralNetworkLayer):
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 = Dot(point_location, color=point_color, radius=self.point_radius / 2)
dot.set_z_index(-1)
point_dots.add(dot)
@ -82,36 +86,27 @@ class EmbeddingLayer(VGroupNeuralNetworkLayer):
def make_forward_pass_animation(self, layer_args={}, **kwargs):
"""Forward pass animation"""
animations = []
animations = []
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,
**anchor_dist_args
)
anchor_dist = GaussianDistribution(self.axes, **anchor_dist_args)
animations.append(Create(anchor_dist))
positive_dist = GaussianDistribution(
self.axes,
**positive_dist_args
)
positive_dist = GaussianDistribution(self.axes, **positive_dist_args)
animations.append(Create(positive_dist))
negative_dist = GaussianDistribution(
self.axes,
**negative_dist_args
)
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
stroke_width=DEFAULT_STROKE_WIDTH / 2,
)
anchor_positive.set_z_index(3)
animations.append(Create(anchor_positive))
@ -120,7 +115,7 @@ class EmbeddingLayer(VGroupNeuralNetworkLayer):
anchor_dist.get_center(),
negative_dist.get_center(),
color=GOLD,
stroke_width=DEFAULT_STROKE_WIDTH/2
stroke_width=DEFAULT_STROKE_WIDTH / 2,
)
anchor_negative.set_z_index(3)
@ -132,14 +127,13 @@ class EmbeddingLayer(VGroupNeuralNetworkLayer):
if "scale_factor" in layer_args:
scale_factor = layer_args["scale_factor"]
self.latent_distribution = GaussianDistribution(
self.axes,
**layer_args["dist_args"]
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,
# self.axes,
# dist_theme=self.dist_theme,
# cov=np.array([[0.8, 0], [0.0, 0.8]])
# )
pass
@ -153,15 +147,9 @@ class EmbeddingLayer(VGroupNeuralNetworkLayer):
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
positive_dist = GaussianDistribution(
self.axes,
**positive_dist_args
)
positive_dist = GaussianDistribution(self.axes, **positive_dist_args)
self.gaussian_distributions.add(positive_dist)
negative_dist = GaussianDistribution(
self.axes,
**negative_dist_args
)
negative_dist = GaussianDistribution(self.axes, **negative_dist_args)
self.gaussian_distributions.add(negative_dist)
animations.append(Create(positive_dist))
@ -182,8 +170,8 @@ class EmbeddingLayer(VGroupNeuralNetworkLayer):
return point_animation
class NeuralNetworkEmbeddingTestScene(Scene):
class NeuralNetworkEmbeddingTestScene(Scene):
def construct(self):
nne = EmbeddingLayer()
mean = np.array([0, 0])
@ -195,4 +183,4 @@ class NeuralNetworkEmbeddingTestScene(Scene):
gaussian = nne.construct_gaussian_distribution(mean, cov)
nne.add(gaussian)
self.add(nne)
self.add(nne)