Shade the surface and contour parts

Question

I want to shade the surface and the contours of a specific function based on some constraint in the domain of the function. So far I have the following and I want to improve it.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import LinearLocator, FormatStrFormatter

plt.figure(figsize=(8, 6))
plt.axes(projection="3d")
xdata = np.linspace(-3, 3, 20000)
ydata = np.linspace(-3, 3, 20000)
X, Y = np.meshgrid(xdata, ydata)

Z1 = X ** 2 + Y ** 2
Z2 = Z1.copy()
Z3 = Z1.copy()
Z1[np.multiply(X, Y) > 3] = np.nan 
Z2[np.multiply(X, Y) <= 3] = np.nan
Z3[np.multiply(X, Y) == 3] = np.nan

ax3d = plt.axes(projection='3d')
ax3d.plot_surface(X, Y, Z1,  cmap='Greys', antialiased=True, vmin=-np.nanmin(Z1), vmax=np.nanmax(Z1))
ax3d.plot_surface(X, Y, Z2, cmap='YlGnBu', antialiased=True, vmin=-np.nanmin(Z2), vmax=np.nanmax(Z2))

ax3d.contourf(X, Y, Z1, zdir='z', offset=0, cmap='Greys')
ax3d.contourf(X, Y, Z2, zdir='z', offset=0, cmap='Greys')

ax3d.set_title('Surface Plot in Matplotlib')
ax3d.set_xlabel('X')
ax3d.set_ylabel('Y')
ax3d.set_zlabel('Z')

plt.show()

Could you please someone help to solve the following problems:

1. The surface is over-imposed by the contour surface.
2. There are some gaps in the surface.
3. The contours of the two constraints are not continuous.
4. Is it possible to plot a line in the border of the two surfaces and contours?

Any help is highly appreciated.

Answer 1

The code below makes the following changes:

• creating a custom colormap combining the two existing colormaps
• using a TwoSlopeNorm to have the separation at z=3
• setting antialiased=False (otherwise matplotlib creates a plot of antialiased lines instead of polygons)
• xdata and ydata with 300 steps
• using rstride=1, cstride=1 so every x and every y will be considered; this makes the surface smoother (but takes more time)
• calling plt.axes(...) only once to prevent a dummy subplot
• calling contourf before plot_surface; due to the painter's algorithm, matplotlib only minimally supports 3D overlaps

import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm, ListedColormap
import numpy as np

xdata = np.linspace(-3, 3, 300)
ydata = np.linspace(-3, 3, 300)
X, Y = np.meshgrid(xdata, ydata)

Z1 = X ** 2 + Y ** 2

cmap1 = plt.get_cmap('Greys')
cmap2 = plt.get_cmap('YlGnBu')
cmap = ListedColormap(np.r_[cmap1(np.linspace(0, 1, 128)), cmap2(np.linspace(0, 1, 128))])
norm = TwoSlopeNorm(vmin=np.nanmin(Z1), vmax=np.nanmax(Z1), vcenter=3)

plt.figure(figsize=(8, 6))
ax3d = plt.axes(projection='3d')

ax3d.contourf(X, Y, Z1, zdir='z', offset=0, cmap=cmap, norm=norm)
ax3d.plot_surface(X, Y, Z1, cmap=cmap, antialiased=False, norm=norm, rstride=1, cstride=1)

ax3d.set_title('Surface Plot in Matplotlib')
ax3d.set_xlabel('X')
ax3d.set_ylabel('Y')
ax3d.set_zlabel('Z')

plt.show()

Answer 2

xdata = np.linspace(-3, 3, 1000)
ydata = np.linspace(-3, 3, 1000)
X, Y = np.meshgrid(xdata, ydata)

Z1 = X ** 2 + Y ** 2
Z2 = Z1.copy()
Z3 = Z1.copy()

Z2[np.multiply(X, Y) <= 3] = np.nan
Z3[np.multiply(X, Y) == 3] = np.nan


plt.figure(figsize=(8, 6))
ax3d = plt.axes(projection='3d')

ax3d.contourf(X, Y, Z1, zdir='z', offset=0, cmap='Greys')
ax3d.contourf(X, Y, Z2, zdir='z', offset=0, cmap='YlGnBu')


ax3d.plot_surface(X, Y, Z1,  cmap='Greys', antialiased=True, vmin=-np.nanmin(Z1), vmax=np.nanmax(Z1), alpha=.5)
ax3d.plot_surface(X, Y, Z2, cmap='YlGnBu', antialiased=True, vmin=-np.nanmin(Z2), vmax=np.nanmax(Z2), alpha=.5)

ax3d.set_title('Surface Plot in Matplotlib')
ax3d.set_xlabel('X')
ax3d.set_ylabel('Y')
ax3d.set_zlabel('Z')

plt.show()