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
```