Conditional colormap based on some threshold value for 3D surfaces using Matplotlib

Question

I have some three dimensional dataset where each data point is of form (x,y,z). I want to make a 3D plot where z is expressed as a function of both x and y, in the form of a 3D surface. For that, I am using the following code:

import matplotlib.pyplot as plt

fig = plt.figure()
ax = fig.add_subplot(1,1,1,projection='3d')
ax.plot_surface(X,Y,Z,rstride=1,cstride=1,cmap='viridis',edgecolor='none')
plt.show()

Now, I have a baseline point (x0,y0,z0), and I want to have the surface's colormap to be defined as a function of z0. More specifically, I want a different colormap for z<z0 and z>z0, for example an intensifying red colormap for z<z0 and a green one for z>z0.

How can I achieve that using matplotlib?

Answer 1

I found a solution to my case based on this answer: Defining the midpoint of a colormap in matplotlib. We define a colour map norm with z0 as the centre, then we choose the colour scheme RdYlGn which goes from green to red, where X, Y, Z are numpy.ndarrays:

minz = min(Z.flatten()) 
maxz = max(Z.flatten())
cnorm = plt.colors.TwoSlopeNorm(vmin=minz, vcenter=z0, vmax=maxz)

ax.plot_surface(X,Y,Z,rstride=1,cstride=1,cmap='RdYlGn',norm=cnorm,edgecolor='none')