Matplotlib 3d plot_surface make edgecolors a function of z-value

Question

How can we make edgecolors of a 3D surface plot a function of the z value?

The following will make the edgecolors uniformly black:

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

x = np.linspace(-1, 1, 20)
y = np.linspace(-1, 1, 20)
x, y = np.meshgrid(x, y)
z = np.sin(x**2 + y**2)
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.plot_surface(x, y, z, cmap=cm.coolwarm, edgecolors='k')
plt.show()

How can we modify this example to:

1. Make the edgecolors white for z < 0.5 and black for z >= 0.5?
2. Make the edgecolors range smoothly from white when z=0 to black when z=1?

Answer 1

I think it's mainly a case of defining the appropriate colour maps. Then you can feed the data into plot_surface() twice - once for the face colours, and a second time just for the edges (edges will overlay the first plot, rendering the final plot).

import matplotlib.pyplot as plt
import matplotlib

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm, colors

x = np.linspace(-1, 1, 20)
y = np.linspace(-1, 1, 20)
x, y = np.meshgrid(x, y)
z = np.sin(x**2 + y**2)

#Custom colour maps:
# Hard white-black transition
white_black = colors.ListedColormap(['white', 'black'])
# Linear transition
white_black_smooth = colors.LinearSegmentedColormap.from_list('Custom', ['white', 'black'])

#Plot using hard the first colourmap
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
#Colour the faces
ax.plot_surface(x, y, z, facecolors=cm.coolwarm(z), shade=False)
#Overlay edge colours (facecolor alpha set to 0)
ax.plot_surface(x, y, z, facecolors=white_black(z), linewidth=2.2, facecolor=(0,0,0,0))
plt.show()

#Plot using the second colourmap
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
#Colour the faces
ax.plot_surface(x, y, z, facecolors=cm.coolwarm(z), shade=False)
#Overlay edge colours (facecolor alpha set to zero)
ax.plot_surface(x, y, z, facecolors=white_black_smooth(z),
                linewidth=2.2, facecolor=(0,0,0,0), shade=False)
plt.show()