This is perfectly straightforward if you spend the time learning how matplotlib (and 3d axes in particular) work:
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from mpl_toolkits.mplot3d import Axes3D
# compute data to plot
r, theta = np.mgrid[1:16, -2*np.pi:2*np.pi:50j]
z = r * np.exp(1j*theta)
w = np.sqrt(r) * np.exp(1j*theta/2)
# plot data
fig = plt.figure()
for plot_index in [1, 2]:
if plot_index == 1:
z_data, c_data = w.real, w.imag
z_comp, c_comp = 'Re', 'Im'
else:
z_data, c_data = w.imag, w.real
z_comp, c_comp = 'Im', 'Re'
c_data = (c_data - c_data.min()) / c_data.ptp()
colors = cm.viridis(c_data)
ax = fig.add_subplot(f'12{plot_index}', projection='3d')
surf = ax.plot_surface(z.real, z.imag, z_data, facecolors=colors,
clim=[z_data.min(), z_data.max()])
ax.set_xlabel('$Re z$')
ax.set_ylabel('$Im z$')
ax.set_zlabel(f'${z_comp} w$')
cb = plt.colorbar(surf, ax=ax)
cb.set_label(f'${c_comp} w$')
plt.show()
The result:
Some things that should be noted:
- Viridis colormap is good, jet is bad.
- In general there could be rendering issues with complex (interlocking) 3d geometries, because matplotlib has a 2d renderer. Fortunately, in this case the dataset is tightly coupled enough that this doesn't seem to happen, even if you rotate around the figure interactively. (But if you were to plot two intersecting surfaces together, things would probably be different.)
- One might want to enable latex rendering of labels to make the result extra crispy.
- The shading looks a lot better if you use the default option of colouring according to the z component of the data.
If we also want to port the second part of my MATLAB answer you will have to use a trick to stitch together the two branches of the function (which, as I said, is necessary to render interlocking surfaces properly). For the specific example in the above code this will not give you perfect results, since both branches themselves contain discontinuities in the imaginary part, so regardless of our efforts in rendering the two surfaces nicely, the result will look a bit bad:
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from mpl_toolkits.mplot3d import Axes3D
# compute data to plot
r0 = 15
re, im = np.mgrid[-r0:r0:31j, -r0:r0:31j]
z = re + 1j*im
r, theta = abs(z), np.angle(z)
w1 = np.sqrt(r) * np.exp(1j*theta/2) # first branch
w2 = np.sqrt(r) * np.exp(1j*(theta + 2*np.pi)/2) # second branch
# plot data
fig = plt.figure()
for plot_index in [1, 2]:
# construct transparent bridge
re_bridge = np.vstack([re[-1, :], re[0, :]])
im_bridge = np.vstack([im[-1, :], im[0, :]])
c_bridge = np.full(re_bridge.shape + (4,), [1, 1, 1, 0]) # 0% opacity
re_surf = np.vstack([re, re_bridge, re])
im_surf = np.vstack([im, im_bridge, im])
w12 = np.array([w1, w2])
if plot_index == 1:
z_comp, c_comp = 'Re', 'Im'
z12, c12 = w12.real, w12.imag
else:
z_comp, c_comp = 'Im', 'Re'
z12, c12 = w12.imag, w12.real
color_arrays = cm.viridis((c12 - c12.min()) / c12.ptp())
z1,z2 = z12
c1,c2 = color_arrays
z_bridge = np.vstack([z1[-1, :], z2[0, :]])
z_surf = np.vstack([z1, z_bridge, z2])
c_surf = np.vstack([c1, c_bridge, c2])
ax = fig.add_subplot(f'12{plot_index}', projection='3d')
surf = ax.plot_surface(re_surf, im_surf, z_surf, facecolors=c_surf,
clim=[c12.min(), c12.max()],
rstride=1, cstride=1)
ax.set_xlabel('$Re z$')
ax.set_ylabel('$Im z$')
ax.set_zlabel(f'${z_comp} w$')
cb = plt.colorbar(surf, ax=ax)
cb.set_label(f'${c_comp} w$')
plt.show()
The ugly jump in the right figure might be fixed with a lot of work, but it won't be easy: it's an actual discontinuity in both surface datasets occuring at negative real arguments. Since your actual problem is probably more like this, you will probably not need to face this issue, and you can use the above stitching (bridging) trick to combine your surfaces.
