pythonmatplotlibmplot3dcontourf

Matplotlib plot contourf on 3d surface


I am trying to use the colormap feature of a 3d-surface plot in matplotlib to color the surface based on values from another array instead of the z-values. The surface plot is created and displayed as follows:

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

def gauss(x, y, w_0):
    r = np.sqrt(x**2 + y**2)
    return np.exp(-2*r**2 / w_0**2)


x = np.linspace(-100, 100, 100)
y = np.linspace(-100, 100, 100)
X, Y = np.meshgrid(x, y)
Z = gauss(X, Y, 50)
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.plot_surface(X, Y, Z, cmap='jet')

Now instead of coloring based on elevation of the 3d-surface, I am looking to supply the color data for the surface in form of another array, here as an example a random one:

color_data = np.random.uniform(0, 1, size=(Z.shape))

However, I did not find a solution to colorize the 3d-surface based on those values. Ideally, it would look like a contourf plot in 3d, just on the 3d surface.


Solution

  • You can use matplotlib.colors.from_levels_and_colors to obtain a colormap and normalization, then apply those to the values to be colormapped.

    import numpy as np
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D
    import matplotlib.colors
    
    x = np.linspace(-100, 100, 101)
    y = np.linspace(-100, 100, 101)
    X, Y = np.meshgrid(x, y)
    Z = np.exp(-2*np.sqrt(X**2 + Y**2)**2 / 50**2)
    
    c = X+50*np.cos(Y/20)  # values to be colormapped
    N = 11                 # Number of level (edges) 
    levels = np.linspace(-150,150,N)
    colors = plt.cm.get_cmap("RdYlGn", N-1)(np.arange(N-1))
    cmap, norm = matplotlib.colors.from_levels_and_colors(levels, colors)
    color_vals = cmap(norm(c))
    
    fig = plt.figure()
    ax = fig.add_subplot(projection='3d')
    ax.plot_surface(X, Y, Z, facecolors=color_vals, rstride=1, cstride=1)
    plt.show()
    

    enter image description here