3910.sagews

Colormap for implicit_plot3d and parametric_plot3d (#12212)

jvkersch

The class below is based on some code from sage.plot.plot3d.plot3d and allows for a texture attribute of a surface to be varied gradually.  The precise way in which the attribute varies can be specified by means of a scalar function, the most straightforward ones being projection onto the x, y, or z-axis.  In this way, it is possible to apply a colormap to a surface, to create partially translucent objects, or to color in surfaces in a variety of ways (for instance, to visualize Gaussian curvature).

from sage.plot.plot3d.base import Graphics3dGroup
from sage.structure.sage_object import SageObject

class GradualTextureTransform(SageObject):
    def __init__(self, attribute, values, fun=None, bounds=None):
        """
        INPUT::
            
            - ``attribute`` -- texture attribute to transform.  
            If this is not one of ``color``, ``opacity``, 
            ``ambient``, ``diffuse``, ``specular`` or ``shininess``,
            the transformation has no effect.
           
            - ``value_range`` -- range of values over which the 
            attribute should range.  In the case where the color 
            attribute is transformed, this can be a colormap.
            
            - ``fun`` -- scalar-valued function specifying the 
            change in attribute.  If this keyword is omitted, 
            the projection x, y, z -> z will be used.  This 
            could be either a symbolic expression, or one of 
            the following: 'height_x', 'height_y', 'height_z', 
            denoting respectively the height above the x-, y-, 
            or z-plane.
            
            - ``bounds`` -- tuple representing the range for the 
            scalar function ``fun``.  If this keyword is omitted, 
            some preset defaults will be used.    
            
        """

        self.attribute = attribute
        self.fun       = fun
        self.bounds    = bounds
        self.values    = list(values)

        self.projection_axis = 2

        if fun is None:
            self.fun = lambda x, y, z: z
            self.projection_axis = 2
        elif isinstance(fun, str):
            if fun is "height_x":
                self.fun = lambda x, y, z: x
                self.projection_axis = 0
            elif fun is "height_y":
                self.fun = lambda x, y, z: y
                self.projection_axis = 1
            elif fun is "height_z":
                self.fun = lambda x, y, z: z
                self.projection_axis = 2
            else:
                raise ValueError, "Not a valid texture function:", fun


    def apply(self, group):
        """
        Apply this texture transformation to a graphical object.
        
        INPUT::
            
            - ``group`` -- a graphical object (``Graphics3d``) or 
            group of objects (``Graphics3dGroup``).
            
        OUTPUT::
            
            - A ``Graphics3dGroup`` representing the transformed object.     
        """

        if not isinstance(group, Graphics3dGroup):
            group = Graphics3dGroup([group])

        # Code below adapted from plot3d_adaptive
        if self.bounds is None:
            # Determine limits for function from bounding box
            bounds = group.bounding_box()

            min_z = bounds[0][self.projection_axis]
            max_z = bounds[1][self.projection_axis]
        else:
            min_z = self.bounds[0]
            max_z = self.bounds[1]

        if max_z == min_z:
            span = 0
        else:
            span = (len(self.values)-1) / (max_z - min_z)  

        def normalized_fun(x, y, z):
            val = self.fun(x, y, z) 
            if val > max_z: val = max_z
            if val < min_z: val = min_z
            return int((val-min_z)*span)

        new_all = []
        kwds = {}
        for obj in group.all:
            obj.triangulate()  # Faces have to exist before object can be partitioned
            parts = obj.partition(normalized_fun) 
     
            for k, G in parts.iteritems():
                kwds[self.attribute] = self.values[k]
                G.set_texture(kwds)
                new_all.append(G)

        return Graphics3dGroup(new_all)



class ColormapTransform(GradualTextureTransform):
    """
    Transformation to apply a colormap to a given surface.  
    
    See ``GradualTextureTransform`` for more details.
    """
    def __init__(self, cmap, fun=None, bounds=None):
        GradualTextureTransform.__init__(self, 'color', cmap, fun, bounds)
        

class OpacityTransform(GradualTextureTransform):
    """
    Transformation to change the opacity of given surface in a non-uniform way. 
    
    See ``GradualTextureTransform`` for more details.
    """

    def __init__(self, values, fun=None, bounds=None):
        GradualTextureTransform.__init__(self, 'opacity', values, fun, bounds)

# Predefined colormaps
cmsel1 = [colormaps['autumn'](i) for i in sxrange(0,1,0.05)]
cmsel2 = [colormaps['spectral'](i) for i in sxrange(0,1,0.02)]

# The original example from the ticket, where using a colormap worked out of the box

var('r v')
p = plot3d(0.2*(r**2 + v**2) + cos(2*r)*sin(2*v),(r,-2,2), (v,-2,2), adaptive=True, color=cmsel1, plot_points=10, opacity=0.9)
p2 = sphere((0,0,0),1,color='black',opacity=0.5)
(p+p2).show(aspect_ratio=(1,1,1))

# The same example, but with the functionality of the ColormapTransform class:

var('r v')
c = ColormapTransform(cmsel1)
p1 = plot3d(0.2*(r**2 + v**2) + cos(2*r)*sin(2*v),(r,-2,2), (v,-2,2), opacity=0.9)
p2 = sphere((0,0,0),1,color='black',opacity=0.5)
p3 = c.apply(p1)
(p2+p3).show(aspect_ratio=(1,1,1))

# An example of an implicit surface with a color map (this is the non-functioning example from the ticket):

c = ColormapTransform(cmsel1)
var('x,y,z')
p1 = implicit_plot3d(x^2+y^2+z^2==4, (x, -3, 3), (y, -3,3), (z, -3,3))
p2 = c.apply(p1)
p2.show()

# By default, the color values are selected according to the z-values of the points on the surface.  Other functions are possible, too:

# Uniform gradient along the x-axis
var('r v')
p = plot3d(0.2*(r**2 + v**2) + cos(2*r)*sin(2*v),(r,-2,2), (v,-2,2))
c = ColormapTransform(cmsel1, 'height_x')
c.apply(p).show(aspect_ratio=(1,1,1))

# Non-uniform coloring.  Here a function is specified.

var('x,y,z')
p = plot3d(0.2*(r**2 + v**2) + cos(2*r)*sin(2*v),(r,-2,2), (v,-2,2))
c = ColormapTransform(fun=lambda x, y, z: x+y, bounds=(-4, 4), cmap=cmsel1)
c.apply(p)

# Another non-uniform coloring, this time with a nonlinear function (cos)

var('r v')
p = implicit_plot3d(x^2+y^2+z^2==4, (x, -3, 3), (y, -3,3), (z, -3,3))
c = ColormapTransform(cmsel1, fun=lambda x, y, z: cos(2*pi*z/3-pi/3), bounds=(-1, 1))
c.apply(p).show(aspect_ratio=(1,1,1))

# A more complicated surface (adapted from the implicit_plot docs)

u, v = var('u,v')
fx = u -u^3/3  + u*v^2
fy = v -v^3/3  + v*u^2
fz = u^2 - v^2
p = parametric_plot3d([fx, fy, fz], (u, -2, 2), (v, -2, 2))
c = ColormapTransform(cmsel2)
c.apply(p)

# We can also gradually vary other texture attributes, such as opacity:
# NOTE: I have no idea why the outer sphere is so coarse after the transformation is applied.  Increasing plot_points in the definition for p2 does not seem to help
    
p1 = sphere((0,0,0), .5, color='red')
p2 = sphere((0,0,0),  1)

opacity_values = sxrange(0, 1, .05)
c = OpacityTransform(opacity_values)
p3 = c.apply(p2)
(p1 + p3).show()

# The next three examples show implicit surfaces colored by Gaussian curvature

# Quick 'n' dirty routine to compute the Gaussian curvature of an implicit surface:
def gaussian_curvature_implicit_surface(f):    
    fx = diff(f, x); fy = diff(f, y); fz = diff(f, z)
    fxx = diff(fx, x); fxy = diff(fx, y); fxz = diff(fx, z)
    fyy = diff(fy, y); fyz = diff(fy, z); fzz = diff(fz, z)

    T1 = (fz*(fxx*fz - 2*fx*fxz) + fx^2*fzz).simplify_full()
    T2 = (fz*(fyy*fz - 2*fy*fyz) + fy^2*fzz).simplify_full()
    T3 = (fz*(-fx*fyz + fxy*fz - fxz*fy) + fx*fy*fzz).simplify_full()
    T4 = (fz^2*(fx^2 + fy^2 + fz^2)^2).simplify_full()

    K = ((T1*T2 - T3^2)/T4).simplify_full()
    return fast_float(K, 'x', 'y', 'z')

# Ellipsoid 
f = x^2/4 + y^2 + z^2 - 1
p = implicit_plot3d(f, (x, -2, 2), (y, -1, 1), (z, -1, 1)) 
K = gaussian_curvature_implicit_surface(f)

c = ColormapTransform(cmsel1, fun=K)
c.apply(p).show(aspect_ratio=(1, 1, 1))

# The previous picture suggest that the highest curvature can be found at the edges of the ellipsoid, whereas the lowest values for the curvature are in the center:

K_high = K(2, 0, 0); K_high

4.0

K_low = K(0, 1, 0); K_low

0.25

# Taking these bounds into account, we can produce a more accurate curvature plot:
c = ColormapTransform(cmsel1, fun=K, bounds=(K_low, K_high))
c.apply(p).show(aspect_ratio=(1, 1, 1))

# This seems to suggest that the Gaussian curvature is quite low over most of the ellipsoid, and increases pretty quickly at the end points.  Let's try to verify this in a more pedestrian way: we take the intersection of the ellipsoid with the positive y-plane, which gives a curve, we compute the Gaussian curvature along this curve, and we plot this.

from numpy import linspace
points = []
for cx in linspace(-2, 2, 20):
    cy = sqrt(1-cx^2/4)
    K_val = K(cx, cy, 0)
    points.append([cx, K_val])

# The curvature indeed increases quickly at the end points
list_plot(points, plotjoined=True)

var('x, y, z')

(x, y, z)

# A saddle surface (negative curvature at the origin)
f = z+x^2-y^2
p = implicit_plot3d(f, (x, -2, 2), (y, -2, 2), (z, -2, 2), plot_points=60, contour=0)

f

x^2 - y^2 + z

K = gaussian_curvature_implicit_surface(f)

# The origin has negative gaussian curvature
K(0, 0, 0)

-4.0

# To have reasonable bounds on the gaussian curvature, we use the fact that the curvature is most negative at the origin, and goes to zero as x, y go infinity.
c = ColormapTransform(cmsel1, fun=K, bounds=(-4, 0))

# Again we see that the curvature falls off pretty quickly away from the origin
c.apply(p)

# A more complicated example, taken from the implicit_plot3d docs

T = RDF(golden_ratio)
f = 2 - (cos(x + T*y) + cos(x - T*y) + cos(y + T*z) + cos(y - T*z) + cos(z - T*x) + cos(z + T*x)) 
r = 4.77
p = implicit_plot3d(f, (x, -r, r), (y, -r, r), (z, -r, r), plot_points=40); p

K = gaussian_curvature_implicit_surface(f) # Very slow!

c = ColormapTransform(cmsel1, fun=K)
c.apply(p)