如何用水平面在Matplotlib中平滑裁剪3D曲面图？

我想要创建带z范围和等高线的3D曲面图，使用matplotlib的plot_surface函数已有初步成果，但始终无法干净地裁剪曲面顶部，尝试了多种网上推荐的方法都存在问题，以下是我在PyCharm的Jupyter Notebook中测试的几种方法及遇到的问题：

方法一：基础plot_surface函数 #

绘制的曲面z=f(x,y)上升极快，z轴刻度数值很大，导致我关注的局部最小值区域难以看清。

import matplotlib
import matplotlib.pyplot as plt
import jax.numpy as jnp

def f(x, y, a, b):
    return x**4 + y**4 + y**3 - (4 * x**2 * y) + y**2 - (a * x) + (b * y)

X, Y = jnp.meshgrid(jnp.arange(-3, 3, 0.01), jnp.arange(-3, 3, 0.01))
Z = f(X, Y, 0, 0)  # 为参数a和b的单一选择生成每个(x,y)处的z值

matplotlib.use("Qt5Agg")

fig, ax = plt.subplots(subplot_kw={"projection": "3d"}, figsize=(16, 12))

surf = ax.plot_surface(X, Y, Z, linewidth=0, antialiased=False, cmap="viridis")
ax.view_init(elev=10, azim=15, roll=0)

plt.tight_layout()
plt.show()

输出1：曲面整体完整显示，但z轴跨度极大，局部最小值区域细节无法分辨。

方法二：设置z范围并启用axlim_clip=True #

此方法放大了目标区域，但颜色缩放仍沿用全局z值范围，导致关注区域内颜色几乎一致，且顶部边缘因整块曲面片被移除而呈现锯齿状。

import matplotlib
import matplotlib.pyplot as plt
import jax.numpy as jnp

def f(x, y, a, b):
    return x**4 + y**4 + y**3 - (4 * x**2 * y) + y**2 - (a * x) + (b * y)

X, Y = jnp.meshgrid(jnp.arange(-3, 3, 0.01), jnp.arange(-3, 3, 0.01))
Z = f(X, Y, 0, 0)  # 为参数a和b的单一选择生成每个(x,y)处的z值
Zmin, Zmax = jnp.min(Z), jnp.percentile(Z, 15.)

matplotlib.use("Qt5Agg")

fig, ax = plt.subplots(subplot_kw={"projection": "3d"}, figsize=(16, 12))

surf = ax.plot_surface(X, Y, Z, linewidth=0, antialiased=False, cmap="viridis", axlim_clip=True)
ax.view_init(elev=10, azim=15, roll=0)
ax.set_zlim(Zmin, Zmax)

plt.tight_layout()
plt.show()

输出2：目标区域被放大，但颜色区分度极低，曲面顶部边缘呈锯齿状。

方法三：用NaN遮罩超出阈值的z值 #

此方法解决了颜色缩放问题，关注区域内颜色渐变正常，但顶部边缘仍呈锯齿状，无法实现干净的裁剪效果。

import matplotlib
import matplotlib.pyplot as plt
import jax.numpy as jnp

def f(x, y, a, b):
    return x**4 + y**4 + y**3 - (4 * x**2 * y) + y**2 - (a * x) + (b * y)

X, Y = jnp.meshgrid(jnp.arange(-3, 3, 0.01), jnp.arange(-3, 3, 0.01))
Z = f(X, Y, 0, 0)  # 为参数a和b的单一选择生成每个(x,y)处的z值

Zmax = jnp.percentile(Z, 15.)
Zplateau = jnp.minimum(Z, Zmax)
Zmasked = jnp.where(Z > Zmax, jnp.nan, Zplateau)

matplotlib.use("Qt5Agg")

fig, ax = plt.subplots(subplot_kw={"projection": "3d"}, figsize=(16, 12))

surf = ax.plot_surface(X, Y, Zmasked, linewidth=0, antialiased=False, cmap="viridis")
ax.view_init(elev=10, azim=15, roll=0)
ax.set_zlim(jnp.min(Z), Zmax)

plt.tight_layout()
plt.show()

输出3：颜色显示正常，但曲面顶部裁剪边缘锯齿明显。

方法四：将超出阈值的z值设为阈值 #

此方法让顶部边缘平滑，但在曲面顶部形成了不必要的平台，不符合需求。

import matplotlib
import matplotlib.pyplot as plt
import jax.numpy as jnp

def f(x, y, a, b):
    return x ** 4 + y ** 4 + y ** 3 - (4 * x ** 2 * y) + y ** 2 - (a * x) + (b * y)

X, Y = jnp.meshgrid(jnp.arange(-3, 3, 0.01), jnp.arange(-3, 3, 0.01))
Z = f(X, Y, 0, 0)  # 为参数a和b的单一选择生成每个(x,y)处的z值

Zmax = jnp.percentile(Z, 15.)
Zclipped = jnp.where(Z > Zmax, Zmax, Z)

matplotlib.use("Qt5Agg")

fig, ax = plt.subplots(subplot_kw={"projection": "3d"}, figsize=(16, 12))

surf = ax.plot_surface(X, Y, Zclipped, linewidth=0, antialiased=False, cmap="viridis")
contours = ax.contour(X, Y, Z, levels=[Zmax], zdir="z", offset=Zmax, colors="k", linewidths=2)
ax.view_init(elev=10, azim=15, roll=0)
ax.set_zlim(jnp.min(Z), Zmax)

plt.tight_layout()
plt.show()

输出4：顶部边缘平滑，但出现了多余的平台区域。

补充说明 #

我尝试过使用更精细的xy网格，以及调整plot_surface的rstride和cstride参数，这些仅能让曲面下部更平滑，但顶部边缘仍呈锯齿状。

内容的提问来源于stack exchange，提问作者Chance Bowman