Python轮子：matpotlib-3d绘图(2)

原文链接：http://www.juzicode.com/python-module-matplotlib-3d-plotting-2

Matplotlib的三维绘图能力远超基础曲面图，本文将展示几种高级三维可视化技术，适用于量子计算、流体力学、分子生物学等专业领域的数据呈现。

应用场景

• 量子态概率分布可视化
• 计算流体动力学(CFD)模拟
• 蛋白质分子结构渲染
• 三维地理信息系统(GIS)数据
• 电磁场矢量分析

安装与导入

pip install matplotlib numpy
import matplotlib.pyplot as plt

三维绘图方法

1）3D量子概率云图

可视化量子力学中的电子概率分布。

# juzicode.com/VX公众号:juzicode
import matplotlib.pyplot as plt
import numpy as np
plt.rcParams['font.family'] =  ['YouYuan']   
plt.rcParams['axes.unicode_minus'] = False  
from skimage import measure  # 需要额外安装pip install scikit-image


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

# 创建氢原子2p轨道数据
x = np.linspace(-5, 5, 100)
y = np.linspace(-5, 5, 100)
z = np.linspace(-5, 5, 100)
X, Y, Z = np.meshgrid(x, y, z)

# 量子波函数计算
r = np.sqrt(X**2 + Y**2 + Z**2)
psi = X * np.exp(-r/2)  # 2p_x轨道

# 概率密度计算
prob = np.abs(psi)**2
prob_normalized = prob / np.max(prob)

# 创建等值面图
verts, faces, _, _ = measure.marching_cubes(prob_normalized, level=0.2, 
                                          spacing=(0.1, 0.1, 0.1))

# 渲染3D网格
ax.plot_trisurf(verts[:, 0], verts[:,1], faces, verts[:, 2],
               cmap='Spectral', alpha=0.8, 
               antialiased=True)

# 设置光照和视角
ax.view_init(elev=15, azim=45)
ax.set_title('氢原子2p_x轨道电子云', fontsize=16)
ax.set_xlabel('X (a.u.)')
ax.set_ylabel('Y (a.u.)')
ax.set_zlabel('Z (a.u.)')
ax.grid(False)

plt.savefig('3d_quantum_cloud.png', dpi=300, transparent=True)

2）盒体表面等高图

从每个曲面中提取数据，并单独绘制等高线图。

# juzicode.com/VX公众号:juzicode
import numpy as np
import matplotlib.pyplot as plt 
plt.rcParams['font.family'] =  ['YouYuan']   
plt.rcParams['axes.unicode_minus'] = False  

Nx, Ny, Nz = 100, 300, 500
X, Y, Z = np.meshgrid(np.arange(Nx), np.arange(Ny), -np.arange(Nz))
data = (((X+100)**2 + (Y-20)**2 + 2*Z)/1000+1)

kw = {
    'vmin': data.min(),
    'vmax': data.max(),
    'levels': np.linspace(data.min(), data.max(), 10),
}

fig = plt.figure(figsize=(5, 4))
ax = fig.add_subplot(111, projection='3d')
ax.set_title('盒体表面等高图', fontsize=8)

_ = ax.contourf(
    X[:, :, 0], Y[:, :, 0], data[:, :, 0],
    zdir='z', offset=0, **kw
)
_ = ax.contourf(
    X[0, :, :], data[0, :, :], Z[0, :, :],
    zdir='y', offset=0, **kw
)
C = ax.contourf(
    data[:, -1, :], Y[:, -1, :], Z[:, -1, :],
    zdir='x', offset=X.max(), **kw
)
# --


# 设置范围
xmin, xmax = X.min(), X.max()
ymin, ymax = Y.min(), Y.max()
zmin, zmax = Z.min(), Z.max()
ax.set(xlim=[xmin, xmax], ylim=[ymin, ymax], zlim=[zmin, zmax])

# 画边沿
edges_kw = dict(color='0.4', linewidth=1, zorder=1e3)
ax.plot([xmax, xmax], [ymin, ymax], 0, **edges_kw)
ax.plot([xmin, xmax], [ymin, ymin], 0, **edges_kw)
ax.plot([xmax, xmax], [ymin, ymin], [zmin, zmax], **edges_kw)

# 设置坐标值
ax.set(
    xlabel='X [km]',
    ylabel='Y [km]',
    zlabel='Z [m]',
    zticks=[0, -150, -300, -450],
)

ax.view_init(40, -30, 0)
ax.set_box_aspect(None, zoom=0.9)

fig.colorbar(C, ax=ax, fraction=0.02, pad=0.1, label='Name [units]')

plt.savefig('3d_box_surface.png', dpi=300, transparent=True)

3）底部着色3D线图

创建填充线图下方区域的多边形，多边形呈半透明状，形成一种”彩绘玻璃”的视觉效果。

# juzicode.com/VX公众号:juzicode
import matplotlib.pyplot as plt
import numpy as np
import math
plt.rcParams['font.family'] =  ['YouYuan']   
plt.rcParams['axes.unicode_minus'] = False  


