Matplotlib画三维圆柱应力云图，应力值与到圆柱中心的距离有关

最近想用Matplotlib画一个圆柱的应力云图，应力值与到中心的距离相关，可以是线性关系也可以定义非线性关系，最好能画出一个rainbow色带的图，在网上找了好久都没有找到怎么画，最后找到了voxels函数，画体素的，终于能画一个圆柱了，但是怎么表示颜色呢，没有办法，只能找了rainbow色带的rgb值，然后自己写了一个函数来定义应力的颜色。话不多说，直接上代码，欢迎大家批评指正。

先放图：

代码附上： 

import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.ticker as ticker

def map_rate(X:list, to_min:float, to_max:float)->list:
    """区间映射
    
    Attribute:
    
    - X: 需要映射的列表
    - to_min: 要映射到的最小值
    - to_max: 要映射到的最大值
    """
    x_min=min(X)
    x_max=max(X)
    return list([int(round(to_min+((to_max - to_min) / (x_max - x_min)) * (i - x_min),1)) for i in X])
# rainbow色带
def rainbow(x):
    # rainbow色带
    data = [(125,0,255),(99,0,255),(80,0,255),(70,0,255),(60,0,255),(50,0,255),(40,0,255),(30,0,255),(20,0,255),(10,0,255),(4,0,255) ,(0,2,255) ,(0,7,255) ,(0,11,255) ,(0,16,255) 
            ,(0,20,255) ,(0,25,255) ,(0,29,255) ,(0,34,255) ,(0,38,255) ,(0,43,255) ,(0,47,255) ,(0,52,255) ,(0,56,255) ,(0,61,255) ,(0,65,255) 
            ,(0,70,255) ,(0,74,255) ,(0,79,255) ,(0,83,255) ,(0,88,255) ,(0,92,255) ,(0,97,255) ,(0,101,255) ,(0,106,255) ,(0,110,255) ,(0,115,255) ,(0,119,255) 
            ,(0,124,255) ,(0,128,255) ,(0,133,255) ,(0,137,255) ,(0,142,255) ,(0,146,255) ,(0,151,255) ,(0,155,255) ,(0,160,255) ,(0,164,255) ,(0,169,255) ,(0,173,255) 
            ,(0,178,255) ,(0,182,255) ,(0,187,255) ,(0,191,255) ,(0,196,255) ,(0,200,255) ,(0,205,255) ,(0,209,255) ,(0,214,255) ,(0,218,255) ,(0,223,255) ,(0,227,255) 
            ,(0,232,255) ,(0,236,255) ,(0,241,255) ,(0,245,255) ,(0,250,255),(0,255,255),(0,255,245),(0,255,236) 
            ,(0,255,226),(0,255,217),(0,255,208),(0,255,198),(0,255,189),(0,255,179) 
            ,(0,255,175),(0,255,165),(0,255,156),(0,255,146),(0,255,137),(0,255,128) 
            ,(0,255,118),(0,255,109),(0,255,99),(0,255,90),(0,255,76),(0,255,66) 
            ,(0,255,62) ,(0,255,57),(0,255,47),(0,255,38),(0,255,29),(0,255,19),(0,255,10) 
            ,(0,255,0),(7,255,0),(16,255,0),(25,255,0),(34,255,0),(43,255,0) ,(52,255,0),(61,255,0),(70,255,0),(79,255,0),(88,255,0),(97,255,0) ,(106,255,0),(115,255,0),(124,255,0),(133,255,0) ,(142,255,0),(151,255,0) 
            ,(160,255,0),(169,255,0),(178,255,0),(187,255,0),(200,255,0),(209,255,0) ,(218,255,0),(227,255,0),(236,255,0),(245,255,0),(254,255,0) ,(255,252,0) ,(255,247,0) ,(255,243,0) 
            ,(255,238,0) ,(255,234,0) ,(255,229,0) ,(255,225,0) ,(255,220,0) ,(255,216,0) ,(255,211,0) ,(255,206,0) ,(255,202,0) ,(255,197,0) ,(255,193,0) ,(255,188,0) 
            ,(255,184,0) ,(255,179,0) ,(255,175,0) ,(255,170,0) ,(255,165,0) ,(255,161,0) ,(255,156,0) ,(255,152,0) ,(255,147,0) ,(255,143,0) ,(255,138,0) ,(255,134,0) 
            ,(255,129,0) ,(255,124,0) ,(255,120,0) ,(255,115,0) ,(255,111,0) ,(255,106,0) ,(255,102,0) ,(255,97,0) ,(255,93,0) ,(255,88,0) ,(255,83,0) ,(255,79,0) 
            ,(255,74,0) ,(255,70,0) ,(255,65,0) ,(255,61,0) ,(255,56,0) ,(255,52,0) ,(255,47,0) ,(255,42,0) ,(255,38,0) ,(255,33,0) ,(255,29,0) ,(255,24,0) ,(255,20,0)] 
    co = map_rate(x,0,175)
    return np.array(list(data[i] for i in co))
