Using x for the value of a colormap/colorbar instead of z

Question

I have a function x = F(z,y). I want to map the x value to the color. How do I do this?

Partial code:

z = r * np.cos(theta)
y = r * np.sin(theta)
x = V(r, V_avg,R)

# Create a 3D plot
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# Plot the surface
surface = ax.plot_surface(x,y,z, cmap='viridis')

# Add a color bar which maps values to colors
fig.colorbar(surface, ax=ax, label='Paraboloid Function Value')

# Set labels
ax.set_xlabel('X-axis')
ax.set_ylabel('Y-axis')
ax.set_zlabel('Z-axis')

# Show the plot
plt.show()

Output:

Answer 1

You could map colors from your color map to x-values, creating an array of colors like in this example, and use those colors to plot the surface together with its color bar:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm

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

# make some paraboloid
radius = 4
# make a grid in cylindrical coordinates
r = np.linspace(0, radius, 100)
theta = np.linspace(0, 2 * np.pi, 100)
R, THETA = np.meshgrid(r, theta)

# convert to Cartesian coordinates
z_grid, y_grid = R * np.cos(THETA), R * np.sin(THETA)

x_grid = -1 * ((z_grid / 2) ** 2 + (y_grid / 2) ** 2)

x_min = np.min(x_grid)
x_max = np.max(x_grid)
x_range = x_max - x_min
x_values = x_grid.ravel()

# make colors
cmap = cm.viridis
colors = cmap((x_grid - x_min) / x_range)

# plot the paraboloid
surface = ax.plot_surface(x_grid, y_grid, z_grid, facecolors=colors)
# add a color bar
fig.colorbar(surface, ax=ax, label='Paraboloid Function Value', values=sorted(x_values))

# Set labels
ax.set_xlabel('X-axis')
ax.set_ylabel('Y-axis')
ax.set_zlabel('Z-axis')

plt.show()

The result: