From 55da5d5d03be38dbf4322d91d1a11f299abbc657 Mon Sep 17 00:00:00 2001
From: Grant Sanderson <grant@3blue1brown.com>
Date: Mon, 30 Jan 2023 20:49:32 -0800
Subject: [PATCH] Remove use of dv_points and du_points, pass unit normals to
 shader instead

---
 manimlib/mobject/types/surface.py  | 59 ++++++++++++++++++------------
 manimlib/shaders/surface/vert.glsl |  6 +--
 2 files changed, 38 insertions(+), 27 deletions(-)

diff --git a/manimlib/mobject/types/surface.py b/manimlib/mobject/types/surface.py
index faaa764d..b79e95da 100644
--- a/manimlib/mobject/types/surface.py
+++ b/manimlib/mobject/types/surface.py
@@ -12,6 +12,7 @@ from manimlib.utils.images import get_full_raster_image_path
 from manimlib.utils.iterables import listify
 from manimlib.utils.iterables import resize_with_interpolation
 from manimlib.utils.space_ops import normalize_along_axis
+from manimlib.utils.space_ops import cross
 
 from typing import TYPE_CHECKING
 
@@ -27,11 +28,9 @@ class Surface(Mobject):
     shader_folder: str = "surface"
     shader_dtype: np.dtype = np.dtype([
         ('point', np.float32, (3,)),
-        ('du_point', np.float32, (3,)),
-        ('dv_point', np.float32, (3,)),
+        ('normal', np.float32, (3,)),
         ('rgba', np.float32, (4,)),
     ])
-    pointlike_data_keys = ['point', 'du_point', 'dv_point']
 
     def __init__(
         self,
@@ -96,8 +95,14 @@ class Surface(Mobject):
             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
+        self.data["normal"] = normalize_along_axis(cross(
+            du_points - points,
+            dv_points - points,
+        ), 1)
+
+    def apply_points_function(self, *args, **kwargs):
+        super().apply_points_function(*args, **kwargs)
+        self.get_unit_normals()
 
     def compute_triangle_indices(self):
         # TODO, if there is an event which changes
@@ -120,16 +125,27 @@ 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]:
-        return (self.data['point'], 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()
-        normals = np.cross(
-            (du_points - s_points) / self.epsilon,
-            (dv_points - s_points) / self.epsilon,
+        nu, nv = self.resolution
+        indices = np.arange(nu * nv)
+
+        left  = indices - 1
+        right = indices + 1
+        up    = indices - nv
+        down  = indices + nv
+
+        left[0] = indices[0]
+        right[-1] = indices[-1]
+        up[:nv] = indices[:nv]
+        down[-nv:] = indices[-nv:]
+
+        points = self.get_points()
+        crosses = cross(
+            points[right] - points[left],
+            points[up] - points[down],
         )
-        return normalize_along_axis(normals, 1)
+        self.data["normal"] = normalize_along_axis(crosses, 1)
+        return self.data["normal"]
 
     @Mobject.affects_data
     def pointwise_become_partial(
@@ -147,12 +163,11 @@ class Surface(Mobject):
             return self
 
         nu, nv = smobject.resolution
-        for key in ['point', 'du_point', 'dv_point']:
-            self.data[key][:] = self.get_partial_points_array(
-                self.data[key], a, b,
-                (nu, nv, 3),
-                axis=axis
-            )
+        self.data['point'][:] = self.get_partial_points_array(
+            self.data['point'], a, b,
+            (nu, nv, 3),
+            axis=axis
+        )
         return self
 
     def get_partial_points_array(
@@ -263,8 +278,7 @@ class TexturedSurface(Surface):
     shader_folder: str = "textured_surface"
     shader_dtype: Sequence[Tuple[str, type, Tuple[int]]] = [
         ('point', np.float32, (3,)),
-        ('du_point', np.float32, (3,)),
-        ('dv_point', np.float32, (3,)),
+        ('normal', np.float32, (3,)),
         ('im_coords', np.float32, (2,)),
         ('opacity', np.float32, (1,)),
     ]
@@ -306,8 +320,7 @@ class TexturedSurface(Surface):
         surf = self.uv_surface
         nu, nv = surf.resolution
         self.resize_points(surf.get_num_points())
-        for key in ['point', 'du_point', 'dv_point']:
-            self.data[key][:] = surf.data[key]
+        self.data['point'][:] = surf.data['point']
         self.data['opacity'][:, 0] = surf.data["rgba"][:, 3]
         self.data["im_coords"] = np.array([
             [u, v]
diff --git a/manimlib/shaders/surface/vert.glsl b/manimlib/shaders/surface/vert.glsl
index d351878b..7cc0e387 100644
--- a/manimlib/shaders/surface/vert.glsl
+++ b/manimlib/shaders/surface/vert.glsl
@@ -3,8 +3,7 @@
 uniform vec4 clip_plane;
 
 in vec3 point;
-in vec3 du_point;
-in vec3 dv_point;
+in vec3 normal;
 in vec4 rgba;
 
 out vec4 v_color;
@@ -15,8 +14,7 @@ out vec4 v_color;
 
 void main(){
     emit_gl_Position(point);
-    vec3 normal = get_unit_normal(point, du_point, dv_point);
-    v_color = finalize_color(rgba, point, normal);
+    v_color = finalize_color(rgba, point, normalize(normal));
 
     if(clip_plane.xyz != vec3(0.0, 0.0, 0.0)){
         gl_ClipDistance[0] = dot(vec4(point, 1.0), clip_plane);