Python3.10中如何为Matplotlib的plot_trisurf指定各三角面颜色？

你需要在Python 3.10环境下，给plot_trisurf生成的每个三角面设置不同颜色。现有示例代码生成2个三角面，尝试传入颜色列表color = [[1, 0, 0], [0 ,0 ,1]]时触发错误RGBA sequence should have length 3 or 4，自定义颜色映射后也不清楚如何为每个面指定索引。

示例代码：

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

X = [1, 3, 2, 4]
Y = [1, 1, 2, 2]
Z = [2, 2, 2, 2]

ax.scatter(X, Y, Z, s = 0)
ax.plot_trisurf(X, Y, Z)

ax.set_xlim(0, 5)
ax.set_ylim(0, 3)
ax.set_zlim(1.9, 2.1)

ax.set_xlabel('X Axis')
ax.set_ylabel('Y Axis')
ax.set_zlabel('Z Axis')

plt.show()

方法一：直接通过facecolors参数指定每个面的颜色 #

plot_trisurf的color参数仅支持单一颜色值，而**facecolors参数**可以接收一个颜色列表，每个元素对应一个三角面的颜色。你需要将颜色列表转换为N×3或N×4的数组（N为三角面数量），示例修改如下：

import matplotlib.pyplot as plt
import numpy as np

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

X = [1, 3, 2, 4]
Y = [1, 1, 2, 2]
Z = [2, 2, 2, 2]

# 定义每个三角面的颜色，2个面对应2组RGB值
face_colors = np.array([[1, 0, 0], [0, 0, 1]])

ax.scatter(X, Y, Z, s=0)
# 使用facecolors参数传入颜色列表
surf = ax.plot_trisurf(X, Y, Z, facecolors=face_colors)

ax.set_xlim(0, 5)
ax.set_ylim(0, 3)
ax.set_zlim(1.9, 2.1)

ax.set_xlabel('X Axis')
ax.set_ylabel('Y Axis')
ax.set_zlabel('Z Axis')

plt.show()

方法二：自定义颜色映射并为每个面指定索引 #

如果你想用颜色映射，可以先获取plot_trisurf生成的三角面数量，然后为每个面分配一个索引值，再通过颜色映射转换为对应颜色：

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

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

X = [1, 3, 2, 4]
Y = [1, 1, 2, 2]
Z = [2, 2, 2, 2]

# 自定义颜色映射
colors = [(1, 0, 0), (0, 0, 1)]
cm = LinearSegmentedColormap.from_list("Custom", colors)

# 先绘制一次获取三角面数量
surf_temp = ax.plot_trisurf(X, Y, Z, alpha=0)
tri_count = len(surf_temp.triangles)
surf_temp.remove()

# 为每个三角面生成索引值，0到1之间均匀分布
face_indices = np.linspace(0, 1, tri_count)
# 通过颜色映射获取对应颜色
face_colors = cm(face_indices)

# 绘制最终的曲面
surf = ax.plot_trisurf(X, Y, Z, facecolors=face_colors)

ax.set_xlim(0, 5)
ax.set_ylim(0, 3)
ax.set_zlim(1.9, 2.1)

ax.set_xlabel('X Axis')
ax.set_ylabel('Y Axis')
ax.set_zlabel('Z Axis')

plt.show()

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