[Matplotlib-users] color in plot3d

[Pablo A. <pab...@ua...>]

   Hello!
   I wonder if there is a way to make 3d plots specifying arbitrary
colors, instead of having the color be a function of the height. I was
able to achieve this making minimal changes to the plot_surface method
of Axes3D, adding as an optional keyword argument a function *cfun*
which specifies the color (it specifies a real number that is mapped
into a color by the color map *cmap*). But is there a standard way?
   Regard
Pablo Angulo


from matplotlib.colors import Normalize, colorConverter

def plot_surface(self, X, Y, Z, *args, **kwargs):

    '''

    Create a surface plot.

    By default it will be colored in shades of a solid color,

    but it also supports color mapping by supplying the *cmap*

    argument.

    ==========  ================================================

    Argument    Description

    ==========  ================================================

    *X*, *Y*,   Data values as numpy.arrays

    *Z*

    *rstride*   Array row stride (step size)

    *cstride*   Array column stride (step size)

    *color*     Color of the surface patches

    *cmap*      A colormap for the surface patches.

    *cfun*      The function giving the color

    ==========  ================================================

    '''

    had_data = self.has_data()

    rows, cols = Z.shape

    tX, tY, tZ = np.transpose(X), np.transpose(Y), np.transpose(Z)

    rstride = kwargs.pop('rstride', 10)

    cstride = kwargs.pop('cstride', 10)

    color = kwargs.pop('color', 'b')

    color = np.array(colorConverter.to_rgba(color))

    cmap = kwargs.get('cmap', None)

*    cfun = kwargs.pop('cfun', None)*

    polys = []

    normals = []

    avgz = []

*    if not cfun:*

*        cfun = lambda p:p[2]*

    for rs in np.arange(0, rows-1, rstride):

        for cs in np.arange(0, cols-1, cstride):

            ps = []

            corners = []

            for a, ta in [(X, tX), (Y, tY), (Z, tZ)]:

                ztop = a[rs][cs:min(cols, cs+cstride+1)]

                zleft = ta[min(cols-1, cs+cstride)][rs:min(rows, rs+rstride+1)]

                zbase = a[min(rows-1, rs+rstride)][cs:min(cols, cs+cstride+1):]

                zbase = zbase[::-1]

                zright = ta[cs][rs:min(rows, rs+rstride+1):]

                zright = zright[::-1]

                corners.append([ztop[0], ztop[-1], zbase[0], zbase[-1]])

                z = np.concatenate((ztop, zleft, zbase, zright))

                ps.append(z)

            # The construction leaves the array with duplicate points, which

            # are removed here.

            ps = zip(*ps)

            lastp = np.array([])

            ps2 = []

            avgzsum = 0.0

            for p in ps:

                if p != lastp:

                    ps2.append(p)

                    lastp = p

                    *avgzsum += cfun(p)*

            polys.append(ps2)

            avgz.append(avgzsum / len(ps2))

            v1 = np.array(ps2[0]) - np.array(ps2[1])

            v2 = np.array(ps2[2]) - np.array(ps2[0])

            normals.append(np.cross(v1, v2))

    polyc = art3d.Poly3DCollection(polys, *args, **kwargs)

    if cmap is not None:

        polyc.set_array(np.array(avgz))

        polyc.set_linewidth(0)

    else:

        colors = self._shade_colors(color, normals)

        polyc.set_facecolors(colors)

    self.add_collection(polyc)

    self.auto_scale_xyz(X, Y, Z, had_data)

    return polyc