如何使用第四变量的取值为Matplotlib绘制的3D表面图上色

实现逻辑 #

Matplotlib的plot_surface接口提供了facecolors参数，支持自定义每个网格面的颜色，不需要绑定z轴的取值。只需将自定义的第四变量归一化到[0,1]区间，再通过色卡映射为对应颜色传入即可实现需求。

完整修改后代码 #

import pylab as pl
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.colors import LinearSegmentedColormap, Normalize
from matplotlib import cm
import numpy as np

def makeData ():
    x = np.arange (-10, 10, 0.1)
    y = np.arange (-10, 10, 0.1)
    xgrid, ygrid = np.meshgrid(x, y)
    zgrid = np.sin (xgrid) * np.sin (ygrid) / (xgrid * ygrid)
    return xgrid, ygrid, zgrid

x, y, z = makeData()

# ----------------------新增第四变量处理逻辑开始----------------------
# 第四变量示例：此处用坐标点到原点的距离作为第四变量，可替换为你实际的第四变量（如复函数的幅角等）
w = np.sqrt(x**2 + y**2)
# 归一化第四变量取值到[0,1]区间
norm = Normalize(vmin=w.min(), vmax=w.max())
# 将归一化后的第四变量映射为对应色卡的颜色值
face_colors = cm.jet(norm(w))
# ----------------------新增第四变量处理逻辑结束----------------------

fig = pl.figure()
axes = Axes3D(fig, auto_add_to_figure=False)
fig.add_axes(axes)

# 替换原plot_surface调用，新增facecolors参数，移除原cmap参数
axes.plot_surface(x, y, z, rstride=4, cstride=4, facecolors=face_colors)

# 新增颜色条，展示第四变量和颜色的对应关系
mappable = cm.ScalarMappable(norm=norm, cmap=cm.jet)
mappable.set_array(w)
fig.colorbar(mappable, ax=axes, label='第四变量取值')

pl.show()

注意事项 #

• 第四变量w的维度必须和x、y、z完全一致，均为相同形状的二维数组。
• 如果是绘制复函数，可将复函数的模值作为z轴变量，幅角作为第四变量w传入即可。
• 可自行替换cm.jet为你需要的其他色卡。

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