color_func is now back-compatible with sympy.plotting

Author: Davide-sd

spb/backends/k3d.py

@@ -252,15 +252,15 @@ def _process_series(self, series):
                     x, y, z, u, v = s.get_data()
                     vertices, indices = get_vertices_indices(x, y, z)
                     vertices = vertices.astype(np.float32)
-                    attribute = s.color_func(vertices[:, 0], vertices[:, 1], vertices[:, 2], u.flatten().astype(np.float32), v.flatten().astype(np.float32))
+                    attribute = s.eval_color_func(vertices[:, 0], vertices[:, 1], vertices[:, 2], u.flatten().astype(np.float32), v.flatten().astype(np.float32))
                 else:
                     x, y, z = s.get_data()
                     x = x.flatten()
                     y = y.flatten()
                     z = z.flatten()
                     vertices = np.vstack([x, y, z]).T.astype(np.float32)
                     indices = Triangulation(x, y).triangles.astype(np.uint32)
-                    attribute = s.color_func(vertices[:, 0], vertices[:, 1], vertices[:, 2])
+                    attribute = s.eval_color_func(vertices[:, 0], vertices[:, 1], vertices[:, 2])
 
                 self._high_aspect_ratio(x, y, z)
                 a = dict(
@@ -493,11 +493,11 @@ def _update_interactive(self, params):
                     if s.is_parametric:
                         x, y, z, u, v = s.get_data()
                         x, y, z, u, v = [t.flatten().astype(np.float32) for t in [x, y, z, u, v]]
-                        attribute = s.color_func(x, y, z, u, v)
+                        attribute = s.eval_color_func(x, y, z, u, v)
                     else:
                         x, y, z = s.get_data()
                         x, y, z = [t.flatten().astype(np.float32) for t in [x, y, z]]
-                        attribute = s.color_func(x, y, z)
+                        attribute = s.eval_color_func(x, y, z)
 
                     vertices = np.vstack([x, y, z]).astype(np.float32)
                     self._fig.objects[i].vertices = vertices.T

spb/backends/matplotlib.py

@@ -413,10 +413,10 @@ def _process_series(self, series):
             elif (s.is_3Dsurface and (not s.is_domain_coloring) and (not s.is_implicit)):
                 if not s.is_parametric:
                     x, y, z = self.series[i].get_data()
-                    facecolors = s.color_func(x, y, z)
+                    facecolors = s.eval_color_func(x, y, z)
                 else:
                     x, y, z, u, v = self.series[i].get_data()
-                    facecolors = s.color_func(x, y, z, u, v)
+                    facecolors = s.eval_color_func(x, y, z, u, v)
                 skw = dict(rstride=1, cstride=1, linewidth=0.1)
                 norm, cmap = None, None
                 if s.use_cm:
@@ -855,10 +855,10 @@ def _update_interactive(self, params):
                 elif s.is_3Dsurface and (not s.is_domain_coloring) and (not s.is_implicit):
                     if not s.is_parametric:
                         x, y, z = self.series[i].get_data()
-                        facecolors = s.color_func(x, y, z)
+                        facecolors = s.eval_color_func(x, y, z)
                     else:
                         x, y, z, u, v = self.series[i].get_data()
-                        facecolors = s.color_func(x, y, z, u, v)
+                        facecolors = s.eval_color_func(x, y, z, u, v)
                     # TODO: by setting the keyword arguments, somehow the
                     # update becomes really really slow.
                     kw, is_cb_added, cax = self._handles[i][1:]

spb/backends/plotly.py

@@ -349,10 +349,10 @@ def _process_series(self, series):
             elif s.is_3Dsurface and (not s.is_domain_coloring) and (not s.is_implicit):
                 if not s.is_parametric:
                     xx, yy, zz = s.get_data()
-                    surfacecolor = s.color_func(xx, yy, zz)
+                    surfacecolor = s.eval_color_func(xx, yy, zz)
                 else:
                     xx, yy, zz, uu, vv = s.get_data()
-                    surfacecolor = s.color_func(xx, yy, zz, uu, vv)
+                    surfacecolor = s.eval_color_func(xx, yy, zz, uu, vv)
 
                 # create a solid color to be used when s.use_cm=False
                 col = next(self._cl)
@@ -658,10 +658,10 @@ def _update_interactive(self, params):
                 elif s.is_3Dsurface and (not s.is_domain_coloring) and (not s.is_implicit):
                     if not s.is_parametric:
                         x, y, z = s.get_data()
-                        surfacecolor = s.color_func(x, y, z)
+                        surfacecolor = s.eval_color_func(x, y, z)
                     else:
                         x, y, z, u, v = s.get_data()
-                        surfacecolor = s.color_func(x, y, z, u, v)
+                        surfacecolor = s.eval_color_func(x, y, z, u, v)
                         self.fig.data[i]["x"] = x
                         self.fig.data[i]["y"] = y
 

spb/functions.py

@@ -626,10 +626,16 @@ def plot_parametric(*args, show=True, **kwargs):
     backend : Plot, optional
         A subclass of `Plot`, which will perform the rendering.
         Default to `MatplotlibBackend`.
-    
+
     color_func : callable, optional
-        A function of 3 variables, x, y, parameter (the points computed by
-        the internal algorithm) which defines the line color. Default to None.
+        A function defining the line color. The arity can be:
+
+        * 1 argument: ``f(t)``, where ``t`` is the parameter.
+        * 2 arguments: ``f(x, y)`` where ``x, y`` are the coordinates of the
+          points.
+        * 3 arguments: ``f(x, y, t)``.
+
+        Default to None.
 
     label : str or list/tuple, optional
         The label to be shown in the legend or in the colorbar. If not
@@ -844,10 +850,16 @@ def plot3d_parametric_line(*args, show=True, **kwargs):
     backend : Plot, optional
         A subclass of `Plot`, which will perform the rendering.
         Default to `MatplotlibBackend`.
-    
+
     color_func : callable, optional
-        A function of 4 variables, x, y, z, parameter (the points computed by
-        the internal algorithm) which defines the line color. Default to None.
+        A function defining the line color. The arity can be:
+
+        * 1 argument: ``f(t)``, where ``t`` is the parameter.
+        * 3 arguments: ``f(x, y, z)`` where ``x, y, z`` are the coordinates of
+          the points.
+        * 4 arguments: ``f(x, y, z, t)``.
+
+        Default to None.
 
     label : str or list/tuple, optional
         The label to be shown in the legend or in the colorbar. If not
@@ -1306,9 +1318,16 @@ def plot3d_parametric_surface(*args, show=True, **kwargs):
         Default to `MatplotlibBackend`.
 
     color_func : callable, optional
-        A function of 5 variables, x, y, z, u, v (the points computed by the
-        internal algorithm and the parameters) which defines the surface color
-        when ``use_cm=True``. Default to None.
+        A function defining the surface color when ``use_cm=True``. The arity
+        can be:
+
+        * 1 argument: ``f(u)``, where ``u`` is the first parameter.
+        * 2 arguments: ``f(u, v)`` where ``u, v`` are the parameters.
+        * 3 arguments: ``f(x, y, z)`` where ``x, y, z`` are the coordinates of
+          the points.
+        * 5 arguments: ``f(x, y, z, u, v)``.
+
+        Default to None.
 
     label : str or list/tuple, optional
         The label to be shown in the colorbar. If not provided, the string
@@ -1401,7 +1420,7 @@ def plot3d_parametric_surface(*args, show=True, **kwargs):
        ...          r * cos(v)
        ... )
        >>> plot3d_parametric_surface(*expr, (u, 0, 2 * pi), (v, 0, pi), "u",
-       ...      n=200, use_cm=True, color_func=lambda x, y, z, u, v: u)
+       ...      n=200, use_cm=True, color_func=lambda u, v: u)
        Plot object containing:
        [0]: parametric cartesian surface: ((sin(7*u + 5*v) + 2)*sin(v)*cos(u), (sin(7*u + 5*v) + 2)*sin(u)*sin(v), (sin(7*u + 5*v) + 2)*cos(v)) for u over (0.0, 6.283185307179586) and v over (0.0, 3.141592653589793)
 

