diff --git a/manimlib/mobject/types/surface.py b/manimlib/mobject/types/surface.py index 3e726f1a..92cb45f4 100644 --- a/manimlib/mobject/types/surface.py +++ b/manimlib/mobject/types/surface.py @@ -5,6 +5,7 @@ import numpy as np from manimlib.constants import GREY from manimlib.constants import OUT +from manimlib.constants import ORIGIN from manimlib.mobject.mobject import Mobject from manimlib.utils.bezier import integer_interpolate from manimlib.utils.bezier import interpolate @@ -31,6 +32,13 @@ class Surface(Mobject): ('dv_point', np.float32, (3,)), ('color', np.float32, (4,)), ] + data_dtype: Sequence[Tuple[str, type, Tuple[int]]] = [ + ('points', np.float32, (3,)), + ('du_point', np.float32, (3,)), + ('dv_point', np.float32, (3,)), + ('rgba', np.float32, (4,)), + ] + pointlike_data_keys = ['points', 'du_point', 'dv_point'] def __init__( self, @@ -85,16 +93,21 @@ class Surface(Mobject): # - Points generated by pure uv values # - Those generated by values nudged by du # - Those generated by values nudged by dv - point_lists = [] - for (du, dv) in [(0, 0), (self.epsilon, 0), (0, self.epsilon)]: - uv_grid = np.array([[[u + du, v + dv] for v in v_range] for u in u_range]) - point_grid = np.apply_along_axis(lambda p: self.uv_func(*p), 2, uv_grid) - point_lists.append(point_grid.reshape((nu * nv, dim))) - # Rather than tracking normal vectors, the points list will hold on to the - # infinitesimal nudged values alongside the original values. This way, one - # can perform all the manipulations they'd like to the surface, and normals - # are still easily recoverable. - self.set_points(np.vstack(point_lists)) + uv_grid = np.array([[[u, v] for v in v_range] for u in u_range]) + uv_plus_du = uv_grid.copy() + uv_plus_du[:, :, 0] += self.epsilon + uv_plus_dv = uv_grid.copy() + uv_plus_dv[:, :, 1] += self.epsilon + + points, du_points, dv_points = [ + np.apply_along_axis( + lambda p: self.uv_func(*p), 2, grid + ).reshape((nu * nv, dim)) + for grid in (uv_grid, uv_plus_du, uv_plus_dv) + ] + self.set_points(points) + self.data["du_point"][:] = du_points + self.data["dv_point"][:] = dv_points def compute_triangle_indices(self): # TODO, if there is an event which changes @@ -117,12 +130,8 @@ class Surface(Mobject): def get_triangle_indices(self) -> np.ndarray: return self.triangle_indices - def get_surface_points_and_nudged_points( - self - ) -> tuple[Vect3Array, Vect3Array, Vect3Array]: - points = self.get_points() - k = len(points) // 3 - return points[:k], points[k:2 * k], points[2 * k:] + def get_surface_points_and_nudged_points(self) -> tuple[Vect3Array, Vect3Array, Vect3Array]: + return (self.data['points'], self.data['du_point'], self.data['dv_point']) def get_unit_normals(self) -> Vect3Array: s_points, du_points, dv_points = self.get_surface_points_and_nudged_points() @@ -147,10 +156,12 @@ class Surface(Mobject): return self nu, nv = smobject.resolution - self.set_points(np.vstack([ - self.get_partial_points_array(arr.copy(), a, b, (nu, nv, 3), axis=axis) - for arr in smobject.get_surface_points_and_nudged_points() - ])) + for key in ['points', 'du_point', 'dv_point']: + self.data[key][:] = self.get_partial_points_array( + self.data[key], a, b, + (nu, nv, 3), + axis=axis + ) return self def get_partial_points_array( @@ -238,7 +249,7 @@ class Surface(Mobject): return shader_data def fill_in_shader_color_info(self, shader_data: np.ndarray) -> np.ndarray: - self.read_data_to_shader(shader_data, "color", "rgbas") + self.read_data_to_shader(shader_data, "color", "rgba") return shader_data def get_shader_vert_indices(self) -> np.ndarray: