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 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: