Enabled filling color by value for :class:.OpenGLSurface, replaced colors keyword argument of :meth:.Surface.set_fill_by_value with colorscale

manim/mobject/graphing/coordinate_systems.py

@@ -28,7 +28,9 @@
 from manim.mobject.graphing.number_line import NumberLine
 from manim.mobject.graphing.scale import LinearBase
 from manim.mobject.opengl.opengl_compatibility import ConvertToOpenGL
+from manim.mobject.opengl.opengl_surface import OpenGLSurface
 from manim.mobject.text.tex_mobject import MathTex
+from manim.mobject.three_d.three_dimensions import Surface
 from manim.mobject.types.vectorized_mobject import (
     VDict,
     VectorizedPoint,
@@ -858,6 +860,90 @@ def construct(self):
         graph.underlying_function = r_func
         return graph
 
+    def plot_surface(
+        self,
+        function: Callable[[float], float],
+        u_range: Sequence[float] | None = None,
+        v_range: Sequence[float] | None = None,
+        colorscale: Sequence[[color], float] | None = None,
+        colorscale_axis: int = 2,
+        **kwargs,
+    ):
+        """Generates a surface based on a function.
+
+        Parameters
+        ----------
+        function
+            The function used to construct the :class:`~.Surface`.
+        u_range
+            The range of the ``u`` variable: ``(u_min, u_max)``.
+        v_range
+            The range of the ``v`` variable: ``(v_min, v_max)``.
+        colorscale
+            Colors of the surface. Passing a list of colors will color the surface by z-value.
+            Passing a list of tuples in the form ``(color, pivot)`` allows user-defined pivots
+            where the color transitions.
+        colorscale_axis
+            Defines the axis on which the colorscale is applied (0 = x, 1 = y, 2 = z), default
+            is z-axis (2).
+        kwargs
+            Additional parameters to be passed to :class:`~.Surface`.
+
+        Returns
+        -------
+        :class:`~.Surface`
+            The plotted surface.
+
+        Examples
+        --------
+        .. manim:: PlotSurfaceExample
+            :save_last_frame:
+
+            class PlotSurfaceExample(ThreeDScene):
+                def construct(self):
+                    resolution_fa = 42
+                    self.set_camera_orientation(phi=75 * DEGREES, theta=-60 * DEGREES)
+                    axes = ThreeDAxes(x_range=(-3, 3, 1), y_range=(-3, 3, 1), z_range=(-5, 5, 1))
+                    def param_trig(u, v):
+                        x = u
+                        y = v
+                        z = 2 * np.sin(x) + 2 * np.cos(y)
+                        return z
+                    trig_plane = axes.plot_surface(
+                        param_trig,
+                        resolution=(resolution_fa, resolution_fa),
+                        u_range = (-3, 3),
+                        v_range = (-3, 3),
+                        colorscale = [BLUE, GREEN, YELLOW, ORANGE, RED],
+                        )
+                    self.add(axes, trig_plane)
+        """
+        if config.renderer != "opengl":
+            surface = Surface(
+                lambda u, v: self.c2p(u, v, function(u, v)),
+                u_range=u_range,
+                v_range=v_range,
+                **kwargs,
+            )
+            if colorscale:
+                surface.set_fill_by_value(
+                    axes=self.copy(),
+                    colorscale=colorscale,
+                    axis=colorscale_axis,
+                )
+        else:
+            surface = OpenGLSurface(
+                lambda u, v: self.c2p(u, v, function(u, v)),
+                u_range=u_range,
+                v_range=v_range,
+                axes=self.copy(),
+                colorscale=colorscale,
+                colorscale_axis=colorscale_axis,
+                **kwargs,
+            )
+
+        return surface
+
     def input_to_graph_point(
         self,
         x: float,

manim/mobject/opengl/opengl_surface.py

@@ -12,12 +12,44 @@
 from manim.utils.bezier import integer_interpolate, interpolate
 from manim.utils.color import *
 from manim.utils.config_ops import _Data, _Uniforms
+from manim.utils.deprecation import deprecated
 from manim.utils.images import change_to_rgba_array, get_full_raster_image_path
 from manim.utils.iterables import listify
 from manim.utils.space_ops import normalize_along_axis
 
 
 class OpenGLSurface(OpenGLMobject):
+    r"""Creates a Surface.
+
+    Parameters
+    ----------
+    uv_func
+        The function that defines the surface.
+    u_range
+        The range of the ``u`` variable: ``(u_min, u_max)``.
+    v_range
+        The range of the ``v`` variable: ``(v_min, v_max)``.
+    resolution
+        The number of samples taken of the surface.
+    axes
+        Axes on which the surface is to be drawn. Optional
+        parameter used when coloring a surface by z-value.
+    color
+        Color of the surface. Defaults to grey.
+    colorscale
+        Colors of the surface. Optional parameter used when
+        coloring a surface by values. Passing a list of
+        colors and an axes will color the surface by z-value.
+        Passing a list of tuples in the form ``(color, pivot)``
+        allows user-defined pivots where the color transitions.
+    colorscale_axis
+        Defines the axis on which the colorscale is applied
+        (0 = x, 1 = y, 2 = z), default is z-axis (2).
+    opacity
+        Opacity of the surface from 0 being fully transparent
+        to 1 being fully opaque. Defaults to 1.
+    """
+
     shader_dtype = [
         ("point", np.float32, (3,)),
         ("du_point", np.float32, (3,)),
@@ -35,7 +67,10 @@ def __init__(
         # each coordinate is one more than the the number of
         # rows/columns of approximating squares
         resolution=None,
+        axes=None,
         color=GREY,
+        colorscale=None,
+        colorscale_axis=2,
         opacity=1.0,
         gloss=0.3,
         shadow=0.4,
@@ -55,6 +90,9 @@ def __init__(
         # each coordinate is one more than the the number of
         # rows/columns of approximating squares
         self.resolution = resolution if resolution is not None else (101, 101)
+        self.axes = axes
+        self.colorscale = colorscale
+        self.colorscale_axis = colorscale_axis
         self.prefered_creation_axis = prefered_creation_axis
         # For du and dv steps.  Much smaller and numerical error
         # can crop up in the shaders.
@@ -206,22 +244,106 @@ def index_dot(index):
 
     # For shaders
     def get_shader_data(self):
+        """Called by parent Mobject to calculate and return
+        the shader data.
+
+        Returns
+        -------
+        shader_dtype
+            An array containing the shader data (vertices and
+            color of each vertex)
+        """
         s_points, du_points, dv_points = self.get_surface_points_and_nudged_points()
         shader_data = np.zeros(len(s_points), dtype=self.shader_dtype)
         if "points" not in self.locked_data_keys:
             shader_data["point"] = s_points
             shader_data["du_point"] = du_points
             shader_data["dv_point"] = dv_points
-        self.fill_in_shader_color_info(shader_data)
+            if self.colorscale:
+                shader_data["color"] = self._get_color_by_value(s_points)
+            else:
+                self.fill_in_shader_color_info(shader_data)
         return shader_data
 
     def fill_in_shader_color_info(self, shader_data):
+        """Fills in the shader color data when the surface
+        is all one color.
+
+        Parameters
+        ----------
+        shader_data
+            The vertices of the surface.
+
+        Returns
+        -------
+        shader_dtype
+            An array containing the shader data (vertices and
+            color of each vertex)
+        """
         self.read_data_to_shader(shader_data, "color", "rgbas")
         return shader_data
 
+    def _get_color_by_value(self, s_points):
+        """Matches each vertex to a color associated to it's z-value.
+
+        Parameters
+        ----------
+        s_points
+           The vertices of the surface.
+
+        Returns
+        -------
+        List
+            A list of colors matching the vertex inputs.
+        """
+        if type(self.colorscale[0]) in (list, tuple):
+            new_colors, pivots = [
+                [i for i, j in self.colorscale],
+                [j for i, j in self.colorscale],
+            ]
+        else:
+            new_colors = self.colorscale
+
+            pivot_min = self.axes.z_range[0]
+            pivot_max = self.axes.z_range[1]
+            pivot_frequency = (pivot_max - pivot_min) / (len(new_colors) - 1)
+            pivots = np.arange(
+                start=pivot_min,
+                stop=pivot_max + pivot_frequency,
+                step=pivot_frequency,
+            )
+
+        return_colors = []
+        for point in s_points:
+            axis_value = self.axes.point_to_coords(point)[self.colorscale_axis]
+            if axis_value <= pivots[0]:
+                return_colors.append(color_to_rgba(new_colors[0], self.opacity))
+            elif axis_value >= pivots[-1]:
+                return_colors.append(color_to_rgba(new_colors[-1], self.opacity))
+            else:
+                for i, pivot in enumerate(pivots):
+                    if pivot > axis_value:
+                        color_index = (axis_value - pivots[i - 1]) / (
+                            pivots[i] - pivots[i - 1]
+                        )
+                        color_index = max(min(color_index, 1), 0)
+                        temp_color = interpolate_color(
+                            new_colors[i - 1],
+                            new_colors[i],
+                            color_index,
+                        )
+                        break
+                return_colors.append(color_to_rgba(temp_color, self.opacity))
+
+        return return_colors
+
     def get_shader_vert_indices(self):
         return self.get_triangle_indices()
 
+    @deprecated(
+        since="v0.16.0",
+        message="Use colorscale attribute instead.",
+    )
     def set_fill_by_value(self, axes, colors):
         # directly copied from three_dimensions.py with some compatibility changes.
         """Sets the color of each mobject of a parametric surface to a color relative to its z-value

manim/mobject/three_d/three_dimensions.py

@@ -22,7 +22,7 @@
 import numpy as np
 from colour import Color
 
-from manim import config
+from manim import config, logger
 from manim.constants import *
 from manim.mobject.geometry.arc import Circle
 from manim.mobject.geometry.polygram import Square
@@ -31,6 +31,7 @@
 from manim.mobject.opengl.opengl_mobject import OpenGLMobject
 from manim.mobject.types.vectorized_mobject import VGroup, VMobject
 from manim.utils.color import *
+from manim.utils.deprecation import deprecated_params
 from manim.utils.iterables import tuplify
 from manim.utils.space_ops import normalize, perpendicular_bisector, z_to_vector
 
@@ -165,19 +166,21 @@ def set_fill_by_checkerboard(self, *colors, opacity=None):
             face.set_fill(colors[c_index], opacity=opacity)
         return self
 
+    @deprecated_params("colors", since="v0.16.0")
     def set_fill_by_value(
         self,
         axes: Mobject,
-        colors: Union[Iterable[Color], Color],
+        colorscale: Union[Iterable[Color], Color] | None = None,
         axis: int = 2,
+        **kwargs,
     ):
         """Sets the color of each mobject of a parametric surface to a color relative to its axis-value
 
         Parameters
         ----------
         axes :
             The axes for the parametric surface, which will be used to map axis-values to colors.
-        colors :
+        colorscale :
             A list of colors, ordered from lower axis-values to higher axis-values. If a list of tuples is passed
             containing colors paired with numbers, then those numbers will be used as the pivots.
         axis :
@@ -210,16 +213,32 @@ def param_surface(u, v):
                         u_range=[0, 5],
                         )
                     surface_plane.set_style(fill_opacity=1)
-                    surface_plane.set_fill_by_value(axes=axes, colors=[(RED, -0.5), (YELLOW, 0), (GREEN, 0.5)], axis=2)
+                    surface_plane.set_fill_by_value(axes=axes, colorscale=[(RED, -0.5), (YELLOW, 0), (GREEN, 0.5)], axis=2)
                     self.add(axes, surface_plane)
         """
+        if "colors" in kwargs and colorscale is None:
+            colorscale = kwargs.pop("colors")
+            if kwargs:
+                raise ValueError(
+                    "Unsupported keyword argument(s): "
+                    f"{', '.join(str(key) for key in kwargs)}"
+                )
+        if colorscale is None:
+            logger.warning(
+                "The value passed to the colorscale keyword argument was None, "
+                "the surface fill color has not been changed"
+            )
+            return self
 
         ranges = [axes.x_range, axes.y_range, axes.z_range]
 
-        if type(colors[0]) is tuple:
-            new_colors, pivots = [[i for i, j in colors], [j for i, j in colors]]
+        if type(colorscale[0]) is tuple:
+            new_colors, pivots = [
+                [i for i, j in colorscale],
+                [j for i, j in colorscale],
+            ]
         else:
-            new_colors = colors
+            new_colors = colorscale
 
             pivot_min = ranges[axis][0]
             pivot_max = ranges[axis][1]

tests/test_graphical_units/test_coordinate_systems.py

@@ -37,6 +37,46 @@ def test_line_graph(scene):
     scene.add(plane, first_line, second_line)
 
 
+@frames_comparison(base_scene=ThreeDScene)
+def test_plot_surface(scene):
+    axes = ThreeDAxes(x_range=(-5, 5, 1), y_range=(-5, 5, 1), z_range=(-5, 5, 1))
+
+    def param_trig(u, v):
+        x = u
+        y = v
+        z = 2 * np.sin(x) + 2 * np.cos(y)
+        return z
+
+    trig_plane = axes.plot_surface(
+        param_trig,
+        u_range=(-5, 5),
+        v_range=(-5, 5),
+        color=BLUE,
+    )
+
+    scene.add(axes, trig_plane)
+
+
+@frames_comparison(base_scene=ThreeDScene)
+def test_plot_surface_colorscale(scene):
+    axes = ThreeDAxes(x_range=(-3, 3, 1), y_range=(-3, 3, 1), z_range=(-5, 5, 1))
+
+    def param_trig(u, v):
+        x = u
+        y = v
+        z = 2 * np.sin(x) + 2 * np.cos(y)
+        return z
+
+    trig_plane = axes.plot_surface(
+        param_trig,
+        u_range=(-3, 3),
+        v_range=(-3, 3),
+        colorscale=[BLUE, GREEN, YELLOW, ORANGE, RED],
+    )
+
+    scene.add(axes, trig_plane)
+
+
 @frames_comparison
 def test_implicit_graph(scene):
     ax = Axes()