如何让3D曲面的colormap与对应下水平集的colormap匹配？

我想让左侧三维曲面的颜色和右侧对应的下水平集颜色匹配，要求右侧图中间小值显示为蓝色、大值显示为红色。但当前代码的输出不符合预期，怀疑contourf没有生成正确的下水平集（下水平集定义为${x\in\mathbb{R}: f(x)\le a}$）。另外想问这个平台支持LaTeX吗？

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d

# Define the function to plot
def f(x, y):
    return x**2 + y**2

# Generate data for the x, y, and z coordinates
x = np.linspace(-6, 6, 100)
y = np.linspace(-6, 6, 100)
X, Y = np.meshgrid(x, y)
Z = f(X, Y)

# Create a 3D figure and a contour plot side by side
fig = plt.figure(figsize=(10, 4))
ax1 = fig.add_subplot(121, projection='3d')
ax2 = fig.add_subplot(122)

# Plot the surface on the left subplot
surf = ax1.plot_surface(X, Y, Z, cmap='jet')

# Plot the contour on the right subplot
for a in range(1,30,5):
    contour_levels = np.arange(-a, a, 0.1)
    ax2.contourf(X, Y, Z, levels=contour_levels, cmap='jet')
    plt.pause(.001)

# Add labels and title to both subplots
ax1.set_xlabel('x')
ax1.set_ylabel('y')
ax1.set_zlabel('z')
ax1.set_title('3D Plot of f(x, y) = sin(sqrt(x^2 + y^2))')
ax2.set_xlabel('x')
ax2.set_ylabel('y')
ax2.set_title('Level Set Contour Plot of f(x, y) = x^2 + y^2')
surf.set_edgecolors(surf.to_rgba(surf._A))

# Show the plot
plt.show()

1. 下水平集生成错误 #

函数f(x,y)=x²+y²的取值范围是非负的，你代码里的contour_levels包含-a到0的无效区间，且每次循环叠加绘制导致颜色混乱。要生成f(x,y)≤a的下水平集，只需取0到a的区间，且每次循环前清除轴内容：

for a in range(1,30,5):
    ax2.clear()  # 清除上一次的绘制内容
    contour_levels = np.linspace(0, a, 50)  # 仅保留有效取值区间
    ax2.contourf(X, Y, Z, levels=contour_levels, cmap='jet', vmin=Z.min(), vmax=Z.max())
    # 重新设置子图元素，避免被clear清除
    ax2.set_xlabel('x')
    ax2.set_ylabel('y')
    ax2.set_title(f'Level Set f(x,y) ≤ {a}')
    plt.pause(0.001)

2. 颜色匹配问题 #

要让3D曲面和轮廓图颜色完全对齐，需统一两者的颜色映射范围（vmin和vmax）：

surf = ax1.plot_surface(X, Y, Z, cmap='jet', vmin=Z.min(), vmax=Z.max())

3. 标题错误修正 #

左侧3D图标题与实际函数不符，修正为：

ax1.set_title('3D Plot of f(x, y) = x^2 + y^2')

4. LaTeX支持说明 #

Stack Exchange平台支持LaTeX语法，只需用美元符号包裹公式即可，比如你写的${x\in\mathbb{R}: f(x)\le a}$可以正常渲染，确保使用标准LaTeX语法即可。

内容的提问来源于stack exchange，提问作者The Mad Scientist