spb/series.py

@@ -1,6 +1,7 @@
 from spb.defaults import cfg
 from sympy import latex
 from sympy.core.containers import Tuple
+from sympy.core.function import arity
 from sympy.core.symbol import symbols
 from sympy.core.sympify import sympify
 from sympy.solvers.solvers import solve
@@ -376,6 +377,43 @@ def _correct_size(a, b):
                 a = a.reshape(b.shape)
         return a
 
+    def eval_color_func(self, *args):
+        """Evaluate the color function.
+
+        Parameters
+        ==========
+
+        args : tuple
+            Arguments to be passed to the coloring function. Can be coordinates
+            or parameters or both.
+
+        Notes
+        =====
+
+        The backend will request the data series to generate the numerical
+        data. Depending on the data series, either the data series itself or
+        the backend will eventually execute this function to generate the
+        appropriate coloring value.
+        """
+        nargs = arity(self.color_func)
+        if nargs == 1:
+            if self.is_2Dline and self.is_parametric:
+                if len(args) == 2:
+                    # ColoredLineOver1DRangeSeries
+                    return self.color_func(args[0])
+                # Parametric2DLineSeries
+                return self.color_func(args[2])
+            elif self.is_3Dline and self.is_parametric:
+                return self.color_func(args[3])
+            elif self.is_3Dsurface and self.is_parametric:
+                return self.color_func(args[3])
+            return self.color_func(args[0])
+        elif nargs == 2:
+            if self.is_3Dsurface and self.is_parametric:
+                return self.color_func(*args[3:])
+            return self.color_func(*args[:2])
+        return self.color_func(*args[:nargs])
+
     def get_data(self):
         """Compute and returns the numerical data.
 
@@ -723,7 +761,7 @@ def get_points(self):
             Color associated to each point.
         """
         x, y = super().get_points()
