Added DotCloud type

This commit is contained in:
Grant Sanderson 2021-01-04 13:26:58 -08:00
parent 6c0f00b4a1
commit a4ea4791e7
5 changed files with 209 additions and 40 deletions

View file

@ -0,0 +1,90 @@
import numpy as np
import moderngl
import numbers
from manimlib.constants import GREY_C
from manimlib.constants import ORIGIN
from manimlib.mobject.types.point_cloud_mobject import PMobject
from manimlib.mobject.geometry import DEFAULT_DOT_RADIUS
from manimlib.utils.bezier import interpolate
from manimlib.utils.color import color_to_rgba
from manimlib.utils.iterables import stretch_array_to_length
class DotCloud(PMobject):
CONFIG = {
"radii": DEFAULT_DOT_RADIUS,
"color": GREY_C,
"opacity": 1,
"vert_shader_file": "true_dot_vert.glsl",
"geom_shader_file": "true_dot_geom.glsl",
"frag_shader_file": "true_dot_frag.glsl",
"render_primitive": moderngl.POINTS,
"shader_dtype": [
('point', np.float32, (3,)),
('radius', np.float32, (1,)),
('color', np.float32, (4,)),
],
}
def __init__(self, points=[[ORIGIN]], radii=DEFAULT_DOT_RADIUS, **kwargs):
self.radii = np.repeat([radii], len(points))
self.points = np.array(points)
super().__init__(**kwargs)
def init_points(self):
pass
def init_colors(self):
self.rgbas = np.zeros((len(self.points), 4))
self.rgbas[:] = color_to_rgba(self.color, self.opacity)
return self
def set_points(self, points):
super().set_points(points)
self.radii = stretch_array_to_length(self.radii, len(points))
return self
def set_points_by_grid(self, n_rows, n_cols, height=None, width=None):
new_points = np.zeros((n_rows * n_cols, 3))
new_points[:, 0] = np.tile(range(n_cols), n_rows)
new_points[:, 1] = np.repeat(range(n_rows), n_cols)
new_points[:, 2] = 0
self.set_points(new_points)
radius = self.radii[0]
if height is None:
height = n_rows * 3 * radius
if width is None:
width = n_cols * 3 * radius
self.set_height(height, stretch=True)
self.set_width(width, stretch=True)
self.center()
return self
def set_radii(self, radii):
if isinstance(radii, numbers.Number):
self.radii[:] = radii
else:
self.radii = stretch_array_to_length(radii, len(self.points))
return self
def scale(self, scale_factor, scale_radii=True, **kwargs):
super().scale(scale_factor, **kwargs)
if scale_radii:
self.radii *= scale_factor
return self
def interpolate(self, mobject1, mobject2, alpha, *args, **kwargs):
super().interpolate(mobject1, mobject2, alpha, *args, **kwargs)
self.radii = interpolate(mobject1.radii, mobject2.radii, alpha)
return self
def get_shader_data(self):
data = self.get_blank_shader_data_array(len(self.points))
data["point"] = self.points
data["radius"] = self.radii.reshape((len(self.radii), 1))
data["color"] = self.rgbas
return data

View file

