Trac #21329: add documentation to exhibit colored revolution plot3d

This works, and nobody knows..

Let us give a simple example.

I have also cleaned up the code (just some pep8 changes)

URL: https://trac.sagemath.org/21329
Reported by: chapoton
Ticket author(s): Frédéric Chapoton
Reviewer(s): Paul Masson

Author: Release Manager

src/sage/plot/plot3d/revolution_plot3d.py

@@ -104,6 +104,23 @@ def revolution_plot3d(curve,trange,phirange=None,parallel_axis='z',axis=(0,0),pr
         sage: u = var('u')
         sage: curve=(sin(3*u),.8*cos(4*u),cos(u))
         sage: revolution_plot3d(curve,(u,0,pi),(0,pi/2),show_curve=True,parallel_axis='z',opacity=0.5).show(aspect_ratio=(1,1,1),frame=False)
+
+    One can also color the surface using a coloring function of two
+    parameters and a colormap as follows::
+
+        sage: u, phi = var('u,phi')
+        sage: def cf(u,phi): return sin(phi+u) ** 2
+        sage: curve = (1+u**2/4, 0, u)
+        sage: revolution_plot3d(curve,(u,-2,2),(0,2*pi),parallel_axis='z',color=(cf, colormaps.PiYG)).show(aspect_ratio=(1,1,1))
+
+    The first parameter of the coloring function will be identified with the
+    parameter of the curve, and the second with the angle parameter.
+
+    .. WARNING::
+
+        This kind of coloring using a colormap can be visualized using
+        Jmol, Tachyon (option ``viewer='tachyon'``) and Canvas3D
+        (option ``viewer='canvas3d'`` in the notebook).
     """
     from sage.symbolic.ring import SR
     from sage.symbolic.constants import pi
@@ -112,77 +129,75 @@ def revolution_plot3d(curve,trange,phirange=None,parallel_axis='z',axis=(0,0),pr
     from sage.functions.trig import cos
     from sage.functions.trig import atan2
 
-
-    if parallel_axis not in ['x','y','z']:
+    if parallel_axis not in ['x', 'y', 'z']:
         raise ValueError("parallel_axis must be either 'x', 'y', or 'z'.")
 
-    vart=trange[0]
+    vart = trange[0]
 
-
-    if str(vart)=='phi':
+    if str(vart) == 'phi':
         phi = SR.var('fi')
     else:
         phi = SR.var('phi')
 
-
-    if phirange is None:#this if-else provides a phirange
-        phirange=(phi,0,2*pi)
-    elif len(phirange)==3:
-        phi=phirange[0]
+    if phirange is None:  # this if-else provides a phirange
+        phirange = (phi, 0, 2 * pi)
+    elif len(phirange) == 3:
+        phi = phirange[0]
         pass
     else:
-        phirange=(phi,phirange[0],phirange[1])
-
-    if isinstance(curve,tuple) or isinstance(curve,list):
+        phirange = (phi, phirange[0], phirange[1])
+        
+    if isinstance(curve, (tuple, list)):
         #this if-else provides a vector v to be plotted
         #if curve is a tuple or a list of length 2, it is interpreted as a parametric curve
         #in the x-z plane.
         #if it is of length 3 it is interpreted as a parametric curve in 3d space
 
-        if len(curve) ==2:
-            x=curve[0]
-            y=0
-            z=curve[1]
-        elif len(curve)==3:
-            x=curve[0]
-            y=curve[1]
-            z=curve[2]
+        if len(curve) == 2:
+            x = curve[0]
+            y = 0
+            z = curve[1]
+        elif len(curve) == 3:
+            x = curve[0]
+            y = curve[1]
+            z = curve[2]
     else:
-        x=vart
-        y=0
-        z=curve
-
-    if parallel_axis=='z':
-        x0=axis[0]
-        y0=axis[1]
+        x = vart
+        y = 0
+        z = curve
+
+    phase = 0
+    if parallel_axis == 'z':
+        x0 = axis[0]
+        y0 = axis[1]
         # (0,0) must be handled separately for the phase value
-        phase=0
-        if x0!=0 or y0!=0:
-            phase=atan2((y-y0),(x-x0))
-        R=sqrt((x-x0)**2+(y-y0)**2)
-        v=(R*cos(phi+phase)+x0,R*sin(phi+phase)+y0,z)
-    elif parallel_axis=='x':
-        y0=axis[0]
-        z0=axis[1]
+        if x0 != 0 or y0 != 0:
+            phase = atan2(y - y0, x - x0)
+        R = sqrt((x-x0)**2 + (y-y0)**2)
+        v = (R*cos(phi+phase)+x0, R*sin(phi+phase)+y0, z)
+    elif parallel_axis == 'x':
+        y0 = axis[0]
+        z0 = axis[1]
         # (0,0) must be handled separately for the phase value
-        phase=0
-        if z0!=0 or y0!=0:
-            phase=atan2((z-z0),(y-y0))
-        R=sqrt((y-y0)**2+(z-z0)**2)
-        v=(x,R*cos(phi+phase)+y0,R*sin(phi+phase)+z0)
-    elif parallel_axis=='y':
-        x0=axis[0]
-        z0=axis[1]
+        if z0 != 0 or y0 != 0:
+            phase = atan2(z - z0, y - y0)
+        R = sqrt((y-y0)**2 + (z-z0)**2)
+        v = (x, R*cos(phi+phase)+y0, R*sin(phi+phase)+z0)
+    elif parallel_axis == 'y':
+        x0 = axis[0]
+        z0 = axis[1]
         # (0,0) must be handled separately for the phase value
-        phase=0
-        if z0!=0 or x0!=0:
-            phase=atan2((z-z0),(x-x0))
-        R=sqrt((x-x0)**2+(z-z0)**2)
-        v=(R*cos(phi+phase)+x0,y,R*sin(phi+phase)+z0)
+        if z0 != 0 or x0 != 0:
+            phase = atan2(z - z0, x - x0)
+        R = sqrt((x-x0)**2 + (z-z0)**2)
+        v = (R*cos(phi+phase)+x0, y, R*sin(phi+phase)+z0)
 
     if print_vector:
         print(v)
+
     if show_curve:
-        curveplot=parametric_plot3d((x,y,z),trange,thickness=2,rgbcolor=(1,0,0))
-        return parametric_plot3d(v,trange,phirange,**kwds)+curveplot
-    return parametric_plot3d(v,trange,phirange,**kwds)
+        curveplot = parametric_plot3d((x, y, z), trange, thickness=2,
+                                      rgbcolor=(1, 0, 0))
+        return parametric_plot3d(v, trange, phirange, **kwds) + curveplot
+
+    return parametric_plot3d(v, trange, phirange, **kwds)