diff --git a/manimlib/mobject/types/surface.py b/manimlib/mobject/types/surface.py index 9f4f3afe..e59c7987 100644 --- a/manimlib/mobject/types/surface.py +++ b/manimlib/mobject/types/surface.py @@ -14,7 +14,9 @@ class ParametricSurface(Mobject): CONFIG = { "u_range": (0, 1), "v_range": (0, 1), - "resolution": (100, 100), + # Resolution counts number of points sampled, which is off by + # 1 from the number of actual approximating squares seen + "resolution": (101, 101), "color": GREY, "opacity": 1.0, "gloss": 0.3, @@ -33,20 +35,15 @@ class ParametricSurface(Mobject): ] } - def __init__(self, function=None, **kwargs): - self.passed_function = function + def __init__(self, uv_func, **kwargs): + self.uv_func = uv_func super().__init__(**kwargs) - def uv_func(self, u, v): - # Typically to be implemented by a subclass - if self.passed_function is not None: - return self.passed_function(u, v) - return [u, v, 0] - def init_points(self): + dim = self.dim nu, nv = self.resolution - u_range = np.linspace(*self.u_range, nu + 1) - v_range = np.linspace(*self.v_range, nv + 1) + u_range = np.linspace(*self.u_range, nu) + v_range = np.linspace(*self.v_range, nv) # Get three lists: # - Points generated by pure uv values @@ -56,28 +53,33 @@ class ParametricSurface(Mobject): 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(self.get_triangle_ready_array_from_grid(point_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.points = np.vstack(point_lists) - def get_triangle_ready_array_from_grid(self, grid): - # Given a grid, say of points or normals, this returns an Nx3 array - # whose rows are elements from this grid in such such a way that successive - # triplets of points form triangles covering the grid. - nu, nv, dim = grid.shape - nu -= 1 - nv -= 1 - arr = np.zeros((nu * nv * 6, dim)) + def get_triangle_ready_array(self, array): + # Given an array of points representing a flattened grid, say of points or + # normals, this returns an Nx3 array whose rows are elements from this grid + # in such such a way that successive triplets of points form triangles covering + # the grid. + dim = array.shape[1] + nu, nv = self.resolution + assert(array.shape == (nu * nv, dim)) + grid = array.reshape((nu, nv, dim)) + + new_length = (nu - 1) * (nv - 1) + arr = np.zeros((6 * new_length, dim)) + shape = (new_length, dim) # To match the triangles covering this surface - arr[0::6] = grid[:-1, :-1].reshape((nu * nv, dim)) # Top left - arr[1::6] = grid[+1:, :-1].reshape((nu * nv, dim)) # Bottom left - arr[2::6] = grid[:-1, +1:].reshape((nu * nv, dim)) # Top right - arr[3::6] = grid[:-1, +1:].reshape((nu * nv, dim)) # Top right - arr[4::6] = grid[+1:, :-1].reshape((nu * nv, dim)) # Bottom left - arr[5::6] = grid[+1:, +1:].reshape((nu * nv, dim)) # Bottom right + arr[0::6] = grid[:-1, :-1].reshape(shape) # Top left + arr[1::6] = grid[+1:, :-1].reshape(shape) # Bottom left + arr[2::6] = grid[:-1, +1:].reshape(shape) # Top right + arr[3::6] = grid[:-1, +1:].reshape(shape) # Top right + arr[4::6] = grid[+1:, :-1].reshape(shape) # Bottom left + arr[5::6] = grid[+1:, +1:].reshape(shape) # Bottom right return arr def init_colors(self): @@ -120,7 +122,10 @@ class ParametricSurface(Mobject): return self def get_shader_data(self): - s_points, du_points, dv_points = self.get_surface_points_and_nudged_points() + s_points, du_points, dv_points = [ + self.get_triangle_ready_array(array) + for array in self.get_surface_points_and_nudged_points() + ] data = self.get_blank_shader_data_array(len(s_points)) data["point"] = s_points data["du_point"] = du_points @@ -163,17 +168,22 @@ class TexturedSurface(ParametricSurface): path = get_full_raster_image_path(filename) self.image = Image.open(path) self.texture_path = path + self.uv_surface = uv_surface - super().__init__(**kwargs) + self.uv_func = uv_surface.uv_func + self.u_range = uv_surface.u_range + self.v_range = uv_surface.v_range + self.resolution = uv_surface.resolution + super().__init__(self.uv_func, **kwargs) def init_points(self): self.points = self.uv_surface.points # Init im_coords nu, nv = self.uv_surface.resolution - u_range = np.linspace(0, 1, nu + 1) - v_range = np.linspace(1, 0, nv + 1) # Reverse y-direction - uv_grid = np.array([[[u, v] for v in v_range] for u in u_range]) - self.im_coords = self.uv_surface.get_triangle_ready_array_from_grid(uv_grid) + u_range = np.linspace(0, 1, nu) + v_range = np.linspace(1, 0, nv) # Reverse y-direction + uv_grid = np.array([[u, v] for u in u_range for v in v_range]) + self.im_coords = uv_grid def init_colors(self): self.opacity = self.uv_surface.rgbas[:, 3] @@ -187,7 +197,7 @@ class TexturedSurface(ParametricSurface): return self def fill_in_shader_color_info(self, data): - data["im_coords"] = self.im_coords + data["im_coords"] = self.get_triangle_ready_array(self.im_coords) data["opacity"] = self.opacity data["gloss"] = self.gloss return data