Update Surface to have a data array

This commit is contained in:
Grant Sanderson
2023-01-15 17:40:46 -08:00
parent ca5e119425
commit d267c00761

View File

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