@ -8,10 +8,6 @@ from manimlib.utils.iterables import stretch_array_to_length
class PMobject(Mobject):
CONFIG = {
"stroke_width": DEFAULT_STROKE_WIDTH,
}
def reset_points(self):
self.rgbas = np.zeros((0, 4))
self.points = np.zeros((0, 3))
@ -20,6 +16,11 @@ class PMobject(Mobject):
def get_array_attrs(self):
return Mobject.get_array_attrs(self) + ["rgbas"]
def set_points(self, points):
self.points = points
self.rgbas = stretch_array_to_length(self.rgbas, len(points))
return self
def add_points(self, points, rgbas=None, color=None, alpha=1):
"""
points must be a Nx3 numpy array, as must rgbas if it is not None
@ -40,23 +41,25 @@ class PMobject(Mobject):
self.rgbas = np.vstack([self.rgbas, rgbas])
return self
def set_color(self, color=YELLOW_C, family=True):
def set_color(self, color, family=True):
rgba = color_to_rgba(color)
mobs = self.family_members_with_points() if family else [self]
for mob in mobs:
mob.rgbas[:, :] = rgba
self.color = color
return self
def get_stroke_width(self):
return self.stroke_width
def set_stroke_width(self, width, family=True):
def set_opacity(self, opacity, family=True):
mobs = self.family_members_with_points() if family else [self]
for mob in mobs:
mob.stroke_width = width
mob.rgbas[:, 3] = opacity
return self
def get_color(self):
return rgba_to_color(self.rgbas[0, :])
def get_all_rgbas(self):
return self.get_merged_array("rgbas")
# def set_color_by_gradient(self, start_color, end_color):
def set_color_by_gradient(self, *colors):
self.rgbas = np.array(list(map(
@ -87,9 +90,8 @@ class PMobject(Mobject):
)
return self
def match_colors(self, mobject):
Mobject.align_data(self, mobject)
self.rgbas = np.array(mobject.rgbas)
def match_colors(self, pmobject):
self.rgbas[:] = stretch_array_to_length(pmobject.rgbas, len(self.points))
return self
def filter_out(self, condition):
@ -123,15 +125,6 @@ class PMobject(Mobject):
mob.apply_over_attr_arrays(lambda arr: arr[indices])
return self
def fade_to(self, color, alpha):
self.rgbas = interpolate(self.rgbas, color_to_rgba(color), alpha)
for mob in self.submobjects:
mob.fade_to(color, alpha)
return self
def get_all_rgbas(self):
return self.get_merged_array("rgbas")
def ingest_submobjects(self):
attrs = self.get_array_attrs()
arrays = list(map(self.get_merged_array, attrs))
@ -140,9 +133,6 @@ class PMobject(Mobject):
self.set_submobjects([])
return self
def get_color(self):
return rgba_to_color(self.rgbas[0, :])
def point_from_proportion(self, alpha):
index = alpha * (self.get_num_points() - 1)
return self.points[index]
@ -157,23 +147,14 @@ class PMobject(Mobject):
)
def interpolate_color(self, mobject1, mobject2, alpha):
self.rgbas = interpolate(
mobject1.rgbas, mobject2.rgbas, alpha
)
self.set_stroke_width(interpolate(
mobject1.get_stroke_width(),
mobject2.get_stroke_width(),
alpha,
))
self.rgbas = interpolate(mobject1.rgbas, mobject2.rgbas, alpha)
return self
def pointwise_become_partial(self, mobject, a, b):
lower_index, upper_index = [
int(x * mobject.get_num_points())
for x in (a, b)
]
def pointwise_become_partial(self, pmobject, a, b):
lower_index = int(a * pmobject.get_num_points())
upper_index = int(b * pmobject.get_num_points())
for attr in self.get_array_attrs():
full_array = getattr(mobject, attr)
full_array = getattr(pmobject, attr)
partial_array = full_array[lower_index:upper_index]
setattr(self, attr, partial_array)

View file

@ -0,0 +1,31 @@
#version 330
uniform vec3 light_source_position;
uniform float gloss;
uniform float shadow;
uniform float anti_alias_width;
in vec4 color;
in float radius;
in vec2 center;
in vec2 point;
out vec4 frag_color;
#INSERT add_light.glsl
void main() {
vec2 diff = point - center;
float dist = length(diff);
vec3 normal = vec3(diff / radius, sqrt(1 - (dist * dist) / (radius * radius)));
float signed_dist = dist - radius;
frag_color = add_light(
color,
vec3(point.xy, 0.0),
normal,
light_source_position,
gloss,
shadow
);
frag_color.a *= smoothstep(0.5, -0.5, signed_dist / anti_alias_width);
}

View file

@ -0,0 +1,44 @@
#version 330
layout (points) in;
layout (triangle_strip, max_vertices = 4) out;
// Needed for get_gl_Position
uniform vec2 frame_shape;
uniform float focal_distance;
uniform float is_fixed_in_frame;
uniform float anti_alias_width;
in vec3 v_point[1];
in float v_radius[1];
in vec4 v_color[1];
out vec4 color;
out float radius;
out vec2 center;
out vec2 point;
// Imports
#INSERT get_gl_Position.glsl
void main() {
color = v_color[0];
radius = v_radius[0];
center = v_point[0].xy;
radius = v_radius[0] / max(1.0 - v_point[0].z / focal_distance / frame_shape.y, 0.0);
float rpa = radius + anti_alias_width;
for(int i = 0; i < 4; i++){
// To account for perspective
int x_index = 2 * (i % 2) - 1;
int y_index = 2 * (i / 2) - 1;
vec3 corner = v_point[0] + vec3(x_index * rpa, y_index * rpa, 0.0);
gl_Position = get_gl_Position(corner);
point = corner.xy;
EmitVertex();
}
EndPrimitive();
}

View file

@ -0,0 +1,23 @@
#version 330
uniform vec2 frame_shape;
uniform float anti_alias_width;
uniform mat4 to_screen_space;
uniform float is_fixed_in_frame;
uniform float focal_distance;
in vec3 point;
in float radius;
in vec4 color;
out vec3 v_point;
out float v_radius;
out vec4 v_color;
#INSERT position_point_into_frame.glsl
void main(){
v_point = position_point_into_frame(point);
v_radius = radius;
v_color = color;
}