-        return x, y, self.color_func(x, y)
+        return x, y, self.eval_color_func(x, y)
 
 
 class AbsArgLineSeries(LineOver1DRangeSeries):
@@ -878,7 +916,7 @@ def get_points(self):
 
         if callable(self.color_func):
             coords = list(coords)
-            coords[-1] = self.color_func(*coords)
+            coords[-1] = self.eval_color_func(*coords)
         return coords
 
 
@@ -1818,7 +1856,7 @@ def get_points(self):
             Color associated to each point.
         """
         x, y = super().get_points()
-        return x, y, self.color_func(x, y)
+        return x, y, self.eval_color_func(x, y)
 
 
 class AbsArgLineInteractiveSeries(LineInteractiveSeries):

tests/test_series.py

@@ -2364,3 +2364,82 @@ def test_is_polar_3d():
     x22, y22 = x1 * np.cos(y1), x1 * np.sin(y1)
     assert np.allclose(x2, x22)
     assert np.allclose(y2, y22)
+
+
+def test_color_func():
+    # verify that eval_color_func produces the expected results in order to
+    # maintain back compatibility with the old sympy.plotting module
+
+    x, y, z, u, v = symbols("x, y, z, u, v")
+
+    s = LineOver1DRangeSeries(sin(x), (x, -5, 5), adaptive=False, n=10,
+        color_func=lambda x: x)
+    xx, yy, col = s.get_data()
+    assert np.allclose(col, xx)
+    s = LineOver1DRangeSeries(sin(x), (x, -5, 5), adaptive=False, n=10,
+        color_func=lambda x, y: y)
+    xx, yy, col = s.get_data()
+    assert np.allclose(col, yy)
+
+    s = Parametric2DLineSeries(cos(x), sin(x), (x, 0, 2*pi),
+        adaptive=False, n=10, color_func=lambda t: t)
+    xx, yy, col = s.get_data()
+    assert (not np.allclose(xx, col)) and (not np.allclose(yy, col))
+    s = Parametric2DLineSeries(cos(x), sin(x), (x, 0, 2*pi),
+        adaptive=False, n=10, color_func=lambda x, y: x * y)
+    xx, yy, col = s.get_data()
+    assert np.allclose(col, xx * yy)
+    s = Parametric2DLineSeries(cos(x), sin(x), (x, 0, 2*pi),
+        adaptive=False, n=10, color_func=lambda x, y, t: x * y * t)
+    xx, yy, col = s.get_data()
+    assert np.allclose(col, xx * yy * np.linspace(0, 2*np.pi, 10))
+
+    s = Parametric3DLineSeries(cos(x), sin(x), x, (x, 0, 2*pi),
+        adaptive=False, n=10, color_func=lambda t: t)
+    xx, yy, zz, col = s.get_data()
+    assert (not np.allclose(xx, col)) and (not np.allclose(yy, col))
+    s = Parametric3DLineSeries(cos(x), sin(x), x, (x, 0, 2*pi),
+        adaptive=False, n=10, color_func=lambda x, y, z: x * y * z)
+    xx, yy, zz, col = s.get_data()
+    assert np.allclose(col, xx * yy * zz)
+    s = Parametric3DLineSeries(cos(x), sin(x), x, (x, 0, 2*pi),
+        adaptive=False, n=10, color_func=lambda x, y, z, t: x * y * z * t)
+    xx, yy, zz, col = s.get_data()
+    assert np.allclose(col, xx * yy * zz * np.linspace(0, 2*np.pi, 10))
+
+    s = SurfaceOver2DRangeSeries(cos(x**2 + y**2), (x, -2, 2), (y, -2, 2),
+        adaptive=False, n1=10, n2=10, color_func=lambda x: x)
+    xx, yy, zz = s.get_data()
+    col = s.eval_color_func(xx, yy, zz)
+    assert np.allclose(xx, col)
+    s = SurfaceOver2DRangeSeries(cos(x**2 + y**2), (x, -2, 2), (y, -2, 2),
+        adaptive=False, n1=10, n2=10, color_func=lambda x, y: x * y)
+    xx, yy, zz = s.get_data()
+    col = s.eval_color_func(xx, yy, zz)
+    assert np.allclose(xx * yy, col)
+    s = SurfaceOver2DRangeSeries(cos(x**2 + y**2), (x, -2, 2), (y, -2, 2),
+        adaptive=False, n1=10, n2=10, color_func=lambda x, y, z: x * y * z)
+    xx, yy, zz = s.get_data()
+    col = s.eval_color_func(xx, yy, zz)
+    assert np.allclose(xx * yy * zz, col)
+
+    s = ParametricSurfaceSeries(1, x, y, (x, 0, 1), (y, 0, 1), adaptive=False,
+        n1=10, n2=10, color_func=lambda u:u)
+    xx, yy, zz, uu, vv = s.get_data()
+    col = s.eval_color_func(xx, yy, zz, uu, vv)
+    assert np.allclose(uu, col)
+    s = ParametricSurfaceSeries(1, x, y, (x, 0, 1), (y, 0, 1), adaptive=False,
+        n1=10, n2=10, color_func=lambda u, v: u * v)
+    xx, yy, zz, uu, vv = s.get_data()
+    col = s.eval_color_func(xx, yy, zz, uu, vv)
+    assert np.allclose(uu * vv, col)
+    s = ParametricSurfaceSeries(1, x, y, (x, 0, 1), (y, 0, 1), adaptive=False,
+        n1=10, n2=10, color_func=lambda x, y, z: x * y * z)
+    xx, yy, zz, uu, vv = s.get_data()
+    col = s.eval_color_func(xx, yy, zz, uu, vv)
+    assert np.allclose(xx * yy * zz, col)
+    s = ParametricSurfaceSeries(1, x, y, (x, 0, 1), (y, 0, 1), adaptive=False,
+        n1=10, n2=10, color_func=lambda x, y, z, u, v: x * y * z * u * v)
+    xx, yy, zz, uu, vv = s.get_data()
+    col = s.eval_color_func(xx, yy, zz, uu, vv)
+    assert np.allclose(xx * yy * zz * uu * vv, col)