# 求中点
def midpoints(x):
   sl = ()
   for i in range(x.ndim):
       x = (x[sl + np.index_exp[:-1]] + x[sl + np.index_exp[1:]]) / 2.0
       sl += np.index_exp[:]
   return x
# 归一化函数
def normalize(data):
    mx = np.max(data)*np.ones(data.shape)
    mn = np.min(data)*np.ones(data.shape)
    return (data-mn)/(mx-mn)
# 定义应力与半径的关系
def Mises(r):
    return np.round(r*2,2) #计算von mises应力，并保留小数点后两位


# 定义圆柱的参数
R = 10 #圆柱的半径
H = 10 #圆柱的高
# 网格点数量
nr = 19j       # 沿半径分几层
ntheta = 25j  # >=4
nh = 9j
# 转换坐标系，并求中点
r, theta, z = np.mgrid[0:R:nr, 0:np.pi*2:ntheta, 0:H:nh]
x = r*np.cos(theta)
y = r*np.sin(theta)

rc, thetac, zc = midpoints(r), midpoints(theta), midpoints(z)

# 填充网格
a,b,c = rc.shape
rr = list(rc[:,0,0])
sphere = np.zeros((a,b,c)) == 0

# 设置颜色
hsv = np.zeros(sphere.shape + (3,))
r_color1 = rainbow(rr)
r_color2 = normalize(r_color1)
rgb_r = r_color2[:,0]
rgb_g = r_color2[:,1]
rgb_b = r_color2[:,2]
for i in range(a):
    hsv[i,..., 0] = rgb_r[i]*np.ones((b,c))
    hsv[i,..., 1] = rgb_g[i]*np.ones((b,c))
    hsv[i,..., 2] = rgb_b[i]*np.ones((b,c))

# 求应力
mises_r = np.linspace(0,R,a)
mises = Mises(mises_r)
# 画图
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.voxels(x, y, z, sphere,
        facecolors=hsv,
        edgecolors=np.clip(2*hsv - 0.5, 0, 1),
        linewidth=0.5)
ax.set_xlabel('mm')
ax.set_ylabel('mm')
ax.set_zlabel('mm')
# 画 colorbar
ax1 = fig.add_axes([0.02, 0.10, 0.03, 0.80])
cmap = mpl.cm.rainbow
norm = mpl.colors.Normalize(vmin=min(mises), vmax=max(mises))
cb = mpl.colorbar.ColorbarBase(ax1, cmap=cmap,
                                norm=norm,
                                orientation='vertical')   

tick_locator = ticker.MaxNLocator(nbins=len(r))  # colorbar上的刻度值个数                         
cb.locator = tick_locator
cb.set_ticks(mises)
cb.update_ticks()
cb.set_label('von Mises')          # colorbar标签
plt.show()