Better ThreeDVMobject implementation

camera/three_d_camera.py

@@ -7,6 +7,10 @@ from constants import *
 from camera.camera import Camera
 from mobject.types.point_cloud_mobject import Point
 from mobject.three_dimensions import ThreeDVMobject
+from mobject.three_d_utils import get_3d_vmob_start_corner
+from mobject.three_d_utils import get_3d_vmob_start_corner_unit_normal
+from mobject.three_d_utils import get_3d_vmob_end_corner
+from mobject.three_d_utils import get_3d_vmob_end_corner_unit_normal
 from mobject.value_tracker import ValueTracker
 
 from utils.color import get_shaded_rgb
@@ -53,9 +57,7 @@ class ThreeDCamera(Camera):
     def modified_rgbas(self, vmobject, rgbas):
         if not self.should_apply_shading:
             return rgbas
-        is_3d = isinstance(vmobject, ThreeDVMobject)
-        has_points = (vmobject.get_num_points() > 0)
-        if is_3d and has_points:
+        if vmobject.shade_in_3d and (vmobject.get_num_points() > 0):
             light_source_point = self.light_source.points[0]
             if len(rgbas) < 2:
                 shaded_rgbas = rgbas.repeat(2, axis=0)
@@ -63,14 +65,14 @@ class ThreeDCamera(Camera):
                 shaded_rgbas = np.array(rgbas[:2])
             shaded_rgbas[0, :3] = get_shaded_rgb(
                 shaded_rgbas[0, :3],
-                vmobject.get_start_corner(),
-                vmobject.get_start_corner_unit_normal(),
+                get_3d_vmob_start_corner(vmobject),
+                get_3d_vmob_start_corner_unit_normal(vmobject),
                 light_source_point,
             )
             shaded_rgbas[1, :3] = get_shaded_rgb(
                 shaded_rgbas[1, :3],
-                vmobject.get_end_corner(),
-                vmobject.get_end_corner_unit_normal(),
+                get_3d_vmob_end_corner(vmobject),
+                get_3d_vmob_end_corner_unit_normal(vmobject),
                 light_source_point,
             )
             return shaded_rgbas
@@ -92,7 +94,7 @@ class ThreeDCamera(Camera):
         def z_key(vmob):
             # Assign a number to a three dimensional mobjects
             # based on how close it is to the camera
