如何借助Matplotlib高亮SymPy Plotting Backends生成的3D曲面图切片

我想用Matplotlib给SymPy Plotting Backends（SPB）生成的3D曲面图添加切片高亮效果，目前已经用SPB的MB后端画出了拉普拉斯变换的3D曲面，但尝试的切片高亮代码无法实现预期效果，预期效果是在3D曲面上突出显示某一平面的切片（类似红色切片的视觉效果）。

原代码如下：

from sympy import *
from spb import * # pip install sympy_plot_backends
import matplotlib.pyplot as plt
init_printing()

var("t, omega, s, alpha")
expr = exp(-t) * sin(t)
res = laplace_transform(expr, t, s, noconds=True)
print(res)

colorscale = [
    [0, "green"],
    [0.01, "orange"],
    [0.05, "red"],
    [1, "red"]
]
print('拉普拉斯变换3D曲面图：')
plot3d(abs(res.subs(s, alpha + I * omega)), (alpha, -2, 2), (omega, -2, 2), 
    use_cm=True, colorbar=False, backend=MB)
print('带切片高亮的图：')
#Z = abs(res.subs(s, alpha + I * omega))
#y0 = 10
#C[y0] = plt.cm.Reds_r(Z[y0]/2)
#plt.figure(figsize = (15,15))
#plt.plot((alpha, -2, 2), (omega, -2, 2),Z[y0,:], color=plt.cm.Reds(.7))
#plt.show()

问题分析 #

原注释代码的核心问题是直接操作符号表达式Z，没有将其转换为可计算的数值数组，同时也没有正确获取SPB绘图后的Matplotlib轴对象（Axes），无法在已有的3D曲面上叠加切片元素。

解决方案 #

要实现切片高亮，需要先将符号表达式转换为数值数据，再获取SPB绘图的Axes对象，最后在3D轴上添加高亮的切片图形：

from sympy import *
from spb import *
import matplotlib.pyplot as plt
import numpy as np

init_printing()

var("t, omega, s, alpha")
expr = exp(-t) * sin(t)
res = laplace_transform(expr, t, s, noconds=True)
print("拉普拉斯变换结果：", res)

# 定义要可视化的绝对值表达式
abs_expr = abs(res.subs(s, alpha + I * omega))

# 生成数值网格，用于计算切片数据
alpha_vals = np.linspace(-2, 2, 100)
omega_vals = np.linspace(-2, 2, 100)
alpha_mesh, omega_mesh = np.meshgrid(alpha_vals, omega_vals)

# 将符号表达式转换为可计算的数值函数
f = lambdify((alpha, omega), abs_expr, "numpy")
Z_mesh = f(alpha_mesh, omega_mesh)

# 绘制3D曲面，先不显示，获取Matplotlib的3D轴对象
print("带切片高亮的3D拉普拉斯变换曲面：")
p = plot3d(abs_expr, (alpha, -2, 2), (omega, -2, 2), 
           use_cm=True, colorbar=False, backend=MB, show=False)
ax = p.ax  # 获取3D绘图轴

# 选择要高亮的切片：固定omega=0，遍历所有alpha值
target_omega = 0
# 找到目标omega在数值数组中的最接近索引
slice_idx = np.argmin(np.abs(omega_vals - target_omega))
# 提取切片对应的alpha和Z值
slice_alpha = alpha_vals
slice_Z = Z_mesh[slice_idx, :]

# 绘制高亮的切片曲面（带填充）
slice_alpha_mesh, slice_omega_mesh = np.meshgrid(slice_alpha, [target_omega, target_omega])
slice_Z_mesh = np.tile(slice_Z, (2, 1))
ax.plot_surface(slice_alpha_mesh, slice_omega_mesh, slice_Z_mesh, 
                color="red", alpha=0.7, zorder=5)  # zorder确保切片在曲面上方

# 绘制切片的边界线，增强突出效果
ax.plot(slice_alpha, [target_omega]*len(slice_alpha), slice_Z, 
        color="darkred", linewidth=3, zorder=6)

# 显示最终图形
p.show()

代码说明 #

1. 数值化处理：用lambdify将符号表达式转换为NumPy可计算的函数，生成网格数据，方便提取切片的数值。
2. 获取绘图轴：通过plot3d的show=False参数延迟显示，获取Matplotlib的3D轴对象，用于后续添加切片元素。
3. 切片绘制：选择固定某一变量（如omega=0）的切片，生成对应的网格数据，用plot_surface绘制填充切片，plot绘制边界线，通过zorder确保切片显示在曲面上方。

内容的提问来源于stack exchange，提问作者Freya the Goddess