[Pyplot] 绘制3D曲面+自定义面片颜色

一、背景

使用python+matplotlib实现绘制3D曲面（由多个小面片组成），支持自定义面片颜色；
实现效果如图(a),(b)所示：

(a) . 使用面片法向作为面片颜色

(b) . 使用默认jet类型colormap作为面片颜色

二、代码

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

# 生成曲面各点 (x,y,z)
# 曲面点位置：(x=x, y=y, z=x*y)
# 曲面点法向: (y, x, 1.0)

x = np.linspace(-0.5, 0.5, 10)
y = np.linspace(-0.5, 0.5, 10)
z = np.zeros((len(x), len(y)))

nor = np.zeros((len(x), len(y),3))
nor_x_min = np.float64('inf')
nor_x_max = -1.0*np.float64('inf')
nor_y_min = np.float64('inf')
nor_y_max = -1.0*np.float64('inf')

for i in range(len(x)):
    for j in range(len(y)):
        z[i][j] = x[i]*y[j]
        nor[i][j] = (y[j], x[i], 1.0)
        nor_x_min = np.fmin(nor_x_min, nor[i][j][0])
        nor_x_max = np.fmax(nor_x_max, nor[i][j][0])
        
        nor_y_min = np.fmin(nor_y_min, nor[i][j][1])
        nor_y_max = np.fmax(nor_y_max, nor[i][j][1])

# 绘制曲面 
# 曲面由多个面片组成 各面片颜色根据法向确定
x, y = np.meshgrid(x, y)
fig = plt.figure()
ax = fig.add_subplot(
    111, projection='3d')

f = np.zeros((len(z), len(z[0])), tuple)
for i in range(len(f)):
    for j in range(len(f[i])):
        f[i][j] = ((nor[i][j][0]-nor_x_min)/(nor_x_max-nor_x_min), (nor[i][j][1]-nor_y_min)/(nor_y_max-nor_y_min), 1.0)
# 使用面片法向作为面片颜色
ax.plot_surface(x, y, z, facecolors=f)
# 使用默认jet类型colormap作为面片颜色
# ax.plot_surface(x, y, z, cmap=cm.jet)
plt.show()