From a4ea4791e7c895400bb976200f401b972e0dbffb Mon Sep 17 00:00:00 2001 From: Grant Sanderson Date: Mon, 4 Jan 2021 13:26:58 -0800 Subject: [PATCH] Added DotCloud type --- manimlib/mobject/types/dot_cloud.py | 90 +++++++++++++++++++ manimlib/mobject/types/point_cloud_mobject.py | 61 +++++-------- manimlib/shaders/true_dot_frag.glsl | 31 +++++++ manimlib/shaders/true_dot_geom.glsl | 44 +++++++++ manimlib/shaders/true_dot_vert.glsl | 23 +++++ 5 files changed, 209 insertions(+), 40 deletions(-) create mode 100644 manimlib/mobject/types/dot_cloud.py create mode 100644 manimlib/shaders/true_dot_frag.glsl create mode 100644 manimlib/shaders/true_dot_geom.glsl create mode 100644 manimlib/shaders/true_dot_vert.glsl diff --git a/manimlib/mobject/types/dot_cloud.py b/manimlib/mobject/types/dot_cloud.py new file mode 100644 index 00000000..b39033be --- /dev/null +++ b/manimlib/mobject/types/dot_cloud.py @@ -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 diff --git a/manimlib/mobject/types/point_cloud_mobject.py b/manimlib/mobject/types/point_cloud_mobject.py index 7e607274..d4a12997 100644 --- a/manimlib/mobject/types/point_cloud_mobject.py +++ b/manimlib/mobject/types/point_cloud_mobject.py @@ -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) diff --git a/manimlib/shaders/true_dot_frag.glsl b/manimlib/shaders/true_dot_frag.glsl new file mode 100644 index 00000000..7a78477d --- /dev/null +++ b/manimlib/shaders/true_dot_frag.glsl @@ -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); +} \ No newline at end of file diff --git a/manimlib/shaders/true_dot_geom.glsl b/manimlib/shaders/true_dot_geom.glsl new file mode 100644 index 00000000..3562ae2a --- /dev/null +++ b/manimlib/shaders/true_dot_geom.glsl @@ -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(); +} \ No newline at end of file diff --git a/manimlib/shaders/true_dot_vert.glsl b/manimlib/shaders/true_dot_vert.glsl new file mode 100644 index 00000000..447e93eb --- /dev/null +++ b/manimlib/shaders/true_dot_vert.glsl @@ -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; +} \ No newline at end of file