pythonmatplotlib3d

MatPlotlib colobar with wrong range in 3D surface


I'm trying to plot a value around the unit sphere using surface plot and facecolors in matplotlib, but my colorbar shows the normalized values instead of the real values. How can I fix this so the colorbar has the right range?

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib import cm

fig, ax = plt.subplots(subplot_kw={"projection": "3d"}, figsize = (10, 14))

# Make data
n_points = 500
r = 1
u = np.linspace(0, 2 * np.pi, n_points)
v = np.linspace(0, np.pi, n_points)
x = r * np.outer(np.cos(u), np.sin(v))
y = r * np.outer(np.sin(u), np.sin(v))
z = r * np.outer(np.ones(np.size(u)), np.cos(v))

ax.plot_wireframe(x, y, z, color="grey", alpha = 0.1)

data = np.random.uniform(0.2, 0.5, n_points)
heatmap = np.array(np.meshgrid(data, data))[1]

ax.set_aspect("equal")

ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')


colormap = cm.viridis
normaliser = mpl.colors.Normalize(vmin=np.min(heatmap), vmax=np.max(heatmap))

print(np.min(heatmap))
print(np.max(heatmap))

surf = ax.plot_surface(
    x, y, z, 
    facecolors=colormap(normaliser(heatmap)), shade=False)
    

fig.colorbar(surf, shrink=0.5, aspect=10, label="Singlet yield", pad = 0.05, norm = normaliser)

plt.show()

This outputs 0.20009725794516225 and 0.49936395079063567 as min and max in the prints, but you can see the range of the colorbar is 0 to 1 in the following image. enter image description here

How can I fix this issue and make it so the colorbar has the appropriate colors?


Solution

  • The colorbar function itself doesn't have a norm argument according to the documentation for this function. For minimal alteration, you can pass a matplotlib.cm.ScalarMappable as the first argument of the colorbar call and it works as expected (presuming you also pass the appropriate ax argument).

    Here is a fully runnable code demonstrating this:

    import numpy as np
    import matplotlib.pyplot as plt
    import matplotlib as mpl
    from matplotlib import cm
    
    fig, ax = plt.subplots(subplot_kw={"projection": "3d"}, figsize=(10, 14))
    
    # Make data
    n_points = 500
    r = 1
    u = np.linspace(0, 2 * np.pi, n_points)
    v = np.linspace(0, np.pi, n_points)
    x = r * np.outer(np.cos(u), np.sin(v))
    y = r * np.outer(np.sin(u), np.sin(v))
    z = r * np.outer(np.ones(np.size(u)), np.cos(v))
    
    ax.plot_wireframe(x, y, z, color="grey", alpha=0.1)
    
    data = np.random.uniform(0.2, 0.5, n_points)
    heatmap = np.array(np.meshgrid(data, data))[1]
    
    ax.set_aspect("equal")
    
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    
    colormap = cm.viridis
    normaliser = mpl.colors.Normalize(vmin=np.min(heatmap), vmax=np.max(heatmap))
    
    print(np.min(heatmap))
    print(np.max(heatmap))
    
    surf = ax.plot_surface(
        x, y, z,
        facecolors=colormap(normaliser(heatmap)), shade=False)
    
    mappable = cm.ScalarMappable(norm=normaliser, cmap=colormap)
    fig.colorbar(mappable, ax=ax, shrink=0.5, aspect=10, label="Singlet yield", pad=0.05)
    
    plt.show()
    
    

    Here is the output it generates: colourbar plot with correct bar extent