Surface Plot of a function B(x,y,z) - matplotlib

Question

I have to plot a surface plot which has axes x,y,z and a colormap set by a function of x,y,z [B(x,y,z)].
I have the coordinate arrays:
x=np.arange(-100,100,1)
y=np.arange(-100,100,1)
z=np.arange(-100,100,1)
Moreover, my to-be-colormap function B(x,y,z) is a 3D array, whose B(x,y,z)[i] elements are the (x,y) coordinates at z.
I have tried something like:
Z,X,Y=np.meshgrid(z,x,y) # Z is the first one since B(x,y,z)[i] are the (x,y) coordinates at z.
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
img = ax.scatter(Z, X, Y, c=B(x,y,z), cmap=plt.hot())
fig.colorbar(img)
plt.show()
However, it unsurprisingly plots dots, which is not what I want. Rather, I need a surface plot.
The figure I have obtained:
The kind of figure I want:
where the colors are determined by B(x,y,z) for my case.

Answer 1

You have to:
use plot_surface to create a surface plot.
your function B(x, y, z) will be used to compute the color parameter, a number assigned to each face of the surface.
the color parameter must be normalized between 0, 1. We use matplotlib's Normalize to achieve that.
then, you create the colors by applying the colormap to the normalized color parameter.
finally, you create the plot.
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib.colors import Normalize
t = np.linspace(0, 2*np.pi)
p = np.linspace(0, 2*np.pi)
t, p = np.meshgrid(t, p)
r1, r2 = 1, 3
x = (r2 + r1 * np.cos(t)) * np.cos(p)
y = (r2 + r1 * np.cos(t)) * np.sin(p)
z = r1 * np.sin(t)
color_param = np.sin(x / 2) * np.cos(y) + z
cmap = cm.jet
norm = Normalize(vmin=color_param.min(), vmax=color_param.max())
norm_color_param = norm(color_param)
colors = cmap(norm_color_param)
fig = plt.figure()
ax = fig.add_subplot(projection="3d")
ax.plot_surface(x, y, z, facecolors=colors)
ax.set_zlim(-4, 4)
plt.show()