-            if isinstance(vmob, ThreeDVMobject):
+            if vmob.shade_in_3d:
                 return np.dot(
                     vmob.get_center(),
                     rot_matrix.T

mobject/coordinate_systems.py

@@ -11,7 +11,6 @@ from mobject.number_line import NumberLine
 from mobject.svg.tex_mobject import TexMobject
 from mobject.types.vectorized_mobject import VGroup
 from mobject.types.vectorized_mobject import VMobject
-from mobject.three_dimensions import ThreeDVMobject
 from utils.config_ops import digest_config
 from utils.space_ops import angle_of_vector
 
@@ -137,16 +136,13 @@ class ThreeDAxes(Axes):
         self.set_axis_shading()
 
     def add_3d_pieces(self):
-        for attr in "x_axis", "y_axis", "z_axis":
-            axis = getattr(self, attr)
+        for axis in self:
             axis.add(VGroup(
                 *axis.main_line.get_pieces(self.num_axis_pieces)
             ))
             axis.main_line.set_stroke(width=0, family=False)
-            axis_3d = ThreeDVMobject(axis)
-            self.remove(axis)
-            self.add(axis_3d)
-            setattr(self, attr, axis_3d)
+            for submob in axis.family_members_with_points():
+                submob.shade_in_3d = True
 
     def set_axis_shading(self):
         def make_func(axis):

mobject/three_d_shading_utils.py

@@ -0,0 +1,56 @@
+import numpy as np
+
+from constants import ORIGIN
+from utils.space_ops import get_unit_normal
+
+
+def get_3d_vmob_gradient_start_and_end_points(vmob):
+    return (
+        get_3d_vmob_start_corner(vmob),
+        get_3d_vmob_end_corner(vmob),
+    )
+
+
+def get_3d_vmob_start_corner_index(vmob):
+    return 0
+
+
+def get_3d_vmob_end_corner_index(vmob):
+    return ((len(vmob.points) - 1) // 6) * 3
+
+
+def get_3d_vmob_start_corner(vmob):
+    if vmob.get_num_points() == 0:
+        return np.array(ORIGIN)
+    return vmob.points[get_3d_vmob_start_corner_index(vmob)]
+
+
+def get_3d_vmob_end_corner(vmob):
+    if vmob.get_num_points() == 0:
+        return np.array(ORIGIN)
+    return vmob.points[get_3d_vmob_end_corner_index(vmob)]
+
+
+def get_3d_vmob_unit_normal(vmob, point_index):
+    n_points = vmob.get_num_points()
+    if vmob.get_num_points() == 0:
+        return np.array(ORIGIN)
+    i = point_index
+    im1 = i - 1 if i > 0 else (n_points - 2)
+    ip1 = i + 1 if i < (n_points - 1) else 1
+    return get_unit_normal(
+        vmob.points[ip1] - vmob.points[i],
+        vmob.points[im1] - vmob.points[i],
+    )
+
+
+def get_3d_vmob_start_corner_unit_normal(vmob):
+    return get_3d_vmob_unit_normal(
+        vmob, get_3d_vmob_start_corner_index(vmob)
+    )
+
+
+def get_3d_vmob_end_corner_unit_normal(vmob):
+    return get_3d_vmob_unit_normal(
+        vmob, get_3d_vmob_end_corner_index(vmob)
+    )

mobject/three_d_utils.py

@@ -0,0 +1,61 @@
+import numpy as np
+
+from constants import ORIGIN
+from constants import UP
+from utils.space_ops import get_unit_normal
+from utils.space_ops import get_norm
+
+
+def get_3d_vmob_gradient_start_and_end_points(vmob):
+    return (
+        get_3d_vmob_start_corner(vmob),
+        get_3d_vmob_end_corner(vmob),
+    )
+
+
+def get_3d_vmob_start_corner_index(vmob):
+    return 0
+
+
+def get_3d_vmob_end_corner_index(vmob):
+    return ((len(vmob.points) - 1) // 6) * 3
+
+
+def get_3d_vmob_start_corner(vmob):
+    if vmob.get_num_points() == 0:
+        return np.array(ORIGIN)
+    return vmob.points[get_3d_vmob_start_corner_index(vmob)]
+
+
+def get_3d_vmob_end_corner(vmob):
+    if vmob.get_num_points() == 0:
+        return np.array(ORIGIN)
+    return vmob.points[get_3d_vmob_end_corner_index(vmob)]
+
+
+def get_3d_vmob_unit_normal(vmob, point_index):
+    n_points = vmob.get_num_points()
+    if vmob.get_num_points() == 0:
+        return np.array(UP)
+    i = point_index
+    im1 = i - 1 if i > 0 else (n_points - 2)
+    ip1 = i + 1 if i < (n_points - 1) else 1
+    unit_normal = get_unit_normal(
+        vmob.points[ip1] - vmob.points[i],
+        vmob.points[im1] - vmob.points[i],
+    )
+    if get_norm(unit_normal) == 0:
+        return np.array(UP)
+    return unit_normal
+
+
+def get_3d_vmob_start_corner_unit_normal(vmob):
+    return get_3d_vmob_unit_normal(
+        vmob, get_3d_vmob_start_corner_index(vmob)
+    )
+
+
+def get_3d_vmob_end_corner_unit_normal(vmob):
+    return get_3d_vmob_unit_normal(
+        vmob, get_3d_vmob_end_corner_index(vmob)
+    )

mobject/three_dimensions.py

@@ -6,68 +6,16 @@ from mobject.types.vectorized_mobject import VMobject
 from mobject.types.vectorized_mobject import VGroup
 from mobject.geometry import Square
 
-from utils.config_ops import digest_config
 from utils.iterables import tuplify
 from utils.space_ops import z_to_vector
-from utils.space_ops import get_unit_normal
 
 ##############
 
 
 class ThreeDVMobject(VMobject):
-    CONFIG = {}
-
-    def __init__(self, vmobject=None, **kwargs):
-        VMobject.__init__(self, **kwargs)
-        if vmobject is not None:
-            self.points = np.array(vmobject.points)
-            self.match_style(vmobject, family=False)
-            self.submobjects = [
-                ThreeDVMobject(submob, **kwargs)
-                for submob in vmobject.submobjects
-            ]
-
-    def get_gradient_start_and_end_points(self):
-        return self.get_start_corner(), self.get_end_corner()
-
-    def get_start_corner_index(self):
-        return 0
-
-    def get_end_corner_index(self):
-        return ((len(self.points) - 1) // 6) * 3
-        # return ((len(self.points) - 1) // 12) * 3
-
-    def get_start_corner(self):
-        if self.get_num_points() == 0:
-            return np.array(ORIGIN)
-        return self.points[self.get_start_corner_index()]
-
-    def get_end_corner(self):
-        if self.get_num_points() == 0:
-            return np.array(ORIGIN)
-        return self.points[self.get_end_corner_index()]
-
-    def get_unit_normal(self, point_index):
-        n_points = self.get_num_points()
-        if self.get_num_points() == 0:
-            return np.array(ORIGIN)
-        i = point_index
-        im1 = i - 1 if i > 0 else (n_points - 2)
-        ip1 = i + 1 if i < (n_points - 1) else 1
-        return get_unit_normal(
-            self.points[ip1] - self.points[i],
-            self.points[im1] - self.points[i],
-        )
-
-    def get_start_corner_unit_normal(self):
-        return self.get_unit_normal(
-            self.get_start_corner_index()
-        )
-
-    def get_end_corner_unit_normal(self):
-        return self.get_unit_normal(
-            self.get_end_corner_index()
-        )
+    CONFIG = {
+        "shade_in_3d": True,
+    }
 
 
 class ParametricSurface(VGroup):

mobject/types/vectorized_mobject.py

@@ -3,6 +3,7 @@
 from colour import Color
 
 from mobject.mobject import Mobject
+from mobject.three_d_utils import get_3d_vmob_gradient_start_and_end_points
 from constants import *
 from utils.bezier import bezier
 from utils.bezier import get_smooth_handle_points
@@ -44,6 +45,7 @@ class VMobject(Mobject):
         "pre_function_handle_to_anchor_scale_factor": 0.01,
         "make_smooth_after_applying_functions": False,
         "background_image_file": None,
+        "shade_in_3d": False,
     }
 
     def get_group_class(self):
@@ -280,14 +282,17 @@ class VMobject(Mobject):
         return self.sheen
 
     def get_gradient_start_and_end_points(self):
-        direction = self.get_sheen_direction()
-        c = self.get_center()
-        bases = np.array([
-            self.get_edge_center(vect) - c
-            for vect in [RIGHT, UP, OUT]
-        ]).transpose()
-        offset = np.dot(bases, direction)
-        return (c - offset, c + offset)
+        if self.shade_in_3d:
+            return get_3d_vmob_gradient_start_and_end_points(self)
+        else:
+            direction = self.get_sheen_direction()
+            c = self.get_center()
+            bases = np.array([
+                self.get_edge_center(vect) - c
+                for vect in [RIGHT, UP, OUT]
+            ]).transpose()
+            offset = np.dot(bases, direction)
+            return (c - offset, c + offset)
 
     def color_using_background_image(self, background_image_file):
         self.background_image_file = background_image_file