gamma = np.vectorize(math.gamma)
N = 31
x = np.linspace(0., 10., N)
lambdas = range(1, 9)

ax = plt.figure().add_subplot(projection='3d')
ax.set_title('底部着色3D线图', fontsize=8)
facecolors = plt.colormaps['viridis_r'](np.linspace(0, 1, len(lambdas)))

for i, l in enumerate(lambdas):
    ax.fill_between(x, l, l**x * np.exp(-l) / gamma(x + 1),
                    x, l, 0,
                    facecolors=facecolors[i], alpha=.7)

ax.set(xlim=(0, 10), ylim=(1, 9), zlim=(0, 0.35),
       xlabel='x', ylabel=r'$\lambda$', zlabel='probability')
       
plt.savefig('3d_line_graph.png', dpi=300)

4）3D地形图

地理数据的三维真实感渲染。

# juzicode.com/VX公众号:juzicode
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LightSource  # 导入LightSource
from matplotlib import cm  # 导入cm模块

plt.rcParams['font.family'] = ['YouYuan']
plt.rcParams['axes.unicode_minus'] = False

x = np.linspace(-3, 3, 200)
y = np.linspace(-3, 3, 200)
X, Y = np.meshgrid(x, y)
Z = np.sin(3*X) * np.cos(2*Y) + 0.3*X

# 创建图形
fig = plt.figure(figsize=(16, 12))
ax = fig.add_subplot(111, projection='3d')

# 创建光源效果 - 修正后
ls = LightSource(azdeg=225, altdeg=45)
rgb = ls.shade(Z, cmap=cm.terrain, 
              vert_exag=0.2, 
              blend_mode='soft')

# 绘制地形表面
surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1,
                      facecolors=rgb,
                      linewidth=0, 
                      antialiased=False)

# 添加等高线
contour = ax.contourf(X, Y, Z, zdir='z', offset=-1.5, 
                     cmap='terrain', alpha=0.4)

# 设置视角和标签
ax.view_init(elev=35, azim=-60)
ax.set_zlim(-1.5, 1.5)
ax.set_title('三维地形高程图', fontsize=18)
ax.set_xlabel('经度')
ax.set_ylabel('纬度')
ax.set_zlabel('海拔 (km)')

# 添加颜色条
cbar = fig.colorbar(surf, shrink=0.5)
cbar.set_label('海拔高度', rotation=270, labelpad=20)

plt.savefig('3d_terrain.png', dpi=300)

5）3D分子结构图

生物大分子的三维空间结构展示。

# juzicode.com/VX公众号:juzicode
import matplotlib.pyplot as plt
import numpy as np
plt.rcParams['font.family'] =  ['YouYuan']   
plt.rcParams['axes.unicode_minus'] = False  

# 定义水分子坐标 (H2O)
atoms = {
    'O': [(0, 0, 0)],
    'H': [(0.76, 0.59, 0), (-0.76, 0.59, 0)]
}

# 原子半径和颜色
radii = {'O': 0.6, 'H': 0.3}
colors = {'O': '#FF3030', 'H': '#1E90FF'}

# 创建图形
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(111, projection='3d')

# 绘制原子球体
for element, positions in atoms.items():
    for pos in positions:
        u = np.linspace(0, 2 * np.pi, 30)
        v = np.linspace(0, np.pi, 30)
        x = radii[element] * np.outer(np.cos(u), np.sin(v)) + pos[0]
        y = radii[element] * np.outer(np.sin(u), np.sin(v)) + pos[1]
        z = radii[element] * np.outer(np.ones(np.size(u)), np.cos(v)) + pos[2]
        ax.plot_surface(x, y, z, color=colors[element], alpha=0.9)

# 绘制化学键
bonds = [
    [atoms['O'][0], atoms['H'][0]],
    [atoms['O'][0], atoms['H'][1]]
]

for bond in bonds:
    start, end = bond
    xs = [start[0], end[0]]
    ys = [start[1], end[1]]
    zs = [start[2], end[2]]
    ax.plot(xs, ys, zs, 'k-', lw=2, alpha=0.7)

# 设置显示属性
ax.set_axis_off()
ax.set_title('水分子三维结构', fontsize=16)
ax.view_init(elev=20, azim=30)

# 添加图例
from matplotlib.patches import Patch
legend_elements = [
    Patch(facecolor=colors['O'], label='氧原子'),
    Patch(facecolor=colors['H'], label='氢原子')
]
ax.legend(handles=legend_elements, loc='upper right')

plt.savefig('3d_molecule.png', dpi=300, transparent=True)

专业三维绘图技巧

• 物理模拟：使用NumPy向量运算加速场方程计算
• 真实感渲染：应用LightSource创建光照效果，调整vert_exag增强地形起伏
• 视角控制：记录关键视角坐标实现论文图表一致性

通过组合使用这些高级三维可视化技术，可以在物理、化学、地理和生物医学等领域创建具有发表质量的科学图表。实际应用中需根据数据特征和展示目的选择最合适的可视化形式。