如何修改plot_surface()的默认cmap行为，使颜色渐变由x值决定而非z值？

默认情况下，plot_surface()会根据z轴数值生成颜色渐变。要改成基于x值（也就是你的dist变量）配色，需要手动生成颜色数组并传给facecolors参数，具体实现代码如下：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize

n = 20
s = 0.05

y = np.linspace(-s, s, n)
z = np.linspace(-s, s, n)
Y, Z = np.meshgrid(y, z)
dist = (Y**6 * Z**2) ** 0.1

# 把x值（dist）归一化到[0,1]区间，适配颜色映射的输入范围
norm = Normalize(vmin=dist.min(), vmax=dist.max())
cmap = plt.cm.viridis
# 根据归一化后的x值生成RGBA颜色数组
face_colors = cmap(norm(dist))

fig, ax = plt.subplots(subplot_kw=dict(projection='3d'))
# 用facecolors指定曲面颜色，替代默认的z值映射
ax.plot_surface(dist, Y, Z, facecolors=face_colors)
ax.set(aspect='equal', xlabel='x', ylabel='y', zlabel='z')

# 手动添加对应x值的颜色条
mappable = plt.cm.ScalarMappable(norm=norm, cmap=cmap)
mappable.set_array(dist)
fig.colorbar(mappable, ax=ax, label='x 值')

plt.show()

关键说明 #

• Normalize负责把x值的实际范围缩放到颜色映射需要的[0,1]区间，保证颜色渐变均匀
• 调用cmap对象传入归一化后的x值，直接生成每个曲面网格对应的RGBA颜色
• 颜色条需要手动创建ScalarMappable关联x值的归一化规则，这样颜色条才能正确对应x值的范围

内容的提问来源于stack exchange，提问作者mins