python matplotlib 三维实体圆柱图

news2025/2/24 2:57:02

环境
python:python-3.12.0-amd64

包:
matplotlib 3.8.2

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 = 9000  # 圆柱的半径
H = 10000  # 圆柱的高
# 网格点数量
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()

注:

自定义颜色名称

'Accent', 'Accent_r', 'Blues', 'Blues_r', 'BrBG', 'BrBG_r', 'BuGn', 'BuGn_r', 'BuPu', 'BuPu_r',
'CMRmap', 'CMRmap_r', 'Dark2', 'Dark2_r', 'GnBu', 'GnBu_r', 'Grays', 'Greens', 'Greens_r', 'Greys', 'Greys_r', 'OrRd', 'OrRd_r', 'Oranges', 'Oranges_r', 'PRGn', 'PRGn_r',
'Paired', 'Paired_r', 'Pastel1', 'Pastel1_r', 'Pastel2', 'Pastel2_r', 'PiYG', 'PiYG_r', 'PuBu', 'PuBuGn', 'PuBuGn_r', 'PuBu_r', 'PuOr', 'PuOr_r', 'PuRd', 'PuRd_r', 'Purples',
'Purples_r', 'RdBu', 'RdBu_r', 'RdGy', 'RdGy_r', 'RdPu', 'RdPu_r', 'RdYlBu', 'RdYlBu_r', 'RdYlGn', 'RdYlGn_r', 'Reds', 'Reds_r', 'Set1', 'Set1_r', 'Set2', 'Set2_r',
'Set3', 'Set3_r', 'Spectral', 'Spectral_r', 'Wistia', 'Wistia_r', 'YlGn', 'YlGnBu', 'YlGnBu_r', 'YlGn_r', 'YlOrBr', 'YlOrBr_r', 'YlOrRd', 'YlOrRd_r', 'afmhot',
'afmhot_r', 'autumn', 'autumn_r', 'binary', 'binary_r', 'bone', 'bone_r', 'brg', 'brg_r', 'bwr', 'bwr_r', 'cividis', 'cividis_r', 'cool', 'cool_r', 'coolwarm',
'coolwarm_r', 'copper', 'copper_r', 'cubehelix', 'cubehelix_r', 'flag', 'flag_r', 'gist_earth', 'gist_earth_r', 'gist_gray', 'gist_gray_r', 'gist_grey', 'gist_heat',
'gist_heat_r', 'gist_ncar', 'gist_ncar_r', 'gist_rainbow', 'gist_rainbow_r', 'gist_stern', 'gist_stern_r', 'gist_yarg', 'gist_yarg_r', 'gist_yerg', 'gnuplot', 'gnuplot2',
'gnuplot2_r', 'gnuplot_r', 'gray', 'gray_r', 'grey', 'hot', 'hot_r', 'hsv', 'hsv_r', 'inferno', 'inferno_r', 'jet', 'jet_r', 'magma', 'magma_r', 'nipy_spectral',
'nipy_spectral_r', 'ocean', 'ocean_r', 'pink', 'pink_r', 'plasma', 'plasma_r', 'prism', 'prism_r', 'rainbow', 'rainbow_r', 'seismic', 'seismic_r', 'spring', 'spring_r',
'summer', 'summer_r', 'tab10', 'tab10_r', 'tab20', 'tab20_r', 'tab20b', 'tab20b_r', 'tab20c', 'tab20c_r', 'terrain', 'terrain_r', 'turbo', 'turbo_r', 'twilight', 'twilight_r',
'twilight_shifted', 'twilight_shifted_r', 'viridis', 'viridis_r', 'winter', 'winter_r'