如何绘制以函数B(x,y,z)为色标的3D曲面图？

问题说明 #

需要绘制以x、y、z为坐标轴的3D曲面图，色标由3D数组B(x,y,z)控制（其中B[:, :, i]对应z=i处的(x,y)坐标点颜色）。现有坐标数组：

x = np.arange(-100, 100, 1)
y = np.arange(-100, 100, 1)
z = np.arange(-100, 100, 1)

原代码使用scatter得到了散点图，而非预期的曲面效果：

Z,X,Y=np.meshgrid(z,x,y) # Z是第一个参数，因为B(x,y,z)[i]对应z处的(x,y)坐标。

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

img = ax.scatter(Z, X, Y, c=B(x,y,z), cmap=plt.hot())
fig.colorbar(img)
plt.show()

解决方案 #

要实现曲面效果，需针对每个z值绘制平行于x-y平面的曲面，并将对应z层的B数组作为颜色映射。具体实现如下：

代码示例 #

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

# 替换为你的实际3D色标数组B，形状需为(len(x), len(y), len(z))
x = np.arange(-100, 100, 1)
y = np.arange(-100, 100, 1)
z = np.arange(-100, 100, 1)
B = np.random.rand(len(x), len(y), len(z))  # 模拟示例数据

# 生成x-y平面的2D网格，所有z层共用该网格
X, Y = np.meshgrid(x, y)

# 统一颜色归一化范围，确保所有曲面使用相同色标
norm = plt.Normalize(B.min(), B.max())
cmap = cm.hot()

fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# 循环绘制每个z层的曲面
for i, z_val in enumerate(z):
    # 将当前z层的B数据转换为RGBA颜色格式
    face_colors = cmap(norm(B[:, :, i]))
    # 绘制z=z_val处的x-y平面曲面
    ax.plot_surface(X, Y, np.full_like(X, z_val), 
                    facecolors=face_colors, 
                    rstride=1, cstride=1, 
                    shade=False)

# 添加统一色标
sm = cm.ScalarMappable(norm=norm, cmap=cmap)
sm.set_array([])
fig.colorbar(sm, ax=ax, label='B(x,y,z)')

# 设置坐标轴标签
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')

plt.show()

关键说明 #

• X, Y = np.meshgrid(x, y)：生成x-y平面的2D网格，所有z层复用该网格，避免冗余计算
• norm = plt.Normalize(B.min(), B.max())：统一色标的取值范围，确保不同z层的颜色映射一致
• np.full_like(X, z_val)：生成与X/Y同形状的数组，值为当前z值，作为曲面的z坐标
• facecolors=face_colors：将当前z层的B数据转换为RGBA颜色，直接赋予曲面
• rstride=1, cstride=1：控制曲面网格的密度，值越小曲面越平滑（计算量相应增大）
• shade=False：关闭光照阴影，保证颜色完全由B数据决定

若需求是绘制连续的三维曲面（非分层平面），需明确z与x、y的函数关系（如z=f(x,y)），此时可直接使用plot_surface(X, Y, Z, facecolors=...)，其中Z是x-y对应的z值数组，facecolors对应每个(x,y,z)点的B值。

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