Matplotlib: Set cmap in plot_surface to x and y-axes

Question

How can I set the colormap in relation to the radius of the figure?

And how can I close the ends of the cylinder (on the element, not the top and bottom bases)?

My script:

import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from math import sin, cos, pi

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

h, w = 60,30
znew = np.random.randint(low=90, high=110, size=(60,30))
theta = np.linspace(0,2*pi, h)
Z = np.linspace(0,1,w)

Z,theta = np.meshgrid(Z, theta)

R = 1
X = (R*np.cos(theta))*znew
Y = (R*np.sin(theta))*znew


ax1 = ax.plot_surface(X,Y,Z,linewidth = 0, cmap="coolwarm",              
vmin= 80,vmax=130, shade = True, alpha = 0.75)

fig.colorbar(ax1, shrink=0.9, aspect=5)

plt.show()

Answer 1

First you need to use the facecolors keyword argument of plot_surface to draw your surface with arbitrary (non-Z-based) colours. You have to pass an explicit RGBA colour four each point, which means we need to sample a colormap object with the keys given by the radius at every point. Finally, this will break the mappable property of the resulting surface, so we will have to construct the colorbar by manually telling it to use our radii for colours:

import numpy as np
from matplotlib import pyplot as plt
import matplotlib.cm as cm
from matplotlib.colors import Normalize
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

h, w = 60,30
#znew = np.random.randint(low=90, high=110, size=(h,w))
theta = np.linspace(0,2*np.pi, h)
Z = np.linspace(0,1,w)

Z,theta = np.meshgrid(Z, theta)
znew = 100 + 10*np.cos(theta/2)*np.cos(2*Z*np.pi)

R = 1
X = (R*np.cos(theta))*znew
Y = (R*np.sin(theta))*znew
true_radius = np.sqrt(X**2 + Y**2)
norm = Normalize()
colors = norm(true_radius) # auto-adjust true radius into [0,1] for color mapping

cmap = cm.get_cmap("coolwarm")
ax.plot_surface(X, Y, Z, linewidth=0, facecolors=cmap(colors), shade=True, alpha=0.75)

# the surface is not mappable, we need to handle the colorbar manually
mappable = cm.ScalarMappable(cmap=cmap)
mappable.set_array(colors)
fig.colorbar(mappable, shrink=0.9, aspect=5)

plt.show()

Note that I changed the radii to something smooth for a less chaotic-looking result. The true_radius arary contains the actual radii in data units, which after normalization becomes colors (essentially colors = (true_radius - true_radius.min())/true_radius.ptp()).

The result:

Finally, note that I generated the radii such that the cylinder doesn't close seamlessly. This mimicks your random example input. There's nothing you can do about this as long as the radii are not 2π-periodic in theta. This has nothing to do with visualization, this is geometry.